Sigmoid Attention as a Better Substrate for Learned KV Cache Eviction

AI summary

Overview

• Research area: Efficient Transformer inference, specifically KV-cache compression and learned cache eviction, with a focus on how the attention normalization function (softmax vs. sigmoid) interacts with trainable eviction gates.
• Technical level: Advanced. The paper assumes familiarity with attention mechanics, KV caching, prefill/decode protocols, perplexity (PPL) evaluation, and rotary positional embeddings.
• Scope: A single-author, controlled 2 × 2 × 2 factorial study of eight GPT-2-scale models trained on OpenWebText, testing whether sigmoid attention makes learned soft-to-hard KV eviction transfer more cleanly than softmax attention. Code and notebooks are available at https://github.com/IsaacLi74/sigmoid-kv-eviction.

What This Paper Is About

Learned KV-cache eviction has a "soft-to-hard mismatch": during training a model can only attenuate a token's contribution through a differentiable gate, but at inference memory is actually freed only when the key/value entry is physically deleted. The paper asks whether the attention substrate itself (softmax versus sigmoid normalization) changes how cleanly that transition from soft attenuation to hard deletion transfers, holding the learned gating mechanism fixed. Using a controlled 2 × 2 × 2 design over attention type, learned gating, and positional encoding, the authors find that sigmoid attention is a worse dense language model but a better substrate for learned hard eviction.

Key Contributions

1. A controlled 2 × 2 × 2 factorial comparison over attention type (softmax / sigmoid), gating (dense / learned), and rotary position treatment (RoPE / NoPE), yielding eight experimental cells that all share the same GPT-2-scale backbone (n_layer = 12, n_head = 12, d_model = 768, d_h = 64, d_ff = 3072, sequence length 512).
2. A matched live-cache evaluation protocol that measures a learned gate's realized final live-cache size and then runs post-hoc H2O and KeyDiff baselines on the paired dense no-gate backbone at that same cache size, isolating learned soft-to-hard eviction against post-hoc eviction on the same substrate.
3. A mechanism analysis with three diagnostics — a random support-removal probe, attention row-sum variation and gate binarization statistics, and token-level gate selectivity — arguing that the sigmoid advantage is specific to learned deletion, not to deletion tolerance in general.
4. An open reproduction package (code and notebooks) for all eight experimental cells at https://github.com/IsaacLi74/sigmoid-kv-eviction.

Main Findings

• Softmax wins when no eviction happens. Dense no-eviction test PPL is 22.079 (SM RoPE), 22.440 (Sig RoPE), 23.398 (SM NoPE), and 24.603 (Sig NoPE). Softmax is better by 0.361 PPL with RoPE and 1.205 PPL without RoPE.

• Sigmoid wins once deletion is learned. Against their own no-eviction references, Sig+G RoPE deletes 19.8% of KV entries with a slightly lower measured PPL (22.424 vs. 22.440), and Sig+G NoPE deletes 32.2% with only a small increase (24.637 vs. 24.603).

• The substrate effect reverses the comparison to post-hoc baselines. Under the matched live-cache protocol, Sig+G RoPE (19.8% compression) obtains PPL 22.424 versus H2O 22.480 and KeyDiff 22.474. Sig+G NoPE (32.2% compression) obtains 24.637 versus H2O 25.107 and KeyDiff 25.786. On softmax backbones the pattern reverses: H2O is stronger than the learned gate (SM+G RoPE: 22.145 vs. 22.514; SM+G NoPE: 23.423 vs. 23.576).

• Sigmoid is not generically robust to arbitrary deletion. In a random support-removal probe on the two dense RoPE backbones (averaged over 50 test sequences and 12 layers), softmax and sigmoid are nearly tied at p = 0.1 (0.574 vs. 0.579), but sigmoid is more perturbed at larger rates: 0.943 vs. 1.051 at p = 0.3, and 1.131 vs. 3.666 at p = 0.5.

• Row normalization is the structural difference. Softmax row sums are 1 by construction, so suppressing one token forces renormalization over the others. Sigmoid row-sum coefficient of variation is nonzero: 0.484 for Sig+G RoPE and 0.444 for Sig+G NoPE, compared with 0 for both softmax models, measured across query positions on 1953 test sequences of length 512.

• Gate binarization alone does not explain the result. Aggregate binarization (fraction of gates with g < 0.05 or g > 0.95) is 43.3% for Sig+G NoPE and 36.8% for SM+G RoPE, but Sig+G RoPE is lower at 24.4% yet still transfers more cleanly to hard deletion. The paper argues both a thresholdable gate structure and a non-row-normalized substrate are needed.

• The gates are token-selective, not uniform cache shrinkers. Across all four gated models, stopwords and newlines tend to receive the lowest gates while whitespace tokens consistently receive the highest. For Sig+G NoPE, stopwords receive a mean gate of 0.178, content tokens 0.281, and whitespace 0.455.

• Sigmoid models assign higher mean gates than softmax models in every token bucket, so the advantage does not come from more aggressive soft suppression during training.

Methodology in Plain English

The authors train eight small (GPT-2-scale, 123M-parameter family) decoder-only Transformers with AdamW for three epochs on a 1B-token OpenWebText training split — roughly 3B token exposures — with 100K held-out validation tokens and 1M test tokens. The eight cells cross three binary choices: softmax versus sigmoid attention, no learned gate versus a learned gate, and RoPE versus NoPE (RoPE rotation removed from queries and keys, with no learned absolute position embeddings added). All variants apply per-head QK-RMSNorm before attention scores.

For sigmoid attention the authors add a per-query negative bias b(i) = −log(i + 1) to stabilize the expected row mass, noting this would be a mathematical no-op under softmax. For gated models, each layer except the last produces a token-wise gate in (0, 1), initialized so that g = σ(5) ≈ 0.993 at step 0 and early training behaves like the no-gate model. The gate is injected in two channels — added to the attention logits as log(g + ε) with ε = 10⁻⁸, and multiplied onto the values as g·v — a differentiable simulation of dropping the token. Training minimizes cross-entropy plus a mean budget term pulling gates toward small values, with λ = 0.03 across all gated variants.

At evaluation, each test sequence is split into a 384-token prefill prefix and a 128-token decode window, and PPL is computed on the decode window. A threshold grid from 0 to 0.9 is swept on 100 validation sequences, and the highest-compression threshold satisfying ΔPPL < 0.1 against that model's own no-eviction reference is frozen before test. Physical eviction then removes every cached entry with g < τ* after prefill, applying the same rule to each newly written entry during decode; reported compression counts only physically removed entries. Post-hoc baselines H2O (attention-score heavy-hitter, splitting the matched cache equally between heavy-hitter and recent tokens without tuning) and KeyDiff (attention-free key-similarity ranking against an unnormalized cached-key anchor) are run on the corresponding dense backbone at the matched final live-cache size.

Why This Matters

The paper reframes KV-cache eviction from a purely inference-time heuristic into an architectural co-design question: which attention normalization makes a trainable retention signal actually convertible into physical memory savings. It also gives a concrete counterexample to the intuition that the better dense language model is automatically the better substrate for compression.

Real-world applications:

• Long-context chat assistants, where persistent multi-turn KV caches dominate GPU memory and a threshold-stable eviction policy could shrink memory without measurable quality loss.
• Retrieval-augmented generation and document QA, where long retrieved contexts create large prefills whose cached entries could be selectively dropped by token class (the paper finds stopwords and newlines receive the lowest gates).
• On-device and edge inference, where the paper's NoPE result is most relevant: Sig+G NoPE deletes 32.2% of KV entries with only a small PPL increase, and dropping rotary embeddings removes a positional-encoding computation.
• High-throughput serving of many concurrent sequences, where live-cache size rather than tokens processed is often the binding constraint on batch size.

Industry relevance: the paper reports PPL and compression rates but explicitly not wall-clock speed, and the authors note they did not implement a production fused kernel for the specific gated sigmoid attention path studied. They therefore frame the results as evidence about cache quality and learnability, not as an end-to-end inference-speed claim — an important caveat for anyone considering deployment.

Future Directions

1. Scaling and stress testing. The authors state it remains open whether the effect persists at larger model scales, longer contexts beyond 512 tokens, multiple seeds (only one seed per cell was used), other text distributions, and downstream long-context tasks.
2. Gate parameterization ablations. The study uses one gate design — token-wise gates with both logit attenuation and value scaling at λ = 0.03. Ablating logits-only versus values-only injection, alternative λ values, per-head gates, and per-layer thresholds is needed to separate the sigmoid substrate's contribution from the gate parameterization's.
3. Better threshold and comparison protocols. Finer threshold grids and alternative operating criteria could select different cache sizes; the authors also note that matched final live-cache size rather than matched PPL does not answer every deployment question, and that sigmoid threshold sweeps can have sharp cliffs (visible in the full validation sweep table, where e.g. SM+G NoPE reaches 995.48 PPL at τ = 0.100) making matched-PPL comparison less stable.
4. Stronger baselines and real kernels. The H2O and KeyDiff comparisons use a fixed H2O heavy-hitter/recent split with no tuning of post-hoc hyperparameters, so they should be read as matched-cache reference points rather than fully optimized state-of-the-art benchmarking; and a fused kernel for the gated sigmoid path would be needed to turn these cache-quality findings into speed claims.

Target Audience

This paper is most useful to researchers and engineers working on efficient Transformer inference and KV-cache compression, and to architecture researchers interested in replacing or modifying softmax attention. It will also interest practitioners building long-context or memory-constrained serving systems who want to understand the tradeoff between dense language-modeling quality and trainable eviction. Readers need working knowledge of attention, KV caching, and perplexity-based evaluation; the controlled factorial design makes the results relatively easy to follow once that background is in place.

Authors’ abstract

Learned KV-cache eviction often faces a soft-to-hard mismatch: during training, differentiable gates typically attenuate token contributions, whereas inference saves memory only when KV entries are physically removed. We ask whether the attention substrate affects this soft-to-hard transition. Using GPT-2-scale Transformers trained on OpenWebText, we run a controlled $2\times2\times2$ comparison over attention type, learned gating, and positional encoding. Although sigmoid attention is worse as a dense language model, learned hard eviction changes the useful operating points: sigmoid-gated models delete KV entries with negligible PPL change relative to their own no-eviction references. Under a matched live-cache protocol on the same dense backbones, learned sigmoid gates obtain lower PPL than our H$_2$O and KeyDiff implementations, whereas softmax gates do not uniformly beat these post-hoc methods. The results suggest that attention normalization can substantially affect whether a training-time soft gate transfers cleanly to hard KV deletion.