From Attention to Activation: Unravelling the Enigmas of Large Language Models

📄 arXiv: 2410.17174v1 📥 PDF

作者: Prannay Kaul, Chengcheng Ma, Ismail Elezi, Jiankang Deng

分类: cs.CL

发布日期: 2024-10-22

备注: 10 pages

💡 一句话要点

针对LLM中Attention集中和激活异常问题，提出Softmax-1和OrthoAdam优化器

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大型语言模型 注意力机制 异常激活 优化器 模型量化

📋 核心要点

  1. 大型语言模型（LLM）中，注意力机制过度集中于首个token，导致信息利用不充分，影响模型性能。
  2. 论文提出Softmax-1来缓解注意力集中问题，并设计OrthoAdam优化器来抑制异常激活，提升模型稳定性。
  3. 实验表明，新方法显著降低了注意力集中度和激活峰度，并提升了量化后模型的性能，改善了模型鲁棒性。

📝 摘要（中文）

本文研究了自回归Transformer中的两个异常现象：(1)注意力头中首个token的支配地位；(2)隐藏状态中出现的大幅度异常激活。研究发现，诸如Llama等流行的LLM中，98%的注意力头都最大程度地关注第一个token，这种行为归因于softmax函数。为了缓解这个问题，我们提出将softmax函数重新定义为softmax-1。此外，我们确定自适应优化器（如Adam）是导致大幅度异常激活的主要因素，并引入OrthoAdam，一种利用正交矩阵变换梯度的新型优化器来解决此问题。最后，我们的方法不仅可以防止这些现象的发生，而且还使Transformer能够在通过基本算法进行量化时保持其性能，这是标准方法无法做到的。总而言之，我们的方法将第一个token上的注意力比例从65%降低到3.3%，将隐藏状态中的激活峰度从1657降低到3.1，并将4比特权重量化下的困惑度惩罚从3565降低到0.3。

🔬 方法详解

问题定义：大型语言模型（LLM）在自回归Transformer结构中存在两个主要问题：一是注意力头过度关注第一个token，导致后续token的信息利用不足；二是隐藏层中出现幅度过大的异常激活值，影响模型的稳定性和泛化能力。现有方法难以有效解决这些问题，尤其是在模型量化后，性能下降更为明显。

核心思路：论文的核心思路是分别针对注意力集中和异常激活问题，提出不同的解决方案。对于注意力集中问题，通过修改softmax函数，降低首个token的权重；对于异常激活问题，通过设计新的优化器，约束梯度的更新方向，从而抑制异常激活的产生。

技术框架：论文主要涉及两个方面的改进：一是注意力机制的改进，通过将softmax函数替换为softmax-1，减少对第一个token的过度关注。二是优化器的改进，提出了OrthoAdam优化器，该优化器利用正交矩阵来变换梯度，从而限制梯度的更新方向，防止出现过大的激活值。整体框架是在标准的Transformer模型上进行改进，无需修改模型的整体结构。

关键创新：论文的关键创新在于提出了softmax-1函数和OrthoAdam优化器。softmax-1通过在softmax结果中减去1，降低了首个token的权重，从而缓解了注意力集中问题。OrthoAdam优化器通过正交矩阵变换梯度，限制了梯度的更新方向，从而抑制了异常激活的产生。这两种方法分别针对LLM中的两个关键问题，提供了有效的解决方案。

关键设计：softmax-1函数的具体实现是在标准的softmax函数计算结果中减去1，并重新归一化。OrthoAdam优化器的关键在于正交矩阵的生成和应用，论文中可能涉及正交矩阵的生成方法、更新策略以及与Adam优化器的结合方式等技术细节。损失函数和网络结构保持不变，主要是在注意力机制和优化器层面进行改进。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，提出的Softmax-1和OrthoAdam方法能够显著降低注意力集中度和激活峰度。具体来说，第一个token上的注意力比例从65%降低到3.3%，隐藏状态中的激活峰度从1657降低到3.1。更重要的是，在4比特权重量化下，困惑度惩罚从3565降低到0.3，表明该方法能够有效提升量化后模型的性能。

🎯 应用场景

该研究成果可应用于各种基于Transformer的大型语言模型，尤其是在资源受限的场景下，例如移动设备或边缘计算平台。通过降低注意力集中度和抑制异常激活，可以提高模型的稳定性和量化后的性能，从而降低部署成本，并提升用户体验。此外，该研究也有助于更好地理解LLM的内部机制，为未来的模型设计提供指导。

📄 摘要（原文）

We study two strange phenomena in auto-regressive Transformers: (1) the dominance of the first token in attention heads; (2) the occurrence of large outlier activations in the hidden states. We find that popular large language models, such as Llama attend maximally to the first token in 98% of attention heads, a behaviour we attribute to the softmax function. To mitigate this issue, we propose a reformulation of softmax to softmax-1. Furthermore, we identify adaptive optimisers, e.g. Adam, as the primary contributor to the large outlier activations and introduce OrthoAdam, a novel optimiser that utilises orthogonal matrices to transform gradients, to address this issue. Finally, not only do our methods prevent these phenomena from occurring, but additionally, they enable Transformers to sustain their performance when quantised using basic algorithms, something that standard methods are unable to do. In summary, our methods reduce the attention proportion on the first token from 65% to 3.3%, the activation kurtosis in the hidden states from 1657 to 3.1, and perplexity penalty under 4-bit weight quantisation from 3565 to 0.3.