Improved DDPM 논문 리뷰(Improved Denoising Diffusion Probabilistic Models)

Improved DDPM은 DDPM이 log-likelihood 관점에서도 좋은 성능을 보일 수 있도록 DDPM을 발전시킨 논문으로 역방향 과정에서의 분산을 학습하여 샘플 품질에는 차이가 없으면서도 순방향 과정을 10배 더 적게 사용하여 샘플링할 수 있도록 하였다. 또한 precision and recall을 사용해서 DDPM과 GAN을 목표 분포 관점에서 비교하였다.

 

Introduction

DDPM이 높은 품질의 샘플을 생성할 수 있다는 것은 DDPM의 논문에서 밝혀졌지만, likelihood-based model(e.g. autoregressive models, VAE)와 비교했을 때 경쟁할만한 log-likelihood를 갖지는 못한다. log-likelihood는 생성 모델링에서 널리 사용되는 지표로, 일반적으로 log-likelihood를 최적화하면 생성 모델이 데이터 분포의 모든 모드를 포착한다고 여겨진다.  이는 실제로 DDPM이 모든 모드의 분포를 캡처할 수 있는지에 대한 근본적인 의문을 가져왔다. 또한 DDPM이 CIFAR-10, LSUN등의 데이터셋이 아닌 ImageNet과 같은 더 높은 다양성을 가진 데이터셋에 대해서도 잘 확장할지도 불확실했다. 

 

그래서 Improved DDPM에서는 DDPM이 ImageNet같은 높은 다양성을 가진 데이터셋에서도 다른 모델들과 경쟁할 수 있는 log-likelihood를 달성할 수 있음을 보여준다. 또한 Variational Lower-Bound(VLB)을 더 최적화하기 위해 새로운 objective를 제안한다. 이 hybrid objective는 더 좋은 log-likelihood를 얻게 해주었지만, 더 많은 gradient noise가 발생한다. 이 논문에서는 Importance sampling을 통해 noise를 줄이고, 더 나은 log-likelihood를 얻게한다. 

 

또한 분산을 학습하는 방식은 더 적은 단계에서 모델로부터 샘플을 추출할 수 있게 해준다. 이는 샘플링 능력을 크게 향상시킬 수 있음을 의미한다. 추가로 improved precision and recall metrics를 통해 GAN과 improved DDPM을 비교하고, 모델의 크기와 학습 계산량 증가를 통한 모델의 성능 향상 정도를 측정했다. Improved DDPMS은 학습 계산량이 증가할 때 성능이 예측 가능한 방식으로 개선된다.

 

Denoising Diffusion Probabilistic Models(DDPM)

* DDPM에 대한 설명은 원본만 첨부

 

DDPM의 기본적인 설명이다. 우리는 $L_{vlb}$를 고려할 때 $L_0$와 $L_T$는 신경쓰지 않고 $L_t$만 고려하면 된다는 것을 알 수 있다.

 

위 식에서 (8)번 식 덕분에 정방향 과정은 임의의 단계로 직접 이동(즉, 임의의 $t$ 값을 샘플링)할 수 있게 된다. DDPM 원본 논문에서는 각 미니배치의 각 이미지에 대해 $t$의 값을 균일하게 선택했다.

 

여기서 눈 여겨 볼 점은 $L_{simple}$은 ${\Sigma}_{\theta}(x_t, t)$의 학습에 아무런 시그널을 주지 않는다는 점이다. 다만 DDPM 논문에서는 분산을 ${\sigma}^2_tI$로 고정하였으므로 문제가 되지 않았다. DDPM에서는 ${\sigma}^2 = {\beta}_t$($q(x_0)$이 iostropic Gaussian noise) 또는 ${\sigma}^2 = \tilde{{\beta}_t}$( $q(x_0)$이 delta function)를 사용해도 품질에 큰 상관이 없음을 보였다. 

 

Improving the Log-likelihood

Learning ${\Sigma}_{\theta}(x_t, t)$

우선 ${\sigma}_t$의 선택($ {\beta}_t}$와 $\tilde{\beta}_t}$)이 샘플 품질에 영향을 미치지 않는지에 대해 고찰해보자.

 

아래의 그림을 보면  $ {\beta}_t}$와 $\tilde{\beta}_t}$는 $t = 0$ 근처를 제외하고는 거의 동일하다. 즉, 모델이 인지할 수 없는 세부사항을 처리하는 부분( $t = 0$)에서만 차이가 난다. 이 차이는 확산 단계의 수를 늘리면 더욱 작아지며, 확산의 단계가 무한대로 가까워지면  ${\sigma}_t$의 선택은 샘플 품질에 전혀 영향을 주지 않는다는 사실을 시사한다. 즉, 확산 단계를 추가할수록 분포를 모델의 평균인 ${\mu}_{\theta}(x_t, t)$가 분산인 ${\Sigma}_{\theta}(x_t, t)$보다 훨씬 더 많이 결정한다는 사실을 알 수 있다.

위 그림은 샘플 품질을 위해 ${\sigma}_t$를 고정하는 것이 합리적인 선택임을 보여주지만, log-likelihood와 아무런 관련이 없는 선택이다. 아래 그림을 보면 확산 과정의 초기 몇 단계가 VLB에 가장 크게 기여한다는 것을 보여준다. 따라서 더 나은  ${\Sigma}_{\theta}(x_t, t)$ 선택을 통해 log-likelihood를 개선할 수 있다.

 

다시 그림 1로 돌아가보자. $ {\beta}_t}$와 $\tilde{\beta}_t}$가 각각 극단적인 상황을 가정하고 있는데, 그 값이 $t = 0$이 아닌 구간에서 매우 비슷하다는 말은 ${\Sigma}_{\theta}(x_t, t)$이 취할 수 있는 범위가 매우 작다는 것이다. 따라서 Improved DDPM에서는 아래와 같이 $ {\beta}_t}$와 $\tilde{\beta}_t}$를 log-domain에서 interpolation하고, 모델의 출력 $v$가 각 차원마다 하나의 성분을 포함한 벡터 $v$를 출력하여 아래와 같이 이용하기로 한다. $v$에는 제약이 없기 때문에 이론적으로는 모델이 보간된 범위를 벗어나 분산을 예측할 수 있으나, 실제 네트워크가 그렇게 행동하지는 않았다.

$L_{simple}$ 텀은 ${\Sigma}_{\theta}(x_t, t)$에 의존하지 않으므로 이 논문에서는 아래와 같은 새로운 hybrid obejctive를 제안한다.

논문에서 $lambda = 0.0001$로 설정하여 $L_{vlb}$가 $L_{simple}$을 압도하지 않도록 했다. 같은 논리로 $L_{vlb}$항에 대해 ${\mu}_{\theta}(x_t, t)$출력에 stop-gradient를 적용했다. 이를 통해 $L_{vlb}$가 ${\Sigma}_{\theta}(x_t, t)$에 영향을 미치는 동안  $L_{simple}$는 ${\mu}_{\theta}(x_t, t)$에만 영향을 미치도록 한다.

* stop-gradient : 모델 학습 과정에서 특정 변수에 대해 역전파를 중지하는 방법

 

Improving the Noise Schedule

DDPM에서 사용한 선형 노이즈 스케쥴은 고해상도 이미지($t$가 작을 때)는 효과적이지만, 해상도가 낮은 이미지($t$가 $T$에 가까울 때)에서는 정보 손실이 과도하게 발생하여 최적의 결과를 제공하지 못한다. 이 논문에서는 선형 노이즈 스케쥴로 학습된 모델의 역확산 과정의 최대 20%를 생략해도 FID score가 크게 나빠지지 않음을 보여준다.

 

이를 해결하기 위해 아래와 같이 제안된 새로운 코사인 노이즈 스케쥴링은 실전에서는 ${\beta}_t = 1 - \frac{\bar{alpha}_t}{\bar{alpha}_{t-1}}$가 0.999보다 커지지 않도록 하여 $t = T$근처에서 singularity가 생기지 않도록 한다. 또한 $t = 0$근처에서 ${\beta}_t$가 너무 작아지는 것을 방지하기 위해 offset $s = 0.008$을 사용한다. 

 

코사인 스케쥴은 중간 과정에서는 $ {\bar{alpha}_{t}$가 선형적으로 감소하고, $t = 0$와 $t = T$에서는 거의 변하지 않도록 설계되었다. 아래 그림은 DDPM에서의 선형 노이즈 스케쥴링과 Improved DDPM에서 제안한 코사인 노이즈 스케쥴링의 진행상황을 보여준다. DDPM의 스케쥴러는 0에 빠르게 수렴하여 필요 이상의 정보를 빠르게 파괴하는 것을 볼 수 있다.

Reducing Gradient Noise

실험을 통해 $L_{vlb}$가 상당히 학습하기 어렵다는 것을 알게 되었는데, 이 논문에서는 $L_{vlb}$의 gradient가 $L_{hybrid}$보다 훨씬 불안정하다는 가설을 세운다. 이를 확인하기 위해 두 목표로 훈련된 모델의 gradient noise scale을 평가하였으며, 아래 그림 7에 나와있다. 따라서 이 논문에서는 log-likelihood를 최적화하기 위해 $L_{vlb}$의 분산을 줄이는 방향으로 고민해보았다.

 

그림2로 돌아가보자. $L_{vlb}$는 각각의 항들이 서로 다른 크기를 가진다. 하지만 $t$를 균일하게 샘플링하는 것은 각각의 항의 크기를 고려하지 않은 계산으로 불필요하게 노이즈를 유발한다고 볼 수 있다. 이를 해결하기 위해 이 논문에서는 Importance Sampling을 사용한다. 

다만 $\mathbb{E}[L_t^2]$는 사전에 알 수 없으며, 학습 중 변할 수 있기 때문에, 각 손실항에 대해 이전 10개의 값을 기록하고, 이를 학습 중 동적으로 업데이트한다. 학습 초기에 우리는 $t \in [0, T-1]$의 모든 $t$에 대해 10개의 샘플을 뽑을 때까지 $t$를 균일하게 샘플링한다. 이렇게 importance samplong된 obejective $L_{vlb}$를 통해 최고의 log-likelihood를 얻을 수 있었다. 이는 $L_{vlb} (resampled)$를 통해 확인할 수 있다. 

 

Resultsand Ablations

$L_{hybrid}$와 코사인 스케쥴을 사용하면 DDPM과 유사한 FID를 유지하면서도 log-likelihood를 개선할 수 있다. $L_{vlb}$를 최적화하면 log-likelihood가 향상되는 대신 FID가 높아진다. 따라서 우리는 샘플의 품질은 희생하지 않으면서도 log-likelihood를 향상시키는 $L_{hybrid}$를 $L_{vlb}$보다 선호한다. 

 

Improving Sampling Speed

$T$ 확산 단계로 훈련된 모델의 경우, 훈련에 사용된 동일한 $t$의 sequence(1, 2, ..., T)를 사용해 샘플링할 수도 있지만, 주어진 임의의 subsequence $S$를 사용해 샘플링하는 것이 가능하다. 주어진 sequence $S$에 대해 훈련 노이즈 스케쥴 $\bar{\alpha}_t$를 기반으로 샘플링 노이즈 스케쥴 $ bar{\alpha}_{S_t}$를 얻을 수 있으며, 아래와 같이 샘플링 분산을 계산할 수 있다.

여기서 ${\Sigma}_{\theta}(x_t, t)$는  $ {\beta}_{S_t}}$와 $\tilde{\beta}_{S_t}}$사이의 범위에 매개변수화되기 때문에 더 짧은 확산 과정에서 자동으로 재조정된다. 따라서 $p(x_{S_{t-1}}|x_{S_t})$를 $N({\mu}_{\theta}(x_{S_t}, S_t)|{\Sigma}_{\theta}(x_{S_t}, S_t))$로 계산할 수 있다. 샘플링 단계를 $T$에서 $K$로 줄이기 위해, 이 논문에서는 1과 T사이를 K개의 균등한 간격으로 나눴다. 아래 그림은 $T= 4000$일 때, $K= = 25, 50, 100, 200, 400, 1000, 4000$ 샘플링 단계를 사용해 FID를 평가했다. 결과는 아래와 같다. (결과는 $L_{hybrid}$를 사용했을 때, 샘플링 단계를 줄여도 좋은 FID를 달성)

DDIM과의 비교에서도 DDIM이 50개 미만의 샘플링 단계에서는 더 나은 샘플을 생성하지만, 50개 이상의 샘플링 단계를 사용하면 샘플 품질이 저하된다는 것을 발견했다. 또한 DDIM은 훈련 초기에는 성능이 떨어지지만, 훈련이 진행됨에 따라 다른 샘플러들과 격차를 좁힌다. 또한 Improved DDPM에서 제안된 striding 기법은 DDIM의 성능은 크게 저하시키지만, 다른 샘플러들의 성능은 향상시킨다. 

 

Comparision to GANs

likelihood는 mode-coverage를 확인하는 좋은 지표이지만, GAN과의 비교 척도로 쓰기엔 부적절하다. 대신 이 논문에서는 precision과 recall을 이용해 GAN과 비교하는데 precision은 생성된 샘플의 품질을, recall은 생성된 샘플이 데이터 분포의 다양한 모드를 얼마나 잘 포괄하는지를 측정하는 척도다. GAN에서는 class-conditional model이 일반적이므로, Improved DDPM에서도 class 정보를 시간 $t$에 맞게 주입하기로 한다. 구체적으로, class embedding $v_i$를 시간 단계 embedding $e_t$에 추가하고, 이 임베딩을 모델 전체의 residual blocks에 전달한다. 논문에서는 $L_{hybrid}$ 목표를 사용해 250개의 샘플링 단계를 통해 파라마터가 큰 모델과 작은 모델을 학습했다. 

 

50K 샘플을 사용해 ImageNet 64 x 64 FID를 계산하였고, 그 결과는 다음 표와 같다. BigGAN-deep이 FID 측면에서는 더 뛰어나지만, recall측면에서는 어려움을 겪는 것을 알 수 있다. 이는 확산 모델이 유사한 GAN보다 분포의 모드를 더 잘 포괄한다는 것을 시사한다.

Scaling Model Size

이번에는 학습 계산량에 따른 FID와 NLL의 변화를 살펴보자. 이 논문의 결과는 예비적이기는 하지만, DDPMs가 학습 계산량이 증가함에 따라 예측 가능한 방식으로 개선된다는 점을 보여준다. 이 실험에서는 ImageNet 64x64에서 $L_{hybrid}$목표를 사용해 4가지 모델(first layer : 64, 96, 128, 192 channel)을 학습시켰다. 아래 그림은 결과를 타나난 것으로 FID 곡선은 log-log plot에서 대략 선형으로 나타나며, 멱법칙(power law)에 따라 스케일링된다는 것을 시사한다. NLL곡선은 멱법칙에 정확히 맞지 않으며, 검증 NLL이 FID보다 덜 유리하게 스케일링됨을 알 수 있다. 

이러한 현상의 원인으로는 (1) 예기치 않게 높은 비가역 손실 (2) 모델이 학습 데이터에 과적합됨을 꼽을 수 있다. 

 

Conclusion

We have shown that, with a few modifications, DDPMs can sample much faster and achieve better log-likelihoods with little impact on sample quality. The likekihood is improved by learning ${\Sigma}_{\theta}$ using our parametrization and $L_{hybrid}$ objective. This brings the likelihood of these models much closure to other likelihood-based models. We surprisingly discover that this change also allows sampling from these models with many fewer steps.

 

We have also found the DDPMs can match the sample quality of GANs while achieving much better mode coverage as measured by recall. Furthermore, we have investigated how DDPMs scale with the amount of available training compute, and found that more training compute trivially leads to better sample quality and log-likelihoods.

 

The combination of these results makes DDPMs an attractive choice for generative modeling, since they combine good log-likelihoods, high-quality samples, and reasonably fast sampling with a well-grounded, stationary training objective that scales easily with training compute. These results indicate that DDPMs are a promising direction for future research.