PaperlayerView exact revision on arXiv

Source-linked AI summary

ResidualQuant: KV Cache Quantization for Looped Transformers with 2-Bit Residuals

Heejun Kim, Junyoung Lee, SangLyul Cho, Dongsu Han, Insu Han, Sehoon Kim

arXiv:2610.10381v1cs.LG

TL;DR

Looped Transformers’ KV cache grows with recurrent depth, and existing low-bit KV quantization can substantially degrade accuracy. ResidualQuant represents non-anchor loop states as low-precision residuals from final-loop KV states, combining scaling, rotations, and loop-wise mixed precision. Across models and benchmarks it improves the accuracy–memory tradeoff over rotation-based quantization, with accuracy close to BF16 and substantial throughput gains in the reported settings.

  • Problem

    Distinct KV states from each loop make cache memory grow with recurrent depth, while existing quantization methods often lose substantial accuracy below 4 bits.

    Source: §1 · pp. 1–2

  • Method

    ResidualQuant uses final-loop KV states as a reference and stores other loops’ differences as low-precision residuals, combined with least-squares scaling, residual rotations, and loop-wise mixed precision.

    Source: §1 · p. 2, §7 · p. 12

  • Results

    Across models and reasoning and code-generation benchmarks, ResidualQuant retains accuracy close to BF16 while reducing theoretical KV storage by up to 80.7% and outperforming rotation-based quantization at comparable memory budgets.

    Source: §1 · p. 2, §7 · p. 12

  • Takeaways & Limitations

    On an RTX 5090, ResidualQuant improves fixed-batch decode throughput by up to 2.73× and peak throughput by up to 4.15× by enabling larger batches.

    Source: §7 · p. 12

Abstract

from arXiv · show

Looped Transformers improve parameter efficiency by repeatedly applying shared Transformer blocks over multiple recurrent loops, increasing computational depth without increasing the parameter count. However, KV cache memory still scales with the number of loops, becoming a key memory bottleneck that limits batch size and inference throughput. KV cache quantization can alleviate this bottleneck, but existing methods often suffer substantial accuracy degradation at aggressive low-precision regimes. We observe that looped Transformers offer a unique opportunity: KV states across loops are highly similar. Based on this observation, we propose ResidualQuant, which uses the final-loop KV states as a reference and represents the remaining loops with low-precision residuals. Our method further combines least-square scaling and rotations applied to the residuals, as well as loop-wise mixed precision, to enable accurate quantization down to INT2 while retaining efficient reconstruction. Across multiple looped Transformer models and mathematical reasoning and code generation benchmarks, ResidualQuant consistently improves the accuracy-memory tradeoff over state-of-the-art rotation-based KV quantization. In particular, our method retains accuracy close to BF16 under mixed-precision settings while reducing theoretical KV storage by 80.7%, achieving up to 13.0% higher accuracy than the rotation-based baseline at the same memory budget. On an RTX 5090, the reduced KV memory traffic improves fixed-batch decode throughput by up to 2.73x, while the smaller memory footprint enables up to 2x larger batches, improving peak throughput by up to 4.15x.

1. Introduction

Looped Transformers gain computational depth through shared blocks, but their distinct per-loop KV states make cache memory grow with recurrent depth. ResidualQuant exploits inter-loop similarity to compress those states while preserving loop-specific information.

1. Introduction

  • Looped Transformers use fewer weights than models with distinct parameters at the same computational depth, but distinct keys and values from each loop make KV cache memory grow with recurrent depth.That growth is in addition to scaling with context length and batch size.Source: §1 · p. 1
  • KV quantization preserves distinct loop states while reducing storage precision, but existing methods often lose substantial accuracy below 4 bits.Sharing or reusing KV states across loops can instead discard loop-specific information and degrade accuracy when representations differ.Source: §1 · p. 2
  • Across loops, KV states are highly similar, allowing ResidualQuant to use smaller residuals that are more robust to low-bit quantization.This approach aims to preserve loop-specific information rather than sharing states outright.Source: §1 · p. 2
  • 80.7% lower theoretical KV storage relative to BF16 with INT2/4, while retaining accuracy close to BF16; ResidualQuant also improves accuracy–memory tradeoffs over rotation-based quantization.The reported evaluation spans Ouro-1.4B and Huginn-3.5B on reasoning and code-generation benchmarks.Source: §1 · p. 2
  • ResidualQuant uses final-loop KV states as an anchor and stores other loops’ differences as low-precision residuals, with least-squares scaling, learned rotations, and loop-wise INT4/INT2 mixed precision.These components balance low-bit accuracy with runtime efficiency.Source: §1 · p. 2

2. Preliminary

Group-wise affine quantization adapts quantization levels to each group’s value distribution, while orthogonal rotations redistribute large magnitudes across channels before quantization. The paper uses learned OptR-H rotations for INT2 KV quantization.

2.1. Group-wise Affine Quantization

  • Affine quantization uses a group-shared scale and offset to adapt quantization levels to each group’s distribution.At INT2, the offset permits four reconstruction levels rather than the three levels of zero-centered symmetric quantization.Source: §2.1 · pp. 2–3

2.2. Rotation-based Quantization

  • Orthogonal rotations redistribute large magnitudes across channels to mitigate outlier-driven increases in group scale and the resulting loss of resolution for smaller values.The rotation and inverse rotation preserve the vector norm.Source: §2.2 · p. 3
  • For INT2 KV cache quantization, the paper uses OptR-H, a Hadamard-initialized OptR variant that learns head-wise orthogonal rotations to minimize attention output error.The learned rotations are used for inter-loop residual quantization.Source: §2.2 · p. 3

3. Methodology

ResidualQuant exploits similarity between loop KV states by quantizing residuals relative to a reference, then combines a Last-loop anchor, least-squares scaling, residual rotations, and loop-wise mixed precision. The resulting method reaches BF16-level MATH500 accuracy in the described Ouro-1.4B mixed-precision setting.

3. Methodology

  • ResidualQuant combines Last-loop references, least-squares scaling, residual rotations, and loop-wise mixed precision to support accurate quantization down to INT2 with efficient reconstruction.The methodology section reports BF16-level performance across the evaluated models and benchmarks.Source: §3 · p. 3, §3.4 · p. 5

3.1. Motivation

  • Looped Transformers apply the same key and value projections across loops, and similar loop inputs yield smaller KV differences that motivate quantizing residuals instead of each state independently.The paper connects this motivation to observed key patterns and residuals with fewer outliers than original keys.Source: §3.1 · pp. 3–4

3.2. Residual Quantization for KV Cache

  • The Previous-loop reference chains reconstructions, but a shared anchor lets each non-anchor loop load only the anchor and its own residual; First-loop shared-anchor reconstruction improves decode throughput by up to 1.62× over chained reconstruction.The Last-loop policy instead uses the final-loop KV state as the reference, reducing early-loop reconstruction error and improving Ouro-1.4B MATH500 accuracy from 27.8% to 57.2% under uniform INT2.Source: §3.2 · pp. 4–5
  • Least-squares scaling reduces key prediction MSE relative to α_r=1 and raises Last-loop uniform-INT2 MATH500 accuracy from 57.2% to 60.6%.The reference and scalar coefficient define a residual, which is quantized and combined with the reference for on-demand reconstruction.Source: §3.2 · pp. 4–5

3.3. Rotation for KV Cache Residual

  • OptR-H rotations are applied to residuals, not KV states, and improve Ouro-1.4B INT2 MATH500 accuracy from 60.6% to 71.6%.Separate key and value rotations are learned for each layer and head, then shared across loops.Source: §3.3 · p. 5

3.4. Loop-wise Mixed Precision

  • For four-loop Ouro-1.4B, assigning INT4 to the final-loop anchor and INT2 to residuals yields [2, 2, 2, 4] and, including quantization metadata, reduces KV memory footprint 5.2× versus BF16.The anchor stores the full KV state used as the shared reference for reconstructing all other loops.Source: §3.4 · p. 5
  • 76.0% MATH500 accuracy with loop-wise mixed precision, matching BF16 with only a 19% increase in KV cache storage over uniform INT2.The method combines LS scaling, Last-loop references, residual rotations, and anchor-focused mixed precision.Source: §3.4 · p. 5

4. Implementation

ResidualQuant’s implementation reconstructs compressed KV inside attention kernels and integrates cache handling with vLLM. Prefill stores the final-loop anchor and residuals after the final loop, while decode temporarily retains current-token KV until that anchor is available.

4.1. Kernel Implementation

  • Custom CUDA attention kernels reconstruct anchors and residuals tile by tile and use them immediately, avoiding writing reconstructed KV to GPU memory and reading it back.With Last-loop references, each residual loop needs only its anchor and its own residual, without sequential reconstruction of intermediate loops.Source: §4.1 · p. 6
  • Group-wise shared-memory lookup tables, asynchronous tile loads, and reuse of rotated queries reduce repeated reconstruction work; softmax and attention accumulation use FP32.For INT2 residuals, the lookup table contains four dequantized V residual values for codes 0–3.Source: §4.1 · p. 6

4.2. vLLM Integration

  • The vLLM integration retains request scheduling and paged KV management, using Triton kernels for KV storage and separate prefill and decode paths.Prefill quantizes the final-loop anchor and computes non-anchor residuals after the final loop; decode buffers current-token KV until its anchor is available, then stores the anchor and residuals.Source: §4.2 · p. 6

5. Experiments

The experiments evaluate ResidualQuant across two looped Transformer models, mathematical reasoning and code-generation benchmarks, component ablations, decode throughput, and activation quantization. They report improved accuracy–memory tradeoffs, throughput gains, and improvements to low-bit activation baselines.

5.1. Main Results

  • 52.37 vs. 52.94 average score under mixed INT2/4 with g=32, while reducing KV memory by 80.7%, brings ResidualQuant close to BF16.At the same setting, OptR-H is 5.77% below BF16; the table summarizes model–benchmark results and effective KV bits.Source: Table 2 · p. 7, §5.1 · p. 8

5.2. Ablation Studies

  • 3.6% gain in MATH500 accuracy, from 70.2% to 73.8%, is reported for LS scaling with INT2/4 and the Last-loop reference.LS scaling improves accuracy in most tested reference and precision configurations and reduces key prediction MSE by up to 24.3%.Source: §5.2 · p. 8, Table 3 · p. 8
  • 1.62× throughput improvement is reported for shared-anchor reconstruction; Last-loop accuracy is comparable to or better than Previous-loop while avoiding chained reconstruction.First-loop has lower accuracy than Last-loop, while Next-loop retains chained reconstruction dependency and its memory-traffic overhead.Source: §5.2 · pp. 8–9
  • At INT2, residual quantization reaches 67.8% MATH500 accuracy without rotation and 75.2% with OptR-H; rotation gains vary by precision and method.At INT4, residual quantization reaches 74.2% without rotation and 77.2% with OSCAR.Source: §5.2 · p. 9, Table 4 · p. 9
  • 73.8% vs. 61.8% MATH500 accuracy without rotation at g=32 favors Ascending over Descending precision allocation for residual quantization.Ascending assigns higher precision to the shared anchor and also outperforms uniform INT2 with g=16 at the same effective bitwidth of 3.1.Source: §5.2 · p. 9, Table 5 · p. 9

5.3. Throughput Improvement

  • 2.73× decode throughput at 16k context and batch size 2 rises to a 4.15× peak-throughput gain when the largest feasible tested batch increases from 2 to 4.At 8k context, the largest feasible tested batch rises from 4 to 16, yielding a 3.27× peak-throughput gain over BF16.Source: §5.3 · p. 10, Figure 4 · p. 10

5.4. Extension to Activation Quantization

  • 49.51% to 56.86% GSM8K accuracy for Ouro-1.4B’s Naive baseline accompanies improvements in nearly all evaluated method–benchmark combinations across both models.The experiment applies residual quantization with LS scaling to W4A4 activation quantization while keeping KV caches in BF16.Source: §5.4 · pp. 10–11, Table 6 · p. 11

6. Related Work

Related work covers KV quantization, looped Transformer architectures, and methods for reducing the costs of looping. ResidualQuant is presented as a quantization approach that uses inter-loop redundancy while retaining loop-specific KV states.

KV Cache Quantization

  • KV quantization methods address distribution outliers through distribution-tailored quantization or rotations; INT2 methods OSCAR and OptR optimize rotations while retaining recent KV in BF16.ResidualQuant combines INT2 residuals with INT4 anchors without a separate BF16 buffer for recent past tokens.Source: §6 · p. 11

Looped Transformers

  • Looped Transformer research explores shared blocks, adaptive computation, and iterative refinement, with prior work linking repeated computation to length generalization and latent reasoning.SMELT studies sparse MoE looping under matched FLOPs, parameter, and KV-cache budgets.Source: §6 · p. 11

Optimizing Looped Transformers

  • Methods for loop costs include speculative decoding, KV reuse or sharing, and W4A4 quantization; ResidualQuant preserves loop-specific states without model retraining or pretrained-weight updates.The paper positions its residual quantization as complementary to existing approaches and identifies Last-loop references, LS scaling, shared rotations, and mixed precision as design choices.Source: §6 · pp. 11–12

7. Conclusion

ResidualQuant reconstructs loop-specific KV states from a final-loop anchor and low-precision residuals. Its combined design maintains accuracy close to BF16 while reducing KV storage and improving decode throughput.

7. Conclusion

  • Combining final-loop references, least-square scaling, residual rotations, and loop-wise mixed precision enables accurate quantization down to INT2 with efficient reconstruction.ResidualQuant reconstructs each loop’s KV states from the final-loop KV states and low-precision residuals.Source: §7 · p. 12
  • Across mathematical reasoning and code-generation benchmarks, ResidualQuant remains close to BF16 accuracy, reduces theoretical KV storage by up to 80.7%, and outperforms rotation-based quantization at comparable memory budgets.On an RTX 5090, fixed-batch decode throughput improves by up to 2.73× and peak throughput by up to 4.15×.Source: §7 · p. 12
  • The residual reconstruction principle also improves accuracy when extended to W4A4 activation quantization with FlatQuant and LoopQ.The paper concludes that inter-loop similarity offers a path to aggressive low-bit inference for looped Transformers.Source: §7 · p. 12

A. Experimental Setup Details

The experiments evaluate Ouro-1.4B and Huginn-3.5B on single-GPU Blackwell systems, using mathematical-reasoning and code-generation benchmarks. Effective KV bitwidth accounts for quantized codes and associated metadata, including residual scaling coefficients when applicable.

A.1. Models and Hardware

  • Ouro-1.4B runs four loops and Huginn-3.5B runs 32, with residual references confined to groups of four loops; evaluations use one GPU, BF16 weights, vLLM, and decode CUDA Graphs.The hardware is NVIDIA RTX 5090 or RTX PRO 6000 Blackwell, the latter with 96 GB of memory.Source: §A.1 · p. 16

A.2. Benchmarks and Prompts

  • The evaluation covers 1,319 GSM8K, 500 MATH500, 164 HumanEval, and 378 MBPP problems, using one greedy-decoded response per problem with model-specific prompts and stopping conditions.GSM8K uses exact match, MATH500 uses Math-Verify, and code uses base pass@1 on original tests with EvalPlus tools.Source: §A.2 · p. 16

A.3. Effective KV Bitwidth

  • Effective KV bitwidth divides the bits for token-dependent quantized codes and metadata by the number of original KV elements; residual quantization with LS scaling also stores BF16 scaling coefficients.For Ouro-1.4B, those coefficients add 0.094 bits per KV element; residual quantization with LS scaling uses approximately 2.6, 3.1, 3.1, and 3.5 bits across the four listed configurations.Source: §A.3 · pp. 16–17

B. Analysis Details

The analysis compares key-residual prediction error under alternative scaling choices and measures key reconstruction error under First-loop and Last-loop policies. It also specifies the data and settings used to visualize key distributions relative to the last-loop anchor.

B. Analysis Details

  • Key prediction MSE compares unit scaling with least-squares scaling against the same reconstructed last-loop anchor, using uniform INT2 or INT4 at group size 32 without rotation.The residual is measured before residual quantization, and MSE is averaged equally across 512 GSM8K training examples.Source: B. Analysis Details · p. 17
  • Key-reconstruction NMSE evaluates First-loop and Last-loop on BF16 post-RoPE keys from all 500 problems, with teacher forcing on full prompts and gold solutions, using uniform INT2, group size 32, αᵣ=1, and no rotation.Both policies use the same BF16 states and packed-cache writer and reader, without feedback from quantized attention.Source: B. Analysis Details · p. 17

C. Key Distributions Relative to the Last-loop Anchor

  • Figure 5 shows absolute key values for loops 1–3, the loop-4 anchor, and their differences on shared scales, using one Ouro-1.4B layer and KV head.The post-RoPE keys come from a 573-token input; differences are shown before LS adjustment or quantization.Source: C. Key Distributions Relative to the… · p. 17

D. KV Reconstruction for Four Reference Choice Configurations

Shared-anchor reconstruction reduces the number of compressed components needed to recover loop KV states, while anchor interval can be lengthened with similar GSM8K accuracy. The section compares reconstruction dependencies and reports storage and accuracy changes for longer Huginn loop groups.

D. KV Reconstruction for Four Reference Choice Configurations

  • Figure 5’s rows represent loops 1–3 and its columns show absolute K_r, K_4, and K_r−K_4 on shared scales.The Last-loop policy uses loop 4 as the reference for loops 1–3.Source: Figure 5 · p. 18, D. KV Reconstruction for Four… · p. 18
  • First-loop and Last-loop policies reconstruct each non-anchor loop from a shared anchor and its own residual in two reads; Previous-loop and Next-loop chains can require four reads in a four-loop model.Reconstructing loop 1 under Last-loop requires the anchor followed by residuals of loops 3, 2, and 1, while the model still executes loops 1 through 4.Source: D. KV Reconstruction for Four… · pp. 18–19
  • For L loops, shared-anchor reconstruction uses 2L−1 logical component accesses versus L(L+1)/2 for neighboring-loop reconstruction; the ratio is 10/7 for L=4 and 528/63 for L=32.These are logical accesses, not byte counts accounting for bitwidth, metadata, cache hits, or actual HBM traffic.Source: D. KV Reconstruction for Four… · p. 19

E. Huginn-3.5B Anchor Interval

  • Increasing Huginn-3.5B anchor groups from four to 16 loops changes GSM8K accuracy from 42.38% to 42.68% at 3.5 to 3.3 effective bits, and from 42.61% to 42.30% at 3.1 to 2.8 bits.The 16-loop configurations reduce logical KV storage by 79.7% and 82.6% relative to BF16, respectively; each group uses a final INT4 anchor.Source: E. Huginn-3.5B Anchor Interval · p. 19

F. Full KV Quantization Ablations

The ablations compare reference, scaling, and rotation choices for KV quantization, while the appendices describe rotation calibration and inference and activation-residual quantization. They also document the relevant baselines and calibration procedures.

F. Full KV Quantization Ablations

  • Table 9 compares reference, scaling, and rotation choices across uniform INT2, uniform INT4, and mixed INT2/4; neither LS coefficient choice uniformly dominates in downstream accuracy.The table marks Last-loop references with LS scaling and OptR-H as the selected configuration.Source: Table 9 · p. 20, F. Full KV Quantization Ablations · p. 20

G.1. Calibration

  • Calibration collects BF16 Q/K/V states, uses reconstructed quantized references for residual targets, and selects per-head matrices by holdout loss.For Ouro-1.4B, rotation learning uses FP32 Adam for 80 steps, with two calibration examples and their last 64 queries at each step.Source: §G.1 · p. 21

G. Rotation Calibration and Inference

  • OSCAR derives rotations from covariance statistics, while OptR-H optimizes attention error; both use fixed rotation matrices during inference.OSCAR uses covariance eigenvectors, Hadamard mixing, and permutation; OptR-H maintains orthogonality through its parameterization.Source: §G.1 · p. 21

G.2. Inference

  • After calibration, the matrices are fixed; decode reconstructs past rotated KV from compressed codes and computes attention without storing fully reconstructed past KV.The cache stores integer codes, group-wise FP8 scales and offsets, and BF16 LS coefficients for residuals.Source: §G.2 · pp. 21–22

H.1. Baseline Methods and Implementation

  • The activation experiments compare Naive, FlatQuant, and LoopQ baselines, with LoopQ reproduced from its published methodology because the authors’ code was unavailable.FlatQuant learns invertible affine transformations; LoopQ uses loop-aware scaling and transformations, transition adapters, and trajectory-aware calibration.Source: §H.1 · p. 22

H.2. Residual Formulation

  • Activation residual quantization directly quantizes the first loop, then quantizes later loops’ residuals from reconstructed previous-loop activations using token-wise LS scaling.It is applied to QKV, attention output, MLP gate/up, and down projections.Source: §H.2 · p. 22

H.3. Calibration

  • FlatQuant’s residual variant receives 15 additional calibration epochs, while the described LoopQ calibration procedures differ by model and available records.For Ouro-1.4B, the archived direct LoopQ bundle lacks its original calibration record, so the Huginn-3.5B schedule is not attributed to it.Source: §H.3 · pp. 22–23

H.4. Evaluation Protocol

The evaluation protocol specifies benchmark prompts, scoring, and generation limits for the activation experiments. Additional INT2 throughput results report decode speed and feasible batch size across context lengths.

H.4. Evaluation Protocol

  • Table 6 evaluates 1,319 GSM8K, 500 MATH500, 164 HumanEval, and 378 MBPP problems, with benchmark-specific prompts, scoring, and generation limits.GSM8K uses flexible-extract exact match; HumanEval and MBPP report base pass@1 on original tests.Source: §H.4 · p. 23

I. Decode Throughput with INT2 Group Size 16

  • At 8k context, INT2 group-size-16 decode throughput rises from 113.1 to 367.4 tokens/s (3.25×), while the largest tested feasible batch rises from 4 to 16.Figure 6 covers Ouro-1.4B on an RTX 5090 at 2k, 4k, 8k, and 16k context.Source: §I · p. 23, Figure 6 · p. 23
  • Figure 6 measures 128 generated tokens over the 127 decode steps after the first output token, averaging three trials after one warm-up.The plotted comparison uses INT2 residuals at group size 16 and an INT4 anchor at group size 32.Source: Figure 6 · p. 23

J. Comparison with FlashLoop

The comparison examines ResidualQuant against FlashLoop and evaluates combining ResidualQuant with FlashLoop’s sparse updates and attention. The reported results show improved accuracy and lower KV bitwidth in that combined setting.

J. Comparison with FlashLoop

  • ResidualQuant is described as orthogonal to FlashLoop’s sparse-update and attention mechanisms, and the comparison applies those mechanisms to ResidualQuant.The evaluation uses Ouro-1.4B on all 500 MATH500 problems with matched model revision, prompts, decoding, and generation limit.Source: J. Comparison with FlashLoop · p. 24
  • Under the same quantization setting, the Previous-loop policy achieves only 62% of Last-loop decode throughput because it requires larger KV memory traffic.Previous-loop reconstruction loads the anchor and intermediate residuals, whereas Last-loop reconstruction refers each residual loop directly to the final-loop anchor.Source: J. Comparison with FlashLoop · p. 24
  • With uniform INT2 and no rotation, Previous-loop and Last-loop unit scaling each reach 69.2% MATH500 accuracy; Last-loop with LS scaling and OptR-H reaches 75.2%, versus 64.2% for OptR-H alone and 75.0% for BF16.FlashLoop reports below 20% accuracy when its residual quantization is reduced to INT2 in its ablation setting.Source: J. Comparison with FlashLoop · p. 24
  • With sparse updates and attention, ResidualQuant at g=16 reaches 74.4% MATH500 accuracy versus FlashLoop’s 72.8%, while reducing prefill and decode average KV bitwidth from 2.64 to 2.18 and 4.50 to 3.47.At g=32, bitwidth falls to 2.01 for prefill and 3.09 for decode, with 73.4% accuracy.Source: Table 10 · p. 24, J. Comparison with FlashLoop · pp. 24–25
Loading 2610.10381v1…