[KongchangAI]
· 2 min read· 1,148 words

PACE: A New 4-Bit Recurrent State Quantization Method That Makes Linear Attention Models More Memory-Efficient

PACE: A New 4-Bit Recurrent State Quantization Method That Makes Linear Attention Models More Memory-Efficient

PACE achieves post-hoc 4-bit quantization of linear attention model states via periodic rounding and channel-aligned encoding, greatly reducing memory usage.

PACE is a post-hoc quantization method for recurrent states in linear attention language models, addressing precision loss in low-precision state storage. Existing quantizers round after every token, causing significant cumulative errors at 4-bit. PACE proposes periodic rounding to defer rounding every few tokens and channel-aligned encoding to share quantization scales within columns, plus mixed 8-bit precision for sensitive heads. Experiments show mixed-precision PACE matches 8-bit baselines at two-thirds the memory, and all-4-bit PACE outperforms Q-Mamba at equal memory. PACE requires no fine-tuning, offering practical memory savings for long-context multi-sequence inference.

Linear attention language models store context in a fixed-size recurrent state, and inference systems keep one copy of that state per concurrently decoding sequence. Lowering state storage precision allows more sequences to share an accelerator, but existing quantizers round after every generated token, causing 4-bit precision loss unless bit width is increased or the model is fine-tuned. The recently proposed PACE method achieves post-hoc 4-bit state quantization through periodic rounding and channel-aligned encoding, delivering competitive results across multiple models.

Why Do We Need Low-Precision Recurrent States?

The core advantage of linear attention models (such as Mamba and RWKV) is that inference only requires maintaining a fixed-size state, unlike standard Transformers which need to cache the entire history of key-value pairs. However, when decoding multiple sequences simultaneously, each sequence requires an independent copy of the state. State precision directly determines memory footprint: 4-bit states can save half the memory compared to 8-bit, meaning the same accelerator can serve nearly twice as many sequences.

Existing quantization schemes (such as DAMP and Q-Mamba) round the state immediately after each generated token. This high-frequency rounding introduces cumulative errors, causing a noticeable drop in model accuracy at 4-bit precision. PACE's starting point is precisely to resolve this contradiction: how to maintain low-precision storage without frequent rounding.

Linear attention models replace the key-value cache of standard Transformers with a recurrent state. Representative architectures include Mamba (based on selective state space models) and RWKV (based on a linear attention variant). During inference, these models only need to update a fixed-dimensional state vector or matrix at each decoding step, rather than storing all historical key-value pairs, so memory usage is independent of sequence length. However, when decoding multiple sequences concurrently, each sequence must independently maintain a state copy, and state precision directly affects the number of sequences that can run concurrently. For example, quantizing the state from FP16 to INT4 reduces memory usage to one quarter, theoretically allowing concurrency to increase fourfold. However, low-precision quantization introduces rounding errors, and errors accumulate in recurrent updates, potentially degrading model output quality.

PACE's Two Core Designs

Periodic Rounding: Reducing Error Accumulation

PACE's key observation is that the chunked kernel already rounds by block during the prefill phase, rather than rounding every token. Inspired by this, PACE defers the rounding operation to execute once every few generated tokens, with intermediate generated tokens temporarily stored in a small buffer. The effect is a significant reduction in cumulative rounding errors during state updates. Experiments show that this "less rounding" strategy effectively mitigates the precision loss of 4-bit quantization while maintaining decoding speed improvements.

The chunked kernel is an implementation approach for parallel computation in linear attention models. It divides long sequences into fixed-size chunks, computes within chunks in parallel, and passes state between chunks. During the prefill phase, chunked kernels typically update state and perform rounding by chunk rather than token-by-token. PACE borrows this idea and adopts a similar deferred rounding strategy during the decoding phase: state updates for multiple tokens are temporarily stored in a buffer and rounded uniformly after accumulating to a certain amount. This reduces the frequency of rounding operations, thereby lowering the cumulative effect of rounding errors while maintaining computational parallelism.

Channel-Aligned Encoding: Exploiting State Column Structure

The research team further discovered that in gated delta states, larger entries tend to cluster in columns. This is because updates for certain channels consistently maintain larger values across multiple tokens. PACE exploits this structural property by sharing each quantization scale within a column, rather than independently assigning a scale to each element. This channel-aligned encoding better matches the actual distribution of the state and improves quantization efficiency.

Additionally, PACE uses 8-bit precision for attention heads where rounding errors persist longer, forming a mixed-precision scheme. This flexibility allows the model to maintain higher precision in critical parts while overall memory usage remains low.

PACE method illustration

The gated delta state is a state representation in linear attention models, commonly found in architectures such as GLA (Gated Linear Attention). Its update rule includes a gating mechanism and a delta rule, causing certain columns of the state matrix (corresponding to specific channels) to consistently receive larger updates across multiple tokens, leading to large values clustering in columns. PACE exploits this structural property by sharing a quantization scale for each column (i.e., all elements within a column use the same scaling factor), rather than quantizing each element independently. This column-wise scale sharing better matches the actual distribution of the state, reducing quantization errors while also lowering the overhead of storing quantization scales.

Experimental Results: The Trade-off Between 4-Bit and 8-Bit

PACE was evaluated on four hybrid models with parameter sizes ranging from 4B to 35B. Key results include:

  • PACE with mixed 4-bit and 8-bit heads performs comparably to 8-bit DAMP on three larger models, while using only about two-thirds of the memory.
  • PACE with all 4-bit heads scores higher than Q-Mamba at the same memory usage.
  • The strategy of rounding once every few tokens preserves the decoding speedup from buffering.

These results demonstrate that PACE provides a practical post-hoc solution for state quantization in linear attention models, applicable without fine-tuning. For scenarios requiring deployment of long-context or multi-sequence inference on limited hardware, this memory saving has practical value.

Implications for Inference Deployment

PACE's significance lies not just in being a quantization algorithm, but in revealing an optimizable dimension in recurrent state storage: the trade-off between rounding frequency and precision. Traditional approaches assume every token requires precise state updates, but PACE proves that moderate deferred rounding combined with structure-aware encoding can substantially reduce memory requirements with almost no loss in accuracy.

As linear attention architectures become more prevalent in long-sequence tasks, such compression techniques for recurrent states will become increasingly important. In the future, we may see more quantization schemes that incorporate hardware characteristics, further pushing the boundaries of efficient inference.

One-line Summary

PACE achieves post-hoc 4-bit quantization of linear attention model states through periodic rounding and channel-aligned encoding, significantly reducing memory usage.

Paragraph Summary

PACE is a post-hoc quantization method for recurrent states in linear attention language models, aimed at solving the precision loss problem in low-precision state storage. Existing quantizers round immediately after each generated token, causing significant cumulative errors at 4-bit precision. PACE proposes two core designs: first, periodic rounding, which defers the rounding operation to execute every few tokens, reducing error accumulation; second, channel-aligned encoding, which exploits the property that large entries in gated delta states cluster in columns, sharing quantization scales within columns. Additionally, PACE uses mixed 8-bit precision for heads sensitive to rounding errors. Experiments show that mixed-precision PACE performs comparably to 8-bit baselines across multiple models while using only two-thirds of the memory; all-4-bit PACE outperforms Q-Mamba at the same memory usage. PACE can be applied without fine-tuning, providing a practical memory-saving solution for long-context multi-sequence inference.

#PACE#4-bit quantization#linear attention#recurrent state#memory-efficient inference#Mamba#RWKV
Share:

Related articles