Variable-Width Transformers（>＜former）とは？×形設計でLLMのFLOPs 22%削減

×字型（早期・後期を幅広く、中間を幅狭く）の非均一設計で、同パラメータ数モデル比でFLOPs 22%・KVキャッシュ 15%削減を達成
パラメータ不要の残差リサイジング機構が次元の増減を吸収し、追加の学習パラメータなしに実装可能
200M〜2BパラメータのデコーダモデルとMoEで有効性を検証し、言語モデリング損失をすべてのスケールで一貫して改善

研究の背景

現在の大規模言語モデル（LLM）のほとんどは、すべてのレイヤーで同じ幅（隠れ次元数）を使用する「均一幅」設計を採用しています。実装がシンプルで既存の演算ライブラリとの相性がよい反面、各レイヤーが担う役割の違いを無視した設計ともいえます。

MIT・IBMの研究チームは、この均一幅の前提を問い直しました。初期・後期レイヤーが多くの情報を保持・変換するのに対し、中間レイヤーはそれほど広い幅を必要としないかもしれません。もしそうなら、中間層の幅を絞ることで、計算量やメモリを削減しつつ同等以上の性能を引き出せるはずです。

提案手法：×字型アーキテクチャ

この考えを具体化した提案が「>＜former（ボウタイフォーマー）」です。蝶ネクタイ（英語でボウタイ）のような×字形状が名前の由来で、早期と後期のレイヤーを幅広くし、中間レイヤーにかけて幅が絞られていきます。

図1: >＜formerの全体構造。各レイヤーが異なる幅を持ち、非アクティブな次元は残差ストリームにコピーアップされる

この設計の核心はパラメータ不要の残差リサイジング機構です。通常、次元数を変えるには学習可能な射影行列が必要ですが、>＜formerは異なるアプローチを取ります。次元を縮小する際は余分な次元を切り捨て、次元を拡大する際は以前に退避させた値をそのまま復元します。これにより追加のパラメータなしに次元の増減を実現しています。

設計のハイパーパラメータは主に2つです。ボトルネック（幅が最も狭くなる）レイヤーの位置と、そのときの幅の比率です。複数スケールの実験から、ボトルネット位置を全レイヤー数の75%付近、ボトルネック幅をベース幅の30%程度にした設定が安定した改善をもたらすことが示されています。

図2: ボトルネック位置（rℓ）と幅比率（rd）が言語モデリング損失に与える影響。rℓ≈0.75、rd≈0.3付近が最適領域となる

実験結果

研究チームは200M、500M、1B、2Bパラメータの密なデコーダモデルと3B MoE（Mixture-of-Experts、複数の専門サブネットワークを組み合わせた手法）モデルで検証を行いました。学習データにはDCLMデータセットを使用し、Chinchilla最適の約2.5倍のトークン数で学習しています。

比較対象は同じパラメータ数の均一幅トランスフォーマーです。結果として、FLOPs（計算量）22%削減・KVキャッシュ 15%削減を達成しながら、言語モデリング損失を全スケールで一貫して下回りました。相対的なパープレキシティ改善は約3%です。

図3: 事前学習FLOPsと言語モデリング損失の関係。>＜formerは同FLOPs・同平均レイヤーサイズのベースラインを一貫して下回る損失を示す

KVキャッシュとは、LLMの推論時に各レイヤーのKey・Value行列を保存しておく仕組みで、長い文章の処理に必要なメモリの大部分を占めます。このキャッシュを15%削減できることは、推論サービングのコスト削減に直結します。KVキャッシュを活用した別のアプローチとして、KVEraserによるキャッシュ局所消去の手法も参照してください。

なぜ×字型が効くのか

研究チームは2BパラメータのモデルでMLPの活性化パターンや内部表現を詳しく分析しました。均一幅モデルでは中間層のMLP次元の多くがほとんど活性化しない「死んだ次元」となっているのに対し、>＜formerは中間層の幅をそもそも絞ることで、使われている次元を均等かつ密に活用することがわかりました。

正規化行列エントロピー（表現の次元利用率の均一さを示す指標）でも同様の傾向が確認されています。均一幅モデルは中間層で表現が崩壊する傾向があるのに対し、>＜formerは中間から後期レイヤーにかけてエントロピーが高い状態を維持し、多様な情報を安定して保ち続けます。

さらにLogit Lens分析（各レイヤーの出力を直接語彙空間に射影して次のトークン予測を確認する手法）では、>＜formerがネットワーク全体を通じてターゲットトークンの確率を高く保ちながら、中間層での分布変化が緩やかであることが示されました。情報処理がより段階的かつ安定的に進んでいる様子が見て取れます。