InfLLM-V2: Dense-Sparse Switchable Attention for Seamless Short-to-Long Adaptation
Weilin Zhao, Zihan Zhou, Zhou Su, Chaojun Xiao, Yuxuan Li, Yanghao Li, Yudi Zhang, Weilun Zhao, Zhen Li, Yuxiang Huang, Ao Sun, Xu Han, Zhiyuan Liu
长序列处理是现代大语言模型的关键能力。然而，标准Transformer架构中的自注意力机制在处理长序列时面临严重的计算和内存瓶颈。虽然可训练的稀疏注意力方法提供了一个有前景的解决方案，但现有方法如NSA引入了过多的额外参数，并破坏了传统的短序列预训练、长序列微调工作流程，导致收敛缓慢且难以加速。为了克服这些限制，我们引入了稠密-稀疏可切换注意力框架，称为InfLLM-V2。InfLLM-V2是一种可训练的稀疏注意力，能够无缝地将模型从短序列适应到长序列。具体而言，InfLLM-V2通过无参数架构修改重用稠密注意力参数，保持短序列和长序列处理之间的一致性。此外，InfLLM-V2通过使用稠密注意力处理短输入并平滑过渡到稀疏注意力处理长序列，确保在所有序列长度上的计算效率。为了实现实际加速，我们进一步引入了InfLLM-V2的高效实现，显著降低了计算开销。我们在长上下文理解和思维链推理上的实验表明，InfLLM-V2比稠密注意力快4倍，同时保持98.1%的性能。
Long-sequence processing is a critical capability for modern large language models. However, the self-attention mechanism in the standard Transformer architecture faces severe computational and memory bottlenecks when processing long sequences. While trainable sparse attention methods offer a promising solution, existing approaches such as NSA introduce excessive extra parameters and disrupt the conventional pretrain-on-short, finetune-on-long workflow, resulting in slow convergence and difficul...