[review] Denoising Diffusion Probabilistic Models
UC Berkeley의 Jonathan HO, Ajay Jain, Pieter Abbeel의 논문입니다. 본문과 코드는 아래 github에서 확인할 수 있습니다.
GitHub - hojonathanho/diffusion: Denoising Diffusion Probabilistic Models
GitHub - hojonathanho/diffusion: Denoising Diffusion Probabilistic Models
Denoising Diffusion Probabilistic Models. Contribute to hojonathanho/diffusion development by creating an account on GitHub.
github.com
Abstract
Diffusion probabilistic model(이하 diffusion model)은 비평형 열역학(nonequilibrium thermodynamics)에서 영감을 받은 latent variable model의 일종입니다. 본 논문에서는 diffusion model을 사용하여 고품질의 이미지 합성을 수행했습니다. 또한 diffusion model과 Langevin 열역학 기반의 denosising score matching 를 연결하여 weighted variational bound 디자인에서 최선의 결과를 도출해냈습니다. 본 논문의 모델은 자동회귀(autoregressive) decoding의 일반화로 여겨지는 progressive lossy decompression을 받아들입니다.
1. Instroduction
심층 생성 모델들은 다양한 종류의 데이터 모델에서 고품질의 이미지들을 만들어내고 있습니다. 대표적으로 GAN, 자동회귀모델, flow, VAE 등이 energy-based 모델링에서 눈에 띄는 발전을 보여주고 있습니다.
Diffusion model은 Markov chain을 variational inference를 이용해 학습시킨 모델로, 원본 데이터에 noise를 조금씩 추가해가면서 만든 확산과정(diffusion process)를 역행하여 원본 데이터를 찾아냅니다. 확산이 작은 크기의 Gaussian noise로 구성되어 있다면, sampling 또한 조건부 Gaussian으로 설정해도 무방할 것입니다. 이는 신경망의 매개변수화를 더 간단하게 할 수 있습니다.
Diffusion model은 간단하게 정의할 수 있고 효율적으로 학습시킬 수 있지만, 아직 이 모델이 고품질의 sample을 만들 수 있다는 것은 증명되지 않았습니다. 본 논문에서는 diffusion model이 고품질의 sample을 만들 수 있으며, 다른 생성형 모델들 보다도 높은 성능을 보일 수 있음을 보입니다. 여기서 나아가 diffusion model에 특정 매개변수를 취했을 때, 훈련단계에서의 다중 noise level denoising score 모델과 동일하며 sampling 단계에서의 annealed Langevin 동역학 모델과 동일하다는 것 또한 보일 것입니다. 이 매개변수를 이용하면 최적의 sample을 얻을 수 있게 됩니다. 따라서 이 동치성은 본 논문의 주요한 기여 중 하나가 될 것입니다.
이 모델의 무손실 코드길이(lossless codelength)의 대부분은 이미지 내의 인지하기 어려운 세부사항을 묘사하는데 사용됩니다. 이 현상에 대해서는 lossy compression의 측면에서 분석할 것입니다. 또한 diffusion model의 sampling 정차가 preogressive decoding의 일종인 것도 보일 것입니다. 이는 autroregressive decoding과 유사합니다.
2. Background
Diffusion model은 $p_{\theta}(x_0) = \int p_{\theta}(x_{0:T}) dx_{1:T}$ 로 표현되는 latent variable model입니다. 여기서 $x$는 $x_0 \sim q(x_0)$와 동일한 차원을 가지는 latent들입니다. 여기서 $p_{\theta}(x_{0:T})$는 reverse process라고 불리며, $p(x_T) = N(x_T; 0, I)$에서 시작하여 Gaussian 변형이 적용된 Marcov chain 형태를 가집니다.
Diffusion model은 다른 모델들과는 달리 forward process(혹은 diffusion process)라고 하는 approximate posterior $q(x_{1:T} | x_0)$가 존재합니다. 이는 variance schedule $\beta$의 값에 따라 원본 데이터에 Gaussian noise를 조금씩 추가해나가 변형된 이미지를 만들어냅니다.
학습은 NLL(Negative Log Likelihood)에 대한 variational bound L을 최적화 하는 방향으로 수행합니다.
$\beta$는 reparametrization을 통해 학습되거나 hyperparameter로서 그냥 상수를 가지기도 합니다. Reverse process의 경우 $p_{\theta}(x_{t-1} | x_t)$의 조건부 Gaussian의 선택에 따라 그 표현이 부분적으로 보장됩니다. 두 과정 모두 $\beta$가 작을 때 같은 함수 형태를 가지고 있기 때문입니다. Forward process에서 주목할 점은 임의의 단계 t에 대해서 $x_t$를 sampling 할 수 있다는 점입니다. $\alpha_t = 1 - \beta_t$, $\overline{\alpha_t} = \prod_{s=1}^{t} {\alpha_s}$라고 할 때, $x_t$는 다음과 같이 나타낼 수 있습니다.
따라서 [식 3]의 L을 stochastic gradient descent를 이용해 최적화 한다면 효과적으로 학습할 수 있습니다. [식 3]은 다음과 같이 나타낼 수 있습니다.
[식 5]는 KL divergence를 이용해 $p_{\theta}(x_{t-1} | x_t)$와 forward posterior $q_(x_{t-1} | x_t, x_0)$ 사이의 차이를 계산합니다. 여기서 forward posterior $q_(x_{t-1} | x_t, x_0)$는 $x_0$가 주어져 있을 때 다음 분포를 따르게 됩니다.
3. Diffusion models and denoising autoencoders
Diffusion model은 latent variable model의 제한된 형태로 보일 수 있지만, 구현할 때 높은 자유도(degree of freedom)를 가지는 편입니다. Forward process에서 $\beta$를 정해야 하며, reverse process에서도 Gaussian 분포의 매개변수들을 정해야 합니다. 본 논문에서는 매개변수의 선택을 위해 diffusion model과 denoising score matching을 연결해 간단하고 가중치가 있는 variational bound를 만들었습니다. 이를 통해 본 논문의 모델은 간단하고 경험적인 결과를 얻을 수 있다고 합니다.
3.1. Forward process and $L_T$
본 논문에서는 $beta$를 reparametrization으로 학습시키지 않고 상수로 고정합니다. 따라서 approximate posterior $q(x_T | x_0)$는 학습할 매개변수를 가지지 않으므로, $L_T$는 상수로 취급해 무시할 수 있습니다.
3.2. Reverse process and $L_{1:T-1}$
이제 $q(x_{t-1} | x_t)$에 대해 논해야 합니다. 먼저 variance를 학습시키지 않기 때문에 $\Sigma_t(x_t, t) = \sigma^2I$로 둡니다. $\sigma^2$는 $\beta_t$($x_0 \sim N(0, I)$일 때) 또는 $ \frac{ 1 - \overline{a_{t-1}}}{ 1 - \overline{a_t} $($x_0$가 주어져 있을 때)가 될 수 있습니다.
그 다음은 평균 $\mu_{\theta}(x_t, t)$을 나타내야 합니다. $q(x_t{t-1} | x_t, x_0)$와 $p_{\theat}(x_{t-1} | x_t)$를 기반으로 $L_{t-1}$를 다음과 같이 나타낼 수 있습니다.
위 식을 통해 가장 적합한 $\mu_{\theta}(x_t, t)$의 형태는 forward posterior의 평균인 $\widetilde{\mu}_t(x_t, x_0)$일 때라는 것을 알 수 있습니다. 또한 [식 4]를 통해 $x_t$를 $x_t = \sqrt{ \overline{ \alpha_t } } x_0 + \sqrt{ 1 - \overline{ \alpha_t } } \epsilon (\epsilon \sim N(0, I))$로 표현하면 [식 8]을 다음과 같이 나타낼 수 있습니다.
[식 10]을 통해 $\mu_{\theta}(x_t(x_0, \epsilon), t)$는 다음과 같은 형태를 가질 때 가장 적합하다는 것을 알 수 있습니다.
여기서 $\epsilon_{\theta}$는 $x_t$가 주어져 있을 때 $\epsilon$의 예측값 함수입니다. 따라서 $p_{\theta}(x_{t-1} | x_t)$에서의 $x_t$ sampling은 $\dfrac{1}{\sqrt{\alpha _{t}}}\left( x_{t}-\dfrac{1-\alpha _{t}}{\sqrt{1-\overline{\alpha }_{t}}}\epsilon _{\theta}\left( x_{t},t\right) \right) +\sigma _{t}z (z \sim N(0, I))$로 나타낼 수 있습니다. 이 값은 Algorithm2에 사용되는 Langevin형태와 유사합니다. 이제 [식 11]을 이용해 [식 10]을 다음과 같이 단순화시킬 수 있습니다.
이는 t에 대한 다중 noise scaled denoising score matching 형태와 유사합니다. [식 12]는 Langevin 형태와 유사한 [식 11]을 포함하고 있기 때문에, denoising score matching과 유사한 형태를 최적화 하는 것은 Langevin와 유사한 sampling chain의 marginal을 구하는 것과 같다고 볼 수 있습니다.
이를 통해 reverse porcess의 평균 함수 근사치 $\mu_{theta}$를 학습시켜 $ \widetilde{\mu}_{theta}$나 $\epsilon$을 예측할 수 있습니다. $x_0$를 예측하도록 할 수도 있지만, 이전 실험들에서 이 식이 더 안 좋은 sample들을 만들어냄을 확인했습니다. 이$\epsilon$-prediction의 매개변수가 Langevin 동역학 모델과 유사하고, diffusion model의 variational bound가 denoising score matching과 유사하도록 단순화시킨다는 것을 보였습니다. 이후 4절에서는 $\epsilon$의 예측값과 $\widetilde{\mu}_t$의 예측값을 비교함으로서 그 효율성을 확인할 것입니다.
3.3. Data scaling, reverse process decoder, and $L_0$
본 논문에서는 이미지 데이터가 0~255의 정수를 [-1, 1] 구간으로 재조정한 것이라고 가정하고 있습니다. 이를 통해 reverse process는 prior인 $p(x_T)$에서부터 항상 일정한 크기의 입력값을 받는다고 보장할 수 있습니다. 이산 log likelihood를 계산하기 위해서 우리는 reverse process의 마지막 항을 독립적인 discrete decoder로 설정할 것입니다. 이 decoder는 $N(x_0; \mu_{\theta}(x_1, 1), \sigma_1^2I)$에서부터 다음과 같이 유도할 수 있습니다.
VAE decoder와 자동회귀모델에서 쓰이는 연속분포와 같이 이 모델은 variational bound가 데이터의 무손실 코드길이가 됨을 보장하며, 여기에는 데이터의 noise나 log likelihodd에서 크기 조절을 위한 Jacobian이 추가될 필요가 없습니다. 이를 통해 최종적으로 noise가 없는 $\mu_{\theta}(x_1, 1)$를 sampling 할 수 있습니다.
3.4. Simplified training objective
앞서 정의된 reverse process와 decoder를 기반으로 [식 12]와 [식 13]에서 유도된 variational bound는 $\theta$에 대해 미분 가능하며 바로 학습시키는데 바로 사용할 수 있습니다. 그러나 다음 식에서의 학습이 더 좋은 품질의 sample을 만들어냄을 밝혀냈습니다.
여기서 t = 1인 경우는 discrete decoder에서 $L_0$이 적분된 경우입니다. 이 적분은 Gaussian 확률밀도함수와 미소 너비의 곱을 수행하며, $\sigma_1^2$와 edge effect는 무시합니다. t > 1인 경우는 [식 12]에서 가중치가 없는 형태에 해당하며, Algorithm 1에서 이 단순화된 목표함수를 이용하고 있습니다.
[식 14]가 [식 12]에서의 가중치가 없기 때문에, 가중치가 있는 variational bound는 재구성(reconstruction)을 측면에서 표준 variational bound와 다른 점을 강조합니다.
4. Experiments
본 논문에서 행한 실험에서는 $T = 1000, \beta_1 = 10^{-4}, $\beta_T = 0.02$로 두었습니다. 이 상수들은 [-1, 1] 구간에 비해 작게 되도록 선택되어 reverse process와 forward process의 함수형태가 거의 같도록 만듭니다. 반면, $x_T$에서의 신호 대비 noise 비율은 가능한 작게 만들어야 합니다. 여기서는 $L_T \approx 10^{-5}$이 되게 합니다.
Reverse process를 나타내기 위해서 본 논문에서는 group normalization throughout이 있는 U-net backbone을 사용합니다.
4.1. Sample quality
[표 1]은 CIFAR10에 대한 inception, FID, NLL 수치를 나타냅니다. 여기어 본 논문에서 제시한 모델(Ours)이 training set과 test set에서 모두 다른 모델들 보다 더 뛰어난 성능을 발휘하는 것을 볼 수 있습니다.
예상하는 바와 같이 실제 variational bound에서 이 모델을 학습시킨다면 단순한 형태의 bound보다 더 좋은 codelength가 생성되지만, 후자가 제일 좋은 sample을 만들어냅니다.
4.2. Reverse process parameterization and training objective ablation
[표 2]는 reverse process의 매개변수와 학습 목표가 성능에 미치는 영향을 보여줍니다. $\widetilde{\mu}$를 예측하는 기본 모델의 경우 실제 variational bound에서 학습할 때는 잘 작동하지만, [식 14]와 같은 MSE(Mean Square Error) 형태의 경우 학습이 되지 않습니다. 또한 variance $\Sigma$를 학습시키는 것이 그 값이 주어져 있는 경우보다 학습능력이 떨어지는 것을 알 수 있습니다. $\epsilon$을 예측하는 새로운 모델의 경우 실제 variational bound에서 $\mu$의 경우와 비슷하면서 MSE 형태에서 훵씬 높은 성능을 보여줍니다.
4.3. Progressive coding
[표 1]은 CIFAR10모델의 codelength에 대해서도 보여줍니다. Train과 test간의 차이는 차원 당 0.03bit로, 본 논문의 모델이 과적합(overfitted)하지 않음을 뜻합니다. 이 lossless codelength는 energy based model과 score matching에 대해 보고된 추정치보다 나은 반면, 다른 likelihood based 생성 모델들과는 경쟁하기 않습니다.
본 논문에서 생성한 sample들은 좋은 품질을 보여주지만, 이 diffusion model에는 lossy compressor가 되도록 하는 어떤 편향이 존재한다고 예상합니다. Variational bound 항 $L_1 + ... + L_T$를 비율(rate), $L_0$를 왜곡(distortion)으로 정의할 때, CIFAR10 모델은 rate = 1.78 bits/dim, distortion = 1.97 bits/dim을 가집니다. 이는 0 ~ 255 범위에서의 MSE인 0.95값에 해당합니다. 전체 lossless codelength의 절반 이상이 distortion을 묘사한다고 볼 수 있습니다.
4.3.1. Progressive lossy compression
[식 5]와 유사한 형태의 progressive lossy code를 사용하면 모델에서의 rate-distortion의 경향에 대해 더 알아볼 수 있습니다. Algorithm 3, 4를 보면 둘 다 minimal random coding 함수에 접근하고 있습니다. Minimal random coding은 어떤 분포 p와 q에 대해 평균적으로 $D_{KL}(q(x) || p(x))$bits를 사용해 sample $x \sim q(x)$를 보냅니다(이 때 receiver는 q만 사용할 수 있습니다). 이를 $x_0 \sim q(x_0)$인 경우에 적용하면 algorithm 3과 4는 $x_T$부터 $x_0$까지를 차례대로 계산하며 총 codelength의 기대값은 [식 5]와 같아집니다. 그 결과 임의의 단계 t에 해당하는 receiver는 $x_t$값을 가지게 되며, 또한 [식 4]를 이용해 $x_0$을 점진적으로 예측할 수 있습니다. $x_0 \sim p(x_0 | x_t)$ 를 통해 바로 구할 수도 있지만, distortion으로 인해 사용하기 어려워 여기서는 다루지 않습니다.
[그림 5]는 CIFAR10의 test set에 대한 rate-distortion 그래프입니다. 각 단계 t에 대해서 distortion은 RMSE로 계산되었고, rate는 t에서 얻게되는 bit의 누적 함수로 계산됩니다. Distortion은 rate가 낮을 때 가파르게 감소하는데, 주요 bit들은 눈에 보이지 않는 수준의 distortion에 할당되기 때문입니다.
4.3.2. Progressive lossy compression
본 논문에서는 무작위 bit에서 점진적인 decompression이 주어졌을 때의 점진적 무조건적 생성(progressive unconditional generative process)과정을 통해 Reverse process의 결과값인 $\widehat{x}_{0}$를 예측했습니다. Reverse process에서의 sampling은 algorithm 2를 사용했습니다. [그림 6]과 [그림 10]은 reverse process의 결과값($\widehat{x}_{0}$)을 보여줍니다. [그림 7]은 모든 t에서 고정된 $x_t$에 대한 $x_0 \sim p(x_0 | x_t)$의 확률적인 예측을 보여줍니다. t 값이 작으면 자세한 부분은 제외한 모든 부분을 보존시키고, t값이 크면 대략적인 특징만 보존됩니다. 이것이 conceptual compression의 실마리가 될 수도 있을 것입니다.
4.3.3. Connection to autoregressive decoding
[식 5]의 variational bound는 다음과 같이 쓸 수도 있습니다.
Forward process에서 $q(x_t | x_0)$는 처음 t개의 좌표를 가린 상태로 $x_0$가 주어졌을 때의 모든 확률질량을 구할 수 있는데, 이를 통해 T를 데이터의 차원으로 설정할 수 있습니다. $p(x_T)$는 빈 이미지에 대한 질량을 할당하며, $p_{theta}(x_{t-1} | x_t)$는 fully expressive conditional distribution으로 설정합니다. 위 조건에서 $D_{KL}(q(x_T) || p(x_T)) = 0$이 되며, $D_{KL}(q(x_{t-1} | x_t) || p_{theta}(x_{t-1} | x_t))$를 최소화 하는 것은 $p_{theta}$가 변경되지 않은 t+1에서 T 까지의 좌표를 복사하고 t 번째 좌표를 예측하도록 힉습시킬 것입니다. 따라서, $p_{theta}$를 학습시키는 것은 자동회귀모델을 학습시키는 것이라고 볼 수 있습니다.
따라서 Gaussian diffusion model는 일반화된 bit ordering을 가진 자동회귀모델의 일종으로 볼 수 있습니다. 이 bit ordering은 데이터 좌표를 재배치해서는 얻을 수 없습니다. 선행 연구들은 이러한 재배치 작업이 smaple의 품질에 영향을 주는 의도적 편향을 유도한다는 것을 입증했습니다. 때문에 Gaussian diffusion의 noise가 masking noise보다 더 자연스럽게 적용된다는 점에서 자동회귀와 비슷하거나 그 이상의 역할을 수행한다고 추측합니다. 또한, Gaussian diffusion의 길이는 데이터의 차원과 항상 같은 필요가 없습니다. 예를 들어 T=1000일 때, 데이터의 차원은 32 * 32 * 3이나 256 * 256 * 3과 같이 더 큰 값을 가질 수 있습니다. Gaussian diffusion은 빠른 sampling을 위해 더 짧게 만들거나 모델의 표현력을 위해 더 길게 만들 수 있습니다.
4.4. Interpolation
원본 이미지 $x_0$와 ${x'}_0 \sim q(x_0)$가 있을 때, $x_t, x'_t \sim q(x_t | x_0)$를 encoder로 사용하면 이미지 사이를 보간할 수 있습니다. 이 선형 보간으로 만들어지는 latent vector $\bar {x_t} = (1 - \lambda)x_0 + \lambda x'_0$는 reverse process $\bar {x_0} \sim p(x_0 | \bar {x_t})$를 사용하면 다시 이미지 형태로 decode 할 수 있습니다. [그림 8-1]는 reverse process를 사용해 선형 보간된 원본 이미지에서부터 일부 요소들을 제거하는 과정을 보여줍니다. 서로 다른 $\lambda$의 값들을 조절하면 $x_t$와 $x'_t$가 같도록 만들 수 있는데, 이를 통해 단계 t에 대한 두 이미지 사이의 보간값을 찾을 수 있습니다.
[그림 8-2]는 original CelebA-HQ의 256 * 256 이미지의 보간과 재구축 과정을 보여줍니다. Reverse process는 고품질의 재국축과 그럴듯한 보간과정을 수행하면서 자세, 피부톤, 머리, 배경 등을 자연스럽게 변화시킬 수 있습니다. 이 때 t의 크기를 크게 할 수록 더 거칠고 다양한 변형이 일어나게 됩니다.
5. Related work
Diffusion model이 flow와 VAE기법과 유사해보이지만, q에 매개변수가 없고 $x_T$가 $x_0$와 거의 배타적 관계를 이룬다는 차이점이 있습니다. 본 논문에서 소개한 $\epsilon$-prediction reverse process의 매개변수화는 Langevin 동역학을 이용한 다중 noise level를 이용해서 diffusion model과 denoising score matching의 연결점을 만들어줍니다. 하지만, 직관적인 log likelihood 평가를 사용하며, variational inference를 사용해 Langevin 동역학 sampling 과정을 명시적으로 학습시킵니다. Diffusion model과 denoising score matching의 접점에 대한 reverse process는 가중치가 있는 denoising score matching 형태와 Langevin samplier를 학습시키는 variational inference가 같도록 합니다. Markov chain을 이용한 다른 기법으로는 infusion training, variational walkback, generative stochastic networks 등이 있습니다.
6. Conclusion
본 논문은 diffusion model을 통해서 고품질의 이미지 sample을 보여주었습니다. 그리고 diffusion model과 variational inference, denoising score matching, Langevin, energy-based model, 자동회귀모델, 그리고 점진적 lossy compression과의 접점을 찾아 냈습니다. Diffusion model은 이미지 데이터에 대해 훌륭한 의도편향을 가지고 있는 것 처럼 보입니다. 때문에 앞으로 다른 데이터 모델과 다른 생성형 및 기계학습 모델에 대한 유용성을 시험해 볼 것이라고합니다.