Diffusion・Flow Matchingの統一的理解とone-step生成への考察

はじめに

画像生成モデルの研究は急速に進展しており、Diffusion ModelやFlow Matchingといった手法が高レベルな品質を達成しています。特にDiffusionベースのモデルはその高精度な生成能力により実用化されている多くのモデルでの採用されています。しかしながらDiffusionモデルは推論時に多数のステップを要するという問題を抱えており、推論時の計算コストと画像の品質は画像生成モデル研究における中心的な問題の一つと言えます。
本記事では、以下の大枠に沿って近年の画像生成モデルの知見に関する考察を行っていきます。

• Diffusion / Flow Matching における統一的な設計パターンの議論
• x-prediction を採用するモチベーション
• FD-loss によるポストトレーニングと、それによる one-step repurposing

1. 潜在空間生成と pixel-space 生成 の動機

現在実用化されている高品質な画像生成モデルの多くは、ピクセル空間ではなく 潜在空間 上で拡散・フロー過程を学習しています。典型的には、画像 $x$ を エンコーダー によって latent $z$ に圧縮し、その 潜在空間上で Diffusion / Flow Matching を行い、最後に デコーダー によって画像空間へ戻す構成です。

潜在空間 を用いる利点は明確です。ピクセル空間と比べて次元が小さく、計算コストを大幅に削減できます。また、多様体仮説の観点から、自然画像の低次元構造をより扱いやすい空間に押し込められる(エンコーダがそのような低次元空間に落とし込んでくれる)という期待もあります。

一方で潜在空間を用いるスキームには原理的な代償もあります。デコーダー を通す以上、そこに伴う誤差の蓄積 は避けられません。また、画像の詳細な特徴、高周波成分、局所的な構造の忠実度は、潜在空間と デコーダー の表現能力に制限されます。したがって、pixel-space 生成 には、そのような根源的な問題 を回避できるという動機があるといえます。
また、画像の高周波成分をうまく扱うための研究は様々存在しますが、画像空間を直接扱う試みはその問題に対するアプローチの一つと考えても良いでしょう。

---

1.1. Prediction Space × Loss Space: デザインパターンとしての整理

ここではまず、Li et al.[1] にて提案された考え方について紹介します。Diffusion / Flow Matchingモデルの設計には、「何を予測するか（Prediction Space）」と「何を損失にするか（Loss Space）」という2つの自由度に基づいて設計を整理するというものです。
説明のため、noise-to-data convention として

$$z_t=(1-t)ϵ+tx,t\in [0,1],$$

を考えます。ここで $x\sim p_{\mathrm{data}}$ 、 $ϵ\sim \mathcal{N}(0,I)$ であり、 $t=0$ が noise、 $t=1$ が data に対応します。自然に考えられるtarget parameterization は次の三つです。

• $x\text{(clean data / x-prediction)}$

• $ϵ\text{(noise / epsilon-prediction)}$

• $v=x-ϵ\text{(velocity / v-prediction)}$

これらは $z_t=(1-t)ϵ+tx$ の線形関係から相互に変換できます。(下表は[1]より引用)

したがって、無限容量・完全最適化・数値誤差なしという理想化のもとでは、どの parameterization を使っても同じ情報を表しているように見えます。しかし、実際にtrainingしてみると、その性能はこの設計選択に大きく依存するということが[1]での主張の肝と言えます。

---

1.2. x-prediction の有効性

近年の議論で重要なのは、x-prediction が単なる parameterization の一つではなく、高次元画像空間において実質的に学習しやすい可能性があるという点です。これも[1]において議論されています。

自然画像は高次元ピクセル空間に存在しますが、その分布は低次元的な構造を持つと考えられています。もしデータが高次元空間に埋め込まれた低次元 manifold 近傍に集中しているなら、clean data $x$ はその manifold 上の量です。一方、 $ϵ$ や $v$ は、データ manifold に対して off-manifold な方向成分を多く含む可能性があります。

この観点では、x-prediction は「データらしい点そのもの」を予測するタスクであり、少なくとも高次元埋め込み設定では、ノイズや速度を直接予測するよりも構造を捉えやすいです。実際、スパイラルなどの低次元構造を高次元空間へ埋め込む toy setting では、埋め込み次元が上がるほど v-prediction / $ϵ$ -prediction は崩れやすく、x-prediction は構造を保ちやすいという挙動が観察されます。

この主張は厳密な定理というほどでもないのですが、理論的に筋が通っているように聞こえますし直感的な納得感もあるので、VAEの計算オーバーヘッド分速くなる、などの副次的なメリットも含め高解像度な画像生成を目指す上で良いアプローチのように思えます。

2. pixel Mean Flow：one-step 生成 と x-prediction の接点

この文脈で興味深いのが pixel Mean Flow (pMF) です。pMF は、pixel-space で one-step 生成 を目指すモデルであり、直感的には MeanFlow的な平均速度場 $u$ を考えるという制約のもとでx-prediction とiMF[2]的なv-lossを組み合わせた設計 として理解できます。

設計の数学的に込み入った部分にはここでは立ち入らないことにします(必要に応じて[3]を参照して下さい)が、先程の直感と以前のデザインパターンの議論を踏まえると次の2点に集約されます。第一には、x-prediction によって高次元画像空間における on-manifold target を直接扱います。第二に、MeanFlow の速度制約によってone-step transport としての整合性を与えます。
pMFはlatent freeであることによる計算コストの低さも主張の軸の一つにしていますが、これは先程副次的なメリットとして紹介した「VAEの計算オーバーヘッドの回避」と、x-predictionによる画像のon-manifoldな回帰がちょうどよく噛み合ったモデルと言えるでしょう。

ちなみに先程のJiTの設定で行ったtoy実験はpMFの設定でも主張に矛盾しない結果を得られます。先程の50NFEの設定よりは多少精度は劣るように見えますが、それでも一貫してx-predictionが高次元設定で優れていることがわかります。

---

2.1. pMF の再現検証

インターン期間中には、pMF の検証実験を行いました。ImageNet 256×256 に対し、論文では バッチサイズ1024で320 epochまで行われていたのに対しバッチサイズ256 の設定で行いました。

epoch	40	80	120	160	200	240	280
FID ↓	6.83	5.33	4.72	4.43	4.21	3.997	3.934
IS ↑	191.8	212.1	226.1	229.8	233.2	236.8	237.0

諸都合により論文よりもかなり少ないepochで打ち切っていますが、少なくとも再現が取れそうなことは確認できたと言えるでしょう。下の画像は獲得されたモデルによる出力画像の例です。

3. FD-lossによるポストトレーニング

ここからはYang el al.[4]において提案された、FD-lossというものについての考察と検証になります。
中心的なアイデアは、長らく評価指標として採用され続けてきたFIDを損失に組み込むことで、アイデア自体はシンプルかつ過去に類似の研究があるとのことですが、FID計算には多くのサンプルが必要なため実用的でないとされてきました。

3.1. 評価指標としてのFID

FD-loss の話に入る前に、FID の定義を正確に整理しておきます。FID は実画像と生成画像のfeatureについての距離であり、定義は次のようになっています。

$$\varphi (x_{\mathrm{real}})\sim \mathcal{N}({\mu }_r,{\mathrm{\Sigma }}_r),\varphi (x_{\mathrm{gen}})\sim \mathcal{N}({\mu }_g,{\mathrm{\Sigma }}_g)$$

としたとき

$$\Vert {\mu }_r-{\mu }_g{\Vert }_2^2+\mathrm{Tr}({\mathrm{\Sigma }}_r+{\mathrm{\Sigma }}_g-2({\mathrm{\Sigma }}_r{\mathrm{\Sigma }}_g{)}^{1/2})$$

これが Inception feature 上で計算されたとき、Fréchet Inception Distance (FID) と呼ばれます。

これは最適輸送コストの特殊な場合として説明されます。

最適輸送コストとして $\mathcal{X}$ 上の距離関数 $\mathcal{X}\times \mathcal{X}\to \mathbb{R}$ についてコスト関数を $C(x,y)=d(x,y{)}^p$ と定義したとき、

$$W_p(\alpha ,\beta ):=\text{OT}(\alpha ,\beta ,C{)}^{1/p}={(\underset{\gamma \in \mathrm{\Gamma }(\mu ,\nu )}{inf}{\int }_{\mathcal{X}\times \mathcal{X}}d(x,y{)}^{p}d\gamma (x,y))}^{1/p}$$

をp-Wasserstein距離と呼び、 $p=2$ の時の $W_2$ はそもそも有限な値として定まります。その上で両分布がガウス(平均と共分散以外の情報を落とす)であると仮定すると、最適輸送が陽に解けて先程の閉形式で書けることが知られています。細かい証明は標準的な最適輸送の図書等を参考にしていただければと思います。

---

3.2. FD-loss：評価指標から distribution-level training objective へ

FID は長らく評価指標として使われてきました。すると自然な問いが生じます。評価に使う距離を、そのまま training objective として最適化できないか、という問いです。

この問いに対する障害は、主に統計的スケールですであり、FD / FID を安定に推定するには、典型的には数万枚規模の生成サンプルが必要になります。一方で、各 step で数万枚すべてに対して gradient を流すことは実用的ではありません。small batch のみから FD を推定すると、特に高次元 feature の covariance 推定が不安定になり、loss として使うにはノイズが大きすぎるため、過去のこれを試みた研究ではかえって性能が下がることが報告されてきました。

FD-loss の核心は、この問題を population size for estimation と batch size for gradient computation の分離によって解く点にあります。

FD-loss では、生成画像を feature extractor $\varphi$ に通し、生成 feature の平均と共分散 ${\mu }_g,{\mathrm{\Sigma }}_g$ を推定します。ただし、この統計量は現在のバッチのみではなく、queue や EMA によって過去 batch の feature も用いて推定します。一方、gradient は current batch にのみ流し、queue や EMA に保持された過去 feature は detach します。
何らかのトリックによりFD計算を実用的な形で損失に組み込めるようにした、と考えていただければ本記事の議論においては十分かと思います。

形式的には

$${\mathcal{L}}_{\mathrm{FD}}={\mathrm{FD}}_{\varphi }(({\mu }_g,{\mathrm{\Sigma }}_g),({\mu }_r,{\mathrm{\Sigma }}_r))$$

を最小化します。ただし $({\mu }_r,{\mathrm{\Sigma }}_r)$ は real data feature から事前計算される参照統計であり、 $({\mu }_g,{\mathrm{\Sigma }}_g)$ は生成 feature の queue / EMA により推定されます。

この設計は、分布の統計量だけを見て全体の形を整えるような調整を行うpost-trainingのような形で理解ができるでしょう。
このイメージは、輸送後の分布の形を直接目標に近づけていくことに実質的に同じと思えば、one-step repurposing が可能であることも納得できるのではないでしょうか。

---

3.3. FD-loss による既存モデル の改善

Throughputの計算は使うGPUやバッチの組み方などエンジニアリングの側面の影響を大きく受けるので既存の報告と平等な比較を行うのはかなり難しいと言えます。

しかしながら、MeanFlowよりはpMFは計算コストが少ないことは[1]にて報告されたとおりであるという期待に基づくとFD-lossによるポストトレーニングモデルはこの図ではかなり右下に位置することになるでしょう。

pMFのポストトレーニングモデル、JiTのポストトレーニング前後を同じ条件でリアルタイム生成するデモを行うと以下のようになりました。

Summary

この記事では画像生成における中心的なトピックとして、one-step 推論および x-predictionの有効性に関する議論を行いました。これらがうまくいく理由の直感的な説明を提供できたかと思います。
現状商用化されている画像生成サービスは言語モデルとの統合により更なる精度向上という恩恵を受けていますが、特にFD-lossによるポストトレーニングは小規模なモデルの安価な性能向上方策として注目されたりするのではないでしょうか。

一方でFD系の指標は現状のFIDが最適解というわけではないでしょうし、他にも様々なOpen Questionは残っていると言えます。ここで紹介した論文はどれもこれからの画像生成モデルの研究の方向性を指し示す重要な論文ではないでしょうか。

参考文献

[1] Li et al., Back to Basics: Let Denoising Generative Models Denoise , 2025.

[2] Geng et al., Improved Mean Flows: On the Challenges of Fastforward Generative Models . 2025

[3] Lu et al., One-step Latent-free Image Generation with Pixel Mean Flows , 2026.

[4] Yang et al., Representation Fréchet Loss for Visual Generation , 2026.

[5] Yu et al., Autoregressive Image Generation with Masked Bit Modeling , 2026.