Looped transformers cut key value cache size with vector compression
LoopVQ: Sub 1-bit KV Cache Quantization for Looped Transformers
Machine LearningArtificial IntelligenceComputation and Language
Summary
Language AI models use repeating layers called looped transformers to save memory on their parameters. But these still store lots of information during each loop, which takes up a lot of space. The authors found that this stored info is quite repetitive and can be compressed without losing much detail. They created LoopVQ, a method that shrinks this stored data using clever coding and fine-tuning, allowing models to run faster and handle longer texts efficiently.
What this means in practice
- •For language model engineers: Reduce memory and improve speed when running looped transformer models with LoopVQ compression on KV caches.
- •For cloud ai service operators: Increase throughput and lower latency of large language model serving by integrating LoopVQ vector quantization.
Authors
Minsoo Cheong, Sungjoo Yoo
Abstract
Looped transformers are a parameter-efficient alternative to standard transformers for language modeling. By reapplying the same layers across multiple loop steps, recent looped language models match the performance of non-looped models several times their size. While this architecture effectively reduces parameter memory, each loop step still stores its own keys and values, so the key-value (KV) cache grows linearly with the number of loop steps and becomes the main memory bottleneck. We observe that the KV cache of looped models is highly redundant along the loop axis and that the loop steps differ substantially in their sensitivity to quantization. Based on these observations, we propose $\textbf{LoopVQ}$, a KV cache vector quantization method for looped transformers that codes channels of multiple loop steps jointly, and fine-tunes the codebooks to match the output of the full-precision model. At 0.5 bits per element, LoopVQ outperforms existing KV cache vector quantization methods for standard LLMs on general, long chain-of-thought reasoning, and long-context retrieval tasks with Ouro and Huginn models. Its design is orthogonal to these methods, and combined with TaSQ, it retains much of the BF16 accuracy even at 0.25 bits per element. We also implement LoopVQ in SGLang, where it nearly halves the decode latency at a 32K-token context and raises the peak serving throughput by up to $4.3\times$ and $6.1\times$ on the 1.4B and 2.6B Ouro models compared to the BF16 baseline.