*Flow Matching Guide and Code( 2412.06264) 를 읽으면서 정리한 내용입니다
<Intro>
1. Flow Matching(FM)이란?
Flow Matching은 생성 모델의 한 종류로, 데이터를 변형하는 과정을 학습하는 방식이다.
이 과정에서 velocity field을 학습하여, 주어진 데이터 분포(source 분포)에서 원하는 데이터 분포(target 분포)로 변환하는 흐름(flow)를 만든다.
Flow Matching의 핵심 과정:
- probability path( pₜ)를 설정하여 소스 분포 → 타겟 분포로 변하는 과정을 정의한다.
- Velocity Field를 학습하여 이 변환 과정을 모방하는 flow를 생성한다.
2. Diffusion Models과는 뭐가 다를까
Diffusion 을 잠깐 복습하고 가보자. Diffusion 모델의 핵심 아이디어는 그냥 쉽게 설명하자면 다음과 같다.
- 데이터에 점점 더 많은 노이즈를 추가하는 forward noising process을 수행한다.
- 이 과정을 거꾸로 되돌리는 방법(reverse process)을 학습하여 깨끗한 데이터를 복원한다.
Flow Matching과 Diffusion Model의 차이점
| Flow Matching(FM) | Diffusion Models (DM) | |
|---|---|---|
| 핵심 아이디어 | 데이터 흐름(Flow) 학습 | 노이즈 추가 & 제거 과정 학습 |
| 학습 방식 | 확률 경로를 직접 학습 | 스코어 함수(Score Function) 학습 |
| 계산 효율성 | 상대적으로 가볍다. | 상대적으로 계산 부담이 크다. |
3. Diffusion Model 발전 과정 & Flow matching과의 관계?
Flow Matching 하는데 왜 자꾸 Diffusion을 언급하냐면, 결국 Flow matching은 디퓨전의 general한 형태이기 때문이다. 추후 블로그에서 더 자세히 다루겠다.
- 처음에는 Discrete-time Gaussian Process로 제안됨 (Sohl-Dickstein et al., 2015; Ho et al., 2020)
- 이후, Continuous-time Stochastic Differential Equation(SDE)으로 정리됨 (Song et al., 2021)
- Flow Matching 관점에서 보면, Diffusion Models도 결국 소스-타겟을 연결하는 경로 pₜ를 학습한다
Flow Matching Key Concept (본격 스타트..)
1. Flow Matching의 목표
어떤 목표 분포 \( q \)가 있다고 가정하자. 이 분포에서 새로운 샘플을 생성할 수 있는 모델을 만들고자 한다.
이를 위해 Flow Matching(FM)은 probability path \( (p_t)_{0 \leq t \leq 1} \)를 구성한다.
- \( p_0 \) (초기 분포) = \( p \) (알고 있는 소스 분포)
- \( p_1 \) (최종 분포) = \( q \) (목표하는 분포)
- 시간 \( t \)에 따라 분포가 \( p_t \)로 변형되며, 이를 Probability Path라 한다.
즉, Flow Matching의 목표는 source 분포 \( p_0 \)를 목표 분포 \( p_1 \)로 변환하는 방법을 학습하는 것이다.
2. Flow&Velocity Field?
Flow랑 Velocity Field 개념이 얼핏 봤을때 헷갈릴수 있지만, Flow Matching의 최종 목표는 Probability Path를 학습하는 것이고, 이걸 그대로 학습할 수는 없으니까, 결국 학습하는거는 시간에 따라 변화하는 Vector Field \( u_t^\theta \) 를 학습하는것이라고 보면된다.
Flow Matching에서는 Vector Field를 사용하여 Flow \( \psi_t \)을 만든다.
\[ \frac{d}{dt} \psi_t(x) = u_t^\theta(\psi_t(x)) \]
즉, 현재 상태 \( x \) 가 vector field \( u_t^\theta(x) \) 에 의해 이동한다.
이 flow를 통해 샘플은 시간이 지나면서 목표 분포로 점진적으로 변화한다.
이 그림이 도움이 될진 모르겠지만, P0(x) 라는 Known distribution에서 P1(x), 즉 데이터 Distribution으로 가기 위해, Flow는 t시점 후에 어디 위치해있는지? Vector Field는 t 시점후에 위치해있는 분포가 어떻게 움직일지? 를 나타낸다고 해석하면 편하다.
Flow \( \psi_t \) 가 샘플을 변형하면, 그에 따라 확률 분포도 변화한다.
이를 통해 Probability Path \( p_t \) 가 정의된다.
\[ X_t := \psi_t(X_0) \sim p_t, \quad X_0 \sim p_0 \]
즉, \( p_0 \) 에서 시작한 샘플이 시간이 지나면서 점점 \( p_1 \) 로 변하게 된다.
근데 그러면 velocity field를 어떻게 학습할까? 위의 그림에서 나와있듯이 두가지 step을 따른다.
- Posibility path \( p_t \) 설계
- \( u_t^\theta \) 학습
1. \( p_t \) 는 뭘로 설계할까
(1) 먼저 Source 분포를 표준정규분포로 가정한다
\[ p = \mathcal{N}(x | 0, I) \]
(2) Posibility Path \( p_t(x) \) 정의
\[ p_t(x) = \int p_{t|1}(x | x_1) q(x_1) dx_1 \]
여기서,
- \( p_{t|1}(x | x_1) \) : \( X_1 = x_1 \)을 조건으로 한 확률 분포
- \( q(x_1) \) : 목표 분포에서 샘플링한 데이터의 분포
이 positblity path는 Conditional Optimal Transport 또는 Linear Path 라고도 한다.
\[ p_{t|1}(x | x_1) = \mathcal{N}(x | t x_1, (1 - t)^2 I) \]
(3) \( X_t \) 정의
\[ X_t = t X_1 + (1 - t) X_0, \quad X_t \sim p_t \]
즉, \( X_0 \sim p \) (초기 분포), \( X_1 \sim q \) (목표 분포)이며 두 샘플을 \( t \)-가중 합(linear combination)으로 구성하여 생성한다.
Flow Matching Loss
\[ \mathcal{L}_{FM}(\theta) = \mathbb{E}_{t, X_t} \left[ \left\| u_t^\theta (X_t) - u_t(X_t) \right\|^2 \right] \]
즉, 모델 예측하는 velocity path \( u_t^\theta(X_t) \) 가 실제 velocity path \( u_t(X_t) \) 와 가까워지도록 학습한다.
근데 여기서 문제가 실제 velocity path \( u_t(X_t) \) 를 어떻게 구할것인가? 이다.
정답은, 데이터 전체를 transform 시키는 velocity path 가 아니라 샘플 하나에 대한 변화를 직접 고려해서 계산을 단순화하는 conditional velocity path를 계산하는 방식으로 해결된다. (이부분 증명은 추후 다루겠다)
Conditional Velocity field \( u_t(x | x_1) \)
\[ X_t | X_1 = x_1 = t x_1 + (1 - t) X_0 \] \[ u_t(x | x_1) = \frac{x_1 - x}{1 - t} \]
Contitional Flow Matching Loss
\[ \mathcal{L}_{CFM}(\theta) = \mathbb{E}_{t, X_t} \left[ \left\| u_t^\theta (X_t) - u_t(X_t | X_1) \right\|^2 \right] \]
이게 아주 신기하게도, gradient를 구해보면, 전체 데이터 velocity field의 gradient를 구하는거나 하나를 잡고 conditional 하게 구하는거나, gradient가 같다고 한다. 그래서 하나만 잡고 conditional하게 구하는 방법이 통하나보다.
최종적으로, 다음과 같은 식이 나온다.
\[ \mathcal{L}_{CFM}^{OT, Gauss}(\theta) = \mathbb{E}_{t, X_0, X_1} \left\| u_t^\theta(X_t) - (X_1 - X_0) \right\|^2, \]
where \( t \sim U[0,1] \), \( X_0 \sim \mathcal{N}(0, I) \), \( X_1 \sim q \).
이 식은 Flow Matching의 가장 간단한 형태로, \( X_0 \) 와 \( X_1 \) 사이의 차이를 직접 모델링하는 방식이다.
여기까지의 내용에 대한 코드를 살펴보자. 전체 코드를 보려면 2412.06264 7p 참고
Probability Path \( p_t(x) \)
Flow Matching에서 확률 경로는 다음과 같이 정의된다:
\[ X_t = (1 - t)X_0 + tX_1 \]
이제 이를 코드로 구현하면 다음과 같다:
x_1 = Tensor(make_moons(256, noise=0.05)[0]) # X_1 ~ q (목표 분포 샘플링)
x_0 = torch.randn_like(x_1) # X_0 ~ N(0, I) (가우시안 샘플링)
t = torch.rand(len(x_1), 1) # t ~ U[0,1] (균일 분포 샘플링)
x_t = (1 - t) * x_0 + t * x_1
Velocity Field
Velocity Field는 데이터를 변환하는 속도를 정의하며, 실제 목표 값은 다음과 같다:
\[ X_1 - X_0 \]
dx_t = x_1 - x_0 # 실제 velocity field
Velocity Field training
Flow Matching에서는 velocity field를 학습하기 위해 다음과 같은 손실 함수를 사용한다:
\[ \mathbb{E} \left[ \| u_t^\theta(X_t) - (X_1 - X_0) \|^2 \right] \]
optimizer.zero_grad()
loss_fn(flow(x_t, t), dx_t).backward()
optimizer.step()
ODE Solver (시간에 따른 샘플 이동)
Flow Matching에서는 ODE(Ordinary Differential Equation) Solver를 사용하여 데이터를 점진적으로 변환한다:
\[ \frac{d}{dt} X_t = u_t^\theta(X_t) \] \[ X_{t+\Delta t} = X_t + \Delta t \cdot u_t^\theta(X_t) \]
def step(self, x_t: Tensor, t_start: Tensor, t_end: Tensor) -> Tensor:
t_start = t_start.view(1, 1).expand(x_t.shape[0], 1)
# Midpoint ODE solver
return x_t + (t_end - t_start) * self(x_t + self(x_t, t_start) * (t_end - t_start) / 2, t_start + (t_end - t_start) / 2)
sampling
x = torch.randn(300, 2)
n_steps = 8
fig, axes = plt.subplots(1, n_steps + 1, figsize=(30, 4), sharex=True, sharey=True)
time_steps = torch.linspace(0, 1.0, n_steps + 1)
axes[0].scatter(x.detach()[:, 0], x.detach()[:, 1], s=10)
axes[0].set_title(f't={time_steps[0]:.2f}')
axes[0].set_xlim(-3.0, 3.0)
axes[0].set_ylim(-3.0, 3.0)
for i in range(n_steps):
x = flow.step(x, time_steps[i], time_steps[i + 1])
axes[i + 1].scatter(x.detach()[:, 0], x.detach()[:, 1], s=10)
axes[i + 1].set_title(f't={time_steps[i + 1]:.2f}')
plt.tight_layout()
plt.show()
