[22’ JMLR] CDM: Cascaded Diffusion Models for High Fidelity Image Generation

0. Abstract

• SR3의 후속 논문으로 CDM(Cascaded Diffusion Models)이라는 cascade 구조를 제안하여 high-fidelity image generation을 수행한다. 저자들은 classifier와 같은 추가 모델을 사용하지 않고, diffusion model만으로 문제를 해결했음을 강조한다.
• Cascading pipeline의 sample quality는 conditioning augmentation이라고 하는 저자들이 제안한 data augmentation 기법에 의해 크게 향상된다. Low-resolution upsampling에서는 Gaussian augmentation을, high-resolution upsampling에서는 Gaussian blurring을 사용한다.

1. Cascaded Diffusion Model

기본적인 구조는 improved DDPM을 사용하였다. Conditional setting에서 class label $y$나 diffusion timestep $t$와 같은 scalar conditioning은 embedding으로 처리하였고 low resolution image $z$ conditioning은 upsampling 후 channelwise concatenation을 사용하였다.

High resolution data를 $x_0$이라 하고, low resolution data를 $z_0$이라 하자. 이때 low resolution diffusion model $p_\theta (z_0)$와 high resolution diffusion model (i.e. super-resolution model) $p_\theta (x_0 \vert z_0)$를 연속적으로(cascading) 학습한다. 이때 cascading pipeline은 다음과 같은 latent variable model으로 볼 수 있다.

\[p_\theta (x_0) = \int p_\theta (x_0 \vert z_0) p_\theta (z_0) dz_0\]

이는 한 번의 cascading이고, 이를 여러 번 반복할 수도 있다. Condition $c$를 사용하여 $p_\theta (z_0 \vert c)$ 및 $p_\theta (x_0 \vert z_0, c)$를 학습할 수도 있다. 이를 나타낸 것이 아래 Figure 4이다.

2. Conditioning Augmentation

지금까지 나온 이야기는 거의 SR3에서 나온 내용과 같다. CDM 논문에서의 가장 큰 contribution은 지금부터 등장할 conditioning augmentation이라고 생각한다. 이는 low-resolution과 high-resolution에서 다르게 적용된다. 결론부터 말하면 low resolution에서는 Gaussian noise augmentation을, high resolution에서는 Gaussian blurring augmentation을 사용한다.

2.1. Blurring Augmentation

먼저 high resolution에서의 blurring augmentation을 살펴보자. $128 \times 128$, $256 \times 256$으로 upsampling하는 과정에서 효과적이었다고 한다. 이는 filter size $k=(3,3)$, sigma $\sigma$를 가진 Gaussian filter를 적용하여 $z$를 $z_b$로 변환하여 사용하는 것이다. 학습 시에 50%의 샘플에서 적용했다고 한다. 50% 외에도 여러 비율을 실험했으나 큰 성능 향상이 없었다고 한다.

2.2. Truncated Conditioning Augmentation

Conditioning augmentation은 low-resolution에서 적용된다. $128 \times 128$보다 작을 때 효과적이었다고 한다. 그러나 이 방법을 적용하는 것은 아이러니하게도 실제 low-resolution model에는 변화가 없고, high-resolution model의 학습 과정 및 구조에 영향을 준다. 무슨 뜻인지 조금 더 알아보자.

Low resolution sample $z_0$로부터 high resolution sample $x_0$을 얻는 ancestral sampling 과정을 다음과 같이 나타낼 수 있다.

\[\begin{aligned} p_\theta (x_0) &= \int p_\theta (x_0 \vert z_0) p_\theta (z_0) dz_0 \\ &= \int p_\theta (x_0 \vert z_0) p_\theta (z_{0:T}) dz_{0:T} \\ \end{aligned}\]

Truncated conditional augmentation은 $z_0$이 아니라 $z_s$라는 early timestep에서 reverse process를 멈추고, $z_s$로부터 $x_0$를 얻는 것이다.

\[\begin{aligned} p_\theta ^ s (x_0) &= \int p_\theta (x_0 \vert z_s) p_\theta (z_s) dz_s \\ &= \int p_\theta (x_0 \vert z_s) p_\theta (z_{s:T}) dz_{s:T} \\ \end{aligned}\]

이제 base model과 super-resolution model을 각각 다음과 같이 나타낼 수 있다.

\[\begin{aligned} \text{Base} \quad& p_\theta (z_s) = \int p_\theta (z_{s:T}) dz_{s+1:T} \\ \text{Super-resolution} \quad& p_\theta (x_0 \vert z_s) = \int p(x_T) \prod_ {t=1} ^ T p_\theta (x_{t-1} \vert x_t, z_s) dx_{1:T} \\ \text{where} \quad& p_\theta(x_{t-1} \vert x_t, z_s) = \mathcal{N}(x_{t-1}; \mu_\theta (x_t, z_s), \Sigma_\theta(x_t, z_s)) \\ \end{aligned}\]

이 과정을 해석하면, $z_0$으로부터 바로 $x_0$을 생성하는 것이 아니라, $z_0$을 Gaussian noise augmentation 시킨 $z_s$로부터 $x_0$을 생성하는 것이다.

저자들은 이 방법을 조금 더 수학적으로 전개하였다. $p_\theta ^s (x_0)$의 ELBO를 구해 본 것이다. 즉, 이를 diffusion model $p_\theta (z_0)$을 prior로 가지고 diffusion decoder $p_\theta(x_0 \vert z_s)$를 가진 VAE로 간주했다. 그 결과는 다음과 같다.

이때 주목할 점은 $\sum {t \gt s}$를 사용하여 truncation을 적용했다는 것이다. 그리고 decoder $p\theta(x_0 \vert z_s)$ 또한 diffusion이므로 다음과 같은 diffusion의 ELBO를 또 구할 수 있다.

이를 합하면 전체 모델의 objective는 다음과 같다.

해석하면 항들이 서로 독립적이므로, low resolution model을 최적화하는 것과, high resolution model을 최적화하는 것은 분리될 수 있다는 것이다. 따라서 고정된 $s$ 값이 있다면 low resolution process를 truncation timestep $s$까지만 학습하고, super-resolution model은 이로부터 생성된 $z_s$로부터 학습한다.

그러나 저자들은 위의 ELBO objective로 학습하지는 않았다. 단지 해석을 위해 식을 전개했을 뿐 improved DDPM의 hybrid objective $L_{\text{hybrid}}$를 사용했다고 한다. 그리고 모든 $s$에 대하여 다른 모델을 만드는 것이 아니라 amortization을 사용하여 super-resolution model은 $z_s$와 $s$도 input으로 받을 수 있도록 모델을 설계했다.

2.3. Non-truncated Conditioning Augmentation

위 과정은 수학적으로는 좋지만 sampling 시에는 조금 불편하다. 만약 서로 다른 $s$에 대하여 upsampling을 진행한다면 각각의 $s$에 대하여 $z_s$를 모두 저장해두어야 하기 때문이다. 대신 저자들은 $z_0$를 만든 뒤 $z_s ^\prime \sim q(z_s \vert z_0)$로 샘플링하여 사용하는 방법을 제안한다. 저자들은 이를 non-truncated conditioning augmentation이라고 부르며, sampling 과정에서만 차이가 난다. 만약 diffusion model이 잘 학습되었다면 $z_s$나 $z_s ^\prime$는 비슷할 것이다. 실제로 이후 실험에서 결과가 비슷함을 보였다.

최종 알고리즘은 다음과 같다. Training 알고리즘에서는 12번째 줄이 $z_t$가 아니라 $z_s$로 바뀌어야 할 것 같다. Sampling 알고리즘을 보면 truncated conditioning augmentation과 non-truncated conditioning augmentation이 구분되어 있다. 학습할 때는 여러 truncation timestep $s$에 대하여 학습하고, 샘플링할 때는 정해진 $s$에 대하여 $z_s$를 생성한다.

3. Experiments

3.1. Comparison with SOTA

다른 주요 모델인 BigGAN-deep, Improved DDPM, ADM(Classifier Guidance 논문의 모델) 등과 비교하여 SOTA 성능을 보였다. CDM이 BigGAN-deep보다 FID는 좋았으나 IS는 낮았다. Figure 7은 classification accuracy를 plot한 것인데 굉장히 직관적인 방식이라 인상적이어서 가져와 보았다. Figure 8은 세 모델 사이의 비교인데 솔직히 자세히 보지 않고서는 큰 차이를 못 느끼겠다.

3.2. Baseline Model Improvement

Baseline model, 즉 class-conditional ImageNet $64 \times 64$ generation model의 성능 향상을 위하여 두 가지를 시도하였다. 첫 번째는 dropout을 사용한 것인데, 수렴 속도는 더 느려졌지만 결국 학습을 더 오래 하여 성능 향상이 있었다고 한다. 그리고 두 번째는 sampling step extension이다. 조금 웃긴 방식인데 $z_{T=4000}$으로부터 $z_0$을 만든 뒤 다시 $z_0$에서 forward process를 통해 $z_{t=2000} ^\prime$을 만들어내고, 이로부터 다시 $z_0 ^ \prime$을 생성하는 것이다. 이러한 방식으로도 성능 향상이 있었다고 한다.

3.3. Experiments at $64 \times 64$

다음은 $64 \times 64$로의 super-resolution model에서의 conditioning augmentation 실험 결과이다. 적절한 정도의 conditioning augmentation이 성능을 향상시킨다는 것을 보여주고 있다. 즉 너무 $s$가 큰 경우에는 너무 noisy한 이미지를 얻어 benefit을 얻지 못하고, 너무 $s$가 작은 경우에는 train-test mismatch가 발생하여 성능이 떨어진다.

Train-test mismatch에 대해 조금 더 설명해보겠다. Low-resolution model이 $z_0$를 생성하는 과정에서 train set에 대한 bias를 가지고 있을 수 있다. 따라서 bias가 포함된 $z_0$으로 upsampling을 진행하면 OOD(out-of-distribution) details를 가진 $z_0$으로부터 $x_0$가 생성되는 것이다. 즉 오류가 누적되는 것이다. 그러나 noisy $z_s$로부터 upsampling을 진행하면 이러한 noise로 인해 train-test mismatch가 줄어들게 된다. 저자들은 이를 뒷받침하는 증거로 ground truth를 사용하여 upsampling을 하는 경우 오히려 conditioning augmentation이 성능을 떨어뜨리는 것을 보였다. 즉, 적당한 conditioning augmentation은 train-test mismatch로부터 오는 오류를 경감시키기 때문에 성능이 향상된다고 주장한다.

3.4. Experiments at $128 \times 128$ and $256 \times 256$

$128 \times 128$ 및 $256 \times 256$으로의 super-resolution의 경우 앞에서 말한 것처럼 50%의 샘플에 대하여 blurring augmentation을 적용했다. 다른 비율도 시도했지만 성능 향상이 없었다. 이때 Gaussian filter의 size $k=(3,3)$이고, $\sigma$는 아래와 같이 hyperparameter tuning을 통해 결정했다. 그리고 class conditioning, large batch training, flip left-right augmentation으로 성능을 조금 더 끌어올렸다.

Sampling cost가 image resolution에 따라 quadratic하게 증가하기 때문에, sampling step 수를 줄이는 방법도 고민하였다. 결과적으로 $T=100$ 정도로도 충분한 성능을 보였다.

💡 Summary

• SR3의 후속 논문으로 CDM(Cascaded Diffusion Models)이라는 cascade 구조를 제안하여 high-fidelity image generation을 수행하였다.
• Conditioning augmentation은 low-resolution에서는 Gaussian noise augmentation을, high-resolution에서는 Gaussian blurring augmentation을 사용하여 성능을 향상시키는 data augmentation 기법이다. Gaussian noise augmentation을 통하여 train-test mismatch로부터 오는 오류를 줄일 수 있다. 이러한 방법은 cascade model의 오류 누적의 문제를 해소하기 위해 noise를 추가하는 것이 적절하다는 것을 시사한다.
• 성능 향상을 위해 dropout, sampling step extension, class conditioning, large batch training, flip left-right augmentation 등을 사용할 수 있다.
• Opinion: 본 논문을 처음 읽었을 때는 SR3과 별반 다른 것이 없어 왜 논문을 낸 것인가 의아했다. 그러나 conditioning augmentation이라는 새로운 방법을 제안하고, 이를 통해 train-test mismatch로부터 오는 오류를 줄일 수 있다는 해석은 꽤나 흥미롭다고 생각한다. 이 방법은 diffusion이 아니더라도 다양한 cascade model이 필요한 상황에서 유용하게 사용될 수 있을 것 같다.

📃 Reference

• [22’ JMLR] CDM: Cascaded Diffusion Models for High Fidelity Image Generation

Generation 카테고리 내 다른 글 보러가기