技術要約：Transformer言語モデルにおける形状慣習の再検討

問題提起

高密度（Dense）なTransformer言語モデルは、歴史的に厳格なアーキテクチャ上の慣習に従ってきました。それは、フィードフォワードネットワーク（FFN）が「狭い–広い–狭い」というマルチレイヤーパーセプトロン（MLP）の形状を用いるというものです。この設計では、モデル次元（$d_{model}$）が通常2から4の拡張比率でより広い中間次元（$d_h$）へと拡張された後、再び投影されます。この構造はパラメータ割り当ての大部分を占め、モデルのパラメータおよび浮動小数点演算の多くを占めることが一般的です。

著者らは、この特定の形状が不可欠なものなのか、それとも単に便利なインスタンス化に過ぎないのかという疑問を投げかけています。近年の理論的研究は、残差を持つ「広い–狭い–広い」（アワーグラス／砂時計型）MLPが最適な関数近似器として機能し、より少ないパラメータで優れた表現力を提供できる可能性を示唆しています。さらに、生成ビジョンモデルやMixture-of-Experts（MoE）Transformerにおける経験的な証拠は、アワーグラス構造が標準的な設計よりも優れた性能を示すことを示しています。核心となる問題は、従来の「狭い–広い–狭い」FFNが必須の制約なのか、それとも、より深く、より狭いアワーグラスFFNが、アテンションモジュールとFFNモジュールの間の容量配分をより効率的に行えるのかという点です。

手法

本論文は、従来のFFNを**アワーグラスFFN（Hourglass FFN）**に置き換えたTransformerの変種を提案しています。その手法は以下の通りです。

1. アーキテクチャの再設計:

   • アワーグラスFFNは、単一の広い拡張レイヤーを、残差経路によって接続された$K$個のアワーグラス形状のサブMLPのスタックに置き換えます。
   • 各サブMLP（$MLP_i$）は、圧縮・拡張構造に従います。すなわち、入力を$d_{model}$からボトルネック次元$d_h$（ここで$d_h < d_{model}$）へと投影し、非線形活性化関数（SwiGLU）を適用した後、$d_{model}$へと投影します。
   • レイヤーの出力は、入力をこれら$K$個の残差サブブロックに通した結果となります。

2. パラメータ再配分戦略:

   • 中間次元$d_h$を小さくすることで（FFNを「より狭く」することで）、FFNの幅に割り当てられていたであろうパラメータを節約します。
   • 節約されたパラメータは、以下のいずれかに再配分されます：
     • モデル次元（$d_{model}$）およびアテンション容量の増加。
     • FFNスタックの深さ（$K$）の増加。
     • Transformerレイヤー数（$L$）の増加。

3. 実験設定:

   • ベースライン: LLaMAスタイルのTransformerおよび最先端のOLMo-2アーキテクチャとの比較を行います。
   • スケール: 113Mから1Bパラメータの範囲のモデルサイズにわたって実験を行います。
   • 制約: FFNの形状による影響を分離するため、総パラメータ予算を一致させ（0.001%以内の差）、アテンションサイズを固定した条件下で実験を行います。
   • 探索空間: $K$（アワーグラススタックの深さ）、$L$（Transformerレイヤー数）、$d_{model}$、および比率$d_h/d_{model}$を変化させるアーキテクチャ探索を実施します。

主な貢献

本論文には主に4つの貢献があります：

1. 標準への挑戦: 従来の「狭い–広い–狭い」FFNの形状に異を唱え、スタックされた残差接続を持つサブMLPを通じて表現力を維持する、効率的な「広い–狭い–広い」アワーグラスFFNを提案しています。
2. パラメータの再配分: パラメータをFFNの幅からアテンションの幅や内部FFNの深さ（$K > 1$）へと転換することで、特に小規模なスケール（113M）において優れた性能が得られることを示しています。
3. アーキテクチャに関する洞察: $d_{model}/L$および$d_h/d_{model}$の最適な比率を見出すことで、U字型の幅・深さのトレードオフを特定し、深いアワーグラスFFN構造が堅牢であることを確立しました。
4. スケーラビリティの検証: 113Mから1Bパラメータまでの範囲でアワーグラスFFNを検証し、それが標準的なベースラインに対して競争力のある代替案であり、特定の構成が従来のモデルを凌駕し、大規模なモデルではそれらに匹敵することを示しました。

結果

経験的な検証により、以下の知見が得られました。

• 固定アテンションにおける性能: アテンションパラメータを固定した113Mスケールにおいて、アワーグラスFFNは従来のFFNよりも低い検証損失およびパープレキシティ（PPL）を達成しました。例えば、$K=8$の場合、PPLは36.441（従来型）から36.179へと低下しました。
• 最適構成:
  • 深さ vs 幅: $d_h/d_{model} \approx 0.4$、かつ$K=4$レイヤーの構成が最も低い検証損失をもたらしました。
  • パラメータのシフト: 113Mにおいて、アテンションモジュールを支配的にするようパラメータをシフトさせた場合（例：アテンション66M対FFN 47M）、最高の性能（PPL 35.101）が得られました。これは、小規模モデルにおいては、静的な事実検索よりも文脈処理の方が価値が高いことを示唆しています。
  • 幅と深さのトレードオフ: $d_{model}/L$の比率に関して、明確なU字型の曲線が観察されました。最適なスイートスポットは100から250の間でした。
• スケーラビリティ:
  • 小規模から中規模（113M – 906M）: アワーグラス・バリアントは、検証損失、パープレキシティ、およびダウンストリームの推論タスク（Arc Easy、HellaSwag、PIQAなど）において、一貫して従来のベースラインを上回りました。
  • 大規模（1B）: 1Bパラメータスケールにおいて、アワーグラス・バリアントは強力な従来型ベースラインに匹敵する性能（検証PPL 20.082 vs 20.002）を達成し、OLMo-2の構成（20.466）を上回りました。
• アブレーション研究: アワーグラスの深さ$K$を増やすことは一貫して性能を向上させましたが、$K=6$を超えると収穫逓減が見られました。$d_h/d_{model}$を0.4未満に下げると大幅な性能低下を招き、FFN圧縮における下限が存在することを示しました。

意義と主張

本論文は、Transformerにおける長年の「狭い–広い–狭い」MLP形状は不可欠ではなく、冗長である可能性があると主張しています。アワーグラスFFNを採用することで、著者らは以下のことを実証しました：

1. 効率性: モデルは、より効率的なパラメータ分布によって、競争力のある、あるいは優れた性能を達成できます。
2. リソースの再配分: アワーグラス構造は、アテンション機構が主要なパラメータ消費量となるような根本的なシフトを可能にし、従来のFFN主導の予算配分に疑問を投げかけます。
3. 設計の再評価: これらの知見は、アテンションとFFNのバランスの再考を促すものであり、効率的かつ表現力豊かな現代の言語モデルのためには、より深く、より狭いFFNと、より広いアテンション機構を組み合わせることが優れた設計戦略である可能性を示唆しています。

著者らは、数十億または数百億のパラメータを持つモデルへのスケーラビリティについては、113Mスケールからの外挿であること、および「深さの呪い」（信号伝播の問題）がより深いアーキテクチャにおいてさらなる調査を必要とすることを挙げ、慎重な姿勢を保っています。また、アワーグラスFFNと高度なアテンションメカニズム（Group-Query Attentionなど）との相互作用については、未探索のままであることも述べています。