Sigmoid Attention as a Better Substrate for Learned KV Cache Eviction
2026-08-24 • Machine Learning
Machine LearningArtificial Intelligence
AI summaryⓘ
The authors studied how learned methods for deciding which stored tokens in a language model to delete (KV-cache eviction) work differently during training versus real use. They tested different types of attention mechanisms, gating methods, and positional encodings on GPT-2-scale models. They found that models using sigmoid attention, although weaker in normal use, handled deleting token entries with little loss in prediction quality compared to others. This shows that how attention is normalized affects whether a soft, training-time gate can effectively translate into actually removing entries during inference. Overall, attention design influences the success of learned KV eviction.
KV-cache evictionsoft-to-hard mismatchattention mechanismsigmoid attentionsoftmax attentionpositional encodinggatinglanguage modelperplexity (PPL)inference
Authors
Isaac, Li
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.