主动-静默注意力头：解构 LLM 中极端 token 现象的机制
机器学习
2024-11-08 v2
摘要
实践者始终观察到基于转换器的大语言模型（LLM）中存在三种令人费解的现象：注意力沉淀、值状态消耗和残差状态峰值，统称为极端 token 现象。这些现象由称为“沉淀 token”的特定 token 接收不成比例的注意力权重，显示显著较小的值状态，且其残差状态范数远大于其他 token。这些极端 token 引发 LLM 推理、量化和可解释性方面的各种挑战。我们阐明了极端 token 现象背后的机制。首先，我们展示这些现象在极简单的架构——仅有一至三个层的转换器——训练于 Bigram-Backcopy (BB) 任务时就会出现。在此 setting 下，我们识别出一种主动-静默机制，即注意力头对特定输入域成为沉淀点，而对其他输入域保持非沉淀状态。我们的训练动力学理论分析揭示，这些现象是由一种相互强化机制驱动的。基于这些洞见，我们提出了在预训练期间缓解极端 token 现象的策略，包括用 ReLU 替代 softmax、用 SGD 替代 Adam。随后，我们将分析扩展到预训练的 LLM，包括 Llama 和 OLMo，发现许多注意力头在 BB 任务中表现出类似的主动-静默机制，并且相互强化机制也支配着极端 token 现象在 LLM 预训练中的产生。我们的结果表明，BB 任务预测的许多静态和动态极端 token 现象属性与预训练 LLM 中的观察结果一致。
引用
@article{arxiv.2410.13835,
title = {Active-Dormant Attention Heads: Mechanistically Demystifying Extreme-Token Phenomena in LLMs},
author = {Tianyu Guo and Druv Pai and Yu Bai and Jiantao Jiao and Michael I. Jordan and Song Mei},
journal= {arXiv preprint arXiv:2410.13835},
year = {2024}
}