Add DeMemWM segment-aware temporal mask

algorithms/dememwm/models/attention.py

@@ -44,7 +44,56 @@ class TemporalAxialAttention(nn.Module):
 
         self.reference_length = reference_length
 
-    def forward(self, x: torch.Tensor):
+    def _frame_memory_attn_bias(
+        self,
+        B,
+        T,
+        H,
+        W,
+        dtype,
+        device,
+        frame_memory_segments,
+        frame_memory_masks,
+    ):
+        allow = torch.zeros((B, T, T), dtype=torch.bool, device=device)
+
+        target_frames = frame_memory_segments["target"] if frame_memory_segments else T
+        cursor = 0
+        target_slice = slice(cursor, cursor + target_frames)
+        target_idx = torch.arange(target_frames, device=device)
+        allow[:, target_slice, target_slice] = target_idx[:, None] >= target_idx[None, :]
+        cursor += target_frames
+
+        segment_slices = {"target": target_slice}
+        for segment in ("anchor", "dynamic", "revisit"):
+            length = frame_memory_segments.get(segment, 0)
+            segment_slice = slice(cursor, cursor + length)
+            segment_slices[segment] = segment_slice
+            if length > 0:
+                if segment == "dynamic":
+                    allow[:, segment_slice, segment_slice] = True
+                else:
+                    idx = torch.arange(cursor, cursor + length, device=device)
+                    allow[:, idx, idx] = True
+            cursor += length
+
+        if frame_memory_masks is not None:
+            valid = torch.ones((B, T), dtype=torch.bool, device=device)
+            for segment, segment_slice in segment_slices.items():
+                mask = frame_memory_masks.get(segment)
+                if mask is not None:
+                    valid[:, segment_slice] = mask.to(device=device, dtype=torch.bool)
+            allow = allow & valid[:, :, None] & valid[:, None, :]
+
+        # Invalid padded rows keep only self-attention finite so SDPA never sees
+        # an all -inf query row.
+        diag_idx = torch.arange(T, device=device)
+        allow[:, diag_idx, diag_idx] = True
+        attn_bias = torch.zeros((B, T, T), dtype=dtype, device=device)
+        attn_bias = attn_bias.masked_fill(~allow, float("-inf"))
+        return attn_bias.repeat_interleave(H * W, dim=0)[:, None]
+
+    def forward(self, x: torch.Tensor, frame_memory_segments=None, frame_memory_masks=None):
         B, T, H, W, D = x.shape
 
         # if T>=9:
@@ -75,7 +124,18 @@ class TemporalAxialAttention(nn.Module):
 
         q, k, v = map(lambda t: t.contiguous(), (q, k, v))
 
-        if self.is_temporal_independent:
+        if frame_memory_segments is not None:
+            attn_bias = self._frame_memory_attn_bias(
+                B,
+                T,
+                H,
+                W,
+                q.dtype,
+                q.device,
+                frame_memory_segments,
+                frame_memory_masks,
+            )
+        elif self.is_temporal_independent:
             attn_bias = torch.ones((T, T), dtype=q.dtype, device=q.device)
             attn_bias = attn_bias.masked_fill(attn_bias == 1, float('-inf'))
             attn_bias[range(T), range(T)] = 0

algorithms/dememwm/models/dit.py

@@ -241,7 +241,8 @@ class SpatioTemporalDiTBlock(nn.Module):
             self.parallel_map = nn.Linear(hidden_size, hidden_size)
 
     def forward(self, x, c, current_frame=None, timestep=None, is_last_block=False, 
-        pose_cond=None, mode="training", c_action_cond=None, reference_length=None):
+        pose_cond=None, mode="training", c_action_cond=None, reference_length=None,
+        frame_memory_segments=None, frame_memory_masks=None):
         B, T, H, W, D = x.shape
 
         # spatial block
@@ -256,7 +257,11 @@ class SpatioTemporalDiTBlock(nn.Module):
         else:
             t_shift_msa, t_scale_msa, t_gate_msa, t_shift_mlp, t_scale_mlp, t_gate_mlp = self.t_adaLN_modulation(c).chunk(6, dim=-1)
         
-        x_t = x + gate(self.t_attn(modulate(self.t_norm1(x), t_shift_msa, t_scale_msa)), t_gate_msa)
+        x_t = x + gate(self.t_attn(
+            modulate(self.t_norm1(x), t_shift_msa, t_scale_msa),
+            frame_memory_segments=frame_memory_segments,
+            frame_memory_masks=frame_memory_masks,
+        ), t_gate_msa)
         x_t = x_t + gate(self.t_mlp(modulate(self.t_norm2(x_t), t_shift_mlp, t_scale_mlp)), t_gate_mlp)
 
         if self.ref_mode == 'sequential':
@@ -540,7 +545,9 @@ class DiT(nn.Module):
         
         for i, block in enumerate(self.blocks):
             x = block(x, c, current_frame=current_frame, timestep=t, is_last_block= (i+1 == len(self.blocks)), 
-                pose_cond=pc, mode=mode, c_action_cond=c_action_cond, reference_length=reference_length)  # (N, T, H, W, D)
+                pose_cond=pc, mode=mode, c_action_cond=c_action_cond, reference_length=reference_length,
+                frame_memory_segments=frame_memory_segments,
+                frame_memory_masks=frame_memory_masks)  # (N, T, H, W, D)
         x = self.final_layer(x, c) # (N, T, H, W, patch_size ** 2 * out_channels)
         # unpatchify
         x = rearrange(x, "b t h w d -> (b t) h w d")

tests/test_dememwm_temporal_attention.py

@@ -0,0 +1,139 @@
+import unittest
+
+import torch
+from torch import nn
+
+from algorithms.dememwm.models.attention import TemporalAxialAttention
+from algorithms.dememwm.models.dit import DiT, SpatioTemporalDiTBlock
+
+
+class IdentityRotary:
+    freqs = None
+
+    def rotate_queries_or_keys(self, x, freqs):
+        return x
+
+
+def _averaging_temporal_attention(reference_length=0):
+    attn = TemporalAxialAttention(
+        dim=1,
+        heads=1,
+        dim_head=1,
+        reference_length=reference_length,
+        rotary_emb=IdentityRotary(),
+    )
+    with torch.no_grad():
+        attn.to_qkv.weight.zero_()
+        attn.to_qkv.weight[2, 0] = 1.0
+        attn.to_out.weight.fill_(1.0)
+        attn.to_out.bias.zero_()
+    return attn
+
+
+class DeMemWMTemporalAttentionTests(unittest.TestCase):
+    def test_without_frame_memory_keeps_causal_reference_mask(self):
+        attn = _averaging_temporal_attention(reference_length=2)
+        x = torch.tensor([10.0, 20.0, 30.0, 100.0, 200.0]).view(1, 5, 1, 1, 1)
+
+        out = attn(x)
+
+        expected = torch.tensor([10.0, 15.0, 20.0, 100.0, 200.0])
+        self.assertTrue(torch.allclose(out.flatten(), expected))
+
+    def test_frame_memory_segments_mask_temporal_streams(self):
+        attn = _averaging_temporal_attention()
+        x = torch.tensor([10.0, 20.0, 30.0, 100.0, 1.0, 3.0, 200.0]).view(1, 7, 1, 1, 1)
+        segments = {"target": 3, "anchor": 1, "dynamic": 2, "revisit": 1}
+
+        out = attn(x, frame_memory_segments=segments)
+
+        expected = torch.tensor([10.0, 15.0, 20.0, 100.0, 2.0, 2.0, 200.0])
+        self.assertTrue(torch.allclose(out.flatten(), expected))
+
+        masks = {
+            "target": torch.ones((1, 3), dtype=torch.bool),
+            "anchor": torch.tensor([[False]]),
+            "dynamic": torch.tensor([[True, False]]),
+            "revisit": torch.ones((1, 1), dtype=torch.bool),
+        }
+        out = attn(x, frame_memory_segments=segments, frame_memory_masks=masks)
+
+        expected = torch.tensor([10.0, 15.0, 20.0, 100.0, 1.0, 3.0, 200.0])
+        self.assertTrue(torch.allclose(out.flatten(), expected))
+
+    def test_block_threads_frame_memory_metadata_to_temporal_attention(self):
+        class ZeroAttention(nn.Module):
+            def forward(self, x):
+                return torch.zeros_like(x)
+
+        class SpyTemporalAttention(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.calls = []
+
+            def forward(self, x, frame_memory_segments=None, frame_memory_masks=None):
+                self.calls.append((frame_memory_segments, frame_memory_masks))
+                return torch.zeros_like(x)
+
+        block = SpatioTemporalDiTBlock(
+            hidden_size=4,
+            num_heads=1,
+            reference_length=0,
+            spatial_rotary_emb=None,
+            temporal_rotary_emb=None,
+        )
+        spy = SpyTemporalAttention()
+        block.s_attn = ZeroAttention()
+        block.t_attn = spy
+        x = torch.zeros((1, 5, 1, 1, 4))
+        c = torch.zeros((1, 5, 4))
+        segments = {"target": 2, "anchor": 1, "dynamic": 1, "revisit": 1}
+        masks = {"dynamic": torch.ones((1, 1), dtype=torch.bool)}
+
+        block(x, c, frame_memory_segments=segments, frame_memory_masks=masks)
+
+        self.assertIs(spy.calls[0][0], segments)
+        self.assertIs(spy.calls[0][1], masks)
+
+    def test_dit_threads_frame_memory_metadata_to_blocks(self):
+        class SpyBlock(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.calls = []
+
+            def forward(self, x, c, **kwargs):
+                self.calls.append(kwargs)
+                return x
+
+        model = DiT(
+            input_h=2,
+            input_w=2,
+            patch_size=1,
+            in_channels=1,
+            hidden_size=8,
+            depth=1,
+            num_heads=1,
+            mlp_ratio=1.0,
+            action_cond_dim=3,
+            pose_cond_dim=0,
+            reference_length=0,
+            use_memory_attention=False,
+        )
+        spy = SpyBlock()
+        model.blocks[0] = spy
+        model.eval()
+        x = torch.zeros((1, 5, 1, 2, 2))
+        t = torch.zeros((1, 5), dtype=torch.long)
+        action_cond = torch.zeros((1, 5, 3))
+        segments = {"target": 2, "anchor": 1, "dynamic": 1, "revisit": 1}
+        masks = {"dynamic": torch.ones((1, 1), dtype=torch.bool)}
+
+        with torch.no_grad():
+            model(x, t, action_cond, frame_memory_segments=segments, frame_memory_masks=masks)
+
+        self.assertIs(spy.calls[0]["frame_memory_segments"], segments)
+        self.assertIs(spy.calls[0]["frame_memory_masks"], masks)
+
+
+if __name__ == "__main__":
+    unittest.main()