Rectified Flow③

論文の続きを読みます.

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

今回もベースとなっている1つ目の論文で, 実験パートです. arXiv版だと22ページのsection5です.

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

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

書籍情報

参考文献も同様です

Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. The Eleventh International Conference on Learning Representations, 2023

関連リンク

• OpenReview
• arXiv
• 公式実装
• 公式blog？のようなもの

Experiments

さて, ここまでの定式では $\pi_0$ と $\pi_1$ という2つの分布を考えていました. これらに特別な条件はなく, 分布がどのようなものでも成り立つというのがrectified flowでの主張でもあります (例えばGANや拡散モデルでは $\pi_0$ として正規分布を仮定したりします). なので, 様々な分布で $\pi_0, \pi_1$ の組み合わせを実験したいです. この論文では内容の厚薄はありますが4つの内容に対して実験を行っています. ですが, 主に画像ドメインに対しての実験です.

1. Toy-Data
2. Unconditional Image Generation
3. Image-to-Image Translation
4. Domain Adaption

それぞれ見ていくことにします.

Toy Examples

これは前回までtoy-dataと言ってきたもので, 主に性質確認として用いています. ただ, その詳細などには触れなかったのでここで触れておきます. toy-dataの実験では, non-parametricな $v^{X, h}(z, t)$ を用いています. ここで, $v^{X, h}(z, t)$ とは

$$v^{X, h}(z, t)=\mathbb{E}\left[\dfrac{X_1-z}{1-t}\omega_h(X_t, z)\right]$$

で登場する $v^{X, h}(z, t)$ です (これ自体は前回のSmooth function approximationで出てきます). 実装上ではこの期待値をearest neighbor estimatorで近似します. データ $\{x_0^{(i)},x_1^{(i)}\}_i$ が $(X_0, X_1)$ から得られるとき,

$$v^{X, h}(z, t)\approx\sum_{i\in\mathrm{knn}(z, m)}\dfrac{x_1^{(i)}-z}{1-t}\omega_h(x_t^{(i)}, z)\ /\ \sum_{i\in\mathrm{knn}(z, m)}\omega_h(x_t^{(i)}, z)$$

で近似します. ここで, $x_t^{(i)}=tx_1^{(i)}+(1-t)x_0^{(i)}$ で, $\mathrm{knn}(z, m)$ は $z\in\{x_t^{(i)}\}_i$ のtop $m$ nearest neighborsです. 結果は $m, h$ の選び方に影響を受けます. ここでは $h=1, m=100$ を用います. 生成時はEuler samplerを使用します. ステップ数を $N$ として step size $\dfrac{1}{N}$ と, 一定にします. $N=100$ です. toy-dataで使用するNNは2つの隠れ層を持ち, それぞれ64個のニューロンで構成されるMLPです. 結果を見てみます. 今回の設定では完全には直線にならないことが見てわかりますが, L2 penaltyを導入することでより直線的にすることができます.

①の回で出した図を見てみます.

この図では, 直線性はシミュレーションされた軌跡に基づいて

$$S(\boldsymbol{Z})=\int_0^1\mathbb{E}[\|(Z_1-Z_0)-\dot{Z}_t\|^2]\,\mathrm{d}t$$

の経験的推定として計算されています. 輸送コストはflowをシミュレーションして得られる $(Z_0, Z_1)$ から集められた $\{z_0^{(i)}, z_1^{(i)}\}_{i=1}^n$ を基に

$$\frac{1}{n} \sum_{i=1}^n \left( \|z_1^{(i)} - z_0^{(i)}\|^2 - \|z_1^{(i^*)} - z_0^{(i)}\|^2 \right),$$

で計算されます. ここで, $z_1^{(i^*)}$ は $\{z_0^{(i)}\}$ と $\{z_1^{(i)}\}$ の間の離散L2最適輸送問題を解いて得られる $z_0^{(i)}$ のoptimal L2 assignmentです. この指標は低次元の場合のみ有効で, 高次元の場合は $v^{X}$ が学習できていなくても指標が0になる傾向があります. 著者らはこの現象が「DDIMがL2最適輸送をもたらす」という誤った仮説を生む原因と主張しています.

『この現象が「DDIMがL2最適輸送をもたらす」という誤った仮説を生む原因と主張しています』について

ここでは以下の論文が参考文献として挙げられています.

https://openreview.net/forum?id=6PIrhAx1j4i

(この論文をちゃんと読んだわけではないですが,) 確かにこの論文の表1では合成データの実験に用いている次元は最大でも7で非常に小さいです (が, 一応画像ドメインでも実験しています. サンプル数が少ないことは述べられていることに注意が必要です).

この論文の査読を見ると似たようなことが指摘されています.

It was very briefly mentioned that recently Lavenant & Santambrogio found a counterexample to the central hypothesis. In my opinion, the paper needs to do a lot more in explaining whether such a counterexample is a singular case or not, other than mentioning the two maps disagree at a single point (how can two maps disagree at only one point if both maps are continuous enough?) and somehow "finite number of samples in higher dimensions" this is not a problem.

Unconditioned Image Generation

さて, 実験パートのメインの部分です. CIFAR-10での実験を軸にLSUN Bedroom, LSUN Church , CelebA HQ, AFHQ Catでの結果まであります. ただし, 高解像度データでの実験は定性評価にとどまっています. 実験設定として, $\pi_0$ に標準正規分布, $\pi_1$ にデータ分布を使用します. U-NetベースであるDDPM+のアーキテクチャを採用し, 特別に述べない限りはEuler samplerです.

訓練ではexponential moving average ratioに $0.999999$ を用いています. Adamで最適化を行い, 学習率は $2e-4$, dropout rateは0.15です. Reflowを行う際は, 400万の $(z_0, z_1)$ のペアを作り, 新しいデータセット $D$ とします. $i$-rectified flow modelを300,000steps更新して $(i+1)$-rectified flow modelを得ます. 蒸留の際は, $k$ step samplingをしたいとき, $t\in[0, 1]$ ではなく $t\in\{0, 1/k, \ldots, (k-1)/k\}$ とします. 特に, 精度向上のため $k=1$ のときはL2 lossではなくLPIPS lossを用います.

では結果を見てみます.

まずは, 表(a)の方に着目します. すると, 1-rectified flowのRK45 samplerがODEモデルの中では最良の結果となっています. 論文では具体的な数値を出して色々主張をしていますが, それは表を見ればわかることなので割愛します. 表(b)はその他のモデルの結果ですが, GAN with U-Netが同様のアーキテクチャで蒸留を行うことで $N=1$ でもGANを上回ることが可能になります. 続いて, Reflowモデルの結果を詳しく見てみます.

蒸留していないreflowを行ったモデルは, 図 (a)からわかるように, 推論回数 $N$ が小さい領域ではFIDおよびrecallが良いことがわかりますが, $N$ を大きくすると性能はサチっているように見えます. これは $v^X$ の推定誤差が蓄積するためと論文では述べられていますが, それだと横ばいになる理由を説明できていないように見えます. $N=1$ では, 蒸留を行うことで大幅な性能改善が見られます. reflowをすると性能改善はあまり感じにくいかもしれませんが, 表(a)のEuler solverの結果をみるとreflowなしの場合は蒸留が非常に (例えばFIDは60倍くらいの性能改善です)効果的であることがわかります. 次に図 (b)を見てみます. ここからは, reflowを行うことによってFIDとrecallが大幅に改善していることがわかります. 理論的な部分の主張では, $K\to\infty$ のとき完全に真っ直ぐなパスになるとの話でしたが, $k=4, 5, \ldots$ の結果がないのが気になります. 実用的には $k\geq4$ ではあまり改善が見られない (コスパが悪い)ということなのでしょうか.

直線軌道になる話題を出したので, その結果も確認します.

この図では確かに直線度合いを表す指標がreflowによって低下しています (2回目のreflowで単調減少ではなくスコアが一度悪くなっているのが気になります). また, random samplingしたpixel valueを見ても同じことが言えそうです. 軌道ばかり見ても実感が湧かないので, 実際にAFHQ catの結果も見てみます.

各点 $z_t$ において, $\hat{z}_1^t=z_t+(1-t)v(z_t, t)$ としています. もし軌跡が直線だとすると, $t$ を変化させても $\hat{z}_1^t$ は変化しないです. 上の図を見ても, reflowの有無によってそれが確認できます. これによってreflowで直線的になることが画像生成でもわかります. また, 1-rectified flowでも $t\approx0.1$ で鮮明な猫を視認できます. この鮮明度をsub-VP ODEで確認するには左下にあるように $t\approx0.6$ が必要で, 直線ではないものの近づいていることがわかります.

最後に, 高解像度のデータセットについての結果を確認します. 高解像度とは言ってもCIFAR-10と比較して高解像度であるという意味で, $256\times256$ です. 比較実験がないのでいい感じの生成ですねくらいしか言うことがありません. 比較可能なLSUN Churchを拡大してみてみると, 個人的にはStyleGAN2の方が上かな〜と思います.

続いて簡単な画像編集の結果を見てみます. まず, 2つの自然画像の上部と下部を繋げた不自然な画像 $z_1$ を作成します. その後, 1-rectified flowを逆方向に実行して潜在変数 $z_0$ を取得します. これを使って $\pi_0=\mathcal{N}(0, I)$ の下で尤度を増加させてより自然に見えるものを生成します. 確かに自然な画像になっています.

Image-to-Image Translation

$\pi_0, \pi_1$ の両方が画像の場合を考えてみます. ここでは教師なし (ペアなし)を考えて, 画像を変換することを考えます. 古典的なアプローチにCycleGANがあります.

https://openaccess.thecvf.com/content_iccv_2017/html/Zhu_Unpaired_Image-To-Image_Translation_ICCV_2017_paper.html

2017年やGANを使ったアプローチが古典的かどうかはさておき, CycleGANでは 2つの生成器 $F, G$ について $F(G(x))\approx x$ かつ $G(F(y))\approx y$ となるような学習を行います. CycleGANではGAN特有の話や, cycle consistencyの話を組み合わせていたため, lossは生成器のlossが2つとcycle consistencyのlossが1つありました. ところがrectified flowではそのようなことを考えないで行います. ここで大事なのは $\pi_0$ に従う　$X_0\sim\pi_0$ を　$\pi_1$ に従う　$X_1\sim\pi_1$ に忠実に変換することが目的ではないということです. 重要なのは, 画像のスタイルを変換しながら主要なオブジェクトの形状などのアイデンティティは保持することです. 例として人間の顔画像を猫の顔画像に変換することを考えると, 元の人間の顔に似ている」人間と猫のハイブリッドのような非現実的な顔を得ることが目標になります.

これを達成するには, $h(x)$ を画像 $x$ のfeature mappingとし, 転送したいstyleを表すとします. このとき $H_t=h(X_t)$ はODE $\mathrm{d}H_t=\nabla h(X_t)^\top(X_1-X_0)\mathrm{d}t$ に従います. すると, styleを正しく転送するために $H'_t=h(Z_t)$ となるように $v$ を学習します. $\mathrm{d}Z_t=v(Z_t, t)\mathrm{d}t$ が $H_t$ を限りなく近似するようにします. これは $\mathrm{d}H_t=\nabla h(Z_t)^\top v(Z_t, t)\mathrm{d}t$ が成り立つことから言えます. そのため, 目的関数は

$$\min_{v}\int_0^1\mathbb{E}\left[\|\nabla h(X_t)^\top(X_1-X_0-v(X_t, t))\|_2^2\right]\mathrm{d}t$$

実装上は $h(x)$ を2つのdomain $\pi_0, \pi_1$ を分類するモデルの潜在表現として設定します. これはImageNetで事前学習済みのモデルをfine-tuningします. 直感的には, $\nabla_x h(x)$ はsaliency scoreとして機能し, 目的関数のペナルティになります.

$\pi_0, \pi_1$ としてAFHQ, MetFace, CelebA-HQのペアを使用します. 各データセットについてtraining dataとtest dataの比率を $8:2$ とします. モデルのアーキテクチャは同じですが, 今度はAdamWを $(\beta_1, \beta_2)=(0.9, 0.999)$ で使用します. weight decayは0.1, dropout rateも0.1です. 訓練はbatch size4で1000 iteration行います. 学習率は $\{5\times10^{-4}, 2\times10^{-4}, 5\times10^{-5}, 2\times10^{-5}, \times10^{-6}\}$ を探索空間としてgrid searchで決めます. では, 結果を見てみます.

これは1-rectified flowのモデルに対して $N=100$ でEuler samplingをした結果です. いい感じに変換できていることがわかります. 続いてreflowの有無と $N$ の変化をみます.

あまり論文内では記述がないので以降では私見が多分に含まれますが, $N=100$ のときは違いがないように見えます. しかし, $N=1$ になるとreflowなしのモデルは変換ができていません. それに対してreflowを1回行ったモデルは $N=1$ でも変換ができており, 軌道がより直線的になっていることがわかります. 最後に $t$ を動かした時の変化をみます.

1-rectified flow (左図)ではスタイルの変化がはっきり視認できるのは右から2枚目でしょうか. 特に記述はないので0.2ずつ刻んでいると仮定して $t=0.6$ になります. 1個左の $t=0.4$ でも若干の変化を確認できますが, まだ $\pi_0$ の要素の方が多いように見えます. これはAFHQのノイズからの生成の時の結果とはかなり違います. 分布依存なのか人間の視覚能力依存なのかはわかりませんが, 定量的に測った結果があるとよさそうです. また, 1- and 2- rectified flowを並べた右図では, より直線軌道的になっていることが実感できます.

Domain Adaption

最後の実験です. 一般に機械学習ではdomain shiftが大きな課題となります. 十分に汎化していない場合はdomain shiftによって大幅な性能低下も考えられます. 著者らはrectified flowによってtraining domain $\pi_1$ をtest domain $\pi_0$ に変換してdomain shiftを軽減できると主張しています.

実験ではDomainNetとOffice-Homeを用います. DomainNetは6つの異なるドメインを持つcommon objectのデータセット (345カテゴリのオブジェクトがあります)で, Office-Homeは65のカテゴリを持つ4つのドメインで構成されています.

まず, 訓練データとテストデータの両方を事前学習済みモデルから潜在表現にし, その潜在表現を用いてrectified flowを学習します. アーキテクチャはこれまでと同じものを用いますが, 訓練時はAdamWをbatch size16で使用します. また, 学習率 $10^{-4}$ で50000iterationの更新を行い, OneCycle の学習率スケジューラーを使います. weight decayは0.1です. 評価は訓練データで学習した分類モデルにおける, 転送されたテストデータの精度です. 結果を見てみます.

基本的には既存手法と比較しても遜色ない, あるいは上回るパフォーマンスです. ただ, 大幅な性能向上ではないのでわざわざrectified flowを持ち出す場面ではなさそうです. 個人的には

A key challenge of applying machine learning to real-world problems is the domain shift between the training and test datasets:

で始まるsubsectionの割に内容が薄いなと思います. 特に, rectified flowにかなり近いスコアを出しているCORALは2016年のモデルで, これがSoTAなわけはないと思いますが...

https://link.springer.com/chapter/10.1007/978-3-319-49409-8_35

CORALについて

この論文ではDeep CORALをCORALと書いているのですが, CORALは別に論文があり (deep CORALでもベースラインとして引用されている),

https://ojs.aaai.org/index.php/AAAI/article/view/10306

2つの違いを明確にする必要があると思いますが, やっつけ仕事みたいだなと思ってしまいました (論文中ではDeep CORALと書いているので表だけ反映されてないのはミスなのか不明です).

おわりに

今回はrectified flowの論文の実験パートを読みました. Appendixにはさらなる結果がありますが, 省略します. 全体的にunconditional image generationが肝のパートで, 他はおまけといった感じです. 特に, domain adaptionについてはあってもなくても変わらないくらいの記述量のため, 全くといっていいほど参考にならないです. 査読でも自動運転を例に似たことが指摘されていますが著者ら的にはこれで十分だとしている様です. あとは, 分散の小さいデータセットでの実験だったのでImageNetとかconditional image generationなどでの結果も気になりますが, それはtextを条件としてStable Diffusion 3でやってくれてるのでそれを参考にすることとします.

https://proceedings.mlr.press/v235/esser24a.html

ただ, Stable Diffusion 3はreflowはしてない点に注意が必要です.

次回は2つ目の論文である

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

を読みます.