【DL輪読会】Energy-Based Transformers are Scalable Learners and Thinkers

DEEP LEARNING JP [DL Papers] Energy-Based Transformers are Scalable Learners and Thinkers Fumiya Uchiyama, Matsuo Lab http://deeplearning.jp/ 1

• http://deeplearning.jp/

書誌情報 ▪ 論文名 Energy-Based Transformers are Scalable Learners andThinkers ▪ 著者 Alexi Gladstone, Ganesh Nanduru, Md Mofijul Islam, Peixuan Han, Hyeonjeong Ha, Aman Chadha, Yilun Du, Heng Ji, Jundong Li, Tariq Iqbal ▪ 出典 arXiv（2507.02092） PDFに査読中と表記されている（2025/9/15） ▪ サイト https://energy-based-transformers.github.io/ ※特に明記のない限り引用番号は論文に従う。また特に断りのない限り画像は論文からの出典である。 2

• https://energy-based-transformers.github.io/

概要 ▪ 貢献 – EBMをTransformerで実装したEnergy-Based Transformer(EBT)を提案 • EBMをスケーラブルにして言語モデリングや動画予測タスクなどでの性能を確認 – EBTの性能を向上させるエネルギー地形正則化手法を特定 ▪ 実験的結果 – System 2的タスクにおいてスケーラビリティを確認 – Out of Dist.（OoD）に対する強力な性能 3

イントロ| System1 / 2 心理学における人間の思考の分類 ▪ System 1 (速い思考) シンプルな問題に対して、過去の経験に基づき、速く直感的 ▪ System 2 (遅い思考) 複雑な問題に対して、慎重で解析的であり、論理的推論を要する 現在のモデルはSystem 1的タスクが得意であるが[7]、System 2的タス クについては不得意[8-10] System 2をあらゆる対象で実装するには？ 4

イントロ|「推論モデル」はSystem 2が可能か？ 近年、o1やR1などの「推論モデル」は思考時間を増やすことで数学や コーディングでの性能が向上している これではいけないのか？ 先行研究から示唆される現在の（強化学習ベース）推論モデルの問題点 ▪ ルールベースの報酬が不可欠であり、System 2が適用可能なタスク の範囲が狭まっている[12] ▪ 新しい推論パターンを誘発しているのではなく、ベースモデルの学習 した推論パターンの確率を高めているだけであり、探索の必要がある タスクでの性能に限界がある[18] 5

イントロ| 「反復的処理のあるモデル」はSystem 2が可能か？ 反復的な推論が、System 2実現の鍵ではないかという見方がある ▪ RNN 再帰的な状態更新 新しい情報が入るタイミングと状態更新のタイミングが同じ ▪ Diffusion 逆拡散過程でタイムステップが増加するほど精度向上 訓練時以上のタイムステップを課しても性能に寄与しない さらに、両者ともに予測の検証メカニズムが外部依存であったり明白で ないという問題がある 動的に処理することが前提で、あらゆるタスクに適用可能な検証メカニ ズムを有することができないか？ 6

イントロ| 本研究のリサーチクエスチョン あらゆる問題に対して、教師なし学習でSystem 2を獲得できるか？ このような能力があれば、現在のSystem 2アプローチをあらゆる問題、 あらゆるモダリティに一般化することができ、外部の人間、報酬、その 他教師信号への依存を回避できる 人間のSystem 2思考の特性と現在のモデルを比較すると、いくつかの 重要な違いがあると主張： 1. 計算の動的割り当て 2. 連続状態空間における不確実性のモデリング 3. 予測の検証 7

イントロ| 人間のSystem 2思考の特性との差分 1. 計算の動的割り当て タスクの難易度に応じた労力の割り当て e.g. キャリアの意思決定と昼食に何を食べるかの意思決定 2. 連続状態空間における不確実性のモデリング 現実世界には本質的に予測不可能な要素が多い ・LLMはトークンレベルの確率を通じてシミュレート可能[33]。 ・視覚などの連続状態空間では、一般のTransformer、RNN、 Diffusionは非対応（参照：離散化スキームや代理損失[34-35] ） ・Energy-Based Model（EBM）は相対的な尤度[42]をモデル化 3. 予測の検証 予測の検証は予測の生成よりも指数関数的に簡単[48] 予測の品質がわかると動的に推論時間を調整できる 8

イントロ| Energy-Based Models (EBMs) ▪ 思考を、学習した検証器に関する最適化手続きと解釈 ▪ 入力と予測のペアの空間についてエネルギー地形を学習させるために EBMを訓練したい – 低いエネルギーが高い確度（不確実性の評価/検証） – ランダムな予測を、エネルギーの地形を下ることで洗練させられる 9

イントロ| Energy-Based Models (EBMs) ▪ 低いエネルギーが高い確度（不確実性の評価/検証） – 検証は生成よりも容易である – 今日のパラダイムである直接生成を学習するのではなく、学習した検証器（エ ネルギー関数）に沿って予測を最適化する – 生成器が検証器の勾配で暗黙的に定義されスケーラビリティと敵対問題が解決 10

手法 | Energy-Based Transformers (EBTs) ▪ 既存のEBMはスケーラビリティが弱い[51] – EBMの学習安定性[51-54] – 長い訓練時間[53,54] ▪ そこで、EBMに特化したTransformer、Energy-Based Transformerを実装する 本研究では以下の２つのモデルを開発した – GPTライクなデコーダ型EBT – BERTやDiffusion Transformers（DiT）ライクな双方向EBT 11

手法 | EBMのBackground ▪ 確率的EBMにおいて、予測確率はボルツマン分布を用いて規定される ▪ 正規化項を無視することで、データ多様体のあるところに低いエネル ギーを、そうでないところに高いエネルギーを割り当てる問題に帰着 12

手法 | EBMをスケーラブルに学習する ▪ EBMの学習を最適化問題として定式化[48,71] – エネルギー地形を暗黙的に正則化することで次元の呪いを回避 – EBMを勾配降下法によって初期予測値を真の解に最適化するように学習 • エネルギー地形は真の解を中心に凸状となり、結果としてエネルギー地形は真のデータ多 様体上のみで低エネルギー状態を保つように正則化される ▪ EBMの学習に対する３つの工夫（エネルギー地形正則化手法） – リプレイバッファ[48,51で既出] – 予測の更新にランダムノイズを添加 • ランジュバン力学的な摂動を最適化時に加える • 探索される領域を広げる – 勾配降下のステップサイズとステップ数のランダム化 • これにより学習中のパスを変える 13

実験 ▪ 事前学習データにどの程度迅速に適合できるか？ – EBMをスケーラブルにしたEBTがどのくらいスケールするか ▪ 思考をスケーリングさせるとモデルの性能がどのように変化するか？ – 既存のTransformerとくらべてSystem 2思考ができるか？ 14

実験 | 自己回帰型言語モデリング ▪ 学習データ RedPajama v2の100Bサンプル ▪ トークナイザ GPT-NeoX ▪ 評価方法 – 論理推論などを要するベンチマークのテキストについてのPerplexity（PPL） – タスクのaccuracyなどと比べて線形に結果が表れやすいから ▪ 比較対象 – Transformer++ • LLaMa2など近年のLLM学習レシピ • arXivで出てくる「Transformer++」[arXiv 2003.04974]とは異なることに注意 15

実験 | 自己回帰型言語モデリング ▪ EBTは複数の軸でTransformer++より速くスケールする 16

実験 | 自己回帰型言語モデリング ▪ エネルギー地形正則化手法を組み合わせると思考と検証の効果が出る – System 2思考によるPPLの改善率が記載されている – “Thinking Longer”はエネルギー勾配降下による予測最適化 – “… and Self-Verification”は上記に加えBest-of-N（BoN）予測の採用 – ランジュバン力学（ノイズ添加）の正則化手法を除去すると、エネルギー地形 の探索領域が減少するため、単一経路のパフォーマンス（Thinking Longer） は向上するもののSelf-Verificationのパフォーマンスは低下 17

実験 | 自己回帰型言語モデリング ▪ 最適化数を増やすことでOoDに対するPPLを下げられる ▪ 事前学習の計算量で検証性能がスケールされる 18

実験 | 自己回帰型言語モデリング ▪ OoDデータほど思考の効果がある ▪ OOD Magnitude Shiftは事前学習データのPPLと対象データのPPLの 比率 ▪ OoDデータほど、思考と検証でのPPL改善率が高い 19

実験 | 自己回帰型言語モデリング ▪ 難しい単語はエネルギーが小さくなる – “is”, ”a”などは低いエネルギーの一方で、”quick”, “brown”などは高い 20

実験 | 自己回帰型動画予測タスク ▪ 言語と同様スケールする（左） ▪ シーンによってエネルギーが上下する（右） 21

実験 | 双方向型画像デノイジング ▪ OoDデータに対して思考することで性能が大幅に改善 ▪ 線形プルービング（Imagenet-1k）でも良好な結果 ▪ DiT（ベースライン）と比べて非常に少ないステップでデノイズ可能 22

結論 ▪ EBTは連続型（画像）・離散型（テキスト）双方のモダリティにおい て、Transformer++を上回る速度でスケールする初の事例 ▪ EBTのSystem 2思考の能力はOoDデータに対する優れた汎化性能を 実現する ▪ 将来の基礎モデルの性能向上に向けた有望なパラダイムシフト？ 23

個人的感想 ▪ スケーラビリティの結果には要注意 – Figure5の結果はほぼすべてTransformer++の方が良い – 小パラメータサイズでスケール率がわずかに高いからといって性能差が逆転す るかは不明 – スケールの速度が速いといっても絶対的なperplexityの差はそこまでないよう に見える ▪ 難易度に合わせて動的に思考時間を調整することは今後の課題 ▪ 関連すると感じた研究 – arXiv 2502.05171 Test Time Scalingの初期値の問題が議論されている（リプレイバッファ） – AKOrN [ICLR 2025] エネルギー関数としての条件を緩和することで性能向上 参考：https://x.com/hillbig/status/1941657099567845696 24

• https://x.com/hillbig/status/1941657099567845696