PaperlayerView exact revision on arXiv

Source-linked AI summary

Locret: Enhancing Eviction in Long-Context LLM Inference with Trained Retaining Heads on Consumer-Grade Devices

Yuxiang Huang, Binhang Yuan, Xu Han, Chaojun Xiao, Zhiyuan Liu

arXiv:2410.01805v2cs.CL

TL;DR

Long-context scaling raises attention computation and KV-cache memory demands, while existing compression methods provide limited memory reduction for consumer-grade inference. LOCRET uses learnable retaining heads to estimate each cache unit’s causal importance and evict low-importance units during chunked prefill. LOCRET outperforms reported baselines in memory efficiency and generation quality, achieving over 20× and 8× KV-cache compression for Phi-3-mini-128K and Llama-3.1-8B-instruct.

  • Problem

    Long-context scaling raises attention computation and KV-cache memory demands, while existing compression methods provide limited memory reduction for consumer-grade inference.

  • Method

    LOCRET uses learnable retaining heads to estimate each cache unit’s causal importance and evict low-importance units during chunked prefill.

  • Results

    LOCRET outperforms reported baselines in memory efficiency and generation quality, achieving over 20× and 8× KV-cache compression for Phi-3-mini-128K and Llama-3.1-8B-instruct.

  • Takeaways & Limitations

    LOCRET enables full comprehension of long contexts on consumer-grade devices without compromising generation quality, with less than 1 GPU hour of offline training.

  • Takeaways & Limitations

    Future evaluations are planned for encoder-decoder and multi-latent architectures and devices such as NVIDIA Jetson.

Abstract

from arXiv · show

Scaling the input context length of a large language model (LLM) incurs a significant increase in computation cost and memory footprint to maintain the attention key-value (KV) cache. Existing KV cache compression methods suffer from inefficient compression strategies and limited memory reduction effects, making it difficult for LLMs to conduct long-context inference on consumer-grade devices, especially when inferring long-context stream input. Such obstacles prevent consumer-grade devices from supporting more complex applications, creating challenges for the democratization of LLMs. To overcome this, we propose Locret, the first framework to create an eviction policy compatible with chunked prefill. By evaluating the causal importance of KV cache units by learnable retaining heads, Locret enables precise eviction of cache units, facilitating efficient long-context inference. In our extensive empirical studies, Locret outperforms the recent popular and competitive approaches in terms of memory efficiency and generation quality -- Locret achieves up to 20x of KV cache compression ratio within less than 10% performance loss. Furthermore, Locret achieves 128K+ long-context inference on a single NVIDIA 4090 GPU without compromising generation quality and only costs <1 GPU hour of additional training.

1. Introduction

Long-context LLM inference is constrained by attention computation and KV-cache memory growth, motivating LOCRET’s trained eviction policy for efficient inference on consumer-grade devices. LOCRET estimates token importance with retaining heads, integrates eviction into chunked prefill, and reports strong compression and generation-quality results.

  • Motivation: Long-context scaling increases both attention computation and KV-cache memory, while backbone optimizations do not address attention efficiency or cache growth.The KV cache continues growing linearly with context length, motivating attention- and cache-focused methods.
  • Method: LOCRET uses learnable retaining heads to estimate causal importance and select cache units for eviction during chunked prefill.The framework is trained offline at a cost below 1 GPU hour.
  • System: The inference implementation evicts low-CIS cache units during chunked prefill, limiting GPU memory usage and accelerating prefill with minimal additional hardware requirements.The system is designed to remain compatible with decoder-only LLMs.
  • Evaluation: LOCRET is reported to achieve comparable performance with a full KV cache while maintaining inference efficiency.The evaluation covers different models and multiple datasets.
  • Results: LOCRET achieves over 20× KV cache compression for Phi-3-mini-128K and 8× for Llama-3.1-8B-instruct.It also enables query-driven prefill acceleration through LOCRET-Q without significant performance degradation.

2. Related Work

Related work reduces long-context inference costs through algorithmic compression and system-level memory management. These approaches include quantization, sparsity, token dropping, pooling, and CPU offloading, but incur implementation overheads or leave cache-growth challenges.

  • Algorithm Optimizations: Algorithm optimizations reduce KV-cache size through quantization, sparsity, or token dropping.Quantization can add operator-customization and inverse-quantization overhead, while sparsity methods exploit attention-head patterns.
  • Open Challenge: Quantization and sparsity methods show promising efficiency gains but do not remove the challenge that KV-cache size grows with context length.This motivates methods focused on selective cache eviction and other compression strategies.
  • Token Compression: Eviction and pooling methods maintain a static cache size by dropping tokens or merging adjacent cache units, but may struggle to identify accurate victims.Pooling methods use predesigned transformations and assign higher weights to more important units.
  • System Optimizations: Offloading methods store KV cache in CPU memory and retrieve relevant chunks to GPU memory, reducing peak GPU usage at the cost of CPU–GPU communication overhead.They are system-level methods that account for hardware features.

3. Methodology of LOCRET

LOCRET combines chunked prefill with trained retaining heads that estimate causal importance and evict lower-value KV cache units while preserving stabilizers. Its training and inference design aims to maintain generation quality under constrained cache capacity.

  • Framework: LOCRET appends a retaining head to each attention module, trains these heads while freezing the LLM backbone, and predicts cache-unit importance during chunked-prefill inference.The retaining heads are small MLPs that estimate causal importance scores for cache units.
  • Importance estimation: Causal importance scores estimate each cache unit’s influence on comprehending subsequent context using only preceding units and the unit itself.The approximation is intended to preserve core information while minimizing attention errors introduced by eviction.
  • Training Retaining Heads: The retaining heads are trained on question-answer supervised data with attention-derived CIS targets, combining Smooth-L1 regression with an adjacent-value L2 smoothing loss.For GQA models, ground truth uses the maximum pre-softmax attention score across query heads within each KV-head group.
  • Robustness and framework notation: LOCRET’s training is reported as robust to changes in retaining-head dimension and training dataset, with minimal performance variation.The framework’s notation includes retaining heads, chunked-prefill time steps, cache budget size, and stabilizer length; the figure omits layer structure for simplicity.
  • Cache eviction: During inference, the system stores predicted causal importance values and evicts lower-scoring cache units whenever the cache reaches its budget.Eviction occurs during chunked prefill, using the predicted scores to select units for retention.
  • Stabilizers: Stabilizers retain the last ns cache units of each chunk to preserve local continuity, because absent or insufficient stabilizers destabilize CIS prediction and can severely degrade performance.Larger stabilizer lengths can also reduce space available for other cache units in some NLU settings.

4. Experiments

Experiments evaluate LOCRET across long-context benchmarks, memory usage, consumer-device speed, and query-driven tasks. LOCRET generally achieves the strongest overall performance with reduced memory, while LOCRET-Q addresses query-dependent retrieval more effectively than competing eviction methods.

  • End-to-end Benchmark: LOCRET uses the least memory while achieving the best overall performance, whereas INFLLM and MINFERENCE retain the full KV cache.Quantization methods can reduce memory in some settings but introduce severe degradation, while MINFERENCE does not compress the KV cache.
  • End-to-end Benchmark: LOCRET achieves the highest overall scores on both ∞Bench and L-Eval while using less GPU memory than competing methods.On ∞Bench, it performs strongly on retrieval and comprehensive tasks; on L-Eval, it achieves the best overall performance and highest scores on most tasks.
  • LOCRET-Q: Supporting Query-Driven Tasks: LOCRET-Q outperforms other eviction-based methods and even MINFERENCE on 128K-context RULER, while eviction methods with chunked prefill exceed 2× prefill speedup over FULLATTN.Standard LOCRET fails on RULER, whereas LOCRET-Q’s query-aware eviction preserves query-correlated critical regions.

5. Conclusion

LOCRET enables memory-efficient long-context inference by predicting cache-unit causal importance during chunked prefill and accurately evicting KV entries. Its query-aware variant extends this approach to query-centric tasks, while future work targets broader architectures, devices, and cache-budget integration.

  • LOCRET predicts cache-unit causal importance during chunked prefill, enabling accurate KV-cache eviction for memory-efficient long-context inference.The method is evaluated across models and datasets against major efficient-inference techniques without CPU-memory offloading.
  • LOCRET outperforms the evaluated baselines using less GPU memory without noticeable generation-quality degradation.
  • LOCRET-Q processes query-centric tasks without significant performance degradation.
  • Future work will test LOCRET on encoder-decoder and multi-latent models, additional devices such as NVIDIA Jetson, and integrated KV-cache budget allocation methods.

A.1. Training

The appendix specifies LOCRET’s training, inference settings, hardware environments, and baseline implementations. These details cover retaining-head training, cache-budget parameters, and comparisons against full attention, offloading, quantization, token dropping, and sparse attention.

  • A.1. Training: Retaining heads are trained with AdamW on LongAlpaca for 3,000 steps using batch size 1, maximum sequence length 10,240, and learning rate 5e-4.Each retaining head is a small two-linear-layer FFN with intermediate size 1,024 and model-aligned nonlinearities.
  • A.2. Inference: LOCRET inference uses cache budget b, chunk size B, stabilizer length ns, and local-token length nloc as its principal hyperparameters.
  • A.3. Hardware: Experiments use multi-GPU workstations for most evaluations, while the Section 4.3 experiments run on a single NVIDIA 4090 GPU.
  • A.4. Baselines: LOCRET is compared with full attention, CPU-offloaded INFLLM, 2-bit KV quantization, entropy-based SIRLLM token dropping, and MINFERENCE sparse attention.INFLLM, HF-2BITS, and SIRLLM required implementations for Phi-3-mini-128K because official implementations were unavailable.

B. The Global and Local Discrepancy of Scoring Functions

Causal scoring avoids the future-information dependence that makes non-causal methods inconsistent under chunked prefill. Experiments show that these discrepancies can cause severe retrieval failures, whereas LOCRET maintains superior end-to-end performance.

  • The Global and Local Discrepancy of Scoring Functions: Causal scoring functions do not rely on subsequent cache units, making them compatible with chunked prefill, unlike non-causal H2O and SNAPKV.SIRLLM is causal but still shows notable inaccuracies and associated performance degradation.
  • The Global and Local Discrepancy of Scoring Functions: Future-dependent H2O and SNAPKV scores differ substantially between prefixes and the full context, causing severe chunked-prefill performance drops on ∞Bench.The discrepancy prevents these methods from identifying some cache units’ importance when those units are first encountered.
  • Evaluation Setup: On RULER’s 128K-token contexts, evaluation compares eviction methods with full attention and MINFERENCE using performance, prefill speed, and decode speed.
  • LOCRET Inference: LOCRET processes prefix chunks, protects stabilizer caches, and retains the top-b cache units by score after each chunk.The algorithm leaves the final nloc tokens out of eviction and assigns stabilizer scores of +∞ before top-b selection.
  • Training Effectiveness: Training retaining heads is necessary: randomly initialized heads make random predictions, evict arbitrary cache units, and fail on all evaluated tasks.

F. Evaluation on LongBench

On English LongBench subtasks with Phi-3-mini-128K, LOCRET achieves the best overall performance while using less memory than the evaluated alternatives. The comparisons also show distinct trade-offs for MInference, SirLLM, and InfLLM.

  • F. Evaluation on LongBench: The evaluation excludes Chinese LongBench subtasks because Phi-3-mini-128K was not specifically trained on Chinese corpora.
  • F. Evaluation on LongBench: LOCRET achieves the best overall LongBench performance and excels on most subtasks without noticeable degradation while consuming less memory.
  • F. Evaluation on LongBench: MInference avoids performance drops but requires more GPU memory than LOCRET.
  • F. Evaluation on LongBench: SirLLM has comparable memory usage but shows performance decline relative to full attention and LOCRET.
  • F. Evaluation on LongBench: InfLLM exhibits the largest performance drop and highest total memory consumption because of its CPU-offloading mechanism.

G. Hyperparameter Analysis

LOCRET remains robust across budget, stabilizer-length, and chunk-size settings, with performance improving faster than SNAPKV as more cache budget is available.

  • Budget: LOCRET improves faster than SNAPKV as the cache budget increases.The comparison uses chunked prefill on LongBench.
  • Stabilizers Length: Performance remains consistent for small stabilizer lengths, while larger values degrade performance by reducing space for other cache units.
  • Chunk Size: LOCRET maintains stable performance across chunk sizes from 256 to 4096 on NQ from L-Eval.

H. Orthogonality to Other Methods

LOCRET is designed to combine with quantization, token merging, head-wise budget allocation, and other compression methods, while retaining strong performance and robustness across settings.

  • KV cache quantization: LOCRET’s quantization-induced performance drop is only slightly higher than full attention’s, indicating compatibility with KV cache quantization.The evaluation uses Quanto and Llama-3.1-8B-instruct on average L-Eval scores.
  • Token merging: LOCRET can combine with token merging, allowing evicted tokens to be merged into a small cache pool retained in GPU memory.
  • Head-wise Budget Allocation: LOCRET is compatible with head-wise budget allocation and is evaluated in combination with PYRAMIDKV on ∞Bench.
  • Retaining-head size: Across retaining-head intermediate sizes from 256 to 4096, performance variations are minimal and surpass all baselines in Table 1.
  • Training settings: Training-dataset changes have minor effects on overall performance, supporting LOCRET’s robustness to training settings.

I.2. Training Data Insensitivity

LOCRET’s performance is largely insensitive to the choice of training dataset, achieving competitive results without delicate data selection.

  • Training Data Insensitivity: Different training-dataset recipes have minor effects on overall performance.The comparison uses L-Eval with Phi-3-mini-128K and aligned training hyperparameters.
  • Training Data Insensitivity: LOCRET can obtain competitive performance without carefully selecting the training data.

J. Extremely Long Context Evaluation

LOCRET supports extremely long-context processing through causal importance-based cache selection, with retained patterns varying by task and layer while retaining-head inference adds no significant latency.

  • Extremely Long Context Evaluation: 1747.6× compression enables LOCRET to process extremely long contexts with a cache budget of 6000.The evaluation uses contexts averaging 10 million tokens and a chunk size of 10240.
  • Query-independent eviction: LOCRET’s cache importance is based on the cache itself rather than subsequent queries, making it suitable for multi-turn conversation scenarios.The method is compared with query-aware eviction methods such as SNAPKV on Rock-Paper-Scissors.
  • Stabilizers: Retaining stabilizers is necessary because cache eviction causes context discontinuity, unstable CIS prediction, and inaccurate later-token calculations.
  • Inference Speed: Retaining heads introduce no significant inference latency increase.Observed numerical differences are attributed to systematic variation rather than additional retaining-head overhead.
  • Causal Importance Score: Causal importance scoring lets LOCRET’s top-b selection simulate a cache problem under a fixed budget.The formulation separates causal calculation into selection and subsequent computation functions.
  • Retained Patterns: LOCRET’s retained patterns depend on task, head, and layer, with inserted answer-related positions receiving strong CIS signals in R.Number.Shallow layers show periodic patterns, whereas middle layers more strongly highlight semantic positions.
Loading 2410.01805v2…