Improving KV Cache Eviction with Sigmoid Gates
Listen to the summary
Uses a voice available on your device
Audio options
On this page 5 sections
Related concepts 2 concepts
Key Takeaways
- Sigmoid-gated models outperform existing H2O and KeyDiff methods when evaluated at matched cache sizes.
- Softmax-gated models fail to beat H2O performance at equivalent levels of cache compression.
- The findings are currently limited to a GPT-2 scale model family and a 1B-token training set.
- The study focuses on perplexity and compression metrics rather than measuring actual wall-clock inference speed.
Summary & Methodology Analysis
This study investigates the structural requirements for learned Key Value (KV) cache eviction, where specific attention entries are dropped to save memory. The authors compare different attention substrates by testing softmax versus sigmoid gates within a controlled GPT-2 scale architecture. While standard attention mechanisms often rely on softmax normalization, the authors find that sigmoid-gated models provide a more effective foundation for determining which tokens to retain during cache compression, consistently achieving lower perplexity than established baseline methods like H2O and KeyDiff.
Interactive System Flowchart
Illustrative Implementation
A short sketch of the paper's core idea, not the authors' own code.
# Illustrative sketch (not from the paper)
import torch
import torch.nn as nn
import torch.nn.functional as F
class SigmoidGatedAttention(nn.Module):
def __init__(self, dim, heads):
super().__init__()
self.dim = dim
self.heads = heads
self.qkv = nn.Linear(dim, dim * 3, bias=False)
self.out = nn.Linear(dim, dim, bias=False)
# learned gate network (dense)
self.gate_proj = nn.Linear(dim, 1)
def forward(self, x, cache_keys, cache_vals, threshold=0.5, training=True):
# x: [B, T, dim]
B, T, _ = x.shape
qkv = self.qkv(x).reshape(B, T, 3, self.heads, -1)
q, k, v = qkv.unbind(dim=2) # each: [B, T, heads, head_dim]
# compute attention scores
scores = torch.einsum('bthd,bshd->bths', q, k) / (self.dim ** 0.5)
# learned gate per token (training only)
if training:
g = torch.sigmoid(self.gate_proj(x)).squeeze(-1) # [B, T]
# attenuate logits and values
scores = scores + torch.log(g.unsqueeze(1).unsqueeze(-1) + 1e-9)
v = v * g.unsqueeze(-1).unsqueeze(-1)
else:
# inference: drop cache entries with low gate values
g = torch.sigmoid(self.gate_proj(x)).squeeze(-1)
keep_mask = g >= threshold
# physically drop entries from cache (illustrative)
cache_keys = cache_keys[keep_mask]
cache_vals = cache_vals[keep_mask]
attn = F.softmax(scores, dim=-1)
context = torch.einsum('bths,bshd->bthd', attn, v)
out = self.out(context.reshape(B, T, -1))
return out, cache_keys, cache_vals
// Illustrative sketch (not from the paper)
const tf = require('@tensorflow/tfjs-node');
class SigmoidGatedAttention {
constructor(dim, heads) {
this.dim = dim;
this.heads = heads;
this.qkv = tf.layers.dense({units: dim * 3, useBias: false});
this.out = tf.layers.dense({units: dim, useBias: false});
// learned gate network (dense)
this.gateProj = tf.layers.dense({units: 1, activation: 'sigmoid'});
}
// x: [batch, seqLen, dim]
// cacheKeys, cacheVals: tensors of cached entries
// threshold: scalar, training: boolean
forward(x, cacheKeys, cacheVals, threshold = 0.5, training = true) {
const [batch, seqLen, _] = x.shape;
const qkv = this.qkv.apply(x).reshape([batch, seqLen, 3, this.heads, -1]);
const [q, k, v] = tf.unstack(qkv, 2);
// attention scores
const scores = tf.matMul(q, k, false, true).div(Math.sqrt(this.dim));
if (training) {
const g = this.gateProj.apply(x).squeeze([-1]); // [batch, seqLen]
const gExp = g.expandDims(1).expandDims(-1);
const scoresAtt = scores.add(tf.log(gExp.add(1e-9)));
const vAtt = v.mul(gExp);
return this._applyAttention(scoresAtt, vAtt, cacheKeys, cacheVals);
} else {
const g = this.gateProj.apply(x).squeeze([-1]);
const keepMask = g.greaterEqual(tf.scalar(threshold));
// illustrative physical drop
const keptKeys = tf.booleanMaskAsync(cacheKeys, keepMask);
const keptVals = tf.booleanMaskAsync(cacheVals, keepMask);
return this._applyAttention(scores, v, keptKeys, keptVals);
}
}
_applyAttention(scores, v, cacheKeys, cacheVals) {
const attn = tf.softmax(scores, -1);
const context = tf.matMul(attn, v);
const out = this.out.apply(context.reshape([scores.shape[0], scores.shape[1], -1]));
return {out, cacheKeys, cacheVals};
}
}
Cross-Examination & FAQs
A deeper dive clarifying mechanics, constraints, and baseline evaluations.
Q1. What is the main finding regarding cache eviction?
Sigmoid-gated models are better at cache eviction than softmax-gated models, outperforming H2O and KeyDiff at matched cache sizes.
Q2. Which models did the researchers test?
The researchers used a GPT-2 scale model family for all experiments.
Q3. Does this method improve inference speed?
The paper does not report wall-clock speed improvements, as no production fused kernel was implemented for the evaluated architecture.
Q4. How does this approach compare to H2O?
Under a matched live-cache protocol, the learned sigmoid gates achieve lower perplexity than H2O implementations, whereas softmax-gated models perform worse than H2O at matched compression.
Q5. What is KeyDiff?
KeyDiff is an attention-free, key-similarity-based method that ranks cached keys by their similarity to a key-space anchor and retains distinctive keys.
Q6. What is H2O?
H2O maintains accumulated attention-based heavy-hitter scores and retains high-score tokens together with recent tokens.
Q7. Are these results ready for deployment?
The study uses a 1B-token training set and a single model family, which the authors note is insufficient to confirm deployment-scale behavior.
Q8. What metrics were used to evaluate success?
The authors report perplexity and compression rates rather than end-to-end inference speed.
Q9. Was a production-ready implementation provided?
No, the authors did not implement a production fused kernel for the specific gated sigmoid attention path studied.