[논문 리뷰] Tackling the generative learning trilemma with denoising diffusion GANs

수업 때 발표할 논문으로, 간단하게 정리겸 리뷰하려고 한다. 예전에 논문스터디에서 다른 분이 발표하셨던 논문이기도 하다. generative model의 단점들을 해결할 수 있는 방법으로, Diffusion과 GAN을 결합한 방식이다. 부분적으로는 이해가 가는데, 전체적인 그림이 잘 안그려지는 듯하다. Training pipeline이 복잡해서 그런 것도 같다. 

 

diffusion GAN

 

Training 구조는 4가지로 이루어져 있다. 첫번째는 noise를 추가하는 diffusion process, 두번째는 real image와 fake image를 판별하는 Discriminator, fake image를 만드는 generator, 또한 fake image를 만드는 Posterior sampling이다. 아직 Diffusion도 GAN도 잘 몰라서 그런지 돌아가는 순서가 헷갈린다. 리뷰를 적으면서 생각을 정리해보겠다. 

Backgrounds

forward diffusion process는 real image $x_0 \sim q(x_0)$에 noise를 gradually add한다. 이때 T step까지 predefined variance schedule $\beta_t$로 수행된다. 기존 diffusion model의 backward process는 denoising하는 과정이다. 

forward process :

$$q(x_{1:T}|x_0) = \Pi_{t \geq 1}q(x_t|x_{t-1}), q(x_t|x_{t-1}) = \mathcal{N}(x_0; \sqrt{1-\beta_t},\beta_t)$$

backward process :

$$p_\theta(x_{0:T}) = p(x_T)\Pi_{t \geq 1}p_\theta(x_{t-1}|x_t), p_\theta(x_{t-1}|x_t)=\mathcal{x_{t-1};\mu_\theta(x_t,t), {\sigma_t}^2}$$

여기서 $\mu_\theta(x_t,t), {\sigma_t}^2$는 denoising model의 mean과 variance이고, $\theta$는 model parameter이다. 기존 diffusion model의 training의 goal은 ELBO(evidence lower bound)를 maximize하여  likelihood $p_\theta(x_0) = \int{p_\theta(x_{0:T})dx_{1:T}}$ 를 maximizing한다. $(\mathcal{L} \geq \log p)\theta(x_0))$ ELBO는 다음 식을 이용해서 parameterized denoising model $p_\theta(x_{t-1}|x_t)$과 함께 true denoising distribution $q(x_{t-1}|t)$ 으로 매칭한다. (? 여기서 매칭한다는 표현이 어떤 의미인지 모르겠다. 그냥 ELBO와 다음 식이 같다는 말로 생각하겠다.) 

$$\mathcal{L} = -\sum_{t \geq 1}\mathbb{E}_{q(x_t)}[D_{KL}(q(x_{t-1}|x_t) || p_\theta(x_{t-1}| x_t))] + C$$ 

이 Objective function은 intractable하다. 왜냐면 아까 말했듯이, true denoising distribution $q(x_{t-1}|x_t)$가 있기 때문이다. 그래서 기존 diffusion 논문들은 score based model을 이용해 tractable하게 만들었다. 

 

그러나 여기서 이 논문 제목에 따라 3중 딜레마가 발생한다. Variational autoencoder를 사용하건, diffusion을 사용하건, GAN을 사용하건, 항상 어딘가는 compute time, mode collapse, good quality issue를 해결하지 못한다. 왜냐하면, VAE를 사용하면 fast sampling과 mode coverage는 가능하지만 High quality sample을 얻기 어렵고, GAN을 사용하면 High quality sample과 fast sampling은 가능하지만 model collapse가 발생하고, Diffusion을 사용하면 high quality sample과 mode coverage는 가능한데 fast sampling이 어렵다. 따라서 이 논문에서는 diffusion과 GAN을 활용해 이 세가지 issue를 모두 해결하고자 한다. 

Methods

방법은 diffusion의 slow sampling을 관찰하면서 시작된다. diffusion의 key assumption은 (1) denoising distribution $p_\theta(x_{t-1}|x_t)$가 Gaussian distribution을 따른다; (2) denoising step은 매우 커야한다. 그 이유를 분석해보자. 일단 bayes rule에 의해 $q( x_{t-1} |x_t) \propto q(x_t|x_{t-1})q(x_{t-1})$ 이다. 첫번째 항 $ q(x_t|x_{t-1}) $ 은 forward gaussian diffusion의 distribution이고 두번째 항 $ q(x_{t-1}) $ 은 step t에서의 marginal distribution이다. 여기서 우변은 첫번째 항에 dominant 된다. 그리고 reversal process도 forward process와 identical한 form을 가진다. 따라서, $\beta_t$가 충분히 작을 때 (왜 이런 조건이 붙지?),  $q(x_t|x_{t-1})$ 도 gaussian 이기 때문에, $q(x_{t-1}|x_t)$ 도 gaussian이다. 그리고 또한, 만약 $q(x_t)$가 gaussian이라면, $q(x_{t-1}|x_t)$도 마찬가지로 gaussian 이다. 하지만 data를 gaussian으로 transforming하는 것 역시 challenging하다. 

 

따라서 $q(x_{t-1}|x_t)$를 gaussian으로 approximate하지 않고 conditional GAN으로 approximate 한다. 그 이유는 이 방법은 true distribution이 gaussian이어야 한다는 assumption이 필요하며 전체 step수가 매우 많아야한다. 그러니까, stepsize 를 아주 잘게 쪼개야만 1개 mode만 나오는 gaussian으로 근사할 수 있기 떄문이다. 만약 multi-modal을 cover하려 한다면, 이는 불가능하다. 

 

결국 이 논문에서 하고싶은건, reverse process에서 stepsize T를 줄이고 싶은것이다. 따라서 multi mode를 cover하기 위해, conditional GAN을 사용하여  complex conditionial distribution을 modeling한다. 즉, 이를 true denoising distribution $q(x_{t-1}|x_t)$ 를 근사하는데에 사용한다. 특별히 forward process는 T($T \leq 8$)가 작아야한다는 main assumption을 원래 diffusion과 비슷하게 적용하고 each diffusion step은 larget $\beta_t$를 갖는다. training은 conditional GAN generator $p_\theta(x_{t-1}|x_t)$ 와 $q(x_{t-1}|x_t)$를 adversarial loss (denoising step마다 divergence $D_{adv}$를 최소화하는)를 사용하여 같도록 formulate한다. 

$$\min_\theta \sum_{t \geq 1}\mathbb{E}_{q(x_t)}[D_{adv}(q(x_{t-1}|x_t)||p_\theta(x_{t-1}|x_t))]$$

where $D_{adv}$ can be Wasserstein distance, Jenson-shannon divergence or f-divergence  depending on the adversarial training step. 이 논문에서는, non-saturating GAN을 사용하는데, 이 GAN은 styleGAN등에 사용되었다. non-saturating GAN이 뭘까.. 어쨌든 이 경우, $D_{adv}$는 f-divergence(called softened reverse KL)라는 특별한 것을 사용하는데,  KL divergence를 사용하는 original diffusion과는 다르다. 

 

adversarial training을 하기 위해서, time-dependent discriminator $D_\phi(x_{t-1},x,t)$를 설정하고, $\phi$로 parameterized 되어있다. 이는 $x_{t-1}, x_t$ 이라는 N dimensional input을 가지고 $x_{t-1}$이 plausible denoised version of $x_t$인지 결정한다.

$$\min_\phi \sum_{t \geq 1} \mathbb{E}_{q(x_t)} [\mathbb{E}_{q(x_{t-1}|x_t)} [-\log(D_\phi(x_{t-1},x_t,t))] + \mathbb{E}_{p_\theta(x_{t-1}|x_t)} [-\log(1-D_\phi(x_{t-1},x_t,t))]]$$

where fake samples from $p_\theta(x_{t-1}|x_t)$ are contrasted against real samples from $q(x_{t-1}|x_t)$. First expectation은  unknown인 $q(x_{t-1}|x_t)$을 필요로 하지만, 저자들은 identity $q(x_t,x_{t-1}) = \int dx_0 q(x_0)q(x_t, x_{t-1}|x_0) = \int dx_0 q(x_0)q(x_{t-1}|x_0)q(x_t|x_{t-1})$를 사용한다.  rewrite된 first expectation 은 다음과 같다.
$$\mathbb{E}_{q(x_t)q(x_{t-1}|x_t)}[-\log(D_\phi(x_{t-1},x_t, t))] = \mathbb{E}_{q(x_0)q(x_{t-1}|x_0)q(x_t|x_{t-1})}[-\log(D_\phi(x_{t-1},x_t,t))]$$

 Given discriminator, 저자들은 generator를 $\max_\theta \sum_{t\geq 1}\mathbb{E}_{q(x_t)}\mathbb{E}_{p_\theta(x_{t-1}|x_t)}[\log (D_\phi(x_{t-1},x_t,t))]$로 업데이트하는데 이것이 non-saturating GAN objective이다. 

 

마지막으로, diffusion model들은 denoising step에서 directly predicting $x_{t-1}$하는 대신, denoising model을 다음과 같이 parmeterize 한다. $p_\theta(x_{t-1}|x_t):=q(x_{t-1}|x_t,x_0=f_\theta(x_t,t))$ 여기서 first $x_0$은 denoising model $f_\theta(x_t,t)$에 의해 predicted되고, 그런 다음 $x_{t-1}$ 은  posterior distribution $q(x_{t-1}|x_t,x_0)$ 를 이용해 sampling한다 (given $x_t$, predicted $x_0$).  $q(x_{t-1}|x_t,x_0)$는 $x_t$를 $x_0$로 denoising할 때 직관적으로 distribution over $x_{t-1}$ 이다. 그리고 이것은 diffusion process 를 위해 언제나 gaussian form을 가진다. 이와 비슷하게 $p_\theta(x_{t-1}|x_t)$를 정의한다.  

$$p_\theta(x_{t-1}|x_t) :=\int p_\theta(x_0|x_t)q(x_{t-1}|x_t,x_0)dx_0 = \int p(z)q(x_{t-1}|x_t,x_0=G_\theta(x_t,z,t))dz$$

where $p_\theta(x_0|x_t)$ is the implicit distribution imposed by the GAN generator $G_\theta(x_t,z,t,)$ that outputs $x_0$ given $x_t$ and L dimensional latent variable $z\sim p(z):=\mathcal{N}(z;0,I)$

 

이 논문의 parameterization은 몇가지 advantages가 있다. 첫번째로, $p_\theta(x_{t-1}|x_t)$는 DDPM과 같이 formulated 되었다. 따라서 DDPM 에서 사용하는 network design과 같은 inductive bias를 빌려올수 있었다. 가장 다른점은, DDPM에서는 $x_0$는  deterministic mapping으로 predicted되었다면, 이들의  $x_0$는 random latent variable $z$를 가지고 generator에서 나왔다. 이것은 $p_\theta(x_{t-1}|x_t)$를 multimodal하고 complex하게 만들어줄 수 있게 되었다. 두번째로, 서로다른 t에 대하여, $x_t$는 다른 levels of perturbation 을 가진다. 그러므로 different t에 single network를 이용하여 directly $x_{t-1}$을 predict하는 것은 아마도 어려울 것이다. 하지만, 이들의 경우 generator가 unperturbed $x_0$만 예측하면 되고 그다음 $q(x_{t-1}|x_t,x_0)$를 이용하여 back perturbation을 더한다.