Flow Matching vs DDPM: 왜 ODE가 더 빠르고 직관적인가

Flow Matching vs DDPM: 왜 ODE가 더 빠르고 직관적인가
DDPM은 1000스텝, Flow Matching은 10스텝. 수식으로 증명하는 직선 경로의 힘.
TL;DR
- DDPM: 노이즈를 조금씩 제거하는 확률적(Stochastic) 방식. 매 스텝 랜덤성이 개입해 경로가 휘어짐
- Flow Matching: 데이터로 직접 향하는 결정론적(Deterministic) 방식. 직선 경로로 빠르게 도달
- 핵심 차이: DDPM은 "노이즈 예측", Flow Matching은 "속도장(velocity field) 예측"
1. 문제 설정: 노이즈에서 데이터로 가는 길
생성 모델의 목표는 단순합니다:
$$\text{노이즈 } z \sim \mathcal{N}(0, I) \quad \longrightarrow \quad \text{데이터 } x \sim p_{\text{data}}$$
이 변환을 어떻게 할 것인가? 여기서 두 패러다임이 갈립니다.
DDPM의 접근: "천천히, 확률적으로"
DDPM은 Markov chain을 정의합니다:
$$x_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1 - \bar{\alpha}_t} \epsilon, \quad \epsilon \sim \mathcal{N}(0, I)$$
여기서 $\bar{\alpha}_t = \prod_{s=1}^{t} \alpha_s$이고, $\alpha_t = 1 - \beta_t$입니다.
시간이 지날수록($t \to T$) 데이터 $x_0$의 정보는 사라지고, 순수한 노이즈 $\epsilon$만 남습니다.
Flow Matching의 접근: "직선으로, 결정론적으로"
Flow Matching은 선형 보간(linear interpolation)을 사용합니다:
$$x_t = (1 - t) x_0 + t \epsilon, \quad t \in [0, 1]$$
$t=0$이면 $x_0$ (데이터), $t=1$이면 $\epsilon$ (노이즈). 그 사이는 직선입니다.
2. 학습 목표의 차이
DDPM: 노이즈 예측
DDPM은 신경망 $\epsilon_\theta$가 추가된 노이즈를 예측하도록 학습합니다:
$$\mathcal{L}_{\text{DDPM}} = \mathbb{E}_{x_0, \epsilon, t} \left[ \| \epsilon - \epsilon_\theta(x_t, t) \|^2 \right]$$
왜 노이즈를 예측할까요? Reverse process에서 $x_{t-1}$을 복원하려면:
$$x_{t-1} = \frac{1}{\sqrt{\alpha_t}} \left( x_t - \frac{\beta_t}{\sqrt{1 - \bar{\alpha}_t}} \epsilon_\theta(x_t, t) \right) + \sigma_t z$$
여기서 $z \sim \mathcal{N}(0, I)$는 매 스텝 추가되는 새로운 랜덤성입니다. 이것이 경로를 휘게 만드는 원인입니다.
Flow Matching: 속도장 예측
Flow Matching은 신경망 $v_\theta$가 속도장(velocity field)을 예측하도록 학습합니다:
$$\mathcal{L}_{\text{FM}} = \mathbb{E}_{x_0, \epsilon, t} \left[ \| v_t(x_t) - v_\theta(x_t, t) \|^2 \right]$$
목표 속도장(target velocity field)은 조건부 경로의 시간 미분입니다:
$$v_t(x_t | x_0, \epsilon) = \frac{d}{dt} x_t = \frac{d}{dt} \left[ (1-t)x_0 + t\epsilon \right] = \epsilon - x_0$$
이 속도장은 상수입니다! 시간 $t$에 관계없이 항상 $\epsilon - x_0$ 방향으로 일정한 속도로 이동합니다.
3. 샘플링 과정의 차이
DDPM 샘플링: SDE 기반
DDPM의 reverse process는 Stochastic Differential Equation (SDE)를 따릅니다:
$$dx = \left[ f(x, t) - g(t)^2 \nabla_x \log p_t(x) \right] dt + g(t) d\bar{w}$$
여기서:
- $f(x, t)$: drift coefficient
- $g(t)$: diffusion coefficient (노이즈 강도)
- $d\bar{w}$: reverse-time Brownian motion
문제점: $g(t) d\bar{w}$ 항 때문에 매 스텝 랜덤성이 추가됩니다. 경로가 브라운 운동처럼 휘어지면서 목표에 도달하므로, 작은 스텝으로 많이 가야 합니다.
Flow Matching 샘플링: ODE 기반
Flow Matching은 Ordinary Differential Equation (ODE)를 따릅니다:
$$\frac{dx}{dt} = v_\theta(x, t)$$
확률 항이 없습니다. 주어진 속도장을 따라 결정론적으로 이동합니다.
샘플링:
# Euler method
x = torch.randn(batch_size, dim) # 시작: 순수 노이즈
dt = 1.0 / num_steps
for t in torch.linspace(1, 0, num_steps):
v = model(x, t) # 속도장 예측
x = x - v * dt # 직선 이동4. 왜 Flow Matching이 더 빠른가?
수학적 직관
DDPM의 경로 길이를 생각해봅시다. 브라운 운동의 특성상:
$$\mathbb{E}\left[ \text{Path Length} \right] = \mathcal{O}(\sqrt{T})$$
$T$는 총 스텝 수입니다. 스텝이 많을수록 경로가 길어집니다.
반면 Flow Matching의 직선 경로:
$$\text{Path Length} = \| \epsilon - x_0 \| = \mathcal{O}(1)$$
스텝 수와 무관하게 최단 거리입니다.
실험적 증거
동일한 품질을 위해 DDPM은 1000스텝이 필요하지만, Flow Matching은 10스텝으로 충분합니다.
5. Rectified Flow: Flow Matching의 진화
Rectified Flow는 Flow Matching을 더 발전시킨 방법입니다.
핵심 아이디어: Reflow
학습된 flow가 완벽히 직선이 아닐 수 있습니다. Reflow는 이를 "펴주는" 과정입니다:
- 학습된 모델로 $(z, x_0)$ 쌍을 생성
- 이 쌍으로 다시 직선 경로를 학습
- 반복하면 경로가 점점 직선에 가까워짐
$$\mathcal{L}_{\text{reflow}} = \mathbb{E}_{(z, x_0) \sim \pi_k} \left[ \| (x_0 - z) - v_\theta(x_t, t) \|^2 \right]$$
Distillation과의 결합
1-step 생성을 위해 distillation을 적용할 수 있습니다:
$$\mathcal{L}_{\text{distill}} = \mathbb{E}_{z} \left[ \| x_0^{\text{teacher}} - G_\theta(z) \|^2 \right]$$
여기서 $G_\theta(z)$는 단일 forward pass로 데이터를 생성하는 generator입니다.
6. 구현 비교
DDPM Forward Process
def ddpm_forward(x0, t, noise_schedule):
"""
x_t = sqrt(alpha_bar_t) * x0 + sqrt(1 - alpha_bar_t) * epsilon
"""
alpha_bar = noise_schedule.alpha_bar[t]
epsilon = torch.randn_like(x0)
x_t = torch.sqrt(alpha_bar) * x0 + torch.sqrt(1 - alpha_bar) * epsilon
return x_t, epsilonFlow Matching Forward Process
def flow_matching_forward(x0, t):
"""
x_t = (1 - t) * x0 + t * epsilon
"""
epsilon = torch.randn_like(x0)
# t를 올바른 shape으로 변환
t = t.view(-1, 1, 1, 1) # for images
x_t = (1 - t) * x0 + t * epsilon
velocity = epsilon - x0 # target velocity
return x_t, velocityTraining Loop 비교
# DDPM
for x0 in dataloader:
t = torch.randint(0, T, (batch_size,))
x_t, epsilon = ddpm_forward(x0, t, noise_schedule)
epsilon_pred = model(x_t, t)
loss = F.mse_loss(epsilon_pred, epsilon)
# Flow Matching
for x0 in dataloader:
t = torch.rand(batch_size) # uniform [0, 1]
x_t, velocity = flow_matching_forward(x0, t)
velocity_pred = model(x_t, t)
loss = F.mse_loss(velocity_pred, velocity)7. 언제 무엇을 써야 하나?
DDPM/DDIM을 선택할 때
- 기존 pretrained 모델 활용 (Stable Diffusion 등)
- 높은 다양성(diversity)이 필요할 때
- 확률적 샘플링이 필요한 경우
Flow Matching을 선택할 때
- 빠른 추론이 최우선일 때
- 처음부터 새로 학습할 때
- 단순하고 직관적인 구현이 필요할 때
Rectified Flow를 선택할 때
- 1-step 또는 few-step 생성이 목표일 때
- 실시간 애플리케이션
- 모바일/엣지 디바이스 배포
8. 수학적 연결: Score와 Velocity의 관계
DDPM의 score function과 Flow Matching의 velocity는 밀접하게 연관되어 있습니다.
Score function은 로그 확률의 기울기입니다:
$$s_\theta(x, t) = \nabla_x \log p_t(x)$$
DDPM에서 score와 노이즈 예측의 관계:
$$s_\theta(x_t, t) = -\frac{\epsilon_\theta(x_t, t)}{\sqrt{1 - \bar{\alpha}_t}}$$
Flow Matching의 velocity와 score의 관계 (probability flow ODE를 통해):
$$v_\theta(x, t) = f(x, t) - \frac{1}{2} g(t)^2 s_\theta(x, t)$$
따라서 잘 학습된 DDPM 모델에서 Flow Matching 모델을 유도할 수 있습니다. 이것이 Stable Diffusion 3가 Rectified Flow로 전환한 이유 중 하나입니다.
결론
Flow Matching은 "왜 굳이 돌아가나?"라는 질문에서 시작된 방법론입니다. 직선이 최단 거리라는 단순한 통찰이 생성 모델의 효율성을 크게 높였습니다.
References
- Ho, J., et al. "Denoising Diffusion Probabilistic Models" (NeurIPS 2020)
- Lipman, Y., et al. "Flow Matching for Generative Modeling" (ICLR 2023)
- Liu, X., et al. "Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow" (ICLR 2023)
- Song, Y., et al. "Score-Based Generative Modeling through Stochastic Differential Equations" (ICLR 2021)