Add DeMemWM trainability controls

algorithms/dememwm/df_video.py

@@ -28,6 +28,9 @@ import glob
 # Utility Functions
 _DEMEMWM_SEGMENT_KEYS = ("target", "anchor", "dynamic", "revisit")
 _DEMEMWM_STREAM_KEYS = ("anchor", "dynamic", "revisit")
+_DEMEMWM_REFERENCE_ATTN_MARKERS = (".r_attn_anchor.", ".r_attn_dynamic.", ".r_attn_revisit.")
+_DEMEMWM_GEOMETRY_PROJ_MARKERS = (".query_pose_proj.", ".key_pose_proj.")
+_DEMEMWM_ADALN_MLP_MARKERS = (".r_adaLN_modulation.", ".r_mlp.")
 
 
 def _cfg_get(cfg, key: str, default=None):
@@ -38,6 +41,50 @@ def _cfg_get(cfg, key: str, default=None):
     return getattr(cfg, key, default)
 
 
+def _trainability_cfg(cfg):
+    return _cfg_get(cfg, "trainability", {})
+
+
+def _dememwm_full_dit_active(trainability, global_step: int) -> bool:
+    if not bool(_cfg_get(trainability, "train_full_dit", False)):
+        return False
+    start_step = _cfg_get(trainability, "full_dit_start_step", 0)
+    return start_step is None or int(global_step) >= int(start_step)
+
+
+def _is_dememwm_memory_trainable_parameter(name: str, trainability) -> bool:
+    is_geometry = any(marker in name for marker in _DEMEMWM_GEOMETRY_PROJ_MARKERS)
+    if bool(_cfg_get(trainability, "geometry_projections", True)) and is_geometry:
+        return True
+    if bool(_cfg_get(trainability, "reference_attention", True)) and not is_geometry:
+        if any(marker in name for marker in _DEMEMWM_REFERENCE_ATTN_MARKERS):
+            return True
+    if bool(_cfg_get(trainability, "adaln_mlp", True)):
+        return any(marker in name for marker in _DEMEMWM_ADALN_MLP_MARKERS)
+    return False
+
+
+def _apply_dememwm_trainability(diffusion_model, vae, trainability, global_step: int) -> None:
+    full_dit_active = _dememwm_full_dit_active(trainability, global_step)
+    for name, param in diffusion_model.named_parameters():
+        param.requires_grad_(full_dit_active or _is_dememwm_memory_trainable_parameter(name, trainability))
+
+    if vae is not None:
+        freeze_vae = bool(_cfg_get(trainability, "freeze_vae", True))
+        for param in vae.parameters():
+            param.requires_grad_(not freeze_vae)
+
+
+def _dememwm_optimizer_parameters(diffusion_model, vae, trainability):
+    include_full_dit = bool(_cfg_get(trainability, "train_full_dit", False))
+    for name, param in diffusion_model.named_parameters():
+        if include_full_dit or _is_dememwm_memory_trainable_parameter(name, trainability):
+            yield param
+
+    if vae is not None and not bool(_cfg_get(trainability, "freeze_vae", True)):
+        yield from vae.parameters()
+
+
 def _derive_memory_condition_length(cfg) -> int:
     memory_cfg = _cfg_get(cfg, "memory_selection")
     if memory_cfg is None:
@@ -609,6 +656,36 @@ class DeMemWMMinecraft(DiffusionForcingBase):
         if self.require_pose_prediction:
             self.pose_prediction_model = PosePredictionNet()
 
+        self._apply_trainability()
+
+    def _global_step_for_trainability(self) -> int:
+        try:
+            trainer = self.trainer
+        except RuntimeError:
+            return 0
+        return int(getattr(trainer, "global_step", 0) or 0)
+
+    def _apply_trainability(self) -> None:
+        _apply_dememwm_trainability(
+            self.diffusion_model,
+            getattr(self, "vae", None),
+            _trainability_cfg(self.cfg),
+            self._global_step_for_trainability(),
+        )
+
+    def configure_optimizers(self):
+        trainability = _trainability_cfg(self.cfg)
+        self._apply_trainability()
+        params = tuple(_dememwm_optimizer_parameters(self.diffusion_model, getattr(self, "vae", None), trainability))
+        if not params:
+            raise ValueError("DeMemWM trainability selected no optimizer parameters")
+        return torch.optim.AdamW(
+            params, lr=self.cfg.lr, weight_decay=self.cfg.weight_decay, betas=self.cfg.optimizer_beta
+        )
+
+    def on_train_batch_start(self, batch, batch_idx, dataloader_idx=0) -> None:
+        self._apply_trainability()
+
     def _generate_noise_levels(self, xs: torch.Tensor, masks = None) -> torch.Tensor:
         """
         Generate noise levels for training.

configurations/algorithm/dememwm_base.yaml

@@ -14,6 +14,13 @@ memory_noise:
   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
+  reference_attention: true
+  adaln_mlp: true
+  geometry_projections: true
+  train_full_dit: false
+  full_dit_start_step: null
 
 n_tokens: ${dataset.n_frames}
 action_cond_dim: ${dataset.action_cond_dim}

tests/test_dememwm_latent_dataset.py

@@ -5,6 +5,7 @@ from pathlib import Path
 
 import numpy as np
 import torch
+from torch import nn
 from omegaconf import OmegaConf
 
 from datasets.video.memory_selection import select_memory_indices
@@ -327,6 +328,64 @@ 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_trainability_controls_memory_groups_full_dit_ramp_and_vae_freeze(self):
+        from algorithms.dememwm.df_video import (
+            _apply_dememwm_trainability,
+            _dememwm_optimizer_parameters,
+            _trainability_cfg,
+        )
+
+        class TinyReferenceAttention(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.to_q = nn.Linear(2, 2, bias=False)
+                self.query_pose_proj = nn.Linear(6, 2, bias=False)
+                self.key_pose_proj = nn.Linear(6, 2, bias=False)
+
+        class TinyBlock(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.s_mlp = nn.Linear(2, 2, bias=False)
+                self.r_adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(2, 2, bias=False))
+                self.r_mlp = nn.Linear(2, 2, bias=False)
+                self.r_attn_anchor = TinyReferenceAttention()
+
+        class TinyDiffusion(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.model = nn.Module()
+                self.model.blocks = nn.ModuleList([TinyBlock()])
+                self.model.final_layer = nn.Linear(2, 2, bias=False)
+
+        diffusion = TinyDiffusion()
+        vae = nn.Linear(2, 2, bias=False)
+        cfg = OmegaConf.create({"trainability": {"freeze_vae": True, "train_full_dit": True, "full_dit_start_step": 10, "geometry_projections": False}})
+        trainability = _trainability_cfg(cfg)
+
+        _apply_dememwm_trainability(diffusion, vae, trainability, global_step=0)
+        flags = {name: param.requires_grad for name, param in diffusion.named_parameters()}
+        self.assertTrue(flags["model.blocks.0.r_attn_anchor.to_q.weight"])
+        self.assertFalse(flags["model.blocks.0.r_attn_anchor.query_pose_proj.weight"])
+        self.assertFalse(flags["model.blocks.0.r_attn_anchor.key_pose_proj.weight"])
+        self.assertTrue(flags["model.blocks.0.r_adaLN_modulation.1.weight"])
+        self.assertTrue(flags["model.blocks.0.r_mlp.weight"])
+        self.assertFalse(flags["model.blocks.0.s_mlp.weight"])
+        self.assertFalse(flags["model.final_layer.weight"])
+        self.assertFalse(any(param.requires_grad for param in vae.parameters()))
+
+        opt_ids = {id(param) for param in _dememwm_optimizer_parameters(diffusion, vae, trainability)}
+        self.assertEqual(opt_ids, {id(param) for param in diffusion.parameters()})
+
+        _apply_dememwm_trainability(diffusion, vae, trainability, global_step=10)
+        self.assertTrue(all(param.requires_grad for param in diffusion.parameters()))
+        self.assertFalse(any(param.requires_grad for param in vae.parameters()))
+
+        memory_only = OmegaConf.create({"trainability": {"freeze_vae": True}})
+        _apply_dememwm_trainability(diffusion, vae, _trainability_cfg(memory_only), global_step=0)
+        default_flags = {name: param.requires_grad for name, param in diffusion.named_parameters()}
+        self.assertTrue(default_flags["model.blocks.0.r_attn_anchor.query_pose_proj.weight"])
+        self.assertFalse(default_flags["model.final_layer.weight"])
+
     def test_validation_step_builds_online_memory_from_committed_latents(self):
         import algorithms.dememwm.df_video as df_video
         from algorithms.dememwm.df_video import DeMemWMMinecraft, _preprocess_dememwm_latent_batch