Rectified Flow⑤

前回後回しにした実験部分を読みます.

https://openreview.net/forum?id=mSHs6C7Nfa

①
https://zenn.dev/fmuuly/articles/37cc3a2f17138e

②
https://zenn.dev/fmuuly/articles/a062fcd340207f

③
https://zenn.dev/fmuuly/articles/0f262fc003e202

④
https://zenn.dev/fmuuly/articles/ebe448556ae2e9

書籍情報

参考文献も同一です.

Sangyun Lee and Zinan Lin and Giulia Fanti. Improving the Training of Rectified Flows. Advances in Neural Information Processing Systems, 38, 2024.

関連リンク

今回はarXivにあるものを読みます.

• OpenReview
• arXiv
• 公式実装

実験

CIFAR-10, AFHQ, FFHQ, ImageNetを用います. ベースラインはdiffusion models, distilled diffusion models, score distillation, GANs, consistency models, rectified flowsとなっていて全部でCIFAR-10では25個, ImageNetでは20個あります. 提案手法の訓練時間はImageNetの場合は64枚のNVIDIA V100で9日間です.

Unconditional Image Generation

最初に結果を示します.

著者らはこの結果から2つのことを主張しています.

1. 1-2NFEの場合に提案手法は両方のデータセットでoutperformあるいはcompetitiveである.
2. 提案手法である2-rectified flow++は通常の2-rectified flowのFIDを最大75%減らす

2つとも表を見ればわかることなので特に述べることはないですが, 蒸留はしないのでしょうか.

Reflow can be computationally more efficient than other distillation methods

Rectified Flowに似たような手法としてconsistency modelsがあります. ここではconsistency training (CT)とconsistency distillations (CD)を対象に, 提案手法のreflowが計算コストの観点で優れていることを確認します. CTとCDの詳しい内容は元論文を参照ください.

https://proceedings.mlr.press/v202/song23a.html

一見するとReflowはCTやCDより計算コストが高いように感じます. なぜなら, 一回訓練データを生成する必要があるからです. しかし, CDは訓練のiterationごとに4回, CTは2回のforwardがそれぞれ行われますが, Reflowでは1回で済みます. backwardもあるので単純に計算量が1/4, 1/2になるわけではないですが, かなり計算量は削減されそうです. 例えばImageNetでの実験の際のforwardの回数はペア生成に395M, 訓練で1433.6Mなので全部で1828.6M回です. 一方でCD, CTはそれぞれ5734.4M回 (3.1倍), 2867.2M回 (1.5倍)です. さらにReflowは1回生成してしまえば再利用できます. コンピュータのストレージ的には, 解像度が $64×6464\times64$ の場合では合成画像が42GBで, ノイズは乱数生成器の状態を持っておけばいいだけなので無視できます.

Effects of samplers

Rectified FlowはNeural ODEなので, 一般的に出力はNFEが増えると真の解に近づきます. そのサンプラーとNFEを変えてみて品質を測定したのが以下の結果です.

一般的なEuler samplerの場合, NFEの増加に伴って, FIDが下がっています. Heunの場合はその傾向がやや強まります. このことは, 高度なsamplerによって更なる改善が期待できることを示唆しています. new samplerというのは, $zt−Δt=zt−v(zt,t)Δt\bold{z}_{t-\Delta t}=\bold{z}_{t}-\bold{v}(\bold{z}_{t}, t)\Delta t$ であったupdate ruleを $zt−Δt=(1−(t−Δt))xθ(zt,t)+(t−Δt)z1\bold{z}_{t-\Delta t}=(1-(t-\Delta t))\bold{x}_{\theta}(\bold{z}_t, t)+(t-\Delta t)\bold{z}_{1}$ としたものです. NFEが3未満の場合は2つは同等ですが, 5未満ではsampling効率が4倍になるようです.

Inversion

Rectified FlowはNeural ODEなので, 逆方向に積分すればデータからノイズへの反転が可能です. Inversionは特に画像編集でよく使われますが, 推論回数が多いです. 下図(a)では, 同じ推論回数でもRectified Flowが再構成誤差が非常に小さいことがわかります. 特に, Rectified Flow2回分の推論はEDM16回分の推論より再構成誤差が小さいです.

図(b)はInversionで得られたノイズの品質を比較していて, EDMよりも本物のガウス分布に近いことがわかります. 実際に可視化した(c)をみてみると, ノイズも再構成の結果もEDMより良いことがわかります.

終わりに

実験パートだけなのですぐ終わりました. 提案手法を用いるとEDMより少ない推論回数で良い結果が得られそうです. 一方で, EDMを初期値として色々工夫した結果がEDMより良くなるのは当たり前に思えるので他に競合する手法があるのかは良くわかりませんが, これでいいのかな？とは思います. ちなみにEDMで初期化してReflowをやっているので論文のタイトルはImproving the Training of Rectified FlowsよりもImproving the Training of Reflowの方がいい気がします.

Stable Diffusion 3でこのReflow手法を取り入れた結果とかも知れたらいいなと思います. ちなみにInstaFlowはStable Diffusion 1.4に対して実質的にreflowを行なっていい感じの結果です.

https://openreview.net/forum?id=1k4yZbbDqX