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

TL;DR

This paper investigates the low-rank structure of linear attention models, provides a theoretical analysis of its impact, and proposes a structured pruning method to significantly reduce state size with minimal performance loss.

Contribution

It offers a theoretical understanding of low-rank effects in linear attention and introduces a novel pruning approach based on QR decomposition for efficient state reduction.

Findings

Low effective rank amplifies query noise in linear attention.

Pruning key and query matrices can reduce state size by 50%.

Minimal performance degradation observed after pruning.

Abstract

Linear attention offers a computationally efficient yet expressive alternative to softmax attention. However, recent empirical results indicate that the hidden state of trained linear attention models often exhibits a low-rank structure, suggesting that these models underexploit their capacity in practice. To illuminate this phenomenon, 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. In addition to these theoretical insights, we conjecture that the low-rank states can be substantially reduced post-training with only minimal performance degradation, yielding faster and more memory-efficient models. To this end, we propose a novel hardware-aware approach that structurally prunes key and query matrices, reducing the state size while retaining compatibility with existing CUDA kernels. We adapt several existing pruning strategies to fit our framework and, building on our theoretical analysis, propose a novel structured pruning method based on a rank-revealing QR decomposition. Our empirical results, evaluated across models of varying sizes and on various downstream tasks, demonstrate the effectiveness of our state reduction framework. We highlight that our framework enables the removal of 50% of the query and key channels at only a marginal increase in perplexity. The code for this project can be found at https://github.com/camail-official/LinearAttentionPruning.