Equilibrium Matching: Generative Modeling with Implicit Energy-Based Models

著者 Runqian Wang, Yilun Du
所属 MIT, Harvard University
投稿日 2025年10月03日
カテゴリ cs.AI, cs.LG
arXiv 2510.02300v1 ↗

Equilibrium Matching: Generative Modeling with Implicit Energy-Based Models

基本情報

簡単に説明すると

この論文は、従来の拡散モデルやフローベース生成モデルとは根本的に異なるアプローチとして「Equilibrium Matching（EqM）」という新しい生成モデリング手法を提案している。
従来モデルが時間に依存する非平衡動態を学習するのに対し、EqMは時間不変の平衡勾配を学習し、暗黙的なエネルギー景観上での最適化ベースサンプリングを可能にする。
ImageNet 256×256でFID 1.90を達成し、既存の拡散・フローモデルを上回る性能を示している。
また、部分ノイズ画像のデノイジング、分布外検出、画像合成などの従来手法では困難なタスクを自然に処理できる独自の特性を持つ。
プロジェクトサイト（https://raywang4.github.io/equilibrium_matching/）で詳細が公開されている。

1. 研究概要

1.1 背景と動機

生成モデリングは拡散モデルとフローベース手法の発展により急速に進歩しており、これらは簡単なノイズ分布を複雑なデータ分布にマッピングする順方向ノイズ過程とその逆方向を学習することで最先端のサンプル品質を達成している。
しかし、これらのモデルは訓練と推論の両方で非平衡動態を採用しており、入力タイムステップに条件付けられて異なるノイズレベルでの異なる動態を学習する構造になっている。

この非平衡設計は実用上の制約を課している。
具体的には、ノイズレベルスケジュールの設定、サンプリング時の固定積分範囲といった問題がある。
また、時間条件付けを単純に除去したノイズ非条件拡散モデルは生成品質が劣化することが示されており、平衡動態学習の単純な試みは成功していない。

一方、エネルギーベースモデル（EBM）は直接的に平衡エネルギー景観を学習するが、訓練不安定性とサンプル品質の問題に悩まされている。
Energy Matchingのような最近のアプローチも別々の訓練段階を必要とし、大規模データセットでフローベース手法を超える性能を達成できていない。

このような背景から、著者らは非平衡動態から平衡動態への根本的な転換を図り、時間条件付けを排除しながらも高品質な生成を実現する新しいフレームワークの必要性を認識した。
鍵となる洞察は、ノイズ非条件拡散モデルの品質劣化が、目標勾配と平衡動態の非互換性に起因するという仮説である。

1.2 主要な貢献

本研究は、生成モデリングの分野における根本的なパラダイムシフトを提案し、理論と実践の両面で重要な貢献を果たしている。
Equilibrium Matchingは、拡散・フローモデルの時間条件付き非平衡動態を、暗黙的エネルギー景観上の単一時間不変平衡勾配に置き換える革新的なフレームワークである。

最も重要な技術的革新は、暗黙的エネルギー関数と整合する新しい目標勾配ファミリーの導入である。
従来手法の目標勾配が平衡動態と非互換であった問題を解決し、データサンプルでの勾配をゼロにする制約c(1)=0を明示的に課すことで、エネルギー景観構築を可能にしている。

  • 平衡動態による生成モデリング: 時間不変勾配場による単一平衡動態学習で、従来の複雑な時間条件付き学習を置き換え
  • 最適化ベースサンプリング: 学習された景観上での勾配降下法による柔軟なサンプリング、適応的ステップサイズとオプティマイザーの利用が可能
  • 暗黙的エネルギー景観の学習: 明示的エネルギー関数を定義せずに一貫したエネルギー構造を学習する新手法
  • 理論的保証の提供: データ多様体学習と多様体からのサンプリングに関する理論的正当化
  • 優れた生成性能: ImageNet 256×256でFID 1.90を達成し、既存の拡散・フローモデルを上回る性能

2. 提案手法

2.1 手法の概要

Equilibrium Matchingの中核概念は、データ多様体上で勾配が消失し、ノイズに向かって増加する時間不変勾配場の学習にある。
これにより、真のサンプルが静止点となる平衡景観を構築し、Flow Matchingが最終タイムステップでのみ真値に収束する変動速度を学習するのとは対照的に、常に真のデータポイントに収束する時間不変勾配景観を学習する。

訓練過程では、補間係数γ∈[0,1]を用いてデータとノイズ間の遷移を定義し、中間汚染サンプルx_γ = γx + (1-γ)εを生成する。
重要な点として、γは暗黙的であり、モデルには見えない。
目標は、これらの中間サンプルで暗黙的エネルギー景観と整合する目標勾配を定義することである。

勾配方向をノイズからデータへと設定し、Equilibrium Matchingの訓練目的を以下のように記述する：
$$L_{EqM} = (f(x_γ) - (ε-x)c(γ))^2$$

ここで、c(γ)は勾配の大きさを制御し、c(1)=0を明示的に課すことで、実サンプルでのエネルギー景観勾配消失を保証する。
目標勾配(ε-x)c(γ)は、ノイズからデータを指す方向(ε-x)と、データ多様体に近づくにつれて消失する大きさc(γ)を持つ。

2.2 技術的詳細

Equilibrium MatchingとFlow Matchingの根本的違いは、その理論的基盤にある。
EqMは正規化フローではなくEBM観点から導出されるため、異なる目標方向を持つ。
EqMは勾配ε-xを学習し、FMは速度x-εを学習する。
また、異なる基本制約を持ち、EqMはc(1)=0（データ多様体での局所最小値構築）、FMは∫₀¹c(γ)dγ=1（有効積分経路構築）を要求する。

勾配大きさ制御関数c(γ)について、複数の選択肢を提案している：

線形減衰: c_linear(γ) = 1-γ
自然な選択として、ノイズで高勾配を割り当て、真の画像に向かって線形減衰する。

切断減衰: c_trunc(γ) = 1 (γ≤a), (1-γ)/(1-a) (γ>a)
データから遠い場合は目標勾配を一定に保ち、データ多様体に近づくと線形減衰する。

区間線形: さらに複雑な区間線形関数も検討される。

明示的エネルギー学習のため、EqM-E変形も提案している。
内積変形では、追加の回帰頭h_ϕを用いてエネルギーの勾配との内積を学習：
$$L_{inner} = (f(x_γ) \cdot h_ϕ(x_γ) - ∥(ε-x)c(γ)∥²)²$$

L₂ノルム変形では、エネルギー勾配のL₂ノルムを直接学習：
$$L_{L2} = (∥f(x_γ)∥ - ∥(ε-x)c(γ)∥)²$$

2.3 新規性

Equilibrium Matchingの最も革新的な点は、生成モデリングにおける平衡と非平衡動態の概念的転換にある。
従来の拡散・フローモデルが時間依存の複雑な動態を学習するのに対し、EqMは単一の時間不変勾配場による優雅なソリューションを提供する。

最適化ベースサンプリングは、従来の固定軌道積分とは根本的に異なるアプローチである。
予め決められた軌道に沿って積分する既存手法に対し、EqMは学習された景観上での自由な最適化を可能にする。
これにより、Nesterov加速勾配法などの勾配最適化技術の自然な採用が可能になり、より良い生成品質を達成できる。

推論時適応計算も重要な新規性である。
EqMは各サンプルに対して独立してサンプリングステップを調整でき、勾配ノルムに基づいて最大60%の関数評価を節約できる。
これは固定ホライゾンを持つ従来手法では実現困難な特性である。

理論的観点では、データ多様体学習とサンプリングに関する保証を提供している。
平滑性仮定の下で、勾配降下がエネルギー最小値（データポイント）に収束することを証明し、生成モデリングに対する堅固な理論的基盤を提供している。

従来のEBMとの相違点として、EqMは訓練不安定性の問題を回避し、明示的エネルギー関数なしでも一貫したエネルギー構造を学習できる。
また、Energy Matchingのような多段階訓練を必要とせず、単一段階でのエンドツーエンド学習を実現している。

3. 実験結果

3.1 実験設定

評価は主に4つの観点から実施されている。
第一に、ImageNetでの生成品質比較、第二に勾配ベースサンプリング手法の特性と性能検証、第三に勾配景観の有効性に関するアブレーション研究、第四にEqMの独自特性の実証である。

主要な実験はクラス条件付きImageNet 256×256画像生成で行われ、StyleGAN-XL、VDM++、DiT-XL/2、SiT-XL/2との比較が実施されている。
評価指標はFID（Fréchet Inception Distance）を主に使用し、生成品質の定量的評価を行っている。

スケーラビリティ評価では、訓練長、モデルサイズ、パッチサイズの3軸でスケーリング特性を検証している。
モデルサイズはS、B、L、XLの4段階、パッチサイズは8、4、2で評価されている。

アブレーション研究では、異なるc(γ)関数の比較、勾配乗数λの調整、ノイズ条件付けの有無、明示的エネルギー変形の比較が行われている。
すべての実験で適切な統計的有意性とベースライン比較が確保されている。

3.2 主要な結果

EqMは既存手法を大幅に上回る生成性能を達成している。
ImageNet 256×256での評価において、EqM-XL/2はFID 1.90を記録し、これまでの最良結果であるSiT-XL/2の2.06を大幅に改善している。
この結果は、StyleGAN-XL（2.30）、VDM++（2.12）、DiT-XL/2（2.27）といった他の最先端手法もすべて上回っている。

スケーラビリティ検証では、EqMが全ての軸でFlow Matchingを一貫して上回ることが確認されている。
訓練エポック数、パラメータ数、パッチサイズのすべての設定において、EqMは有意なマージンでFMベースラインを超える性能を示している。
これは、EqMの強いスケーリングポテンシャルと将来性を示唆している。

最適化ベースサンプリングの優位性も実証されている。
Nesterov加速勾配法（NAG-GD）を用いることで、従来のODEサンプラーやSDEサンプラーを上回る性能を達成している。
具体的には、Euler ODEサンプラーで1.93、標準勾配降下で1.93、NAG-GDで1.90のFIDを記録している。

視覚化結果では、EqMがFMよりも高速に収束することが確認されている。
また、生成サンプルの最近傍分析により、EqMが訓練データの暗記に依存せず、未見サンプルに対する汎化能力を持つことが示されている。

3.3 既存手法との比較

拡散モデルとの比較では、EqMは時間条件付けの複雑さを排除しながら優れた性能を達成している。
DiT-XL/2やVDM++といった最先端拡散モデルと比較して、EqMはより簡潔なアーキテクチャでありながら生成品質で勝っている。

Flow Matchingとの直接比較では、理論的基盤の違いが実践的な利点に転換されていることが明確である。
FMが複雑な時間依存動態を学習する必要があるのに対し、EqMの単一平衡動態はより効率的で安定した学習を可能にしている。

Energy-based modelとの比較では、EqMが従来のEBMの訓練不安定性問題を回避しながら、エネルギー構造の利点を享受していることが重要である。
明示的エネルギー関数なしでも一貫したエネルギー景観を学習できることは、実用的な大きな利点となっている。

計算効率の観点では、EqMの適応計算機能により、固定ステップ手法と比較して大幅な計算削減が可能である。
勾配ノルムベースの早期停止により、品質を犠牲にすることなく推論時間を短縮できている。

サンプリング柔軟性では、EqMが様々な最適化手法を自然に組み込めることが大きな利点となっている。
従来の積分ベース手法では困難な適応的ステップサイズや加速手法の利用が可能である。

4. 実用性評価

4.1 実装の容易性

Equilibrium Matchingの実装は、概念的にはシンプルでありながら、いくつかの技術的考慮が必要である。
訓練目的関数の変更は比較的単純で、既存のフローマッチング実装から容易に適応できる。
主要な変更点は、目標勾配の方向転換（x-ε から ε-x）と、勾配大きさ制御関数c(γ)の導入である。

しかし、最適な性能を得るためには、c(γ)関数の選択と勾配乗数λの適切な調整が重要である。
実験結果によると、切断減衰（a=0.8）と勾配乗数λ=4の組み合わせが最良の性能を示しており、これらのハイパーパラメータチューニングが実装成功の鍵となる。

明示的エネルギー変形（EqM-E）の実装は、追加的な安定性考慮が必要である。
特にL₂ノルム変形は初期化に敏感であり、事前訓練されたEqMモデルからの初期化が推奨される。
一方、内積変形はより安定しており、実用的な選択として推奨される。

4.2 計算効率

EqMの計算効率は、従来手法と比較して複数の利点を提供している。
最も重要な特徴は、適応計算機能による推論時の効率化である。
勾配ノルムに基づく早期停止により、品質を維持しながら最大60%の関数評価を削減できる。

訓練時の計算オーバーヘッドは最小限である。
目標勾配の計算は既存のフローマッチングと同等の複雑さを持ち、追加的な計算コストは無視できる程度である。
むしろ、時間条件付けの除去により、モデルアーキテクチャが簡略化される場合もある。

最適化ベースサンプリングは、従来の固定ステップ積分と比較して柔軟性を提供する。
Nesterov加速勾配法などの高度な最適化手法の利用により、同じ計算予算でより高品質なサンプルを生成できる。

メモリ使用量の観点では、時間埋め込みの除去により若干の削減が期待される。
また、適応サンプリングにより、バッチサイズや並列化戦略の最適化も可能である。

4.3 応用可能性

EqMは従来の生成モデルでは困難または不可能な多様な応用を可能にする。
最も注目すべき特性は、部分ノイズ画像の直接処理能力である。
従来の拡散・フローモデルが明示的なノイズレベル入力を必要とするのに対し、EqMは任意のノイズレベルから開始して高品質な生成を行える。

分布外（OOD）検出機能は、追加のモジュールなしで実現される重要な特性である。
学習されたエネルギー景観により、分布内サンプルは低エネルギー、分布外サンプルは高エネルギーを持つ傾向があり、効果的なOOD検出を可能にする。

画像合成機能は、EBMと同様の原理により複数モデルの組み合わせを可能にする。
異なるクラス条件でのモデルを加算することで、概念的に意味のある合成画像を生成できる。
これは拡散モデルでは複雑な実装を要する機能である。

研究応用では、エネルギー景観の直接分析が可能になることで、新たな解釈可能性研究の道が開かれる。
また、最適化理論との自然な結合により、生成モデリングと最適化の交差領域での発展が期待される。

商用応用では、適応計算による効率化とハードウェア要件の柔軟性が重要な利点となる。
品質と計算コストのトレードオフを動的に調整できることは、実用的なデプロイメントにおいて大きな価値を持つ。

5. まとめと所感

5.1 論文の意義

この研究は、生成モデリング分野における根本的なパラダイムシフトを実現した極めて重要な貢献である。
拡散モデルとフローベース手法が支配的な現状において、平衡動態による全く新しいアプローチを提示し、理論と実践の両面で優位性を実証している。

最も重要な理論的貢献は、非平衡動態から平衡動態への概念的転換とその数学的定式化である。
従来の時間条件付き複雑な動態学習を、時間不変の単一勾配場学習に置き換えることで、生成モデリングの本質的な簡素化を達成している。
この転換は、エネルギーベースモデルと正規化フローの橋渡しとしても機能し、両分野の統一的理解に寄与している。

実証的な成果として、ImageNet 256×256でのFID 1.90という記録的性能は、提案手法の実用的価値を明確に示している。
特に重要なのは、すべてのスケールで一貫した優位性を示していることで、手法の堅牢性と将来性を裏付けている。

技術的革新として、最適化ベースサンプリングの導入は生成モデリングに新たな次元を加えている。
従来の固定軌道積分から脱却し、適応的最適化手法の利用を可能にすることで、推論時の柔軟性と効率性を大幅に向上させている。

産業的インパクトとして、計算効率の改善と独自機能（OOD検出、画像合成等）の提供は、実用的なAIシステムの発展に直接貢献する。
特に、適応計算による60%の関数評価削減は、大規模デプロイメントにおいて重要な経済的価値を持つ。

5.2 今後の展望

Equilibrium Matchingが開拓した新しい研究領域は、多方面での発展可能性を秘めている。
短期的には、異なるデータドメイン（音声、テキスト、3D）への拡張が期待される。
平衡動態の概念は画像生成に限定されず、より広範な生成タスクへの適用が可能である。

勾配大きさ制御関数c(γ)の設計は、さらなる研究が期待される重要な領域である。
現在の線形・切断減衰を超えて、学習可能な関数や適応的な関数の開発により、性能のさらなる向上が期待される。
また、タスク特有の制約や事前知識を組み込んだ制御関数の設計も有望な方向性である。

明示的エネルギー学習の改善も重要な課題である。
現在のEqM-E変形は暗黙的バージョンを下回る性能を示しているが、より安定した訓練手法や新しいエネルギー定式化により、明示的エネルギーの利点を活用できる可能性がある。

最適化ベースサンプリングの発展により、生成プロセスの細粒度制御が可能になる。
制約付き最適化、多目的最適化、確率的最適化手法の組み込みにより、より複雑で制御可能な生成が実現されるだろう。

長期的には、EqMの平衡動態学習が物理系シミュレーション、分子生成、材料設計などの科学計算分野への応用も期待される。
自然系の平衡状態学習という観点から、科学的発見を支援するAIツールとしての発展も考えられる。

理論的深化として、より一般的な設定での収束保証、最適化景観の幾何学的性質、他の生成手法との理論的関係の解明などが重要な研究課題となる。

ただし、現在の限界として、明示的エネルギー学習の困難さ、ハイパーパラメータ調整の複雑さ、特定ドメインでの性能検証不足などが挙げられる。
これらの課題解決が、EqMの広範な普及と実用化の鍵となるだろう。