揭示大模型中极端标记现象的内在机制,提出缓解方法。
Active-Dormant Attention Heads: Mechanistically Demystifying Extreme-Token Phenomena in LLMs
- 发现注意力头存在'活跃-休眠'机制,导致特定输入下某些标记成为关注焦点。
- 极端标记的注意力权重高、值状态小、残差范数大,且在简单模型中已出现。
- 建议用ReLU替代Softmax、SGD替代Adam,可有效抑制该现象,适合模型优化者阅读。
实践者长期观察到基于Transformer的大语言模型(LLMs)存在三个谜题现象:注意力汇聚点、值状态耗竭和残差状态峰值,统称为极端标记现象。这些现象表现为某些‘汇聚标记’接收远高于其他标记的注意力权重,值状态显著更小,残差状态范数却更大。这些极端标记引发推理、量化和可解释性等多方面挑战。本文揭示其机制:在仅1至3层的简单架构上训练的二元回溯(Bigram-Backcopy, BB)任务中,即可观察到此类现象。我们发现注意力头在特定输入域呈现‘活跃-休眠’模式,且训练动态显示其由相互强化机制驱动。基于此,提出预训练阶段缓解策略,如以ReLU替代Softmax、以SGD替代Adam。进一步扩展分析至预训练模型(包括Llama和OLMo),发现多数注意力头同样表现出类似机制,且相互强化仍主导极端标记的形成。BB任务预测的静态与动态特性与真实大模型观测高度一致。
原文摘要 · Abstract (English)
Practitioners have consistently observed three puzzling phenomena in transformer-based large language models (LLMs): attention sinks, value-state drains, and residual-state peaks, collectively referred to as extreme-token phenomena. These phenomena are characterized by certain so-called "sink tokens" receiving disproportionately high attention weights, exhibiting significantly smaller value states, and having much larger residual-state norms than those of other tokens. These extreme tokens give rise to various challenges in LLM inference, quantization, and interpretability. We elucidate the mechanisms behind extreme-token phenomena. First, we show that these phenomena arise in very simple architectures -- transformers with one to three layers -- trained on a toy model, the Bigram-Backcopy (BB) task. In this setting, we identify an active-dormant mechanism, where attention heads become sinks for specific input domains while remaining non-sinks for others. Our theoretical analysis of the training dynamics reveals that these phenomena are driven by a mutual reinforcement mechanism. Building on these insights, we propose strategies to mitigate extreme-token phenomena during pretraining, including replacing softmax with ReLU and Adam with SGD. Next, we extend our analysis to pretrained LLMs, including Llama and OLMo, showing that many attention heads exhibit a similar active-dormant mechanism as in the BB task, and that the mutual reinforcement mechanism also governs the emergence of extreme-token phenomena during LLM pretraining. Our results reveal that many of the static and dynamic properties of extreme-token phenomena predicted by the BB task align with observations in pretrained LLMs.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。