Simplify DeMemWM memory noise modes

algorithms/dememwm/df_video.py

@@ -237,17 +237,6 @@ def _select_online_dynamic_indices(start_frame: int, count: int) -> np.ndarray:
     return np.arange(start, int(start_frame), dtype=np.int64)[-count:]
 
 
-def _shift_memory_noise_levels(levels, shift: float, diffusion_model, bounds):
-    if levels.numel() == 0 or float(shift) == 0.0 or not hasattr(diffusion_model, "snr"):
-        return levels
-
-    log_snr = diffusion_model.snr.to(device=levels.device, dtype=torch.float32).clamp_min(1e-20).log()
-    shifted_log_snr = log_snr[levels.clamp(0, log_snr.numel() - 1)] + float(shift)
-    table = log_snr.view(*((1,) * shifted_log_snr.ndim), -1)
-    shifted = (table - shifted_log_snr.unsqueeze(-1)).abs().argmin(dim=-1).to(dtype=levels.dtype)
-    return torch.minimum(shifted, bounds.reshape(1, -1).expand_as(shifted))
-
-
 def _memory_noise_levels_for_streams(cfg, diffusion_model, query_noise_levels, stream_lengths, mode: str):
     is_training = str(mode) == "training"
     noise_cfg = _cfg_get(cfg, "memory_noise")
@@ -258,15 +247,26 @@ def _memory_noise_levels_for_streams(cfg, diffusion_model, query_noise_levels, s
             for key in _DEMEMWM_STREAM_KEYS
         }
 
-    if str(_cfg_get(noise_cfg, "mode", "random_cleaner_fraction")) != "random_cleaner_fraction":
-        raise ValueError("memory_noise.mode must be random_cleaner_fraction")
+    noise_mode = str(_cfg_get(noise_cfg, "mode", "random_cleaner_fraction"))
+    if noise_mode not in {"random_cleaner_fraction", "independent"}:
+        raise ValueError("memory_noise.mode must be random_cleaner_fraction or independent")
     noisy_memory = bool(_cfg_get(noise_cfg, "enabled", True)) and (
         is_training or bool(_cfg_get(noise_cfg, "validation_noisy_memory", False))
     )
-    snr_cfg = _cfg_get(noise_cfg, "snr_shift")
-    snr_enabled = is_training and noisy_memory and bool(_cfg_get(snr_cfg, "enabled", False))
     query = query_noise_levels[:, None] if query_noise_levels.ndim == 1 else query_noise_levels
-    bounds = query.to(dtype=torch.long).clamp_min(0).min(dim=0).values
+    query_bounds = query.to(dtype=torch.long).clamp_min(0).min(dim=0).values
+    if noise_mode == "independent":
+        if is_training:
+            max_level = int(getattr(diffusion_model, "timesteps", 0) or 0) - 1
+        else:
+            max_level = int(
+                getattr(diffusion_model, "sampling_timesteps", 0)
+                or getattr(diffusion_model, "timesteps", 0)
+                or 1
+            )
+        bounds = query_bounds.new_full(query_bounds.shape, max(0, max_level))
+    else:
+        bounds = query_bounds
     defaults = {"anchor": 0.25, "dynamic": 0.50, "revisit": 0.25}
     levels_by_stream = {}
     for key in _DEMEMWM_STREAM_KEYS:
@@ -274,14 +274,25 @@ def _memory_noise_levels_for_streams(cfg, diffusion_model, query_noise_levels, s
         if count == 0:
             levels_by_stream[key] = bounds.new_zeros((0, bounds.numel()))
             continue
-        fraction = float(_cfg_get(noise_cfg, f"{key}_max_fraction", defaults[key])) if noisy_memory else 0.0
-        max_levels = torch.floor(bounds.to(dtype=torch.float32) * fraction).to(dtype=torch.long).clamp_min(0)
-        if is_training and noisy_memory:
+        if noise_mode == "independent":
+            max_levels = bounds if noisy_memory else bounds.new_zeros(bounds.shape)
+        else:
+            fraction = float(_cfg_get(noise_cfg, f"{key}_max_fraction", defaults[key])) if noisy_memory else 0.0
+            max_levels = torch.floor(bounds.to(dtype=torch.float32) * fraction).to(dtype=torch.long).clamp_min(0)
+        if is_training and noisy_memory and noise_mode == "independent":
+            levels = torch.randint(
+                0,
+                int(bounds.max().item()) + 1,
+                (count, bounds.numel()),
+                device=bounds.device,
+                dtype=torch.long,
+            )
+            levels = torch.minimum(levels, max_levels.reshape(1, -1).expand_as(levels))
+        elif is_training and noisy_memory:
             levels = torch.floor(torch.rand((count, bounds.numel()), device=bounds.device) * (max_levels[None].float() + 1)).to(dtype=torch.long)
         else:
             levels = max_levels.reshape(1, -1).expand(count, -1).contiguous()
-        shift = float(_cfg_get(snr_cfg, key, 0.0)) if snr_enabled else 0.0
-        levels_by_stream[key] = _shift_memory_noise_levels(levels, shift, diffusion_model, bounds)
+        levels_by_stream[key] = levels
     return levels_by_stream
 
 

algorithms/dememwm/models/dit.py

@@ -171,7 +171,8 @@ class FrameMemoryReferenceAttention(nn.Module):
             relative_rays = geometry_cache["relative_rays"].to(device=target_hidden.device, dtype=target_hidden.dtype)
         else:
             query_rays = None
-        
+            relative_rays = None
+
         q_input = target_hidden
         if query_rays is not None:
             # Values stay visual-only; Plucker features only steer Q/K matching.

configurations/algorithm/dememwm_base.yaml

@@ -8,11 +8,11 @@ memory_selection: ${dataset.memory_selection}
 
 memory_noise:
   enabled: true
+  # Options: random_cleaner_fraction (query-bounded), independent (full schedule).
   mode: random_cleaner_fraction
   anchor_max_fraction: 0.25
   dynamic_max_fraction: 0.50
   revisit_max_fraction: 0.25
-  snr_shift: {enabled: true, anchor: 0.0, dynamic: 0.0, revisit: 0.0}
 noise_route: {anchor: all, dynamic: all, revisit: all}
 trainability:
   freeze_vae: true

tests/test_dememwm_latent_dataset.py

@@ -436,15 +436,14 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
         with self.assertRaisesRegex(TypeError, "latent dict batch contract"):
             DeMemWMMinecraft.training_step(Harness(), (None, None, None, None), 0)
 
-    def test_memory_noise_helpers_route_shift_and_keep_validation_clean(self):
+    def test_memory_noise_helpers_route_and_keep_validation_clean(self):
         from algorithms.dememwm.df_video import _apply_memory_route_masks, _memory_noise_levels_for_streams
 
         class Diffusion:
             stabilization_level, timesteps, sampling_timesteps = 7, 10, 4
-            snr = torch.exp(torch.linspace(5.0, -4.0, 10))
 
         cfg = OmegaConf.create({
-            "memory_noise": {"enabled": True, "anchor_max_fraction": 1.0, "dynamic_max_fraction": 0.5, "revisit_max_fraction": 0.25, "snr_shift": {"enabled": True, "anchor": 100.0}},
+            "memory_noise": {"enabled": True, "anchor_max_fraction": 1.0, "dynamic_max_fraction": 0.5, "revisit_max_fraction": 0.25},
             "noise_route": {"anchor": "high", "dynamic": "low", "revisit": "all"},
         })
         query_noise = torch.tensor([[9, 1], [7, 1]], dtype=torch.long)
@@ -452,7 +451,7 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
         torch.manual_seed(0)
 
         levels = _memory_noise_levels_for_streams(cfg, Diffusion, query_noise, stream_lengths, mode="training")
-        self.assertEqual(levels["anchor"].tolist(), [[0, 0]])
+        self.assertTrue(bool((levels["anchor"] <= torch.tensor([[7, 1]])).all()))
         self.assertTrue(bool((levels["dynamic"] <= torch.tensor([[3, 0]])).all()))
         self.assertTrue(bool((levels["revisit"] <= torch.tensor([[1, 0]])).all()))
         validation_levels = _memory_noise_levels_for_streams(cfg, Diffusion, query_noise, stream_lengths, mode="validation")
@@ -462,6 +461,34 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
         routed = _apply_memory_route_masks(masks, query_noise, cfg, Diffusion, mode="training")
         self.assertEqual({key: routed[key].tolist() for key in stream_lengths}, {"anchor": [[True], [False]], "dynamic": [[False], [True]], "revisit": [[True], [True]]})
 
+    def test_memory_noise_independent_mode_ignores_query_bounds(self):
+        from algorithms.dememwm.df_video import _memory_noise_levels_for_streams
+
+        class Diffusion:
+            stabilization_level, timesteps, sampling_timesteps = 7, 10, 4
+
+        cfg = OmegaConf.create({
+            "memory_noise": {
+                "enabled": True,
+                "mode": "independent",
+            },
+        })
+        query_noise = torch.ones((2, 2), dtype=torch.long)
+        stream_lengths = {"anchor": 3, "dynamic": 0, "revisit": 0}
+
+        def fake_randint(low, high, size, device=None, dtype=None):
+            self.assertEqual((low, high, size), (0, 10, (3, 2)))
+            return torch.full(size, high - 1, device=device, dtype=dtype)
+
+        with mock.patch("algorithms.dememwm.df_video.torch.randint", side_effect=fake_randint):
+            levels = _memory_noise_levels_for_streams(cfg, Diffusion, query_noise, stream_lengths, mode="training")
+
+        self.assertEqual(levels["anchor"].tolist(), [[9, 9], [9, 9], [9, 9]])
+        self.assertTrue(bool((levels["anchor"] > query_noise.max()).all()))
+
+        validation_levels = _memory_noise_levels_for_streams(cfg, Diffusion, query_noise, stream_lengths, mode="validation")
+        self.assertEqual(sum(int(v.sum().item()) for v in validation_levels.values()), 0)
+
     def test_trainability_controls_memory_groups_full_dit_ramp_and_vae_freeze(self):
         from algorithms.dememwm.df_video import (
             _apply_dememwm_trainability,