Moonshot AIが「AttnRes」を発表：残差接続を注意機構に置き換え推論性能を向上

ニューラルネットワークを少しでも深く組んだことがある人なら、残差接続（Residual Connection）のありがたみを知っているはずだ。入力をそのまま出力に足し合わせるこのシンプルな工夫こそが、勾配消失を防ぎ、現在の巨大なTransformerの基礎を作った。

しかし、10年近く使われてきたこの仕組みには弱点があった。すべての層の出力を均等に足し合わせるため、層が深くなるにつれて初期の表現が薄まり、PreNorm特有の隠れ状態の数値が発散していく問題に直面していた。

Moonshot AI（Kimiチーム）はこの課題に対し、「AttnRes（Attention Residuals：注意残差）」という新しいアーキテクチャを提案した。従来の固定的な残差累加を、深さ方向への動的なアテンション機構へと置き換える試みだ。

過去の層の記憶を、必要な分だけ引き出す

従来の残差接続は、いわば単純な書類の積み重ねだ。どれだけ層が進んでも機械的に足し合わされるため、最初の方に記録した重要な特徴量は、あとの層になるほど埋もれてしまう。

AttnResのアプローチは明快だ。トークン間で使っている注意機構を、層の深さ方向に対しても適用する。各層に学習可能な軽量のクエリを1つ用意し、先行するすべての層の出力に対してソフトマックス関数を用いた注意の重みを計算する。これにより、深層のレイヤーであっても、初期層の細かな特徴を直接選び取って利用できるようになる。

全層でこれを愚直に行うとメモリ消費量がO(Ld)に跳ね上がってしまう。そこでMoonshotチームは、実用性を重視した「Block AttnRes」を考案した。

層を約8個などのブロックに分割し、ブロック内では計算負荷の低い従来の残差接続を使い、ブロック間をまたぐ時だけアテンションを適用する。これによりメモリ消費量をO(Nd)に抑え、わずかなオーバーヘッドで全量アテンションと同等の効果を再現することに成功した。

複雑な推論ベンチマークで7.5ポイントの上昇

公開されたデータによると、Block AttnResを導入したモデルは、ベースラインモデルが1.25倍の計算量を投じて到達した損失値と肩を並べた。単純な構造改善だけで、計算資源の効率を25%底上げした計算になる。

性能面で最も恩恵を受けたのは、多段階の思考が求められる推論とプログラミングの領域だ。難関とされるGPQA-Diamondでは、ベースラインの36.9から44.4へと一気に7.5ポイント上昇した。コード生成ベンチマークのHumanEvalでも59.1から62.2へ、MBPPでも72.0から73.9へと着実に伸ばしている。

一般的なベンチマークでもMMLUが73.5から74.6、BBHが76.3から78.0へ、中国語のC-Evalも79.6から82.5へと全体的に底上げされた。さらに、PreNormによる隠れ状態の増大を抑制し、層ごとの勾配ノルムの分布が滑らかになることも確認されている。

なぜ重要か

この研究が巧みなのは、誰もが「当たり前の前提」として触れてこなかった層間の骨格にメスを入れた点にある。ここ数年、業界全体がコンテキスト長の拡張やアテンションのスパース化ばかりに注力してきた中で、垂直方向の残差接続にまだこれだけの伸び代があることを証明してみせた。

基盤モデルを自前で事前学習している開発陣にとって、ごくわずかなコード修正で計算効率を1.25倍に引き上げられる魅力は大きい。特にGPQA-Diamondのような高難度の推論タスクで目に見える成果が出ている点は、推論モデル開発における強力な武器になるはずだ。

ただし、超大規模クラスタでの運用には注意が必要だと考える。Block構造によってメモリを抑えたとはいえ、パイプライン並列（Pipeline Parallelism）を採用する環境では、異なるGPU間でブロック表現のやり取りが発生し、通信のオーバーヘッドやスケジューリングの複雑化を招く恐れがある。この並列化の壁をどうクリアするかが、実戦投入への鍵になりそうだ。