dKV Forcing: Separating Generation and Memory Representations through a Differentiable KV Cache
Organizations: Georgia Institute of Technology
Abstract
A causal video generator produces two key--value (KV) states for every segment: the KV it uses while denoising that segment, and the KV it writes into the cache for later segments to read. In a pretrained generator the two already differ substantially, later segments attend to them differently, and replacing the written KV with the generation KV degrades the continuation. We propose dKV Forcing (dKV denotes differentiable KV) to learn a dedicated memory representation from its future use: a writer-specific low-rank residual on the shared backbone separates memory writing from generation and is trained through the future reads of what it writes. Learning a write from its future use needs gradients through the cache, which the mutable cache in common use cannot provide. dKV Forcing replaces the training-time cache by persistent differentiable pages with local rematerialization and, in one temporal traversal, computes exactly the gradient that a two-pass reconstruction of the same objective computes. One traversal computes the full memory gradient at generator-update overhead over detached-write training, with less resident memory and shorter training cycles than the two-pass reconstruction. Over Self Gradient Forcing (SGF), the same objective trained with the same traversal, writer adaptation raises VBench-Long aesthetic quality by points at 60,s and points at 240,s with the highest subject and background consistency, and rollout dropout of the residual keeps most of this gain while restoring the motion of training without write gradients.
Figures & tables
| time (s) | GPU memory (GiB) | ||||||
|---|---|---|---|---|---|---|---|
| method | cycle | gen. update | critic rollout | vs. two-pass | gen. update peak | whole-run peak | resident |
| two-pass reconstruction | 122.8 | 35.7 | 11.6 | – | 69.9 | 75.4 | 40.7 |
| CF++ (writes detached) | 104.7 | 31.8 | 8.9 | 55.9 | 55.9 | 40.7 | |
| mutable cache (write differentiated) | out of memory in the first generator backward | ||||||
| dKV Forcing (default) | 97.3 | 33.9 | 7.0 | 62.1 | 62.1 | 29.2 | |
| dKV Forcing (fused) | 95.9 | 32.4 | 6.8 | 65.6 | 65.6 | 29.2 | |
| vs. | readout by later frames | |||
|---|---|---|---|---|
| K | V | attention mass | read-out vector | |
| relative difference | 0.35–0.37 | 0.70–0.72 | – | 0.65–0.68 |
| token-wise cosine | 0.93 | 0.73 | – | 0.76–0.77 |
| per-head linear CKA | 0.89–0.91 | 0.72–0.74 | – | – |
| checkpoint | next-frame change (other seed) | consistency, frames 0–2 | consistency, all frames |
|---|---|---|---|
| CD initialization (untrained) | 0.55 (0.97) | ||
| SGF | 0.57 (1.08) | ||
| writer LoRA, | 0.56 (1.07) | ||
| writer LoRA, | 0.57 (1.05) |
| configuration | memory gradient | residual | aesthetic | dynamic | imaging | subject |
|---|---|---|---|---|---|---|
| CF++, detached writes | none | none | ||||
| SGF | backbone | none | ||||
| writer LoRA + rollout dropout (ours) ∗ | backbone + residual | write | ||||
| writer LoRA without rollout dropout | backbone + residual | write | ||||
| separate writer network (second 1.3B copy, backbone learning rate) | backbone + writer network | none (full network) |
| VBench 5 s | VBench-Long 60 s | ||||||||
| configuration | Total | Quality | Semantic | dynamic | aesthetic | dynamic | imaging | subject | backgr. |
| frame-wise (one latent frame per segment) | |||||||||
| CF++ | |||||||||
| SGF | |||||||||
| writer LoRA | |||||||||
| + rollout dropout | |||||||||
| configuration | aesthetic | dynamic | imaging | subject | background | smoothness | flicker |
|---|---|---|---|---|---|---|---|
| frame-wise | |||||||
| CF++ | |||||||
| SGF | |||||||
| writer LoRA | |||||||
| + rollout dropout | ( ) | ||||||
| chunk-wise | |||||||
Appendix figures & tables11 assets
Supplementary material from the paper’s appendix.
Appendix
| Method | Cycle, no LoRA (s) | vs. two-pass | Cycle, LoRA (s) | vs. two-pass | Gen. peak | Critic peak | Resident | Step-0 peak | All-step peak | Reserved |
|---|---|---|---|---|---|---|---|---|---|---|
| Two-pass | 78.2 1.1 | +0.0% | 81.0 1.0 | +0.0% | 36.1 | 29.3 | 20.4 | 59.0 | 59.0 | 60.8 |
| Unoptimized | 88.8 1.2 | +13.6% | 96.3 1.9 | +18.9% | 36.4 | 23.5 | 14.7 | 20.2 | 36.4 | 58.7 |
| default | 66.9 0.9 | 14.5% | 67.9 2.8 | 16.2% | 36.1 | 23.5 | 14.7 | 20.2 | 36.1 | 58.3 |
| fused | 1.2 | 23.9% | 0.5 | 23.7% | 36.6 | 23.5 | 14.7 | 20.2 | 36.6 | 60.6 |
| hidden | 69.4 1.9 | 11.2% | 75.2 1.0 | 7.2% | 23.6 | 14.7 | 20.2 | |||
| fused+hidden | 0.6 | 20.0% | 0.7 | 19.0% | 23.6 | 14.7 | 20.2 |
| configuration | aesthetic | dynamic | imaging | subject | |
|---|---|---|---|---|---|
| CF++ (writes detached) | 4 | ||||
| SGF ( ) | 5 | ||||
| writer LoRA, | 4 | ||||
| SGF, | 4 | ||||
| SGF, | 4 | ||||
| writer LoRA, | 4 |
| configuration | aesthetic | dynamic | imaging | ||
|---|---|---|---|---|---|
| no residual | (CF++, writes detached) | 4 | |||
| 4 | |||||
| 4 | |||||
| (SGF) | 5 | ||||
| writer LoRA | 4 | ||||
| 4 |
Memory-Augmented Video Generation
Memory-augmented video generation gives video models persistent context to preserve identities, scenes, and events across long sequences.
See all 121 papers