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

feat: add person fusion screening variants

Browse Source
This commit is contained in:
Xin Wang
2026-09-06 21:48:32 +08:00
parent 8e70281330
commit b91a24f308
8 changed files with 447 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files

136
easycom_head.py
View File

@@ -61,7 +61,18 @@ LAYER_FUSION_MODES = (
"scalar-mixture",
"concat-mlp",
)
LOCAL_FUSION_MODES = ("pooled", "frame-set-attention")
PERSON_SET_FUSION_MODES = (
"frame-set-attention",
"audio-person-frame-set",
"gated-person-state-adapter",
)
LOCAL_FUSION_MODES = ("pooled", *PERSON_SET_FUSION_MODES)
LOCAL_FUSION_SCOPES = {
"pooled": "pooled-baseline",
"frame-set-attention": "f1-current-frame-person-set",
"audio-person-frame-set": "f5-screening-current-frame-no-person-history",
"gated-person-state-adapter": "f6-screening-post-thinker-not-internal",
}
DatasetName = Literal["easycom", "ego4d"]
@@ -132,11 +143,11 @@ class TurnPredictionHead(nn.Module):
raise ValueError("Person attention dimension/heads must be positive")
if person_attention_dim % person_attention_heads:
raise ValueError("Person attention dimension must be divisible by heads")
if local_fusion == "frame-set-attention" and (
if local_fusion in PERSON_SET_FUSION_MODES and (
local_visual_dim <= 0 or local_audio_dim <= 0
):
raise ValueError(
"Frame-set Attention requires both local visual people and audio"
"Person-set fusion requires both local visual people and audio"
)
self.local_fusion = local_fusion
self.person_attention_dim = person_attention_dim
@@ -160,7 +171,7 @@ class TurnPredictionHead(nn.Module):
fused_size = hidden_size
projected_modalities = (
(("audio", local_audio_dim),)
if local_fusion == "frame-set-attention"
if local_fusion in PERSON_SET_FUSION_MODES
else (("visual", local_visual_dim), ("audio", local_audio_dim))
)
for name, dimension in projected_modalities:
@@ -180,7 +191,11 @@ class TurnPredictionHead(nn.Module):
self.person_audio_query: nn.Module | None = None
self.person_attention: nn.MultiheadAttention | None = None
self.null_person: nn.Parameter | None = None
if local_fusion == "frame-set-attention":
self.audio_person_attention: nn.MultiheadAttention | None = None
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
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),
@@ -204,7 +219,33 @@ class TurnPredictionHead(nn.Module):
torch.zeros(1, 1, person_attention_dim)
)
nn.init.normal_(self.null_person, std=0.02)
fused_size += person_attention_dim + 1
if local_fusion == "audio-person-frame-set":
# Screening version of the F5 reverse read: an audio-only
# query first reads the current person memory, then a
# zero-initialized residual gate adds that evidence to the
# regular G+A floor query. This remains a current-frame
# module; it is not the full history-conditioned F5 block.
self.audio_person_attention = nn.MultiheadAttention(
person_attention_dim,
person_attention_heads,
dropout=dropout,
batch_first=True,
)
self.audio_person_gate = nn.Parameter(torch.zeros(()))
if local_fusion == "gated-person-state-adapter":
# Screening version of F6 for precomputed frozen states. It
# injects selected person evidence after the Thinker readout,
# not into a Thinker transformer block. The scalar gate makes
# the visual residual exactly zero at initialization.
self.person_state_adapter = nn.Sequential(
nn.LayerNorm(person_attention_dim),
nn.Linear(person_attention_dim, person_attention_dim),
nn.GELU(),
nn.Linear(person_attention_dim, hidden_size),
)
self.person_state_adapter_gate = nn.Parameter(torch.zeros(()))
else:
fused_size += person_attention_dim + 1
self.net = nn.Sequential(
nn.LayerNorm(fused_size),
nn.Linear(fused_size, hidden_size // 2),
@@ -250,7 +291,7 @@ class TurnPredictionHead(nn.Module):
@property
def local_input_size(self) -> int:
if self.local_fusion == "frame-set-attention":
if self.local_fusion in PERSON_SET_FUSION_MODES:
return self.local_audio_dim + 1
return sum(d + 1 for d in (self.local_visual_dim, self.local_audio_dim) if d)
@@ -260,12 +301,12 @@ class TurnPredictionHead(nn.Module):
local: torch.Tensor,
person_features: torch.Tensor | None,
person_mask: torch.Tensor | None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
if states.ndim != 2:
raise ValueError("Frame-set Attention expects batched [B,D] states")
raise ValueError("Person-set fusion expects batched [B,D] states")
if local.shape != (len(states), self.local_audio_dim + 1):
raise ValueError(
"Frame-set Attention requires aligned [audio,valid] local features"
"Person-set fusion requires aligned [audio,valid] local features"
)
if (
person_features is None
@@ -274,14 +315,14 @@ class TurnPredictionHead(nn.Module):
or person_features.shape[2] != self.local_visual_dim
):
raise ValueError(
"Frame-set Attention requires person features [B,P,visual_dim]"
"Person-set fusion requires person features [B,P,visual_dim]"
)
if (
person_mask is None
or person_mask.dtype != torch.bool
or person_mask.shape != person_features.shape[:2]
):
raise ValueError("Frame-set Attention requires a boolean person mask [B,P]")
raise ValueError("Person-set fusion requires a boolean person mask [B,P]")
if not torch.isfinite(person_features).all():
raise ValueError("Person features contain NaN/Inf")
assert self.person_input_projection is not None
@@ -293,10 +334,8 @@ class TurnPredictionHead(nn.Module):
audio_values = local[:, : self.local_audio_dim]
audio_valid = local[:, self.local_audio_dim : self.local_audio_dim + 1]
audio_projected = self.local_projections["audio"](audio_values) * audio_valid
query = (
self.person_state_query(states)
+ self.person_audio_query(audio_values) * audio_valid
)
audio_query = self.person_audio_query(audio_values) * audio_valid
query = self.person_state_query(states) + audio_query
people = self.person_input_projection(person_features)
null = self.null_person.expand(len(states), -1, -1)
memory = torch.cat((null, people), dim=1)
@@ -306,6 +345,20 @@ class TurnPredictionHead(nn.Module):
device=person_mask.device,
)
memory_valid = torch.cat((null_mask, person_mask), dim=1)
if self.audio_person_attention is not None:
assert self.audio_person_gate is not None
audio_context, _ = self.audio_person_attention(
audio_query[:, None, :],
memory,
memory,
key_padding_mask=~memory_valid,
need_weights=False,
)
query = query + (
torch.tanh(self.audio_person_gate)
* audio_context[:, 0, :]
* audio_valid
)
context, _ = self.person_attention(
query[:, None, :],
memory,
@@ -314,10 +367,7 @@ class TurnPredictionHead(nn.Module):
need_weights=False,
)
person_valid = person_mask.any(dim=1, keepdim=True).to(states.dtype)
return audio_projected, audio_valid, torch.cat(
(context[:, 0, :], person_valid),
dim=-1,
)
return audio_projected, audio_valid, context[:, 0, :], person_valid
def _shared_features(
self,
@@ -329,19 +379,29 @@ class TurnPredictionHead(nn.Module):
"""Return the activity MLP representation shared by both readouts."""
states = self._fuse_layers(states)
if self.local_fusion == "frame-set-attention":
if self.local_fusion in PERSON_SET_FUSION_MODES:
if local is None:
raise ValueError("Frame-set Attention requires local audio features")
audio, audio_valid, person_context = self._frame_set_context(
raise ValueError("Person-set fusion requires local audio features")
audio, audio_valid, person_context, person_valid = self._frame_set_context(
states,
local,
person_features,
person_mask,
)
states = torch.cat(
(states, audio, audio_valid, person_context),
dim=-1,
)
if self.local_fusion == "gated-person-state-adapter":
assert self.person_state_adapter is not None
assert self.person_state_adapter_gate is not None
states = states + (
torch.tanh(self.person_state_adapter_gate)
* self.person_state_adapter(person_context)
* person_valid
)
states = torch.cat((states, audio, audio_valid), dim=-1)
else:
states = torch.cat(
(states, audio, audio_valid, person_context, person_valid),
dim=-1,
)
elif self.local_input_size:
if local is None or local.shape != (*states.shape[:-1], self.local_input_size):
raise ValueError(
@@ -362,7 +422,7 @@ class TurnPredictionHead(nn.Module):
states = torch.cat(pieces, dim=-1)
elif local is not None and local.shape[-1] != 0:
raise ValueError("Local features supplied to a baseline head")
if self.local_fusion != "frame-set-attention" and (
if self.local_fusion not in PERSON_SET_FUSION_MODES and (
person_features is not None or person_mask is not None
):
raise ValueError("Person sets supplied to a non-frame-set head")
@@ -420,6 +480,20 @@ class TurnPredictionHead(nn.Module):
return None
return self.layer_logits.softmax(dim=0)
def person_interaction_gates(self) -> dict[str, float]:
"""Return bounded residual strengths for screening fusion modes."""
result: dict[str, float] = {}
if self.audio_person_gate is not None:
result["audio_person"] = float(
torch.tanh(self.audio_person_gate.detach()).cpu()
)
if self.person_state_adapter_gate is not None:
result["person_state_adapter"] = float(
torch.tanh(self.person_state_adapter_gate.detach()).cpu()
)
return result
def _validate_loss_inputs(
logits: torch.Tensor,
@@ -1036,10 +1110,10 @@ def load_session_data(
name,
dataset=dataset,
)
if local_fusion == "frame-set-attention":
if local_fusion in PERSON_SET_FUSION_MODES:
if local_modalities != "av":
raise ValueError(
"Frame-set Attention requires --local-modalities av"
"Person-set fusion requires --local-modalities av"
)
from local_turn_features import load_frame_set_features
@@ -1065,7 +1139,7 @@ def load_session_data(
elif local_contract is not None:
raise ValueError("--local-feature-root is required for this checkpoint")
elif local_fusion != "pooled":
raise ValueError("Frame-set Attention requires --local-feature-root")
raise ValueError("Person-set fusion requires --local-feature-root")
return SessionData(
name,
states,