Use key-only DeMemWM pose geometry

Disable DeMemWM timestamp embedding by default and route frame-memory pose geometry to memory keys while leaving query features unconditioned.

algorithms/dememwm/df_video.py

@@ -901,7 +901,8 @@ class DeMemWMMinecraft(DiffusionForcingBase):
         self.relative_embedding = getattr(cfg, "relative_embedding", True)
         self.state_embed_only_on_qk = getattr(cfg, "state_embed_only_on_qk", True)
         self.use_memory_attention = getattr(cfg, "use_memory_attention", True)
-        self.add_timestamp_embedding = getattr(cfg, "add_timestamp_embedding", True)
+        self.add_timestamp_embedding = getattr(cfg, "add_timestamp_embedding", False)
+        self.memory_attention_key_only_geometry = getattr(cfg, "memory_attention_key_only_geometry", True)
         self.ref_mode = getattr(cfg, "ref_mode", 'sequential')
         self.log_curve = getattr(cfg, "log_curve", False)
         self.focal_length =  getattr(cfg, "focal_length", 0.35)
@@ -951,6 +952,7 @@ class DeMemWMMinecraft(DiffusionForcingBase):
             add_timestamp_embedding=self.add_timestamp_embedding,
             ref_mode=self.ref_mode,
             focal_length=self.focal_length,
+            memory_attention_key_only_geometry=self.memory_attention_key_only_geometry,
         )
 
         self.validation_lpips_model = LearnedPerceptualImagePatchSimilarity(sync_on_compute=False)

algorithms/dememwm/models/diffusion.py

@@ -31,6 +31,7 @@ class Diffusion(nn.Module):
         add_timestamp_embedding=False,
         ref_mode='sequential',
         focal_length=0.35,
+        memory_attention_key_only_geometry=True,
     ):
         super().__init__()
         self.cfg = cfg
@@ -60,6 +61,7 @@ class Diffusion(nn.Module):
         self.add_timestamp_embedding = add_timestamp_embedding
         self.ref_mode = ref_mode
         self.focal_length = focal_length
+        self.memory_attention_key_only_geometry = memory_attention_key_only_geometry
 
         self._build_model()
         self._build_buffer()
@@ -75,7 +77,8 @@ class Diffusion(nn.Module):
                                             use_memory_attention=self.use_memory_attention,
                                             add_timestamp_embedding=self.add_timestamp_embedding,
                                             ref_mode=self.ref_mode,
-                                            focal_length=self.focal_length)
+                                            focal_length=self.focal_length,
+                                            memory_attention_key_only_geometry=self.memory_attention_key_only_geometry)
         else:
             raise NotImplementedError
 

algorithms/dememwm/models/dit.py

@@ -143,11 +143,12 @@ class FinalLayer(nn.Module):
 
 
 class FrameMemoryReferenceAttention(nn.Module):
-    def __init__(self, hidden_size, num_heads):
+    def __init__(self, hidden_size, num_heads, key_only_geometry=True):
         super().__init__()
         self.hidden_size = hidden_size
         self.num_heads = num_heads
         self.head_dim = hidden_size // num_heads
+        self.key_only_geometry = key_only_geometry
         self.to_q = nn.Linear(hidden_size, hidden_size, bias=False)
         self.to_k = nn.Linear(hidden_size, hidden_size, bias=False)
         self.to_v = nn.Linear(hidden_size, hidden_size, bias=False)
@@ -167,9 +168,13 @@ class FrameMemoryReferenceAttention(nn.Module):
         T_memory = memory_hidden.shape[1]
         P = H * W
         if geometry_cache is not None:
-            query_rays = geometry_cache["query_rays"].to(device=target_hidden.device, dtype=target_hidden.dtype)
-            relative_rays = geometry_cache["relative_rays"].to(device=target_hidden.device, dtype=target_hidden.dtype)
-            target_time = geometry_cache.get("target_time")
+            query_rays = None if self.key_only_geometry else geometry_cache.get("query_rays")
+            if query_rays is not None:
+                query_rays = query_rays.to(device=target_hidden.device, dtype=target_hidden.dtype)
+            relative_rays = geometry_cache.get("relative_rays")
+            if relative_rays is not None:
+                relative_rays = relative_rays.to(device=target_hidden.device, dtype=target_hidden.dtype)
+            target_time = None if self.key_only_geometry else geometry_cache.get("target_time")
             relative_time = geometry_cache.get("relative_time")
         else:
             query_rays = None
@@ -252,6 +257,7 @@ class SpatioTemporalDiTBlock(nn.Module):
         relative_embedding=False,
         state_embed_only_on_qk=False,
         use_memory_attention=False,
+        key_only_geometry=True,
         ref_mode='sequential'
     ):
         super().__init__()
@@ -304,9 +310,9 @@ class SpatioTemporalDiTBlock(nn.Module):
                 drop=0,
             )
             self.r_adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 6 * hidden_size, bias=True))
-            self.r_attn_anchor = FrameMemoryReferenceAttention(hidden_size, num_heads)
-            self.r_attn_dynamic = FrameMemoryReferenceAttention(hidden_size, num_heads)
-            self.r_attn_revisit = FrameMemoryReferenceAttention(hidden_size, num_heads)
+            self.r_attn_anchor = FrameMemoryReferenceAttention(hidden_size, num_heads, key_only_geometry=key_only_geometry)
+            self.r_attn_dynamic = FrameMemoryReferenceAttention(hidden_size, num_heads, key_only_geometry=key_only_geometry)
+            self.r_attn_revisit = FrameMemoryReferenceAttention(hidden_size, num_heads, key_only_geometry=key_only_geometry)
 
         self.reference_length = reference_length
         self.relative_embedding = relative_embedding
@@ -473,6 +479,7 @@ class DiT(nn.Module):
         add_timestamp_embedding=False,
         ref_mode='sequential',
         focal_length=0.35,
+        memory_attention_key_only_geometry=True,
     ):
         super().__init__()
         self.in_channels = in_channels
@@ -522,6 +529,7 @@ class DiT(nn.Module):
                     relative_embedding=relative_embedding,
                     state_embed_only_on_qk=state_embed_only_on_qk,
                     use_memory_attention=use_memory_attention,
+                    key_only_geometry=memory_attention_key_only_geometry,
                     ref_mode=ref_mode
                 )
                 for _ in range(depth)
@@ -792,7 +800,7 @@ class DiT(nn.Module):
 def DiT_S_2(action_cond_dim, pose_cond_dim, reference_length, 
 use_plucker, relative_embedding, 
 state_embed_only_on_qk, use_memory_attention, add_timestamp_embedding,
-ref_mode, focal_length=0.35):
+ref_mode, focal_length=0.35, memory_attention_key_only_geometry=True):
     return DiT(
         patch_size=2,
         hidden_size=1024,
@@ -806,6 +814,7 @@ ref_mode, focal_length=0.35):
         state_embed_only_on_qk=state_embed_only_on_qk,
         use_memory_attention=use_memory_attention,
         add_timestamp_embedding=add_timestamp_embedding,
+        memory_attention_key_only_geometry=memory_attention_key_only_geometry,
         ref_mode=ref_mode,
         focal_length=focal_length,
     )

configurations/algorithm/dememwm_base.yaml

@@ -36,7 +36,8 @@ use_plucker: true
 relative_embedding: true
 state_embed_only_on_qk: true
 use_memory_attention: true
-add_timestamp_embedding: true
+memory_attention_key_only_geometry: true
+add_timestamp_embedding: false
 ref_mode: sequential
 focal_length: 0.35
 log_video: false

tests/test_dememwm_temporal_attention.py

@@ -502,7 +502,7 @@ class DeMemWMTemporalAttentionTests(unittest.TestCase):
         self.assertTrue(torch.allclose(second_only, torch.full_like(second_only, 10.0)))
         self.assertTrue(torch.equal(empty, torch.zeros_like(empty)))
 
-    def test_frame_memory_reference_attention_uses_geometry_qk(self):
+    def test_frame_memory_reference_attention_uses_key_geometry_only(self):
         attn = FrameMemoryReferenceAttention(hidden_size=2, num_heads=1)
         with torch.no_grad():
             attn.to_q.weight.zero_()
@@ -513,7 +513,7 @@ class DeMemWMTemporalAttentionTests(unittest.TestCase):
             attn.to_v.weight[0, 1] = 1.0
             attn.query_pose_proj.weight.zero_()
             attn.query_pose_proj.bias.zero_()
-            attn.query_pose_proj.weight[0, 5] = 1.0
+            attn.query_pose_proj.weight[0, 5] = 100.0
             attn.key_pose_proj.weight.zero_()
             attn.key_pose_proj.bias.zero_()
             attn.key_pose_proj.weight[0, 5] = 1.0
@@ -522,10 +522,12 @@ class DeMemWMTemporalAttentionTests(unittest.TestCase):
             attn.out_proj.bias.zero_()
 
         target = torch.zeros((1, 1, 1, 1, 2))
+        target[0, 0, 0, 0, 0] = 1.0
         memory = torch.zeros((1, 2, 1, 1, 2))
         memory[0, :, 0, 0, 1] = torch.tensor([1.0, 10.0])
         query_rays = torch.zeros((1, 1, 1, 1, 6))
-        query_rays[..., 5] = 1.0
+        query_rays_changed = query_rays.clone()
+        query_rays_changed[..., 5] = 10.0
         relative_a = torch.zeros((1, 1, 2, 1, 1, 6))
         relative_b = torch.zeros_like(relative_a)
         relative_a[0, 0, 0, 0, 0, 5] = 2.0
@@ -536,12 +538,18 @@ class DeMemWMTemporalAttentionTests(unittest.TestCase):
             memory,
             geometry_cache={"query_rays": query_rays, "relative_rays": relative_a},
         )
+        out_a_query_changed = attn(
+            target,
+            memory,
+            geometry_cache={"query_rays": query_rays_changed, "relative_rays": relative_a},
+        )
         out_b = attn(
             target,
             memory,
             geometry_cache={"query_rays": query_rays, "relative_rays": relative_b},
         )
 
+        self.assertTrue(torch.allclose(out_a, out_a_query_changed))
         self.assertLess(out_a[0, 0, 0, 0, 0].item(), out_b[0, 0, 0, 0, 0].item())
         self.assertFalse(torch.allclose(out_a, out_b))