Diffusion Models#

Diffusion Models 혹은 Score-based Models는 데이터에 점진적으로 Gaussian noise를 주입하고 reverse denoise process로 noise를 제거하여 데이터를 sampling하는 기법이다. 반면 forwad process는 원본 데이터 분포인 $p_{data}(x)$를 주변 확률분포인 $q_t(x_t)$로 변환한다:

여기서 $\alpha (t)$와 $\sigma (t)$는 noise scheduler를 의미한다. 연속적인 timestep의 관점에서 이를 확률미분방정식(Stochastic Differential Equation, SDE)으로 나타낼 수 있는데, 다음과 같다:

또한 주변 확률분포 $q_t(x)$는 **Ptobability Flow ODE(PF-ODE)**라는 상미분방정식(Ordinary Differential Equation, ODE)을 만족하는데 다음과 같다:

이때 Diffusion model은 $-\mathrm{\nabla }\log q_t(x_t)$(score function)를 예측하는 noise 예측 모델(${ϵ}_{\theta }(x_t,t)$)을 학습시킨다. 학습된 모델은 score function의 근사치를 예측하고 sampling하는데 이를 empirical PF-ODE라 한다 (경험적 PF-ODE):

Classifier-Free Guidance (CFG)는 sampling의 퀄리티를 높이기 위해 GLIDE, Stable Diffusion, DALL$\cdot$E2, Imagen 등 다양한 conditional model에서 사용되었다. CFG의 scale $\omega$가 주어졌을 때 원본 noise prediction은 conditional, unconditional noise prediction을 선형적으로 혼합하여 대체된다:

Consistency Models#

Consistenct Model(CM)은 몇 step 혹은 한번의 step 만으로 데이터를 생성할 수 있는 모델이다. CM의 핵심은 PF-ODE의 궤적에 어떤 point와 PF-ODE의 solution에 대해 mapping되는 function ($f:(x_t,t)\mapsto x_{ϵ}$)을 추정하는 것이다.

$ϵ$은 고정된 매우 작은 양수값을 가지며 CM의 function은 자기 자신에 대한 self-consistency를 만족해야한다. 즉 어떠한 time step에 대해서도 $x_{ϵ}$을 sampling 할 수 있어야 한다.

$f_{\theta }(x,ϵ)=x$를 만족하는 모델 $f_{\theta }$는 다음과 같이 정리할 수 있다:

$c_{skip}(t)$와 $c_{out}(t)$는 미분 가능한 함수이며 $c_{skip}=1,c_{out}=0$이기 때문에 $f_{\theta }(x,ϵ)=x$를 만족한다. $\theta$는 학습 가능한 파라미터로 $F_{\theta }$는 심층 신경망을 의미한다.

CM은 pre-trained 모델에 대한 Distillation 방식과 scratch부터 학습하는 방식이 있는데 주로 Distillation 방식을 사용한다. Distillation 방식은 parameter ${\theta }^{-}$가 $\theta$를 통해 학습하며 모델에 대한 self-consistency를 위해 다음과 같이 손실함수를 구성한다:

이 때 ${\theta }^{-}$는 $\theta$에 대한 지수평균이동(Exponential Moving Average, EMA)이며 ${\theta }^{-}\leftarrow \mu {\theta }^{-}+(1-\mu )\theta$이다. $d(\cdot ,\cdot )$은 두 sample 사이의 거리를 측정하는 지표이다. ${\hat{x}}_{t_n}^{\varphi }$은 $x_{t_{n+1}}$에 대한 $x_{t_n}$을 추정한 값으로 다음과 같다:

$\mathrm{\Phi }$는 PF-ODE에 사용되는 ODE Solver로 Euler나 Heun Method등의 수치적인 ODE solver를 사용할 수 있다. 즉 Consistency Distillation은 ODE Solver로 예측한 ${\hat{x}}_{t_n}^{\varphi }$과 $x_{t_{n+1}}$을 입력으로 $f_{{\theta }^{-}}$와 $f_{\theta }$로 예측한 값의 Consistency를 비교하는 방식으로 Distillation을 수행한다.

3.1 Consistency Distillation in the Latent Space#

본 논문에서는 pre-trained 된 Stable Diffusion에 Consistency Distillation을 적용한 Latent Consistency Distillation (LCD)을 제안한다. LCMs는 LDM(SD)을 기반으로 설계되었기 때문에 $z=\epsilon (x)$를 통해 $x$를 latent vector로 임베딩하고 $\hat{x}=\mathcal{D}(z)$를 통해 원본 영상으로 복원한다. latent space 상에서 연산이 이뤄지기 때문에 Computation Cost를 크게 줄일 수 있어 high-resolution 영상을 laptop GPU에서 생성할 수도 있다.

condition을 추가한 PF-ODE의 reverse process는 다음과 같이 정의된다:

$z_t$는 t step의 image latents, ${ϵ}_{\theta }(z_t,c,t)$는 noise 예측 모델, c는 text와 같은 conditional prompt를 의미한다. PF-ODE상에서 모든 t step에 대해 consistency function $f_{\theta }:(z_t,c,t)\mapsto z_0$이기 때문에 이를 수식으로 정리하자면 다음과 같이 나타낼 수 있다 (${\widehat{ϵ}}_{\theta }$는 noise prediction model.):

수식을 살펴보면 ddpm 등의 reparameterization trick인 $x_t:=\sqrt{{\overline{\alpha }}_t}{x}_0+\sqrt{1-{\overline{\alpha }}_t}ϵ$을 변형하여 식에 대입한 것을 알 수 있음. ($x$→$z$로 치환)

CM과 마찬가지로 $c_{skip}(0)=1,c_{out}(0)=0$이고 ${\widehat{ϵ}}_{\theta }(z,c,t)$는 teacher diffusion model과 유사한 noise 예측 모델 parameter이다. $f_{\theta }$는 $ϵ-Prediction$ 외에도 $x-Prediction$이나 $v-Prediction$을 사용할 수도 있다. ($x-Prediction$은 DDPM, $v-prediction$은 PD에서 나온 개념)

$\psi (z_t,t,x,c)$는 ODE solver이며 특정한 time step $t\sim s$ 사이에 대한 Eq. 8의 우항을 근사한 값이다. ODE Solver이기 때문에 DDIM, DPM-Solver, DPM-Solver++ 등을 사용할 수 있다. 또한 $\psi$는 학습 및 Distillation시에만 사용한다. 이때 $t_n$은 EDM을 토대로 CM에서 나오는 값이다. 기존 timestep $[t,T]$에 대한 하위 간격으로 $t_1=ϵ<t_2<\cdots <t_N=T$인 어떠한간격을 의미한다. $t_i$는 다음과 같이 나타낼 수 있다:

Eq, 8을 $t_{n+1}\sim t_n$까지 t에 대해 적분 했을 때 다음과 같은 수식을 얻을 수 있다:

3.2 One-Stage Guided Distillation by solving augmented PF-ODE#

Clasifier-free Guidance(CFG)는 high-quality의 conditional 이미지 생성을 가능하게 했다. 다만 CFG는 2개의 Diffusion Model을 훈련해야하기 때문에 효율적이지 못하며, LCMs와 같은 few-step sampling method에 사용하기 힘들다. 따라서 이를 해결하기 위해 본 논문에서는 CFG를 Distillation 과정에서 통합하였다.

Guided-Distill의 경우 two-stage Distillation을 통해 few-step sampling에 CFG를 통합하였으나 학습시간이 길고 2단계를 거치며 손실이 누적되기 때문에 최적의 성능을 내기 힘들다.

t이에 반해 LCMs는 augmented PF-ODE를 해결하는 방식으로 one-stage의 guided Distillation을 제안했다. 일단 CFG에 대한 reverse diffusion process는 다음과 같다:

CFG는 conditional noise 예측값과 unconditional noise 예측값을 선형 결합하여 사용한다. 즉 noise 값이 $\omega$에 따라 변형되므로 augmented PF-ODE라고 한다. augmented PF-ODE는 다음과 같이 나타낼 수 있다:

consistency function도 $\omega$를 변수로 받아오기 때문에 $f_{\theta }:(z_t,\omega ,c,t)\mapsto z_0$로 다시 정의된다. Consistency Distillation Loss 또한 다음과 같이 나타낼 수 있다:

$\omega$와 $n$는 각각 $[{\omega }_{min},{\omega }_{max}]$, $\{1,\dots ,N-1\}$에서 sampling된다. ${\hat{z}}_{t_n}^{\psi ,\omega }$는 이전과 마찬가지로 CFG가 추가된 ODE-Solver를 사용하여 근사한 값을 의미한다. 이때 사용되는 새로운 noise 예측모델 ${\widetilde{ϵ}}_{\theta }(z_t,\omega ,c,t)$는 Eq. 11처럼 $t_{n+1}\sim t_n$까지 t에 대해 적분 했을 때 다음과 같이 나타낼 수 있다:

마찬가지로 PF-ODE Solver $\psi (\cdot ,\cdot ,\cdot ,\cdot )$에는 DDIM, DPM-Solver, DPM-Solver++ 등을 사용할 수 있다.

3.3 Accelerating Distillation with Skipping Time Steps#

Stable Diffusion 등 보통의 Diffusion Model들은 매우 큰 step을 전체 time step으로 잡고 학습한다. 그러나 이같이 촘촘한 time step은 각 $t_n$과 $t_{n+1}$의 변화량을 감소시키기 때문에 자연스럽게 Consistency Distillation Loss도 작아지게 된다. Loss가 작아지면 학습의 수렴속도도 느려지게 된다. 따라서 LCMs는 학습 수렴의 속도를 높이기 위해 time step을 수천에서 수십으로 크기 단축시키는 SKIPPING-STEP 방법을 제안하였다.

기존 CMs 모델의 경우 time scheduler로 EDM을 사용하고 ODE-Solver로 Euler 방법이나 Heun 방법을 사용한다. 그러나 LCMs는 Eq. 8을 통해 DDIM, DPM-Solver, DPM-Solver++와 같은 효율적인 solver도 효과적으로 데이터를 생성할 수 있다는 것을 증명했다. 따라서 SKIPPING-STEP 방법은 $t_{n+1}\to t_n$ 사이의 Consistency를 비교하는것이 아니라 특정 k-step만큼 거리가 있는 time step에 대한 Consistency를 비교한다. ($t_{n+k}\to t_n$)

이때 $k$값의 크기는 trade-off 관계를 가진다. 너무작으면 ($k=1$) 기존과 같이 느린 수렴속도를 갖게되며, 너무 큰 값일 때는 ODE solver 를 통해 근사할 때 오차가 매우 커질수 있다. 논문의 저자는 $k=20$을 사용해 time step을 수천에서 수십으로 대폭 줄여 학습을 Accelerating 할 수 있었다. Eq. 14에 k값을 추가해 SKIPPING-STEP을 표현할 수 있다.

${\hat{z}}_{t_n}^{\psi ,\omega }$에 대한 수식도 다음과 같이 변경할 수 있다.

4.1 Text-To-Image Generation#

3가지 데이터셋에 대한 평가를 진행했다. (LAION-5B, LAION-Aesthetics-6+(12M), LAION-Aesthetics-6.5+(650k)) 앞서말한것처럼 하나의 Resolution이 아닌 512x512, 768x768의 high resolution을 생성했다. 512 size는 $ϵ$-prediction, 768 size는 $v$-prediction을 사용했고 ODE-Solver로는 DDIM을 사용했다. 앞서말한것처럼 SKIPPING-STEP은 20의 값을 가진다.

DDIM, DPM-Solver, DPM-Solver++, Guided-Distill 4가지 모델에 대해 LCM과 성능비교를 했는데 이때 Guided-Distill은 오픈소스 코드가 없기 때문에 논문의 내용과 동일하게 Implementation 해서 성능을 비교하였다. LCM은 같은 메모리 Cost 대비 더 빠르게 수렴하고 더 좋은 품질의 영상을 생성하였다. 특히 Guided-Distill은 2 stage Distillation이지만 LCM은 1 Stage만 사용해도 이같은 성능을 보여줬다.

4.2 Abulation Study#

ODE Solvers & Skipping-Step Schedule#

augmented PF-ODE를 푸는 solver들(DDIM, DPM, DPM++)을 LCM에 사용할 때 성능 비교와 SKIPPING-STEP schedule의 $k$값에 따른 성능 변화를 비교하였다. 모든 모델은 2,000 iteration에서의 4-step inference로 고정해서 비교했다.

Skipping step의 경우 $k$ 값을 올렸을 때 훨씬더 빠르게 수렴하며 때때로 더 좋은 FID 값을 보여주었다. 또한 DPM과 DPM++은 $k$가 50일 때 DDIM보다 더 좋은 성능을 보였다. 이는 $k$ 값이 클수록 더 큰 ODE approximation error를 가지는 DDIM에 비해 오차가 적기 때문이다.

$k=20$일 때, 3가지 모델 모두 좋은 성능이 보였다.

The Effect of Guidance Scale $\omega$#

일반적으로 $\omega$값이 클수록 CLIP score 같은 품질의 지표는 좋아지지만 작을수록 다양성이 떨어져 FID Score가 떨어진다. 즉 $\omega$의 크기는 Quality와 Diversity에 대한 trade-off가 있다.

그래프를 보면 2~8 step inference는 성능에 큰 차이를 가지지는 않는것으로 확인된다. 그러나 1 step inference는 아직 개선의 여지가 있는것을 확인할 수 있다.

$\omega$에 따른 실제 생성 이미지를 비교해 봤을 때 생성 영상의 Quality 차이가 확연하게 들어난다. 즉 Distillation 시에도 CFG를 적용하는 것이 성능을 크게 개선할 수 있다는 것을 증명한다.

4.3 Downstream Consistency Fine-tuning Results#

포켓몬 데이터셋과 심슨 데이터셋에 LCF를 적용했을 때를 비교하였다. 90%는 학습 데이터로, 10%는 검증 데이터로 사용했다. 완벽하진 않지만 Custom Dataset의 style을 잘 catch한 모습을 보여준다.