MosaiChunk: Compositing Spatio-Temporal Memory for Autoregressive Video Generation

What it claims

A frozen autoregressive video DiT can take a non-contiguous subset of its own evicted KV cache (for example, only the latent positions covering one object) and redraw that content correctly when it comes back into view. MosaiChunk builds on that observation. A small trained router picks fixed-size KV "sections" from across space and time to fill a far-memory budget. On RememBench, median revisit CLIP rises from 0.841 to 0.936 (T2V) and from 0.807 to 0.863 (I2V, after a 180° camera turn away and back), measured against whole-chunk retrieval (MoC).

How it works

• Sections: Each chunk's unrotated KV is split into equal-size groups using balanced k-means. A lightweight descriptor encoder maps each section's pooled keys to a descriptor. The section bank stores the verbatim KV and descriptors in CPU memory.
• Scoring: The sections of the latest chunk act as queries. Each historical section outside the sliding window gets a score: its maximum cosine similarity to any query descriptor.
• Composition: The router takes the global top-N sections, where N is set by the budget. Softmax weights over all scores are renormalised to unit mean over the selected sections and used to scale their values. The selected KV is concatenated and read alongside the sliding window and an attention sink.
• Why train descriptors at all: Raw keys from neighbouring chunks are highly similar, so plain similarity retrieval favours recent content. Training is meant to produce an anti-recency descriptor space.

flowchart LR
  N[Same noisy latent + sliding window] --> T[Teacher: whole historical chunks]
  N --> S[Student: router's MosaiChunk, smaller budget]
  T --> D[Frozen DiT]
  S --> D
  D --> L[MSE on velocity]
  L -->|grads via value weights| E[Descriptor encoder]

Only the descriptor encoder is trained, by self-distillation. Both teacher and student pass through the same frozen DiT. The teacher's far memory is the whole chunks from the earlier visit:

• for T2V, identified from the prompt schedule;
• for I2V, identified from matching camera poses.

Gradients flow through the student's value weights, not through the discrete top-N indices. No ground-truth target videos are needed.

Evidence

The paper's backbones are RAVEN-adapted MiniMax-H3 (H3-AR) and LingBot-World-Infinity. Their prompt-update and camera-control interfaces are used to build the two splits.

T2V split. 100 samples whose scenarios are disjoint from router training. Each uses a four-segment appear / out-of-sight / reappear prompt structure, extended from Ring Forcing. Metrics are medians over 100 marked departure–revisit pairs. Base is sliding-window inference with the same total KV budget.

I2V split. 150 scenes (50 indoor, 100 outdoor). Each starts from a DL3DV frame plus a scene prompt.

Limits

• Not released yet. Code and data are only promised ("will be released").
• Drift is not improved. At budget 2, MosaiChunk's drift is 0.055, slightly worse than Base (0.050). The gains are in revisit similarity, not general rollout stability.
• Training needs revisit metadata. Teacher chunks come from prompt schedules or camera poses, so you need that information for your data.
• Training cost is unreported. Training backpropagates through the frozen DiT, so it is not free even though only a small encoder is updated. The page gives no compute figures.
• CPU-memory overhead is unreported. The page does not quantify the CPU-held section bank's size or transfer overhead.
• Captured text is truncated. It cuts off in the I2V benchmark description. The full I2V table, any ablations (for example, with and without trained descriptors) and the authors' stated limitations are not visible here.