Masaki Saito

Energy-Based Model（EBM）を大規模化し、さまざまな改良を加えることで、少ない計算量で高品質な画像生成を実現するTransformer Energy-Based Model（TEM）が提案された。 “Scaling Transformer Energy-Based Models for High Quality Image Synthesis” 609MパラメータのTransformerを用いて、ImageNet 256×256のクラス条件付き画像生成でFID 3.82を達成した。これはClassifier-Free Guidance（CFG）を使用しないDiffusion Transformerと競争できる水準である。 EBMは、データの確率分布をエネルギー関数E_θ(x)で表現するモデルである。 p_θ(x) = exp(−E_θ(x)) / Z_θ EBMはエネルギー関数を直接学習し、その勾配を使ってMCMCでサンプリングする。 x_{t+1} = x_t − η∇E_θ(x_t) + σε ランジュバンサンプリングと同様に、エネルギーの低い領域に向かって状態を更新しながら、ノイズを加える。 EBMの強みとして、分布外検出ができること（エネルギーが高いほどモデル上の相対的な確率密度が低いこと）、補完ができること、そして複数の確率分布の積に対応する分布からサンプリングできることが挙げられる。 例えば、2つの確率分布p₁(x)、p₂(x)の積に比例する分布を生成したい場合、 p(x) ∝ p₁(x)p₂(x) E(x) = E₁(x) + E₂(x) となる。そのため複数の制約を保った上でのサンプリングは、それらの制約に対応するエネルギーを足し合わせ、その勾配を使ってサンプリングするだけでよい EBMは従来から研究されてきた手法だが、高次元空間でのMCMCサンプリングが難しいことや、大規模モデルの学習が不安定になるという問題があった。 TEMでは、主に以下の5つの工夫によってこれらの問題に取り組んでいる。 ・VAEの潜在空間でエネルギー関数を学習 元の画像空間ではなく、VAEによって変換された32×32×4の潜在変数空間でエネルギー関数を学習し、MCMCサンプリングを行う。これにより、サンプリングする空間の次元数を大幅に削減する。 ・Transformerの複数層からエネルギーを抽出 Transformerの浅い層、中間層、深い層からそれぞれエネルギーを抽出し、それらの和を最終的なエネルギーとする。 浅い層由来のエネルギーはテクスチャ、中間層は物体の配置、深い層は大域的な意味構造に対応する傾向があり、異なる抽象度の制約を組み合わせることで生成品質を改善する。 ・Positive Representation Alignment（PRA） 事前学習済みのDINOv2-Largeの特徴表現と、Transformerの中間表現を一致させる学習を、Positiveサンプル（実データ）に対してのみ適用する。 Negativeサンプルはデータ分布から外れている場合が多く、その特徴表現まで一致させようとすると、本来の対比学習と矛盾する可能性があるためである。 ・Sampler Matching Loss EBMの最尤学習では、実データのエネルギーを下げ、モデルが生成したサンプルのエネルギーを上げるように学習する。 しかし、実際には有限ステップのMCMCで生成される分布q_θは、エネルギー関数が定義する本来の分布p_θと一致しない。 そこで、KL(q_θ || p_θ) に由来する正則化を追加し、エネルギー関数側を固定したうえで、 E_q[E_θ(x)] − H(q_θ) を最小化するようにする。 1つ目の項は、MCMCで生成したサンプルのエネルギーを下げる項である。エネルギー関数の直接の勾配ではなく、サンプリング過程を通じて勾配を流すことで、生成サンプルがより低エネルギーの領域に到達するように学習する。 この計算には二階微分が必要となるが、40ステップすべてを逆伝播するのではなく、最後の1ステップだけを微分することで計算量を抑えている。 2つ目の項は、生成サンプルのエントロピーを最大化することに対応する。Replay Bufferに保存された過去のサンプルとの最近傍距離を利用し、サンプル同士が近すぎないようにする。 これによって、低エネルギー領域に向かうだけでなく、生成サンプルの多様性も維持する。 ・Persistent Replay Buffer 毎回ガウスノイズから生成を開始するのではなく、過去のMCMCで生成した潜在変数を保存して再利用する。 Replay Bufferの容量は16,000で、99%の確率で過去のサンプルを再利用し、1%の確率で新たにガウスノイズからサンプリングを開始する。 また、クラス条件を毎回独立に与えるため、あるクラスの生成途中の潜在変数を、別のクラスの生成に利用することもできる。 学習時には40ステップ、推論時には主に60ステップのサンプリングを行う。 これらの工夫の中でも、特にPRAとSampler Matchingの効果が大きいとみられる。 実際、122Mパラメータのモデルによるアブレーションでは、Contrastive Lossのみの場合のFID 121.11が、Sampler Matchingの導入によって42.76まで改善し、さらにPRAを加えることで19.86まで改善している。 また、エネルギー関数を組み合わせた画像生成、追加学習なしのInpainting、エネルギー値による分布外検出など、EBM本来の特徴も実証している。 コメント === EBMは理論的に柔軟で興味深いモデルであり、非正規化確率密度を直接評価できるなどの利点がある。一方、これまではDiffusionやFlow Matchingほど高い画像生成品質を実現することが難しかった。 今回、さまざまな工夫を組み合わせることで、EBMでも高品質な画像生成が可能であることが示された。 この中で個人的に一番興味があるのはEBMでは、複数の確率分布の積に対応する分布を、エネルギーの加算によって自然に表現できる点である。 例えば、複数の制約を同時に満たすような生成を比較的簡単に実現できるだけでなく、それぞれのエネルギー項を調べることで、生成したものがどの制約をどの程度満たしているかを評価する手がかりにもなる。 ロボット制御や化合物生成などで複数の制約を考慮するといったことが自然に実現できる。 また、今回の手法でPersistent Replay Bufferの役割も注目している。 DiffusionやFlow Matchingは、多峰性を持つ複雑な分布からの生成という問題に対して、ガウス分布から目標分布への経路を学習することで対処している。 一方、EBMでは多峰性を持つエネルギー地形におけるMCMCのMixingの難しさがあったが、これをPersistent Replay Bufferが、学習過程を通じて分布の輸送を実現している可能性がある。 学習初期にはReplay Bufferのサンプルはガウス分布から初期化されるが、学習が進むにつれてエネルギー関数自体が変化し、それに従ってBuffer内のサンプルも徐々に移動していく。 つまり、エネルギー関数が徐々にデータ分布を表現するようになる過程で、Replay Bufferの分布も単峰性のガウス分布から多峰性のデータ分布へと変化していくと考えられる。 このように考えると、DiffusionやFlow Matchingが生成時にガウス分布からデータ分布への輸送を行うのに対し、今回のEBMでは学習過程そのものが輸送の役割を果たしているとも解釈できる。 特に、学習初期にはまだエネルギー障壁が十分に形成されていないため、サンプルがさまざまなモードへ分かれていき、その状態がReplay Bufferに維持されることで、最終的なエネルギー地形におけるMixingの難しさを緩和している可能性がある。 これは今回の論文では直接検証されていないが、興味深い仮説である。