Gated Attention for Large Language Models: Non-linearity, Sparsity, and Attention-Sink-Free 深度解析

研究背景與動機

隨著大型語言模型（Large Language Models, LLMs）規模持續擴增，基於注意力機制（Attention Mechanism）的架構已成為主流。當前主流的 Transformer 架構中，Softmax 點積注意力（Scaled Dot-Product Attention, SDPA）是一個關鍵組件，它通過計算查詢（Query）與鍵（Key）之間的相似度，在輸入序列中分配權重以捕捉長距依賴。然而，傳統的 Softmax 注意力存在一定的瓶頸，例如線性表現能力受限、過於密集的計算導致擴展性問題，以及「注意力陷阱」（Attention Sink）問題，後者會使注意力集中在少數位置，限制模型理解長文本上下文的能力。

另一方面，門控機制（Gating Mechanism）自從 LSTM、Highway Networks 等早期模型便被廣泛應用，以改善信息流動和非線性變換能力。近年來，進階應用如線性注意力（Linear Attention）和狀態空間模型（State Space Models）也引入多種門控策略，然而在大型 Transformer 及其 Softmax 注意力層中，針對門控機制的系統性研究仍較為缺乏，且其實際作用尚未明確。

基於此，本論文提出一套系統化的實驗框架，深入探討門控如何在 Softmax 注意力中發揮作用，藉由大規模模型與龐大訓練資料來驗證設計，期望對提升模型性能與訓練穩定性帶來突破。

核心方法與創新

本論文的核心創新是引入「head-specific sigmoid gate」（頭特定的 Sigmoid 門控）於標準的 SDPA 後段，簡單但顯著地改善了模型表現。具體流程是：在原有計算出來的注意力輸出結果上，進一步乘以一組經 Sigmoid 函數非線性映射的門控向量，此門控向量依賴於輸入 query，且對每個注意力頭（Attention Head）獨立學習。

這看似簡單的結構有兩個關鍵能力：

1. 引入非線性：Softmax 本身是一種將相似度映射到概率分布的變換，本論文指出透過加入 Sigmoid 門控，創造出在低秩映射（low-rank mapping）上的多層次非線性，提升模型表達複雜度和靈活度。
2. 稀疏門控（Sparse Gating）：門控向量本質上是對注意力輸出進行按元素調制，且是對查詢依賴的動態調整，導致部分通道被部分屏蔽（啟用程度降低），自然形成稀疏性，這不只降低計算冗餘，也大幅緩解了「注意力陷阱」問題。

此外，論文廣泛比較了門控位置（例如在 Key、Value 或 Attention Output）、不同類型門控函數（如 ReLU、Tanh）及其計算成本，發現 head-specific Sigmoid 門控在 SDPA 後方的配置是最佳選擇，既保有操作便利性，也兼顧計算效率。

主要實驗結果

本研究基於超大規模訓練集——3.5 兆字元（trillion tokens），對比超過30種門控方法，涵蓋15B 參數的 Mixture-of-Experts（MoE）模型及1.7B 參數的致密密集模型（Dense Model），並系統性分析性能與訓練穩定性。

• 性能提升顯著：在多項下游任務與語言理解基準上，添加門控的模型普遍優於無門控基線，平均提升數個百分點的準確率與損失下降。
• 訓練穩定性增強：引入門控使模型能承受更大學習率，縮短收斂時間，並大幅減少梯度爆炸與消失問題，這對大規模訓練尤為關鍵。
• 長文本表現改善：傳統注意力隨文本長度增加表現退化明顯，但門控機制有效減少「注意力陷阱」現象，穩定模型在超長上下文中的推理能力。
• 測試不同門控變體：從門控位置、門控函數到稀疏剪枝，一致指向 head-specific sigmoid gate 後置配置為最適，且稀疏門控有助實現計算節省，誘發更有效的特徵選擇。

研究團隊同時釋出相對應的模型與程式碼，促進未來社群基於此架構繼續拓展和優化。

對 AI 領域的深遠影響

本論文突破性地展示了門控機制在大型語言模型中不僅是參數引入的輔助，更是提升性能與模型穩定性的關鍵元件。這不僅提供了 Transformer 結構優化的全新視角，也對未來大規模模型設計有深遠啟發：

1. 強化非線性與動態調節能力：透過門控介入注意力內部，模型能更精細地根據不同上下文動態控制信息流，突破單純線性投影的限制，強調了靈活性與泛化能力。
2. 稀疏性與計算效率平衡：可動態稀疏門控不僅避免了「注意力陷阱」這種表現衰退現象，也促使未來模型更易於實現稀疏計算和高效推理，是向更大規模且節能 AI 方向演進的關鍵。
3. 模型穩定性提升推動深度學習大規模化：穩定且可容忍更高學習率的訓練方法，加速模型訓練，降低資源成本，對產業應用和研究都具有激勵作用。
4. 開源促進社群合作：本論文附帶公開的代碼與模型為社群提供基準，推動該領域後續在門控機制與注意力機制創新上持續前進，可能催生更多突破。

總結來說，「Gated Attention for Large Language Models」 不僅用實證驗證了門控機制對於大型 Transformer 的重大正向影響，更為理論與應用層面提供新框架，為未來大型語言模型更全面、高效且穩定的設計奠定重要基石。

---

論文資訊
📄 Gated Attention for Large Language Models: Non-linearity, Sparsity, and Attention-Sink-Free
👥 Qiu, Wang, Zheng, Huang, Wen, Yang, Men, Yu, Huang, Huang, Liu, Zhou, Lin
🏆 NeurIPS 2025 · Best Paper
🔗 arxiv.org/abs/2505.06708