feat(bridge): trajectory-step activation caching and denoising patching for diffusion LMs

.github/workflows/checks.yml

@@ -432,6 +432,7 @@ jobs:
           - "BERT"
           - "Backward_Lens_Demo"
           - "Bridge_Evals_Demo"
+          - "Diffusion_Trajectory_Patching_Demo"
           - "Exploratory_Analysis_Demo"
           # - "Grokking_Demo"
           - "Head_Detector_Demo"

docs/make_docs.py

@@ -40,6 +40,7 @@ def copy_demos(_app: Optional[Any] = None):
     copy_to_dir = GENERATED_DIR / "demos"
     notebooks_to_copy = [
         "Backward_Lens_Demo.ipynb",
+        "Diffusion_Trajectory_Patching_Demo.ipynb",
         "Exploratory_Analysis_Demo.ipynb",
         "Jacobian_Lens_Coordinate_Patch_Benchmark_Demo.ipynb",
         "Jacobian_Lens_Decomposition_Demo.ipynb",

docs/source/content/tutorials.md

@@ -46,6 +46,8 @@ The two core notebooks, which pair with the links in [Where To Start](#where-to-
 
 - [**Natural Language Autoencoders**](https://colab.research.google.com/github/TransformerLensOrg/TransformerLens/blob/main/demos/Natural_Language_Autoencoders_Demo.ipynb) - Runs the [Natural Language Autoencoder](https://transformer-circuits.pub/2026/nla/) round trip with `TransformerBridge` hooks and generation: a verbalizer model explains one residual-stream activation in text, and a reconstructor model predicts the activation back from that text. The mechanics run on CPU with a tiny Qwen2 checkpoint; the opt-in released-model section (`TL_RUN_RELEASED_NLA=1`) downloads about 39 GiB of Qwen2.5-7B-Instruct checkpoints.
 
+- [**Diffusion Trajectory Caching & Denoising Patching**](https://colab.research.google.com/github/TransformerLensOrg/TransformerLens/blob/main/demos/Diffusion_Trajectory_Patching_Demo.ipynb) - Records step-indexed activations (`DiffusionTrajectoryCache`) across iterative masked-diffusion denoising steps via `diffusion_generate(..., return_cache=True)` and patches residual-stream activations across denoising steps (`make_step_patch_hook`) before vs. after the answer commitment step on `LLaDA-8B-Instruct` and `Dream-7B` ARC-Challenge trajectories. Runs on CPU in seconds; the opt-in released-model section (`TL_RUN_RELEASED_LLADA=1`) runs `GSAI-ML/LLaDA-8B-Instruct`.
+
 - [**Patchscopes & Generation with Patching**](https://colab.research.google.com/github/TransformerLensOrg/TransformerLens/blob/main/demos/Patchscopes_Generation_Demo.ipynb) - Implements [Patchscopes](https://arxiv.org/abs/2401.06102) (read a hidden representation from a source prompt and patch it into a target prompt), extends it to generating multiple tokens with patching, and applies it to logit lens, entity description and zero-shot feature extraction on GPT-2 small.
 
 - [**SVD Circuits**](https://colab.research.google.com/github/TransformerLensOrg/TransformerLens/blob/main/demos/SVD_Circuits_Demo.ipynb) - Decomposes one GPT-2 attention head's QK and OV maps into singular directions with `FactoredMatrix.svd`, reads the OV directions out through the unembedding with `logit_readout`, and keeps or ablates single directions on an IOI prompt through `hook_result`. Based on [Ahmad et al. (2025)](https://arxiv.org/abs/2511.20273); runs on CPU.

docs/source/index.md

@@ -71,6 +71,7 @@ generated/demos/SVD_Circuits_Demo
 generated/demos/SwiGLU_Gate_Sign_Demo
 generated/demos/Weight_Edit_Audit_Demo
 generated/demos/Backward_Lens_Demo
+generated/demos/Diffusion_Trajectory_Patching_Demo
 generated/demos/Natural_Language_Autoencoders_Demo
 generated/demos/Main_Demo
 generated/demos/Exploratory_Analysis_Demo

tests/integration/model_bridge/test_llada_adapter.py

@@ -904,3 +904,49 @@ def test_causal_loss_modes_fail_clearly(models: TinyModels, return_type: str) ->
     with pytest.raises(NotImplementedError, match="causal.*loss|shifted.*loss"):
         models.bridge(tokens, return_type=return_type)
     assert len(models.bridge.original_model.output_attentions_requests) == calls_before
+
+
+def test_diffusion_generate_trajectory_cache_and_step_patching(models: TinyModels) -> None:
+    from transformer_lens.ActivationCache import DiffusionTrajectoryCache
+
+    prompt_a = torch.tensor([[5, 7, 9]])
+    prompt_b = torch.tensor([[11, 13, 15]])
+    mask_id = models.bridge.cfg.mask_token_id
+
+    seq_a, traj_a = models.bridge.diffusion_generate(
+        prompt_a,
+        max_new_tokens=3,
+        return_cache=True,
+        names_filter=["blocks.0.attn.hook_pattern", "blocks.1.hook_out"],
+    )
+    assert isinstance(seq_a, torch.Tensor)
+    assert isinstance(traj_a, DiffusionTrajectoryCache)
+    assert seq_a.shape == (1, 6)
+    assert torch.equal(seq_a[:, :3], prompt_a)
+    assert torch.all(seq_a[:, 3:] != mask_id)
+    assert traj_a.steps == (0, 1, 2)
+    # Step 0 starts with all 3 continuation positions masked; each step unmasks 1 token
+    assert (traj_a.step_tokens[0][0, 3:] == mask_id).sum().item() == 3
+    assert (traj_a.step_tokens[1][0, 3:] == mask_id).sum().item() == 2
+    assert (traj_a.step_tokens[2][0, 3:] == mask_id).sum().item() == 1
+    assert traj_a.stack("blocks.0.attn.hook_pattern").shape == (3, 1, 4, 6, 6)
+    assert traj_a.stack("blocks.1.hook_out").shape == (3, 1, 6, 32)
+
+    # Step-selective activation patching across denoising steps
+    patch_fn = traj_a.make_step_patch_hook(
+        "blocks.1.hook_out",
+        steps=[0],
+        positions=slice(3, 6),
+    )
+    with models.bridge.hooks(fwd_hooks=[("blocks.1.hook_out", patch_fn)]):
+        _, patched_traj_b = models.bridge.diffusion_generate(
+            prompt_b,
+            max_new_tokens=3,
+            return_cache=True,
+            names_filter=["blocks.1.hook_out"],
+            cache_steps=[0, 1],
+        )
+    torch.testing.assert_close(
+        patched_traj_b[0]["blocks.1.hook_out"][:, 3:6, :],
+        traj_a[0]["blocks.1.hook_out"][:, 3:6, :],
+    )

tests/unit/model_bridge/supported_architectures/test_llada_adapter.py

@@ -6,7 +6,10 @@
 import pytest
 import torch
 
-from tests.unit.model_bridge.supported_architectures.helpers import make_bridge_cfg
+from tests.unit.model_bridge.supported_architectures.helpers import (
+    DiffusionContractTests,
+    make_bridge_cfg,
+)
 from transformer_lens.config import TransformerBridgeConfig
 from transformer_lens.conversion_utils.conversion_steps import RearrangeTensorConversion
 from transformer_lens.conversion_utils.param_processing_conversion import (
@@ -323,3 +326,22 @@ def test_existing_causal_attention_default_remains_triangular() -> None:
     pattern = cache["blocks.0.attn.hook_pattern"]
     assert pattern.shape == (1, 4, 5, 5)
     assert torch.count_nonzero(torch.triu(pattern, diagonal=1)) == 0
+
+
+class TestLLaDADiffusionContract(DiffusionContractTests):
+    adapter_cls = LLaDAArchitectureAdapter
+    architecture = "LLaDAModelLM"
+    expected_sampler = "diffusion_generate"
+
+    def _make_adapter(self) -> LLaDAArchitectureAdapter:
+        return LLaDAArchitectureAdapter(_bridge_config())
+
+
+def test_prepare_model_binds_diffusion_generate_sampler(
+    adapter: LLaDAArchitectureAdapter,
+) -> None:
+    dummy = SimpleNamespace(config=SimpleNamespace(output_attentions=True, use_cache=True))
+    adapter.prepare_model(dummy)
+    assert dummy.config.output_attentions is False
+    assert dummy.config.use_cache is False
+    assert callable(getattr(dummy, "diffusion_generate", None))

tests/unit/model_bridge/test_diffusion_generate.py

@@ -122,6 +122,12 @@ def test_llada2_block_length_never_exceeds_budget(self) -> None:
         kwargs = self._kwargs("llada2_moe", "LLaDA2MoeArchitectureAdapter", 4, 8)
         assert kwargs["block_length"] <= 4
 
+    def test_llada_rounds_gen_length_to_whole_blocks(self) -> None:
+        kwargs = self._kwargs("llada", "LLaDAArchitectureAdapter", 40, 8)
+        assert kwargs["gen_length"] % kwargs["block_length"] == 0
+        assert kwargs["gen_length"] >= 40
+        assert kwargs["steps"] == kwargs["gen_length"]
+
     def test_gidd_max_length_is_generated_tokens_not_canvas(self) -> None:
         """Gidd's windows start at prompt_length and span max_length, so folding
         the prompt into it would over-generate by prompt_len tokens."""
@@ -266,3 +272,175 @@ def test_shape_mismatch_raises(self) -> None:
         bridge, _ = self._router()
         with _pytest.raises(ValueError, match="does not match"):
             bridge.set_processed_weights({"layer.weight": torch.zeros(3, 3)})
+
+
+class TestDiffusionTrajectoryCache:
+    """Trajectory-step activation caching and step-selective patching for diffusion LMs."""
+
+    def _make_fake_diffusion_bridge(self):
+        import torch
+
+        from transformer_lens.hook_points import HookPoint
+        from transformer_lens.model_bridge.bridge import TransformerBridge
+        from transformer_lens.model_bridge.bridge_core import BridgeCore
+
+        class _StepModel(torch.nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.hp = HookPoint()
+                self.hp.name = "blocks.0.hook_out"
+
+            def forward(self, canvas: torch.Tensor) -> torch.Tensor:
+                act = canvas.to(torch.float32).unsqueeze(-1).repeat(1, 1, 4)
+                return self.hp(act)
+
+            def diffusion_generate(
+                self, tokens: torch.Tensor, max_new_tokens: int = 3, steps: int = 3, **kwargs
+            ) -> torch.Tensor:
+                prompt_len = tokens.shape[-1]
+                canvas = torch.full(
+                    (tokens.shape[0], prompt_len + max_new_tokens),
+                    99,
+                    dtype=tokens.dtype,
+                )
+                canvas[:, :prompt_len] = tokens
+                for s in range(steps):
+                    out = self(canvas)
+                    # Unmask position prompt_len + s using the hooked activation's first channel
+                    unmasked_token = int(out[0, prompt_len + s, 0].item()) % 50 + 10
+                    canvas = canvas.clone()
+                    canvas[:, prompt_len + s] = unmasked_token
+                return canvas
+
+        class _Adapter:
+            native_sampler = "diffusion_generate"
+            supports_generation = False
+
+            def native_sampler_kwargs(self, max_new_tokens: int, prompt_len: int) -> dict:
+                return {"max_new_tokens": max_new_tokens, "steps": max_new_tokens}
+
+        class _Cfg:
+            architecture = "FakeDiffusionLM"
+            device = "cpu"
+            n_layers = 1
+
+        class _FakeBridge(BridgeCore, torch.nn.Module):
+            diffusion_generate = TransformerBridge.diffusion_generate
+            current_diffusion_step = TransformerBridge.current_diffusion_step
+
+            def __init__(self) -> None:
+                torch.nn.Module.__init__(self)
+                self.original_model = _StepModel()
+                self.adapter = _Adapter()
+                self.cfg = _Cfg()
+                self.tokenizer = None
+                self.compatibility_mode = False
+                self.context_level = 0
+
+            @property
+            def hook_dict(self):
+                return {"blocks.0.hook_out": self.original_model.hp}
+
+            def _gated_hook_reason(self, name: str):
+                return None
+
+            def _check_hook_fireable(self, caller_name: str, canonical_name: str) -> None:
+                if canonical_name not in self.hook_dict:
+                    raise KeyError(caller_name)
+
+        return _FakeBridge()
+
+    def test_return_cache_preserves_every_denoising_step_without_clobbering(self) -> None:
+        import torch
+
+        from transformer_lens.ActivationCache import DiffusionTrajectoryCache
+
+        bridge = self._make_fake_diffusion_bridge()
+        prompt = torch.tensor([[1, 2]])
+        sequences, traj_cache = bridge.diffusion_generate(
+            prompt, max_new_tokens=3, return_cache=True, names_filter="blocks.0.hook_out"
+        )
+        assert sequences.shape == (1, 5)
+        assert isinstance(traj_cache, DiffusionTrajectoryCache)
+        assert traj_cache.steps == (0, 1, 2)
+        assert traj_cache.num_steps_run == 3
+        assert len(traj_cache) == 3
+        # Step 0 has all 3 continuation positions masked (99); step 1 has position 2 unmasked
+        assert traj_cache.step_tokens[0].tolist() == [[1, 2, 99, 99, 99]]
+        assert traj_cache.step_tokens[1][0, 2].item() != 99
+        assert traj_cache.step_tokens[1][0, 3].item() == 99
+        stacked = traj_cache.stack("blocks.0.hook_out")
+        assert stacked.shape == (3, 1, 5, 4)
+        assert not torch.equal(
+            traj_cache[0]["blocks.0.hook_out"], traj_cache[2]["blocks.0.hook_out"]
+        )
+        assert bridge.current_diffusion_step is None
+
+    def test_cache_steps_filters_recorded_denoising_steps(self) -> None:
+        import torch
+
+        bridge = self._make_fake_diffusion_bridge()
+        prompt = torch.tensor([[1, 2]])
+        _, traj_cache = bridge.diffusion_generate(
+            prompt,
+            max_new_tokens=4,
+            return_cache=True,
+            names_filter="blocks.0.hook_out",
+            cache_steps=[0, 2],
+            pos_slice=slice(2, 4),
+            remove_batch_dim=True,
+        )
+        assert traj_cache.steps == (0, 2)
+        assert traj_cache.num_steps_run == 4
+        assert 1 not in traj_cache
+        assert traj_cache[0]["blocks.0.hook_out"].shape == (2, 4)
+        assert traj_cache.stack("blocks.0.hook_out").shape == (2, 2, 4)
+
+    def test_make_step_patch_hook_patches_only_selected_steps_and_positions(self) -> None:
+        import torch
+
+        from transformer_lens.ActivationCache import (
+            ActivationCache,
+            DiffusionTrajectoryCache,
+        )
+
+        bridge = self._make_fake_diffusion_bridge()
+        prompt = torch.tensor([[1, 2]])
+        # Build a source trajectory cache whose step-1 activation at pos 3 is set to 7.0
+        _, clean_traj = bridge.diffusion_generate(
+            prompt, max_new_tokens=3, return_cache=True, names_filter="blocks.0.hook_out"
+        )
+        custom_act = clean_traj[1]["blocks.0.hook_out"].clone()
+        custom_act[:, 3, :] = 7.0
+        source_traj = DiffusionTrajectoryCache(
+            {1: ActivationCache({"blocks.0.hook_out": custom_act}, bridge)},
+            num_steps_run=3,
+            model=bridge,
+        )
+        patch_fn = source_traj.make_step_patch_hook(
+            "blocks.0.hook_out", steps=1, positions=slice(3, 4)
+        )
+        with bridge.hooks(fwd_hooks=[("blocks.0.hook_out", patch_fn)]):
+            patched_seq, patched_traj = bridge.diffusion_generate(
+                prompt, max_new_tokens=3, return_cache=True, names_filter="blocks.0.hook_out"
+            )
+        # Step 1 position 3 was patched to 7.0 -> unmasked token at step 1 is (7 % 50 + 10) == 17
+        assert patched_seq[0, 3].item() == 17
+        assert torch.all(patched_traj[1]["blocks.0.hook_out"][:, 3, :] == 7.0)
+        # Step 0 was not in steps=1 and remains unpatched (99.0 at masked position 3)
+        assert torch.all(patched_traj[0]["blocks.0.hook_out"][:, 3, :] == 99.0)
+
+    def test_invalid_cache_steps_rejected(self) -> None:
+        from typing import Any, cast
+
+        import pytest as _pytest
+        import torch
+
+        bridge = self._make_fake_diffusion_bridge()
+        prompt = torch.tensor([[1, 2]])
+        with _pytest.raises(ValueError, match="non-negative"):
+            bridge.diffusion_generate(prompt, max_new_tokens=2, return_cache=True, cache_steps=-1)
+        with _pytest.raises(TypeError, match="cache_steps"):
+            bridge.diffusion_generate(
+                prompt, max_new_tokens=2, return_cache=True, cache_steps=cast(Any, True)
+            )

tests/unit/test_make_docs.py

@@ -39,6 +39,7 @@ def test_copy_demos_creates_generated_dir_when_absent(tmp_path, monkeypatch):
     copied = sorted(p.name for p in (generated / "demos").iterdir())
     assert copied == [
         "Backward_Lens_Demo.ipynb",
+        "Diffusion_Trajectory_Patching_Demo.ipynb",
         "Exploratory_Analysis_Demo.ipynb",
         "Jacobian_Lens_Coordinate_Patch_Benchmark_Demo.ipynb",
         "Jacobian_Lens_Decomposition_Demo.ipynb",