wx44wx/icassp2027-experiment
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add final-layer mixture readout

Browse Source
This commit is contained in:
Xin Wang
2026-08-26 09:56:54 +08:00
parent ca6dd88923
commit dbef1fc173
3 changed files with 54 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
run_extract_easycom_direct.sh
View File

@@ -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}" \

45
streaming_turn_model.py
View File

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

10
streaming_turn_model_directly.py
View File

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