大規模言語モデルで標準となっているプレノルム[PreNorm]を用いたレジidual接続[Residual connections with PreNorm]が抱える課題を解決することが抱える課題を解決するために,北京月之暗面科技[Beijing Moonshot AI Technology Co., Ltd.,中]のKimi Teamにより提案されたのが,先行する層の出力を学習可能な重みで選択的に集計する仕組みであるAttnRes[Attention Residuals].
従来のモデルは各層の出力を一律に加算するため,層が深くなるほど情報が希釈される問題があった.本手法はソフトマックス・アテンションを用いることにより,過去の表現を適応的に選択して統合することを可能にしている.計算コストやメモリ負荷を抑える工夫として,層をグループ化して処理するBlock AttnResも導入されており,実用的な実装が実現されている.
実験の結果,この手法は出力の安定化や勾配分布の改善に寄与し,幅広いタスクにおいてモデルの精度を向上させることが確認された.最終的に,480億パラメータを持つKimi Linearアーキテクチャに統合され,大規模な学習データを通じてその有効性が実証されている.
課題
LLMでは,プレノルム[PreNorm]を用いたレジidual接続が一般的であるが,2つの課題がある.
- 固定された重み: 従来のレジidual接続は,すべての層の出力を,固定された単位重み(1.0)で累積していく.
- 隠れ状態の無制御な増大: 層が深くなるにつれて隠れ状態のマグニチュードが無制御に増大し,結果として「各層の貢献度が希釈される[PreNorm dilution]という現象が起こる.
AttnResの仕組み
AttnResは,単に前の層の出力を足し合わせるのではなく,ソフトマックス・アテンション[softmax attention]を用いて,過去のすべての層の出力を集計する.
- 入力依存の動的な選択: 各層は,先行する層の表現をどの程度取り込むかを,入力に応じて学習された重みで動的に決定.
- 選択的集計: これにより,モデルは現在の処理に必要な情報を過去の特定の層から選択的に抽出できるようになる.
実装上の工夫:Block AttnRes
全層のアテンション計算は,大規模モデルの訓練においてメモリや通信のオーバーヘッドを増大させる可能性がある.これを解決するために, Block AttnRes という手法が導入されている.
- ブロック化: 層をいくつかのブロックに分割し,ブロック単位の表現に対してアテンションを行う.
- 効率化: キャッシュベースのパイプライン通信や,二段階の計算戦略と組み合わせることにより,従来のレジidual接続と置き換え可能な,オーバーヘッドを最小限に抑えた実装を実現.
期待される効果と実証結果
研究チームがKimi Linearアーキテクチャ[総パラメータ数48B / 活性パラメータ数3B]に AttnRes を統合して1.4兆トークンで事前学習を行った結果,以下の効果が確認されたとする.
- 勾配と出力の安定化: 深さ方向に対する出力マグニチュードと勾配の分布がより均一になり,プレノルム[PreNorm]による希釈問題が緩和.
- パフォーマンスの向上: スケーリング則[Scaling law]の実験により,モデルサイズに関わらず一貫した性能向上が確認され,評価されたすべてのタスクでダウンストリーム性能が向上.
- 内容依存の選択: アブレーション研究により, 内容に依存した深さ方向の選択[content-dependent depth-wise selection]が有効であることが検証されている
つまり,AttnResは「層の積み重ね」を「過去の知見からの動的な選択」へと進化させることにより, モデルの深層化に伴う効率低下を防ぎ,性能を引き出す技術と言うことができる.
用語解説
- プレノルム:Transformerにおいては,各層は典型的に残差接続[residual connection]と正規化を伴う.初期の設計では,サブレイヤ[自己注意やフィードフォワード]の出力に対して正規化を行うポストノルム[Post-Norm]が採用されていた.\[x \mapsto \mathrm{Norm}\bigl(x + \mathrm{SubLayer}(x)\bigr)\]しかし,この方法では,深いモデルでは学習が不安定になるという問題があった.これに対し,プレノルム[Pre-Norm]は各サブレイヤに入力する前に正規化を行う方式であり,\[x \mapsto x + \mathrm{SubLayer}(\mathrm{Norm}(x))\]という形を取る.この構造により,勾配が残差経路を通じてより直接的に伝播するため,深層化においても安定した学習が可能になる.
ここで関数解析的な意味でのプレノルムとのゆるやかな類似性が見出される
参考文献
- Kimi Team(2026). Attention Residuals
arXiv: https://arxiv.org/abs/2603.15031
2026-03-21.
