[논문 리뷰] Consistency Models 리뷰

안녕하세요, 오늘은 OpenAI에서 발표한 Consistency Models을 소개합니다. 이 모델은 기존 Diffusion models이 노이즈로부터 원본 이미지 복원을 위해 수백번 ~ 수천번의 iteration을 반복하는 과정을 획기적으로 줄였다고 합니다.

저는 생성 모델의 전문가가 아니라서 심도 있는 내용은 없지만 본 설명이 이해에 도움이 된다면 좋겠습니다.

• 논문: https://arxiv.org/pdf/2303.01469.pdf
• Github: https://github.com/openai/consistency_models

Introduction

최근 득세하고 있는 Diffusion Models... 그 이전에 한참을 유행했던 GAN(Generative Adversarial Network)에는 없는 장점을 가졌습니다. 대충 살펴보면

• do not rely on adversarial training and are thus less prone to issues such as unstable training and mode collapse
• do not impose the same strict constraints on model architectures as in autoregressive models

대충 정리해보면

VAE, GAN, Diffusion의 특징

이렇듯 Diffusion models은 느리게 동작한다는 것만 빼면 참 나무랄 데 없는 좋은 모델입니다. 하지만 관점에 따라 그 단점이 치명적일 수 있습니다. 요즘은 모델도 큼지막하고 그로 인한 환경 파괴도 심하기 때문에 다른 장점들을 그대로 가져가면서 빠른 샘플링이 가능하다면 정말 혁신이겠죠?

그런데 Diffusion models은 왜 Fast Sampling이 불가능할까요?

수학적 자세한 내용은 여기 또는 여기를 참고하시면 되지만 짧게 말하면 노이즈로부터 원본 방향으로 Denoising 하는 과정

$$q\left (\mathbf{x}_{t-1} | \mathbf{x}_{t} \right )$$

에서 noise \(\beta\)가 매우 작다는 가정 하에 Normal distribution을 사용할 수 있기 때문입니다. 때문에 기존의 Diffusion models은 수백 ~ 수천번의 iteration을 돌며 학습을 했습니다.

그럼 Consistency models은 이걸 어떻게 해결한건가요?

Probability Flow (PF) ODE가 주어진다면 해당 flow 내의 어떠한 샘플도 원본으로 복원할 수 있다!

위의 그림을 보면 느낌을 알 수 있는데, Probability Flow (PF) ODE가 주어졌을 때 해당 flow 내의 모든 샘플(\(x_{t}, t), t \in [1, T]\)은 원본 (\(x_{0}, 0)\)으로 한 번에 복원할 수 있게끔 학습하는게 포인트입니다.

Probability Flow (PF) ODE

이 논문에서 참 많이 나오는 개념인데, 여기서 처음 제안된 내용은 아닙니다. 이 논문의 저자 Yang Song이 이미 이전 논문에서 제안했던 내용인데요, Stochastic Differential Equation (SDE) of diffusion models는 아래와 같고,

$$d\mathbf{x}_{t} = \mu\left (\mathbf{x}_{t}, t \right )\mathup{d}t + \sigma \left ( t \right ) \mathup{d}\mathbf{w}_{t} $$

여기서

• \(\mu\): Drift coefficient
• \(\sigma\): Diffusion coefficient
• \(\mathbf{w}_{t}\): Standard brownian motion

입니다. 그래서 결국 시간에 따른 물리적 확산(diffusion)을 수식으로 표현하면

$$ \mathup{d}\mathbf{x}_{t} = \left [ \mu\left (\mathbf{x}_{t}, t \right )\mathup{d}t - \frac{1}{2}\sigma(t)^{2}\nabla\mathup{log}p_{t}(\mathbf{x}_{t})\right ] \mathup{d}t$$

입니다. 그리고 이걸 Probability Flow (PF) ODE라고 합니다. 아무튼 diffusion을 시간에 대한 식으로 표현한 형태입니다.

 

아무튼 이 연구 consistency models은 아래와 같은 특징을 갖고 있습니다.

• Single-step generation
• 하지만 여전히 Iterative generation 또한 가능

---

Methods

그래서 Consistency model은 trajectory가 주어졌을 때 모든 \(t\)에 대해 \(\mathbf{\mathit{f}} : (\mathbf{x}_{t}, t) \mapsto \mathbf{x}_{\epsilon}\)를 찾는 과정입니다. 참고로 모델의 학습은 Distillation / Isolation 두 가지 방법이 제안됐습니다. 서로 접근이 좀 상이한데요... 그럼 어떻게 학습을 시키는 지 살펴보겠습니다.

Distillation Training

Distillation Training(Consistency Distillation, CD)은 이미 학습된 diffusion models이 있다는 것을 전제 하에 진행합니다. 이 글을 읽는 분들께서는 Diffusion models에 대한 기본적인 개념이 있다는 것을 가정하고 바로 Loss function부터 들어가면,

$$\mathcal{L}_{CD}^{N}(\mathbf{\theta, \theta^{-};\phi}) := \mathbb{E}[\lambda(t_{n})d(\mathbf{f_{\theta}}(\mathbf{x}_{t_{n+1}}, t_{n+1}), \mathbf{f_{\theta^{-}}}(\mathbf{\hat{x}}_{t_{n}}^{\phi}, t_{n}))]$$

입니다.

너무 어렵죠?? 하나씩 풀어보면

 

• \(\theta\): Consistency models의 파라미터
• \(\theta^{-}\): 학습 과정에서 얻는 모델 파라미터 \(\theta\)의 running average
• \(\phi\): 미리 학습해놓은 Diffusion models의 파라미터
• \(\lambda\): 결과가 항상 양수인 가중치 함수 \(\lambda(\cdot) \in \mathbb{R}^{+}\)
• \(\mathbf{\hat{x}}_{t_{n}}^{\phi}\): Diffusion models에서의 \(t_{n}\)번째 timestep feature
• \(d(\cdot, \cdot)\): 두 값 사이의 거리를 의미하는 metric function (L1, L2 distance 등을 말합니다)

한마디로 Diffusion models의 결과(\(\mathbf{\hat{x}}_{t_{n}}^{\phi}\))가 어떤 trajectory 내에 있다는 것을 알고 이것을 직접 활용하여 Consistency models의 그 다음 step을 알아내는데 활용한다는 뜻입니다.

Consistency Distillation 학습 과정

Isolation Training

Isolation Training(Consistency Training, CT)은 다른 모델 없이 처음부터 학습하는 것을 말합니다. 이 부분은 논문에 나와 있는 것만 읽어보면

사실 CD는 Score-based diffusion models을 활용했는데, 이 score를 대체할 수 있다면 누구의 힘도 빌리지 않고 단독으로 학습이 가능합니다. 그 방법으로 제안하는게 Monte Carlo Estimate입니다.

$$\nabla\mathup{log}p_{t}(\mathbf{x}_{t}) = -\mathbb{E}\left [ \frac{\mathbf{x}_{t} - \mathbf{x}}{t^{2}} \bigg| \mathbf{x}_{t}\right ]$$

즉, 랜덤하게 많은 수의 샘플링을 한다면 충분히 위의 score를 대체할 수 있다고 합니다.

Isolation 학습 과정

---

Results

논문에서는 정말 다양한 실험을 진행했습니다. Metric function을 바꿔보거나, step의 수를 바꿈에 따라서 달라지는 결과를 첨부했습니다. 세로 축의 FID는 Frechet Inception Distance로써 낮을 수록 좋습니다.

특히 CD의 경우 아무리 기존 diffusion models을 활용했다고는 하나 정말 적은 수의 step임에도 좋은 결과를 보여줬다는 것을 확인할 수 있습니다.

 

또한 기존 모델과의 정량적/정성적 비교 또한 잊지 않았습니다.

FID나 IS를 보면 "뭐야 성능이 너무 낮잖아?"라고 생각하실 수 있습니다. 하지만 이 논문의 의의는 NFE(Network Function Evaluation)에 있습니다. NFE는 forward pass 중에 전체 모델 파라미터가 몇번이나 계산되었는가를 의미하는 지표인데, 우리가 알고 있는 diffusion models은 하나의 모델을 수없이 iteration 돌기 때문에 숫자가 매우 높은 것을 알 수 있습니다. 하지만 CT, CD는 고작 한두번 돌고 저정도 성능을 보인다니 정말 매력적입니다.

여타 diffusion models처럼 색칠도 해주고 upsampling도 해주고, 빈 공간도 채워주는 기특한 기능을 갖고 있습니다.

Reference

https://developer.nvidia.com/blog/improving-diffusion-models-as-an-alternative-to-gans-part-1/

https://happy-jihye.github.io/diffusion/diffusion-1

https://openreview.net/pdf/ef0eadbe07115b0853e964f17aa09d811cd490f1.pdf