PaperlayerView exact revision on arXiv

Source-linked AI summary

Sigmoid Attention as a Better Substrate for Learned KV Cache Eviction

Isaac, Li

arXiv:2608.23296v1cs.LGcs.AI

TL;DR

Learned KV-cache eviction must bridge soft training-time attenuation and hard inference-time deletion, raising whether the attention substrate affects that transition. The paper evaluates this with a controlled 2×2×2 GPT-2-scale comparison and matched-cache baselines. Sigmoid-gated models transfer more effectively to hard deletion, while the effect is not explained by generic robustness to arbitrary removal.

  • Problem

    Learned KV-cache eviction faces a soft-to-hard mismatch, and prior work leaves open which attention substrate makes this transition easier to learn.

  • Method

    The study compares softmax and sigmoid attention, dense and learned eviction, and RoPE and NoPE in GPT-2-scale Transformers, using thresholded physical deletion and matched live-cache evaluations.

  • Results

    Sigmoid-gated models delete 19.8% of RoPE KV entries with PPL 22.424 versus 22.440 without eviction, and outperform H2O and KeyDiff at matched cache sizes; softmax gates do not uniformly do so.

  • Takeaways & Limitations

    The same learned soft-to-hard rule is weak under softmax but useful under sigmoid attention, with RoPE improving quality near matched compression and NoPE expanding the safe compression range.

  • Takeaways & Limitations

    The study uses one 123M GPT-2-scale model family, one seed per cell, a 1B-token OpenWebText split, and context length 512, leaving broader deployment behavior open.

Abstract

from arXiv · show

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.

1. Introduction

The paper studies whether attention normalization affects the transfer from learned soft gating to hard KV-cache deletion. A controlled comparison finds that sigmoid attention, despite weaker dense modeling, provides more useful learned-eviction operating points than softmax.

  • Learned eviction must be threshold-stable because training attenuates tokens softly, whereas inference frees memory only through physical KV-entry deletion.The desired transition is hard deletion with little or no next-token prediction loss.
  • The study asks whether attention normalization affects how learnable hard KV deletion is, complementing frozen-model post-hoc retention methods such as H2O and KeyDiff.The comparison holds the learned eviction mechanism fixed while varying the attention substrate.
  • A controlled 2×2×2 experiment varies softmax versus sigmoid attention, dense versus learned eviction, and RoPE versus NoPE using GPT-2-scale models.NoPE removes rotary embeddings from queries and keys, and gated models share the same soft-to-hard rule.
  • 19.8% of KV entries are deleted by Sig+G RoPE with negligible PPL change, 22.424 versus 22.440, relative to its dense no-eviction reference.Sig+G NoPE deletes 32.2% with PPL changing from 24.603 to 24.637.
  • At matched live-cache sizes, learned sigmoid gates obtain lower PPL than the H2O and KeyDiff implementations, whereas softmax-gated models lose to H2O.This separates dense language-model quality from learned eviction quality.
  • Random deletion perturbs sigmoid outputs more than softmax at larger removal rates, indicating that sigmoid’s advantage appears for learned rather than arbitrary deletion patterns.The associated mechanism is a cleaner local attenuation channel from non-row-normalized attention.

2. Method

The method compares softmax and sigmoid attention under learned gating and RoPE or NoPE, then evaluates whether differentiable gate signals support physical KV-cache eviction. It combines gate stabilization, two-channel token suppression, threshold selection, and matched-cache comparisons with post-hoc baselines.

  • Controlled model design: Eight models form a 2 × 2 × 2 design over attention type, learned gating, and RoPE versus NoPE positional treatment.NoPE removes RoPE rotation and does not add learned absolute position embeddings.
  • Attention and position treatment: All variants use per-head QK-RMSNorm before computing attention scores, with RoPE applied only in RoPE cells.The NoPE condition omits the query and key rotations.
  • Sigmoid-attention stabilization: A per-query bias b(i) = −log(i + 1) stabilizes sigmoid attention row mass, while the same row-wise bias would be inert under softmax.The bias accounts for causal queries attending to only i + 1 keys.
  • Learned gating: Gates enter both attention logits and values, so g →0 drives a token’s contribution toward zero and approximates hard deletion during training.The final layer has no gate because it has no subsequent KV cache to evict; gates start near one with σ(5) ≈0.993.
  • Physical eviction: Inference physically removes cached entries when g_j < τ, including newly written decode entries, while retained entries keep learned logit and value scaling.Compression counts only physically removed KV entries, not soft attenuation of retained entries.
  • Evaluation protocol: Thresholds maximize compression subject to ΔPPL < 0.1 against each gated model’s no-eviction reference, and post-hoc baselines are evaluated at the learned gate’s realized live-cache size.H2O and KeyDiff operate on paired dense no-gate backbones at matched final live-cache size.

3. Results

The results distinguish dense language-modeling quality from learned hard-eviction behavior: softmax is stronger without eviction, but sigmoid gates yield better threshold-selected and matched-cache outcomes. Sigmoid primarily improves quality at similar RoPE compression and expands safe compression without RoPE.

  • Dense baselines: Softmax is better than sigmoid by 0.361 PPL with RoPE and 1.205 PPL without RoPE in dense no-eviction evaluation.
  • Validation threshold selection: The experiment selects frozen hard-deletion thresholds by maximizing compression subject to validation ∆PPL < 0.1 versus each gated model’s τ = 0 reference.Thresholds are swept over a 15-point validation grid before test evaluation.
  • Validation threshold selection: Both softmax models select τ ⋆ = 0.01, whereas both sigmoid models select τ ⋆ = 0.05.
  • Validation threshold selection: Without RoPE, Sig+G reaches 31.5% compression compared with 9.7% for SM+G.
  • Matched live-cache size results: At matched final live-cache sizes, learned sigmoid gates outperform H2O and KeyDiff, while learned gates are not uniformly better on softmax backbones.Sig+G RoPE deletes 19.8% with PPL 22.424 versus 22.440 densely; Sig+G NoPE deletes 32.2% with PPL 24.637 versus 24.603.
  • Matched live-cache size results: The substrate-specific effect differs by positional treatment: RoPE mainly improves quality at nearly matched compression, while NoPE mainly expands the safe compression range.

4. Mechanism Analysis

The diagnostics indicate that learned eviction benefits from both thresholdable gates and an attention substrate whose row sums are not forced to renormalize. Sigmoid attention shows this structural freedom, while learned gates selectively suppress token classes rather than uniformly shrinking the cache.

  • Not arbitrary-deletion robustness: Random deletion perturbs sigmoid attention more than softmax at larger removal rates, so sigmoid’s eviction advantage is not generic deletion robustness.At p=0.5, average perturbation is 1.131 for softmax versus 3.666 for sigmoid; the probe averages over 12 layers and 50 test sequences.
  • Row-sum freedom and gate stability: Softmax row sums are fixed at one, causing suppressed-token gates to redistribute attention mass across the remaining tokens.Sigmoid attention permits query-, layer-, and sequence-dependent row mass, avoiding this forced compensation.
  • Row-sum freedom and gate stability: Sigmoid gated models retain nonzero row-sum variation, with row-sum CVs of 0.484 for Sig+G RoPE and 0.444 for Sig+G NoPE versus 0 for softmax.The coefficient of variation is measured across query positions on 1953 test sequences of length 512.
  • Row-sum freedom and gate stability: Gate binarization alone is insufficient: Sig+G RoPE transfers more cleanly despite lower aggregate binarization than SM+G RoPE, at 24.4% versus 36.8%.The useful regime appears to require both thresholdable gates and a substrate that makes soft attenuation locally approximate hard deletion.
  • Token selectivity: Learned gates are token-selective rather than uniform cache shrinkers, generally assigning the lowest values to stopwords and newlines and the highest to whitespace.For Sig+G NoPE, mean gates are 0.178 for stopwords, 0.281 for content tokens, and 0.455 for whitespace.
  • Token selectivity: Sigmoid models assign higher mean gates than softmax models in every token bucket, so their advantage does not arise from more aggressive soft suppression during training.The advantage appears after hard thresholding, when sigmoid gate values form a more useful keep/delete partition.

5. Discussion

The discussion frames sigmoid attention as a potentially better substrate for learned physical KV eviction, not as a universally better language-modeling mechanism. It also identifies evaluation and scope limitations that constrain deployment conclusions.

  • 5. Discussion: Sigmoid attention is presented as better suited to converting learned soft retention signals into hard KV deletion, while softmax remains stronger for dense modeling.The paper distinguishes learned eviction from post-hoc cache removal and treats the substrate effect as an architectural co-design question.
  • 5. Discussion: Sig+G RoPE deletes 19.8% of KV entries with PPL 22.424 versus 22.440 for its dense no-eviction reference.The small negative PPL difference is not interpreted as a reliable modeling gain; its operational importance is cache reduction without measurable quality loss.
  • 5. Discussion: The proposed future direction is to jointly design attention mechanisms and KV-retention policies if the substrate effect persists at larger scale and longer context.The current evidence motivates, rather than establishes, this broader architectural direction.
  • 5.3. Limitations: Threshold selection uses a discrete validation grid and ∆PPL < 0.1, so finer grids, per-layer thresholds, or alternative criteria could choose different cache sizes.Sigmoid threshold sweeps can have sharp cliffs, making matched-PPL comparisons less stable.
  • 5.3. Limitations: The study uses one 123M GPT-2-scale family, one seed per cell, context length 512, and an OpenWebText split, leaving broader deployment-scale behavior open.Unresolved settings include larger models, longer contexts, multiple seeds, other text distributions, and downstream long-context tasks.
  • 5.3. Limitations: Only one token-wise gate design with logit attenuation, value scaling, and λ = 0.03 is tested, so substrate and parameterization effects are not separated.The paper does not ablate injection channels, λ, per-head gates, or per-layer thresholds.
  • 5.3. Limitations: H2O and KeyDiff comparisons use fixed or untuned post-hoc settings, making them matched-cache reference points rather than fully optimized SOTA benchmarks.The baselines represent complementary attention-importance and key-similarity signals.
  • 5.3. Limitations: The experiments report PPL and compression rates rather than wall-clock speed, so they do not establish an end-to-end inference-speed advantage.No production fused kernel was implemented for the studied gated sigmoid-attention path.

A. Additional Validation Sweep Results

The appendix reports the complete validation threshold sweep used to select hard-deletion operating points. Bold entries identify the selected points under the paper’s compression-quality criterion.

  • A. Additional Validation Sweep Results: Table 5 gives decode-window PPL and realized hard-KV compression for each threshold on 100 validation sequences.The sweep covers the full validation threshold grid.
  • A. Additional Validation Sweep Results: Bold entries mark the operating points selected under the ∆PPL < 0.1 rule.Selection favors the highest compression satisfying the quality constraint.

B. Qualitative Gate Heatmaps

The qualitative appendix visualizes learned gate patterns across prompt types. These heatmaps illustrate non-uniform variation across layers, token positions, and attention substrates rather than serving as quantitative metrics.

  • B. Qualitative Gate Heatmaps: Heatmap rows correspond to gate layers and columns correspond to token positions.Green indicates larger gate values, while red indicates smaller gate values.
  • B. Qualitative Gate Heatmaps: The five prompt categories are narrative, technical, dialogue, factual, and question answering.Each category has a separate gate heatmap.
  • B. Qualitative Gate Heatmaps: The heatmaps qualitatively suggest that learned gates vary across layers, token positions, and substrates instead of shrinking caches uniformly.The figures are not used as quantitative metrics.
Loading 2608.23296v1…