kvbudget
A mass-budgeted KV cache for autoregressive video generation, and the bound that makes it a knob rather
than a heuristic. Its sibling result kvexact showed that a bounded
KV cache is lossy by exactly the attention mass it drops. This turns that equality into an eviction
policy: instead of keeping the last W frames, keep the smallest set of past tokens whose attention mass
reaches 1 - tau, and the output error is bounded by tau times the value range, by construction. The
policy adapts its memory to the attention, keeps the keyframe sinks a fixed window throws away, and beats
the fixed window by about an order of magnitude at matched memory on video-structured attention.
The policy and its bound
At each generation step the query attends over the cached past tokens. A fixed window keeps the last W;
the mass budget keeps the heavy hitters, the smallest set carrying mass at least 1 - tau. Because the
output error equals the dropped attention mass times the value range, keeping 1 - tau of the mass caps
the error at tau times that range for every query. The bound is exact and holds for any attention
distribution:
=== mass-budget obeys the kvexact bound (error <= tau * Vmax) ===
tau= 0.1: cache 35.5 max error 0.1993 tau*Vmax 0.983 within bound True
tau=0.05: cache 48.2 max error 0.1211 tau*Vmax 0.492 within bound True
tau=0.02: cache 65.1 max error 0.0634 tau*Vmax 0.197 within bound True
tau=0.01: cache 77.4 max error 0.0338 tau*Vmax 0.098 within bound True
Tightening tau spends more cache for less error, and the worst-case error never crosses tau * Vmax.
The Pareto
Run over a rollout on attention with the structure video attention is documented to have, temporal recency plus a few keyframe sinks, the mass budget's error-versus-memory curve sits about an order of magnitude below the fixed window's at every matched cache size:
window (cache,err): [(15.9, 1.70), (31.5, 1.10), (46.5, 0.69), (60.9, 0.40), (88.4, 0.10), (113.8, 0.014)]
mass (cache,err): [(35.5, 0.128), (48.2, 0.067), (57.6, 0.041), (65.1, 0.028)]
error ratio window/mass at matched cache: median 9.7x, min 7.7x; mass better at every size: True
The reason is the sinks. A fixed window drops a distant keyframe whenever the rollout moves past it, and on those queries its error spikes; the mass budget never evicts a heavy hitter, so its worst-case error stays bounded while the window's does not:
=== the win tracks sink strength (matched cache ~48) ===
sink_boost= 0.0: window mean/max 0.124/0.282 | mass mean/max 0.076/0.293
sink_boost= 4.0: window mean/max 0.247/2.488 | mass mean/max 0.073/0.194
sink_boost= 8.0: window mean/max 1.464/7.310 | mass mean/max 0.059/0.089
worst-case error: window 5.258 vs mass 0.121 (43x tail)
With no sinks the two policies nearly converge, because then the heavy hitters are exactly the recent
tokens and the window keeps them anyway. As the sinks strengthen the fixed window's worst-case error
climbs to 7.3 while the mass budget's stays near 0.1, a forty-three-times tighter tail at matched
memory. The advantage is precisely the distant heavy hitters the window cannot see.
Why it works and where the rigor sits
The bound is weight-independent: for any attention distribution, an unkept set of mass m changes the
output by at most m times the value range, so the mass budget's tau caps the error whatever the model
does. The size of the benefit is not weight-independent, it grows with how concentrated the attention is
and how much mass sits on distant tokens, so it is reported here as a function of the recency and sink
strength rather than asserted. Video attention is documented to carry both kinds of structure, which is
why the policy is worth its bookkeeping there.
Falsifier
The bound fails if a mass-budget query's error exceeds tau times the value range, or if the kept mass
falls below 1 - tau. The Pareto claim fails if, at matched cache size, a fixed window ever matches or
beats the mass budget on attention with distant heavy hitters. Neither happens here; with the sinks
removed the two converge, as the mechanism predicts.
Layout
cache.py: video-structured causal attention with tunable recency and sinks, the three eviction policies (full, fixed window, mass budget), and the per-query error and cache size over a rollout.pareto.py: thetaubound, the matched-cache Pareto against the fixed window, the advantage tracking sink strength, and the bounded worst case.test_kvbudget.py: full cache exact, thetaubound, tightertautrading cache for error, beating the window at matched cache, the sink dependence, and the bounded tail.
Reproduce
python cache.py
python pareto.py
python test_kvbudget.py
The attention structure and value vectors are deterministic under fixed seeds on this machine. The point
is that kvexact's equality, output error equals dropped mass, makes a heavy-hitter cache into a policy
with a tau knob and a bounded worst case, which a fixed window cannot offer.