Last Layer KV Cache is All You Need

One history for every layer

Transformers store separate key–value (KV) states at every layer, causing memory to scale with network depth while preventing deep representations of the past from informing shallower future layers. In LastKV, all layers instead share the last-layer KV states of completed context. This reduces long-term KV storage by approximately the number of layers and creates recurrence across both time and depth.

Scroll horizontally to view the full figure. Select the image to open it at full size.

With L layers, T tokens, and C tokens per chunk, the persistent KV ratio relative to a standard Transformer is (T + (L − 1)C) / (LT), which approaches 1/L as the number of chunks grows.

The limit assumes equal KV widths and deduplicated storage. It counts persistent KV entries; temporary buffers and training activations are excluded. Dense attention remains quadratic.

Recurrence and gradient flow

Scroll horizontally to view the full figure. Select the image to open it at full size.

Long-context language modeling

We pretrain 124M and 760M models at 32K context. Tail loss measures the 30K+ portion of the sequence; RULER accuracy averages evaluations at 4K, 8K, 16K, and 32K. LastKV improves average retrieval at both scales.

Selected baselines; RULER in %. Bold marks the best value in each column. At 124M, LastKV’s tail loss is higher than Transformer; at 760M, its tail loss is lower and the average retrieval advantage is small at the displayed precision. The paper reports all baselines and per-length RULER results.

State tracking and memory retention

LastKV is best or tied-best on five of the seven synthetic task families in the full comparison. The paired results below report token accuracy at 2K context.

Token accuracy (%), higher is better. Bold marks the best value among all methods in the paper’s full comparison.

Novel view synthesis

Sharing last-layer KV improves reconstruction PSNR on both evaluated datasets.

PSNR in dB. Bold marks the best value in each column. † The Objaverse LastKV result uses logged in-training evaluation with random view draws. Other NVS results use fixed-view full-test evaluation.

Loss and reconstruction curves

Scroll horizontally to view the full figure. Select the image to open it at full size.

Scroll horizontally to view the full figure. Select the image to open it at full size.

Autoregressive video generation

LastKV improves VBench total scores at both 1.3B and 14B after teacher-forced fine-tuning.

VBench scores (%), from 4,730 generated videos: 946 prompts × five samples. Bold marks the better value at each size. The paper provides the full VBench breakdown and user-study protocol.

Less cache, with measured training costs

LastKV substantially reduces physical self-attention KV allocation after prefill and cached decoding. The comparisons use matched workloads, precision, and A100-SXM4-80GB hardware. LLMs train at 32K and use a 128K inference prefix with 4K chunks.

Scroll horizontally to view the full figure. Select the image to open it at full size.

Absolute training time and physical self-attention KV allocation; bold marks the lower value in each pair. LLM decode KV follows 128 additional tokens. The 1.5B and 3B training fixtures use matched random initialization to measure runtime.

What the results establish

Sharing the last-layer history reduces persistent KV storage toward 1/L while LastKV matches or outperforms standard Transformers across the four task domains in the reported settings. The gains are not uniform: at 124M, LastKV’s long-context tail loss is higher than Transformer’s, and at 760M its average retrieval advantage is small at the displayed precision. The Objaverse LastKV result uses in-training evaluation with random view draws; the other NVS results use the fixed-view protocol.

Physical KV savings are not total GPU-memory savings, and training time increases in the reported LLM and video configurations. Dense attention remains quadratic, and broader gradient connectivity alone does not establish better optimization.

Citation

@misc{cai2026lastkv,
  title  = {Last Layer {KV} Cache is All You Need},
  author = {Cai, Zefan and Bi, Sai and Zhang, Lin and Hong, Yicong and
            Zhu, Yuchen and Li, Lu and Zhao, Haozhe and Jiang, Hanwen and
            Hu, Qinzhe and Huang, Zeyi and Lin, Shawn and Luo, Cheng and
            Ge, Chongjian and Jian, Jinlin and Zhu, Ruijie and Huang, Xun and Gu, Jiuxiang and
            Hu, Junjie and Tan, Hao},
  year   = {2026},
  note   = {Preprint},
  url    = {https://zefan-cai.github.io/LastKV.github.io/}
}