技術的概要：NRGPT – GPT のエネルギーベースの代替案

問題定義

生成事前学習トランスフォーマー（GPT）アーキテクチャは現在、自己回帰言語モデリングを支配しており、トランスフォーマー層の一系列を通じてトークンの系列をシフトされた系列（次のトークンの予測）にマッピングすることで動作します。効果的である一方、この標準的な設定はエネルギーベースモデル（EBM）とは明確に異なります。EBM は、推論をエネルギー地形上の動的プロセスとして捉えるパラダイムであり、低エネルギー状態が高尤度データに対応します。

トランスフォーマーと勾配降下法（例えば、コンテキスト内学習におけるもの）の間に深い関連性があることを示す証拠が増加していること、およびこれらの分野を統合しようとする以前の試み（Energy Transformers など）にもかかわらず、標準的な因果的 GPT 設定をエネルギーベースモデルとして定式化する明確な枠組みは依然として見出されていません。既存のエネルギーベースアプローチは、しばしばマスクされたトークンの再構成（因果的生成には不適切）に依存するか、トランスフォーマーを動的システムではなく単なるエネルギー計算機として扱っています。著者らは、言語モデリングの自己回帰性を犠牲にすることなく、標準的な GPT 推論メカニズムと EBM 枠組みを統合するギャップを特定しました。

手法

著者らは、GPT アーキテクチャを EBM 枠組みと統合する最小限の変更であるNRGPT（eNeRgy-GPT）を提案します。中核的な革新は、推論ステップを静的なフォワードパスではなく、定義されたエネルギー関数の勾配によって駆動される更新則に従ってエネルギー地形上でトークンを探索するプロセスとして概念化することです。

1. エネルギー関数の定式化

モデルは、アテンションエネルギー（$E_{AT}$）とフィードフォワードエネルギー（$E_{FF}$）の 2 つの項から構成される総エネルギー関数 $E$ を定義します。

• アテンションエネルギー: マルチヘッドアテンション機構から導出され、$E_{AT}$ は、トークン状態に関するその勾配がアテンション出力を生成するように定式化されます。著者らは、勾配降下ステップが標準的なアテンション機構を模倣する学習可能なアテンションエネルギーを定義し、学習可能な重み $\alpha$ と推論率行列 $\eta$ を利用します。
• フィードフォワードエネルギー: フィードフォワード（FF）ネットワークは、密結合連想記憶（Dense Associative Memory）の定式化から導出されます。エネルギー $E_{FF}$ は、その勾配が標準的な 2 層 MLP 更新を生成するように構築されます。著者らは、$E_{FF} = -\|\sigma(W g_A)\|^2$（NRGPT_H_FF1）や、2 つの重み行列を含むより複雑な変種（NRGPT_H_FF2W）など、特定の変種を探求します。

2. 更新則とダイナミクス

推論プロセスは、トークン状態 $x$ が時間ステップ $t$ において進化するとモデル化される動的システムとして記述されます：
$$x^{(t+1)} = x^{(t)} - \eta(t) \nabla E(x^{(t)})$$
ここで、$\eta(t)$ は推論率行列（前処理行列として機能）です。

• 因果的制約: トークン $A$ のエネルギーは、トークン $B \leq A$ のみに依存し、因果的生成を確保します。
• 漸近安定性: 著者らは、$\eta$ に関する特定条件（例えば、$\eta = c \cdot \text{diag}(\gamma)$ ここで $\gamma$ はレイヤーノーマリゼーションのスケーリング因子）および正規化の下で、システムが漸近安定性を示すことを証明します。過渡期間の後、各トークンのエネルギーは単調に減少し、トークン状態は固定点に収束します。これにより、「フォワードパス」が有効なエネルギー降下であるという理論的保証が提供されます。

3. アーキテクチャ変種

本論文はいくつかの構成を調査します：

• 正規化: モデルはレイヤーノーマリゼーションと RMS ノーマリゼーションをサポートし、これらの正規化がエネルギーを有界にし収束を確保するのに役立つことを示す理論的証明があります。
• FF 変種: 著者らは、パフォーマンスと標準 GPT へのアーキテクチャ的類似性のバランスを取るために、$E_{FF}$ の異なる形式をテストします。
• 再帰的推論: 層が明確に区別される標準 GPT と異なり、NRGPT は単一のブロックを再帰的に使用します。再帰ステップの数はネットワークの「深さ」に対応します。

主要な貢献

1. 統合された枠組み: GPT 設定を記述するエネルギー関数と更新則の設計により、自己回帰言語モデリングとエネルギーベースモデリングを統合します。
2. 理論的保証: 推論率行列と正規化に関する特定条件の下で、推論ダイナミクスが過渡期間後に証明的にエネルギーを減少させることを証明し、漸近安定性の基礎を確立します。
3. 実証的検証: NRGPT が以下で競争力のあるパフォーマンスを発揮することを示します：
   • ListOps: 入れ子になった算術および論理タスク。
   • Shakespeare: 文字レベルの言語モデリング。
   • OpenWebText (OWT): 大規模自然言語モデリング。
4. 体系的比較: 再帰型トランスフォーマーと NRGPT の間のパフォーマンススケーリング、パラメータ効率、FLOP コストを含む詳細な分析。
5. 動的解析: エネルギー地形上でのトークン軌跡の実証的研究により、安定した低エネルギー状態への収束を観察します。

実験結果

• ListOps: NRGPT 変種（特に NRGPT_H_FF1 と NRGPT_H_FF2W）は、精度と訓練損失の面で再帰型 GPT ベースライン（GPT_Rec_parallel）と同等のパフォーマンスを達成しました。モデルは、同程度のパラメータ数（約 $2.3 \times 10^4$ から $3.0 \times 10^4$ パラメータ）で学習遷移（80% 以上の精度への到達）を達成しました。
• Shakespeare: NRGPT は、さまざまなパラメータサイズにおいて、深い GPT ベースラインと同等の検証損失を達成しました。特に、NRGPT は標準 GPT に比べて、より大きなパラメータサイズにおいて訓練セットでの過学習が少なかったことが注目されます。
• OpenWebText: OWT データセットにおいて、NRGPT（特に NRGPT_H_FF2W 変種）は、標準 GPT（1 億 2400 万パラメータ）と競争力のある生成品質指標（ペルプレキシティ、文法スコア、多様性）を達成しながら、約 3400 万パラメータ少ない（9000 万対 1 億 2400 万）パラメータ数で構成されました。
• 下流タスク（MMLU）: MMLU データセットにおける 5 ショット設定において、1 億 2800 万パラメータの NRGPT モデルは、同程度のサイズの GPT（28.0%）および再帰型 GPT（26.0%）ベースラインと比較して、より高い平均精度（29.3%）を達成しました。著者らは、これらの結果はモデル規模が小さいため慎重に解釈すべきであると指摘しつつも、潜在的な利点を示唆しています。

限界と観察

• 計算コスト: NRGPT はパラメータ効率が高い一方、FF モジュールの設計とアテンションにおけるフルランク重み行列の使用により、同等の再帰型 GPT に比べて反復あたり 1〜2 倍の FLOP コストが発生します。
• ハイパーパラメータ感受性: NRGPT は、標準 GPT 変種と比較してハイパーパラメータの選択（学習率など）に対してより高い感受性を示します。
• 過学習の挙動: モデルは過学習に対してより抵抗性があり、これは小規模データ領域では有益ですが、過学習がしばしば複雑なパターンの学習の代理となる大規模データセットでは、必ずしも最適ではない可能性があります。
• 推論率の制約: エネルギー降下を保証するために、推論率行列 $\eta$ は特定の制約（例えば、正規化スケーリングに比例すること）を満たす必要があります。しかし、実証結果は、制約のない $\eta$ でも実際には収束につながる可能性を示唆しています。

意義と主張

本論文は、NRGPT が自己回帰言語モデリングとエネルギーベースモデリングを成功裏に統合する GPT アーキテクチャの最小限の変更を表すと主張しています。その主な意義は以下の点にあります：

1. 理論的統合: 因果的言語モデルの「フォワードパス」が明示的にエネルギー地形上の勾配ベースの動的システムであるという、原理的な基盤の提供。
2. 漸近安定性: この動的システムが自然に安定状態へ収束することを示し、トランスフォーマー推論に対する新たな視点を提供すること。
3. 代替パラダイム: エネルギーベースアプローチが、過学習への抵抗性や体系的な解空間の探索（早期停止や正則化子によるなど）といった固有の特性を提供しつつ、標準 GPT と競争力のあるパフォーマンスを発揮できることを示すこと（ただし、本論文はこれらの高度な応用は将来の課題であると述べています）。

著者らは、NRGPT をすべての GPT の代替としてではなく、エネルギーと最適化のレンズを通じてトランスフォーマーアーキテクチャを理解するための意味のある一歩として位置づけており、静的な深層学習モデルと動的システムの間のギャップを埋めています。