From Flow Matching to Optimal Transport: A Physics-based View of Generative Models

Hwan Heo (허환)·2025년 7월 16일
2

3D AI

목록 보기
25/25
post-thumbnail
  • Normalizing Flow 부터 Continuous Normalizing FlowConditional Flow Matching, 그리고 Rectitifed Flow 가 왜 Linear Path 를 채택하는지 수학적, 물리적으로 상세하게 이해해보자.

  • Flow Matching 은 그 이론적 깊이 때문에 쉽게 접근할 수 있는 설명 자료가 아직 부족합니다. 많은 연구자, 개발자들이 '왜 하필 직선 경로가 효과적인가?' 라는 근본적인 질문에 대한 깊이 있는 직관을 갖기는 어렵습니다. 이 글은 그런 실정에서 수학, 물리학, 그리고 최신 논문에 흩어져 있는 Flow Matching 의 핵심 아이디어들을 하나의 유기적인 이야기로 엮어내어, Flow Matching 의 'Why' 에 대한 직관을 얻어 가는 것을 목표로 삼고 씌여졌습니다.


Introduction

최근 생성 모델의 흐름은 크게 두 가지 접근법으로 나눌 수 있다.

하나는 노이즈에서부터 점차 이미지를 생성해나가는 Diffusion Model 이고, 다른 하나는 노이즈와 이미지를 직접적으로 잇는 Flow 기반 모델이다.

특히 Diffusion Model 은 복잡한 Stochastic Differential Equation (SDE) 을 기반으로 한다. 이는 마치 입자가 무작위적인 힘에 의해 이리저리 흔들리며 움직이는 'Brownina Motion' 처럼, 구불구불하고 예측 불가능한 경로를 통해 샘플을 생성하는 방식이다. 이는 뛰어난 성능을 보였지만, 본질적으로 수많은 단계를 거쳐야 하기 때문에 느린 생성 속도라는 단점을 안고 있었다.

이러한 배경 속에서, "더 빠르고, 더 효율적인 생성 방법은 없을까?"라는 질문에 대한 대답으로 Ordinary Differential Equation (ODE) 기반의 framework 가 주목받기 시작했다. ODE 는 SDE 와 달리 무작위성 (stochastic) 이 배제된 deterministic path 를 따른다. 즉, 출발점과 도착점이 정해지면 그 사이의 경로는 단 하나로 명확하게 결정된다.

  • Top: Diffusion <-> Middle & Bottom: Flow-Matching. Stochastic 과 deterministic 의 차이를 잘 볼 수 있다.

이 단순하고 강력한 아이디어를 기반으로 한 Rectified Flow 는 등장과 함께 가장 주목받는 state-of-the-art 생성 프레임워크 중 하나로 자리 잡았다. Rectified Flow 는 노이즈 z0z_0 와 target z1z_1 사이의 무수히 많은 가능한 ODE 경로 중, 가장 단순한 직선 경로를 가정한다. 이 덕분에 학습은 안정화되고 샘플링 속도는 비약적으로 빨라졌다.

이번 글에서는 이 ODE 기반 생성 모델의 한 축을 담당하는 Flow Matching 에 대해 먼저 알아보며, 그 중 한 형태인 Rectified Flow 가 왜 그토록 단순하고 강력한 '직선 경로' 를 채택했는지 파헤치기 위해 Optimal Transport 이론까지 깊이 살펴보도록 하겠다.

Part A. Flow Matching


1. Normalizing Flow

1.1. What is Flow

Flow 기반 모델의 아이디어는 단순하고 우아하다.

"쉽게 다룰 수 있는 단순한 확률 분포 (e.g., gaussian distribution) 를, 어떤 복잡하지만 근사하고 싶은 실제 데이터 분포로 '변환'하는 함수를 학습할 수 있을까?"

이 '변환 함수' ϕ\phi 를 학습하는 과정 자체를 Normalizing Flow 라고 부른다.

  • Input: noise distribution p0p_0 에서 뽑은 샘플 x0x_0
  • Output: 변환 함수를 통과한 결과 x1=ϕ(x0)x_1 = \phi(x_0)
  • Goal: 이 결과물 x1x_1의 분포가 real data distribution p1p_1 과 같아지도록, 즉 그럴듯한 이미지가 생성되도록 ϕ\phi 를 학습시킨다.

따라서, 학습 목표는

ϕ(x0)=x1p1\phi(x_0) = x_1 \sim p_1

관계를 만족하는 continuous, differentiable, invertible 한 mapping function ϕ\phi 를 찾는 것이다.

이 목표를 달성하려면, 변환된 샘플 x1x_1이 데이터 분포 p1p_1에 속할 확률, 즉 likelihood p1(x1)p_1(x_1) 을 계산하고 이를 최대화해야 한다.


1.2. Likelihood of Normalizing Flow

이제 Change of Variables Theorem 를 이용하여 x1x_1x0x_0 의 관계를 정리하면, 변환 ϕ\phi 전후의 infinitesimal probability 이 동일해야한다는 관계를 얻는다 (probability mass).

p1(x1)dx1=p0(x0)dx0p_1(x_1)|dx_1| = p_0(x_0)|dx_0|

이를 p1(x1)p_1(x_1) 에 대해 정리하고, x0=ϕ1(x1)x_0 = \phi^{-1}(x_1) 관계를 이용해 정리하면 다음과 같다.

p1(x1)=p0(ϕ1(x1))det(ϕ1(x1)x1)p_1(x_1) = p_0(\phi^{-1}(x_1)) \left| \det\left( \frac{\partial \phi^{-1}(x_1)}{\partial x_1} \right) \right|
  • non-linear function ϕ1\phi^{-1}x1x_1 근방에서 국소적으로 linear transformation 으로 근사할 수 있을 것이다. 이 국소적 선형 변환을 나타내는 행렬이 바로 Jacobian

    Jϕ1(x1)=ϕ1(x1)x1J_{\phi^{-1}}(x_1) = \frac{\partial \phi^{-1}(x_1)}{\partial x_1}

    이다.

  • 여기서 det()\det(\cdot) 는 변환에 따른 space 의 부피 변화를 보정해주는 값인데, 선형 대수학에서 한 선형 변환이 공간의 부피를 얼마나 변화시키는지는, 그 변환을 나타내는 행렬의 determinant 값으로 측정되기 때문이다 (cf: Geometric Meaning of the Determinant).

이 식의 양변에 로그를 취하면, 우리가 loss function 으로 사용하여 최대화해야 할 log-likelihood 를 얻게 된다.

logp1(x1)=logp0(ϕ1(x1))+logdet(ϕ1(x1)x1)\log p_1(x_1) = \log p_0(\phi^{-1}(x_1)) + \log \left| \det\left( \frac{\partial \phi^{-1}(x_1)}{\partial x_1} \right) \right|

하지만 이 접근법에는 두 가지 큰 난관이 존재한다.

  1. Invertibility: 변환 함수 ϕ\phi 의 역함수 ϕ1\phi^{-1} 를 계산할 수 있어야 한다.
  2. 계산 비용: Jacobian determinant 의 계산 비용O(D3)\mathcal{O}(D^3) 이 너무 커서 고차원 데이터에는 적용하기 어렵다.

2. Continuous Normalizing Flow

2.1. Continuous Flow

이 문제를 해결하기 위해,

"하나의 거대한 변환 함수 ϕ\phi 대신, 아주 작은 변화들을 무한히 연결한 'Flow' 로 생각하면 어떨까?"

라는 아이디어가 등장했다.

이는 변환 과정을 시간의 흐름에 따른 continuous trajectory 으로 모델링하는 것을 의미한다.

이 trajectory 위에서 샘플 x0x_0 은 시간 t=0t=0 에서의 위치이고, 우리가 원하는 데이터 x1x_1 는 시간 t=1t=1 에서의 최종 위치가 된다.

즉, transformation ϕ\phi 는 다음과 같이 multiple function 의 composition 으로 표현되고,

ϕ=ϕ(T1)t ϕ1tϕ0t\phi = \phi_{(T-1)\cdot \triangle t} \ \circ \dots \circ \phi_{1 \cdot \triangle t} \circ \phi_{0 \cdot \triangle t}

각 function 에는 다음과 같은 관계가 성립합다.

ϕt(xt)=xt+tut(xt),t=1T\phi_t(x_t) = x_t + \triangle t \cdot u_t(x_t), \\ \triangle t = \frac{1}{T}

이제

TT \rightarrow \infin

를 생각해보면, 샘플의 변화는 더 이상 단일 함수가 아닌, Ordinary Differentiable Equation (ODE) 으로 기술된다.

dxtdt=ut(xt)\frac{dx_t}{dt} = u_t(x_t)

여기서 ut(xt)u_t(x_t)는 시간 tt 와 위치 xtx_t 가 주어졌을 때, 해당 위치의 확률 밀도 '입자'가 어느 방향과 속도로 움직여야 하는지 정의하는 velocity vector field 이다. 이제 우리의 목표는 변환 함수 ϕ\phi 대신, 이 velocity field utu_t를 신경망으로 학습하는 것이 된다.


2.2. Likelihood of Continuous Normalizing Flow

Continuous Flow 을 도입하면서 얻는 가장 극적인 이점은 likelihood 계산 방식의 근본적인 변화이다. 기존의 이산적 변환에서는 계산이 거의 불가능했던 Jacobian determinant 이, 연속적인 흐름에서는 훨씬 다루기 쉬운 divergence 의 적분 으로 대체된다.

이 원리를 이해하기 위해, 먼저 확률을 하나의 'fluid' 로 바라보는 관점 에서 시작해보자.


2.2.1. Probability Fluid and Flow

시간에 따라 변화하는 어떠한 확률 분포 pt(x)p_t(x) 를 상상해보자.

이는 공간에 퍼져있는 수많은 미세한 'probability particle' 들의 density distribution 로도 생각할 수 있다 (입자들이 움직이면 특정 지역의 밀도는 높아지고 다른 지역은 낮아질 것). 이제 이 관점에서 몇 가지 용어를 먼저 짚고 넘어가겠다.

Probability Current:

  • 이 probability particle 들의 움직임을 나타내는 것이 바로 velocity vector field ut(xt)u_t(x_t) 이다. 하지만 단순히 속도만으로는 '흐름의 양' 을 알 수 없다.

  • 밀도가 0 인 곳에서는 속도가 아무리 빨라도 흐르는 입자가 없으며, 특정 지점 xx 에서의 실제 '확률의 흐름의 양' 은 density pt(x)p_t(x)와 velocity ut(x)u_t(x)의 곱 으로 정의하는 것이 자연스럽다. 이를 Probability Current

    Jt(x)=pt(x)ut(x)J_t(x) = p_t(x) u_t(x)

    라고 부르며, 이 벡터는 해당 지점에서 확률이 어느 방향으로 얼마나 많이 흐르고 있는지를 나타낸다.

Net Outflow:

  • 어떤 infinitesimal volume (미소 부피) 를 생각해보자. 이 부피 안의 probability density 가 시간에 따라 변하는 이유 (pt/t\partial p_t / \partial t) 는 단 하나, 경계를 통해 확률이 들어오거나 나갔기 때문이다.

  • 이 '들어오고 나가는 양의 차이', 즉 net outflow 을 수학적으로 표현하는 도구가 바로 Divergence \nabla 이다.

  • Divergence 은 특정 지점에서 vector field (여기서는 probability flow JtJ_t) 이 얼마나 '뻗어 나오는지 (source)' 혹은 '사라지는지 (sink)'를 측정하는 값이다.

    • Jt>0\nabla \cdot J_t > 0: 해당 지점에서 확률이 밖으로 나오고 있음 ( → density 감소).
    • Jt<0\nabla \cdot J_t < 0: 해당 지점으로 확률이 들어가고 있음 ( → density 증가).
  • 수학적으로는 각 방향의 편미분의 합으로 정의된다.
    ut=u1x1+u2x2++uDxD=i=1Duixi\nabla \cdot u_t = \frac{\partial u_1}{\partial x_1} + \frac{\partial u_2}{\partial x_2} + \dots + \frac{\partial u_D}{\partial x_D} = \sum_{i=1}^D \frac{\partial u_i}{\partial x_i}

따라서, 한 지점에서의 density 변화율은 Net Outflow 에 음수 부호를 붙인 것과 같으며, 이를 수식으로 나타낸 것이 바로 Continuity Equation (연속 방정식) 이다.

pt(x)t=Jt(x)=(pt(x)ut(x))\boxed{\frac{\partial p_t(x)}{\partial t} = - \nabla \cdot J_t(x) = - \nabla \cdot (p_t(x) u_t(x))}

이 방정식은 "확률은 갑자기 생기거나 사라지지 않고, 오직 흐름을 통해서만 그 밀도가 변한다" 는 확률 질량 보존 법칙을 수학적으로 표현한 것이다.

Deterministic drift (ODE) 와 더불어, 브라운 운동 같은 무작위적인 노이즈 (diffusion) 을 고려하는 일반적인 편미분 방정식이 Fokker-Planck Equation 로 알려져 있는데,

pt(x)t=[f(x,t)pt(x)]Drift (flow)+12i,j2xixj[[g(t)g(t)T]ijpt(x)]Diffusion\frac{\partial p_t(x)}{\partial t} = \underbrace{-\nabla \cdot [f(x,t)p_t(x)]}_{\text{Drift (flow)}} + \underbrace{\frac{1}{2}\sum_{i,j} \frac{\partial^2}{\partial x_i \partial x_j} [[g(t)g(t)^T]_{ij} p_t(x)]}_{\text{Diffusion}}

Continuous Normalizing Flow (CNF)는 diffusion 항이 없는 (g=0g=0), 순수한 deterministic 형태이므로, CNF 가 따르는 확률 분포의 dynamics 은 Fokker-Planck Equation 의 특수한 경우인 Continuity Equation 라고도 볼 수 있다.


2.2.2. Derivation of Log-Likelihood

이제 우리는 한 입자를 따라 움직이면서 그 입자가 느끼는 log probability 의 변화율, 즉

ddtlogpt(xt)\frac{d}{dt} \log p_t(x_t)

를 계산해야 한다.

이는 두 가지 효과의 합으로 나타낼 수 있는데,

  1. 시간이 흘러 분포 자체가 변하는 효과 (Eulerian 관점): 내가 가만히 있어도 (xx 고정) 시간이 흐름에 따라 변하는 내 위치의 density
    tlogpt\frac{\partial}{\partial t} \log p_t
  2. 내가 움직여서 다른 밀도의 지역으로 가는 효과 (Lagrangian 관점): 내가 속도 utu_t로 움직이면서 density 가 다른 곳으로 이동하기 때문에 변하는 density
    dxtdtlogpt\frac{dx_t}{dt} \cdot \nabla \log p_t

의 둘을 합친 Total Derivative 로 나타낼 수 있다.

dlogpt(xt)dt=logpt(xt)t+dxtdtxlogpt(xt)\frac{d \log p_t(x_t)}{dt} = \frac{\partial \log p_t(x_t)}{\partial t} + \frac{dx_t}{dt} \cdot \nabla_x \log p_t(x_t)

여기서 dxt/dt=ut(xt){dx_t}/{dt} = u_t(x_t) 이므로, 우리는 다음을 얻는다.

dlogpt(xt)dt=logpt(xt)t+ut(xt)xlogpt(xt)(eqn. 1)\frac{d \log p_t(x_t)}{dt} = \frac{\partial \log p_t(x_t)}{\partial t} + u_t(x_t) \cdot \nabla_x \log p_t(x_t) \quad \cdots \text{(eqn. 1)}

i. First Term

이제 첫 번째 항인 logpt/t\partial \log p_t / \partial t 를 살펴보면

logptt=1ptptt\frac{\partial \log p_t}{\partial t} = \frac{1}{p_t} \frac{\partial p_t}{\partial t}

임을 알 수 있다 (로그 미분),

여기서 pt/t\partial p_t / \partial t 이 위에서 살펴본 continuity equation 인 것이 보이는가? 따라서 continuity equation 을 여기에 대입하면 우리는

logptt=1pt(ptut)\frac{\partial \log p_t}{\partial t} = -\frac{1}{p_t} \nabla \cdot (p_t u_t)

의 관계를 얻는다.

이제 divergence 의 곱셈을 다음과 같이 분배한 후,

(ptut)=(pt)ut+pt(ut)\nabla \cdot (p_t u_t) = (\nabla p_t) \cdot u_t + p_t (\nabla \cdot u_t)

아래와 같은 결과를 얻을 수 있다.

logptt=1pt[(pt)ut+pt(ut)]=(1ptpt)ut(ut)\frac{\partial \log p_t}{\partial t} = -\frac{1}{p_t} [(\nabla p_t) \cdot u_t + p_t (\nabla \cdot u_t)] = - \left( \frac{1}{p_t} \nabla p_t \right) \cdot u_t - (\nabla \cdot u_t)

이제 logpt=pt/pt\nabla \log p_t = \nabla p_t / p_t (score function) 관계를 이용하면,

logptt=(logpt)ut(ut)(eqn. 2)\frac{\partial \log p_t}{\partial t} = -(\nabla \log p_t) \cdot u_t - (\nabla \cdot u_t) \quad \cdots \text{(eqn. 2)}

첫번째 항은 위와 같이 정리된다.


ii. Final Equation

이제 마지막으로, eqn1 에,

dlogpt(xt)dt=logpt(xt)t+ut(xt)xlogpt(xt)(eqn. 1)\frac{d \log p_t(x_t)}{dt} = \frac{\partial \log p_t(x_t)}{\partial t} + u_t(x_t) \cdot \nabla_x \log p_t(x_t) \quad \cdots \text{(eqn. 1)}

방금 얻은 결과인 eqn2 를 대입해보자.

dlogpt(xt)dt=[(logpt)ut(ut)]from (eqn. 2)+utlogpt=ut(xt)\begin{aligned} \frac{d \log p_t(x_t)}{dt} &= \underbrace{\left[-(\nabla \log p_t) \cdot u_t - (\nabla \cdot u_t)\right]}_{\text{from (eqn. 2)}} + u_t \cdot \nabla \log p_t \\ &= - \nabla \cdot u_t(x_t) \end{aligned}

결과적으로 우리는 다음과 같은 매우 간결한 최종 eqn 을 얻게 되었다.

dlogpt(xt)dt=ut(xt)\frac{d \log p_t(x_t)}{dt} = - \nabla \cdot u_t(x_t)

2.2.3. Divergence & Trace

이 결과가 혁신적인 이유는 divergence ut\nabla \cdot u_t 가 velocity field utu_tJacobian Matrix 의 대각합 (trace)과 같기 때문이다.

ut=i=1Dut,ixi=Tr(utxt)\nabla \cdot u_t = \sum_{i=1}^D \frac{\partial u_{t,i}}{\partial x_i} = \text{Tr}\left( \frac{\partial u_t}{\partial x_t} \right)

Divergence 와 Jacobian 의 정의를 다시 살펴보면 이것이 자명하다.

  • Divergence:

    ut=u1x1+u2x2++uDxD=i=1Duixi\nabla \cdot u_t = \frac{\partial u_1}{\partial x_1} + \frac{\partial u_2}{\partial x_2} + \dots + \frac{\partial u_D}{\partial x_D} = \sum_{i=1}^D \frac{\partial u_i}{\partial x_i}
  • Jacobian:

    Ju=ux=(u1x1u1x2u1xDu2x1u2x2u2xDuDx1uDx2uDxD)J_u = \frac{\partial u}{\partial x} = \begin{pmatrix} \frac{\partial u_1}{\partial x_1} & \frac{\partial u_1}{\partial x_2} & \cdots & \frac{\partial u_1}{\partial x_D} \\ \frac{\partial u_2}{\partial x_1} & \frac{\partial u_2}{\partial x_2} & \cdots & \frac{\partial u_2}{\partial x_D} \\ \vdots & \vdots & \ddots & \vdots \\ \frac{\partial u_D}{\partial x_1} & \frac{\partial u_D}{\partial x_2} & \cdots & \frac{\partial u_D}{\partial x_D} \end{pmatrix}

즉 여기서 Jacobian matrix 의 trace 를 구해보면,

Tr(Ju)=u1x1+u2x2++uDxD\text{Tr}(J_u) = \frac{\partial u_1}{\partial x_1} + \frac{\partial u_2}{\partial x_2} + \dots + \frac{\partial u_D}{\partial x_D}

divergence 의 정의와 동일한 것을 알 수 있다.

직관적으로도 당연한 결과인데, off-diagonal 원소 (ui/xj({\partial u_i}/{\partial x_j} , ij)i \neq j) 의 의미를 살펴보면, 이는 jj 축 방향으로 움직일 때 ii 방향 속도가 어떻게 변하는지를 나타내는 값이다.

즉, 유체의 부피 변화 (팽창/수축) 가 아닌, rotation 이나 shear 을 유발한다. 애초에 우리가 determinant 로 구하고 싶은 것은 선형 변환 간의 부피의 변화율 이었으므로 이러한 off-diagonal 들은 여기에 기여하지 않는다는 것을 알 수 있다.

결론적으로, 이산적 변환에서 필요했던 Jacobian determinant 계산이 continuous flow 에서는 Jacobian trace 계산으로 바뀌었다. Determinant 은 모든 원소에 의존하며 계산량이 O(D3)\mathcal{O}(D^3)이지만, trace 은 대각 원소들만의 합이므로 계산이 훨씬 간단하다.


3. Flow Matching

Continuous Normalizing Flow (CNF)는 Jacobian determinant 을 trace 의 적분으로 대체하여 계산을 용이하게 만들었지만, 여전히 심각한 문제, 바로 느린 학습 속도 가 남아있다. Log-likelihood 를 계산하기 위해서는 매 학습 단계마다 ODE 를 풀어야 했고, 이는 본질적으로 diffusion model 의 느린 샘플링 과정을 매번 반복하는 것과 유사했다.

Neural Network 로 velocity field vθ(xt,t)v_\theta(x_t, t) 를 학습시키기 위한 이상적인 loss function 은 다음과 같다.

L=EtU,xtpt[vθ(xt,t)ut(xt)2]\mathcal{L} = \mathbb{E}_{t \sim \mathcal{U}, x_t \sim p_t} [\|v_\theta(x_t, t) - u_t(x_t)\|^2]

이 loss function 은 시간 tt의 샘플 xtx_t 를 보고, 그 지점에서의 실제 속도장 ut(xt)u_t(x_t) 를 예측하는 regression 문제이다.

  • vθ(xt,t)v_\theta(x_t, t): 우리가 학습시킬 Network 가 예측한 velocity vector field
  • ut(xt)u_t(x_t): 우리가 알아야 할 'GT' marginal vector field

이는 구조적으로 DDPM 의 GT noise regression 문제와도 유사 하지만, 결정적인 차이가 존재한다. "전체 데이터 분포 p0p_0p1p_1 으로 변환하는 'GT' marginal velocity field ut(xt)u_t(x_t) 를 애초에 알지 못한다"는 사실이다.

이 velocity field 는 시간 tt 에서의 전체 데이터 분포 pt(x)p_t(x) 의 변화를 나타내는 복잡한 양으로, 직접 계산하는 것이 불가능하다.


3.1. Conditional Flow Matching

이 막다른 길을 뚫기 위해 Conditional Flow Matching, (CFM) 라는 트릭이 등장한다. CFM 의 핵심 아이디어는 다음과 같이 요약된다.

"분포 전체의 경로 pt(x)p_t(x)를 따라가는 복잡한 'marginal vector field' ut(x)u_t(x)를 직접 학습하는 대신, 개별 샘플 쌍 (z0z_0, z1z_1)을 잇는 간단한 'conditional vector field' ut(xz1)u_t(x|z_1)을 모방하도록 Network 를 학습시키자."

이것이 어떻게 가능한지 단계별로 살펴보자.

3.1.1. Conditional Probability Path

전체 분포의 알 수 없는 경로 pt(x)p_t(x) 대신, 다루기 쉬운 Conditional Probability Path pt(xx1)p_t(x|x_1) 를 직접 정의할 수 있을 것이다. 이는 데이터 샘플 x1x_1이 주어졌을 때, 시간 tt 에 입자가 어디에 있을지에 대한 확률 분포이다.

가장 대표적인 방법은 경로를 gaussian distribution 의 trajectory 으로 정의하는 것인데,

pt(xx1)=N(xμt(x1),σt2(x1)I)p_t(x|x_1) = \mathcal{N}(x | \mu_t(x_1), \sigma_t^2(x_1)I)

시간 t=0t=0 에서는 standard normal distribution (완전한 노이즈) 가 되고, 시간 t=1t=1 에서는 target data x1x_1 에 도달하도록 경계 조건을 설정하면된다.

  • t=0t=0: μ0(x1)=0,σ0(x1)=1\mu_0(x_1) = 0, \sigma_0(x_1) = 1
  • t=1t=1: μ1(x1)=x1,σ1(x1)=σmin0\mu_1(x_1) = x_1, \sigma_1(x_1) = \sigma_{min} \approx 0

이는 diffusion 의 Forward Process q(xtx0)q(x_t|x_0) 와 정확히 같은 개념적 구조를 가진다. DDPM 이 x0x_0 에서 noise 로 가는 경로를 정의했다면, CFM 은 반대로 target x1x_1을 고정하고 그곳으로 향하는 가상의 경로를 정의하는 것이라고 해석할 수 있다.

이제 우리는 직접 설계한 path pt(xx1)p_t(x|x_1) 으로부터, 이 path 를 생성하는 conditional vector fields ut(xx1)u_t(x|x_1) 역시 continuity equation 을 통해 유도할 수 있다.

pt(xt)t=(pt(xt)ut(xt)){\frac{\partial p_t(x_t)}{\partial t} = - \nabla \cdot (p_t(x_t) u_t(x_t))}

위의 continuity equation 으로부터 아래와 같은 conditional vector fields 가 유도된다.

pt(xtx1)t=(pt1(xtx1)ut(xtx1)){\frac{\partial p_t(x_t | x_1)}{\partial t} = - \nabla \cdot (p_{t|1}(x_t | x_1) u_t(x_t | x_1))}

이 때, 주어진 conditional probability 를 만족하는 무수히 많은 conditional vector fields 가 존재함을 기억하자.

path 1path 2

3.1.2. Conditional Flow Matching

이제 CFM의 핵심인 두 가지 중요한 수학적 사실을 살펴보자.

[1] Marginal vector fields 은 conditional vector fields 의 기댓값과 같다.

marginal vector fields ut(xt)u_t(x_t) 는 conditional vector fields ut(xtx1)u_t(x_t|x_1) 을 통해 다음과 같이 나타낼 수 있고 (bayes rule),

ut(xt)=ut(xtx1)pt1(xtx1)p1(x1)pt(xt)dx1u_t(x_t) = \int u_t(x_t|x_1) \frac{p_{t|1}(x_t | x_1) p_1 (x_1)}{p_t(x_t)} dx_1

앞서 살펴본 continuity equation

pt(xtx1)t=(pt1(xtx1)ut(xtx1)){\frac{\partial p_t(x_t | x_1)}{\partial t} = - \nabla \cdot (p_{t|1}(x_t | x_1) u_t(x_t | x_1))}

또한 bayes rule 을 이용해 정리해보면,

pt(xt)t=tpt1(xtx1) p1(x1)dx1=(pt1(xtx1)ut(xtx1)) p1(x1)dx1=(ut(xtx1))pt1(xtx1)p1(x1)pt(xt)dx1 pt(xt))=[ut(xt) pt(xt)]\begin{aligned} \frac{\partial p_t(x_t)}{\partial t} &= \boxed{\frac{\partial}{\partial t} \int p_{t|1}(x_t | x_1)} \ p_1 (x_1) dx_1 \\ &= \boxed{ - \int \nabla \cdot (p_{t|1}(x_t | x_1) u_t(x_t | x_1)) }\ p_1 (x_1) dx_1 \\ &= - \nabla \cdot \left ( \boxed{ \int u_t(x_t | x_1)) \frac{p_{t|1}(x_t | x_1) p_1 (x_1)}{p_t(x_t)} dx_1 } \ p_t(x_t) \right ) \\ &= - \nabla \cdot \bigg [ \boxed{u_t(x_t)} \ p_{t}(x_t ) \bigg ] \end{aligned}

가 성립하는 것을 알 수 있다.

즉, 어떤 지점 xtx_t에서의 marginal vector fields ut(xt)u_t(x_t) 는, 그 지점을 통과하는 모든 가능한 conditional vector fields ut(xtx1)u_t(x_t|x_1) 을 평균낸 것과 같다.

ut(xt)=Ep(x1xt)[ut(xtx1)]=ut(xtx1)p(x1xt)dx1u_t(x_t) = \mathbb{E}_{p(x_1|x_t)}[u_t(x_t|x_1)] = \int u_t(x_t|x_1) p(x_1|x_t) dx_1

[2] 따라서, 전체 loss function 은 conditional loss function 과 동등하다.

첫 번째 사실로부터, 우리가 풀고 싶었던 이상적이지만 계산 불가능한 loss function 이, 계산 가능한 conditional loss function 과 동등하다는 것이 증명된다.

Lmarginal=Et,xt[vθut2]    LCFM=Et,x1,xtx1[vθut(x1)2]+C\mathcal{L}_{\text{marginal}} = \mathbb{E}_{t, x_t} [\|v_\theta - u_t\|^2] \\ \iff \\ \mathcal{L}_{\text{CFM}} = \mathbb{E}_{t, x_1, x_t|x_1} [\|v_\theta - u_t(\cdot|x_1)\|^2] + C

두 loss 가 동등한 이유는 loss function 을 전개했을 때 나타나는 inner product 항을 살펴보면 명확해진다.

Extpt[vθ(xt),ut(xt)]=[vθ(xt)ut(xt) ]pt(xt)dxt=[vθ(xt)ut(xtx1)pt(x1xt)dx1]pt(xt)dxt=vθ(xt)ut(xtx1)pt(x1xt)pt(xt)dx1dxt=vθ(xt)ut(xtx1)p(x1,xt)dx1dxt=E(x1,xt)p(x1,xt)[vθ(xt),ut(xtx1)]\begin{aligned} \mathbb{E}_{x_t \sim p_t} [\langle v_\theta(x_t), u_t(x_t) \rangle] &= \int \left[ v_\theta(x_t) \cdot \boxed{u_t(x_t)} \ \right ]p_t(x_t) dx_t \\ &= \int \left[ v_\theta(x_t) \cdot \boxed{\int u_t(x_t|x_1) p_t(x_1|x_t) dx_1 } \right ]p_t(x_t) dx_t \\ &= \iint v_\theta(x_t) \cdot u_t(x_t|x_1) \boxed{ p_t(x_1|x_t) p_t(x_t)} dx_1 dx_t \\ &= \iint v_\theta(x_t) \cdot u_t(x_t|x_1) \boxed{p(x_1, x_t)} dx_1 dx_t \\ &= \mathbb{E}_{(x_1, x_t) \sim p(x_1, x_t)} [\langle v_\theta(x_t), u_t(x_t|x_1) \rangle] \end{aligned}

이는 알 수 없는 utu_t에 대한 regression problem 을, 우리가 직접 설계하여 정답을 알고 있는 ut(x1)u_t(\cdot|x_1)에 대한 regression problem 로 완벽하게 대체할 수 있음을 의미한다.


3.1.3. Sampling w/ Conditional Flow

이제 마지막으로, 학습 데이터 xtpt(xx1)x_t \sim p_t(x|x_1)는 어떻게 샘플링할까 고민해보자.

Gaussian Distribution 등에서 직접 샘플링하는 대신, CFM 은 더 효율적인 Conditional Flow Map ϕt\phi_t 을 사용한다.

xt=ϕt(x0x1)=σt(x1)x0+μt(x1)x_t = \phi_t(x_0|x_1) = \sigma_t(x_1)x_0 + \mu_t(x_1)

이는 standard normal distribution (base distribution, p0p_0)에서 샘플링한 노이즈 x0x_0 를 시간 tt의 target distribution pt(xx1)p_t(x|x_1) 위의 한 점 xtx_t 로 deterministic 하게 매핑하는 함수이다.

이 덕분에 우리는 복잡한 분포에서 샘플링할 필요 없이, 간단한 덧셈과 곱셈만으로 학습에 필요한 데이터 xtx_t를 생성할 수 있다.


3.2. Final Training Algorithm

이 모든 과정을 종합하면, 최종 학습 알고리즘은 다음과 같이 단순해진다.

  1. Sampling: base distribution 에서 노이즈 x0p0x_0 \sim p_0 를, 실제 data distribution 에서 x1p1x_1 \sim p_1를 각각 샘플링한다.
  2. Time Sampling: 시간 tUt \sim \mathcal{U} 를 균등하게 샘플링한다.
  3. Trajectory Path: Conditional Flow Map 을 이용해 경로 위의 점 xt=ϕt(x0x1)x_t = \phi_t(x_0|x_1) 를 계산한다.
  4. Target Velocity: 해당 경로의 '정답' velocity ut(xtx0,x1)u_t(x_t|x_0, x_1) 를 계산한다.
  5. Training: Neural Network vθ(xt,t)v_\theta(x_t, t) 가 target velocity 를 예측하도록 다음 loss 를 최소화한다:
    L=vθ(xt,t)ut(xtx0,x1)2\mathcal{L} = \|v_\theta(x_t, t) - u_t(x_t|x_0, x_1)\|^2

결론적으로 Flow Matching 은 학습 과정에서 ODE 를 풀 필요가 전혀 없으며, 두 점을 잇는 path 의 기울기를 예측하는 간단한 regression 문제로 변환된다.

'두 점을 잇는 deterministic path' 라는 아이디어는 DDIM 의 핵심 통찰과도 맞닿아 있으며, Flow Matching 은 이를 학습 패러다임 자체에 녹여내어 생성 모델의 학습 효율성을 극적으로 끌어올렸다.


3.3. Rectified Flow

CFM 의 틀 안에서, 우리는 어떤 'Conditional Path' 든 자유롭게 설계할 수 있다. 여기서 Rectified Flow는 가장 단순하고 직관적인 선택을 한다.

"노이즈 x0x_0 와 데이터 x1x_1 사이를 잇는 가장 간단한 경로는 직선이다."

이 직선 경로는 확률적으로

pt(xtx0,x1)=δ(xt((1t)x0+tx1))p_t(x_t|x_0, x_1) = \delta(x_t - ((1-t)x_0 + tx_1))

로 표현되며, 샘플링 관점에서는 다음과 같다.

  • 경로: xt=(1t)x0+tx1x_t = (1-t)x_0 + t x_1
  • 속도: 이 경로를 만들기 위한 속도는 경로의 시작과 끝에만 의존하는 상수 벡터, 즉
    ut(xtx0,x1)=dxtdt=x1x0u_t(x_t|x_0, x_1) = \frac{dx_t}{dt} = x_1 - x_0
    이다.

이제 CFM 의 loss function 에 이를 대입하면 Rectified Flow 의 최종 목표 함수가 완성된다.

L=EtU,x0p0,x1p1[vθ((1t)x0+tx1,t)(x1x0)2]\mathcal{L} = \mathbb{E}_{t \sim U, x_0 \sim p_0, x_1 \sim p_1} \left[ \left\| v_\theta((1-t)x_0 + tx_1, t) - (x_1 - x_0) \right\|^2 \right]

신경망은 중간 지점 xtx_t를 보고, 최종 목적지에서 출발지를 뺀 방향 벡터 (x1x0x_1 - x_0) 를 맞추는 아주 간단한 regression 문제 를 풀게 된다. 이 단순함이 비약적인 학습 속도와 안정성을 가져온다.

마지막으로 Rectified Flow 은 또다른 장점이 존재한다.

한 번의 학습으로 얻어진 coupling (Z0,Z1)(Z_0, Z_1) 은 초기 노이즈와 데이터의 독립적인 coupling (X0,X1)(X_0, X_1) 보다 이미 훨씬 정돈된 (less entangled) 상태이다. 만약 우리가 이 (Z0,Z1)(Z_0, Z_1) 을 새로운 '데이터'로 삼아 다시 한번 Rectified Flow 를 학습시키면 어떻게 될까? Rectified Flow 연구진은 이 'Reflow' 과정을 반복할수록 샘플들의 이동 경로가 기하급수적으로 직선에 가까워짐을 보였다.

이는 단 몇 번의 반복만으로도 거의 완벽한 직선 경로를 학습할 수 있다는 의미이며, 생성 과정에서 필요한 ODE 스텝 수를 획기적으로 줄여 (심지어 단 한 스텝으로도) 매우 빠른 샘플링을 가능하게 한다.

이제 우리는 어떻게 Flow Matching을 통해 직선 경로를 학습할 수 있는지 알게 되었다. 하지만 근본적인 질문이 남는다.

"수많은 가능한 경로 중에, 왜 하필 이 '직선' 경로가 이토록 효과적인가?"

단순히 계산이 편하다는 공학적 타협일까? 아니면 그 이면에 더 깊은 수학적, 물리적 원리가 숨어있는 것일까?

Part B. Rectified Flow and Optimal Transport

앞선 질문에 답하기 위해, 우리는 하나의 가정을 설정하고, 그 가정이 어떤 결론으로 이어지는지 고찰해 보고자 한다. 바로 생성 모델이 노이즈에서 데이터로 가는 경로를 찾는 '비용 (Cost)' 을, 물리학에서 입자들이 움직이는 '총 운동 에너지 (Total Kinetic Energy)' 와 동일하다고 가정 하는 것이다.

이는 수학적으로 두 분포 사이의 L2 Wasserstein 거리를 최소화하는 Dynamic Optimal Transport 문제와 깊은 관련이 있는데, 이를 이해하기 위해 물리학의 Least Action Principle (최소 작용의 원리) 와 수학의 Optimal Transport 을 통해 그 근본 원리를 파헤쳐보자.


4. Optimal Transport

4.1. What is Optimal Transport

먼저 optimal transport 를 직관적으로 이해해보자. 여기 흙이 쌓여있는 언덕 (XX, 소스 분포)이 있고, 이 흙을 모두 파내서 어딘가에 있는 구덩이 (YY, 타겟 분포)를 메워야 한다고 상상해보자.

이때 "가장 효율적으로" 흙을 옮기려면 어떻게 해야 할까?

'효율적'이라는 것은 총 이동 비용(Cost)을 최소화하는 것을 의미한다. 최적 수송은 바로 이 문제, 즉 한 확률 분포 (흙더미)를 다른 확률 분포 (구덩이)로 변환하는 데 필요한 최소 비용의 Transport Plan을 찾는 수학 이론이다.

4.2. Dynamic Optimal Port & Principle of Least Action

전통적인 OT는 '어디의 흙을 어디로 보낼 것인가'라는 최종 매핑 (mapping)에 집중한다. 하지만 생성 모델에서는 샘플이 점차 변해가는 경로 (path) 자체가 더 중요하며, 이는 동적 최적 수송 (Dynamic Optimal Transport) 의 영역이다. Dynamic OT에서는 시간 tt 에 따른 분포의 변화 ptp_t 와 그 경로를 만드는 velocity field vtv_t 를 고려한다.

Rectified Flow의 '직선 경로'를 단순한 수학적 트릭이 아닌, 근본적인 원리로 이해하기 위해서는 물리학의 Principle of Least Action 를 이해할 필요가 있다.


5. Classical Mechnics & Lagrangian Dynamics

5.1. What is Action

고전 역학에서, 어떤 물체가 한 지점에서 다른 지점으로 이동할 때, 가능한 수많은 경로 중에서 Action, SS 이라는 물리량을 최소화하는 경로를 선택하며, 이것이 바로 Principle of Least Action 이다.

이 액션은 Lagrangian, LL 이라는 양을 시간에 따라 적분한 값으로 정의되며,

S=t0t1L(x,x˙,t)dtS = \int_{t_0}^{t_1} L(x, \dot{x}, t) \, dt

라그랑지언 LL은 그 시스템의 운동 에너지(TT)에서 위치 에너지(VV)를 뺀 값이다.

L=TVL = T - V

유튜브 veritasium 에 action 에 대해서 정말 훌륭하게 설명하는 영상이 있으므로 참고하길 바란다: 영상 link

  • Fermat's Principle 로 알려진 빛의 굴절 현상도 (거리가 아닌 시간 최소화) Least Action Principle 의 한 예이다.

이제 가장 간단한 경우, 즉 아무런 힘 (중력, 마찰 등)이 작용하지 않는 자유 공간을 움직이는 단일 입자를 생각해보자. 이 경우 위치 에너지는 V=0V=0 이며, 라그랑지언은 순수하게 운동 에너지만 남게된다.

L=T=12mv2=12mx˙2L = T = \frac{1}{2}m v^2 = \frac{1}{2}m ||\dot{x}||^2

우리는 한 분포에서 다른 분포로의 매핑 자체에만 관심이 있고, 그 과정에서 특정 경로에 패널티를 주는 외부 요인 (Potential) 을 가정하지 않으므로, 시스템의 비용은 순수하게 입자들의 이동 거리와 속도 (Kinetic Energy) 에만 의존한다고 볼 수 있다.

이 때, 이 입자의 액션은 다음과 같으며

S=t0t112mx˙(t)2dtS = \int_{t_0}^{t_1} \frac{1}{2}m ||\dot{x}(t)||^2 \, dt

최소 작용 원리에 따르면, 입자는 이 action SS를 최소화하는 경로를 따라 움직인다. 적분 안의 값이 항상 양수이므로, 이 값을 최소화하는 가장 직관적인 방법은 속도(x˙||\dot{x}||)를 일정하게 유지하며 최단 거리, 즉 직선으로 움직이는 것일 것이다.

실제로 위 액션을 최소화하는 경로를 찾기 위한 미분방정식을 풀어보면, 가속도(x¨\ddot{x})가 0이라는 결과가 나온다 (mx¨=0m\ddot{x} = 0). 이는 곧 등속 직선 운동을 의미한다.


5.2. Euler-Lagrange Equation

Euler-Lagrange Equation 이란, action 을 최소화 (혹은 stationary 하게, 즉, 극값이나 saddle point!) 하게 만드는 조건에 대한 미분방정식이다.

이는 변분법으로 유도할 수 있으며, 아래와 같이 정의된다:

Lxddt(Lx˙)=0\frac{\partial L}{\partial x} - \frac{d}{dt}\left(\frac{\partial L}{\partial \dot{x}}\right) = 0

이 방정식의 각 항이 의미하는 바는 다음과 같은데,

  • L/x{\partial L}/{\partial x} : 경로의 위치(xx)가 변할 때, 라그랑지언 (cost)이 얼마나 변하는가?
  • L/x˙{\partial L}/{\partial \dot{x}} : 경로의 속도(x˙\dot{x})가 변할 때, 라그랑지언 (cost)이 얼마나 변하는가? (참고: 이 값은 운동량 p=mvp=mv와 같다)
  • d/dt(){d}/{dt}(\dots) : 변화량이 시간에 따라 어떻게 변하는가?

즉, 위 방정식을 품으로써 우리는 action 을 최소화하는 조건의 해가 무엇인지를 구할 수 있다. 이제 우리가 1단계에서 구한 라그랑지언 L=12mx˙2L = \frac{1}{2}m\dot{x}^2 을 이 방정식에 넣고 계산해 보자.

첫 번째 항: L/x{\partial L}/{\partial x} 계산

라그랑지언 L=12mx˙2L = \frac{1}{2}m\dot{x}^2 을 보자. 이 식에는 위치를 나타내는 변수 xx가 없으며, 오직 속도 x˙\dot{x} 만 들어있다. 따라서, 위치 xx에 대해 미분하면 그 값은 0 이 된다.

Lx=x(12mx˙2)=0\frac{\partial L}{\partial x} = \frac{\partial}{\partial x}\left(\frac{1}{2}m\dot{x}^2\right) = 0

직관적으로는 자유 공간에서는 입자가 어디에 있든 비용(라그랑지언)은 동일하다고 볼 수 있다. 비용은 오직 얼마나 빠른지 에만 의존하기 때문.

두 번째 항: ddt(Lx˙)\frac{d}{dt}\left(\frac{\partial L}{\partial \dot{x}}\right) 계산
먼저 L/x˙{\partial L}/{\partial \dot{x}} 부터 살펴보자. 라그랑지언 L=12mx˙2L = \frac{1}{2}m\dot{x}^2 을 속도 x˙\dot{x}에 대해 미분하면,

Lx˙=x˙(12mx˙2)=12m(2x˙)=mx˙\frac{\partial L}{\partial \dot{x}} = \frac{\partial}{\partial \dot{x}}\left(\frac{1}{2}m\dot{x}^2\right) = \frac{1}{2}m(2\dot{x}) = m\dot{x}

이 되고, 이 결과를 시간 tt에 대해 다시 미분하면,

ddt(mx˙)\frac{d}{dt}\left(m\dot{x}\right)

질량 mm은 상수이므로 밖으로 나오고, 속도 x˙\dot{x}를 시간으로 미분하면 가속도 x¨\ddot{x}가 된다.

ddt(mx˙)=mx¨\frac{d}{dt}\left(m\dot{x}\right) = m \ddot{x}

이제 두 결과를 대입해보면,

Lxddt(Lx˙)=0\frac{\partial L}{\partial x} - \frac{d}{dt}\left(\frac{\partial L}{\partial \dot{x}}\right) = 0
0mx¨=00 - m\ddot{x} = 0

질량 mm 은 0이 아니므로, 이 식이 성립하려면 반드시 다음 조건이 만족되어야 하는 것을 알 수 있다.

x¨=0\ddot{x} = 0

즉 이는 가속도가 0인, 등속 직선 운동(uniform linear motion) 을 의미하게 된다.


5.3. Action in Optimal Transport

5.3.1. Action in Probability Distribution

이제 이 개념을 단일 입자에서 무수히 많은 입자들의 집합, 즉 확률 분포로 확장해 보자.

  • 단일 입자의 운동 에너지: T=12mv2T = \frac{1}{2}m ||v||^2
  • 확률 분포의 총 운동 에너지: 분포를 구성하는 모든 입자들의 운동 에너지의 기댓값

시간 tt에서의 분포를 pt(x)p_t(x)로, 각 지점 xx에서의 입자들의 속도를 속도장 vt(x)v_t(x)로 표현한다. 특정 시간 tt에서의 총 운동 에너지는 각 지점의 운동 에너지(12vt(x)2\frac{1}{2}||v_t(x)||^2)를 그 지점의 밀도(pt(x)p_t(x))로 가중 평균하여 계산할 수 있다.

Total Kinetic Energy at t=Expt[12vt(x)2]=12vt(x)2pt(x)dx\text{Total Kinetic Energy at } t = \mathbb{E}_{x \sim p_t}\left[\frac{1}{2}||v_t(x)||^2\right] = \int \frac{1}{2}||v_t(x)||^2 p_t(x) \,dx

자유 입자와 마찬가지로, 외부 힘이 없는 분포의 변환 (위치 에너지가 없는 경우) 을 가정하면 이 시스템의 라그랑지언은 총 운동 에너지가 된다. 따라서 전체 변환 과정의 액션은 이 총 운동 에너지를 시간에 대해 적분한 값이 된다.

A(pt,vt)=01(12vt(x)2pt(x)dx)dt\mathcal{A}(p_t, v_t) = \int_0^1 \left( \int \frac{1}{2}||v_t(x)||^2 p_t(x) \,dx \right) dt

최소 작용의 원리와 같이, 우리는 이 액션을 최소화하는 경로 (pt,vtp_t, v_t) 를 찾고자 한다.

이렇게 확률 분포의 변환을 총 운동 에너지 (Action) 를 최소화하는 경로를 찾는 문제로 재정의하는 것은 optimal transport 에서 'Benamou-Brenier formula' 로 알려진 dynamic OT 공식의 핵심 아이디어와 정확히 일치한다. 우리는 지금 생성 모델의 문제를 유체 역학의 언어로 풀어내고 있는 셈이다.


5.3.2. Constrained Euler-Lagrange Equation

여기서도 한 가지 제약 조건이 있다. Flow 절에서 살펴봤던 continuity equation 이 기억나는가? 이는 입자 (probability mass)는 갑자기 생기거나 사라지지 않고 보존되어야 하는 확률 질량 보존을 나타내는 식이었다.

우리는 probability distribution 에 대한 문제를 풀고 있으므로, 여기서는 continuity equation 이 우리가 지켜야하는 constraint 가 된다.

ptt=(ptvt)\frac{\partial p_t}{\partial t} =- \nabla \cdot (p_t v_t)

이러한 제약이 있는 최적화 문제를 풀기 위해, Lagrange Multiplier 함수 ϕ(x,t)\phi(x, t) 를 도입한다. 이 함수는 액션에 제약 조건을 결합하여 새로운 functional L\mathcal{L}을 만든다.

L[p,v,ϕ]=A[p,v]01Rdϕ(x,t)(ptt+(ptvt))dxdt\mathcal{L}[p, v, \phi] = \mathcal{A}[p, v] - \int_0^1 \int_{\mathbb{R}^d} \phi(x, t) \left( \frac{\partial p_t}{\partial t} + \nabla \cdot (p_t v_t) \right) \,dx \,dt

이제 제약 조건 없는 L\mathcal{L} 을 최적화하는 문제로 바꿀 수 있다. 최적의 경로에서 L\mathcal{L}pp, vv, ϕ\phi 의 미소 변화 (variation) 에 대해 안정적이어야 한다. 즉, 각 변수에 대한 functional derivative 이 0이 되어야 한다.

계산을 용이하게 하기 위해, 제약 조건 항에 부분적분 (integration by parts)을 적용하여 미분을 ppvv 에서 ϕ\phi 로 옮겨보자.

L=01Rd(12v2p+pϕt+pvϕ)dxdtRd[ϕp]01dx\mathcal{L} = \int_0^1 \int_{\mathbb{R}^d} \left( \frac{1}{2} \|v\|^2 p + p \frac{\partial \phi}{\partial t} + p v \cdot \nabla \phi \right) dx dt - \int_{\mathbb{R}^d} [\phi p]_0^1 dx

이제 이 식을 이용하여 각 변수에 대한 variation 을 계산한다.

Velocity vv 에 대한 functional derivative
L\mathcal{L} 을 속도 vv 에 대해 functional derivative (δvL\delta_v \mathcal{L})하고 그 결과를 0 으로 둔다.

피적분 함수를 vv 에 대해 미분하면 pv+pϕ=p(v+ϕ)p v + p \nabla \phi = p(v + \nabla \phi) 이고, 이 값이 0 이 되려면 다음이 성립해야 한다.

v+ϕ=0v + \nabla \phi = 0

따라서 첫 번째 핵심 결과, 즉 최적 vector field (vtv_t) 은 어떤 scalar potential function (ϕ\phi) 의 gradient 여야 한다는 것을 얻는다.

vt(x)=ϕ(x,t){v_t(x) = -\nabla \phi(x, t)}

  • Vector field (right) and corresponding scalar potential (left).

Density pp 에 대한 functional derivative

이제 L\mathcal{L} 을 density pp에 대해 functional derivative (δpL\delta_p \mathcal{L}) 하고 그 결과를 0으로 둔다.피적분 함수는 pp 에 대해 선형이므로, 이 변분이 0이 되려면 괄호 안의 항이 0이 되어야 한다.

12v2+ϕt+vϕ=0\frac{1}{2} \|v\|^2 + \frac{\partial \phi}{\partial t} + v \cdot \nabla \phi = 0

Hamilton-Jacobi Equation
이제 첫 번째 결과 (v=ϕv = -\nabla \phi) 를 두 번째 결과에 대입해보자.

12ϕ2+ϕt+(ϕ)(ϕ)=0\frac{1}{2} \|-\nabla \phi\|^2 + \frac{\partial \phi}{\partial t} + (-\nabla \phi) \cdot (\nabla \phi) = 0
12ϕ2+ϕtϕ2=0\frac{1}{2} \|\nabla \phi\|^2 + \frac{\partial \phi}{\partial t} - \|\nabla \phi\|^2 = 0

이를 정리하면, 포텐셜 함수 ϕ\phi가 만족해야 하는 편미분방정식, 즉 Hamilton-Jacobi Equation 을 얻는다.

ϕt12ϕ(x,t)2=0{\frac{\partial \phi}{\partial t} - \frac{1}{2} \|\nabla \phi(x, t)\|^2 = 0}

이 방정식은 자유 입자의 운동을 설명하는 고전 역학의 근본적인 방정식이다.

가속도가 0 임을 증명하기
마지막으로 이 최적의 vector field 를 따라 움직이는 입자의 경로 x(t)x(t)를 생각해 보자. 입자의 속도는 x˙(t)=vt(x(t))\dot{x}(t) = v_t(x(t)) 이다. 우리는 이 입자의 가속도 x¨(t)\ddot{x}(t)가 0임을 보이면 된다.

가속도 x¨(t)\ddot{x}(t)는 속도 vt(x(t))v_t(x(t))의 시간에 대한 total derivative 이다.

x¨=ddtvt(x(t))=vt+(v)v\ddot{x} = \frac{d}{dt} v_t(x(t)) = \frac{\partial v}{\partial t} + (v \cdot \nabla)v

각 항을 포텐셜 ϕ\phi로 표현해 보자.

  1. 첫 번째 항: 위에서 유도한 Hamilton-Jacobi Equation (ϕt=12v2\frac{\partial \phi}{\partial t} = \frac{1}{2}\|v\|^2)과 v=ϕv = -\nabla\phi를 이용하면,
    vt=t(ϕ)=(ϕt)=12(v2)\frac{\partial v}{\partial t} = \frac{\partial}{\partial t}(-\nabla \phi) = -\nabla\left(\frac{\partial \phi}{\partial t}\right) = -\frac{1}{2}\nabla(\|v\|^2)
  2. 두 번째 항: vector field vv는 scalar potential ϕ\phi 의 gradient (v=ϕv = -\nabla\phi) 이므로, 회전 (curl)이 0이다 (×v=×(ϕ)=0\nabla \times v = \nabla \times (-\nabla\phi) = 0). 따라서,
    (v)v=12(v2)(v \cdot \nabla)v = \frac{1}{2}\nabla(\|v\|^2)
    이제 두 항을 조합하여 가속도를 계산하면,
    x¨=vt+(v)v=12(v2)+12(v2)=0\ddot{x} = \frac{\partial v}{\partial t} + (v \cdot \nabla)v = -\frac{1}{2}\nabla(\|v\|^2) + \frac{1}{2}\nabla(\|v\|^2) = 0

cf:

12(AA)=(A)A+A×(×A)\frac{1}{2}\nabla(A \cdot A) = (A \cdot \nabla)A + A \times (\nabla \times A)

5.4. Conclusion: Why Recitifed-Flow is Linear?

x¨(t)=0\ddot{x}(t) = 0

Action 을 최소화하고 continuity equation 을 만족하는 최적의 흐름을 따르는 모든 입자의 가속도는 0 이다. 이는 각 입자가 시작점 z0z_0에서 도착점 z1z_1까지 등속 직선 운동을 한다는 것을 수학적으로 증명한다.

이것이 바로 Rectified Flow가 두 분포의 샘플 z0,z1z_0, z_1을 뽑아 그 사이를 잇는 직선 경로를 학습의 목표로 삼는 강력한 이론적 배경이다. 모델이 학습하는 속도 벡터 v=z1z0v = z_1 - z_0는 이 최적 경로의 속도장에 해당한다.

따라서 Rectified Flow가 채택한 '직선 경로' zt=(1t)z0+tz1z_t = (1-t)z_0 + t z_1 는 임의의 선택이 아니라, 물리 법칙과 수학적 최적화에 근거한 가장 자연스럽고 효율적인 경로라고 해석할 수 있다.


마치며

여기까지 긴 글을 읽어주셔서 감사합니다.

사실 이 글은 요새 집필 중인 시리즈 글 : Building Large 3D Generative Model from Scratch 에서 Model Structure 를 설명하는 부분에서 시작하였습니다.

Flow 는 최근 이미지, 비디오 뿐만 아니라 3D 생성에서도 Network 를 훈련시키는 dominant framework 이기 때문에 flow 에 대한 제반 사항을 명료하게 정리할 필요가 있었고, 그 부분이 너무 길어지는 바람에 따로 글을 분리하였습니다.

최대한 제가 아는 지식 내에서 flow 에 대한 내용을 정리해보았지만, 전공이 수학이나 물리는 아니라 action 을 비롯해 여러 개념을 전개하는 과정에서 오류가 있을 수 있습니다. 언제든 날카로운 지적과 가르침을 댓글로 남겨주시면 감사히 배우고 수정하겠습니다.

다음 글부터는 다시 Building Large 3D Generative Model from Scratch 시리즈로 돌아가서, 예고한대로 multi node 환경에서의 training 에 대해 톺아볼 예정입니다. 블로그 자체가 좀 더 이론을 잘 이해하고 설명하는데 치중한 경향이 있는데, 흔치 않게 CS 적인 글이 되겠네요.

감사합니다 :)


You may also likes

References

profile
기타치는AI Researcher

0개의 댓글