code/attention.md

# Adding an Attention-Based Memory Module to Echo-Memory

This guide explains how to add a **new attention-based memory mechanism** to the Wan 2.1 backbone in this repository, following the exact integration pattern used by the **Block-wise SSM** row (`diffsynth/models/memory/block_wise_ssm.py`). Block-wise SSM is the right template because it is a **per-DiT-block module**: it is instantiated inside each selected transformer block and called during the block's forward pass — precisely where an attention variant would live.

Read `doc/memory_mechanisms.md` and the repo `CLAUDE.md` first for the two-chunk paradigm and the public-repo constraints. This document assumes that background.

---
I want to add some loss, for example, using the model's uncertainty about the current prediction, um: e^(-um)*sg(MSE(VAE(Target View), predicted x0)) + um, as a metric for the model's retrieval, thereby improving memory capabilities. Please return the code implementation of this method in the current codebase. In addition, please provide methods you believe can improve the model's ability to retrieve preceding frames, and compile them into a retrieve.md file and save it.
## 0. How per-block memory is wired (the pattern you will copy)

Unlike the FramePack/Spatial rows (which act in the pipeline on context latents), Block-wise SSM attaches **inside the DiT transformer blocks**. The same block class is defined in **two places that must stay in sync**:

| Location | Role |
| --- | --- |
| `src/model_training/train.py` → `class DiTBlock_w_Action` (~line 160) | **Training** definition of the augmented block |
| `env/loop_utils.py` → `class DiTBlock_w_Action` (~line 43) | **Inference** definition (must match training byte-for-byte in module structure) |

The base Wan blocks (`diffsynth/models/wan_video_dit.py`) are **replaced** at load time by `DiTBlock_w_Action`, copying over `self_attn / cross_attn / norm* / ffn / modulation` weights from the original blocks. The memory module is an *extra* sub-module on selected blocks; its weights are the only newly-trained parameters.

**The forward hook location** (train.py `DiTBlock_w_Action.forward`, ~line 230-238):

```python
input_x = modulate(self.norm1(x), shift_msa, scale_msa)
x = self.gate(x, gate_msa, self.self_attn(input_x, freqs))   # ← self-attention
if num_frames is not None:
    if hasattr(self, "block_wise_ssm"):
        x = self.block_wise_ssm(x, f=num_frames)             # ← MEMORY HOOK (after self-attn, before cross-attn)
x = x + self.cross_attn(self.norm3(x), context)              # ← cross-attention
input_x = modulate(self.norm2(x), shift_mlp, scale_mlp)
x = self.gate(x, gate_mlp, self.ffn(input_x))
```

`x` here is shape `(B, F*S, D)` — batch, (latent frames × spatial tokens) flattened, hidden dim. `num_frames` (`f`) is the number of latent frames, used to reshape per-frame. **This is where your attention-memory module reads/writes.**

**Detection at inference is by checkpoint key name.** `env/loop_utils.py` (~line 288) scans ckpt keys with a regex and decides which blocks get the module:

```python
m = re.match(r"blocks\.(\d+)\.block_wise_ssm\.", key)   # → block_wise_block_ids
```

So the *attribute name* you give your module (e.g. `self.attn_memory = ...`) becomes the checkpoint key prefix and must be matched by a new regex. **Keep the attribute name stable** — `env/memory_baseline_runtime.py` and `inference/unified_inference.py` rely on it.

---

## 1. Files to ADD

### 1a. The module — `diffsynth/models/memory/attn_memory.py`

Mirror the shape contract of `block_wise_ssm.py`: input `(B, F*S, D)`, takes `f` (frames), returns the same shape with a **zero-initialized residual gate** so an untrained module is an identity at step 0 (critical — the base backbone must not be disturbed before training).

```python
import torch
import torch.nn as nn
import torch.nn.functional as F


class AttentionMemory(nn.Module):
    """
    Example attention-based memory module attached per DiT block.

    Reads the within-block hidden state and applies an extra attention
    operation along the TIME axis for each spatial-token trajectory
    (analogous to block_wise_ssm's recurrent time update, but attention).

    Shape contract (must match block_wise_ssm):
      forward(x, f) where x: (B, F*S, D), f = number of latent frames.
    """

    def __init__(self, dim: int, num_heads: int = 8):
        super().__init__()
        self.dim = int(dim)
        self.num_heads = int(num_heads)
        self.q = nn.Linear(dim, dim)
        self.k = nn.Linear(dim, dim)
        self.v = nn.Linear(dim, dim)
        self.o = nn.Linear(dim, dim)
        # Zero-init residual gate => identity at init (do NOT skip this).
        self.gate = nn.Parameter(torch.zeros(1))

    def forward(self, x: torch.Tensor, f: int, **_kwargs):
        if x is None or x.ndim != 3:
            return x
        b, n, d = x.shape
        f = int(f or 0)
        if d != self.dim or f <= 1 or n % f != 0:
            return x

        spatial = n // f
        # (B, F*S, D) -> (B*S, F, D): attend across frames per spatial trajectory.
        x_seq = x.reshape(b, f, spatial, d).permute(0, 2, 1, 3).reshape(b * spatial, f, d)

        h = self.num_heads
        q = self.q(x_seq).reshape(b * spatial, f, h, d // h).transpose(1, 2)
        k = self.k(x_seq).reshape(b * spatial, f, h, d // h).transpose(1, 2)
        v = self.v(x_seq).reshape(b * spatial, f, h, d // h).transpose(1, 2)
        y = F.scaled_dot_product_attention(q, k, v)                 # causal=False -> full temporal memory
        y = y.transpose(1, 2).reshape(b * spatial, f, d)
        y = self.o(y)

        y = y.reshape(b, spatial, f, d).permute(0, 2, 1, 3).reshape(b, n, d)
        return x + torch.tanh(self.gate) * y                        # gated residual
```

> Vary the internal operation to run the experiment you care about — causal vs. bidirectional temporal attention, cross-attention into stored context tokens, windowed/strided attention, etc. The **interface and the gated-residual identity-at-init must stay fixed**; only the body changes.

### 1b. Export it — `diffsynth/models/memory/__init__.py`

Add alongside the existing exports:

```python
from .attn_memory import AttentionMemory
```

### 1c. The training launcher — `train/memory_baselines_basic/run_ablation_attn_memory_two_chunk.sh`

Copy `run_ablation_block_wise_ssm_two_chunk.sh` and swap the memory flags (see §4).

---

## 2. Files to MODIFY for TRAINING (`src/model_training/train.py`)

There are **five** edit sites. Search for `block_wise_ssm` to find each by analogy.

**(1) Module import** (~line 77, plus the `modules_to_clear` list ~line 44 so hot-reload works):

```python
from diffsynth.models.memory.attn_memory import AttentionMemory
# and add 'diffsynth.models.memory.attn_memory' to modules_to_clear
```

**(2) `DiTBlock_w_Action.__init__`** (~line 160) — add a constructor flag and instantiate:

```python
def __init__(self, ..., use_block_wise_ssm=False, use_videossm_hybrid=False,
             use_attn_memory: bool = False, attn_memory_heads: int = 8, ...):
    ...
    self.use_attn_memory = bool(use_attn_memory)
    if use_attn_memory:
        self.attn_memory = AttentionMemory(dim, num_heads=attn_memory_heads)
```

**(3) `DiTBlock_w_Action.forward`** (~line 232) — add the hook next to the SSM hook:

```python
if num_frames is not None:
    if hasattr(self, "block_wise_ssm"):
        x = self.block_wise_ssm(x, f=num_frames)
    if hasattr(self, "attn_memory"):
        x = self.attn_memory(x, f=num_frames)        # ← new
```

**(4) argparse flags** (~line 1479-1510). Add the numeric default tuples near `--ssm_every_n_blocks`, and the boolean flag to the store-true list (~line 1510 alongside `--use_block_wise_ssm`):

```python
("--attn_memory_every_n_blocks", dict(type=int, default=4)),
("--attn_memory_heads",          dict(type=int, default=8)),
# ... and add "--use_attn_memory" to the list of store_true flags
```

Optionally add a mutual-exclusion guard like the one at ~line 1556 (`use_block_wise_ssm and use_videossm_hybrid`) if your module should not co-exist with another per-block memory.

**(5) block instantiation + parameter freezing** (~line 1695-1761):

```python
# instantiation loop (~line 1695)
use_attn_memory = bool(_arg('use_attn_memory', False))
attn_every_n   = max(int(_arg('attn_memory_every_n_blocks', 4) or 4), 1)
...
for block_id, old_block in enumerate(old_blocks):
    attach_block_ssm  = use_block_wise_ssm and (block_id % ssm_every_n == 0)
    attach_attn_mem   = use_attn_memory   and (block_id % attn_every_n == 0)   # ← new
    new_block = DiTBlock_w_Action(
        ...,
        use_block_wise_ssm=attach_block_ssm,
        use_attn_memory=attach_attn_mem,                                       # ← new
        attn_memory_heads=int(_arg('attn_memory_heads', 8) or 8),
    )
```

Then add `"attn_memory"` to **all three** `requires_grad` filters (~line 1744, 1751, 1758) so only your module (plus action MLP / self-attn-with-action) trains and the rest of the DiT stays frozen:

```python
if "action_mlp" in name or "self_attn_with_action" in name \
   or "block_wise_ssm" in name or "videossm_hybrid" in name \
   or "attn_memory" in name:          # ← new
    param.requires_grad = True
else:
    param.requires_grad = False
```

> ⚠️ If you forget the freeze filter, your module will not be trained (it'll be frozen with everything else), or the whole DiT becomes trainable — both break the controlled-ablation premise.

---

## 3. Files to MODIFY for INFERENCE / EVAL

### 3a. `env/loop_utils.py` — mirror the block and detect from ckpt keys

This file re-defines `DiTBlock_w_Action` and rebuilds blocks at load time. **It must structurally match `train.py`.** Three edits:

1. **Import + module attribute** in its `DiTBlock_w_Action` (~line 43-75): add the `use_attn_memory` ctor arg, `self.attn_memory = AttentionMemory(dim, ...)`, and the same forward hook (~line 111).
2. **`_build_action_blocks`** (~line 202): add an `attn_memory_block_ids` set param and pass `use_attn_memory=block_id in attn_memory_block_ids` into each block.
3. **`load_pipeline_and_ckpt`** (~line 288): add a detection regex so the right blocks are reconstructed before loading weights:

```python
m = re.match(r"blocks\.(\d+)\.attn_memory\.", key)
if m:
    attn_memory_block_ids.add(int(m.group(1)))
```

and thread `attn_memory_block_ids` into the `_build_action_blocks(...)` call (~line 304).

> Because detection is **automatic from checkpoint keys**, no inference flag is needed to turn the module on — the presence of `blocks.N.attn_memory.*` keys in the safetensors reconstructs the slots. This mirrors how `block_wise_ssm` / `videossm_hybrid` work today.

### 3b. `env/memory_baseline_runtime.py` — register the profile

Three edits (search `block_wise_ssm`):

1. **`MemoryProfile` dataclass** (~line 46): add `use_attn_memory: bool = False`.
2. **`MEMORY_PROFILE_REGISTRY`** (~line 308): add a spec. The `ckpt_substrings` must match your launcher's `--output_path` folder name:

```python
MemoryProfileSpec(
    profile_id="attn_memory_two_chunk",
    ckpt_substrings=("memory_baselines_basic_abl_attn_memory_two_chunk",),
    paper_tag="attention_memory",
    train_flags=("--use_attn_memory", "--context_memory_frames 5"),
    infer_flags=("load_pipeline_and_ckpt auto-infers Attn-memory from attn_memory.* ckpt keys",),
    eval_flags=("--context_frames 5",),
    profile=MemoryProfile(use_attn_memory=True, context_override=5),
),
```

3. **`profile_to_argv`** (~line 393) and **`apply_memory_baseline_pipe`** (~line 427): add the `--use_attn_memory` emission and `pipe.use_attn_memory = ...` line, mirroring the block_wise_ssm entries.

### 3c. `inference/unified_inference.py` — add the alias

In `_REGISTRY_ALIAS` (~line 72) add:

```python
"attn_memory": "attn_memory_two_chunk",
```

and (~line 138) `pipe.use_attn_memory = bool(getattr(profile, "use_attn_memory", False))`. Update the `--help` memory-type list near line 172.

---

## 4. The training launcher

`train/memory_baselines_basic/run_ablation_attn_memory_two_chunk.sh` — copy the block-wise SSM launcher and change only the memory flags, run name, and output folder. The critical line:

```bash
#!/bin/bash
set -euo pipefail
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
source "${SCRIPT_DIR}/common_env.sh"
source "${SCRIPT_DIR}/common_sampling_two_chunk.sh"
accelerate launch src/model_training/train.py \
  --dataset_base_path "${dataset_base_path}" --dataset_metadata_path "${dataset_base_path}/${METADATA_NAME}" \
  --context_source replay --prev_chunk_frames 81 --context_memory_frames 5 --use_rt_relative --height 352 --width 640 \
  --dataset_repeat 1 --per_device_train_batch_size 1 --gradient_accumulation_steps 1 --num_workers 16 \
  --model_paths "${model_paths}" --learning_rate 5e-5 --num_epochs 1 --remove_prefix_in_ckpt "${remove_prefix_in_ckpt}" \
  --output_path "${output_base}_abl_attn_memory_two_chunk" --trainable_models dit --ckpt_interval "${CKPT_INTERVAL:-1000}" --save_full_model \
  --wandb_run_name "abl_attn_memory_two_chunk" \
  --enable_context_memory --training_mode context --context_drop_prob 0.1 --cfg_target_only \
  --train_cam_pose --add_action_attn --action_use_temporal_attention \
  --use_moc --moc_temperature 1.0 \
  --use_attn_memory --attn_memory_every_n_blocks 4 --attn_memory_heads 8 \
  --timestep_shift "${TIMESTEP_SHIFT:-15}" \
  "${SAMPLING_TWO_CHUNK_FLAGS[@]}" \
  2>&1 | tee "${LOG_DIR}/abl_attn_memory_two_chunk_$(date +%Y%m%d_%H%M%S).log"
```

> The `--output_path` folder name (`..._abl_attn_memory_two_chunk`) **must contain the `ckpt_substrings` token** you registered in §3b, so eval/inference auto-resolve the profile from the checkpoint path. Keep all the shared hyperparameters (LR 5e-5, 352×640, 81 frames, shift 15) identical to the other rows — that is what makes this a *controlled* ablation; only the memory pathway should differ.

---

## 5. Running the experiment

**Environment** (repo root):

```bash
export WAN_BASE_MODEL=/path/to/Wan2.1-T2V-1.3B
export DATASET_BASE_PATH=data/Context-as-Memory-Dataset
export PYTHONPATH=$PWD:${PYTHONPATH:-}
export OUTPUT_BASE_ROOT=$PWD/outputs
```

**Sanity-check the plumbing first** — the existing standalone tests cover context/two-chunk wiring; run them after your edits:

```bash
PYTHONPATH=. python3 tests/test_context_chunk_utils.py
PYTHONPATH=. python3 tests/test_two_chunk_anchor_readout.py
```

Add a tiny shape test for your module (recommended), asserting identity at init and `(B,F*S,D)` round-trip:

```python
import torch
from diffsynth.models.memory.attn_memory import AttentionMemory
m = AttentionMemory(64, num_heads=8)
x = torch.randn(2, 5*16, 64)            # B=2, F=5, S=16, D=64
assert torch.allclose(m(x, f=5), x)     # zero gate => identity at init
```

**Train** (optionally do a quick run with a small metadata index first):

```bash
# Optional: 1000-row index for a fast smoke run
OUTPUT_CSV="${DATASET_BASE_PATH}/metadata_1000.csv" METADATA_MAX_ROWS=1000 bash scripts/run_generate_metadata.sh

bash train/memory_baselines_basic/run_ablation_attn_memory_two_chunk.sh
```

Checkpoints land under `outputs/memory_baselines_basic_abl_attn_memory_two_chunk/epoch-*.safetensors`. **Verify the new weights are present and gradients flowed:**

```bash
python -c "
from safetensors.torch import load_file
k = load_file('outputs/memory_baselines_basic_abl_attn_memory_two_chunk/epoch-0.safetensors')
hits = [x for x in k if 'attn_memory' in x]
print(len(hits), 'attn_memory tensors; sample:', hits[:3])
"
```

**Inference** (auto-detected from ckpt keys; `--memory_type auto` works once §3 is done):

```bash
python inference/unified_inference.py \
  --ckpt outputs/memory_baselines_basic_abl_attn_memory_two_chunk/epoch-0.safetensors \
  --memory_type attn_memory \
  --context_image assets/opendomain_revisit/1774363417.png \
  --action_path env/action_rotation_left_45.json \
  --prompt "A toy bear on a table" \
  --output_path attn_memory_test.mp4
```

**Evaluate** (the row folder name in `CKPT` drives profile selection):

```bash
export CKPT=outputs/memory_baselines_basic_abl_attn_memory_two_chunk/epoch-0.safetensors
bash eval/v2/run_basic_replay_gt.sh                          # fast fidelity check (~5 min)
bash eval/v2/run_static_consistency_loop_and_revisit.sh      # full paper bundle (loop closure + revisit)
PHASE=stage1 OOD_DIR=assets/opendomain_revisit \
  bash eval/v2/revisit_suite/run_one_click_revisit_eval.sh   # open-domain generalization
```

Compare the resulting MSE / PSNR / SSIM / LPIPS (and revisit-tail consistency) against the `block_wise_ssm`, `spatial_mem`, and `context_k*` rows to place your attention-memory variant in the paper matrix.

---

## 6. Checklist & pitfalls

- [ ] `diffsynth/models/memory/attn_memory.py` added; exported in `__init__.py`.
- [ ] `train.py`: import + `modules_to_clear` + ctor + forward hook + argparse + instantiation + **all three freeze filters**.
- [ ] `env/loop_utils.py`: `DiTBlock_w_Action` mirrors training **exactly** (structure, attribute name, hook order), + detection regex + `_build_action_blocks` thread-through.
- [ ] `env/memory_baseline_runtime.py`: `MemoryProfile` field + registry spec + `profile_to_argv` + `apply_memory_baseline_pipe`.
- [ ] `inference/unified_inference.py`: alias + pipe flag + help text.
- [ ] Launcher `--output_path` folder name contains the registry `ckpt_substrings` token.

**Common failure modes**

- **Module attribute name mismatch** between train/inference, or vs. the detection regex → checkpoint keys won't reconstruct the slots; weights load as "unexpected" and silently do nothing. Keep `attn_memory` consistent everywhere.
- **Non-zero gate at init** → corrupts the frozen backbone before training and destabilizes early steps. Always zero-init the residual gate.
- **Forgotten freeze filter** → either the module isn't trained, or the whole DiT trains (breaks the controlled comparison and the `--save_full_model` checkpoints diff against base unexpectedly).
- **`loop_utils.DiTBlock_w_Action` drifting from `train.py`** → load-time block replacement mismatches and `strict=False` hides it as missing/unexpected keys. When you edit one, edit both.

**Public-repo constraints** (from `CLAUDE.md`): no machine-local absolute paths, no upload scripts, keep diffs minimal and match existing bash/Python patterns. Don't commit `outputs/`, `data/`, or weights.