Add learnable last-four-layer query fusion

easycom_head.py

@@ -46,12 +46,32 @@ class TurnPredictionHead(nn.Module):
         num_horizons: int = NUM_HORIZONS,
         num_classes: int = NUM_CLASSES,
         dropout: float = 0.1,
+        num_input_layers: int = 1,
+        learnable_layer_mixture: bool = False,
     ) -> None:
         super().__init__()
+        if num_input_layers <= 0:
+            raise ValueError("num_input_layers must be positive")
+        if learnable_layer_mixture and num_input_layers < 2:
+            raise ValueError(
+                "learnable_layer_mixture requires at least two input layers"
+            )
+        if not learnable_layer_mixture and num_input_layers != 1:
+            raise ValueError(
+                "Multiple input layers require learnable_layer_mixture=True"
+            )
         self.hidden_size = hidden_size
         self.num_horizons = num_horizons
         self.num_classes = num_classes
         self.dropout = dropout
+        self.num_input_layers = num_input_layers
+        self.learnable_layer_mixture = learnable_layer_mixture
+        if learnable_layer_mixture:
+            # Zero initialization exactly reproduces a uniform layer mean at
+            # step zero while allowing training to select useful query layers.
+            self.layer_logits = nn.Parameter(torch.zeros(num_input_layers))
+        else:
+            self.register_parameter("layer_logits", None)
         self.net = nn.Sequential(
             nn.LayerNorm(hidden_size),
             nn.Linear(hidden_size, hidden_size // 2),
@@ -61,11 +81,36 @@ class TurnPredictionHead(nn.Module):
         )
 
     def forward(self, states: torch.Tensor) -> torch.Tensor:
+        if states.shape[-1] != self.hidden_size:
+            raise ValueError(
+                f"Expected hidden size {self.hidden_size}, got {states.shape[-1]}"
+            )
+        if self.learnable_layer_mixture:
+            if states.ndim < 3 or states.shape[-2] != self.num_input_layers:
+                raise ValueError(
+                    "Layer-mixture input must end in "
+                    f"[{self.num_input_layers},{self.hidden_size}], got "
+                    f"{tuple(states.shape)}"
+                )
+            assert self.layer_logits is not None
+            weights = self.layer_logits.softmax(dim=0).to(states.dtype)
+            view_shape = (1,) * (states.ndim - 2) + (
+                self.num_input_layers,
+                1,
+            )
+            states = (states * weights.view(view_shape)).sum(dim=-2)
         logits = self.net(states)
         return logits.view(
             *states.shape[:-1], self.num_horizons, self.num_classes
         )
 
+    def layer_mixture_weights(self) -> torch.Tensor | None:
+        """Return normalized oldest-to-newest layer weights, if enabled."""
+
+        if self.layer_logits is None:
+            return None
+        return self.layer_logits.softmax(dim=0)
+
 
 def _validate_loss_inputs(logits: torch.Tensor, targets: torch.Tensor) -> None:
     if logits.shape != targets.shape:
@@ -219,6 +264,10 @@ class SessionData:
     def hidden_size(self) -> int:
         return int(self.states.shape[-1])
 
+    @property
+    def num_input_layers(self) -> int:
+        return 1 if self.states.ndim == 2 else int(self.states.shape[-2])
+
     def future_examples(
         self,
         num_horizons: int = NUM_HORIZONS,
@@ -256,8 +305,12 @@ def load_session_data(
     if not isinstance(payload, dict) or "states" not in payload:
         raise ValueError(f"Invalid feature payload: {feature_path}")
     states = payload["states"]
-    if not isinstance(states, torch.Tensor) or states.ndim != 2:
-        raise ValueError(f"states must be [T,D] in {feature_path}")
+    if not isinstance(states, torch.Tensor) or states.ndim not in {2, 3}:
+        raise ValueError(f"states must be [T,D] or [T,L,D] in {feature_path}")
+    if states.ndim == 3 and states.shape[1] < 2:
+        raise ValueError(
+            f"Stacked states must contain at least two layers: {feature_path}"
+        )
     if not torch.is_floating_point(states):
         raise ValueError(f"states must be floating point in {feature_path}")
     if not torch.isfinite(states).all():
@@ -351,6 +404,8 @@ def checkpoint_head_config(payload: dict[str, Any]) -> dict[str, Any]:
     config.setdefault("num_horizons", NUM_HORIZONS)
     config.setdefault("num_classes", NUM_CLASSES)
     config.setdefault("dropout", 0.1)
+    config.setdefault("num_input_layers", 1)
+    config.setdefault("learnable_layer_mixture", False)
     return config
 
 
@@ -360,4 +415,4 @@ def validate_disjoint_sessions(
 ) -> None:
     overlap = set(first) & set(second)
     if overlap:
-        raise ValueError(f"Session sets overlap: {sorted(overlap)}")
+        raise ValueError(f"Session sets overlap: {sorted(overlap)}")

infer_turn_head.py

@@ -283,6 +283,11 @@ def main() -> None:
                 f"{name} feature dim is {data.hidden_size}, checkpoint expects "
                 f"{head_config['hidden_size']}"
             )
+        if data.num_input_layers != int(head_config["num_input_layers"]):
+            raise ValueError(
+                f"{name} has {data.num_input_layers} feature layers, checkpoint "
+                f"expects {head_config['num_input_layers']}"
+            )
         probabilities = _predict_states(
             head,
             data.states,
@@ -323,14 +328,29 @@ def main() -> None:
             "checkpoint_epoch": checkpoint_payload.get("epoch"),
             "checkpoint_train_sessions": checkpoint_payload.get("train_sessions"),
             "checkpoint_dev_sessions": checkpoint_payload.get("dev_sessions"),
+            "layer_mixture_weights_oldest_to_newest": (
+                head.layer_mixture_weights().detach().float().cpu().tolist()
+                if head.layer_mixture_weights() is not None
+                else None
+            ),
         }
     )
     metrics_path = args.output_dir / "metrics.json"
     metrics_path.write_text(json.dumps(metrics, indent=2) + "\n")
     _print_metrics(metrics)
+    if metrics["layer_mixture_weights_oldest_to_newest"] is not None:
+        print(
+            "\nLayer mixture weights (oldest->newest): "
+            + ", ".join(
+                f"{value:.6f}"
+                for value in metrics[
+                    "layer_mixture_weights_oldest_to_newest"
+                ]
+            )
+        )
     print(f"\nmetrics: {metrics_path}")
     print(f"EgoSpeak-compatible predictions: {prediction_dir}")
 
 
 if __name__ == "__main__":
-    main()
+    main()

run_extract_easycom_direct.sh

@@ -17,6 +17,7 @@ MEMORY_SECONDS="${MEMORY_SECONDS:-20}"
 NATIVE_BLOCK_SECONDS="${NATIVE_BLOCK_SECONDS:-2}"
 VIDEO_FPS="${VIDEO_FPS:-2}"
 READOUT_LAST_N_LAYERS="${READOUT_LAST_N_LAYERS:-1}"
+READOUT_LAYER_FUSION="${READOUT_LAYER_FUSION:-mean}"
 MODE="${MODE:-open-stream-replay}"
 DTYPE="${DTYPE:-bfloat16}"
 DECODER="${DECODER:-decord}"
@@ -74,6 +75,7 @@ run_one() {
     echo "output=${output}"
     echo "memory_seconds=${MEMORY_SECONDS} video_fps=${VIDEO_FPS} mode=${MODE}"
     echo "readout_last_n_layers=${READOUT_LAST_N_LAYERS}"
+    echo "readout_layer_fusion=${READOUT_LAYER_FUSION}"
     echo "decoder=${DECODER} decoder_gpu_id=${DECODER_GPU_ID} frame_cache_size=${FRAME_CACHE_SIZE}"
 
     if ! CUDA_VISIBLE_DEVICES="${gpu}" PYTHONUNBUFFERED=1 \
@@ -87,6 +89,7 @@ run_one() {
         --native-block-seconds "${NATIVE_BLOCK_SECONDS}" \
         --video-fps "${VIDEO_FPS}" \
         --readout-last-n-layers "${READOUT_LAST_N_LAYERS}" \
+        --readout-layer-fusion "${READOUT_LAYER_FUSION}" \
         --frame-cache-size "${FRAME_CACHE_SIZE}" \
         --decoder "${DECODER}" \
         --decoder-gpu-id "${DECODER_GPU_ID}" \

streaming_turn_model.py

@@ -43,6 +43,7 @@ from typing import Any, Literal
 import torch
 from easycom_head import (
     TurnPredictionHead,
+    checkpoint_head_config,
     egospeak_multihorizon_loss,
     miniroad_multihorizon_loss,
     multihorizon_cross_entropy,
@@ -141,6 +142,7 @@ class StreamConfig:
     video_min_pixels: int = 64 * 28 * 28
     video_max_pixels: int = 128 * 28 * 28
     readout_last_n_layers: int = 1
+    readout_layer_fusion: str = "mean"
 
     def validate(self) -> None:
         if self.step_seconds <= 0:
@@ -157,6 +159,15 @@ class StreamConfig:
             raise ValueError("video_max_pixels must be >= video_min_pixels")
         if self.readout_last_n_layers <= 0:
             raise ValueError("readout_last_n_layers must be positive")
+        if self.readout_layer_fusion not in {"mean", "stack"}:
+            raise ValueError("readout_layer_fusion must be 'mean' or 'stack'")
+        if (
+            self.readout_layer_fusion == "stack"
+            and self.readout_last_n_layers < 2
+        ):
+            raise ValueError(
+                "stack readout requires readout_last_n_layers >= 2"
+            )
         _require_integer_ratio(
             self.native_block_seconds, self.step_seconds, "native block / step"
         )
@@ -170,7 +181,7 @@ class ExtractionResult:
     """One fused Thinker state for every causal prediction time."""
 
     times: torch.Tensor  # [N], seconds in the original session
-    states: torch.Tensor  # [N, D]
+    states: torch.Tensor  # [N,D] or preserved query layers [N,L,D]
     verification: list[dict[str, Any]]
 
 
@@ -956,6 +967,7 @@ def _read_state(
     next_rope_offset: int,
     query_ids: torch.LongTensor,
     last_n_layers: int = 1,
+    layer_fusion: str = "mean",
 ) -> torch.Tensor:
     """Read one q_t without committing the artificial assistant query."""
 
@@ -989,25 +1001,38 @@ def _read_state(
         return_dict=True,
         cache_position=cache_position,
     )
-    return _readout_hidden_state(outputs, last_n_layers=last_n_layers)
+    return _readout_hidden_state(
+        outputs,
+        last_n_layers=last_n_layers,
+        layer_fusion=layer_fusion,
+    )
 
 
-def _readout_hidden_state(outputs: Any, *, last_n_layers: int) -> torch.Tensor:
-    """Return the query state, optionally averaging the final decoder layers."""
+def _readout_hidden_state(
+    outputs: Any,
+    *,
+    last_n_layers: int,
+    layer_fusion: str,
+) -> torch.Tensor:
+    """Return one query vector or an ordered stack of final-layer vectors."""
 
+    if layer_fusion not in {"mean", "stack"}:
+        raise ValueError(f"Unknown readout layer fusion: {layer_fusion}")
     if last_n_layers == 1:
-        return outputs.last_hidden_state[:, -1, :]
-    hidden_states = outputs.hidden_states
-    if hidden_states is None or len(hidden_states) < last_n_layers:
-        available = 0 if hidden_states is None else len(hidden_states)
-        raise RuntimeError(
-            f"Requested the last {last_n_layers} hidden states, but the model "
-            f"returned {available}"
+        stacked = outputs.last_hidden_state[:, -1:, :]
+    else:
+        hidden_states = outputs.hidden_states
+        if hidden_states is None or len(hidden_states) < last_n_layers:
+            available = 0 if hidden_states is None else len(hidden_states)
+            raise RuntimeError(
+                f"Requested the last {last_n_layers} hidden states, but the model "
+                f"returned {available}"
+            )
+        stacked = torch.stack(
+            [state[:, -1, :] for state in hidden_states[-last_n_layers:]],
+            dim=1,
         )
-    return torch.stack(
-        [state[:, -1, :] for state in hidden_states[-last_n_layers:]],
-        dim=0,
-    ).mean(dim=0)
+    return stacked.mean(dim=1) if layer_fusion == "mean" else stacked
 
 
 @torch.inference_mode()
@@ -1019,6 +1044,7 @@ def _read_open_state(
     closure_position_ids: torch.LongTensor,
     query_ids: torch.LongTensor,
     last_n_layers: int = 1,
+    layer_fusion: str = "mean",
 ) -> torch.Tensor:
     """Temporarily close an open media turn, read q_t, and discard the branch."""
 
@@ -1078,7 +1104,11 @@ def _read_open_state(
         return_dict=True,
         cache_position=cache_position,
     )
-    return _readout_hidden_state(outputs, last_n_layers=last_n_layers)
+    return _readout_hidden_state(
+        outputs,
+        last_n_layers=last_n_layers,
+        layer_fusion=layer_fusion,
+    )
 
 
 @torch.inference_mode()
@@ -1195,6 +1225,7 @@ class CausalPrefixExtractor:
             next_rope_offset=next_offset,
             query_ids=self.query_ids,
             last_n_layers=self.config.readout_last_n_layers,
+            layer_fusion=self.config.readout_layer_fusion,
         )
 
     def extract(
@@ -1358,6 +1389,7 @@ class OpenStreamReplayExtractor:
             closure_position_ids=closure_position_ids,
             query_ids=self.query_ids,
             last_n_layers=self.config.readout_last_n_layers,
+            layer_fusion=self.config.readout_layer_fusion,
         )
 
     def _cache_from_intervals(
@@ -1567,6 +1599,7 @@ class OpenStreamReplayExtractor:
                 closure_position_ids=closure_position_ids,
                 query_ids=self.query_ids,
                 last_n_layers=self.config.readout_last_n_layers,
+                layer_fusion=self.config.readout_layer_fusion,
             )
             states.append(state.cpu())
 
@@ -1808,6 +1841,7 @@ class OverlappingWorkingMemoryExtractor:
             closure_position_ids=closure_position_ids,
             query_ids=self.query_ids,
             last_n_layers=self.config.readout_last_n_layers,
+            layer_fusion=self.config.readout_layer_fusion,
         )
 
     def extract(
@@ -2002,12 +2036,16 @@ def _compare_states(
 ) -> dict[str, Any]:
     cached_f = cached.float()
     exact_f = exact.float()
+    cached_flat = cached_f.reshape(1, -1)
+    exact_flat = exact_f.reshape(1, -1)
     return {
         "kind": kind,
         "time": float(current_time),
         "max_abs": float((cached_f - exact_f).abs().max().item()),
         "mean_abs": float((cached_f - exact_f).abs().mean().item()),
-        "cosine": float(nn.functional.cosine_similarity(cached_f, exact_f).item()),
+        "cosine": float(
+            nn.functional.cosine_similarity(cached_flat, exact_flat, dim=-1).item()
+        ),
     }
 
 
@@ -2104,9 +2142,18 @@ def _load_head(
     hidden_size: int,
     device: torch.device,
 ) -> TurnPredictionHead:
-    head = TurnPredictionHead(hidden_size=hidden_size)
     payload = torch.load(checkpoint, map_location="cpu", weights_only=True)
-    state_dict = payload.get("head", payload) if isinstance(payload, dict) else payload
+    if isinstance(payload, dict) and "head" in payload:
+        config = checkpoint_head_config(payload)
+        if int(config["hidden_size"]) != hidden_size:
+            raise ValueError(
+                f"Head hidden size {config['hidden_size']} != model {hidden_size}"
+            )
+        head = TurnPredictionHead(**config)
+        state_dict = payload["head"]
+    else:
+        head = TurnPredictionHead(hidden_size=hidden_size)
+        state_dict = payload
     head.load_state_dict(state_dict)
     return head.to(device).eval()
 
@@ -2160,6 +2207,15 @@ def parse_args() -> argparse.Namespace:
             "causal readout token (default: 1 preserves the original result)."
         ),
     )
+    parser.add_argument(
+        "--readout-layer-fusion",
+        choices=("mean", "stack"),
+        default="mean",
+        help=(
+            "Fuse selected query layers by a fixed mean, or preserve [N,D] "
+            "for a trainable downstream layer mixture."
+        ),
+    )
     parser.add_argument("--verify-every", type=int, default=0)
     parser.add_argument(
         "--verify-continuous-every",
@@ -2228,6 +2284,7 @@ def main() -> None:
         video_min_pixels=args.video_min_pixels,
         video_max_pixels=args.video_max_pixels,
         readout_last_n_layers=args.readout_last_n_layers,
+        readout_layer_fusion=args.readout_layer_fusion,
     )
     config.validate()
 

streaming_turn_model_directly.py

@@ -1061,6 +1061,15 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
             "causal readout token (default: 1 preserves the original result)."
         ),
     )
+    parser.add_argument(
+        "--readout-layer-fusion",
+        choices=("mean", "stack"),
+        default="mean",
+        help=(
+            "Fuse selected query layers by a fixed mean, or preserve [N,D] "
+            "for a trainable downstream layer mixture."
+        ),
+    )
     parser.add_argument("--frame-cache-size", type=int, default=8)
     parser.add_argument(
         "--audio-feature-device",
@@ -1341,6 +1350,7 @@ def main(
         video_min_pixels=args.video_min_pixels,
         video_max_pixels=args.video_max_pixels,
         readout_last_n_layers=args.readout_last_n_layers,
+        readout_layer_fusion=args.readout_layer_fusion,
     )
     config.validate()
 

train_omni_turn_lora.py

@@ -726,12 +726,7 @@ def _load_initial_head(
         raise ValueError(
             f"Head hidden size {config['hidden_size']} != cache hidden size {hidden_size}"
         )
-    head = TurnPredictionHead(
-        hidden_size=hidden_size,
-        num_horizons=int(config["num_horizons"]),
-        num_classes=int(config["num_classes"]),
-        dropout=float(config["dropout"]),
-    )
+    head = TurnPredictionHead(**config)
     head.load_state_dict(payload.get("head", payload))
     if verbose:
         print(f"Initialized classification head from {checkpoint}")

train_turn_head.py

@@ -91,6 +91,14 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--warmup-epochs", type=float, default=10.0)
     parser.add_argument("--warmup-factor", type=float, default=0.3)
     parser.add_argument("--dropout", type=float, default=0.2)
+    parser.add_argument(
+        "--learnable-layer-mixture",
+        action="store_true",
+        help=(
+            "Learn softmax weights over stacked [T,L,D] query features before "
+            "the anticipation head."
+        ),
+    )
     parser.add_argument(
         "--loss",
         choices=("lstr-mce", "miniroad-normalized"),
@@ -151,11 +159,12 @@ def _load_examples(
     expected_mode: str | None,
     dataset: str,
     selected_indices: dict[str, set[int]] | None = None,
-) -> tuple[torch.Tensor, torch.Tensor, int, list[dict[str, Any]]]:
+) -> tuple[torch.Tensor, torch.Tensor, int, int, list[dict[str, Any]]]:
     all_states: list[torch.Tensor] = []
     all_targets: list[torch.Tensor] = []
     metadata: list[dict[str, Any]] = []
     hidden_size: int | None = None
+    num_input_layers: int | None = None
     for name in sessions:
         data = load_session_data(
             feature_root,
@@ -186,6 +195,13 @@ def _load_examples(
                 f"Feature dimensions differ: {name} has {data.hidden_size}, "
                 f"expected {hidden_size}"
             )
+        if num_input_layers is None:
+            num_input_layers = data.num_input_layers
+        elif data.num_input_layers != num_input_layers:
+            raise ValueError(
+                f"Feature layer counts differ: {name} has "
+                f"{data.num_input_layers}, expected {num_input_layers}"
+            )
         all_states.append(states)
         all_targets.append(targets)
         metadata.append(
@@ -207,12 +223,13 @@ def _load_examples(
             f"valid_10_horizon_examples={dense_example_count}, "
             f"selected_examples={len(states)}"
         )
-    if hidden_size is None:
+    if hidden_size is None or num_input_layers is None:
         raise ValueError("No sessions selected")
     return (
         torch.cat(all_states, dim=0),
         torch.cat(all_targets, dim=0),
         hidden_size,
+        num_input_layers,
         metadata,
     )
 
@@ -318,6 +335,7 @@ def _checkpoint(
     train_sessions: list[str],
     dev_sessions: list[str],
     hidden_size: int,
+    num_input_layers: int,
     feature_metadata: list[dict[str, Any]],
     history: list[dict[str, Any]],
 ) -> dict[str, Any]:
@@ -329,6 +347,8 @@ def _checkpoint(
             "num_horizons": NUM_HORIZONS,
             "num_classes": NUM_CLASSES,
             "dropout": args.dropout,
+            "num_input_layers": num_input_layers,
+            "learnable_layer_mixture": args.learnable_layer_mixture,
         },
         "alignment": {
             "state_label_offset": 0,
@@ -363,6 +383,7 @@ def _checkpoint(
             "scheduler": "linear_warmup_times_cosine",
             "seed": args.seed,
             "expected_mode": args.expected_mode,
+            "learnable_layer_mixture": args.learnable_layer_mixture,
             "supervision_policy": (
                 "transition-sparse-manifest"
                 if args.sparse_anchor_cache_dir
@@ -440,7 +461,7 @@ def main() -> None:
     print(f"device: {device}")
     print(f"train sessions: {', '.join(train_sessions)}")
     print(f"dev sessions: {', '.join(dev_sessions) if dev_sessions else 'none'}")
-    states, targets, hidden_size, feature_metadata = _load_examples(
+    states, targets, hidden_size, num_input_layers, feature_metadata = _load_examples(
         args.feature_root,
         args.target_root,
         train_sessions,
@@ -448,11 +469,29 @@ def main() -> None:
         args.dataset,
         sparse_anchor_indices,
     )
-    print(f"training examples: {len(states)}, feature_dim={hidden_size}")
+    print(
+        f"training examples: {len(states)}, feature_dim={hidden_size}, "
+        f"input_layers={num_input_layers}"
+    )
+    if num_input_layers > 1 and not args.learnable_layer_mixture:
+        raise ValueError(
+            f"Features contain {num_input_layers} stacked layers; pass "
+            "--learnable-layer-mixture"
+        )
+    if num_input_layers == 1 and args.learnable_layer_mixture:
+        raise ValueError(
+            "--learnable-layer-mixture requires stacked [T,L,D] features"
+        )
     dev_states: torch.Tensor | None = None
     dev_targets: torch.Tensor | None = None
     if dev_sessions:
-        dev_states, dev_targets, dev_hidden_size, _ = _load_examples(
+        (
+            dev_states,
+            dev_targets,
+            dev_hidden_size,
+            dev_num_input_layers,
+            _,
+        ) = _load_examples(
             args.feature_root,
             args.target_root,
             dev_sessions,
@@ -465,6 +504,11 @@ def main() -> None:
                 f"Dev feature dimension {dev_hidden_size} differs from "
                 f"training dimension {hidden_size}"
             )
+        if dev_num_input_layers != num_input_layers:
+            raise ValueError(
+                f"Dev feature layer count {dev_num_input_layers} differs from "
+                f"training count {num_input_layers}"
+            )
         print(f"development examples: {len(dev_states)}")
 
     dataset = TensorDataset(states, targets)
@@ -483,9 +527,30 @@ def main() -> None:
         num_horizons=NUM_HORIZONS,
         num_classes=NUM_CLASSES,
         dropout=args.dropout,
+        num_input_layers=num_input_layers,
+        learnable_layer_mixture=args.learnable_layer_mixture,
     ).to(device)
+    if head.layer_logits is None:
+        optimizer_parameters: Any = head.parameters()
+    else:
+        optimizer_parameters = [
+            {
+                "params": [
+                    parameter
+                    for name, parameter in head.named_parameters()
+                    if name != "layer_logits"
+                ],
+                "weight_decay": args.weight_decay,
+            },
+            {
+                "params": [head.layer_logits],
+                # Weight decay on softmax logits artificially pulls the layer
+                # distribution back toward uniform.
+                "weight_decay": 0.0,
+            },
+        ]
     optimizer = torch.optim.Adam(
-        head.parameters(),
+        optimizer_parameters,
         lr=args.learning_rate,
         weight_decay=args.weight_decay,
     )
@@ -546,6 +611,21 @@ def main() -> None:
             "train_loss": train_loss,
             "learning_rate": scheduler.get_last_lr()[0],
         }
+        layer_weights = head.layer_mixture_weights()
+        if layer_weights is not None:
+            record["layer_mixture_weights_oldest_to_newest"] = (
+                layer_weights.detach().float().cpu().tolist()
+            )
+        weight_text = (
+            ""
+            if layer_weights is None
+            else ", layer_weights=["
+            + ",".join(
+                f"{value:.4f}"
+                for value in record["layer_mixture_weights_oldest_to_newest"]
+            )
+            + "]"
+        )
 
         if dev_sessions:
             assert dev_states is not None and dev_targets is not None
@@ -560,7 +640,7 @@ def main() -> None:
             record["dev_valid_mean_map"] = dev_map
             print(
                 f"epoch {epoch:02d}: train_loss={train_loss:.6f}, "
-                f"dev_valid_mean_mAP={dev_map * 100:.3f}"
+                f"dev_valid_mean_mAP={dev_map * 100:.3f}{weight_text}"
             )
             if dev_map > float(best_dev_map):
                 best_dev_map = dev_map
@@ -576,6 +656,7 @@ def main() -> None:
                         train_sessions=train_sessions,
                         dev_sessions=dev_sessions,
                         hidden_size=hidden_size,
+                        num_input_layers=num_input_layers,
                         feature_metadata=feature_metadata,
                         history=history + [record],
                     ),
@@ -583,7 +664,9 @@ def main() -> None:
                 )
                 print(f"saved new best checkpoint: {args.output}")
         else:
-            print(f"epoch {epoch:02d}: train_loss={train_loss:.6f}")
+            print(
+                f"epoch {epoch:02d}: train_loss={train_loss:.6f}{weight_text}"
+            )
         history.append(record)
 
     if not dev_sessions:
@@ -598,6 +681,7 @@ def main() -> None:
                 train_sessions=train_sessions,
                 dev_sessions=dev_sessions,
                 hidden_size=hidden_size,
+                num_input_layers=num_input_layers,
                 feature_metadata=feature_metadata,
                 history=history,
             ),
@@ -614,6 +698,13 @@ def main() -> None:
     summary = {key: value for key, value in payload.items() if key != "head"}
     summary_path.write_text(json.dumps(summary, indent=2) + "\n")
     print(f"training summary: {summary_path}")
+    if args.learnable_layer_mixture:
+        saved_logits = payload["head"]["layer_logits"].float()
+        saved_weights = saved_logits.softmax(dim=0).tolist()
+        print(
+            "saved layer mixture weights (oldest->newest): "
+            + ", ".join(f"{value:.6f}" for value in saved_weights)
+        )
 
 
 if __name__ == "__main__":