pith. sign in

arxiv: 2605.26266 · v1 · pith:UTKIMQGGnew · submitted 2026-05-25 · 💻 cs.LG · cs.AI· cs.CV· cs.GR· eess.IV

Quantized Keys Steal Attention: Bias Correction for KV-Cache Compression in Video Diffusion

Tuna Tuncer , Felix Becker , Thomas Pfeil This is my paper

Pith reviewed 2026-06-29 22:47 UTC · model grok-4.3

classification 💻 cs.LG cs.AIcs.CVcs.GReess.IV
keywords KV cache quantizationvideo diffusionattention biasJensen biasmemory compressionautoregressive generationlow-bit inferencesoftmax correction
0
0 comments X

The pith

Quantization of KV caches creates Jensen bias that steals attention from current queries in video diffusion, but a per-score correction removes it in expectation.

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

Chunk-wise autoregressive video diffusion models store past computations in a KV cache that grows with video length and quickly exhausts memory. Quantizing this cache to low bit-widths reduces the footprint but degrades output quality because the exponential inside softmax is convex, so quantization noise systematically inflates attention weights on the cached keys at the expense of the unquantized current chunk. The paper derives an exact correction term for each attention score that cancels this bias on average, using only the quantization step sizes of the stored keys and the norm of the query. A second-order Taylor expansion makes the correction cheap enough to compute on the fly with no added memory. On three different video diffusion models the fix recovers most of the lost quality at INT2, bringing results close to full BF16 precision and allowing INT2 to beat conventional INT4 while using half the memory.

Core claim

A key driver of quality degradation under KV-cache quantization is a systematic bias in attention weights caused by the convexity of the exponential in softmax: quantization noise inflates the contribution of cached keys, an effect termed the Jensen bias that causes those keys to steal attention mass from the unquantized current chunk. A per-attention-score correction removes this bias in expectation and is computed from the quantization step sizes of the cached keys together with the query norm. Using a second-order Taylor approximation, the correction adds negligible overhead and requires no extra memory alongside the cache.

What carries the argument

The Jensen bias correction: a per-attention-score adjustment computed from cached-key quantization step sizes and query norm that removes the expected inflation of attention weights due to convexity of the exponential.

If this is right

  • At INT2 the correction brings video quality close to BF16 on MAGI-1, SkyReels-V2 and HY-WorldPlay.
  • INT2 with correction can exceed the quality of standard INT4 quantization while using 50 percent less memory.
  • The correction needs no additional storage and incurs only negligible extra compute from the Taylor approximation.
  • The approach applies to any chunk-wise autoregressive video diffusion model that maintains a growing KV cache.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

  • If the same convexity mechanism dominates in other attention-based generators, the correction could be ported to language-model KV caches with similar memory-quality trade-offs.
  • The method might allow even lower bit-widths than INT2 if the approximation error remains small on longer sequences.
  • Testing the correction on video lengths beyond those reported would show whether error accumulation eventually limits the gains.

Load-bearing premise

That the Jensen bias arising from convexity of the exponential is the primary driver of quality loss under KV quantization and that the second-order Taylor approximation removes it accurately enough without new artifacts.

What would settle it

Measure the distribution of attention weights on quantized versus full-precision keys with and without the correction; if the bias is not removed and quality metrics do not approach BF16 levels, the central claim is falsified.

Figures

Figures reproduced from arXiv: 2605.26266 by Felix Becker, Thomas Pfeil, Tuna Tuncer.

Figure 1
Figure 1. Figure 1: Qualitative comparison on MAGI-1 for two representative prompts. Columns show [PITH_FULL_IMAGE:figures/full_fig_p002_1.png] view at source ↗
Figure 2
Figure 2. Figure 2: Attention weights for MAGI-1 for the prompt “a person” under INT2 KV-cache quantization. [PITH_FULL_IMAGE:figures/full_fig_p002_2.png] view at source ↗
Figure 3
Figure 3. Figure 3: Illustration of the Jensen bias and its correction on a single attention score. Left: Quanti￾zation noise δ ∼ Uniform[−∆/2, ∆/2] with zero mean produces a noisy score sˆ = s + δ centered at s. Center: After exponentiation the distribution becomes right-skewed: its mean E[e sˆ ] strictly exceeds e s by the so-called Jensen bias. Right: Subtracting a correction b shifts the mean E[e sˆ−b ] closer to e s , la… view at source ↗
Figure 4
Figure 4. Figure 4: Shift in attention mass assigned to the cached block of tokens before (purple) and after [PITH_FULL_IMAGE:figures/full_fig_p009_4.png] view at source ↗
Figure 5
Figure 5. Figure 5: Trade-off between image quality, measured by PSNR, and memory footprint of the KV [PITH_FULL_IMAGE:figures/full_fig_p009_5.png] view at source ↗
read the original abstract

Chunk-wise autoregressive video diffusion models rely on a KV cache of previously generated chunks to avoid redundant computation, but this cache quickly becomes a memory bottleneck as videos grow longer. Methods that quantize the KV cache to low bitwidths reduce memory pressure but degrade video quality. We show that a key driver of this degradation is a systematic bias in attention weights: due to the convexity of the exponential in softmax attention, quantization noise inflates the contribution of cached keys, a phenomenon we call the Jensen bias. This effect causes quantized keys to steal attention mass from the unquantized current chunk. We derive a per-attention-score correction that removes this bias in expectation, computed on the fly from the quantization step sizes of the cached keys and the query norm. Using a second-order Taylor approximation, the additional computational overhead is negligible, and no additional memory is needed alongside the cache. Evaluated on MAGI-1, SkyReels-V2, and HY-WorldPlay at INT2 quantization, our correction recovers most of the quality lost to aggressive quantization, reaching near-BF16 video quality, and can outperform INT4 quantization while using 50% less memory.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit. Tearing a paper down is the easy half of reading it; the pith above is the substance, this is the friction.

Referee Report

2 major / 2 minor

Summary. The manuscript identifies Jensen bias in softmax attention as a primary cause of quality loss when quantizing the KV cache in chunk-wise autoregressive video diffusion models. It derives an on-the-fly per-score correction term based on the quantization step sizes of cached keys and the query norm, replaces the exact correction with a second-order Taylor expansion in the step size and query norm to keep overhead negligible, and reports that the method restores most of the quality lost at INT2 on MAGI-1, SkyReels-V2, and HY-WorldPlay, reaching near-BF16 levels and sometimes surpassing INT4 while halving memory.

Significance. If the derivation is correct and the second-order approximation remains accurate at INT2 step sizes, the approach offers a parameter-free, zero-extra-memory way to compress KV caches aggressively without retraining or auxiliary storage. This directly targets a scaling bottleneck for long video generation and could be adopted in production inference pipelines. The use of only observable quantities (step sizes and query norm) and the claim of negligible compute cost are notable strengths.

major comments (2)
  1. [Abstract; derivation of the correction term (presumably §3)] The central claim rests on the second-order Taylor expansion of the Jensen-bias correction being sufficiently accurate at INT2. The abstract and derivation section provide no explicit remainder bound or empirical measurement of the approximation error (e.g., comparison of the exact correction versus the Taylor term across the observed range of ||q|| and Δ). At INT2 the step size is large relative to typical key norms, so higher-order moments of the quantization error can dominate; without this check the “near-BF16 recovery” and “outperforms INT4 at half memory” statements cannot be verified.
  2. [Evaluation section and associated tables/figures] Table or figure reporting the main INT2 results (the three-model comparison) does not include an ablation that isolates the Taylor approximation itself—e.g., exact correction versus second-order approximation versus no correction. Such an ablation is required to confirm that the reported quality recovery is attributable to the proposed correction rather than other implementation details.
minor comments (2)
  1. [Abstract] Notation for the quantization step size Δ and the per-head query norm should be introduced once and used consistently; the abstract uses both without prior definition.
  2. [Evaluation section] The manuscript should state the precise video-quality metric (e.g., VBench, FVD, or human preference) and the number of samples used for the “near-BF16” claim so that the magnitude of recovery can be assessed.

Simulated Author's Rebuttal

2 responses · 0 unresolved

We thank the referee for the thoughtful and constructive report. The two major comments both concern validation of the second-order Taylor approximation at INT2. We agree these checks strengthen the paper and will add the requested empirical measurements and ablation in the revision.

read point-by-point responses

  1. Referee: [Abstract; derivation of the correction term (presumably §3)] The central claim rests on the second-order Taylor expansion of the Jensen-bias correction being sufficiently accurate at INT2. The abstract and derivation section provide no explicit remainder bound or empirical measurement of the approximation error (e.g., comparison of the exact correction versus the Taylor term across the observed range of ||q|| and Δ). At INT2 the step size is large relative to typical key norms, so higher-order moments of the quantization error can dominate; without this check the “near-BF16 recovery” and “outperforms INT4 at half memory” statements cannot be verified.

    Authors: We agree that an empirical measurement of the approximation error is necessary to support the claims at INT2. In the revised manuscript we will add a dedicated figure that directly compares the exact per-score Jensen-bias correction against the second-order Taylor term over the observed ranges of query norm and quantization step size drawn from the three evaluated models. This will quantify the residual error and confirm that higher-order terms remain negligible in practice. revision: yes

  2. Referee: [Evaluation section and associated tables/figures] Table or figure reporting the main INT2 results (the three-model comparison) does not include an ablation that isolates the Taylor approximation itself—e.g., exact correction versus second-order approximation versus no correction. Such an ablation is required to confirm that the reported quality recovery is attributable to the proposed correction rather than other implementation details.

    Authors: We concur that isolating the contribution of the Taylor approximation is important. The revised evaluation section will include an ablation table (or extended figure) reporting video quality metrics for the three models under three conditions: (i) no correction, (ii) exact Jensen-bias correction, and (iii) the second-order Taylor approximation. This will demonstrate that the quality recovery is driven by the bias correction and that the approximation introduces only minimal additional error relative to the exact term. revision: yes

Circularity Check

0 steps flagged

No circularity: derivation from Jensen inequality and Taylor expansion on observables

full rationale

The paper's central derivation starts from the known convexity of exp(·) and applies Jensen's inequality to the attention scores, then computes an expectation correction using the quantization step size Δ and query norm ||q||. The second-order Taylor expansion is introduced solely to reduce compute, with all terms expressed directly in terms of these observables. No fitted parameters, no self-referential predictions, and no load-bearing self-citations appear in the described chain. The result is therefore independent of its own outputs and self-contained against external benchmarks.

Axiom & Free-Parameter Ledger

0 free parameters · 2 axioms · 0 invented entities

The approach rests on the convexity of the exponential function and the validity of a second-order Taylor expansion for the correction; no free parameters or new entities are introduced in the abstract.

axioms (2)
  • standard math The exponential function inside softmax attention is convex, so Jensen's inequality applies to quantized keys.
    Invoked to explain why quantization inflates cached-key contributions.
  • domain assumption A second-order Taylor approximation suffices to compute the bias correction on the fly with negligible error.
    Used to claim negligible computational overhead.

pith-pipeline@v0.9.1-grok · 5751 in / 1350 out tokens · 38598 ms · 2026-06-29T22:47:58.050161+00:00 · methodology

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Reference graph

Works this paper leans on

6 extracted references · 6 canonical work pages · 4 internal anchors

  1. [1]

    SkyReels-V2: Infinite-length Film Generative Model

    URLhttps://arxiv.org/abs/2504.13074. Shuo Chen, Cong Wei, Sun Sun, Ping Nie, Kai Zhou, Ge Zhang, Ming-Hsuan Yang, and Wenhu Chen. Context forcing: Consistent autoregressive video generation with long context, 2026a. URL https://arxiv.org/abs/2602.06028. Ziyang Chen, Xing Wu, Junlong Jia, Chaochen Gao, Qi Fu, Debing Zhang, and Songlin Hu. Long- bench pro: ...

    work page internal anchor Pith review Pith/arXiv arXiv doi:10.48550/arxiv 2022

  2. [2]

    URLhttps://arxiv.org/abs/2311.17982. Ziqi Huang, Fan Zhang, Xiaojie Xu, Yinan He, Jiashuo Yu, Ziyue Dong, Qianli Ma, Nattapol Chanpaisit, Chenyang Si, Yuming Jiang, Yaohui Wang, Xinyuan Chen, Ying-Cong Chen, Limin Wang, Dahua Lin, Yu Qiao, and Ziwei Liu. Vbench++: Comprehensive and versatile benchmark suite for video generative models, 2024. URLhttps://ar...

    work page arXiv 2024

  3. [3]

    URLhttps://arxiv.org/abs/2310.06825. Weijie Kong, Qi Tian, Zijian Zhang, Rox Min, Zuozhuo Dai, Jin Zhou, Jiangfeng Xiong, Xin Li, Bo Wu, Jianwei Zhang, Kathrina Wu, Qin Lin, Junkun Yuan, Yanxin Long, Aladdin Wang, Andong Wang, Changlin Li, Duojun Huang, Fang Yang, Hao Tan, Hongmei Wang, Jacob Song, Jiawang Bai, Jianbing Wu, Jinbao Xue, Joey Wang, Kai Wang...

    work page internal anchor Pith review Pith/arXiv arXiv 2025

  4. [4]

    Mistral AI

    Accessed: 2026-05-07. Mistral AI. Mistral-7B-Instruct-v0.3 model card. https://huggingface.co/mistralai/ Mistral-7B-Instruct-v0.3, 2024. Accessed: 2026-05-07. Nilesh Prasad Pandey, Marios Fournarakis, Chirag Patel, and Markus Nagel. Softmax bias correction for quantized generative models, 2023. URLhttps://arxiv.org/abs/2309.01729. Qwen. Qwen2.5-32B-Instru...

    work page arXiv 2026

  5. [5]

    Wan: Open and Advanced Large-Scale Video Generative Models

    URLhttps://arxiv.org/abs/2503.20314. Zhou Wang, Alan C. Bovik, Hamid R. Sheikh, and Eero P. Simoncelli. Image quality assessment: From error visibility to structural similarity.IEEE Transactions on Image Processing, 13(4): 600–612, 2004. B. Widrow, I. Kollar, and Ming-Chang Liu. Statistical theory of quantization.IEEE Transactions on Instrumentation and M...

    work page internal anchor Pith review Pith/arXiv arXiv doi:10.1109/19.492748 2004

  6. [6]

    URLhttps://arxiv.org/abs/2602.02958. Zhuoyi Yang, Jiayan Teng, Wendi Zheng, Ming Ding, Shiyu Huang, Jiazheng Xu, Yuanming Yang, Wenyi Hong, Xiaohan Zhang, Guanyu Feng, Da Yin, Yuxuan Zhang, Weihan Wang, Yean Cheng, Bin Xu, Xiaotao Gu, Yuxiao Dong, and Jie Tang. Cogvideox: Text-to-video diffusion models with an expert transformer, 2025. URLhttps://arxiv.or...

    work page internal anchor Pith review Pith/arXiv arXiv 2025