YOCO++: Enhancing YOCO with KV Residual Connections for Efficient LLM Inference

1. Problem Statement

Large Language Models (LLMs) based on the Transformer architecture face a significant bottleneck during inference: the Key-Value (KV) cache. As sequence lengths grow, the memory required to store the KV cache becomes a limiting factor for batch size and throughput.

To address this, Cross-layer KV compression methods have been proposed. These methods reduce memory consumption by retaining KVs for only a subset of layers and reusing them in others. A prominent method, YOCO (You Only Cache Once), shares the KVs of the middle layer with all top-half layers, achieving a 50% reduction in KV cache size.

The Core Issue: While cross-layer sharing improves efficiency, it introduces non-negligible performance degradation. Directly sharing KVs from one layer to another weakens the model's capacity because distinct layers lose their ability to learn unique representations. Existing solutions like FusedKV attempt to mitigate this by using weighted combinations of KVs from multiple layers but introduce significant cache I/O overhead, which degrades inference throughput.

2. Methodology: YOCO++

The authors propose YOCO++, an enhanced version of YOCO that improves model performance without sacrificing inference efficiency or introducing extra I/O overhead.

Core Mechanism: Weighted KV Residual Connections

Instead of directly sharing KVs, YOCO++ introduces a weighted residual connection between the KVs of the bottom layer and the KVs of each subsequent bottom-half layer.

• Architecture: The model consists of $L$ layers. The bottom $L/2$ layers are "self-decoders," and the top $L/2$ are "cross-decoders."
• KV Combination: In each bottom-half layer $i$ (where $i \in \{2, ..., L/2\}$), the cached KVs ($K'_i, V'_i$) are computed as a weighted sum of the current layer's KVs ($K_i, V_i$) and the bottom layer's KVs ($K_1, V_1$):
  $$K'_i = \lambda(\alpha^K_{i,1}K_1 + \alpha^K_{i,2}K_i)$$
  $$V'_i = \lambda(\alpha^V_{i,1}V_1 + \alpha^V_{i,2}V_i)$$
  Where $\alpha$ are learnable residual weights, and $\lambda$ is a scaling factor.
• Top-Half Reuse: The top-half layers reuse the modified KVs from the middle layer ($K'_{L/2}, V'_{L/2}$), which now carry richer information than the original YOCO middle-layer KVs.

Key Technical Innovations

1. Cache Strategy: To avoid the I/O overhead seen in FusedKV (which requires fetching multiple layers' KVs during decoding), YOCO++ caches the combined KVs ($K', V'$) instead of the original ones. This ensures that during the decoding stage, the model only accesses one set of KVs per layer, maintaining the same I/O profile as YOCO.
2. Scaling Factor ($\lambda$): The authors observed that without a scaling factor, the model struggles to learn the residual weights effectively (they remain close to initialization). They introduce a fixed hyperparameter $\lambda > 1$ (empirically set to 35) to amplify the residual signal, facilitating better learning of the weights.
3. Initialization: The weights are initialized such that $\alpha_{i,1}=0$ and $\alpha_{i,2}=1/\lambda$, ensuring the model starts in a state identical to standard YOCO.

3. Efficiency Analysis

The paper rigorously analyzes the efficiency of YOCO++ compared to YOCO and FusedKV:

• Parameters: Negligible increase (only 4 learnable weights per bottom-half layer).
• Computation: Minimal overhead; the weighted combination is trivial compared to the attention mechanism.
• Memory: Identical to YOCO. Both cache $L/2$ layers.
• I/O Overhead:
  • Training/Prefill: Slightly higher I/O due to accessing $K_1, V_1$ and $K_i, V_i$, but this stage is compute-bound, so latency is unaffected.
  • Decoding: No additional overhead. Since the combined KVs are pre-computed and cached, the decoder only reads one KV set per layer, matching YOCO's throughput.

4. Experimental Results

Experiments were conducted on a 1.1B parameter model (TinyLlama configuration) using the SlimPajama dataset.

• Inference Efficiency (Figure 3):

  • Prefill Latency: YOCO++ matches YOCO and FusedKV, achieving 50% lower latency than the standard Transformer due to early exit.
  • Decoding Throughput: YOCO++ maintains high throughput comparable to YOCO and significantly outperforms the standard Transformer. Unlike FusedKV, whose throughput drops to Transformer levels due to I/O bottlenecks, YOCO++ remains efficient.

• Model Performance (Table 1 & 2):

  • Zero-Shot Accuracy: YOCO++ achieves State-of-the-Art (SOTA) performance among cross-layer compression methods at a 50% compression rate.
  • Comparison: It outperforms the standard Transformer, YOCO, FusedKV, and FusedKV-Lite on commonsense reasoning benchmarks (HellaSwag, OBQA, WinoGrande, ARC, PIQA).
  • Training Loss: YOCO++ achieves the lowest validation loss (2.4194), outperforming even the standard Transformer (2.4266).

• Ablation Studies:

  • Scaling Factor: Removing $\lambda$ causes performance to drop to YOCO levels, confirming its necessity for learning.
  • Key Residuals: Removing the key residual connection ($K$) leads to performance degradation, proving that residual connections for both Keys and Values are beneficial.

5. Key Contributions

1. YOCO++ Architecture: A novel method that enhances cross-layer KV compression by introducing weighted residual connections between the bottom layer and bottom-half layers.
2. Efficiency Preservation: Successfully eliminates the I/O overhead associated with previous "fusion" methods (like FusedKV) by caching the combined representation, maintaining the high throughput of YOCO.
3. Scaling Factor: The introduction of a scaling factor $\lambda$ to solve the optimization difficulty of learning residual weights in deep networks.
4. SOTA Performance: Demonstrates that it is possible to compress KV caches by 50% while outperforming the standard uncompressed Transformer in both training loss and downstream task accuracy.

6. Significance

This work addresses the critical trade-off between LLM inference efficiency and model capability. By showing that cross-layer compression can not only match but exceed the performance of standard Transformers without incurring I/O penalties, YOCO++ provides a highly practical solution for deploying large models on resource-constrained hardware. It suggests that "residual connections" can be effectively applied to the KV cache domain to preserve information flow across layers, offering a new direction for efficient LLM design.