【DL輪読会】Attention Residuals

Attention Residuals Hiroyuki Matsushima, Matsuo・Iwasawa Lab 1

書誌情報 Attention Residuals ● Author: Kimi Team, Guangyu Chen,et al. (Moonshot) ● Preprint ( Arxiv: https://arxiv.org/abs/2512.24880) TL;DR ○ 固定重みの残差接続 hₗ = hₗ₋₁ + fₗ₋₁(hₗ₋₁) を「深さ方向の softmax attention」に置き換えることで、 PreNormの希釈問題を緩和し、下流タスクで一貫した性能向上を達成する。 2

• https://arxiv.org/abs/2512.24880

背景：なぜ残差接続を「再考」するのか モチベーション ■ ■ ■ 残差接続とPreNormは現代LLMの標準構成だが、各層の 出力を固定の単位重みで足し合わせることしかできな い。 PostNorm ノルム有界・勾配消失 その結果、隠れ状態のノルムが深さに比例して O(L) で増 大し（PreNorm希釈問題）、後段の層は影響力を保つため に出力を不必要に大きく学習せざるを得なくなる。 PreNorm 既存の改良手法（DeepNet、Highway、mHC等）はいずれ も単一状態の再帰か固定重みにとどまり、前の層の出力 を内容に応じて選択的に取り出す仕組みが欠けている。 → 個別の層出力に選択的アクセスする機構が必要 恒等経路を保持・ノルム発散(希釈) 両者ともに 「層l は hₗ₋₁ という単一の圧縮状態にしかアクセスできない」 Attention Residuals ｜ Kimi Team, 2026 ｜ Slide 2 / 16

目的・仮説： RNNからの着想に基づく置換 研究目的 ■ 残差接続を「時間方向の RNN 的再帰」と捉え、その置換手法を提案する。 ■ 深さ方向にも softmax attention を適用し、各層が任意の層出力を 選択できるようにする。 中心的仮説 ■ RNN→Transformer で「時間」に対し起きた linear→softmax 移行を、「深さ」 に対しても実施することで性能が向上するはずである。 ■ 層数 L < 1000 程度であるため、 O(L²) attention は実装可能な計算量であ る はずである。 主要評価項目 ■ Validation loss、出力ノルム・勾配ノルム、 14種の下流ベンチマーク Attention Residuals ｜ Kimi Team, 2026 ｜ Slide 3 / 16

アプローチ全体像： 3つのアーキテクチャ Fig. 1 ● 従来の標準残差（ a）は、全層の出力を均 等な重みで足し合わせるだけの単純な構 造である。 ● Full AttnRes（b）は、層ごとの学習済みベ ク トル を用いて、全ての先行層出力に softmax attentionを適用し、重要な層を選 択的に集約する。 ● Block AttnRes（c）は、L層をN個のブロッ ク に分割し、ブロック内は通常の残差で足し 合わせ、ブロック間でのみ attentionを適用 することで、メモリ・通信コストを O(Ld) から O(Nd)に削減する。 ● Full AttnResとBlock AttnResの本質的な 違いはattentionの適用範囲のみであり、 Block AttnResはdrop-inで置き換えが可能 である。 Attention Residuals ｜ Kimi Team, 2026 ｜ Slide 4 / 16

手法 (1)：Full AttnRes の定式化 数式定義 (Eq.2) softmax over depth ● 計算量は O(L*2d)だが、ネットワークの深さ Lはシーケンス長よりはるかに小さいため、 負荷はわずかである。 ● 通常の学習ではバックプロパゲーション用 のメモリを再利用できるため、 追加メモリはほぼ不要だが、大規模訓練 （パイプライン並列など）では層を跨ぐ O(Ld) の通信が課題となる。 ● 追加パラメータは各層にベクトル w_l と RMSNormを1つずつ足すのみで、モデル全 体のサイズに対して無視できるほど 小さい。 (Eq.3) Query/Key 設計 qₗ = wₗ ∈ ℝᵈ (層ごとに 1 本だけ持つ学習可能ベクトル) kᵢ = vᵢ = h₁ (i=0) / fᵢ(hᵢ) (i≥1) (Eq.4) 層 l の入力 Attention Residuals ｜ Kimi Team, 2026 ｜ Slide 5 / 16

インフラ： Cross-stage cache と Two-phase 計算 システム最適化のポイント ● Cross-stage caching: パイプライン並列時に 重複する送信を省き、通信量を仮想ステージ 数（V）倍削減する。 ● Two-phase computation : ブロック内の全層 をまとめて計算する Phase 1と、逐次計算す るPhase 2に分け、メモリ読み出し（ I/O）を大 幅に効率化する。 ● 推論効率 : ブロック版のメモリアクセス量は 約 5.5d で、類似手法（ mHC）の約 1/6 と極 めて低い。 ● 実測負荷 : 訓練時の負荷は 4% 未満、推論 時の遅延は 2% 未満と、無視できるほど小さ い。 ● 長文対応 : 128Kトークンの入力でも、テンソ ル並列を活用してデバイスあたりのメモリ消 費を約 1.9GB まで抑制できる。 ★ ここに Fig. 3 を貼り付け Cache-based pipeline communication 例 (P=4 ranks, V=2 virtual stages) (Each rank caches previously received blocks; transitions transmit only incremental blocks (+[b1,b2]).) ▸ 配置メモ: 論文 p.6 の Figure 3 全体を貼付（左右のVirtual stage 0/1 のダイアグラム）。 ▸ 縦横比: 横長 (おおよそ 16:7) (元論文の比率を維持すること) Attention Residuals ｜ Kimi Team, 2026 ｜ Slide 7 / 16

結果 (1)：Scaling law Fig. 4 ● Full AttnResおよびBlock AttnResのいずれも、すべての 計算リソース範囲においてBaselineを上回る低い損失 （Loss）を達成した。 ● Block AttnResの性能はFull AttnResに極めて近く、 モデル規模が大きくなるほどその差は 0.001 まで縮小す る。 ● 同じ計算量（5.6 PFLOP/s-days）で比較すると、Block AttnResはBaselineが1.25倍の計算量を使った場合に匹 敵する低い損失を達成した。 ● 両者のスケーリング曲線の傾きはほぼ同一であり、 AttnResは学習効率の指数を変えるのではなく、 ベースとなる性能を底上げしている。 Attention Residuals ｜ Kimi Team, 2026 ｜ Slide 8 / 16

結果 (2)：出力ノルム・勾配の改善 Fig. 5 ★ ここに Fig. 5 を貼り付け (a) Validation loss / (b) Output magnitude / (c) Gradient magnitude across blocks の3パネル (AttnRes consistently lower val loss; bounded output magnitudes; uniform gradient distribution across depth.) ▸ 配置メモ: 論文 p.10 の Figure 5 全体（3パネル）を貼付。 ▸ 縦横比: 横長 (3:1 程度) (元論文の比率を維持すること) (a) Validation loss (b) Output magnitude (c) Gradient magnitude AttnResは学習の全期間を通じて Baseline より低い損失を維持し、特に学習終盤の減 衰フェーズでその差がさらに拡大して、最終 的に明確に優れた精度に到達する。 従来のBaselineでは深い層ほど出力が 指数関数的に増大する「 PreNorm希釈 問題」が発生するが、 Block AttnResは ブロック境界で蓄積をリセットするため、出 力の大きさを一定範囲内に 抑えられる Baselineでは勾配が最初期の層に極端に 集中してしまうが、 AttnResはsoftmaxによ る層間の重み競合を通じて、勾配を全層に わたって均一に分布させる。 Attention Residuals ｜ Kimi Team, 2026 ｜ Slide 9 / 16

結果 (3)：下流タスクでの比較 図から読み取れること Table 3 ■ 評価した全 14タスクにおいて、 AttnResはBaselineと同等以上 の成績を収めた。この改善効果は一部のタスクに限定され ず、全方位に及んでいる。 ■ 特に、多段階の思考が必要なタスクで顕著な伸びが見ら れ、GPQA-Diamondで+7.5、Mathで+3.6、HumanEvalで+3.1の 大幅な向上を記録した。 ■ 特に、多段階の思考が必要なタスクで顕著な伸びが見ら れ、GPQA-Diamondで+7.5、Mathで+3.6、HumanEvalで+3.1の 大幅な向上を記録した。 ■ CMMLUやC-Evalといった中国語ベンチマークも向上しており、 この手法が言語を問わず普遍的な効果を発揮することを示して いる。 深さ方向の情報伝達が改善されたことで、深い層が初期の層の 表現を必要に応じて再利用できるようになった。 このことが、複数の情報を組み合わせて解く「合成的・複雑な」タスクに おいて、特に高い効果を発揮したという仮説を裏付けている。 Attention Residuals ｜ Kimi Team, 2026 ｜ Slide 10 / 16

結果 (4)：アブレーション実験 Fig. 6 Table 4 Ablationのポイント ■ 入力に応じて動的に重みを変える仕組みが不可欠であり、固定重みの手法（DenseFormer等）を上回る。重みの決定には softmaxが最適であり、各層のスケールを揃えるRMSNormも性能維持に極めて重要である

分析 (1)：AttnRes はより深いネットワークを好む ポイント ● 計算量とパラメータ数を一定に保った条件 下で、25通りの構成による総当たり実験 （グリッドサーチ）を実施した。 ● どちらの手法も モデルの次元数が多いほ ど・Multi-Head Attentionのヘッド数を少な くするほど性能が向上するという共通の傾 向が見られた。 ● 同じパラメータ数であれば、AttnResは Baselineよりも層が深くてパラメータが少な いモデルでより高い効果を発揮することを 示している。 Attention Residuals ｜ Kimi Team, 2026 ｜ Slide 12 / 16

分析 (2)：学習された depth-wise attention の構造 Fig. 8 3つの学習現象 ■ すべての残差接続のバリエーションは、過去の 出力をどう混ぜるかを示す行列を用いた数式 として統一的に記述できる。 ■ 標準的な残差やHighwayなどの従来手法は、 この行列 が単純な低ランク構造であり、実質 的には「深さ方向のLinear Attention」を行って いるに等しい。 ■ mHCなどの手法は、情報を複数の ストリームで持つことで行列のランクを 拡張しており、これは深さ方向における「状態 の拡張」と見なせる。 ■ これらに対し、AttnResは密なsoftmax重みを 採用することで、深さ方向における「softmax Attention」を完成させたモデルであるといえ る。 Attention Residuals ｜ Kimi Team, 2026 ｜ Slide 13 / 16

限界と強み Strengths ● 48Bパラメータの MoEモデルを用い、 1.4兆トークンの膨 大なデータで訓練してその効果を証明している。 ● 全14種のベンチマークで Baselineを上回り、言語理解か ら数学、コード生成まで分野を問わず一貫して改善す る。 ● シーケンスと深さの類似性に基づき、あらゆる残差手法 を共通の理論モデルの中に位置づけた。 ● 既存モデルにベクトル 1本とRMSNormを足すだけで置き 換え可能であり、訓練・推論への負荷も極めて小さい ● 動的な重み付け、 softmax、ブロックサイズなど、 各構成要素の重要性が詳細な比較実験で個別に裏付 けられている Limitations ■ 比較対象が PreNormに偏っており、他の正規化手法との対 比は十分ではない ■ ブロック数の決定は経験則に頼っており、理論的な根拠の 解明が課題 ■ 性能と負荷の兼ね合いから、クエリは入力に依存しない設 計としている Attention Residuals ｜ Kimi Team, 2026 ｜ Slide 15 / 16