技術的概要：因果的エネルギー最小化を通じたトランスフォーマー層のパラメータ化の再検討

問題定義

トランスフォーマーアーキテクチャ、特にマルチヘッドアテンション（MHA）とゲート付き多層パーセプトロン（MLP）の組み合わせは、現代の大規模言語モデル（LLM）の基盤となる構築ブロックである。しかし、層のパラメータ化に関する具体的な選択（重み共有、射影構造、更新則など）は、明示的な計算原理ではなく、主に実証的研究によって駆動されている。エネルギーベースモデル（EBM）は、計算を低エネルギー状態の探索として捉える視点（例えば、ホップフィールドネットワークにおけるもの）を以前から提供してきたが、既存の研究はこれらのエネルギーベースの解釈を標準的なトランスフォーマー層の特定のパラメータ化と完全に結びつけておらず、またこの視点がより効率的かつ効果的な層のバリアントの設計をどのように導くことができるかを体系的に探求してきたわけではない。

手法：因果的エネルギー最小化（CEM）

著者らは、トランスフォーマー層を条件付きエネルギー関数上の最適化ステップとして再構成するフレームワークである**因果的エネルギー最小化（CEM）**を導入する。その核心となる前提は、層が入力系列 $h_{1:J}$ を出力 $h'_{1:J}$ に写像する際、各位置 $i$ を最適化変数 $x_i$（初期値は $h_i$）として扱い、因果的な履歴に依存する条件付きエネルギー $\epsilon(x_i | h_{1:i})$ を最小化するように手続き $\mathcal{A}$ を通じて更新する、というものである。

本論文はこのフレームワークから標準的なトランスフォーマー構成要素を導出し、拡張を提案する：

1. 重み共有アテンションの導出：
   現代のホップフィールドネットワークに類似した log-sum-exp 形式に基づく相互作用エネルギーを定義することで、著者らは単一の勾配降下ステップが標準的な MHA 更新を回復することを示す。重要なのは、この導出が特定の重み共有制約を必然的に必要とする点である：

   • キー射影（$W^K$）は値射影（$W^V$）と共有される。
   • クエリ射影（$W^Q$）は出力射影（$W^O$）と共有される。
   • 温度パラメータ $\tau$ はヘッド次元に関連する。

2. 重み共有ゲート付き MLP の導出：
   各トークンの特徴ベクトルに独立してエネルギーを割り当てる要素ごとのエネルギーを定義することで、勾配更新がゲート付き MLP の構造を回復する。この解釈は、MLP 内の「アップ」および「ダウン」射影行列が共有されるべき（具体的には、ダウン射影の転置がアップ射影に等しい）ことを示唆する。

3. 提案される拡張：
   エネルギー最適化の視点に基づき、著者らは 3 つの具体的なアーキテクチャの強化を提案する：

   • 対角＋低ランク（D+LR）パラメータ化： 相互作用行列 $A_k$ に対して純粋な低ランク射影を行う代わりに、$A_k = \text{diag}(d_k) + W^{Q\top}_k W^K_k$ を提案する。これは低ランク行列が見逃す相互作用を捉えつつ、計算効率を維持する。また、パラメータ削減のために対角項をヘッド間で共有することも提案する。
   • 学習可能な軽量前処理行列： 単一の勾配ステップを超えて収束を改善するために、D+LR 構造（$P = \text{diag}(\text{softplus}(d)) + UV^\top + VU^\top$）を持つ前処理行列 $P$ を導入する。これらは曲率情報の学習された代理として機能し、完全な 2 次手法のコストなしで勾配を再スケーリングする。
   • 層内再帰的更新： 単一の更新ではなく、フレームワークは単一の層内で $T$ 回の再帰ステップを可能にする（$x^{(t+1)}_i = x^{(t)}_i - \eta \nabla \epsilon$）。これにより、層は新しいパラメータを追加することなく、エネルギー関数をより深く最小化できるようになり、ブロック全体を再利用する「ループ型」トランスフォーマーとは区別される。

主要な貢献

• 理論的統合： 本論文は、特定の重み共有スキーム（アテンションにおける Q/O および K/V の共有、MLP におけるアップ/ダウンの共有）を持つ標準的なトランスフォーマー層が、特定の相互作用および要素ごとのエネルギー関数上の単一の勾配ステップと数学的に等価であることを確立する。
• 設計空間の拡大： CEM の視点は、D+LR 相互作用、軽量前処理行列、再帰的更新を含むトランスフォーマー層の新たな設計空間を特定し、これらはヒューリスティックではなく最適化理論に基づいて動機付けられている。
• 実証的検証： 著者らは、言語モデリングタスクにおいて中規模（数億パラメータ規模）で CEM 由来の層を評価し、これらの制約されたパラメータ化が標準的なトランスフォーマーのベースラインと同等かそれ以上の性能を達成しうることを示した。

実験結果

実験は、SlimPajama データセットを用いて行われ、モデルサイズは約 86M から約 162M パラメータの範囲で、計算最適化されたトークン予算（Chinchilla スケーリング）でトレーニングされた。

• パラメータ効率： 重みを共有した CEM モデル（Llama ベースラインの MHA パラメータの約半分、MLP パラメータの約 3 分の 2 を使用）は、競争力のあるテストパープレキシティを達成した。CEM アテンションにおいて、パラメータの大幅な削減にもかかわらず、性能の低下は最小限であった。
• 拡張の影響：
  • 再帰： 再帰ステップ数を $T=1$ から $T=2$ に増やすことは、CEM アテンションと MLP の両方の性能を一貫して向上させた。特に、再帰型 CEM アテンション変種は、より少ないパラメータを使用しながら Llama ベースラインを上回った。
  • 前処理行列： 学習された前処理行列は、再帰と比較してわずかな利益しか提供せず、改善の主な駆動力は再帰構造であることを示唆している。
  • 対角項： 相互作用行列に対角項を含めること（特にヘッド間で共有される対角項）は良好な性能に不可欠であり、純粋な低ランクまたは対角項なしの変種を上回った。
• 安定性： 完全に CEM 由来の層（$T=2$ および D+LR 前処理行列を使用）で構成されたトランスフォーマーは、エンドツーエンドで安定してトレーニングされ、より少ないパラメータで Llama ベースラインと同等のパープレキシティを達成した。
• 合成データ： 合成ガウス過程データでの実験は、より深い再帰（$T=4, T=8$）が RMSE のさらなる向上をもたらすことを確認したが、標準的なトレーニング設定下での完全なトランスフォーマー設定における利益は $T=2$ に限られた。

意義と主張

本論文は、CEM がトランスフォーマー層のパラメータ化を理解するための有用な理論的レンズを提供し、トランスフォーマーアーキテクチャをエネルギーベースモデルと効果的に結びつけることを主張する。その意義は以下の点にある：

1. 原理的な設計： 層の設計を、経験的な試行錯誤から、エネルギー最小化と最適化ダイナミクスに基づいたフレームワークへと移行させる。
2. 効率性： このフレームワークから導き出された重み共有と再帰的更新が、競争力のある性能を維持しながらパラメータ効率の高いモデルを可能にすることを示す。
3. 将来の方向性： 著者らは、この視点が以下のさらなる探求を動機付けることを示唆する：
   • カスタムカーネル： 構造化された重み共有と再帰的更新のランタイムを最適化すること（例えば、操作の融合）。
   • 推論時スケーリング： 推論時の計算スケーリングのための追加の軸として、層内再帰を利用すること。
   • ハードウェア共同設計： 新しいハードウェアアーキテクチャに特化した層パラメータ化の設計。

著者らは謙虚であり、評価が数億パラメータ規模に限定されていること、およびこれらの利点が最先端モデルへの転移可能性を完全に評価するためには、より大規模な研究と下流タスクの評価が必要であることを指摘している。また、現在の実装はランタイム性能のために最適化されていないことも認めている。