Summary (Overview)

• LinQuant is a framework for low-precision inference of linear recurrent neural networks, based on seven functional invariances that preserve full-precision computation while making operands more amenable to quantization.
• The invariances comprise four mixer-external changes of basis (key-side rotation, key-side rescaling, value-side rotation, value-side rescaling) and three mixer-internal absorptions (scalar-decay, diagonal-decay, and write-factor absorptions).
• The framework's applicability is characterized across common recurrence families: Linear Attention, Mamba-2, GLA, DeltaNet, Gated DeltaNet, and KDA, with formal proofs of validity conditions.
• LinQuant improves perplexity at matched latency over the state-of-the-art SSDi8 baseline on Mamba-2 (1.3B, 2.7B, 8B) and Nemotron-H (8B) models in production W4A8 settings.
• INT8 Triton kernel implementations demonstrate that query–key rotation adds no measurable prefill latency, and the framework generalizes beyond Mamba-2 to other modern recurrent architectures.

---

Introduction and Theoretical Foundation

Background and Motivation

Linear recurrences are emerging as practical alternatives to softmax attention in large language models (LLMs), appearing in standalone architectures (e.g., Mamba-2) and hybrid designs. They summarize past tokens in a fixed-size state, avoiding the growing key–value cache of softmax attention. During prefill, chunkwise parallel algorithms process blocks of tokens through dense matrix multiplications, making them naturally amenable to low-precision execution.

While low-precision (e.g., INT8) attention kernels are mature in the Transformer inference stack, low-precision kernels for linear recurrences remain largely unexplored. Prior work (e.g., MambaQuant, Quamba2) focused on quantizing projections and activations around the recurrence, while SSDi8 executes the recurrent computation itself in INT8—but only for Mamba-2 specifically.

Theoretical Foundation

The general linear recurrence is defined by the state update:

$$\mathbf{S}_t = \mathbf{S}_{t-1} \boldsymbol{\Lambda}_t \left(\mathbf{I} - \delta \beta_t \mathbf{k}_t \mathbf{k}_t^\top\right) + \beta_t \mathbf{v}_t \mathbf{k}_t^\top, \tag{1}$$

where $\beta_t > 0$ is the write gate, $\|\mathbf{A}_t\| \leq 1$ denotes the decay, $\delta \in \{0, 1\}$ distinguishes linear writes from delta-rule corrections, and $\mathbf{S}_t \in \mathbb{R}^{D_v \times D_k}$ is the matrix-valued recurrent state.

Key structural choices:

• Decay: Scalar decay $\Lambda_t = \alpha_t \mathbf{I}$ vs. diagonal decay $\Lambda_t = \mathrm{diag}(\alpha_t)$
• Write gating: Scales the new association $\mathbf{v}_t \mathbf{k}_t^\top$
• Delta correction: Removes the component of the existing state aligned with the current key before writing

Invariance-Based Quantization

Functional invariances preserve the full-precision operator while changing the numerics presented to a quantizer. Two broad classes are distinguished:

• Mixer-agnostic invariances: Act on the residual stream or individual layers (e.g., SmoothQuant, QuaRot-like rotations).
• Mixer-specific invariances: Depend on the sequence operator itself, further divided into:
  • Mixer-external invariances: Changes of basis applicable outside the sequence mixer
  • Mixer-internal invariances: Functions of the recurrence's algorithmic realization

The paper characterizes which attention-inspired invariances survive in recurrences, what new restrictions arise, and what invariances are unique to linear recurrences.

---

Methodology

Low-Precision Chunkwise Realization

The chunkwise computation is organized into four matrix contractions plus a state pass. For Mamba-2, the low-precision realization is:

$$\widehat{\mathbf{H}}_{[i]} = \mathcal{Q}(\mathbf{B}_{[i]} \mathbf{V}_{[i]})^\top \mathcal{Q}(\mathbf{\Gamma}_{[i]}^\rightarrow \odot \mathbf{K}_{[i]}) \quad \text{(state write)}$$ $$\widehat{\mathbf{S}}_{[i+1]} = \mathcal{Q}(\gamma_{[i]}^C \mathbf{S}_{[i]} + \widehat{\mathbf{H}}_{[i]}) \quad \text{(state pass)}$$ $$\widehat{\mathbf{O}}_{[i]}^{\text{inter}} = \mathbf{\Gamma}_{[i]}^\leftarrow \odot \big(\mathcal{Q}(\mathbf{Q}_{[i]}) \widehat{\mathbf{S}}_{[i]}^\top\big) \quad \text{(state read)}$$ $$\widehat{\mathbf{A}}_{[i]} = (\mathcal{Q}(\mathbf{Q}_{[i]}) \mathcal{Q}(\mathbf{K}_{[i]})^\top) \odot \mathbf{\Gamma}_{[i]} \quad \text{(local score)}$$ $$\widehat{\mathbf{O}}_{[i]}^{\text{intra}} = \widehat{\mathbf{A}}_{[i]} \mathcal{Q}(\mathbf{B}_{[i]} \mathbf{V}_{[i]}) \quad \text{(intra-chunk output)}$$

where $\mathcal{Q}$ denotes quantization with implicit scales.

The Seven Functional Invariances

Changes of Basis (Mixer-External):

1. Key-side rotation $(1/7)$:

   $$(\mathbf{q}_t, \mathbf{k}_t) \rightarrow (\mathbf{R}_{QK} \mathbf{q}_t, \mathbf{R}_{QK} \mathbf{k}_t), \quad \mathbf{R}_{QK} \in \mathcal{O}(D_k).$$

   Valid when $\mathbf{R}_{QK} \boldsymbol{\Lambda}_t = \boldsymbol{\Lambda}_t \mathbf{R}_{QK}$ (holds for no/scalar decay; for diagonal decay, restricted to local contractions).

2. Key-side rescaling $(2/7)$:

   $$(\mathbf{q}_t, \mathbf{k}_t) \rightarrow (\mathbf{D}_{QK} \mathbf{q}_t, \mathbf{D}_{QK}^{-1} \mathbf{k}_t), \quad \mathbf{D}_{QK} \succ 0 \text{ diagonal.}$$

   Valid for linear writes under no, scalar, or diagonal decay; commutes with diagonal decay but does not preserve delta-corrected writes. Parameterized by $d_n = (\max_t |k_{t,n}| / \max_t |q_{t,n}|)^\kappa$ with $\kappa \in [0,1]$; $\kappa = \frac{1}{2}$ balances query/key magnitudes.

3. Value-side rotation $(3/7)$:

   $$\mathbf{v}_t \rightarrow \mathbf{R}_{V} \mathbf{v}_t, \quad \mathbf{R}_{V} \in \mathcal{O}(D_v).$$

   Valid across all recurrences; introduces a choice of state storage basis (canonical vs. rotated), affecting quantization error.

4. Value-side rescaling $(4/7)$:

   $$\mathbf{v}_t \rightarrow \mathbf{D}_{V} \mathbf{v}_t, \quad \mathbf{D}_{V} \succ 0 \text{ diagonal.}$$

   State transforms as $\mathbf{S}_t \to \mathbf{D}_{V} \mathbf{S}_t$; parameterized by $d_p = \big((\prod_j \max_t |v_{t,j}|)^{1/D_v} / \max_t |v_{t,p}|\big)^\kappa$.

Absorptions (Mixer-Internal):

5. Scalar-decay absorption $(5/7)$: Splits tail-decay between value and key sides:

   $$\mathbf{H}_{[i]} = ((\mathbf{\Gamma}_{[i]}^\rightarrow)^\eta \odot (\mathbf{B}_{[i]} \mathbf{V}_{[i]}))^\top ((\mathbf{\Gamma}_{[i]}^\rightarrow)^{1-\eta} \odot \mathbf{K}_{[i]}), \quad \eta \in [0,1].$$

   $\eta = 0$: key-side decay; $\eta = 1$: value-side decay (used by SSDi8).

6. Diagonal-decay absorption $(6/7)$: Factors each token's decay into scalar magnitude and relative channel pattern:

   $$e_{[i]}^r := \max_n (\mathbf{\Gamma}_{[i]}^\rightarrow)_{rn}, \quad \mathbf{E}_{[i]} := \mathrm{diag}(e_{[i]}^1,\dots,e_{[i]}^C), \quad \overline{\mathbf{D}}_{[i]}^{\text{tail}} := \mathbf{E}_{[i]}^{-1} \mathbf{\Gamma}_{[i]}^\rightarrow.$$

   Then $\mathbf{H}_{[i]} = (\mathbf{E}_{[i]}^\eta \mathbf{B}_{[i]} \mathbf{V}_{[i]})^\top (\mathbf{E}_{[i]}^{1-\eta} \overline{\mathbf{D}}_{[i]}^{\text{tail}} \odot \mathbf{K}_{[i]})$.

7. Write-factor absorption $(7/7)$: Splits write gate between operands:

   $$\mathbf{H}_{[i]} = (\mathbf{B}_{[i]}^\xi \mathbf{V}_{[i]})^\top (\mathbf{\Gamma}_{[i]}^\rightarrow \odot \mathbf{B}_{[i]}^{1-\xi} \mathbf{K}_{[i]}), \quad \xi \in [0,1].$$

   For delta-rule writes, uses two parameters $(\sigma, \tau)$ to split the gate inside the Gram construction and in the complementary write factor.

Experimental Setup

• Controlled "mixer zoo": Linear Attention, Mamba-2, GLA, DeltaNet, Gated DeltaNet, KDA at 340M (15B tokens) and 1.3B (100B tokens) parameters on SlimPajama.
• INT6 quantization for controlled studies (to amplify effects), INT8 for production.
• Scale granularities: $\langle H, C\rangle$ (per-head/per-channel), $\langle T, C\rangle$ (per-tensor/per-channel), $\langle T, T\rangle$ (per-tensor/per-tensor).
• Production setup: W4A8 models with INT8 recurrent states, augmenting SSDi8.

---

Empirical Validation / Results

Controlled Study Findings

Key-axis transformations (Figure 2):

• $\mathbf{R}_{QK}$ benefit depends on recurrence structure: effective with no/scalar decay, little benefit with diagonal decay (restricted to local contractions).
• QK normalization reduces quantization error independently.
• Coarser key scales increase $\mathbf{R}_{QK}$ benefit; $\mathbf{D}_{QK}$ remains useful with per-head scales.
• $\kappa$ sweep shows best rescaling is generally shared between queries and keys ($\kappa = 1/2$).

Value-axis transformations (Figure 3):

• Both $\mathbf{R}_V$ and $\mathbf{D}_V$ most useful with per-tensor value scales, recovering large portions of excess NLL.
• Quantizing the state in a rotated basis incurs more error than canonical basis storage.
• $\mathbf{D}_V$ favored over $\mathbf{R}_V$ since it folds into quantization scales without extra cost.

Absorptions (Figure 4):

• Write-gate absorption recovers significant excess NLL under coarse scales ($\langle T,T\rangle$).
• Preferred write-gate placement depends on scale granularity: $\xi = 1$ (gate on values) preferred under $\langle T,C\rangle$; splitting favored in other regimes.
• For delta writes, non-standard placements $( \sigma \in \{0, \frac{1}{2}\}, \tau = 1 )$ outperform the standard WY placement of Flash Linear Attention.
• Scalar decay placement matters most: moving decay from values to keys increases excess NLL under per-tensor value scales.
• Diagonal-decay absorption provides little gain.

Combined transformations: Reduce quantization error more than individual transformations, recovering roughly twice as much excess NLL as the best single transformation at INT6 under $\langle T,T\rangle$ for 340M Mamba-2 and LA.

Production Results

Table 2: Production WikiText-2 perplexity (lower is better)

Kernel Performance Results

• Query–key rotation adds no measurable end-to-end prefill latency on evaluated Mamba-2 models.
• Fused rescaling leaves measured layer latency unchanged.
• Value-side decay absorption is faster for target Mamba-2 geometries, though the gap narrows (and can reverse) as value dimension $D_v$ increases.
• Canonical $\mathbf{R}_V$ storage adds overhead; neither value-side transformation is applied in the production configuration.

Applicability Summary (Table 1)

---

Theoretical and Practical Implications

Theoretical Implications

• New degrees of freedom: Linear recurrences exhibit invariance structures qualitatively different from softmax attention—decay, gating, and delta-rule updates impose constraints but also create unique degrees of freedom (e.g., decay/write absorption) with no attention analogue.
• Characterization of validity: The paper formally derives when each invariance holds across recurrence families, showing that diagonal decay breaks global rotations but admits rescalings, while delta-rule updates restrict key-side transformations but allow novel write-factor splits with two parameters.
• Scale-granularity interaction: The effectiveness of each invariance depends critically on quantization scale granularity, revealing that optimal invariant selection is coupled with the quantization scheme.

Practical Implications

• Production-ready: LinQuant improves perplexity on Mamba-2 (1.3B, 2.7B, 8B) and Nemotron-H (8B) over SSDi8 at matched latency, with no measurable overhead from query–key rotation.
• Kernel efficiency: The invariances with zero incremental cost (folding into quantization scales or register-level operations) are preferred in production; canonical storage under value rotation is rejected due to overhead.
• Broader design space: The framework generalizes beyond integer prefill to floating-point and other future low-precision implementations of recurrent models.
• Guidance for practitioners: The controlled study provides concrete guidance—e.g., use value-side decay absorption for speed, apply $\mathbf{D}_{QK}$ for Mamba-2-style recurrences, and prefer $\mathbf{D}_V$ over $\mathbf{R}_V$ to avoid storage-basis overhead.

---

Conclusion

LinQuant provides a systematic framework for low-precision linear recurrent inference based on seven functional invariances: four changes of basis (key/value rotations and rescalings) and three absorptions (scalar-decay, diagonal-decay, and write-factor redistributions). The paper characterizes applicability across common recurrence families and demonstrates through controlled experiments that the benefits of each invariance depend on:

• Recurrence structure (decay type, delta correction, write gating)
• Query–key normalization
• Quantization scale granularity
• State-storage basis

Key Findings

1. Changes of basis and write-gate allocation are particularly effective under coarse scales.
2. $\mathbf{R}_{QK}$ benefits no/scalar-decay recurrences and is restricted for diagonal decay.
3. $\mathbf{D}_{QK}$ commutes with diagonal decay and improves quantization error where applicable.
4. Value-side transformations are most useful with per-tensor scales; $\mathbf{D}_V$ is preferred over $\mathbf{R}_V$ due to storage-basis costs.
5. Standard floating-point placements of decay/write factors are not always optimal after quantization.

Production Impact

On production models, LinQuant improves perplexity over the SSDi8 baseline on all evaluated Mamba-2 and Nemotron-H models, with query–key rotation adding no measurable prefill latency—demonstrating practical utility of invariance-based quantization for recurrent models.

Future Directions

• Exploring alternative parameterizations of diagonal-decay absorption
• Extending the invariance framework to floating-point low-precision formats
• Investigating additional functional invariances unique to other recurrent architectures
• Studying the interaction of invariances with other quantization techniques (e.g., weight quantization, activation reconstruction)