From dbef1fc173d8e60d00b14eee8be0fb3cc93815a5 Mon Sep 17 00:00:00 2001 From: Xin Wang Date: Wed, 26 Aug 2026 09:56:54 +0800 Subject: [PATCH] Add final-layer mixture readout --- run_extract_easycom_direct.sh | 3 +++ streaming_turn_model.py | 45 +++++++++++++++++++++++++++++--- streaming_turn_model_directly.py | 10 +++++++ 3 files changed, 54 insertions(+), 4 deletions(-) diff --git a/run_extract_easycom_direct.sh b/run_extract_easycom_direct.sh index 50d6c20..ec7aff4 100644 --- a/run_extract_easycom_direct.sh +++ b/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}" \ diff --git a/streaming_turn_model.py b/streaming_turn_model.py index 00c40e3..c7d26fb 100644 --- a/streaming_turn_model.py +++ b/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() diff --git a/streaming_turn_model_directly.py b/streaming_turn_model_directly.py index 0f50b0d..0acff60 100644 --- a/streaming_turn_model_directly.py +++ b/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()