[論文メモ] Pretraining is All You Need for Image-to-Image Translation

学習済みの生成モデルを利用してImage-to-Image(I2I) translationを実現。

I2I translationは基本的にアーキテクチャ設計が大変だったり、スクラッチからの学習が難しかったりする。
そこで既存の生成モデルを利用するI2I translationのフレームワークを提案。

手法

システムの全体像は以下の図。

生成モデルとしてdiffusion modelを利用する。

画像 $\boldsymbol{x}_0$ に $T$ 回ガウシアンノイズを加えていく。 $\boldsymbol{x}_T$ は完全なノイズ画像( $\boldsymbol{x}_T \sim \mathcal{N}(\textbf{0}, \boldsymbol{I})$ )。

$\alpha_t := \Pi^t_{i=1}(1 - \beta_i)$ として、 $\boldsymbol{x}_t$ は $\boldsymbol{x}_0$ から直接求められる。

もしくは $\boldsymbol{\epsilon}$ をガウシアンとして $\boldsymbol{x}_t = \sqrt{\alpha_t}\boldsymbol{x}_0 + \sqrt(1 - \alpha_t)\boldsymbol{\epsilon}$ 。

Generative pretraining

このdiffusion modelについて学習済みのものを利用するが、I2I で利用するのでよい潜在空間(近い点からは似た画像生成されるなど)になっていてほしい。

これらを満たす既存モデルとしてGLIDEを利用。GLIDEはテキストから画像を生成するdiffusion modelで大量のテキストと画像ペアで学習されていて生成能力も高い。
arxiv.org

Downstream adaptation

Base model finetuning

text2imageのモデルができたら、テキストが変換される潜在空間に様々な条件(スケッチ画像やセグメンテーションマスク)を埋め込む必要がある。
そこで条件に合わせたエンコーダを用意する。デコーダ(diffusion model)を $\tilde{\mathcal{D}}$ 、エンコーダを $\tilde{\mathcal{E}}$ 、条件を $\boldsymbol{y}$ として $\boldsymbol{x}_t = \tilde{\mathcal{D}}(\tilde{E}(\boldsymbol{x}_0, \boldsymbol{y}))$ という感じ。

学習は2ステージにする。
1ステージ目はエンコーダのみを学習する。これでもそこそこの画像が得られる。
2ステージ目はデコーダも合わせてfine tuningする。

Adversarial diffusion upsample

また、学習時の画像にBreaking stretching ratio (BSR)による劣化を施すことでテスト時とのサンプルのドメインシフトを防ぐ。
$L_0$ フィルタによる擬似的なオーバースムージングも導入。

これらのaugmentationでも生成結果がオーバースムージングされたようになる。
これはdiffusion denoisingがガウシアンノイズを仮定しているのが原因と推測。
ノイズ予測だけでなく、perceptual lossとadversarial lossを課すことで画像局所的な構造を改善する。
$\hat{\boldsymbol{x}}_0^t = (\boldsymbol{x}_t - \sqrt{1 - \alpha_t} \boldsymbol{\epsilon}_{\theta}(\boldsymbol{x}_t, \boldsymbol{y}, t)) / \sqrt{\alpha_t}$ にたいして、

$D_{\theta}$ はDiscriminator、 $\boldsymbol{\psi}_m$ はVGGの複数階層の特徴。

Normalized classifier-free guidance

diffusion modelが条件を無視するかもしれないのでclassifier-free guidanceを導入。
$p(\boldsymbol{x}_t|\boldsymbol{y}) \propto \frac{p(\boldsymbol{y}|\boldsymbol{x}_t)}{p(\boldsymbol{x}_t)}$ なので $\log p(\boldsymbol{y}|\boldsymbol{x_t})$ の勾配は以下のようになる。

これは $\boldsymbol{y}$ に識別する方向なのでこの方向に持っていけばよい。
条件なしを $\boldsymbol{\emptyset}$ として

このあたりはGLIDEと同じ。

ただ、このままだと分散と平均がシフトするらしい。
そこで式(7)のガイドノイズの平均を計算する。 $\hat{\boldsymbol{\epsilon}}_{\theta}(\boldsymbol{x}_t|\boldsymbol{y})$ の平均を $\hat{\mu}$ とするとこれがclassifier-free guidanceの平均のシフトになる。
$\boldsymbol{\epsilon}_{\theta}(\boldsymbol{x}_t|\boldsymbol{y})$ の平均を $\mu$ 、分散も同様に計算しこれらを使って標準化する。