技術 約8分で読めます

MoonshotAI(Kimi)がTransformerの残差接続をAttentionで置き換えるAttnResを提案、1.25倍の計算効率

MoonshotAI(Kimi)チームが2026年3月16日、Transformerの残差接続を深さ方向のAttentionに置き換える手法 Attention Residuals(AttnRes)を arXiv に公開した(arXiv:2603.15031)。実装はGitHubに公開されている。

そもそも残差接続とは

ニューラルネットワークは層を重ねて深くするほど表現力が上がるが、深くしすぎると勾配消失(学習信号が途中の層で消えてしまう現象)が起きて逆に性能が落ちる。2015年にResNet(残差ネットワーク)で提案されたのが残差接続(skip connection)で、各層の処理結果に入力をそのまま足し算するというシンプルな仕組みだ。

graph TD
    IN[層の入力] --> ATT[Attention<br/>or FFN]
    ATT --> ADD((+))
    IN -->|残差接続| ADD
    ADD --> OUT[出力]

数式で書くと h_l = h_{l-1} + f(h_{l-1})f が各層の処理(AttentionやFFN)で、その出力に入力 h_{l-1} をそのまま足す。これにより勾配が「ショートカット」を通って深い層まで届くようになり、100層以上のネットワークでも安定して学習できるようになった。現在のTransformerは全層でこの残差接続を使っている。

通常のTransformerでは何が起きているか

残差接続は深層学習のブレイクスルーだったが、ここで「通常のTransformerが情報をどう扱っているか」を整理しておく。

48層のTransformerを考えよう。入力トークンの埋め込みベクトルが第1層に入ると、各層で以下が繰り返される。

  1. 前の層の出力を受け取る
  2. Attention(またはFFN)で処理する
  3. 処理結果 + 前の層の出力をそのまま足す(残差接続)
  4. 次の層に渡す

ポイントは3番目。すべての層で重み1の固定足し算なので、第48層の出力は48個の層の処理結果が均等に混ぜ合わさった状態になっている。各層が「この層の情報は重要」「この層の情報は不要」と判断する仕組みがない。

graph TD
    subgraph 通常のTransformer
    X0[入力埋め込み] --> L1[Layer 1]
    L1 -->|"+1"| A1((+))
    X0 -->|"+1"| A1
    A1 --> L2[Layer 2]
    L2 -->|"+1"| A2((+))
    A1 -->|"+1"| A2
    A2 --> L3[Layer 3 ...]
    L3 -->|"+1"| A3((+))
    A2 -->|"+1"| A3
    A3 --> OUT1[最終出力]
    end

これは画像認識のCNNでも言語モデルのTransformerでも同じ構造で、2015年のResNet以来ほぼ変わっていない。

希薄化の問題

均等足し算が具体的に何を引き起こすか。

たとえば第30層が「この文脈ではプログラミング言語の型情報が重要」と学習して強いシグナルを出したとする。しかし最終出力に到達するまでに、第31層〜第48層の出力が同じ重みで上乗せされていく。第30層のシグナルは18回の足し算で薄まる。

逆方向も問題で、第5層が出した低レベルの構文情報が最後まで同じ重みで残り続ける。第48層が高度な推論をしている段階で、もう不要な初期層の情報がノイズとして混入している状態だ。

これが希薄化(dilution) で、hidden state のノルムが層数 L に比例して O(L) で増大する。PreNorm(LayerNormをAttention/FFNの前に適用する方式)が標準化した現代のLLMでは特に顕著で、各層の相対的な寄与がどんどん小さくなっていく。

通常のディープラーニングとAttnResの根本的な違い

通常のディープラーニング(ResNet系のすべてのアーキテクチャ)では、層間の情報伝搬は静的なハードコーディングだ。設計者が「前の層の出力を重み1で足す」と決めたら、どんな入力が来てもその配分は変わらない。数学の問題でもコード生成でも翻訳でも、全く同じ均等配分で情報を混ぜ合わせる。

通常のディープラーニングAttnRes
層間の情報伝搬固定(重み1の足し算)動的(softmax attention)
参照先直前の層のみ全前層
入力への適応なし入力ごとに重みが変化
不要な層の情報そのまま蓄積重みを下げて抑制できる
遠い層の情報中間層を経由してしか届かない直接参照できる

AttnResは「どの前層の情報をどれだけ使うか」を入力に応じて動的に決める。コード生成タスクなら構文解析に強い層の重みを上げ、数学の推論なら論理構造を捉えた層の重みを上げる、といった適応が可能になる。

Transformerの成功から得た着想

論文が指摘する重要な類比がある。Transformerは時系列方向で、RNNの固定重み集約をsoftmax attentionに置き換えて成功した。

固定重み集約Attention集約
時系列方向RNN(前ステップを固定重みで引き継ぐ)Transformer(全時刻をsoftmaxで選択参照)
深さ方向残差接続(前層を重み1で足す)AttnRes(全前層をsoftmaxで選択参照)

RNNが長距離の情報を保持できなかった問題を、Transformerがattentionで解決したのと同じ構図。深さ方向でも固定重みの限界をattentionで突破するのがAttnResの着想だ。

Attention Residualsの仕組み

通常の残差接続が「前層の出力を均等に足す」のに対し、AttnRes は「全前層の出力をsoftmax attentionで重み付けして合成する」。

hl=i=0l1αilvih_l = \sum_{i=0}^{l-1} \alpha_{i \to l} \cdot v_i

viv_i は第 ii 層の出力、αil\alpha_{i \to l} はsoftmax正規化された注意重み。通常の残差接続では α\alpha が全部均等だったところを、学習可能な重みに置き換えた形になる。

graph TD
    subgraph AttnRes
    Y0[入力埋め込み] --> M1[Layer 1]
    M1 --> S1{softmax<br/>attention}
    Y0 -->|"α₀→₁"| S1
    S1 --> M2[Layer 2]
    M2 --> S2{softmax<br/>attention}
    Y0 -->|"α₀→₂"| S2
    M1 -->|"α₁→₂"| S2
    S2 --> M3[Layer 3 ...]
    M3 --> S3{softmax<br/>attention}
    Y0 -->|"α₀→₃"| S3
    M1 -->|"α₁→₃"| S3
    M2 -->|"α₂→₃"| S3
    S3 --> OUT2[最終出力]
    end

注意重みの計算式は以下。

αil=ϕ(wl,ki)j=0l1ϕ(wl,kj),ϕ(q,k)=exp(qRMSNorm(k))\alpha_{i \to l} = \frac{\phi(w_l, k_i)}{\sum_{j=0}^{l-1} \phi(w_l, k_j)}, \quad \phi(q,k) = \exp(q^\top \text{RMSNorm}(k))

重要な設計上の選択が2つある。

クエリは層ごとの固定ベクトル wlw_l(入力非依存)。入力依存クエリを試した実験では検証損失は若干改善するが(1.731 vs 1.737)、デコード時の逐次メモリアクセスが発生してしまう。スループットとのトレードオフで固定クエリを採用した。

RMSNormを kk に適用。大きなマグニチュードを持つ層が注意重みを独占するのを防ぐ。これを外すと損失が悪化する(1.743)。

過去手法との比較

深さ方向の集約を行列 M で表現すると、各手法の特性が整理できる。

手法重みの種類参照範囲
残差接続(標準)固定(均等)直前層のみ
DenseFormer静的スカラー全前層
mHC動的(入力依存)mストリーム
AttnRes Full動的(softmax)全前層
AttnRes Block動的(softmax)ブロック単位で集約した前層

DenseFormerは「全前層を参照するが重みが固定」という構成で、アブレーション実験では baseline の損失(1.766)と変わらなかった(1.767)。重みが静的(入力に関係なく固定)では効果がないということで、入力依存の動的な重み付けが本質という結論の根拠になっている。通常のディープラーニングの「設計時に決めた固定配分」がそもそも限界だったわけだ。

Block AttnRes: 実用的な効率化

Full AttnRes の問題は、全前層の出力を保持するメモリコストとパイプライン並列通信コストが深さに比例して増大する点だ。Block AttnRes はこれを解決する。

L層をN個のブロック(各ブロックS=L/N層)に分割し、各ブロックをその出力の総和 bnb_n で代表させる。結果としてメモリと通信量が O(Ld) から O(Nd) に削減される。

実験によると N≈8 ブロックでほぼ Full AttnRes と同等の性能が得られる。

訓練時の最適化

パイプライン並列学習では各ステージがブロック表現を前ステージから受け取る必要がある。論文は cross-stage caching を導入し、各物理ステージが前の仮想ステージのブロック表現をキャッシュする設計にした。

ナイーブな実装ではパイプラインの累積チャンク数 C に比例して通信量が増えるが(O(C²))、キャッシング後はパイプライン物理ステージ数 P に圧縮され(O(P²))、V倍の改善になる。定常状態(1F1B スケジュール)では計算と完全にオーバーラップできるため、実際の訓練オーバーヘッドは4%未満

通常のTransformerの訓練パイプラインをほぼそのまま使えるのは実用上大きい。学習コードの変更は最小限で済む。

推論時の最適化

デコード時は Two-Phase Computation で対応する。

フェーズ処理目的
Phase 1(並列)ブロック間attentionを全S層で一括計算メモリアクセスを1回に集約
Phase 2(逐次)ブロック内attentionを各層で計算し、Phase 1結果とonline softmaxで合成入力依存の細粒度計算

ブロック表現の読み取りをS回から1回に削減することで、推論レイテンシは2%未満の増加に抑えた。通常のTransformerと比べてほぼ同等の推論速度を維持している。

長文コンテキスト(128Kトークン)のプリフィリングではブロック表現の保持に15GBが必要になるが、テンソル並列のシャーディングとチャンクドプリフィル(16Kチャンク)の組み合わせで0.3GB未満に削減できる。

Kimi Linear 48Bでの実証

Block AttnRes を Kimi Linear アーキテクチャ(48B総パラメータ / 3B活性化、MoEでDeepSeek-V3設計に類似)に統合し、1.4Tトークンで事前学習した結果。

ベンチマークベースラインAttnRes改善
GPQA-Diamond36.944.4+7.5
Math53.557.1+3.6
HumanEval59.162.2+3.1
MMLU73.574.6+1.1
BBH76.378.0+1.7
C-Eval79.682.5+2.9
CMMLU82.082.9+0.9

特に多段階推論(GPQA-Diamond)とコード生成(HumanEval)での改善が大きい。これらのタスクは深い層で構築された抽象的な表現を正確に参照する必要があり、通常のTransformerの均等足し算では情報が希薄化しやすい領域だ。深さ方向の選択的な情報流が効くという仮説と一致している。

スケーリング則

194M〜528M活性化パラメータの5サイズでの検証では、Block AttnRes が baseline と比べて同等の検証損失を1.25倍少ない計算量で達成。別の見方をすれば、同じ計算量を使ったとき性能が上がる。

訓練後の解析では、各層の学習された注意重みの可視化から興味深いパターンが見えている。

  • 直前層への重みが最大だが、選択的な遠距離スキップ接続が出現する。通常の残差接続では不可能な「飛ばし参照」をネットワークが自発的に学習している
  • pre-attentionとpre-MLPで参照パターンが異なる層特化が進む。Attention層は浅い層の構文情報を多く参照し、FFN層は深い層の意味情報を多く参照する、といった分業が生まれている

AttnRes によってネットワークが「深くて細い」アーキテクチャを好むようになるという結果も面白い。固定パラメータ予算の下でのアーキテクチャサーチでは、baseline が dmodel/Lb60d_\text{model}/L_b \approx 60 付近で最適化されるのに対し、AttnRes は 45\approx 45 付近(より深い・細い)に最適が移動した。深さ方向の情報流が改善されたことで、深いネットワークのメリットをより引き出せるようになった形だ。

参照