H₂O: Heavy-Hitter Oracle for Efficient Generative Inference
This research paper shows that attention mass concentrates on a few heavy-hitter tokens, and that an eviction policy balancing those hitters with a local window can slash KV memory while preserving accuracy on OPT, LLaMA, and GPT-NeoX.
Long chat and story generation do not only pay for model weights. Every past token parks a key/value state in the KV cache, so GPU memory grows with sequence length and batch size. Once that cache dominates the working set, throughput stalls and serving systems start swapping or OOM.
Zhang et al. start from a sparse fact about Softmax attention: a small subset of tokens carries most of the accumulated attention mass. They call those tokens Heavy Hitters (). Drop them and quality falls. Keep only a recent window (Local) and you still miss the same sparse peaks. Heavy Hitter Oracle () is the eviction policy that dynamically retains a balance of recent tokens and tokens inside a fixed budget .
What this paper explains
Autoregressive Transformers cache and for every past position so decode steps can attend without recomputing the prompt. For length the cache is per layer and grows again with batch size.
The authors ask: can we evict KV states aggressively — down to a small constant budget — without pretending every past token was equally useful?
Their answer has three layers:
- Empirically show that accumulated attention scores follow a power-law and that the heavy tail aligns with frequent co-occurrence in text.
- Prove that removing tokens hurts; a pure Local baseline that keeps only the newest KVs is not enough.
- Ship : a greedy eviction policy that fills a size- cache with recent tokens plus high accumulated-score tokens, with a dynamic-submodular analysis under mild assumptions.
Prior limits
- Full cache. Exact attention over the entire prefix. Accuracy ceiling — and a memory wall once or the batch is large.
- Local / sliding window. Keep only the most recent KVs. Cheap and simple; the paper’s comparisons show it loses the sparse high-mass positions that still matter later in generation.
- Static Top- without dynamics. Selecting important tokens once (or ignoring score accumulation bias) fails to track which positions remain influential as decoding proceeds.
- Orthogonal serving tricks. Offloading (FlexGen-style), paging (PagedAttention), and quantization shrink how KVs are stored; they do not answer which tokens to keep when the logical cache must shrink.
The mechanism
Heavy Hitters. Sum Softmax attention weights into each past key across the generation trace. The paper finds these accumulated attention scores are highly skewed: a few tokens dominate. They name that set . Scatter plots in the paper further show accumulated score tracks token co-occurrence frequency in the data — is not an accident of one prompt.
Why Local fails. Oldest-first eviction quietly deletes early high-mass positions. Later Softmax steps redistribute over a cache that no longer contains the tokens the model learned to lean on, and task accuracy drops in their OPT / LLaMA / GPT-NeoX sweeps.
eviction. Maintain a cache set of size at step :
- When , admit new KV states.
- When full, admit the newest token and evict the member with lowest score under , where the practical score is accumulated attention (with care for the bias that older tokens have had more steps to accumulate — the paper discusses this and balances with an explicit recent quota).
Informally (Definition 4.3): and at most one membership change per step (). The implementation skeleton keeps a running score vector, protects a local block of recent tokens, and fills the remaining slots with the current set.
Submodular lens. They cast eviction as a dynamic submodular problem and argue the greedy policy inherits approximation-style guidance under mild assumptions — useful as a design rationale, not something you need to re-derive to use the system.
Algorithm / figure walkthrough
- Prefill the prompt; materialize full KV for the prefix until the budget fills (or immediately begin scoring if the prompt already exceeds ).
- For each new decode token, compute attention over the current cached keys and add those Softmax weights into a per-token accumulated score.
- Append the new KV. If size exceeds , evict the lowest-scoring non-protected (non-local) entry so the cache again holds local + slots.
- Attend only over the surviving KVs on the next step.
- Repeat — memory stays while generated length grows.
Contrast the three policies in the first figure: Full never evicts; Local keeps a trailing window only; re-ranks who deserves the non-local slots using accumulated mass.
What to notice when reading
- is defined by attention mass, not by “being the first tokens.” That is the clean contrast with StreamingLLM’s attention sinks, which pin early positions because Softmax needs a denominator home. H₂O may keep early tokens when they score as hitters, and may keep mid-sequence spikes too.
- The default story in the abstract is a heavy-hitter-oriented budget relative to the full cache in their main throughput claims — treat as the authors’ reported operating point, not a universal constant.
- Score accumulation favors tokens that have been present longer; the local + split is how they avoid a pure “oldest always wins” bias.
- H₂O is a content policy for which KVs survive. It composes with quantization and with system stacks (their tables include compressed-weight variants).
Results and evidence
Numbers below are from the paper’s reported experiments (arXiv:2306.14048), not independent re-runs.
Accuracy vs Local / Full. Across zero- and few-shot style evaluations on OPT, LLaMA, and GPT-NeoX families, tracks Full much more closely than Local at matched cache budgets (Figure 4 and Table 1 narrative). Ablations that strip while keeping sparse/Local structure degrade (Table 2: OPT-30B at KV budget).
Throughput (T4). With a heavy-hitter-oriented cache, their implementation improves generation throughput over DeepSpeed ZeRO-Inference, Hugging Face Accelerate, and FlexGen by up to , , and on OPT-6.7B and OPT-30B (abstract / Table 3 setting). Same-batch-size latency improves by up to .
Long streams. Figure 5 shows streaming with on inputs on the order of four million tokens, and compares perplexity against StreamLLM-style baselines under matched start/heavy-hitter vs local budgets (see also Figure 9’s H2O-256-style budget labels).
Compatibility. Table 6 reports composition with weight quantization; compressed -c variants appear in the throughput appendices.
Limitations
Compressed from the paper’s scope and what it does not claim:
- Eviction is irreversible for that decode: once a KV leaves, mid-context facts living only there are gone unless you add retrieval or recomputation.
- Matching Full on their task suite at a given budget is not a claim of parity with infinite dense attention on every downstream metric.
- Accumulated-score bookkeeping and the local/ split need careful implementation; naive Top- each step is not the same algorithm.
- End-to-end speedups depend on the baseline stack (offload vs in-GPU). Gains vs FlexGen () differ in kind from gains vs ZeRO-Inference / Accelerate () in their tables.
- Complementary to — not a replacement for — paged allocators, continuous batching, or speculative decoding.
How to read the paper
- Abstract — budget teaser and the / / claims.
- Figure 1 — Full vs Local vs cartoons and the accuracy–memory sketch.
- §3 + Figure 2 — sparsity, power-law accumulated scores, co-occurrence correlation.
- §4 + Algorithm 1 / Figure 3 — eviction loop and the dynamic-submodular framing.
- §5 Figures 4–5 + Tables 1–5 — accuracy, throughput, long-stream behavior.
- Appendix comparisons with StreamLLM / SpAtten / Top- — place H₂O among sibling sparse-KV ideas.
Knowledge check
What primarily defines a Heavy Hitter (H₂) token in this paper?
Why is a pure Local (recent-only) KV cache a weak eviction baseline here?
In the abstract’s main throughput claim, what cache budget do the authors highlight for H₂O?
Keep reading / Sources
- Paper: H₂O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models (arXiv:2306.14048)
- Related Fanout Daily: StreamingLLM attention sinks (pin early sinks vs score-based hitters)
- Related Fanout Daily: FlexGen (offload stack H₂O compares against)
- Related Fanout Daily: PagedAttention (block allocator for KV memory)
- Code link from the paper abstract (authors’ release pointer on arXiv)