Transformer の深さを「ランク保存」から読み解く

導入

Transformer は深く積み重ねられる一方で、残差接続の強さ、正規化の位置、フィードフォワード層の幅といった細部に大きく依存する。論文「Transforming Rank: How Architecture Navigates the Spectral Pathologies of Depth」は、この問題を「ランクが深さを越えてどれだけ生き残るか」という観点から説明しようとしている。

ここでいう問題は、活性値や勾配のノルムが大きくなる、あるいは小さくなることだけではない。層を重ねることで、高次元空間に存在する独立した方向が失われ、ヤコビアンのランクが低下する可能性がある。著者は Transformer のフィードフォワードブロックに焦点を当て、初期化時の入力出力ヤコビアンや分岐ヤコビアンのランクを分析する。

主要なポイント

• 残差接続はランク低下を迂回する経路になる。 行列積や非線形活性化は表現力を生む一方で、ランクを減らす要因にもなる。残差経路はその分岐を迂回して勾配を流すため、より多くの方向を保ちやすい。ただしスキップ経路が強すぎると、深い合成というより浅い変換の集合のような振る舞いに近づく。

• 正規化の位置は分岐とスキップの比率を変える。 Pre-Norm と Post-Norm の違いは、単に安定性の経験則ではなく、深さに応じて残差分岐とスキップ経路の相対スケールをどう設定するかという問題として捉えられる。論文は、Post-Norm でランク崩壊が起こりやすく、Pre-Norm でランクがプラトーに達しやすい理由をこの観点から説明する。

• 二つの線形行列にはランク保存上の意味がある。 Transformer のフィードフォワード層は通常、次元を拡張し、活性化を通し、再び縮小する。著者は、二つ目の行列が非中心化された活性化から生じる一貫した平均スパイクを分散させ、残差表現が少数方向へ潰れることを防ぐと説明する。

• 中間幅の拡張は分岐ヤコビアンを保つ。 活性化によって一部の方向が失われるなら、より広い空間で活性化を適用することで、元の次元を張るのに十分な方向を残しやすくなる。論文はこの幅の条件を Marchenko–Pastur 法則と結びつけている。

• 初期化時のランクは訓練可能性の手がかりになる。 著者は、入力出力ヤコビアンの初期ランクが CIFAR-10 で訓練できるネットワークの予測に関係すると報告している。

意義と影響

この論文の面白さは、新しいブロックを提案することよりも、既存の設計要素を別の言葉で説明し直す点にある。残差接続、正規化、幅拡張はしばしばスケール安定化の道具として語られるが、本論文ではそれらが深いネットワークで独立方向をどれだけ保つかという共通課題に対応していると見る。

大規模モデルの設計にとって、この視点は Pre-Norm が深い Transformer でよく使われる理由や、フィードフォワード層がいったん幅を広げる理由を理解する助けになる。単に層数を増やすだけではなく、ランク崩壊、深い合成能力、パラメータ数の間でバランスを取る必要があるという示唆も与える。

もっとも、分析の中心は初期化時とフィードフォワードブロックであり、完全な大規模言語モデルの訓練や多様な活性化関数にどこまで一般化できるかは今後の検証が必要だ。それでも、深さを「ランクの伝播」として見る考え方は、Transformer アーキテクチャを理解する有用な補助線になる。

出典：arXiv