PaperlayerView exact revision on arXiv

Source-linked AI summary

Dynamic Memory Compression: Retrofitting LLMs for Accelerated Inference

Piotr Nawrot, Adrian Łańcucki, Marcin Chochowski, David Tarjan, Edoardo M. Ponti

arXiv:2403.09636v2cs.CL

TL;DR

Autoregressive generation stores past-token keys and values in a KV cache whose growth with sequence length and batch size burdens memory. DMC learns to append or merge KV representations during inference and retrofits existing LLMs without adding parameters. It retains original performance at up to 4× compression and achieves up to 7× throughput increase on an H100 GPU.

  • Problem

    During generation, the KV cache grows with sequence length and batch size, making memory-intensive inference especially challenging for long contexts and concurrent requests.

    Source: §1 · p. 1

  • Method

    DMC retrofits existing LLMs to learn, for each new token, whether to append its key–value representations or merge them with the top cache element.

    Source: §7 · p. 10

  • Results

    At up to 4× compression, DMC retains or improves original LLM performance and achieves up to 7× throughput increase on an H100 GPU.

    Source: §7 · p. 10

  • Takeaways & Limitations

    DMC is suitable as a drop-in replacement for vanilla KV caching to achieve higher inference efficiency.

    Source: §5.1 · p. 6

    The paper focuses on retrofitting existing LLMs; preliminary DMC-from-scratch experiments had negative results compared with GQA’s training curve.

    Source: H. Limitations · p. 18

Abstract

from arXiv · show

Transformers have emerged as the backbone of large language models (LLMs). However, generation remains inefficient due to the need to store in memory a cache of key-value representations for past tokens, whose size scales linearly with the input sequence length and batch size. As a solution, we propose Dynamic Memory Compression (DMC), a method for online key-value cache compression at inference time. Most importantly, the model learns to apply different compression ratios in different heads and layers. We retrofit pre-trained LLMs such as Llama 2 (7B, 13B and 70B) into DMC Transformers, achieving up to 7x throughput increase during auto-regressive inference on an NVIDIA H100 GPU. DMC is applied via continued pre-training on a negligible percentage of the original data without adding any extra parameters. DMC preserves the original downstream performance with up to 4x cache compression, outperforming up-trained grouped-query attention (GQA) and key-value eviction policies (H$_2$O, TOVA). GQA and DMC can be even combined to obtain compounded gains. Hence, DMC can serve as a drop-in replacement for KV caching in existing LLMs to fit longer contexts and larger batches within any given memory budget.

1. Introduction

Generation stores past-token keys and values in a cache that grows with sequence length and batch size, creating a substantial memory burden. DMC compresses this cache online and retains downstream performance at high compression while improving inference throughput.

1. Introduction

  • The KV cache stores past-token keys and values to avoid recomputing them, but its growth with sequence length and batch size makes generation memory-intensive.The burden is particularly apparent in long-context generation and when serving many user queries.Source: §1 · p. 1
  • DMC chooses at each step whether to append the current key and value representations or merge them with the top cache item using a weighted average.With an appropriate prior over compression decisions, DMC memory can grow sub-linearly.Source: §1 · p. 2
  • DMC retrofits Llama 2 7B, 13B, and 70B using a negligible share of original pretraining data and adds no parameters.The reported data amounts are approximately 2% for 2× compression and 8% for 8× compression.Source: §1 · p. 2
  • DMC retains downstream performance similar to the original LLM while GQA, H2O, and TOVA incur significant degradation at high compression; GQA and DMC can also combine for 16× compression.The 16× total is for Llama 2 70B pretrained with GQA 8× and then compressed with DMC 2×.Source: §1 · p. 2
  • Up to 700% throughput gains with DMC 8× were observed, with an approximately 5% relative MMLU performance drop for Llama 2 7B and 13B.At DMC 4×, throughput increased by 350%–390% for Llama 2 7B and 13B without performance loss on NVIDIA H100 or A100 GPUs.Source: §1 · p. 2

2. Background

The background explains how multi-head attention operates and how autoregressive inference reuses cached keys and values. It also relates cache size to memory use and generation latency.

2.1. Multi-Head Self-Attention

  • Multi-head self-attention splits embeddings into heads that process attention separately, then concatenates and linearly transforms their outputs.Each head uses its own query, key, and value projections.Source: §2.1 · p. 2

2.2. KV Caching During Inference

  • During autoregressive inference, each step computes keys and values for the current token and retrieves earlier ones from the KV cache; the cache grows with every new token.Queries need not be cached because only the current token’s query is needed at each inference step.Source: §2.2 · p. 2

2.3. Memory-Bound and Compute-Bound Operations

  • Autoregressive generation is typically memory-bound, and attention latency scales linearly with KV-cache size.The attention computation remains memory-bound regardless of batch size, unlike linear layers, which can become compute-bound at sufficiently large batch sizes.Source: §2.3 · pp. 2–3

3. Method: Dynamic Memory Compression

Dynamic Memory Compression (DMC) changes inference-time KV-cache updates and trains pretrained models to reproduce those updates. Its training procedure handles discrete compression decisions and the mismatch between sequential inference and parallel training.

3.1. Inference

  • DMC independently applies its cache-update algorithm to each attention head, allowing KV sequences to have different lengths across heads.At each step, the model predicts whether to append or accumulate the current key and value representations.Source: §3.1 · p. 3

3.2. Training

  • The model uses a continuous relaxation of its discrete append-or-accumulate decisions during training, then rounds the decision score for inference.Stochastic reparameterization supports gradient estimation, and training starts with no compression to match vanilla Transformer behavior.Source: §3.1 · p. 3, §3.2 · p. 4
  • During training, all intermediate key and value states are retained, while an additive mask simulates inference-time access to the evolving cache.The mask is based on the sequence of decision scores and blocks access to stale intermediate states and future states.Source: §3.2 · p. 4, Figure 3 · p. 4, Figure 2 · p. 4
  • A global one-sided loss compares the total decisions with the expected number of KV tokens at the target compression ratio.The compression loss is added to the language-modeling loss.Source: §3.2 · pp. 4–5

3.3. Practical Considerations

  • Window grouping reduces partial-accumulation computation from O(n) sequential steps to O(w), but may require caching the window during inference.This trade-off applies to heads that accumulate more than w tokens.Source: §3.3 · p. 5

4. Experimental Setup

The experiments retrofit Llama 2 models at three sizes and compare DMC with GQA and cache-eviction baselines. They evaluate downstream tasks and use continued pre-training schedules tied to the target compression ratio.

4. Experimental Setup

  • Experiments use Llama 2 7B, 13B, and 70B, comparing DMC with GQA and the TOVA and H2O cache-eviction policies.GQA checkpoint conversion averages key and value projection sub-matrices within each group, yielding fixed compression during training.Source: §4 · p. 5
  • DMC retrofitting reuses the first dimensions of query and key representations to predict segmentation and importance scores without adding parameters.A one-billion-token initial phase trains the model to disregard those dimensions in attention before the main compression phase.Source: §4 · p. 5
  • Continued pre-training uses 24B, 72B, 120B, and 168B tokens for target compression ratios of 2×, 4×, 6×, and 8×, respectively.The target ratio is increased gradually; the 70B model is not up-trained to 4× because it is already pretrained with GQA 8×.Source: §4 · p. 6
  • Evaluation covers MMLU, HumanEval, and common-sense question-answering benchmarks, using 5-shot, pass@1, and average 0-shot scores, respectively.The common-sense set includes PIQA, BoolQ, Arc-C, Arc-E, HellaSwag, and WinoGrande.Source: §4 · p. 6

5. Results

DMC generally preserves Llama 2’s downstream performance under compression and outperforms GQA and cache-eviction baselines in the reported comparisons. Measurements show throughput gains alongside latency trade-offs, while learned compression ratios favor deeper layers.

5.1. Main Results

  • DMC remains close to or exceeds original Llama 2 performance across compression settings; for 7B, MMLU drops are -1.7 at 6× CR and -2.8 at 8× CR.At 2× CR, DMC increases MMLU and CS-QA for 7B and 13B and HumanEval for all tested scales; at 4×, reported drops are small.Source: §5.1 · p. 6, Table 1 · p. 7
  • At equal 2×/4× CR, DMC's MMLU advantage over GQA grows from +5.4 to +9.2 for 7B and from +4.6 to +5.6 for 13B.DMC also outperforms GQA at 2× and 4× across both compared model scales; for CS-QA and HumanEval, gains are comparable across compression ratios and scales.Source: §5.1 · p. 6, Table 1 · p. 7
  • At 4×–8× CR, H2O and TOVA drop substantially in accuracy relative to DMC, especially on HumanEval; their post-softmax scoring requires an n^2 attention-score tensor.At 2×, the eviction policies are almost comparable to DMC on MMLU and CS-QA.Source: §5.1 · p. 7
  • For Llama 2 70B pretrained with 8× GQA, adding 2× DMC yields 16× total cache compression with unchanged performance.The 70B model combines the pre-existing GQA compression with additional DMC compression during retrofitting.Source: §5.1 · pp. 7–8, Table 1 · p. 7
  • At 2× CR, GQA cannot reach DMC's 8K-step performance even after more than 17K fine-tuning steps, for both 7B and 13B.A final 2K-step phase at fixed CR recovers a significant amount of original performance, especially at higher target CRs such as 4×.Source: Figure 4 · p. 7, §5.1 · p. 8

5.2. Throughput and Latency Measurements

  • DMC reaches up to 7× throughput at 8× CR; throughput scaling becomes sub-linear once larger batches make models compute-bound.Reported maxima are 2.0× for DMC 2×, 3.9× for DMC 4×, and 5.5× for DMC 6×.Source: Figure 5 · p. 8, §5.2 · p. 8
  • With DMC 4× at the original model's batch size, longer-context latency improves as the KV cache footprint shrinks.At the largest batch size for either model, inference time begins scaling linearly with context after approximately 2200 generated tokens.Source: §5.2 · pp. 8–9

5.3. Per-Head Learned Compression Ratios

  • DMC favors compressing deeper layers (>16 for 7B, >22 for 13B, >44 for 70B), while compressing the final layers less.Above 4× CR, some early-layer heads also receive very high compression; during training, compression increases first in deeper layers and then in non-contiguous intermediate ranges.Source: §5.3 · p. 9

6. Related Work

Prior work reduces KV-cache memory through shared KV heads, latent projections, token merging, restricted attention, quantization, or eviction. DMC differs through adaptive online merging and retrofitting of pretrained models.

6. Related Work

  • KV-cache reduction approaches include shared KV heads, low-rank latent projections, token merging or pooling, sliding windows, quantization, and cache eviction.GQA shares KV representations across query-head subsets; sliding-window methods cap attention to preceding tokens, while quantization reduces KV precision.Source: §6 · pp. 9–10
  • Unlike fixed-size compression approaches, DMC adapts its compression schema to the input sequence.The paper reports superior results for this adaptation in an ablation study in Appendix F.Source: §6 · p. 10
  • DMC merges tokens rather than discarding them, has constant context-length complexity, and compresses the full sequence online, including generated output.The cited learnable token-discarding method has linear complexity; prompt compression by Mu et al. uses costly generation and applies only to model input.Source: §6 · p. 10
  • Unlike heuristic eviction strategies, DMC learns compression end-to-end; it also extends Dynamic Token Pooling to KV-cache compression and retrofits pretrained LLMs.DMC introduces a continuous relaxation of the pooling decision during training rather than training language models from scratch.Source: §6 · p. 10

7. Conclusions

DMC reduces Transformer KV-cache length through learned online merging, enabling inference-time efficiency gains. Up to 4× compression preserves or improves original-model performance, while higher compression brings limited degradation.

7. Conclusions

  • DMC learns for each new token whether to append its key–value representations to the cache or merge them with the top cache item.The method is applied to Llama 2 models at 7B, 13B, and 70B scales with negligible extra data and no extra parameters.Source: §7 · p. 10
  • DMC reduces Transformer KV-cache length to improve inference efficiency.The conclusion describes retrofitting existing LLMs rather than adding parameters.Source: §7 · p. 10
  • Up to 4× compression retains or improves original LLM performance, while 6× and 8× compression bring limited degradation; DMC reaches up to 7× throughput increase on an H100 GPU.For comparable compression ratios, DMC also has higher downstream performance and sample efficiency than GQA, H2O, and TOVA.Source: §7 · p. 10

Impact Statement

The impact statement presents DMC as a possible route to lower inference costs and broader hardware access. The accompanying experiments report close replication of Llama 2 results, little effect from the retrofitting data mixture, and a corpus assembled from varied sources.

Impact Statement

  • DMC may reduce operational costs and environmental impact while making advanced models suitable for a broader range of hardware.The impact statement links this possibility to higher throughput and lower latency.Source: Impact Statement · p. 10

A. Replicating the Original Results

  • Up-training the vanilla Llama 2 checkpoint for 200 steps, or 1B tokens, nearly replicates the original paper’s results, and the internal data mixture has little downstream benchmark effect.The study averages and reports standard deviations across checkpoints at 50, 100, 150, and 200 steps.Source: A. Replicating the Original Results · p. 14

B. Retrofitting Data

  • The retrofitting corpus draws on sections of The Pile and additional sources including ArXiv, mC4, Common Crawl, PubMed Central, and The Stack.Other listed sources include Reddit, Wikipedia, CC-News, and BigScience Workshop datasets.Source: B. Retrofitting Data · p. 14

C. Analysis of the Compression Schema Learned by DMC

DMC’s learned compression schema varies across layers, sequence lengths, and token positions, with some groupings aligning with linguistic units. The constrained variant makes per-layer compression more similar, but can substantially reduce downstream performance in some settings.

C. Analysis of the Compression Schema Learned by DMC

D. Similar Per-layer Compression Rates

  • DMC-C’s auxiliary loss forces heads in a layer to maintain similar compression ratios, allowing key and value sequences to be stored as padded tensors with minimal padding.The constrained variant adds an extra loss term to Equation (11).Source: D. Similar Per-layer Compression Rates · p. 15
  • At 7B 4×, DMC-C records a -6.4 drop in MMLU compared to the ceiling; standard DMC recovers all performance loss in DMC-C.The authors prefer standard DMC with custom attention implementations that avoid excessive padding, while retaining memory-efficiency benefits.Source: D. Similar Per-layer Compression Rates · p. 15
  • DMC-C’s compression schema differs from DMC’s because its auxiliary loss constrains heads within each layer, though both favor compressing deeper layers.The discussion points to Figure 12 for the DMC-C schema.Source: D. Similar Per-layer Compression Rates · p. 15

E. Training Ablations

The training ablations examine compression-ratio schedules and training duration. They report a trade-off between training time and performance, and show that immediate compression-ratio increases can hurt downstream accuracy despite perplexity recovery.

E. Training Ablations

  • For constrained DMC, increasing the compression ratio by 1 every 3K steps (12B tokens) or 6K steps (24B tokens) exposes a trade-off between training time and performance.DMC also supports arbitrary real-valued compression ratios, unlike GQA’s integer ratios divisible by 2.Source: E. Training Ablations · p. 16
  • DMC-immediate produces an early perplexity spike and degrades MMLU accuracy across the board, even though perplexity recovers to below the annealed DMC point.The result supports avoiding perplexity spikes when retrofitting an LLM.Source: E. Training Ablations · p. 16

F. DMC Ablations

The ablations favor dynamically learned, globally allocated compression over fixed pooling or more constrained variants. Learned importance weighting also outperforms uniform weighting in DMC-C.

Fixed vs Learned Memory Compression

  • Dynamically learned compression decisions achieve lower perplexity and higher downstream accuracy than deterministic Fixed Memory Pooling.Fixed Memory Pooling deterministically averages every n tokens, where n equals the compression ratio.Source: F. DMC Ablations · p. 16

In-layer Relaxation

  • At 1.9× compression, default DMC maintains consistent MMLU performance across compression ratios, while DMC-C and DMC-HardC drop sharply; DMC-C also degrades versus DMC in several settings but remains superior to GQA.Table 3 compares MMLU, commonsense question-answering, and HumanEval across Llama 2 scales and compression ratios.Source: Table 3 · p. 16, F. DMC Ablations · p. 17

Global vs Local (Layer-wise) Compression Prior

  • The Global Prior improves MMLU performance over the Local Prior by allowing layers to use different compression ratios while matching the target average.The Local Prior requires each layer to compress approximately the same amount.Source: F. DMC Ablations · p. 16

Importance Scores

G. Masking Implementation Details

DMC applies a compression-decision mask to unnormalized attention scores and remains compatible with efficient inference libraries. The paper focuses on retrofitting existing LLMs; preliminary from-scratch DMC pre-training performed negatively against GQA.

G. Masking Implementation Details

  • The mask modifies each unnormalized attention score; its log(1 − α_{j+1}) term is computed as log-sigmoid(−α_{j+1}) for numerical precision.Source: G. Masking Implementation Details · p. 17
  • At inference, DMC remains compatible with efficient attention libraries such as Flash Attention despite using arbitrary masks in PyTorch's memory-efficient MHSA implementation.Source: G. Masking Implementation Details · p. 17

H. Limitations

  • The study focuses on retrofitting existing LLMs; preliminary from-scratch DMC pre-training performed worse than GQA, with modeling–segmentation feedback offered as a speculation.The authors suggest that an EM-style alternation might break the proposed feedback cycle.Source: H. Limitations · p. 18
Loading 2403.09636v2…