Last Layer Key-Value Cache is All You Need
Zefan Cai ⋅ Sai Bi ⋅ Lin Zhang ⋅ Yicong Hong ⋅ Chongjian GE ⋅ Xun Huang ⋅ Jiuxiang Gu ⋅ Junjie Hu ⋅ Hao Tan
Abstract
The key-value (KV) cache of a long-context Transformer grows with both sequence length and network depth: every layer stores its own keys and values for every retained token. Cross-layer KV sharing shrinks the depth dimension by letting several layers read one cache, but existing designs either preserve quality only at modest sharing factors among adjacent layers, or condense history at every token position, which breaks exact parallel training or requires a redesigned decoder-decoder architecture. We push cross-layer sharing to its limit under a chunked design. \method{} processes each chunk as a standard per-layer-KV Transformer, writes only the last layer's KV cache when the chunk completes, and lets every layer of every future chunk read this single shared cache through ordinary attention. Older history therefore keeps one cache instead of $L$: the cross-chunk cache drops from $O(TCLd)$ to $O(TCd)$ ($O((T+L)Cd)$ with a nearest-chunk refinement for boundary continuity). Because sharing applies only across completed chunks, within-chunk computation and training parallelism are unchanged: training proceeds chunk by chunk with ordinary backpropagation and no iterative approximation. We separate this cache accounting from the quality question: can all layers actually work from one shared last-layer cache? Across language modeling, novel view synthesis, and autoregressive video generation, controlled comparisons show that maximal cross-layer sharing of older history, with the nearest-chunk refinement enabled for language and video, preserves or improves task quality: it lowers language-model validation loss and stays competitive on RULER retrieval, raises novel-view PSNR, and maintains comparable aggregate VBench scores. Continued pretraining further converts an openly released 8B Transformer into \method{}, which retains one layer of KV out of 32 while scoring like the unconverted baseline across twelve benchmarks.
Chat is not available.
Successful Page Load