ITADN
v-code01/kvbudget
v-code01/kvbudget · 文件
文件最后提交记录最后更新时间
README.md

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: the tau bound, the matched-cache Pareto against the fixed window, the advantage tracking sink strength, and the bounded worst case.
  • test_kvbudget.py: full cache exact, the tau bound, tighter tau trading 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.