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):

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:

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).

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:

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):

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:

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:

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):

("--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):

# 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:

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:

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:

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:

"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:

#!/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):

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:

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:

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):

# 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:

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):

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):

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.