PACE: Periodic and Channel-Aligned Encoding of Four-Bit Recurrent States
Abstract
PACE is a post-hoc four-bit quantizer for linear-attention recurrent states that rounds less frequently and uses column-wise scales to reduce error, enabling lower memory usage without fine-tuning.
Linear-attention language models keep their context in a recurrent state of fixed size, where an inference system stores one copy of this state for every sequence it decodes at the same time. Storing the state at lower precision therefore lets more sequences share an accelerator. Existing state quantizers round the state after every generated token and report that four bits lose accuracy unless the bit width is raised or the model is fine-tuned.\par We introduce PACE, a post-hoc four-bit state quantizer. PACE rounds the state only once every few generated tokens, as chunked kernels already do once per chunk during prefill, and keeps the tokens generated in between in a small buffer. Rounding less often reduces the error that accumulates in the state. Additionally, we find that the large entries of the gated-delta state gather in columns, because some channels of the update written at each token remain large across many tokens. PACE therefore shares each quantization scale within one column. It also keeps the heads that retain rounding error longest at eight bits. On four hybrid models from 4B to 35B parameters, PACE with mixed four- and eight-bit heads is competitive with eight-bit DAMP on the three larger models at about two thirds of its memory, PACE with four-bit heads throughout scores higher than Q-Mamba at equal memory, and rounding every few tokens preserves the decoding speedup of the buffer.
Community
We propose a new quantizer and also a new way on how to quantize states in linear-attention-based layers in language models. As the actual state is not materialized in the kernels for every token step, but at a chunk size, we can also adapt the number of rounding. This gives headroom for lower precision states and also aligns well with kv buffering methods during decoding.
(I hope you notice the nice acronym naming of our approach)
Rounding the state once every few tokens instead of every token is a simple idea, and the column-wise scale follows directly from where you found the large entries. Two questions. How do you pick the buffer length, and how sensitive is accuracy to it on the 4B model? And have you looked at CPU decoding with a single sequence, where state memory matters less than bandwidth? I am curious if the four-bit state still buys speed there, or if the gain is mostly about fitting more concurrent sequences on an accelerator.
Thanks for the questions!
- So buffer length becomes a trade-off hyperparameter. The longer the buffer the more bits you reintroduce back which are stored in memory, but the less rounding you do. Figure 4 in the paper depicts that. We settled for 8 token buffer for our headline numbers, as the best trade off. With this configuration we get close to now accuracy drop for the 4B model. Full precision with 80. 8 accuracy (avg) against our two configurations with 79.1 and 79.6.
- We havent measured CPU decoding. At batch 1 the weights dominate traffic and the recurrent state is a small fraction of that. So we expect four bit state to give little speedup there. The main benefit is memory per sequence.
Get this paper in your agent:
hf papers read 2610.29545 Don't have the latest CLI?
curl -LsSf https://hf.co/cli/install.sh | bash Models citing this paper 0
No model linking this paper
Datasets citing this paper 0
No dataset linking this paper
Spaces citing this paper 0
No Space linking this paper
Collections including this paper 0
No Collection including this paper
