用线性注意力机制让大模型长文本推理更省显存
InAttention: Linear Context Scaling for Transformers
- token只关注初始状态，将注意力复杂度从平方降至线性
- 推理显存降低90%以上，可在消费级显卡跑超长序列
- 微调后长文本效果提升，训练成本低适合实用部署
Transformer模型的显存需求随上下文长度呈平方增长，源于自注意力机制。本文修改解码器仅用的Transformer结构，以InAttention替代自注意力机制，在推理阶段实现与上下文长度线性相关的计算开销，使每个词仅关注初始状态。基准测试显示，InAttention显著降低推理时的显存占用，使在消费级GPU上处理长序列成为可能。我们验证了微调能高效扩展上下文长度，在不增加高训练成本的前提下提升长序列任务性能。InAttention为Transformer模型中的长程依赖问题提供了一种可扩展的解决方案，推动了进一步优化。
原文摘要 · Abstract (English)
VRAM requirements for transformer models scale quadratically with context length due to the self-attention mechanism. In this paper we modify the decoder-only transformer, replacing self-attention with InAttention, which scales linearly with context length during inference by having tokens attend only to initial states. Benchmarking shows that InAttention significantly reduces VRAM usage during inference, enabling handling of long sequences on consumer GPUs. We corroborate that fine-tuning extends context length efficiently, improving performance on long sequences without high training costs. InAttention offers a scalable solution for long-range dependencies in transformer models, paving the way for further optimization.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。