Linear attention models need new methods to quantize recurrent state

When Dynamic Scaling Fails: A Scale-State Feedback Governs Recurrent-State Quantization in Linear-Attention and SSM LLMs

Machine Learning

Summary

Some large language models hold long conversations by remembering information in a fixed-size memory, which can be made smaller by compressing that memory. The authors find that simply using low-bit number storage causes errors to build up as the conversation continues. They discover this is because the scaling used for compression feeds back on the memory and causes instability. Attempting to fix this by using a reference for scaling helps with numeric issues but does not fully restore the model's performance. The paper shows that carefully managing the scale is more important than targeting rare outlying values for better compression.

What this means in practice

  • For machine learning engineers: Optimize memory compression for linear-attention LLM decoding to balance efficiency and output quality over long sequences.
  • For hardware designers: Design quantization hardware that incorporates scale-state feedback control to avoid numeric instability in LLM inference accelerators.

Authors

Yuhang Jiang

Abstract

Linear-attention and state-space LLMs carry long context in a fixed-size recurrent state, and quantizing that state is the natural way to cut its cost during long decode. Unlike a weight tensor, the state is rewritten at every step, so a quantizer applied at each step sits inside the recurrence and its errors compound. INT4 that leaves synthetic recall perfectly intact at the start destroys it a thousand steps later. The mechanism is a feedback between the quantizer scale and the state it is derived from. Yoking that scale to a full-precision reference removes the resulting norm explosion entirely, yet recovers only 1.8 of the 5.4 nats lost, so numerical instability and functional degradation are separable failures. The remainder is a resolution problem that locating outliers does not solve, since every outlier-aware scheme here loses to plain persistent scaling. Allocating bits across layers does help, but under article-disjoint splits a learned allocation collapses onto the trivial rule of protecting the first few layers. Run out to 32K, group-wise INT4 is worth no more than discarding the state. Contain the scale rather than locate outliers, and report the decode horizon at which a bit-width was certified.
View PDFOpen arXiv