Add final-layer mixture readout

run_extract_easycom_direct.sh

@@ -16,6 +16,7 @@ LOG_ROOT="${LOG_ROOT:-${PROJECT_ROOT}/logs/easycom_direct}"
 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}"
 MODE="${MODE:-open-stream-replay}"
 DTYPE="${DTYPE:-bfloat16}"
 DECODER="${DECODER:-decord}"
@@ -72,6 +73,7 @@ run_one() {
     echo "target=${target}"
     echo "output=${output}"
     echo "memory_seconds=${MEMORY_SECONDS} video_fps=${VIDEO_FPS} mode=${MODE}"
+    echo "readout_last_n_layers=${READOUT_LAST_N_LAYERS}"
     echo "decoder=${DECODER} decoder_gpu_id=${DECODER_GPU_ID} frame_cache_size=${FRAME_CACHE_SIZE}"
 
     if ! CUDA_VISIBLE_DEVICES="${gpu}" PYTHONUNBUFFERED=1 \
@@ -84,6 +86,7 @@ run_one() {
         --memory-seconds "${MEMORY_SECONDS}" \
         --native-block-seconds "${NATIVE_BLOCK_SECONDS}" \
         --video-fps "${VIDEO_FPS}" \
+        --readout-last-n-layers "${READOUT_LAST_N_LAYERS}" \
         --frame-cache-size "${FRAME_CACHE_SIZE}" \
         --decoder "${DECODER}" \
         --decoder-gpu-id "${DECODER_GPU_ID}" \

streaming_turn_model.py

@@ -140,6 +140,7 @@ class StreamConfig:
     video_fps: float = 2.0
     video_min_pixels: int = 64 * 28 * 28
     video_max_pixels: int = 128 * 28 * 28
+    readout_last_n_layers: int = 1
 
     def validate(self) -> None:
         if self.step_seconds <= 0:
@@ -154,6 +155,8 @@ class StreamConfig:
             raise ValueError("video_min_pixels must be positive")
         if self.video_max_pixels < self.video_min_pixels:
             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")
         _require_integer_ratio(
             self.native_block_seconds, self.step_seconds, "native block / step"
         )
@@ -952,6 +955,7 @@ def _read_state(
     *,
     next_rope_offset: int,
     query_ids: torch.LongTensor,
+    last_n_layers: int = 1,
 ) -> torch.Tensor:
     """Read one q_t without committing the artificial assistant query."""
 
@@ -981,11 +985,29 @@ def _read_state(
         past_key_values=query_cache,
         inputs_embeds=embeddings,
         use_cache=True,
-        output_hidden_states=False,
+        output_hidden_states=last_n_layers > 1,
         return_dict=True,
         cache_position=cache_position,
     )
-    return outputs.last_hidden_state[:, -1, :]
+    return _readout_hidden_state(outputs, last_n_layers=last_n_layers)
+
+
+def _readout_hidden_state(outputs: Any, *, last_n_layers: int) -> torch.Tensor:
+    """Return the query state, optionally averaging the final decoder layers."""
+
+    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}"
+        )
+    return torch.stack(
+        [state[:, -1, :] for state in hidden_states[-last_n_layers:]],
+        dim=0,
+    ).mean(dim=0)
 
 
 @torch.inference_mode()
@@ -996,6 +1018,7 @@ def _read_open_state(
     closure_ids: torch.LongTensor,
     closure_position_ids: torch.LongTensor,
     query_ids: torch.LongTensor,
+    last_n_layers: int = 1,
 ) -> torch.Tensor:
     """Temporarily close an open media turn, read q_t, and discard the branch."""
 
@@ -1051,11 +1074,11 @@ def _read_open_state(
         past_key_values=branch_cache,
         inputs_embeds=embeddings,
         use_cache=True,
-        output_hidden_states=False,
+        output_hidden_states=last_n_layers > 1,
         return_dict=True,
         cache_position=cache_position,
     )
-    return outputs.last_hidden_state[:, -1, :]
+    return _readout_hidden_state(outputs, last_n_layers=last_n_layers)
 
 
 @torch.inference_mode()
@@ -1171,6 +1194,7 @@ class CausalPrefixExtractor:
             media_cache,
             next_rope_offset=next_offset,
             query_ids=self.query_ids,
+            last_n_layers=self.config.readout_last_n_layers,
         )
 
     def extract(
@@ -1333,6 +1357,7 @@ class OpenStreamReplayExtractor:
             closure_ids=closure_ids,
             closure_position_ids=closure_position_ids,
             query_ids=self.query_ids,
+            last_n_layers=self.config.readout_last_n_layers,
         )
 
     def _cache_from_intervals(
@@ -1541,6 +1566,7 @@ class OpenStreamReplayExtractor:
                 closure_ids=pieces.closure_ids,
                 closure_position_ids=closure_position_ids,
                 query_ids=self.query_ids,
+                last_n_layers=self.config.readout_last_n_layers,
             )
             states.append(state.cpu())
 
@@ -1781,6 +1807,7 @@ class OverlappingWorkingMemoryExtractor:
             closure_ids=working_pieces.closure_ids,
             closure_position_ids=closure_position_ids,
             query_ids=self.query_ids,
+            last_n_layers=self.config.readout_last_n_layers,
         )
 
     def extract(
@@ -2124,6 +2151,15 @@ def parse_args() -> argparse.Namespace:
     )
     parser.add_argument("--video-min-pixels", type=int, default=64 * 28 * 28)
     parser.add_argument("--video-max-pixels", type=int, default=128 * 28 * 28)
+    parser.add_argument(
+        "--readout-last-n-layers",
+        type=int,
+        default=1,
+        help=(
+            "Uniformly average the final N Thinker hidden layers at the "
+            "causal readout token (default: 1 preserves the original result)."
+        ),
+    )
     parser.add_argument("--verify-every", type=int, default=0)
     parser.add_argument(
         "--verify-continuous-every",
@@ -2191,6 +2227,7 @@ def main() -> None:
         video_fps=args.video_fps,
         video_min_pixels=args.video_min_pixels,
         video_max_pixels=args.video_max_pixels,
+        readout_last_n_layers=args.readout_last_n_layers,
     )
     config.validate()
 

streaming_turn_model_directly.py

@@ -1052,6 +1052,15 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
     parser.add_argument("--video-fps", type=float, default=2.0)
     parser.add_argument("--video-min-pixels", type=int, default=64 * 28 * 28)
     parser.add_argument("--video-max-pixels", type=int, default=128 * 28 * 28)
+    parser.add_argument(
+        "--readout-last-n-layers",
+        type=int,
+        default=1,
+        help=(
+            "Uniformly average the final N Thinker hidden layers at the "
+            "causal readout token (default: 1 preserves the original result)."
+        ),
+    )
     parser.add_argument("--frame-cache-size", type=int, default=8)
     parser.add_argument(
         "--audio-feature-device",
@@ -1331,6 +1340,7 @@ def main(
         video_fps=args.video_fps,
         video_min_pixels=args.video_min_pixels,
         video_max_pixels=args.video_max_pixels,
+        readout_last_n_layers=args.readout_last_n_layers,
     )
     config.validate()