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

Add causal person fusion and MuVAP stage-3 baselines

Browse Source
This commit is contained in:
Xin Wang
2026-09-10 06:54:22 +08:00
parent 0534c9d3d0
commit d426b4740a
225 changed files with 27716 additions and 119 deletions
Download Patch File Download Diff File Expand all files Collapse all files

436
easycom_head.py
View File

@@ -66,8 +66,18 @@ CURRENT_FRAME_PERSON_SET_FUSION_MODES = (
"oracle-frame-set-attention",
"audio-person-frame-set",
"gated-person-state-adapter",
"highrate-person-av-attention",
)
HIGHRATE_PERSON_AV_FUSION_MODES = ("highrate-person-av-attention",)
TRACK_PERSON_FUSION_MODES = (
"track-gru-residual",
"causal-person-adapter-v2",
"causal-person-adapter-egospeak-gru",
)
PERSON_ADAPTER_V2_MODES = (
"causal-person-adapter-v2",
"causal-person-adapter-egospeak-gru",
)
TRACK_PERSON_FUSION_MODES = ("track-gru-residual",)
PERSON_SET_FUSION_MODES = (
*CURRENT_FRAME_PERSON_SET_FUSION_MODES,
*TRACK_PERSON_FUSION_MODES,
@@ -79,7 +89,16 @@ LOCAL_FUSION_SCOPES = {
"oracle-frame-set-attention": "causal-oracle-visible-person-speaking-history",
"audio-person-frame-set": "f5-screening-current-frame-no-person-history",
"gated-person-state-adapter": "f6-screening-post-thinker-not-internal",
"highrate-person-av-attention": (
"global-omni-traditional-audio-base-plus-zero-init-25hz-person-residual"
),
"track-gru-residual": "causal-anonymous-track-gru-pooled-baseline-residual",
"causal-person-adapter-v2": (
"causal-anonymous-track-gru-dynamic-gated-logit-residual"
),
"causal-person-adapter-egospeak-gru": (
"causal-anonymous-track-egospeak-style-gru-dynamic-gated-logit-residual"
),
}
DatasetName = Literal["easycom", "ego4d"]
@@ -211,15 +230,36 @@ class TurnPredictionHead(nn.Module):
self.audio_person_gate: nn.Parameter | None = None
self.person_state_adapter: nn.Module | None = None
self.person_state_adapter_gate: nn.Parameter | None = None
self.person_track_input_embedding: nn.Module | None = None
self.person_track_gru: nn.GRU | None = None
self.person_track_output_activation: nn.Module | None = None
self.person_track_delta: nn.Module | None = None
self.person_track_residual_gate: nn.Parameter | None = None
self.person_track_local_state: nn.Module | None = None
self.person_track_dynamic_gate: nn.Module | None = None
self.person_track_residual_output: nn.Linear | None = None
self.highrate_person_residual: nn.Linear | None = None
self._last_person_adapter_diagnostics: dict[str, float] = {}
if local_fusion in PERSON_SET_FUSION_MODES:
self.person_input_projection = nn.Sequential(
nn.LayerNorm(local_visual_dim),
nn.Linear(local_visual_dim, person_attention_dim),
nn.GELU(),
)
if local_fusion == "causal-person-adapter-egospeak-gru":
# Clean-room adaptation of EgoSpeak MiniROADA's temporal stem
# (MiniROAD/model/rnn/rnn.py): Linear -> LayerNorm -> ReLU ->
# Dropout -> one-layer GRU -> ReLU. EgoSpeak applies it to a
# global RGB+audio sequence; this ablation applies the same
# layer order independently to each anonymous visual track.
self.person_track_input_embedding = nn.Sequential(
nn.Linear(local_visual_dim, person_attention_dim),
nn.LayerNorm(person_attention_dim),
nn.ReLU(),
nn.Dropout(dropout),
)
self.person_track_output_activation = nn.ReLU()
else:
self.person_input_projection = nn.Sequential(
nn.LayerNorm(local_visual_dim),
nn.Linear(local_visual_dim, person_attention_dim),
nn.GELU(),
)
self.person_state_query = nn.Sequential(
nn.LayerNorm(hidden_size),
nn.Linear(hidden_size, person_attention_dim),
@@ -263,25 +303,75 @@ class TurnPredictionHead(nn.Module):
nn.Linear(person_attention_dim, hidden_size),
)
self.person_state_adapter_gate = nn.Parameter(torch.zeros(()))
elif local_fusion == "track-gru-residual":
elif local_fusion in TRACK_PERSON_FUSION_MODES:
# Apply one shared causal GRU independently to every anonymous
# face track. The regular pooled G+A+V head remains the exact
# base path; this branch can only add a zero-gated correction.
# base path for every track experiment.
self.person_track_gru = nn.GRU(
person_attention_dim,
person_attention_dim,
num_layers=1,
batch_first=True,
)
delta_input = hidden_size // 2 + person_attention_dim + 1
self.person_track_delta = nn.Sequential(
nn.LayerNorm(delta_input),
nn.Linear(delta_input, hidden_size // 2),
nn.GELU(),
nn.Dropout(dropout),
nn.Linear(hidden_size // 2, num_horizons * num_classes),
if local_fusion == "track-gru-residual":
# Legacy screen: the scalar zero gate is intentionally
# preserved so old checkpoints and the measured null result
# remain reproducible.
delta_input = hidden_size // 2 + person_attention_dim + 1
self.person_track_delta = nn.Sequential(
nn.LayerNorm(delta_input),
nn.Linear(delta_input, hidden_size // 2),
nn.GELU(),
nn.Dropout(dropout),
nn.Linear(hidden_size // 2, num_horizons * num_classes),
)
self.person_track_residual_gate = nn.Parameter(torch.zeros(()))
elif local_fusion in PERSON_ADAPTER_V2_MODES:
# V2: person-track context and direct local audio form a
# floor-evidence state. A sample-dependent vector gate
# controls its correction to the exact pooled baseline.
# Zero-initialize only the final logit projection: unlike a
# zero scalar multiplier, it starts exactly at the baseline
# without permanently starving the branch of gradients.
base_feature_dim = hidden_size // 2
projected_audio_dim = local_projection_dim or local_audio_dim
local_state_input = (
person_attention_dim + projected_audio_dim + 2
)
self.person_track_local_state = nn.Sequential(
nn.LayerNorm(local_state_input),
nn.Linear(local_state_input, base_feature_dim),
nn.GELU(),
nn.Dropout(dropout),
)
self.person_track_dynamic_gate = nn.Sequential(
nn.LayerNorm(base_feature_dim * 2),
nn.Linear(base_feature_dim * 2, base_feature_dim),
)
self.person_track_residual_output = nn.Linear(
base_feature_dim,
num_horizons * num_classes,
)
nn.init.zeros_(self.person_track_residual_output.weight)
nn.init.zeros_(self.person_track_residual_output.bias)
else:
raise AssertionError(
f"Unhandled track fusion mode: {local_fusion}"
)
elif local_fusion in HIGHRATE_PERSON_AV_FUSION_MODES:
# Keep global Omni + shared traditional audio as the complete
# base path. The high-rate ASD
# person set contributes only a zero-initialized logit
# residual, so step zero is exactly the same function as the
# audio-only head with matching shared weights. This gives a
# controlled answer to whether person evidence is useful.
residual_input = hidden_size // 2 + person_attention_dim + 1
self.highrate_person_residual = nn.Linear(
residual_input,
num_horizons * num_classes,
)
self.person_track_residual_gate = nn.Parameter(torch.zeros(()))
nn.init.zeros_(self.highrate_person_residual.weight)
nn.init.zeros_(self.highrate_person_residual.bias)
else:
fused_size += person_attention_dim + 1
self.net = nn.Sequential(
@@ -413,17 +503,20 @@ class TurnPredictionHead(nn.Module):
local: torch.Tensor,
person_histories: torch.Tensor | None,
history_mask: torch.Tensor | None,
) -> tuple[torch.Tensor, torch.Tensor]:
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
"""Encode each causal track with a shared GRU, then query the set."""
if states.ndim != 2:
raise ValueError("Track GRU fusion expects batched [B,D] states")
if local.shape != (len(states), self.local_input_size):
raise ValueError("Track GRU fusion requires pooled [visual,audio,valid] input")
if person_histories is None or person_histories.ndim != 4:
raise ValueError(
"Track GRU fusion requires person histories "
"[B,P,person_history_frames,visual_dim]"
)
if (
person_histories is None
or person_histories.ndim != 4
or person_histories.shape[0] != len(states)
person_histories.shape[0] != len(states)
or person_histories.shape[2] != self.person_history_frames
or person_histories.shape[3] != self.local_visual_dim
):
@@ -431,49 +524,22 @@ class TurnPredictionHead(nn.Module):
"Track GRU fusion requires person histories "
"[B,P,person_history_frames,visual_dim]"
)
if (
history_mask is None
or history_mask.dtype != torch.bool
or history_mask.shape != person_histories.shape[:3]
):
raise ValueError("Track GRU fusion requires a boolean history mask [B,P,L]")
if not torch.isfinite(person_histories).all():
raise ValueError("Person track histories contain NaN/Inf")
if bool((history_mask[:, :, 1:] & ~history_mask[:, :, :-1]).any()):
raise ValueError("Person track histories must be left aligned")
assert self.person_input_projection is not None
assert self.person_state_query is not None
assert self.person_audio_query is not None
assert self.person_attention is not None
assert self.null_person is not None
assert self.person_track_gru is not None
batch, people, length, _ = person_histories.shape
person_valid = history_mask.any(dim=-1)
if people:
projected = self.person_input_projection(person_histories)
flat = projected.reshape(batch * people, length, -1)
lengths = history_mask.sum(dim=-1).reshape(-1)
safe_lengths = lengths.clamp_min(1).to(device="cpu")
packed = nn.utils.rnn.pack_padded_sequence(
flat,
safe_lengths,
batch_first=True,
enforce_sorted=False,
)
_, hidden = self.person_track_gru(packed)
track_states = hidden[-1].reshape(batch, people, -1)
track_states = track_states * person_valid[..., None]
else:
track_states = states.new_zeros(
(batch, 0, self.person_attention_dim)
)
track_states, person_valid = self._encode_person_track_histories(
person_histories,
history_mask,
)
batch = len(states)
visual_width = self.local_visual_dim + 1
audio_values = local[:, visual_width : visual_width + self.local_audio_dim]
audio_valid = local[:, visual_width + self.local_audio_dim :]
if audio_valid.shape != (batch, 1):
raise AssertionError("Unexpected pooled AV local feature layout")
audio_projected = self.local_projections["audio"](audio_values) * audio_valid
query = self.person_state_query(states)
query = query + self.person_audio_query(audio_values) * audio_valid
memory = torch.cat(
@@ -494,7 +560,81 @@ class TurnPredictionHead(nn.Module):
need_weights=False,
)
any_person = person_valid.any(dim=1, keepdim=True).to(states.dtype)
return context[:, 0, :] * any_person, any_person
return (
context[:, 0, :] * any_person,
any_person,
audio_projected,
audio_valid,
)
def _encode_person_track_histories(
self,
person_histories: torch.Tensor,
history_mask: torch.Tensor | None,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Return one shared-GRU state per causal anonymous person track."""
if (
person_histories.ndim != 4
or person_histories.shape[2] != self.person_history_frames
or person_histories.shape[3] != self.local_visual_dim
):
raise ValueError(
"Track GRU fusion requires person histories "
"[B,P,person_history_frames,visual_dim]"
)
if (
history_mask is None
or history_mask.dtype != torch.bool
or history_mask.shape != person_histories.shape[:3]
):
raise ValueError(
"Track GRU fusion requires a boolean history mask [B,P,L]"
)
if not torch.isfinite(person_histories).all():
raise ValueError("Person track histories contain NaN/Inf")
if bool((history_mask[:, :, 1:] & ~history_mask[:, :, :-1]).any()):
raise ValueError("Person track histories must be left aligned")
assert self.person_track_gru is not None
assert (
self.person_input_projection is not None
or self.person_track_input_embedding is not None
)
batch, people, length, _ = person_histories.shape
person_valid = history_mask.any(dim=-1)
if not people:
return (
person_histories.new_zeros(
(batch, 0, self.person_attention_dim)
),
person_valid,
)
if self.person_track_input_embedding is not None:
projected = self.person_track_input_embedding(person_histories)
else:
assert self.person_input_projection is not None
projected = self.person_input_projection(person_histories)
flat = projected.reshape(batch * people, length, -1)
lengths = history_mask.sum(dim=-1).reshape(-1)
safe_lengths = lengths.clamp_min(1).to(device="cpu")
packed = nn.utils.rnn.pack_padded_sequence(
flat,
safe_lengths,
batch_first=True,
enforce_sorted=False,
)
initial_hidden = flat.new_zeros(
self.person_track_gru.num_layers,
batch * people,
self.person_track_gru.hidden_size,
)
_, hidden = self.person_track_gru(packed, initial_hidden)
track_states = hidden[-1].reshape(batch, people, -1)
if self.person_track_output_activation is not None:
track_states = self.person_track_output_activation(track_states)
return track_states * person_valid[..., None], person_valid
def _shared_features(
self,
@@ -506,7 +646,23 @@ class TurnPredictionHead(nn.Module):
"""Return the activity MLP representation shared by both readouts."""
states = self._fuse_layers(states)
if self.local_fusion in CURRENT_FRAME_PERSON_SET_FUSION_MODES:
if self.local_fusion in HIGHRATE_PERSON_AV_FUSION_MODES:
# Do not inject person values into the base representation. The
# person set is selected separately in _activity_logits and can
# only alter the prediction through its explicit residual.
if local is None or local.shape != (
len(states),
self.local_audio_dim + 1,
):
raise ValueError(
"High-rate person AV fusion requires aligned "
"[audio,valid] local features"
)
audio_values = local[:, : self.local_audio_dim]
audio_valid = local[:, self.local_audio_dim :]
audio = self.local_projections["audio"](audio_values) * audio_valid
states = torch.cat((states, audio, audio_valid), dim=-1)
elif self.local_fusion in CURRENT_FRAME_PERSON_SET_FUSION_MODES:
if local is None:
raise ValueError("Person-set fusion requires local audio features")
audio, audio_valid, person_context, person_valid = self._frame_set_context(
@@ -569,26 +725,118 @@ class TurnPredictionHead(nn.Module):
person_mask,
)
flat_logits = self.net[-1](features)
if self.local_fusion == "track-gru-residual":
if self.local_fusion in HIGHRATE_PERSON_AV_FUSION_MODES:
if local is None:
raise ValueError("Track GRU residual requires pooled local AV")
assert self.person_track_delta is not None
assert self.person_track_residual_gate is not None
raise ValueError("High-rate person AV fusion requires local audio")
fused_states = self._fuse_layers(states)
context, person_valid = self._track_gru_context(
_, _, context, person_valid = self._frame_set_context(
fused_states,
local,
person_features,
person_mask,
)
delta = self.person_track_delta(
assert self.highrate_person_residual is not None
delta = self.highrate_person_residual(
torch.cat((features, context, person_valid), dim=-1)
)
flat_logits = flat_logits + (
torch.tanh(self.person_track_residual_gate)
* delta
* person_valid
flat_logits = flat_logits + delta * person_valid
with torch.no_grad():
valid_rows = person_valid[:, 0].bool()
diagnostics = {
"valid_person_rate": float(person_valid.mean().cpu()),
"residual_output_weight_norm": float(
self.highrate_person_residual.weight.norm().cpu()
),
"residual_output_bias_norm": float(
self.highrate_person_residual.bias.norm().cpu()
),
}
if bool(valid_rows.any()):
valid_delta = delta[valid_rows]
diagnostics.update(
{
"delta_logit_mean_abs": float(
valid_delta.abs().mean().cpu()
),
"delta_logit_max_abs": float(
valid_delta.abs().max().cpu()
),
}
)
self._last_person_adapter_diagnostics = diagnostics
elif self.local_fusion in TRACK_PERSON_FUSION_MODES:
if local is None:
raise ValueError("Person-track fusion requires pooled local AV")
fused_states = self._fuse_layers(states)
context, person_valid, audio, audio_valid = self._track_gru_context(
fused_states,
local,
person_features,
person_mask,
)
if self.local_fusion == "track-gru-residual":
assert self.person_track_delta is not None
assert self.person_track_residual_gate is not None
delta = self.person_track_delta(
torch.cat((features, context, person_valid), dim=-1)
)
flat_logits = flat_logits + (
torch.tanh(self.person_track_residual_gate)
* delta
* person_valid
)
elif self.local_fusion in PERSON_ADAPTER_V2_MODES:
assert self.person_track_local_state is not None
assert self.person_track_dynamic_gate is not None
assert self.person_track_residual_output is not None
local_state = self.person_track_local_state(
torch.cat(
(context, audio, person_valid, audio_valid),
dim=-1,
)
)
gate = torch.sigmoid(
self.person_track_dynamic_gate(
torch.cat((features, local_state), dim=-1)
)
)
delta = self.person_track_residual_output(gate * local_state)
flat_logits = flat_logits + delta * person_valid
with torch.no_grad():
valid_rows = person_valid[:, 0].bool()
diagnostics = {
"valid_person_rate": float(person_valid.mean().cpu()),
"residual_output_weight_norm": float(
self.person_track_residual_output.weight.norm().cpu()
),
"residual_output_bias_norm": float(
self.person_track_residual_output.bias.norm().cpu()
),
}
if bool(valid_rows.any()):
valid_gate = gate[valid_rows]
valid_delta = delta[valid_rows]
diagnostics.update(
{
"dynamic_gate_mean": float(valid_gate.mean().cpu()),
"dynamic_gate_std": float(
valid_gate.std(unbiased=False).cpu()
),
"dynamic_gate_min": float(valid_gate.min().cpu()),
"dynamic_gate_max": float(valid_gate.max().cpu()),
"delta_logit_mean_abs": float(
valid_delta.abs().mean().cpu()
),
"delta_logit_max_abs": float(
valid_delta.abs().max().cpu()
),
}
)
self._last_person_adapter_diagnostics = diagnostics
else:
raise AssertionError(
f"Unhandled track fusion mode: {self.local_fusion}"
)
return features, flat_logits
def forward(
@@ -660,6 +908,35 @@ class TurnPredictionHead(nn.Module):
)
return result
def person_adapter_diagnostics(self) -> dict[str, float]:
"""Return V2's latest dynamic/residual diagnostics for audit logs."""
if self.local_fusion not in (
*PERSON_ADAPTER_V2_MODES,
*HIGHRATE_PERSON_AV_FUSION_MODES,
):
return {}
result = dict(self._last_person_adapter_diagnostics)
if self.person_track_residual_output is not None:
result.setdefault(
"residual_output_weight_norm",
float(self.person_track_residual_output.weight.detach().norm().cpu()),
)
result.setdefault(
"residual_output_bias_norm",
float(self.person_track_residual_output.bias.detach().norm().cpu()),
)
if self.highrate_person_residual is not None:
result.setdefault(
"residual_output_weight_norm",
float(self.highrate_person_residual.weight.detach().norm().cpu()),
)
result.setdefault(
"residual_output_bias_norm",
float(self.highrate_person_residual.bias.detach().norm().cpu()),
)
return result
def _validate_loss_inputs(
logits: torch.Tensor,
@@ -1297,16 +1574,27 @@ def load_session_data(
raise ValueError(
"Person-set fusion requires --local-modalities av"
)
from local_turn_features import load_frame_set_features
if local_fusion in HIGHRATE_PERSON_AV_FUSION_MODES:
from local_turn_features import load_highrate_person_av_features
local, local_people, contract = load_frame_set_features(
local_path,
payload,
sample_id=name,
dataset=dataset,
expected_contract=local_contract,
oracle_analysis=local_fusion == "oracle-frame-set-attention",
)
local, local_people, contract = load_highrate_person_av_features(
local_path,
payload,
sample_id=name,
dataset=dataset,
expected_contract=local_contract,
)
else:
from local_turn_features import load_frame_set_features
local, local_people, contract = load_frame_set_features(
local_path,
payload,
sample_id=name,
dataset=dataset,
expected_contract=local_contract,
oracle_analysis=local_fusion == "oracle-frame-set-attention",
)
else:
from local_turn_features import load_local_features