大型语言模型的门控注意力:非线性、稀疏性与无注意力陷阱
计算与语言
2025-05-13 v1
摘要
门控机制在从早期模型如LSTM和高way网络到最新的状态空间模型、线性注意力以及softmax注意力中都得到了广泛应用。然而,现有文献很少考察门控的具体影响。本文我们进行了全面的实验,以系统性地探讨门控增强的softmax注意力变体。具体而言,我们在30个15B规模的Mixture-of-Experts(MoE)模型和1.7B规模的稠密模型上进行了全面比较,这些模型在3.5万亿token的数据集上进行训练。我们的核心发现是,一种简单修改——在缩放点积注意力(Scaled Dot-Product Attention, SDPA)之后施加特定头的Sigmoid门控——始终能提高性能。这种修改还增强了训练稳定性,容许更大的学习率,并提升了扩展性能。通过比较各种门控位置和计算变体,我们将这种有效性归因于两个关键因素:(1)在softmax注意力中的低秩映射上引入非线性,以及(2)应用查询依赖的稀疏门控得分来调制SDPA输出。值得注意的是,我们发现这种稀疏门控机制可缓解“注意力陷阱”(attention sink),并增强了长程上下文外推性能,我们还发布了相关代码和模型以促进未来研究。
引用
@article{arxiv.2505.06708,
title = {Gated Attention for Large Language Models: Non-linearity, Sparsity, and Attention-Sink-Free},
author = {Zihan Qiu and Zekun Wang and Bo Zheng and Zeyu Huang and Kaiyue Wen and Songlin Yang and Rui Men and Le Yu and Fei Huang and Suozhi Huang and Dayiheng Liu and Jingren Zhou and Junyang Lin},
journal= {arXiv preprint arXiv:2505.06708},
year = {2025}
}