層ごとに幅を変える×字形トランスフォーマーが少ない計算で性能を改善
この論文は、トランスフォーマー型の言語モデルで「層ごとの幅(隠れ次元)」を一定にしない設計を試したものです。研究者は、前半と後半の層を広くし、中間層を狭くする×字形(バツ字)プロファイルのモデルを提案しました。こうした非一様な幅割り当てで、同じパラメータ数の均一幅モデルよりも言語モデルの損失が小さくなり、計算量や推論時のメモリも節約できることを示しています。具体的には、実験した範囲では性能の改善とともにKVキャッシュ(注意機構で使うキー・バリューの保存領域)や入出力のコストが減りました。
研究者は「×字形」と呼ぶ設計を詳しく調べました。ここでは初めの層と最後の層を広く取り、中間層を細くします。層ごとの幅を変えるときに問題になるのが層間での次元の増減です。これに対して彼らはパラメータを増やさない「残差ストリームのサイズ変更」方式を使いました。具体的には、各層は固定された大きな残差ベクトルの一部だけを読み書きします。ある座標がその層で使われない場合は、その座標はその層を飛ばし、直前にその座標を扱った層からコピーして戻します。縮めるときは余分な次元を切り捨て、拡張する時は直近で有効だった値を復元し、存在しなければ0で埋めます。
実験はデコーダーのみの言語モデルで行われ、モデル規模は200M、500M、1B、2Bパラメータ(密なモデル)と、合計3B/1BアクティブのMixture-of-Experts(専門家混合)モデルを含みます。論文の要約では、損失に合わせたスケーリング曲線の解析では全体のFLOPs(浮動小数点演算数)を最大で22%削減し、KVキャッシュとI/Oコストを15%削減できるとされています。実際の訓練済みモデル群では、200M〜2Bの範囲でおよそ3%の相対的なパープレキシティ改善を示し、KVキャッシュを約10%、FLOPsを約3%削減したと報告しています。Mixture-of-Expertsの場合も同様の利点が見られ、変形幅モデルは総パラメータを合わせた場合でもベースより優れていました。
なぜ重要かというと、モデルの「幅」を層ごとに最適に配分すると、同じパラメータ予算でより良い性能が得られる可能性があるからです。パラメータ数は幅の二乗で増える一方、注意計算やKV保存は幅に線形に比例します。したがって幅を非一様にすると平均的な層幅が小さくなり、推論時のメモリや計算が節約できます。論文の解析では、×字形のボトルネックが中間層での表現崩壊を緩和し、残差経路に異なる表現戦略を生むことも示されました。
注意点としては、本研究は観察的・経験的な結果に基づいています。得られた改善は実験の設定やモデル規模、学習データ量、ハイパーパラメータに依存します。論文はパラメータ数を揃えた比較を行っていますが、パラメータ数とFLOP数を同時に一致させることは理論的に不可能である点も指摘しています。また、残差ストリームの「コピーによる拡張」など実装上の工夫が利得に重要であるとされ、他の設計や別のタスク・データセットで同じ効果が得られるかは今後の検証が必要です。