The Key to State Reduction in Linear Attention: A Rank-based Perspective

Paper: 2602.04852 Authors: Philipp Nazari, T. Konstantin Rusch Affiliations: Max Planck Institute for Intelligent Systems, ETH Zürich, ELLIS Institute Tübingen, Liquid AI Code: https://github.com/camail-official/LinearAttentionPruning Categories: cs.LG

Abstract

Linear attention offers a computationally efficient yet expressive alternative to softmax attention. However, recent empirical results indicate that the state of trained linear attention models often exhibits a low-rank structure, suggesting that these models underexploit their capacity in practice. We provide a theoretical analysis of the role of rank in linear attention, revealing that low effective rank can affect retrieval error by amplifying query noise.

Key Contributions

• Theoretical insights: Analysis of how rank utilization affects retrieval error in linear attention
• State reduction framework: Novel post-training approach to reduce state size with minimal performance degradation
• Structured pruning method: Based on rank-revealing QR decomposition
• Empirical validation: 50% of query/key channels can be removed with only marginal perplexity increase

Background: Linear Attention

Softmax Attention

$O = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right) V$

• Quadratic time and memory complexity $O(T^2)$
• Problematic for long sequences

Linear Attention

Removes softmax and exploits associativity:

$QK^T V = Q(K^T V)$

Sequential form with hidden state:

$S_t = S_{t-1} + v_t k_t^T$ $o_t = S_t q_t$

Where $S_t$ is a matrix of size $(d_v, d_k)$, the associative memory.

Benefits:

• Linear time complexity
• Constant memory inference
• Chunkwise parallel training

The Low-Rank Problem

Empirical observation: Trained linear attention models exhibit low-rank structure in their hidden states.

Why Does This Matter?

The rank of the associative memory is bounded:

$\text{rank}(S_t) \leq \min(\text{rank}(K_t), \text{rank}(V_t)) \leq t$

But the effective rank (stable rank) is what matters:

$\text{sr}(A) = \frac{\|A\|_F^2}{\|A\|_2^2}$

This measures how evenly energy is distributed across singular values.

Rank Utilization

$u(S) = \frac{\text{sr}(S)}{\min(d_v, d_k)}$

• Low utilization ($u \ll 1$): Rank collapse, energy concentrated in few components, redundant information
• High utilization ($u \approx 1$): Isotropic memory, energy evenly distributed

Impact on Retrieval Error

For a noisy query $\tilde{q} = q^* + n$:

• Low rank utilization amplifies query noise
• High rank utilization provides more robust retrieval

Key insight: The conditioning of keys $\kappa(K_t)$ is tightly connected to the effective rank of the hidden state.

State Reduction Framework

Since models underutilize their capacity, we can prune the state post-training.

Key Observation

Linear attention models are invariant under orthogonal transformations applied jointly to queries and keys:

$(QT)(KT)^T = QK^T$

This allows selecting subsets of columns that contribute most to performance.

Structured vs Unstructured Pruning

Type	Effect	State Size
Unstructured	Sparse weights, same dimensions	Unchanged
Structured	Eliminates basis vectors	Reduced to d'_k<d_k

Structured pruning strictly decreases FLOPs for the recurrence.

Pruning Methods

1. Magnitude-based: Remove columns with smallest norms
2. Gradient-saliency: Remove columns with smallest gradient contributions
3. Rank-revealing QR (proposed): Select columns that maximize rank utilization of remaining memory

Results

Across models of varying sizes and downstream tasks:

• 50% of query/key channels can be removed with only marginal perplexity increase
• Works even before recovery fine-tuning
• Complementary to other compression methods (SliceGPT, SpinQuant, QuaRot)

Comparison

Method	Operates on	Complementary
SliceGPT	Embedding space	Yes
SpinQuant	Embedding space	Yes
This work	State space	Yes

Takeaways

1. Linear attention underutilizes capacity: Trained models exhibit low-rank hidden states
2. Low rank hurts retrieval: Amplifies query noise and reduces robustness
3. Post-training pruning works: 50% reduction with minimal performance loss
4. Rank-revealing QR is effective: Principled selection of important dimensions
5. Faster and more memory-efficient: Reduced state size directly improves inference