用线性RNN替换部分注意力头，提升长文本生成效率
RecurFormer: Not All Transformer Heads Need Self-Attention
- 将部分注意力头换成Mamba结构的线性RNN
- 推理速度提升，缓存占用降低但生成质量不变
- 适合需要处理长序列的LLM部署场景
基于Transformer的大语言模型在建模复杂语言模式方面表现优异，但在推理阶段面对长输入时面临显著计算开销，主要源于注意力机制带来的内存负担。我们观察到某些注意力头的注意力权重集中在查询词附近的词元上，表现出对近期内容的敏感性，即关注局部与短距离依赖。基于此，我们提出RecurFormer，用线性循环神经网络（特别是Mamba架构）替代这些注意力头，从而在不丢弃词元的前提下减少缓存大小，保持生成质量。剩余的注意力头仍可建模长距离依赖，且能复用预训练Transformer模型权重并持续训练。实验表明，RecurFormer在性能上匹配原模型，同时显著提升推理效率。该方法为Transformer类大模型的推理计算挑战提供了实用解决方案，尤其适用于长输入任务。
原文摘要 · Abstract (English)
Transformer-based large language models (LLMs) excel in modeling complex language patterns but face significant computational costs during inference, especially with long inputs due to the attention mechanism's memory overhead. We observe that certain attention heads exhibit a distribution where the attention weights concentrate on tokens near the query token, termed as recency aware, which focuses on local and short-range dependencies. Leveraging this insight, we propose RecurFormer, a novel architecture that replaces these attention heads with linear recurrent neural networks (RNNs), specifically the Mamba architecture. This replacement reduces the cache size without evicting tokens, thus maintaining generation quality. RecurFormer retains the ability to model long-range dependencies through the remaining attention heads and allows for reusing pre-trained Transformer-based LLMs weights with continual training. Experiments demonstrate that RecurFormer matches the original model's performance while significantly enhancing inference efficiency. Our approach provides a practical solution to the computational challenges of Transformer-based LLMs inference, making it highly attractive for tasks involving long inputs.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。