Add learnable last-four-layer query fusion

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()