Recurrent Autoregressive Diffusion: Global Memory Meets Local Attention

Abstract

Method: Recurrent Autoregressive Diffusion (RAD)

1. Overall Architecture

Figure 2: RAD Model Architecture — Each DiT Block contains Spatial Attention, Temporal Attention, RNN Memory Block

RAD extends standard DiT by adding an independent temporal RNN block after spatial & temporal attention in every layer:

• Spatial Axis Attention + Temporal Axis Attention with RoPE positional encoding
• RNN block only operates along temporal dimension (per-frame compression of global history)
• AdaLN inject timestep and action condition into attention and RNN layers for interactive world modeling

2. Two Autoregressive Modes

Chunk-wise Autoregression (-c)

• Attention windows have no overlap between chunks
• All cross-chunk visual information must be stored in RNN hidden state
• Works well on simple low-texture scenes (Maze), fails on complex Minecraft environment

Frame-wise Autoregression (-f)

• Window slides 1 frame each step, full overlapping local attention
• Local pixel continuity transmitted via attention, RNN only stores global memory
• Consistently better PSNR/SSIM/LPIPS on all datasets, no chunk boundary artifacts

3. Hidden State Prefetch (Parallel Training)

Figure 4: Prefetch vs Standard Sliding Window Training

Vanilla RNN blocks force sequential training and break attention parallelism. We precompute all RNN hidden states on clean frames in stage 1; stage 2 runs full parallel attention denoising with pre-fetched memory states.

• RNN only receives clean frames (noised latent not used for memory) → faster convergence
• Attention computation parallelized over entire sequence, huge training speedup
• Train-inference consistency: inference also uses clean hidden states after DDIM denoising steps

4. RNN Module Comparison (LSTM / Mamba2 / TTT)

LSTM separates short-term hidden output yₜ and long-term cell state Cₜ, which makes it superior when attention cannot cross chunk boundaries. With overlapping frame-wise windows, three RNN variants perform nearly identically.

Video Visualization Results

Each sample shows a 4×2 grid: the top row is chunk-wise generation and the bottom row is frame-wise generation, each compared across the Mamba2, TTT, and LSTM recurrent modules, alongside the shared ground-truth reference.

Experimental Results

Datasets & Evaluation Metrics

• Memory Maze: 30k train videos (1000 frames each), test set 200 videos (100–300 frames), agent round-trip path memory test
• Minecraft MineRL: 20k train 1200-frame gameplay videos, test set with rotation & complex geometry scenes
• Metrics: PSNR (↑), SSIM (↑), LPIPS (↓), VBench temporal consistency / flicker / motion smoothness
• Baselines: Diffusion Forcing, Teacher Forcing, VRAG, LaCT (TTT-c)

1. Memory Maze Quantitative Results

2. Minecraft Quantitative Results

3. Computational Resource Analysis

TTT has severe training speed overhead; LSTM achieves balanced memory and runtime cost despite larger parameter count. Prefetch stage overhead only accounts for ~17.5% total FLOPs.

Ablation Studies

1. Clean vs Noised Memory Frames for RNN

Memory Input	PSNR	SSIM	LPIPS
Noised Latent Frames	14.70	0.38	0.52
Clean Frames (Ours Prefetch)	15.30	0.41	0.50

Feeding denoised noisy latents into RNN damages memory encoding; clean frame prefetch is critical for stable long-term memory.

2. Partial Overlap Chunk Ablation

Partially overlapping windows cannot reach performance of full frame-wise stride=1 sliding attention.

Citation

@misc{chen2025recurrentautoregressivediffusionglobal,
      title={Recurrent Autoregressive Diffusion: Global Memory Meets Local Attention}, 
      author={Taiye Chen and Zihan Ding and Anjian Li and Christina Zhang and Zeqi Xiao and Yisen Wang and Chi Jin},
      year={2025},
      eprint={2511.12940},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2511.12940}, 
}