通过秩分析揭示线性注意力状态冗余,实现高效压缩且几乎不损失性能。
The Key to State Reduction in Linear Attention: A Rank-based Perspective
- 从秩的角度解析线性注意力的低效,指出查询噪声被放大导致误差
- 训练后可删减50%的键值通道,困惑度仅轻微上升
- 提出基于秩揭示QR分解的结构化剪枝,适配现有加速库
线性注意力在计算效率和表达能力上优于Softmax注意力,但近期实证发现其隐藏状态常呈现低秩结构,表明模型实际未充分利用容量。本文从理论上分析了秩的作用,揭示低有效秩会放大查询噪声,从而影响检索误差。基于此,我们提出训练后可显著压缩状态空间,仅带来微小性能损失。为此,设计了一种面向硬件的结构化剪枝方法,对键和查询矩阵进行剪枝,减少状态大小的同时保持与现有CUDA内核兼容。将多种已有剪枝策略适配至该框架,并结合理论分析提出一种基于秩揭示QR分解的新剪枝方法。在不同规模模型及多个下游任务上的实验验证了该框架的有效性。结果表明,可移除50%的查询与键通道,仅导致困惑度轻微增加。项目代码见:https://github.com/camail-official/LinearAttentionPruning。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。