LoopVQ: Sub 1-bit KV Cache Quantization for Looped Transformers

Paper

Introduction

A looped transformer is a model that reuses the same group of layers several times. It saves parameters, but each iteration step has to keep its own key/value cache, so the serving bottleneck shifts from the weights to the cache. This post summarizes the observations, design, and measured results of LoopVQ, a method that compresses that cache to under 1 bit per element.

Terms to know first

Looped transformer. An architecture that reuses the same $L_b$ decoder layers $T$ times. The $t$-th loop step takes the previous step’s output $h^{(t-1)}_{1:N}$ as input and computes $h^{(t)}_{1:N} = F_\theta(h^{(t-1)}_{1:N})$, and the final output distribution $p_j$ is read from the hidden state $h^{(T)}_j$ of the last step. Only one thing matters here. Each loop step leaves its own keys and values in the cache, so the cache is $T$ times larger than in a non-looped model that uses the same weights. The evaluated models Ouro-1.4B and 2.6B use $T=4$, and Huginn-3.5B runs 32 iterations.

KV cache. Memory that stores the keys and values of past tokens during autoregressive generation and uses them to compute attention for the next token. The paper defines that token $j$ produces a key $k^{(t,\ell)}_{j,h}\in\mathbb{R}^D$ and a value $v^{(t,\ell)}_{j,h}\in\mathbb{R}^D$ for each loop step $t$, layer $\ell$, and KV head $h$. The longer the context, the more heads and layers, and, as noted above, the more loop steps, the larger this memory grows.

Quantization. Compression that stores values at lower precision to reduce memory. The paper emphasizes that quantization does not discard tokens; it only shrinks the cache.

Vector quantization. A method that splits a $D$-channel vector into $D/G$ chunks of $G$ consecutive channels and stores each chunk as the index of the nearest entry in a codebook $\mathcal{B}=\{e_1,\dots,e_K\}\subset\mathbb{R}^G$. Following equation (2) of the paper exactly, the index is chosen by $a(z)=\arg\min_{r\le K}\|z-e_r\|_2^2$ and reconstruction is $\hat z = e_{a(z)}$. If the codebook size is $K=2^b$, each chunk costs $b$ bits, that is, $b/G$ bits per element. In the sub-1-bit regime this post targets, existing vector quantization loses a great deal of accuracy.

Bits per element. The number of bits used to store one channel in the cache. At 0.5 bits, that is 0.5 bits per element, a 32x compression relative to 16 bits (bf16).

Codebook and calibration. The codebook is initially built from offline calibration data. The paper summarizes that prior work uses $k$-means weighted by each chunk’s importance $\omega_n=\|\partial L/\partial z_n\|_2^2$ (equation 3). $\omega_n$ is the squared gradient norm with respect to the next-token loss $L$ and equals the diagonal trace of the empirical Fisher information. What matters here is that this weight alone cannot capture the problems that arise in a looped architecture.

RoPE (rotary position embedding). A technique that applies a position-dependent rotation to key vectors to represent relative positions. The paper quantizes keys before applying RoPE, reconstructs them in attention, and then applies RoPE. The reason is in Appendix C. RoPE rotates keys in the same channel group differently at each position, so a single code group tied together before rotation scatters across many entries in the codebook.

SGLang. A serving framework that executes structured LLM programs efficiently. This paper attaches a LoopVQ-specific attention kernel to it and measures efficiency.

1. Problem: the cache is the bottleneck, and redundancy remains along the loop axis

1.1 Saving parameters makes the cache larger

Figure 1a of the paper plots each model as a point. The horizontal axis is KV cache size per token (8KB–2MB), and the vertical axis is bf16 weight memory (2GB–32GB). Ouro-1.4B has 6x the KV cache per token of the non-looped Llama-3.1-8B, and at moderate context lengths its cache memory already exceeds its weights. In other words, when serving a looped model, the real bottleneck is the cache, not the weights.

Figure 1b of the paper plots GSM8K accuracy against bits per element (0.25–4), showing that existing vector quantization (CQ, TaSQ) collapses sharply once it goes below 1 bit (bf16 baseline 86.1, Ouro-2.6B).

1.2 The iteration axis still holds untapped redundancy

The $T$ caches of one layer are computed by applying the same projection weights to the hidden states of the same token. So the keys and values across steps resemble each other.

• Fig. 2a: Pre-RoPE keys of one attention head in Ouro-1.4B are plotted across four steps, and the four steps show nearly identical channel patterns. Across all of Ouro-1.4B, the correlation between adjacent steps rises from 0.69 between the first two steps to 0.94 between the last two.
• Fig. 2b: Color shows the mean $|r|$ over all (loop step, channel) coordinate pairs within a 16-channel group. The parts connecting the same channel across adjacent steps form bright lines, while different channels in the same group are dark. Numerically, the mean $|r|$ between adjacent steps of the same channel is 0.85, while between different channels in the same 16-channel group it is 0.16, and the gap is even wider for values.
• Fig. 2c: With bit width and codebook size fixed, it compares the reconstruction error (NMSE) of $k$-means under two grouping schemes: “$G$ channels of one step” and “$G/4$ channels across four steps.” Grouping along the loop axis reduces the error to roughly half at 0.5 bits on Ouro-1.4B, and reduces it substantially at every bit width on both Ouro and Huginn (§3, Fig. 2 of the paper).

The axis existing vector quantization uses is the “channel axis.” A looped model’s cache has one more redundancy axis next to it: the “iteration axis.”

1.3 But the perceived quantization error differs by step

When one grouped vector represents several steps together, it matters which step inside it incurred the error. An error introduced early propagates through all later steps, while an error in the last step stays confined to that step.

Table 1 of the paper quantizes only one step while keeping the others at full precision, then measures the KL divergence between the full-precision output and the quantized output (loop-axis grouping + Fisher-weighted $k$-means codebook, measured on sequences not used for calibration).

Quantizing the first step is the most harmful, and the third step the least. The gap between the two steps is about 5x.

2. A small example: one index stands in for four steps

Below is an illustrative example built from Ouro-1.4B’s actual settings ($D=128$ channels, $T=4$ steps, 8-bit code $K=256$, $G=16$ at 0.5 bits). It is a construction for explaining the mechanism, not measured values from the paper.

In other words, LoopVQ does not use more memory. It keeps the same number of storage slots and puts similar values together in those slots to reduce reconstruction error.

3. Method

3.1 Loop-axis vector quantization

The authors split the $T$ steps into consecutive blocks of $T_c$ steps (Ouro has $T=T_c=4$, so one block; Huginn has 32 steps in 8 blocks). One code covers $c=G/T_c$ channels in each of the $T_c$ steps within the block. Equation (4) of the paper defines this combined vector, and equation (5) is the rule by which each step slices out its own share (coordinate $iT_c+m$ of the combined vector) from that index for reconstruction.

• Because one code represents $T_c$ steps at once, bits per element stay at $b/G$. Memory does not grow.
• Keys are quantized before RoPE, reconstructed in attention, and then RoPE is applied (§2, Appendix C).
• The codebook is initialized with Fisher-weighted $k$-means. Here $\omega_n$ sums the squared gradients over all $T_c$ steps covered by the combined vector (§4.1).

3.2 Loop-aware codebook fine-tuning

Fisher-weighted $k$-means misses something. Because this objective penalizes each code vector’s error independently, it cannot express the cross terms that arise when an earlier step’s error propagates to later steps. Appendix A.1 shows that for small errors the objective can be approximated as $J(B)\approx\frac{1}{2}\mathbb{E}_s[\delta^\top F\delta]$, and that weighted $k$-means only matches the diagonal blocks $F_{nn}$ and always discards cross terms of the form $F_{nn'}$. No reweighting removes this limitation. Moreover, it allocates precision equally to every step inside the combined vector.

So the authors fix the model weights and fit the output directly. For a calibration sequence $s$, they take the full-precision cache $C$ and its output distribution $p_j$, compute the output $q_j(\mathcal{B})$ when all $T_c$ steps attend to $\hat C(\mathcal{B})$, the cache reconstructed from the codebook $\mathcal{B}$, and minimize $J(\mathcal{B})=\mathbb{E}_s[\frac{1}{N}\sum_j \mathrm{KL}(p_j\|q_j(\mathcal{B}))]$ from equation (6). The sum runs over the $N$ positions after the bf16 sink tokens. This objective captures both error propagation across steps and per-step sensitivity differences.

• Codebook entries are updated directly by backprop through this objective. The arg min selection does not need to be differentiated, so codebook entries receive exact gradients and no straight-through estimator is needed (Appendix A.2).
• Assignments are recomputed with the current codebook at every iteration, so the encoder at training time and the encoder at inference time do not diverge.
• As a result, KL divergence relative to the full-precision output drops to roughly half on Ouro-1.4B (Fig. 4 of the paper: starting from a $k$-means codebook and repeating fine-tuning).
• Training setup: fine-tune on 12 of the 16 calibration sequences (Adam, learning rate $10^{-3}$, one sequence per iteration, 4 epochs), then pick the epoch with the lowest KL on the remaining 4 (§5.1).

3.3 Combining with TaSQ

TaSQ handles the channel axis. Equation (10) in Appendix B is a transform that multiplies the key’s channels by weights $w$, divides by a per-token scale $\sigma$, and rearranges them with a permutation $P$. The goal is to make channels that share a code more correlated. LoopVQ handles the loop axis, so the two are orthogonal. TaSQ+LoopVQ applies the TaSQ transform first and then groups along the loop axis, so one code covers (the $c$ channels chosen on the channel axis) × (the $T_c$ steps on the loop axis).

The permutation $P$ is obtained from the $T_cD\times T_cD$ covariance matrix formed by stacking the $T_c$ step keys of a block, tying each step’s RoPE pair $\{i,i+D/2\}$ together so that a group has $G=cT_c$ coordinates. The only cost is storing the per-token key scale $\sigma$ in fp16, which is 16/(HD) bits per key element, about 0.008 bits for Ouro and 0.003 bits for Huginn. That is why Table 2 shows key bit widths of 0.508/0.5 (Ouro) and 0.503/0.5 (Huginn).

Huginn has segments without a loop axis. The two-layer prelude and coda are encoded separately per layer, like CQ. For reference, Huginn-3.5B runs 32 iterations of a four-layer core block and adds the prelude and coda, producing 132 KV caches per token.

3.4 SGLang implementation: when indices are created and how they are read

The flow in Fig. 3 of the paper is as follows. Each loop step writes its own cache, but in the encoding stage $c=G/T_c$ channels are grouped across $T_c$ steps to produce one index. That index is shared by the $T_c$ steps, and at inference each step $t$ reads the same index as its own key and value through its own codebook slice $\mathcal{B}_t$.

The storage and lookup structure in Fig. 5 of the paper reads in three stages.

1. Storage layout. The first $n_s$ tokens (sink) and the most recent $n_r$ tokens stay in bf16 across all steps, and only the tokens in between are stored as codebook indices. The cache is paged in units of 8 tokens, and $n_s$ and $n_r$ are multiples of 8, so a page is entirely bf16 or entirely quantized. Because one index is shared by $T_c$ steps, a quantized page has only one index row per step block, giving the same memory per token as encoding separately per step at the same $G$. The experiments use 8/64 for general benchmarks and 64/256 for reasoning and long-context benchmarks.
2. When indices are created. A combined vector can be encoded only after all $T_c$ steps of the block have been computed. So in prefill, each step’s keys and values are held until the last step of that block, then distances to all codebook entries are computed with tensor-core matrix multiplication and encoded in one go. In decode, the new token’s keys and values are written to the bf16 recent region (once the forward pass finishes, all steps for that token are present), and every 8 generated tokens the oldest page is converted to indices across all layers and steps, keeping the recent region at $n_r$–$n_r$+7 tokens.
3. How attention reads. At loop step $t$, the dedicated kernel loads the codebook slice $\mathcal{B}_t$ into shared memory once and looks up keys and values from registers using each token’s index. bf16 pages are read by SGLang’s Triton decoding kernel, and the two results are merged in Flash-Decoding style by combining the logsumexp of the splits. Quantized keys are reconstructed, after which the kernel applies RoPE using precomputed per-position sine and cosine tables. The number of indices read per head per token is $D/c$.

The implementation is available at https://github.com/mscheong01/LoopVQ .

4. Results

4.1 General benchmarks

The evaluated models are Ouro-1.4B, Ouro-2.6B, and Huginn-3.5B, and the metric is average accuracy (%) on GSM8K, MATH500, HumanEval, and MBPP. All use greedy decoding (Table 2 of the paper). CQ is the baseline vector quantization that uses a separate codebook per loop step and layer.

0.5 bits

0.25 bits

CQ and TaSQ already show a large gap from BF16 at 0.5 bits, and the gap is wider on math and code tasks. For Ouro-1.4B, MATH500 is 17.8 for CQ, 38.8 for TaSQ, and 69.4 for LoopVQ, and HumanEval is 14.6 for CQ, 24.4 for TaSQ, and 63.4 for LoopVQ. LoopVQ’s large lead over both baselines holds at 0.25 bits as well as 0.5 bits, TaSQ+LoopVQ approaches BF16 at 0.5 bits, and the additional gain grows at 0.25 bits. On Huginn-3.5B as well, LoopVQ and TaSQ+LoopVQ are higher than both baselines according to Table 2.

4.2 Long chain-of-thought (CoT) reasoning

Averages of the Thinking variants evaluated on AIME'24, AIME'25, GPQA-Diamond, HumanEval+, and MBPP+ (temperature 1.0, top-p 0.7, max 32,768 generated tokens per model card settings, averaged over 3 seeds, Table 3 of the paper).

0.5 bits

0.25 bits

The gaps are larger than on general benchmarks. At 0.5 bits, CQ drops to about a quarter of BF16 and TaSQ to a third to a half, but LoopVQ preserves most of the accuracy, and combined with TaSQ it narrows the gap to 1.1–1.6 points from BF16. At 0.25 bits, TaSQ+LoopVQ substantially outperforms LoopVQ alone and is the best quantization method on every task for both models.

4.3 Long-context retrieval

Fig. 6 of the paper is split into four panels (Ouro-1.4B and Ouro-2.6B × 0.5 bits and 0.25 bits). The horizontal axis is context length (4K–64K), the vertical axis is accuracy (%) averaged over RULER’s 8 NIAH tasks with 3 seeds, and the curves are BF16, CQ, TaSQ, LoopVQ, and TaSQ+LoopVQ. CQ and TaSQ stay well below BF16 at all lengths, while LoopVQ finds a considerable number of needles at 0.5 bits, and TaSQ+LoopVQ is highest at every length and bit width.

LongBench (16 English tasks) in Appendix D.5 follows the same order. In average accuracy, CQ and TaSQ are 7–13 points below BF16, LoopVQ is higher than both baselines in every setting, and TaSQ+LoopVQ comes within 1.5 points of BF16 at 0.5 bits.

4.4 When quantization lengthens reasoning, the per-token latency gain also fades

Fig. 10 of the paper shows the distribution of tokens generated by Ouro-1.4B-Thinking at 0.5 bits on GPQA-Diamond. The horizontal axis is the number of generated tokens (0–32K), the vertical axis is the fraction of problems (%), and the bars show the shares of 32K cap hits, early failures, and non-terminations. CQ and TaSQ exhaust the 32K budget on most problems, generating about 2x as many tokens as BF16. LoopVQ and TaSQ+LoopVQ, by contrast, follow BF16’s length distribution at 0.5 bits, and at 0.25 bits only TaSQ+LoopVQ still does (Appendix D.4).

Appendix D.4 measures the opposite direction as well: early termination. On BF16, the fraction of generations finishing within 2K tokens is 16% and 22% for Ouro-1.4B-Thinking and Ouro-2.6B-Thinking, while at 0.25 bits it rises to 33% and 59% for CQ and 30% and 47% for TaSQ. LoopVQ and TaSQ+LoopVQ stay at or below 30% in every setting. In other words, the per-token throughput gain carries through to actual time to complete an answer.

4.5 System efficiency

The measurement environment is a single RTX 6000 Ada 48GB, the target is Ouro-1.4B-Thinking, and sink and recent were kept in bf16 at 8 and 64 tokens (§5.3).

Kernel and latency. The decode kernel reads only $D/c$ indices per head per token, so it reads 8x (0.5 bits) to 16x (0.25 bits) fewer bytes than bf16 keys and values. Long-context decode is bound by cache reads, so this saving grows with context length (Fig. 8b of the paper, batch 1).

Prefill, by contrast, gets slightly slower because block-wise encoding cost is added. On Ouro-1.4B at context lengths 4K, 16K, and 32K, it is +9%, +7%, and +5% at 0.5 bits and +8%, +6%, and +4% at 0.25 bits, with the overhead shrinking as context grows (Fig. 8a). On Ouro-2.6B it is 5–9% slower. The time breakdown in Appendix D.3 reports that with batch 16 and 2K input prompts, attention time per generated token falls from 1.87ms for bf16 to 0.63ms (0.5 bits) and 0.54ms (0.25 bits) for LoopVQ, and that encoding accounts for under 1% of decode time and 5–7% of prefill time. At batch 1, attention dominates decode time at 32K for bf16, but for LoopVQ attention stays smaller than the linear layers.

Throughput. For long generations, the KV cache limits batch size.

In Fig. 7 of the paper, the horizontal axis is batch size (up to 128), the vertical axis is throughput (tok/s), and the three panels are for output lengths 4K, 8K, and 16K (all with 2K inputs). The gaps between the vertical line at bf16’s maximum batch and the LoopVQ curve are shown as 3.2x, 3.7x, and 4.3x. On Ouro-2.6B-Thinking, whose cache is heavier, they are 5.0x, 5.2x, and 6.1x, and on an A100 (40GB) they are 2.8x–3.7x (Appendix D.1, D.2). On the A100, 32K decode latency drops to −51% at 0.5 bits and −56% at 0.25 bits relative to bf16.

4.6 Effect of each design element and calibration data

Looking at the two elements separately (Table 4 of the paper, Ouro-1.4B), loop-axis grouping produces most of the gain. Adding fine-tuning (FT) on top has little effect on GSM8K, raises accuracy on HumanEval, and is most pronounced when added to TaSQ. Using both together gives the best or tied-best result in every setting. At 0.25 bits, grouping gives a larger gain over CQ when applied to TaSQ.

Calibration data also influences the results. Fig. 9 of the paper recalibrates while swapping the calibration corpus among WikiText-2, C4, and SlimPajama. CQ varies by up to 11 points on GSM8K, whereas TaSQ+LoopVQ varies by up to 4 points (GSM8K) and 1 point (HumanEval) and stays best on every corpus. Increasing the number of calibration sequences consistently raises accuracy (Table 5 of the paper, at 0.25 bits LoopVQ GSM8K 57.5→60.8, HumanEval 22.6→28.0, from 8 to 32 sequences).

Calibration cost is modest. On a single RTX 6000 Ada, Ouro-1.4B takes 6.1 min for $k$-means plus 2.3 min for fine-tuning, and Ouro-2.6B takes 14.3 min plus 5.9 min (on an A100, 5.5+2.0 min and 11.2+4.6 min), and even with more sequences the cost stays in the minutes range.

5. Limitations and when to apply it

Limitations stated by the authors. First, the efficiency comparison is only against BF16 serving. The serving efficiency of each method depends on kernels optimized for its own cache format, and since it is hard to match all methods to an equal level of optimization, they state that they avoided a comparison that would favor the most heavily optimized method. Second, they did not compare accuracy against FlashLoop and ResidualQuant, KV cache quantization methods designed for looped transformers. Both use scalar quantization at a minimum of 1 bit per element and target 2–4 bits, and they do not present settings below that, so they cannot be matched at the same bit width in the sub-1-bit regime LoopVQ targets.

Conditions for application. This technique assumes a loop axis, so it cannot be used as-is on ordinary LLMs without an iterative structure. Even with an iterative structure like Huginn, segments that are not looped, such as the prelude and coda, must be encoded separately in the existing way. The gain is pronounced for long contexts and long generations, and for short contexts or prefill-heavy workloads it can actually be slower because of encoding cost. The codebook must be calibrated per model and separately for the Thinking variants (depending on the calibration corpus, CQ varies by up to 11 points on GSM8K), and although calibration cost is in minutes, it must be redone each time.

Assessment. For serving looped models that have few parameters but are cache-bound, over long contexts and long generations, LoopVQ is a better option than existing vector quantization in the sub-1-bit regime. However, the efficiency comparison was against BF16 serving, and comparisons under kernel optimization equal to other quantization methods, as well as accuracy comparisons with FlashLoop and ResidualQuant, remain open.

Figures in this post are taken from the original arXiv:2611.05304 (CC BY 4.0). Only size and format were changed.