D2O: Dynamic Discriminative Operations for Efficient Long-Context Inference of Large Language Models

TL;DR

The paper tackles the memory bottleneck of KV caches in long-context LLM inference by introducing Dynamic Discriminative Operations (D2O), which dynamically allocate per-layer KV cache sizes based on attention-density and perform EMA-guided token merging to compensate evicted information. This two-level approach preserves essential context without fine-tuning and delivers substantial memory savings and throughput gains (up to 3×) while maintaining high-quality long-text generation. Comprehensive experiments across multiple model families and long-context benchmarks demonstrate D2O’s superiority over existing eviction-based methods, including under extreme compression scenarios. The work offers a practical, plug-and-play solution for efficient long-context inference with broad applicability and potential for integration with other compression techniques.

Abstract

Generative inference in Large Language Models (LLMs) is impeded by the growing memory demands of Key-Value (KV) cache, especially for longer sequences. Traditional KV cache eviction strategies, which discard less critical KV pairs based on attention scores, often degrade generation quality, leading to issues such as context loss or hallucinations. In this work, we introduce Dynamic Discriminative Operations (D2O), a KV cache compression method that optimizes KV cache size dynamically and discriminatively at two levels without fine-tuning, while preserving essential context. At layer level, D2O leverages the varying densities of attention weights between shallow and deep layers to dynamically determine which layers should avoid excessive eviction via a novel dynamic allocation strategy to minimize information loss. At token level, D2O incorporates a compensation mechanism that maintains a similarity threshold to re-discriminate the importance of currently discarded tokens, determining whether they should be recalled and merged with similar tokens. We conduct experiments on various benchmarks and LLM architectures. Our results show that D2O not only achieves significant memory savings and enhances inference throughput by more than 3$\times$ but also maintains high-quality long-text generation.

Figures (9)

• Figure 1: Attention map density comparisons of shallow layers (layer 0, 1) and deep layers (layer 16, 30, 31) of LLaMA-2-7B on the GSM8K dataset. We use the mean value of heads for each layer.
• Figure 2: Existing methods (a) vs. $\text{D}_{2}\text{O}$ (b) at token level.
• Figure 3: Overview of $\text{D}_{2}\text{O}$. For the layer cache size allocation at the layer level, $\text{D}_{2}\text{O}$ addresses the issue of inconsistent attention density across higher and lower layers by incorporating a dynamic cache at each layer. The size of the cache is determined by the variance metric of attention and dynamic allocation strategy. At the token level, $\text{D}_{2}\text{O}$ addresses long-context information loss by incorporating a combination of a token eviction scheme and a dynamic token merging technique (where compensation mechanism is located).
• Figure 4: Variances of attention score across different layers for various models.
• Figure 5: Performance of $\text{D}_{2}\text{O}$ and other methods for LLama backbones on reasoning datasets including CoQA, GSM8K, and TruthfulQA.