State Rank Dynamics in Linear Attention LLMs

TL;DR

The paper investigates runtime state dynamics in Linear Attention LLMs and uncovers State Rank Stratification, where attention heads bifurcate into persistent low-rank and high-rank regimes that are temporally and data-invariant. It provides a theoretical framework based on Rank Saturation and Norm Accumulation to explain stable rank and norm distribution over long sequences, and demonstrates that low-rank heads are essential for reasoning while high-rank heads often exhibit redundancy. Building on this, the authors introduce Joint Rank-Norm Pruning (JRNP), a training-free method that prunes saturated high-rank heads to reduce KV-cache overhead by about 38.9% with negligible accuracy loss on several tasks, while pruning low-rank heads severely degrades retrieval and long-context capabilities. The work advances understanding of spectral dynamics in recurrent linear attention and suggests architectural biases favoring low-rank subspaces to balance efficiency and retrieval performance.

Abstract

Linear Attention Large Language Models (LLMs) offer a compelling recurrent formulation that compresses context into a fixed-size state matrix, enabling constant-time inference. However, the internal dynamics of this compressed state remain largely opaque. In this work, we present a comprehensive study on the runtime state dynamics of state-of-the-art Linear Attention models. We uncover a fundamental phenomenon termed State Rank Stratification, characterized by a distinct spectral bifurcation among linear attention heads: while one group maintains an effective rank oscillating near zero, the other exhibits rapid growth that converges to an upper bound. Extensive experiments across diverse inference contexts reveal that these dynamics remain strikingly consistent, indicating that the identity of a head,whether low-rank or high-rank,is an intrinsic structural property acquired during pre-training, rather than a transient state dependent on the input data. Furthermore, our diagnostic probes reveal a surprising functional divergence: low-rank heads are indispensable for model reasoning, whereas high-rank heads exhibit significant redundancy. Leveraging this insight, we propose Joint Rank-Norm Pruning, a zero-shot strategy that achieves a 38.9\% reduction in KV-cache overhead while largely maintaining model accuracy.

Figures (5)

• Figure 1: State Rank Dynamics across 32 attention heads in a representative layer. The dashed black line indicates the theoretical rank upper bound. The plot vividly illustrates State Rank Stratification, revealing a clear bifurcation between distinct high-rank and low-rank regimes.
• Figure 2: Visualization of the Effective Rank for all 48 layers in Qwen3-Next. Darker colors indicate a higher effective rank. The visualization reveals that the number of high-rank and low-rank heads is approximately equal within the linear attention layers, with both types appearing to be randomly distributed across all 48 layers. LA: Linear Attention layer; SA: Softmax Attention layer.
• Figure 3: Strong Temporal Consistency of State Norm and Rank. The scatter plots illustrate the stability of Nuclear Norm (Top Row, a-d) and Effective Rank (Bottom Row, e-h) across different time steps ($t$ vs. $t+\Delta t$) in representative layers. Each point represents an attention head. The red dashed line denotes the identity mapping ($y=x$). The extremely high cosine similarity scores ($\text{Cos\_Similarity} > 0.97$) and the tight alignment of points along the diagonal confirm the strong temporal stability and rank-order preservation of the state properties: heads with relatively higher rank or norm values maintain their magnitude relative to others throughout the inference process.
• Figure 4: Representative visualization of rank dynamics under Attack Scenario I.
• Figure 5: Representative visualization of rank dynamics under Attack Scenario II.

Theorems & Definitions (25)

• Theorem 3.1: State Rank Upper Bound
• Theorem 4.2: Recursive Stability of Rank Vector Cosine Similarity
• Theorem 4.4: Directional Stability via Relative Step Size Decay
• Remark 4.5: Asymptotic Stability
• Theorem 1.1: Rank Upper Bound for Linear Attention
• proof
• Lemma 2.1: Structural Rank Bound via Key/Value Subspaces
• proof
• Lemma 2.2: Step-wise Rank Growth Bound
• proof