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

LoopVQ. Sub 1-bit KV Cache Quantization for Looped Transformers. By Minsoo Cheong and one other, from Seoul National University. Published on the third of November.

Looped transformers save parameters by reusing them. The same layers are applied again and again across several loop steps, and recent looped language models match non-looped models several times their size. That solves the parameter memory problem, but it quietly creates another. Every loop step still stores its own keys and values. The key-value cache therefore grows linearly with the number of loop steps, and in these models it becomes the main memory bottleneck. The architecture that was meant to be lean ends up heavy in a different place.

Two observations open the way. First, the cache of a looped model is highly redundant along the loop axis: what one loop step stores overlaps heavily with what the others store. Second, the loop steps are not equally fragile. They differ substantially in how sensitive they are to quantization. LoopVQ is a vector quantization method for the KV cache, designed around both of these facts.

Rather than quantizing each loop step on its own, LoopVQ codes channels from several loop steps jointly, which is how it meets the redundancy along the loop axis. As a vector quantization method, LoopVQ works through codebooks, and it does not leave those codebooks as they are. They are fine-tuned so that the output of the quantized model matches the output of the full-precision model. What the tuning aims at, then, is the model's output as a whole, measured against the model it is compressing.

At half a bit per element, LoopVQ beats existing KV cache vector quantization methods built for standard language models. That holds on general tasks, on long chain-of-thought reasoning, and on long-context retrieval, tested with Ouro and Huginn models. Its design is orthogonal to those methods, so it can be combined with them. Paired with TaSQ, it keeps much of the BF16 accuracy even at a quarter of a bit per element.

The authors also implement LoopVQ in SGLang to measure serving. At a context of thirty-two thousand tokens, decode latency nearly halves. Peak serving throughput, measured against the BF16 baseline, rises by up to four point three times on the one point four billion parameter Ouro model. On the two point six billion parameter Ouro model, the rise is up to six point one times.

Looping cut the weights, and LoopVQ goes after what looping left behind: a cache that repeats itself from one loop step to the next, now held at well under one bit per element. The gains show up at both ends of the stack. On accuracy, it beats quantizers designed for ordinary transformers on reasoning and retrieval. On serving, it shortens decoding and multiplies throughput. And because it is orthogonal to existing methods, it can stack with them rather than replace them.

LoopVQ. Sub 1-bit KV Cache Quantization for Looped Transformers. By Minsoo Cheong and one other, from Seoul National University. Published on the third of November. This is the long version: the whole paper rather than the abstract.

Looped transformers run the same stack of layers several times over. Ouro runs its stack four times, and its one point four billion parameter model performs on par with non-looped models several times its size. The weights are shared across loop steps. The cache is not: every loop step stores its own keys and values. That small Ouro model needs six times the per-token cache of Llama three point one at eight billion parameters, and its cache outgrows its weights even at moderate context lengths. Existing vector quantization methods can compress it, but below one bit per element they lose much of their accuracy, and they treat each loop step's cache as unrelated to the others.

Vector quantization cuts each cached key or value into short sub-vectors of neighbouring channels, and stores each one only as the index of its nearest entry in a learned codebook. With eight-bit codes, a code spanning sixteen channels costs half a bit per element, and one spanning thirty-two channels costs a quarter. The codebooks are fitted by k-means on a small calibration set, and prior work weights each sub-vector by its sensitivity: the squared gradient of the next-token loss, a sum of diagonal Fisher information.

Two observations follow. The keys of a layer are strongly correlated across loop steps, because the same projection weights act on hidden states of the same token. A key channel correlates with the same channel at the adjacent loop step at about point eight five, but with the other channels of its own sub-vector at only point one six. And the loop steps differ in how much their errors matter. Quantizing only the first loop step is roughly five times as harmful as quantizing the third, because an early error perturbs the hidden state handed to every later loop step.

LoopVQ has two parts. The first is loop-axis vector quantization. On Ouro, all four loop steps form one block. Instead of a code spanning sixteen channels of one loop step, it spans four channels at each of four loop steps. One index serves the whole block, and the bitwidth is unchanged. In a controlled test that changes only which coordinates share a code, this regrouping roughly halves the reconstruction error of the smaller Ouro model at half a bit.

The second part is loop-aware codebook fine-tuning. With the model weights frozen, the k-means codebooks are refined so the quantized model's output distribution matches the full-precision one. Why not just reweight k-means? For small errors, the output divergence contains cross terms coupling the codes of one token at different loop steps and layers, because errors propagate through the recurrence. Any weighted k-means objective is a sum of terms each depending on a single coded vector, so it always drops those cross terms. Training the entries is simpler than it sounds. Picking the nearest entry has no useful derivative, but the vectors being coded come from the unquantized cache and do not depend on the codebook, so each entry receives an exact gradient summed over the tokens assigned to it. No straight-through estimator is needed. On the smaller Ouro model, fine-tuning roughly halves the divergence from the full-precision output.

For serving, the method is built into SGLang. A token can be encoded only once every loop step of its block exists, so new tokens wait in a full-precision window, and after every eight generated tokens the oldest page is encoded. At each loop step, attention reads the shared indices through that step's own slice of the codebook.

On general benchmarks with the two point six billion Ouro model, full precision averages eighty-one point five. At half a bit, the baseline CQ averages thirty-seven point seven, while LoopVQ reaches seventy-seven point seven and LoopVQ combined with TaSQ eighty point three. Long reasoning is harsher for the baselines, and LoopVQ again retains most of the full-precision accuracy.

The most direct test of the thesis is the ablation, which adds each component separately to the baselines. Coupling along the loop axis accounts for most of the gain. Applied to CQ at half a bit on grade-school maths, coupling alone lifts accuracy from forty-six point eight to seventy-five point four, while fine-tuning alone reaches only forty-eight point seven. Fine-tuning matters more on code generation. There is also a cost that per-token speed hides. On the graduate science questions, CQ and TaSQ exhaust the thirty-two thousand token budget on most problems and generate about twice as many tokens as full precision, while LoopVQ follows the full-precision lengths. In serving, decode latency roughly halves at a thirty-two thousand token context, and peak throughput rises by three point two to four point three times on the smaller model.

The paper has a limitations section naming two gaps. Efficiency is compared only with full-precision serving, not with the baseline quantizers, because each format depends on its own optimized kernels and the authors wanted to avoid tilting the comparison toward the method they optimized most. And accuracy is never compared with FlashLoop or ResidualQuant, the existing cache quantizers for looped models, which are scalar methods designed for two to four bits per element and cannot reach the sub-one-bit range. Beyond that, the general benchmarks are single greedy runs, the reasoning results show wide spreads across three seeds on the competition maths sets, prefill gets a few percent slower, and serving is measured only on Ouro.

The idea to carry away is that sharing weights creates a new axis of redundancy. When the same layers produce the cache again and again, the cache repeats itself along the loop, and a code that spans loop steps buys accuracy at no extra bits. Fitting those codes to the model's output rather than to reconstruction error then accounts for how errors travel through the loop. Combined with TaSQ, the two keep much of the full-precision accuracy even at a quarter of a bit per element.