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

Add Qwen tail and chunk ablations

Browse Source
This commit is contained in:
Xin Wang
2026-09-17 02:32:08 +08:00
parent a183fe5944
commit a9f748d6a6
7 changed files with 1219 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files
streaming_turn_model.py
+734 -12
View File
@@ -1,12 +1,16 @@
"""Omni-only causal turn-taking features at 5 Hz.
This module implements three inference paths for Qwen2.5-Omni-3B:
This module implements five inference paths for Qwen2.5-Omni-3B:
1. ``causal-prefix``: recompute a causal AV window ending at every 200 ms.
This is deliberately slow, but is the simplest correctness baseline.
2. ``open-stream-replay``: cache completed AV blocks inside one
still-open user/media turn; closure and readout tokens remain temporary.
3. ``overlap-working-memory``: a frozen ablation that serializes cached
3. ``completed-block-hold``: discard the incomplete native block and hold the
most recent completed-block state on the 5 Hz output grid.
4. ``micro-chunk-200ms``: commit each 200-ms interval independently as an
explicitly non-native chunk-size stress test.
5. ``overlap-working-memory``: a frozen ablation that serializes cached
complete blocks and an overlapping ``[t-2,t]`` recent recap.
Neither path uses an RNN. A small head maps one Thinker state to the same
@@ -37,7 +41,7 @@ import logging
import math
import subprocess
from collections import OrderedDict
from dataclasses import asdict, dataclass
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Any, Literal
@@ -63,6 +67,8 @@ from transformers.cache_utils import DynamicCache
Mode = Literal[
"causal-prefix",
"open-stream-replay",
"completed-block-hold",
"micro-chunk-200ms",
"overlap-working-memory",
]
BackboneLoader = Literal["qwen", "roma"]
@@ -389,6 +395,7 @@ class ExtractionResult:
times: torch.Tensor # [N], seconds in the original session
states: torch.Tensor # [N,D] or preserved query layers [N,L,D]
verification: list[dict[str, Any]]
diagnostics: dict[str, Any] = field(default_factory=dict)
@dataclass
@@ -621,6 +628,29 @@ def _time_grid(start: float, end: float, step: float) -> list[float]:
return [start + i * step for i in range(1, count + 1)]
def _periodic_boundary_mask(
times: list[float],
*,
origin: float,
period: float,
) -> torch.BoolTensor:
"""Mark right edges aligned to a positive multiple of ``period``."""
values = []
for current_time in times:
ratio = (current_time - origin) / period
values.append(
ratio > 0
and math.isclose(
ratio,
round(ratio),
rel_tol=0.0,
abs_tol=1e-6,
)
)
return torch.tensor(values, dtype=torch.bool)
def _media_content(
*,
video_path: str | Path | None,
@@ -1076,6 +1106,31 @@ def _tmrope_positions_per_chunk(
return int(rounded)
def _tmrope_positions_per_step(
model: Qwen2_5OmniThinkerForConditionalGeneration,
step_seconds: float,
) -> int:
"""Return the native TMRoPE advance for one micro-chunk step."""
positions_per_second = float(model.config.position_id_per_seconds)
value = positions_per_second * step_seconds
rounded = round(value)
if not math.isclose(value, rounded, rel_tol=0.0, abs_tol=1e-6):
raise ValueError(f"Non-integral TMRoPE positions per step: {value}")
steps_per_chunk = _require_integer_ratio(
float(model.config.seconds_per_chunk),
step_seconds,
"model native chunk / micro step",
)
positions_per_chunk = _tmrope_positions_per_chunk(model)
if rounded * steps_per_chunk != positions_per_chunk:
raise ValueError(
"Micro-step TMRoPE advances do not tile one native chunk: "
f"{rounded} * {steps_per_chunk} != {positions_per_chunk}"
)
return int(rounded)
def _cache_length(cache: Any | None) -> int:
if cache is None:
return 0
@@ -1906,10 +1961,634 @@ class OpenStreamReplayExtractor:
global_step_offset=0,
progress=progress,
)
steps_per_block = _require_integer_ratio(
self.config.native_block_seconds,
self.config.step_seconds,
"native block / step",
)
return ExtractionResult(
times=torch.tensor(all_times, dtype=torch.float64),
states=torch.cat(all_states, dim=0),
verification=all_checks,
diagnostics={
"production_processor_calls": len(all_times),
"production_prepared_media_seconds": float(
sum(
((index % steps_per_block) + 1)
* self.config.step_seconds
for index in range(len(all_times))
)
),
"state_source_times": torch.tensor(
all_times, dtype=torch.float64
),
"initialization_mask": torch.zeros(
len(all_times), dtype=torch.bool
),
"state_update_mask": torch.ones(
len(all_times), dtype=torch.bool
),
"payload_commit_mask": _periodic_boundary_mask(
all_times,
origin=start,
period=self.config.native_block_seconds,
),
"native_boundary_mask": _periodic_boundary_mask(
all_times,
origin=start,
period=self.config.native_block_seconds,
),
},
)
class CompletedBlockHoldExtractor(OpenStreamReplayExtractor):
"""Commit native 2-s blocks and causally hold the last emitted state.
This is the deliberately stale control for partial-tail processing. It
never constructs an incomplete AV interval and never backfills a state
into rows before that block was fully observed.
"""
def _zero_state(self) -> torch.Tensor:
hidden_size = _hidden_size(self.model)
dtype = getattr(self.model, "dtype", torch.float32)
if self.config.readout_layer_fusion == "stack":
shape = (1, self.config.readout_last_n_layers, hidden_size)
else:
shape = (1, hidden_size)
return torch.zeros(shape, dtype=dtype, device="cpu")
def extract(
self,
*,
video_path: str | Path | None,
audio_path: str | Path | None,
start: float,
end: float,
use_audio_in_video: bool,
verify_every: int = 0,
verify_continuous_every: int = 0,
show_progress: bool = True,
) -> ExtractionResult:
if verify_every < 0 or verify_continuous_every < 0:
raise ValueError("verification intervals must be >= 0")
times = _time_grid(start, end, self.config.step_seconds)
committed_cache = None
block_start = start
history_start = start
payload_base = None
positions_per_chunk = _tmrope_positions_per_chunk(self.model)
latest = self._zero_state()
source_time = start
states: list[torch.Tensor] = []
source_times: list[float] = []
updates: list[bool] = []
initialization: list[bool] = []
checks: list[dict[str, Any]] = []
processor_calls = 0
prepared_seconds = 0.0
replay_block_count = max(
0,
_require_integer_ratio(
self.config.memory_seconds,
self.config.native_block_seconds,
"memory / native block",
)
- 1,
)
needs_rolling_rebuild = (
self.config.memory_seconds < end - start - 1e-6
)
completed_blocks = CompletedBlockPieceCache(
origin=start,
block_seconds=self.config.native_block_seconds,
max_blocks=replay_block_count if needs_rolling_rebuild else 0,
)
with tqdm(
total=len(times),
desc="completed-block-hold",
unit="state",
dynamic_ncols=True,
mininterval=0.0,
miniters=1,
disable=not show_progress,
) as progress:
for global_step, current_time in enumerate(times, start=1):
block_end = block_start + self.config.native_block_seconds
update = math.isclose(
current_time,
block_end,
rel_tol=0.0,
abs_tol=1e-6,
)
if update:
if (
committed_cache is not None
and block_start - history_start
>= self.config.memory_seconds - 1e-6
):
history_start = block_start - (
self.config.memory_seconds
- self.config.native_block_seconds
)
history_intervals = _block_intervals(
history_start,
block_start,
self.config.native_block_seconds,
)
if history_intervals:
prepared_history = completed_blocks.get_many(
history_intervals
)
(
committed_cache,
_,
payload_base,
_,
_,
) = self._cache_from_intervals(
video_path=video_path,
audio_path=audio_path,
intervals=history_intervals,
use_audio_in_video=use_audio_in_video,
prepared_blocks=prepared_history,
)
else:
committed_cache = None
payload_base = None
pieces = self._prepare_pieces(
video_path=video_path,
audio_path=audio_path,
interval=(block_start, block_end),
use_audio_in_video=use_audio_in_video,
)
processor_calls += 1
prepared_seconds += block_end - block_start
if payload_base is None:
payload_base = int(
pieces.payload_position_ids.min().item()
)
block_index = round(
(block_start - history_start)
/ self.config.native_block_seconds
)
payload_offset = (
payload_base + block_index * positions_per_chunk
)
(
committed_cache,
_,
closure_position_ids,
) = _append_open_pieces(
self.model,
pieces,
past_key_values=committed_cache,
rope_offset=payload_offset,
include_prefix=committed_cache is None,
)
latest = _read_open_state(
self.model,
committed_cache,
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,
layer_fusion=self.config.readout_layer_fusion,
).cpu()
source_time = current_time
completed_blocks.put(
CachedOpenMediaBlock(
start=block_start,
end=block_end,
pieces=pieces,
)
)
if verify_every > 0 and global_step % verify_every == 0:
exact = self.state_from_scratch(
video_path=video_path,
audio_path=audio_path,
intervals=_block_intervals(
history_start,
current_time,
self.config.native_block_seconds,
),
use_audio_in_video=use_audio_in_video,
)
checks.append(
_compare_states(
current_time,
latest,
exact,
kind="completed_hold_cache_replay",
)
)
if (
verify_continuous_every > 0
and global_step % verify_continuous_every == 0
):
continuous = self.continuous_reference.state_for_intervals(
video_path=video_path,
audio_path=audio_path,
intervals=[(history_start, current_time)],
use_audio_in_video=use_audio_in_video,
)
checks.append(
_compare_states(
current_time,
latest,
continuous,
kind="completed_hold_continuous_semantics",
)
)
block_start = block_end
states.append(latest)
source_times.append(source_time)
updates.append(update)
initialization.append(source_time <= start + 1e-6)
progress.set_postfix_str(
f"source={source_time:.1f}s", refresh=False
)
progress.update(1)
return ExtractionResult(
times=torch.tensor(times, dtype=torch.float64),
states=torch.cat(states, dim=0),
verification=checks,
diagnostics={
"production_processor_calls": processor_calls,
"production_prepared_media_seconds": float(prepared_seconds),
"state_source_times": torch.tensor(
source_times, dtype=torch.float64
),
"initialization_mask": torch.tensor(
initialization, dtype=torch.bool
),
"state_update_mask": torch.tensor(updates, dtype=torch.bool),
"payload_commit_mask": torch.tensor(updates, dtype=torch.bool),
"native_boundary_mask": torch.tensor(updates, dtype=torch.bool),
},
)
class MicroChunkReplayExtractor(OpenStreamReplayExtractor):
"""Commit each new 200-ms payload as a non-native streaming stress test."""
def __init__(
self,
model: Qwen2_5OmniThinkerForConditionalGeneration,
processor: Qwen2_5OmniProcessor,
config: StreamConfig,
fixed_grid_av: FixedGridAV | None = None,
) -> None:
super().__init__(model, processor, config, fixed_grid_av=fixed_grid_av)
self.positions_per_step = _tmrope_positions_per_step(
model, config.step_seconds
)
def _prepare_micro_pieces(
self,
*,
video_path: str | Path | None,
audio_path: str | Path | None,
interval: tuple[float, float],
use_audio_in_video: bool,
) -> OpenMediaPieces:
duration = interval[1] - interval[0]
if not math.isclose(
duration,
self.config.step_seconds,
rel_tol=0.0,
abs_tol=1e-6,
):
raise ValueError(
f"Micro interval is {duration:g}s, expected "
f"{self.config.step_seconds:g}s"
)
pieces = self._prepare_pieces(
video_path=video_path,
audio_path=audio_path,
interval=interval,
use_audio_in_video=use_audio_in_video,
)
temporal_positions = pieces.payload_position_ids[0]
temporal_span = int(
temporal_positions.max().item()
- temporal_positions.min().item()
)
if temporal_span >= self.positions_per_step:
raise RuntimeError(
"A 200-ms payload does not fit its native TMRoPE step: "
f"interval={interval}, temporal_span={temporal_span}, "
f"step_positions={self.positions_per_step}"
)
return pieces
def _micro_cache_from_intervals(
self,
*,
video_path: str | Path | None,
audio_path: str | Path | None,
intervals: list[tuple[float, float]],
use_audio_in_video: bool,
prepared_blocks: list[CachedOpenMediaBlock] | None = None,
) -> tuple[Any, int, torch.LongTensor, torch.LongTensor]:
if not intervals:
raise ValueError("At least one micro interval is required")
if prepared_blocks is not None and len(prepared_blocks) != len(intervals):
raise ValueError("prepared micro-block count does not match intervals")
media_cache = None
payload_base = None
closure_ids = None
closure_position_ids = None
for index, interval in enumerate(intervals):
if prepared_blocks is None:
pieces = self._prepare_micro_pieces(
video_path=video_path,
audio_path=audio_path,
interval=interval,
use_audio_in_video=use_audio_in_video,
)
else:
prepared = prepared_blocks[index]
if not (
math.isclose(
prepared.start,
interval[0],
rel_tol=0.0,
abs_tol=1e-6,
)
and math.isclose(
prepared.end,
interval[1],
rel_tol=0.0,
abs_tol=1e-6,
)
):
raise ValueError("prepared micro-block interval mismatch")
pieces = prepared.pieces
if payload_base is None:
payload_base = int(pieces.payload_position_ids.min().item())
payload_offset = payload_base + index * self.positions_per_step
(
media_cache,
_,
closure_position_ids,
) = _append_open_pieces(
self.model,
pieces,
past_key_values=media_cache,
rope_offset=payload_offset,
include_prefix=index == 0,
)
closure_ids = pieces.closure_ids
if media_cache is None or payload_base is None:
raise RuntimeError("Micro-chunk replay produced no cache")
if closure_ids is None or closure_position_ids is None:
raise RuntimeError("Micro-chunk replay produced no closure")
return media_cache, payload_base, closure_ids, closure_position_ids
def state_from_micro_scratch(
self,
*,
video_path: str | Path | None,
audio_path: str | Path | None,
intervals: list[tuple[float, float]],
use_audio_in_video: bool,
) -> torch.Tensor:
cache, _, closure_ids, closure_positions = (
self._micro_cache_from_intervals(
video_path=video_path,
audio_path=audio_path,
intervals=intervals,
use_audio_in_video=use_audio_in_video,
)
)
return _read_open_state(
self.model,
cache,
closure_ids=closure_ids,
closure_position_ids=closure_positions,
query_ids=self.query_ids,
last_n_layers=self.config.readout_last_n_layers,
layer_fusion=self.config.readout_layer_fusion,
)
def extract(
self,
*,
video_path: str | Path | None,
audio_path: str | Path | None,
start: float,
end: float,
use_audio_in_video: bool,
verify_every: int = 0,
verify_continuous_every: int = 0,
show_progress: bool = True,
) -> ExtractionResult:
if verify_every < 0 or verify_continuous_every < 0:
raise ValueError("verification intervals must be >= 0")
times = _time_grid(start, end, self.config.step_seconds)
committed_cache = None
block_start = start
history_start = start
payload_base = None
states: list[torch.Tensor] = []
checks: list[dict[str, Any]] = []
prepared_seconds = 0.0
needs_rolling_rebuild = (
self.config.memory_seconds < end - start - 1e-6
)
replay_step_count = max(
0,
_require_integer_ratio(
self.config.memory_seconds - self.config.native_block_seconds,
self.config.step_seconds,
"micro replay memory / step",
),
)
prepared_cache = CompletedBlockPieceCache(
origin=start,
block_seconds=self.config.step_seconds,
max_blocks=replay_step_count if needs_rolling_rebuild else 0,
)
with tqdm(
total=len(times),
desc="micro-chunk-200ms",
unit="state",
dynamic_ncols=True,
mininterval=0.0,
miniters=1,
disable=not show_progress,
) as progress:
for global_step, current_time in enumerate(times, start=1):
while (
current_time
> block_start + self.config.native_block_seconds + 1e-6
):
block_start += self.config.native_block_seconds
if (
committed_cache is not None
and block_start - history_start
>= self.config.memory_seconds - 1e-6
):
history_start = block_start - (
self.config.memory_seconds
- self.config.native_block_seconds
)
history_intervals = _block_intervals(
history_start,
block_start,
self.config.step_seconds,
)
if history_intervals:
prepared_history = prepared_cache.get_many(
history_intervals
)
(
committed_cache,
payload_base,
_,
_,
) = self._micro_cache_from_intervals(
video_path=video_path,
audio_path=audio_path,
intervals=history_intervals,
use_audio_in_video=use_audio_in_video,
prepared_blocks=prepared_history,
)
else:
committed_cache = None
payload_base = None
interval_start = current_time - self.config.step_seconds
pieces = self._prepare_micro_pieces(
video_path=video_path,
audio_path=audio_path,
interval=(interval_start, current_time),
use_audio_in_video=use_audio_in_video,
)
prepared_seconds += current_time - interval_start
if payload_base is None:
payload_base = int(pieces.payload_position_ids.min().item())
micro_index = round(
(interval_start - history_start)
/ self.config.step_seconds
)
payload_offset = (
payload_base + micro_index * self.positions_per_step
)
(
committed_cache,
_,
closure_positions,
) = _append_open_pieces(
self.model,
pieces,
past_key_values=committed_cache,
rope_offset=payload_offset,
include_prefix=committed_cache is None,
)
state = _read_open_state(
self.model,
committed_cache,
closure_ids=pieces.closure_ids,
closure_position_ids=closure_positions,
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())
prepared_cache.put(
CachedOpenMediaBlock(
start=interval_start,
end=current_time,
pieces=pieces,
)
)
if verify_every > 0 and global_step % verify_every == 0:
intervals = _block_intervals(
history_start,
current_time,
self.config.step_seconds,
)
exact = self.state_from_micro_scratch(
video_path=video_path,
audio_path=audio_path,
intervals=intervals,
use_audio_in_video=use_audio_in_video,
)
checks.append(
_compare_states(
current_time,
state,
exact,
kind="micro_chunk_cache_replay",
)
)
if (
verify_continuous_every > 0
and global_step % verify_continuous_every == 0
):
continuous = self.continuous_reference.state_for_intervals(
video_path=video_path,
audio_path=audio_path,
intervals=[(history_start, current_time)],
use_audio_in_video=use_audio_in_video,
)
checks.append(
_compare_states(
current_time,
state,
continuous,
kind="micro_chunk_vs_continuous_semantics",
)
)
if math.isclose(
current_time,
block_start + self.config.native_block_seconds,
rel_tol=0.0,
abs_tol=1e-6,
):
block_start = current_time
progress.set_postfix_str(
f"media={current_time:.1f}s", refresh=False
)
progress.update(1)
output_times = torch.tensor(times, dtype=torch.float64)
return ExtractionResult(
times=output_times,
states=torch.cat(states, dim=0),
verification=checks,
diagnostics={
"production_processor_calls": len(times),
"production_prepared_media_seconds": float(prepared_seconds),
"state_source_times": output_times.clone(),
"initialization_mask": torch.zeros(
len(times), dtype=torch.bool
),
"state_update_mask": torch.ones(
len(times), dtype=torch.bool
),
"payload_commit_mask": torch.ones(
len(times), dtype=torch.bool
),
"native_boundary_mask": _periodic_boundary_mask(
times,
origin=start,
period=self.config.native_block_seconds,
),
"tmrope_positions_per_micro_step": self.positions_per_step,
},
)
@@ -2240,8 +2919,8 @@ def _compare_states(
*,
kind: str,
) -> dict[str, Any]:
cached_f = cached.float()
exact_f = exact.float()
cached_f = cached.detach().float()
exact_f = exact.detach().to(device=cached_f.device, dtype=torch.float32)
cached_flat = cached_f.reshape(1, -1)
exact_flat = exact_f.reshape(1, -1)
return {
@@ -2394,7 +3073,12 @@ def parse_args() -> argparse.Namespace:
parser.add_argument("--output", type=Path, required=True)
parser.add_argument(
"--mode",
choices=("causal-prefix", "open-stream-replay"),
choices=(
"causal-prefix",
"open-stream-replay",
"completed-block-hold",
"micro-chunk-200ms",
),
default="open-stream-replay",
)
parser.add_argument("--model-id", default=MODEL_ID)
@@ -2554,7 +3238,12 @@ def main() -> None:
show_progress=not args.no_progress,
)
else:
extractor = OpenStreamReplayExtractor(
extractor_class = {
"open-stream-replay": OpenStreamReplayExtractor,
"completed-block-hold": CompletedBlockHoldExtractor,
"micro-chunk-200ms": MicroChunkReplayExtractor,
}[args.mode]
extractor = extractor_class(
model, processor, config, fixed_grid_av=fixed_grid_av
)
result = extractor.extract(
@@ -2573,11 +3262,16 @@ def main() -> None:
"backbone_loader": args.backbone_loader,
"backbone_metadata": backbone_metadata(model),
"mode": args.mode,
"stream_representation": (
"single_causal_window"
if args.mode == "causal-prefix"
else "single_open_media_turn_with_2s_payload_commits"
),
"stream_representation": {
"causal-prefix": "single_causal_window",
"open-stream-replay": (
"native_2s_completed_blocks_plus_temporary_partial_tail"
),
"completed-block-hold": (
"native_2s_completed_blocks_forward_held_to_5hz"
),
"micro-chunk-200ms": "non_native_200ms_payload_commits",
}[args.mode],
"config": asdict(config),
"start": args.start,
"end": end,
@@ -2614,9 +3308,37 @@ def main() -> None:
if args.mode == "open-stream-replay"
else None
),
"chunk_ablation_policy": (
{
"kind": "native_completed_block_hold",
"native_chunk_seconds": config.native_block_seconds,
"output_step_seconds": config.step_seconds,
"tail_policy": "discard_incomplete_tail_and_hold_last_state",
"initial_state": "zeros_until_first_complete_native_block",
"fresh_at_every_5hz_step": False,
"backbone": "frozen",
}
if args.mode == "completed-block-hold"
else (
{
"kind": "non_native_micro_chunk_stress_test",
"micro_chunk_seconds": config.step_seconds,
"native_chunk_seconds": config.native_block_seconds,
"tmrope_positions_per_micro_step": (
_tmrope_positions_per_step(model, config.step_seconds)
),
"tail_policy": "commit_each_5hz_interval_as_a_payload",
"fresh_at_every_5hz_step": True,
"backbone": "frozen",
}
if args.mode == "micro-chunk-200ms"
else None
)
),
"times": result.times,
"states": result.states,
"verification": result.verification,
"extraction_diagnostics": result.diagnostics,
}
if args.head_checkpoint is not None:
head, head_loss_name = _load_head(