Skip to content

Fix 3D attn_mask head/batch layout in T2V/V2T cross-attention (repeat → repeat_interleave) - #77

Merged
h-munakata merged 1 commit into
line:mainfrom
komatta-san:fix/t2v-attn-mask-repeat-interleave
Aug 30, 2026
Merged

Fix 3D attn_mask head/batch layout in T2V/V2T cross-attention (repeat → repeat_interleave)#77
h-munakata merged 1 commit into
line:mainfrom
komatta-san:fix/t2v-attn-mask-repeat-interleave

Conversation

@komatta-san

Copy link
Copy Markdown
Contributor

Problem

nn.MultiheadAttention expects a 3D attn_mask of shape (N * num_heads, L, S) indexed as batch_idx * num_heads + head_idx. The mask is built with .repeat(nhead, 1, 1), which tiles instead: [M0, M1, M0, M1] rather than the expected [M0, M0, M1, M1]. Heads then attend using the wrong sample's padding mask.

Impact

When num_heads > 1 and samples have different valid lengths (the normal case), a sample's output depends on the other samples in the batch, breaking determinism. At batch_size == 1 the two ops are identical, so single-sample outputs are unchanged.

Fix

Use .repeat_interleave(self.nhead, dim=0) in all affected encoder layers: CIM.py (T2V+V2T), cg_detr_, qd_detr_, taskweave_, tr_detr_transformer.py.

@awkrail

awkrail commented Aug 30, 2026

Copy link
Copy Markdown
Contributor

lgtm. @h-munakata Could you merge it?

@h-munakata
h-munakata merged commit 91f8fe9 into line:main Aug 30, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants