Consistency Model, Yang Song, 2023

---

Background of Diffusion

Diffusion을 처음 접하는 이들을 위하여 간단하게 설명하고자 한다.

Diffusion을 이해하기 전에 다음 논문들을 먼저 보는 것을 추천한다.
VAE와 같은
DDPM 논문
DDIM 논문
VAE 논문

Diffusion의 이해

Diffusion은 노이즈를 점진적으로 제거하면서 데이터를 '생성'하는 확률 생성 모델이다.
완전한 랜덤 노이즈에서 시작해서 Time step에 따라 깨끗한 이미지로 변해가는데, 필자가 이전에 다룬 DnCNN과는 다르다.

노이즈가 자연 이미지 분포에서 더 가능성이 높은 방향으로 이동한다. 즉, 노이즈를 자연 이미지 분포로 끌어당기는 것이다.

Diffusion은 두 가지 과정으로 나눌 수 있다.

• Forward Process
• Reverse Process

Forward Process

Forward Process는 깨끗한 이미지를 완전한 노이즈로 덮어버리는 과정이다.
$x_0 \rightarrow x_1 \rightarrow x_2\rightarrow...\ x_T$ : $x_T \sim N(0, I)$ : 완전한 가우시안 노이즈

• 깨끗한 데이터 $x_0$에서 시작해서, 점점 가우시안 노이즈를 더하여 최종적으로 완전한 노이즈 $x_T$까지 가는 것이다.
• 이 과정은 인위적으로 수식을 통해 노이즈를 주입하는 것이기 때문에, 모델이 학습할 필요가 없는 고정된(Fixed)과정이다.
• SDE 관점에서는 다음과 같이 표현한다.
  $dx_t = f(x_t, t)dt + g(t)dw_t$
  여기서 $dw_t$가 바로 매 순간 추가되는 노이즈이다.

Backward Process

Backward Process는 노이즈로 꽉 찬 이미지에서 의미 있는 이미지를 깎아내는 과정이자, Diffusion의 핵심이다.

• Forward Process를 시간 역으로 뒤집어, 완전한 노이즈 $x_T$에서 시작해 점진적으로 노이즈를 걷어내어 원본 데이터 $x_0$으로 복원하는 과정이다.
• 우리는 이 때 "어느 방향으로 노이즈를 빼야 하는지"알아야 한다. 수학적으로 이 방향을 Score 함수로 알 수 있다. ($\nabla x logp_t(x)$ 이다.)
• 우리는 Score function을 모르기 때문에 데이터로부터 직접 추정해야한다. 때문에 신경망 $s_\theta (x_t,t)$를 학습 시켜서 이 스코어를 예측하도록 한다.

밑에서 이야기 할 Consistency Model이 등장할 수 밖에 없는 한계점을 미리 말하자면, 이 Reverse SDE를 따라가려면 수치 오차를 줄이기 위해 Step을 매우매우 잘게 쪼개야 한다.

결국 기존 Diffusion의 퀄리티는 우수하지만, 수 백번의 Backward Process를 밟아야 해서 실시간 서비스에는 도저히 적용을 할 수 없었다.

Introduction

Diffusion 모델(내지는 Score-based Generative Models)은 최근 이미지 생성 분야에서 큰 성과를 거두고 있다. 누구나 한 번쯤은 GPT와 같은 인공지능에게 "~~ 이미지 만들어줘" 라고 명령을 내려봤을 것이다. (최근 지브리풍 이미지를 생성하는 것처럼 말이다.)

하지만 인공지능에게 이미지를 생성해달라고 명령을 내리면, 수 초 내에 이미지가 생성 되는가?
아니다. 이미지를 생성하는데에는 꽤 오랜 시간이 걸린다.

이렇게 Diffusion 모델이 이미지를 생성하는데에 오랜 시간이 걸리는 이유는 다음과 같다.

• Time step에 따른 PF ODE의 Numerical Integration의 필요성

기존 Diffusion 모델은 노이즈를 '점진적'으로 제거하면서 데이터를 '생성'하는 확률 생성 모델인데, 평균적으로 DDPM은 1000 step, DDIM은 500 step 정도가 걸린다고 한다.

(Consistency Model은 이러한 한계를 극복하기 위해 설계된 새로운 생성 모델로, 노이즈를 데이터로 직접 매핑하는 방식을 통해 설계 단계부터 one-step generation만으로 고품질 샘플 생성을 지원한다.)

---

연구진들은 이러한 Diffusion의 긴 Inference Time을 해결하고자 고민하던 찰나에 무작위적인 SDE와 동일한 확률 분포를 가지면서도 무작위성이 없는 Deterministic 한 ODE가 존재함을 수학적으로 증명해냈다.

* Reverse SDE 방식은 Stochastic한 특성인데, 매 순간 무작위 노이즈 ($dw_t$)가 개입하는 '확률적' process이다. 이러한 '무작위성'은 치명적인 단점이 바로 단번에 정답($x_0$)으로 갈 수 없어서 시간이 매우 오래 걸려 꼭 Stochastic 해야 하는가에 대한 주장이 제기되었다.

여기서 나온 Probability flow (PF) Ordinary differential Equation (ODE)가 바로 Consistency Models의 근간이 되는 이론이다.

Probability Flow (PF) ODE?

PF ODE의 핵심 IDEA는 다음과 같다.


PF ODE는 분포(Probability)의 흐름은 SDE와 똑같이 맞추되, 개별 입자의 Trajectory는 Deterministic한 방정식이다.

[기존 Forward SDE]
$dx_t = f(x_t, t)dt + g(t)dw_t$

• $f(x_t, t)dt$ : 확정적으로 이동하는 방향, Drift
• $g(t)dw_t$ : 무작위로 흔들리는 노이즈
  (Drift는 Deterministic 한데, $dw_t$ 때문에 전체적으로 stochastic 해진다.
  
  [PF ODE]
  $dx_t = [f(x_t, t) - \frac{1}{2}g(t)^2\nabla x logp_t(x_t)]$dt
  
  위 두 식을 비교해보자.
• 바로 $dw_t$가 사라졌다!. 시간에 따른 랜덤한 변화량 $dt$만 있기 때문에 완벽한 Deterministic한 경로가 되었다.
• $dw_t$ 대신 Score Funtion $\nabla x logp_t(x_t)$가 생겼다. 즉, 무작위로 입자를 흩뿌리지 않고 데이터 분포 형태(Score)를 수식에 직접 반영하여 억지로 입자들을 밀어내어 똑같은 분포를 만드는 것이다.
  
  정말 중요한 수학적 성질이 하나 더 있다.
  
  PF ODE 의 Trajectory 위의 adjacent points 한 쌍은 Time step이 0이 될 때까지 적분(Integration)을 하면 같은 Target Point(원본 데이터)를 가리킨다.
  그렇다면 수백 번씩 적분할 필요 없이, 모델에게 궤적의 어떤 점이 들어오든 똑같은 Target으로 output을 내어라 라고 하면, Single-Step 만에 이미지를 만들 수 있지 않을까?
  
  이것이 바로 Consistency Model의 탄생 핵심이다.

Self-Consistency

앞서 우리는 PF ODE를 통해 노이즈 ($x_T$)에서 원본 ($x_0$)으로 가는 Trajectory를 찾았다.

논문의 저자들은 이 궤적의 성질을 이용하여 한 함수를 정의하는데, 이것이 바로 Consistency Function이다.

이 함수의 목표는 다음과 같다.

"PF-ODE Trajectory 위 아무 점 $x_t$를 입력 받으면, 그 Trajectory의 원점인 $x_0$을 바로 출력하기"

수학적 성질 (Consistency Function)

이 함수의 핵심을 수학적으로 보면 다음과 같다.
$f_\theta(x_t, t) = f_\theta(x_{t \prime}, t \prime)$

여기서 논문이 중요하게 말하는 Consistency가 나온다.

"시간이 어떻든, 결과값은 똑같아야 한다."

결국 똑같은 소리다.. (그만큼 중요하다)

Boundary Condition

Consistency Model이 수학적으로 유효하기 위해서는 경계 조건 ($f(x, \epsilon)=x$ 를 반드시 만족해야 한다. 이는 궤적의 종착점인 $t=\epsilon$ 에서 모델의 출력이 입력값 자신으로 수렴해야 함을 의미하며, 이러한 구속력은 모델이 학습 과정에서 모든 출력을 0으로 수렴시키는 Trivial solution에 빠지는 것을 방지하는 수학적 장치이다.

그냥 $t=\epsilon$일 때 원본 이미지랑 같도록 Loss Function에 추가해서 학습 시키면 안되나?
: 신경망($\theta$)에게 이 당연한 사실을 수만 번 가중치 업데이트로 깨닫게 하는것은 자원 낭비가 심하다.

Skip Connection 기반 파라미터화

위와 같은 이유로 신경망이 학습할 필요가 없도록, 모델의 출력 구조 자체를 Parameterization 한다. 단순한 심층 신경망 출력으로는 위에서 살펴본 Condition을 보장할 수 없기 때문이기도 하다.

$f_\theta(x_t, t) = c_{\text{skip}}(t)x_t + c_{\text{out}}(t)F_\theta(x_t, t)$

• $F_\theta$ : 딥러닝 신경망
• $x_t$ : 모델에 들어온 입력 이미지
• $c_{\text skip}(t), c_{\text out}(t)$ : 시간 $t$에 따라 변하는 단순한 가중치 함수 (스케쥴러)

여기서 $c_{\text skip}(t), c_{\text out}(t)$ 는 미분 가능한 함수이며, $c_{\text skip}(\epsilon) = 1$, $c_{\text out}(\epsilon) = 0$ 조건을 엄격히 준수해야 한다.

이러한 설계는 다음과 같은 이점을 제공한다.

1. 경계 조건의 강제 이행 별도의 손실함수 없이 아키텍처 구조 자체로 $t=\epsilon$ 에서의 Identity Mapping을 보장한다.

2. 연속 시간 미분 가능성 (Differentiability) : $t=\epsilon$에서 모델의 미분 가능성을 확보하는 것은 Continuous-time 학습 시 경사항 안정성과 수렴 특성을 개선하는 데에 결정적인 역할을 한다.

3. 기존 아키텍처 수용성 : EDM (Elucidating the Design Space of Diffusion-Based Generative Models)와 같은 기존의 고성능 확산 모델 백본($F_\theta$)을 유연하게 통합할 수 있다.

기존의 Diffusion 모델은 Target을 한 번에 알지 못했다. 떄문에 여러 번의 Step을 Integration 한 것을 기억할 것이다.

하지만, 만약 위 Self-Consistency 조건을 완벽하게 학습한다면, 1-Step Generation이 가능할 것이다.

모델 훈련 방법

앞서 우리는 모델의 가장 중요한 규칙, Self-Consistency를 배웠다. 그렇다면 이 규칙을 어떻게 모델에게 train 할까?

논문에서는 두 가지 방법이 나온다.

방법 1. Consistency Distillation (CD)

가장 확실하고 좋은 방법은 pre-trained Diffusion model을 사용하는 것이다. 이 과정을 증류 (Distillation)이라고 부른다.

1. 같은 궤적 상 두 점 뽑기
   먼저 노이즈가 잔뜩 낀 출발점 A에서 Teacher Model (pre-trained model)이 자신의 지식과 ODE Solver를 사용하여 다음 지점 B를 찍어준다. 이렇게 같은 Trajectory 상의 인접한 한 쌍의 점이 완성된다.

2. Student Model과 Target Model

   • Student Model (우리가 학습하고자 하는 모델) : 출발점 A를 보고 스스로 목적지를 예측한다.
   • Target Model : Teacher Model이 찍어준 점 B를 보고 똑같이 예측한다. (Target 모델은 Student Model의 EMA (Exponential Moving Average)기법을 사용한 복제품이라고 생각하면 된다.)

두 네트워크 사이의 출력 차이를 최소화함으로써 PF ODE 궤적 전반에 걸친 일관성을 확보한다.

왜 두 모델을 사용할까? 그냥 Student 하나만 사용하면?
꼭 Student Model의 EMA 복제품을 사용해야할까?

그렇다. 만약 Student Model 혼자 Teacher Model이 준 두 점을 혼자서 비교하게 되면 학습이 불안정해질 수 있다.
양쪽 점이 동시에 움직여서 Loss의 기준점이 흔들리기 때문에 한쪽은 Gradient를 받지 않는 고정된 Network가 필요한 것이다. (Target)

ODE Solver의 수렴 특성 및 이론적 근거 (Theorem 1)

ODE Solver는 인접한 점을 추정하기 위해 사용되는데, 이는 모델의 정확도에 매우 큰 영향을 미친다.

논문 Theorem 1에 따르면, ODE Solver의 Local Error 가 $O(\Delta t^{p+1}$ 일 때, 최종 Consistency Function의 추정 오차는 $O(\Delta t^p)$로 결정된다.
- Euler Solver ($p=1$) : 1차 근사 방식으로 구현은 용이하나, 오차 누적이 발생하기 쉽다.
- Heun Solver ($p=2$) : 2차 고차 근사 방식으로, 동일한 이산화 단계 ($N$)에서 Euler 대비 훨씬 낮은 Global Error를 달성한다. 따라서 고품질 Distillation을 위해서는 Heun Solver가 우월하다고 볼 수 있다.

방법 2. Consistency Training (CT)

pre-trained model도 좋지만, 단점이 있다.
해당 사전학습 모델이 좋은 모델이어야 하고, 해당 Teacher Model이 궤적을 계산하는 데에 시간과 비용이 든다.
(우리는 시간과 비용을 줄이려고 했는데.. 그 반대의 상황이 나온 것이다.)

때문에 저자들은 Teacher Model 없이 원본 데이터만 가지고 독학하는 방법도 생각해냈다.

Teacher 없이 어떻게 가능할까? 우리는 Forward Noising Process는 알고 있다. 때문에 $x_0$가 Gaussian으로 변환하는 과정을 통해 가능한 것이다.

• Teacher Model이 없으니 PF ODE Trajectory는 모른다. 대신, 단순한 방법으로 가정을 한다.

  그냥 원본 데이터($x_0$)를 향해 일직선으로 노이즈를 더하고 빼자.

• 원본 데이터 하나를 꺼내서, 점 두 개를 인위적으로 만든다.
  A(노이즈 심한 점), B(노이즈 약한 점)
  그리고 이 두 점이 똑같은 원본 데이터를 가리키도록 강제한다.

이 방법은 방법 1 대비 속도가 빠르지만, Distillation 방식 보다는 생성된 이미지의 퀄리티가 약간 떨어지는 경향이 있다.

---

다단계 Sampling (Algorithm 1)

Consistency Model의 핵심은 바로 1-Step 생성이다.

하지만, 속도와 퀄리티는 반비례 하기 마련으로, 디테일이 아쉬울 수 있는데, 이를 위해 다단계 샘플링 (Multistep Consistency Sampling, Algorithm 1)을 제시한다.

1. 완전한 노이즈 $x_T$를 모델에 넣어 1차 완성본 ($x_0$)를 예측한다.
2. 예측된 $x_0$이 완벽하지 않다고 가정하고, 여기에 미리 계획해 둔 시간 ($T_2$)만큼의 노이즈를 다시 섞어준다.
3. 노이즈가 살짝 낀 상태에서 다시 모델에 통과시켜 2차 완성본 ($x_0$)를 예측한다.
4. 이 과정에서 내가 원하는 스텝 수만큼 반복한다. 마지막 스텝에 도달하면 더 이상 노이즈를 섞지 않고 최종 이미지를 반환한다.

Multi-Sampling Trade-off

이 기법은 생성 품질과 계산 비용의 완벽한 Trade-off를 가능하게 했다.

기존 Diffusion은 50번 돌리다 25번에서 멈추면 치지직 거리는 노이즈 화면이 나오는 반면 Consistency Model은 다르다.

• 1-step : 초고속 렌더링이 필요한 실시간 게임 엔진에 사용
• 2~4 step : 적당한 화질과 속도가 필요한 모바일 앱
• 8 step 이상 : 극강의 초고화질이 필요한 예술 작업

즉, Consistency Model은 단순히 빠른 모델이 아니라, 사용자가 상황에 맞게 조절할 수 있는 생성 모델이다.