Improving Model Representation and Reducing KV Cache via Skip Connections with First Value Heads

TL;DR

SkipV1Former addresses the memory bottleneck of KV caches in autoregressive Transformers by introducing cross-layer reuse of the first-layer Value heads: in an $L$-layer decoder, each layer $i\in\{2,\dots,L\}$ injects $H/2$ first-layer Value heads, effectively reducing Value projections and KV cache by about $0.5$ while keeping the total head count $H$ constant. The authors ground this design in a mesa-optimizer perspective, proving that using uncompressed first-layer Values accelerates deeper-layer optimization and reduces training loss by a positive constant $c>0$ on a simplified model. Empirically, SkipV1Former achieves around a $0.25$–$0.30$ KV-cache reduction with perplexity improvements on GPT‑2 and LLaMA scales, and it can be uptrained from pre-trained MHA checkpoints with only ~10–15% extra compute. The approach remains compatible with other KV-saving methods (GQA, MLA, YOCO), enabling further savings and performance gains, and a practical uptraining recipe supports efficient retrofitting to existing checkpoints for real-world deployment.

Abstract

Transformer models have driven breakthroughs across various language tasks by their strong capability to learn rich contextual representations. Scaling them to improve representation, however, often demands substantial memory and compute costs, such as the Key-Value (KV) cache used during auto-regressive decoding. Skip connections offer a promising way to improve representation without bloating resource usage, yet most prior works either improve expressivity while leaving KV costs unchanged, or reduce memory at the cost of weaker representation. In this work, we propose SkipV1Former, a Transformer variant that uses skip connections from the first layer's Value heads to strengthen model representation and reduce KV cache. Specifically, from the second block onward, each layer reuses half of its Value heads from the very first layer, while computing the other half as usual-cutting Value projections and V cache by nearly 50 \%. Theoretically, we show that routing uncompressed first-layer Values into deeper layers restores information lost to compression and accelerates the model's implicit mesa-optimization-a key pattern of Transformer in auto-regressive tasks. Empirically, across different model scales, SkipV1Former delivers consistent reductions of approximately 25 \% in KV cache while improving perplexity relative to standard Multi-Head Attention (MHA) Transformers and some advanced variants. Moreover, we propose a recipe for uptraining existing MHA Transformer checkpoints to SkipV1Former with only 10-15\% additional compute. Finally, SkipV1Former can seamlessly combine advanced methods like Group-Query Attention and Multi-Latent Attention to achieve further KV cache savings and performance improvement. When combined with YOCO, it cuts KV cache size by nearly 50 \% while still improving performance.

Figures (16)

• Figure 1: Overview of the SkipV1Former architecture. Left: A full $L$-layer decoder, where each layer from 2 to $L$ interleaves half of its Value heads with the corresponding heads from layer 1. Center: A zoomed-in view of the first and $L$-th layers, showing both the attention and the FFN sublayers. Right: Detailed illustration of the cross-layer Value skip connection, with orange shading indicating the first-layer Value heads that are injected into subsequent layers’ head sets.
• Figure 2: Standard MHA Transformer (left) compresses token pairs across layers, incurring information loss in mesa-optimization, whereas SkipV1Former (right) reintroduces raw token values into deep layers, preserving information fidelity.
• Figure 3: Left: Uptraining checkpoint conversion: we apply mean pooling over every two $W_V$ heads to reduce dimensional mismatch. Middle: Integration of SkipV1Former with GQA. Right: Integration of SkipV1Former with MLA; shaded regions follow the style of liu2024deepseek and indicate cached components.
• Figure 4: Validation loss curves across four GPT-2 model sizes (125M, 200M, 355M, 550M) for six architectures plotted over training iterations. Solid lines indicate KV-cache–efficient variants (YOCO-V, CLA-V, SkipV1Former), while dashed lines correspond to non-KV-cache-efficient models (MHA Transformer, DenseFormer, ResFormer).
• Figure 5: Validation Loss, Perplexity, and Loss Curve for LLaMA Models scaling from 60M to 3B. SkipV1Former consistently outperformed MHA Transformer with with reduced KV cache and fewer parameters.

Theorems & Definitions (10)

• Theorem 1: Informal
• Theorem 2: Formal
• Lemma 1
• proof
• Lemma 2
• proof
• Lemma 3: Lemma 3 and 4 in DBLP:conf/iclr/MahankaliH024
• Lemma 4
• proof
• proof : Proof of Theorem \ref{['thm main']}