画像生成において高い性能を誇るDiffusionモデルだが、その潜在表現のもつれのなさについての研究が近年行われている。潜在変数には、高レベルのセマンティクスや、離散化、コンパクトさ、潜在空間において意味のある線形補間を実行する能力がほしい。その研究であるDisDiffまでの流れを追う。

Diff-AE

DiffusionベースのAuto Encoder
高レベル（属性や何が写っているか）のセマンティクスと低レベル（細かいディテール）の確率的広がりの両方が必要
→学習するエンコーダーで高レベル、DDIM(Denoising Diffusion Implicit Models)をデコーダーとして低レベルを制御
潜在変数を２分割
1.semantic subcode(CNNエンコーダ)
2.stochastic subcode(semantic subcodeで条件づけされたDDIMの逆過程)

DDIMの画像デコーダを$p(x_{t-1}|x_{t},z_{sem})$と表し、
セマンティックエンコーダを$z_{sem} = \text { Enc } _ {\phi} (x_0)$と表す。
潜在変数は $z = (z_{sem},x_T)$。

無条件生成の際にはlatent DDIM $p_{\omega} (z_{sem,t-1}|z_{sem,t}) $から残差接続付きMLP(10-20 層) を後から学習させて使う。
latent DDIMを後付けでトレーニングするという選択には、いくつかの実用的な理由がある。
第一に、latent DDIMのトレーニングは全トレーニング時間のほんの一部しかかからないので、ポストホックトレーニングにより、同じdiff-AEで異なるlatent DDIMの迅速な実験が可能になる。
もう一つの理由は、VAEの事前損失のような、潜在変数の品質を損なうような制約を課さないことで、$z_{sem}$をできるだけ表現豊かに保つためである。

PDAE

事前学習済み のDiffusionモデルを使って、Diff-AEと同様の表現学習をしたい

逆拡散過程のギャップ

条件付きDPM(Diffusion Probablistic Models)はDPMと比較して最適化後の損失が低くなる。これは、条件付きDPMが求める平均 $\mu_{\theta}(x_t,y,t)$が$\mu_{\theta}(x_t,t)$よりも理想とする$\tilde{\mu} _ t (x_t,x_0)$に近いことを表す。事後平均$\mu_{\theta}$の$\tilde\mu$とのずれは、順拡散過程の情報損失が原因である。ここで条件yのように、DPMに$x_0$の知識を導入すれば、ギャップは小さくなる。yが含む$x_0$の情報が多ければ多いほど、ギャップは小さくなる。
分類器による誘導では、平均を分類器の勾配を用いて誘導しているが、この誘導により、ギャップが小さくなっている。

一方、エンコードされた表現$z$に従い、平均を動かすモデルを用いてギャップを埋めるように学習させれば、エンコーダは$x_0$からできるだけ多くの情報を学習し、ギャップを埋めるのを助ける。ギャップが埋められれば埋められるほど、再構成が完璧になる。この原理に従い、PDAEは事前学習済みのDPMに基づくオートエンコーダを構築する。

ギャップを埋める教師なし表現学習

まず、画像$x_0$から潜在表現$z$を得るエンコーダ$E_\varphi (x_0)$を定義する。そして、事前学習済みの無条件DPMを画像最構成のデコーダに合わせる。
そして、勾配推定器$G_\psi (x_t,z,t)$を用いる。そして、分類器による誘導と同様に平均をシフトさせる。損失関数は以下のようになる。

$$L(\psi, \varphi) = \mathbf{E}_{x_0,t,\epsilon} [|| \epsilon - \epsilon _{\theta} (x_t,t) +\frac{\sqrt{\alpha_t} \sqrt{1-\bar{\alpha_t}}}{\beta_t} \Sigma _{\theta} (x_t,t) \cdot G _ \psi (x_t, E _{\varphi}(x_0) ,t) ||^2]$$

訓練済みDPMを用いるので、$\theta$は固定。$\Sigma _{\theta} = \frac{1-\bar\alpha _{t-1}}{1-\alpha_t} \beta_t I$とすることが多い。上の損失関数は$\Sigma _{\theta} (x_t,t) G _ \psi (x_t, E _{\varphi}(x_0) ,t)$を$\tilde{\mu}_t (x_t, t) - \mu _ {\theta} (x_t,t)$に近づけようとするものだと解釈できる。
また、Diff-AEと同様にlatentDPMを訓練させる。

ネットワークのアーキテクチャ

上の図がPDAEのアーキテクチャを示す。
Diff-AEと異なり、エンコーダ$E_\varphi$には純粋なconvolution layersとlinear layerを用いる。勾配推定器$G_\psi$として、U-Netを用いる。その中で、学習済みDPMのエンコーダ部分とtime-embedding層を再利用することで、学習済みDPMのノウハウをさらに活用できる。zを入力として用いるために以下のように正規化された特徴マップにスケーリングとシフトを2回適用することによりグループ正規化を拡張する。

$$ \text{AdaGN}(h,t,z) = z_s(t_s\text{GroupNorm}(h) + t_b) + z_b $$

ただし、$[t_s, t_b]$と$[z_s, z_b]$は$t,z$の線形写像である。
再利用されるエンコーダーから新しいデコーダーへの接続は、依然としてスキップ接続を使用していることに注意。
このようにして、 $G_\psi$ は事前に訓練されたDPMによって完全に決定され、異なるU-Netアーキテクチャに普遍的に適用することができる。

DisDiff

DPMのdisentanglement(もつれをほどく)をしたい。
事前に学習されたDPMが与えられた場合、DisDiffは画像データの背後にある固有の因子を自動的に発見し、DPMの勾配場を、発見された各因子の表現に条件付けられたいくつかのサブ勾配場に分離することができる。表現とサブ勾配の分離を容易にするために、DisDiffのための新しい分離損失を提案している。

前述したDiff-AEやPDAEにおいて、学習された潜在表現は、事前に定義されたセマンティクスで訓練された線形分類器によってのみ解釈することができる。
一方、分離表現学習は、観測されたデータの背後にある根本的な説明可能な要因の表現を学習することを目的としている。異なる要因は、それぞれ、独立に、異なる種類の画像のバリエーションに対応する。ほとんどの手法は、VAEやGANなどの生成モデルに基づいて、分離された表現を学習する。VAEベースの手法には、異分離能力と生成品質の間のトレードオフが内在している。GANベースの手法はGANの反転(画像→潜在変数)の難しさによる再構成の問題に悩まされる。我々の知る限り、DPMを用いて分離表現を学習する方法はない。

DPMの分離の利点は2つある：
(i)画像データの背後にある固有の意味的要因を自動的に発見することで、完全に教師なしでの画像制御を可能にする。これらの要素は、アノテーション／画像-テキストペアのような人間の定義したものや、CLIPのような教師あり事前学習モデルからDPMの条件情報を拡張するのを助ける。
また、部分勾配場の重ね合わせによって導入される情報の部分的な条件を柔軟にサンプリングすることも可能であり、これは既存のDPMでは斬新である。

(ii)DPMは画像生成品質において顕著な性能を持ち、逆問題、例えばDDIMの逆問題に対して自然に適合する。VAE（逆接能力と生成品質のトレードオフ）やGAN（GAN反転の問題）に比べ、DPMは分離表現学習のフレームワークとして優れている。潜在空間の分離には全相関以外の他の帰納的バイアスが必要である。DPMは、新しいタイプの帰納的バイアスとして、すべての異なるタイムステップからの制約を採用することを可能にする。さらに、データの情報には、分解されたものと分解されていないものがある。DPMは、分解されていない（条件付きでない）情報をサンプリングする能力を持っており、これは当然、分離に適合する。
DPMを分離するタスクに対処するために、DisDiffと名付けられた、事前学習されたDPMの分離のための教師なしソリューションを提案している。
DisDiffは、各因子の分離された潜在表現を学習するエンコーダーと、対応する分離された条件付き部分勾配場を学習するデコーダーを採用する。
さらに、符号化された表現が分離要件を満たし、入力画像を再構成するための新しい分離損失を提案する。

拡散に基づく画像編集は注目されており、主に2つのカテゴリに分類される。まず、DPMの潜在変数と入力画像を混合して画像を編集する画像ガイド型の作品である。しかし、Kwonらが指摘するように、画像を用いて編集属性を指定することは曖昧さを引き起こす可能性がある。
第二に、分類器誘導型は、余分な分類器の勾配を利用して画像を編集する。これらの方法は勾配を計算する必要があり、コストがかかる。一方、これらの手法では、アノテーションや、ラベル付きデータで事前に訓練されたモデルが必要である。教師データ無しで画像を編集するDisDiffを提案している。

一方、拡散モデルに関する文献では、表現学習にはほとんど注意が払われていない。関連する作品として、Diff-AEとPADEがある。Diff-AEは、画像再構成のための拡散ベースのオートエンコーダを提案している。PADEは、画像再構成のためのオートエンコーダを構築するために、事前に訓練されたDPMを用いる。しかし、これら2つの研究で学習された潜在表現は、データセットの根本的な要因に明示的に対応していない。我々の知る限り、我々のDisDiffは、離散表現学習のための最初の拡散ベースのフレームワークである。

Bengioらは、分離表現学習の概念を導入した。離散化表現学習の目的は、観測されたデータの説明因子を発見することである。分離された表現とは、分離された表現の各次元が独立した要因に対応するものとして定義される。の各次元が独立した要因に対応すると定義される。このような定義に基づき、いくつかのVAEに基づく研究は、表現の確率分布の制約によってのみ、離散化（disentanglement）を達成している。
Locatelloらは、これらの制約だけでは離散化には不十分であり、余分な帰納的バイアスが必要であることを証明し、識別可能な問題を指摘している。例えば、Yangら(2021)は、群論でモデル化された対称性を帰納的バイアスとして用いることを提案している。ほとんどの分離手法はVAEに基づいている。また、事前学習された生成モデルを活用するなど、GANに基づく研究もある。
DisDiffは、新しいタイプの帰納的バイアスとして、拡散過程における全ての時間ステップの制約を導入している。さらに、DPMは非因子化（非条件付き）情報のサンプリングが可能であり、分離に自然に適合する。このように、我々はDPMの新しい枠組みに基づいて、分離に光を当てる。

手法

DPMの分離

データセット$D$が$N$個の正しい要素（概念・因子）$C= \{ 1,2,...N \}$によって生成されたとする。例えば、Shape3Dの場合、基礎となる概念因子は、背景色、床色、オブジェクト色、オブジェクト形状、オブジェクトスケール、ポーズを含む。したがって、各サンプルと因子表現 $ h : x_0 \mapsto (f^1,\dots,f^N), \forall x_0 \in D $ の各タプルとの間には1対1のマッピングが存在する。データ分布$p(x)$は$N$個の独立した分 $\{ p(x|f^k) |k=1,\dots,N \}$ に分解でき、それぞれ1つの要因のみを条件とする。下図の(a)の画像空間の局面として示されている。

DPMは$p(x)$に収束する、分布の系列 $\{ p_t(x) | t=0,1,\dots,T \}$ を学習する。
このような収束は、対応する勾配場 $\{ \nabla_x \log{p_t (x)} |t=0,1,\dots,T \} $ を最適化することによって達成される。
分離されたDPMは$N$個の分布系列を含み、$N$個の分布 $\{ p(x|f^c) | c=1,\dots,N \}$ に収束する。

DPMの分離の方法は、それぞれの因子$c$について、$\nabla_x \log{p_t (x|f^c)}$を推定する $ G_\psi^c$を学習することである。これは図1(a)の曲面を指す矢印に対応する。

画像空間のようなデータサンプル空間では、曲面はデータサンプル空間がうまく整理されておらず、因子のばらつきがもつれていることを示している。図1(b)に示すように、画像空間と比較すると、正解の因子空間はよく整理されており、因子の部分空間は互いに直交している。
離散化表現学習は、生データを離散化表現に符号化するエンコーダ$E_{\phi}$を用いて正しい因子空間をモデル化することを目的とする。もつれ解消表現は、分離されたDPMの条件分布$\{ p(x|f^c )|c = 1, . . N \}$の条件を表現するのに理想的である。
分離された表現を条件として、分離されたDPMは柔軟にデータサンプルを生成することができる。
本論文では、DPMのもつれ解消のためのソリューションとして、DisDiffと名付けられた手法を提案する。

DisDiffの全体像

DisDiffの全体像を図2に示す。

$N$個の因子$C = \{1, 2,\dots, N\}$を持つデータセット$D$について事前に訓練された無条件DPMが与えられ、我々の目標は、教師なしでDPMを分離することである。具体的には、$x_0 \in D$と、各因子$c \in C$について分離された潜在表現$z^c$を得ることを指す。分離された勾配場$\nabla_{x_t} \log{p(z^c|x_t)}$とエンコーダ$E_\phi$を用いて、$E_\phi (x_0) = \{ E_\phi^1 (x_0), \dots, E_\phi^N (x_0) \} = \{ z^1, \dots, z^N \}$を求める。

したがって、条件付き逆拡散過程（因子集合$S \in C$によって条件付け、$z^S = \{ z^c | c \in S \}$）はガウシアン分布$p_\theta (x_{t-1} | x_t, z^S)$によって以下のように定式化される。

$$ N(x_{t-1}; \mu_{\theta} (x_t, t) + \Sigma_t \sum_{c \in S} \nabla_{x_t} \log{p(z^c|x_t)}, \sigma_t) $$

$p(z^c|x_t)$は実行不可能なので、推定器$G_\psi^c (x_t, z^c, t)$を用いる。
Sのオプションを変えることで、スコアベースの条件付けのトリックに従った近似器を次のように柔軟に考案することができる：
$$\epsilon_\psi (x_t, z^S, t) = \epsilon_\theta (x_t, t) -\sum_{c \in S} \sqrt{1-\bar{\alpha_t}} G_\psi^c (x_t, z^c, t)$$
そして、Tweedieの公式を用いて、対応するデータサンプルを次のように導き出すことができる：
$$ \hat{x}_0^S = \frac{x_t - \sqrt{1-\bar{\alpha}_t} \epsilon _\psi (x_t, z^S, t)}{\sqrt{\bar{\alpha_t}}} $$

例えば、分離された因子を一つだけ選ぶことができる。言い換えると、$S = c$とし、勾配$\nabla _ {x_t} \log{p(z^c|x_t)}$を使って（$z^c$のみに条件付けられた）事前学習済みDPMの生成を誘導する。

離散化の「完全性」要件によれば、離散化された表現$E_\phi(x_0)$はサンプル$x_0$の完全な情報を含んでいなければならない。すなわち、すべての分離された表現$E_\phi(x_0)$を条件として$x_0$を再構成することができる。

上の式に従い、$S=C$と設定してすべての分離された表現が含まれ、結果としてデータサンプル$x^c_0$が得られる。これはまさにPDAEにおける再構成のケースであることに注意されたい。ここでは、同じ再構成損失を採用する：

$$ L_r = \mathbf{E}_{x_0,t,\epsilon} || \epsilon - \epsilon _{\theta} (x_t,t) + \frac{\sqrt{\alpha_t} \sqrt{1-\bar{\alpha}_t}}{\beta_t} \sigma_t \sum _ {c \in C} G^c _ {\psi} (x_t, z^c, t) || $$

上記の完全性要件に加えて、各分離された表現は対応する1つの要素のみを独立に反映する必要がある。我々は、訓練済みDPMの分離を達成するために、分離損失という新しい損失を考案する。以下では、分離損失と全損失について述べる。

分離損失

上で述べたように、サンプル画像$x_0$と分離表現$E(x_0)$が与えられた際に、ランダムに$c \in C$をサンプリングし、$S=c$として条件付きサンプル画像$\hat{x}^c_0$（表現$z^c$にのみ条件付けられる）を得る。
分離要件によれば、無条件の予測画像 $\hat{x}_0$（事前学習済み無条件DPMにより生成）と比較して、$\hat{x}^c_0$は以下の2つの要件を満たす必要がある。

不変条件
k番目$ (k \neq c, k \in C)$の分離表現について、$E^k _ \phi
( \hat{x} _ 0 ^c)$ は $E^k_\phi (\hat{x}_0)$ と同じになっている必要がある。
変異条件
c番目の分離表現について、条件付き$E^c_ \phi ( \hat{x}_ 0)$は無条件$E^c_ \phi ( \hat{x}_ 0)$よりも$E^c_\phi (x_0)$ に近い必要がある。
以下では、上記2つの条件の詳細な実装について説明する。
上記(i)の不変条件によれば、k番目 $(k \neq c、k \in C)$ の表現は同じに保たれなければならない。
$E_\phi$を用いて2つのサンプルをエンコードし、k番目の表現間の距離を次のように導出する：
$$ d_k = || E_ \phi^k (\hat{x}_ 0^c) - E_ \phi^k (\hat{x}_0) || $$

すると、距離ベクトルは$ d = [d_1, d_2, \dots, d_C] $と表せ、最後に、クロスエントロピー損失を使ってインデックスcを特定し、他は変更せずに抑える
$$ L_{in} = \mathbf{E} _ {x _ 0,t, \epsilon,c} [ CrossEntropy(d,c) ] $$

これを不変性損失と呼ぶ。

(ii)の変異条件として、まず、無条件$\hat{x}^c_0$と$x_0$、条件付き$\hat{x}_0$と$x_0$それぞれの間のk番目の表現の距離スカラーを計算する。

$$ d^n_k = || E_ \phi^k (\hat{x}_ 0) - E_ \phi^k (x_0) || $$

$$ d^p_k = || E_ \phi^k (\hat{x}_ 0^c) - E_ \phi^k (x_0) || $$

上記の条件(ii)によれば、条件付き因子cに対して、$d^n_c - d^p_c $は最大化されるべきであり、他の因子、すなわち$d^n_k -d^p_k (k \neq c, k \in C)$は0に最小化されるべきである。先ほどと同様にエントロピーロスを用いて以下のようにする。

$$ L_{va} = \mathbf{E} _ {x _ 0,t, \epsilon,c} [ CrossEntropy(d^n - d^p ,c) ] $$

ただし、$ d^n = [d_1^n, \dots, d_C^n] , d^p = [d^p_1, \dots, d^p_C] $であり、これを変異性損失と呼ぶ。

不変性損失と変異性損失の2つが分離損失として用いられる。

全体の損失

完全性と分離の両方の要件を満たす必要があるため、全体の損失には上記の再構成損失（$L_r$）と分離損失（$L_{in}$と$L_{va}$）が含まれる。
ただし、2つのロスをバランスさせるための重みは慎重に設定する必要がある。
上記の分離損失 $L_{in}$と$L_{va}$は、拡散過程のサンプリングされた時間ステップを条件としている。しかし、拡散モデルの条件は時間ステップごとに異なる。
例えば、tがTに十分近い場合、$G^c _ {\psi} (x_t, z^c, t)$は主に、$z^c$を条件とする。そして、tが0に近い場合、出力は主に$x_t$を条件とする。

従って、異なる時間ステップでは、分離損失のための異なる重みを使用する必要がある。エンコーダの入力の差は条件によって変化する。具体的には、tがTに近い場合、$\hat{x}_t$と$\hat{x}^c_t$の差は小さくなる。また、tが0に近ければ、その差は大きくなる。以上の議論から、出力がより$z^c$を条件としているほど、重みを大きくする必要がある。そこで、エンコーダの入力間のMSE距離を重み係数として用いることを提案する：

$$ \gamma_d = \lambda || \hat{x}_0 - \hat{x}^c_0 ||^2 $$

ただし、$\gamma$はハイパラ、重みについての勾配計算はしない。

よって、全体の損失は以下のようになる。

$$ L_a = L_r + \gamma_d (L_{in} + L_{va}) $$

総相関との関係

このセクションでは、総相関がないことが私たちの分離損失の必要条件であることを示す。全相関はかつて、表現の分離のための重要な制約と見なされてきた。しかし、Locatelloら（2019）は、全相関の他に、他の帰納的バイアスも考慮すべきであると指摘している。本論文では、付加的な帰納的バイアスとしてDPMの分離損失を導入する。全相関がないことが必要条件であることを証明する。

最構成損失が最小化された場合、以下の等式が成立する。

$$ \nabla_{x_t} \log{p (z^1, \dots, z^N | x_t)} = \sum_{c \in C} G_{\psi} ^c (x_t , z^c, t) $$

一方、分離損失が最小化されたとき、$ G_{\psi} ^c (x_t , z^c, t) = \nabla_{x_t} \log{p (z^c | x_t)} $ となる。よって、結合確率分布と、周辺分布の積が等しくなる。

$$ p (z^1, \dots, z^N | x_t) = \prod_{c \in C} p(z^c|x_t) $$

したがって、全相関がないことはすべての$x_t$について成立する。

まとめ

Diffusionモデルを用いたもつれのない表現を得る手法を追った。Diffusionモデルの解釈性や制御性につながる研究分野だと思う。

Diffusionモデルのもつれのない表現学習

目次