콘텐츠로 이동

11. Flow Matching

이 챕터의 목적

Flow Matching은 간단한 source distribution에서 출발한 sample을 data distribution으로 옮기는 시간에 따른 vector field를 학습한다. 이 장의 목표는 수치 ODE를 유도하거나 최신 변형을 구현하는 것이 아니다. 다음 세 가지를 설명할 수 있으면 충분하다.

  • 무엇이 흐르고 무엇을 학습하는가
  • 학습과 생성이 각각 어떤 계산인가
  • Diffusion과 무엇을 공유하고 어디서 관점이 갈리는가

미리 알아볼 말

아래 항목은 계산 숙련보다 문장에서 역할을 알아보는 수준이면 된다.

  • Probability distribution: 가능한 sample과 그 상대적인 가능성을 나타낸다.
  • Vector field: 각 위치와 시간에서 어느 방향으로 얼마나 움직일지를 지정한다.
  • Ordinary differential equation, ODE: 현재 위치와 velocity가 다음 움직임을 결정한다.
  • Numerical solver: 연속적인 ODE를 유한한 step으로 근사해 계산한다.
  • Diffusion의 noise와 denoising: 10장에서 본 생성 경로와 비교하기 위한 기준이다.

1. 생성 문제를 운반 문제로 다시 보기

표준 Gaussian처럼 sample을 얻기 쉬운 분포를 \(p_0\)라고 하고, 우리가 만들고 싶은 data distribution을 \(p_1\)이라고 하자. 생성은 \(x_0 \sim p_0\)를 뽑은 뒤 \(x_1 \sim p_1\)처럼 보이는 sample로 바꾸는 일이다.

여기서는 \(t=0\)이 source noise이고 \(t=1\)이 data다. 반면 10장의 diffusion 표기에서는 \(x_0\)가 clean data이고 \(x_T\)가 noise다. 두 장의 시간축은 표기상 서로 반대 방향을 가리키므로, 같은 아래첨자를 같은 상태로 해석하지 않는다.

Flow 기반 관점은 이 변환을 한 번의 복잡한 함수로만 보지 않는다. \(t=0\)에서 \(t=1\)까지 중간 분포 \(p_t\)를 놓고 sample이 연속적인 경로를 따라 움직인다고 생각한다. 시간 \(t\)에서 sample \(x_t\)의 움직임은 다음 ODE로 표현할 수 있다.

\[ \frac{d x_t}{d t}=v_\theta(x_t,t) \]

\(v_\theta\)는 신경망으로 표현한 velocity field다. 입력은 현재 sample과 시간이고, 출력은 지금 sample이 움직여야 할 방향과 크기다. 학습이 잘 되면 \(p_0\)의 sample들을 이 field에 따라 이동시켰을 때 최종 분포가 \(p_1\)에 가까워진다.

중요한 구분은 다음과 같다.

  • Trajectory는 sample 하나가 시간에 따라 지나가는 경로다.
  • Probability path는 sample 전체의 분포가 시간에 따라 변하는 경로다.
  • Vector field는 각 위치와 시간에서 trajectory의 velocity를 정한다.

2. 왜 Flow Matching이 필요한가

Continuous Normalizing Flow, CNF도 ODE를 이용해 분포를 변환한다. 하지만 일반적인 likelihood 기반 CNF 학습에는 density 변화와 Jacobian trace를 다루는 계산이 들어간다. Flow Matching은 적절한 target velocity를 구성하고 이를 직접 회귀하는 방식으로 vector field를 학습할 수 있게 한다.

대표적인 형태의 목적함수는 다음처럼 읽을 수 있다.

\[ \mathcal{L}_{FM}=\mathbb{E}_{t,x_t}\left[\lVert v_\theta(x_t,t)-u_t(x_t)\rVert^2\right] \]
  • \(t\)는 중간 시간을 나타낸다.
  • \(x_t\)는 선택한 probability path 위의 중간 sample이다.
  • \(u_t(x_t)\)는 그 위치와 시간에서의 target velocity다.
  • 모델은 \(v_\theta\)가 \(u_t\)에 가까워지도록 학습한다.

실제 data distribution 전체의 target field를 직접 계산하기는 어렵다. Conditional Flow Matching은 data sample과 source sample 등의 조건으로 다루기 쉬운 conditional path와 velocity를 만든 뒤 회귀한다. 적절한 조건에서는 conditional objective를 최적화해 원하는 marginal flow를 얻을 수 있다. 초심자에게 핵심은 증명이 아니라 생성 문제를 velocity prediction 문제로 바꾼다는 점이다.

3. 학습과 sampling은 다르다

학습에서는 시간 \(t\)와 중간 sample \(x_t\)를 만들고 target velocity를 계산해 회귀 loss를 줄인다. 이때 neural network parameter는 2장과 3장에서 본 backpropagation과 optimizer로 갱신된다.

생성에서는 다음 순서를 따른다.

  1. \(x_0 \sim p_0\)를 sampling한다.
  2. 학습된 \(v_\theta(x,t)\)를 ODE의 오른쪽 항으로 사용한다.
  3. Numerical solver로 \(t=0\)부터 \(t=1\)까지 적분한다.
  4. 최종 상태 \(x_1\)을 생성 sample로 사용한다.

따라서 Flow Matching은 backpropagation을 대체하는 학습법이 아니다. 무엇을 예측하도록 학습할지 정하는 objective이며, parameter 학습에는 보통 backpropagation을 사용한다. 생성 시점에는 gradient update가 아니라 ODE integration을 수행한다.

4. Diffusion과 비교하기

Diffusion과 Flow Matching은 완전히 무관한 계열이 아니다. 둘 다 쉬운 분포와 data distribution 사이의 time-dependent path를 사용하고, 시간 정보를 받는 neural network를 학습하며, 생성 시 여러 step을 거칠 수 있다. 또한 stochastic diffusion을 설명하는 SDE에는 같은 marginal distribution을 갖는 probability-flow ODE가 대응될 수 있다.

하지만 입문 단계에서는 다음 관점 차이가 유용하다.

  • Diffusion은 forward noising process와 그 reverse 또는 denoising을 중심으로 설명한다. 생성 과정은 확률적인 reverse SDE 또는 결정론적인 ODE와 sampler로 표현할 수 있다.
  • Flow Matching은 선택한 probability path를 실현하는 velocity field를 회귀하는 문제에서 출발한다. 생성 과정은 보통 ODE integration으로 설명한다.

따라서 Diffusion은 noise prediction, Flow Matching은 velocity prediction이라는 문장은 첫 직관으로는 유용하지만 완전한 구분은 아니다. Diffusion 모델도 parameterization에 따라 velocity를 예측할 수 있고, 두 분야는 probability-flow ODE 등에서 연결된다. 정확한 비교에는 path, objective, stochasticity, parameterization, solver를 함께 봐야 한다.

5. Path와 solver도 모델의 일부다

같은 source와 target을 잇더라도 중간 path는 여러 방식으로 설계할 수 있다. Path가 크게 휘거나 velocity가 급변하면 numerical solver가 작은 step을 요구할 수 있다. 비교적 곧은 trajectory는 sampling step을 줄이는 데 유리할 수 있지만, 실제 성능은 data geometry, 학습 objective, network, solver의 상호작용에 달려 있다.

이 맥락에서 다음 용어를 위치만 알아둔다.

  • Optimal transport path: 운반 비용과 geometry를 고려해 source와 target을 연결하는 관점이다.
  • Rectified flow: trajectory를 더 곧게 만드는 방향의 학습과 재결합을 탐구하는 계열이다.
  • Solver order와 step size: 속도와 discretization error 사이의 trade-off를 만든다.
  • Adaptive solver: 추정 error에 따라 step size를 조정한다. Sample마다 계산량이 달라질 수 있다.

Solver step을 줄이면 언제나 같은 품질을 유지한다거나, 특정 path가 모든 문제에서 가장 낫다고 가정하면 안 된다. Wall-clock time에는 network evaluation 비용과 hardware 효율도 영향을 준다.

6. 조건부 생성

Label, text, property, 구조적 constraint 같은 조건 \(c\)가 있다면 vector field를 \(v_\theta(x,t,c)\)로 만들 수 있다. 이 경우 모델은 source sample을 아무 data sample로 옮기는 것이 아니라 조건과 맞는 영역으로 운반하도록 학습된다.

조건부 생성의 평가는 두 축을 함께 봐야 한다.

  • Fidelity와 diversity: 생성물이 data distribution을 얼마나 잘 반영하고 서로 충분히 다양한가
  • Condition adherence: 생성물이 주어진 조건을 얼마나 충실히 만족하는가

과학 응용에서는 여기에 물리적 유효성, constraint 만족, downstream simulation의 안정성을 추가해야 한다.

핵심으로 가져갈 내용

  • Flow Matching은 source distribution을 target distribution으로 옮기는 time-dependent vector field를 학습한다.
  • 학습은 주로 velocity regression과 backpropagation이고, 생성은 학습된 ODE의 numerical integration이다.
  • Probability path와 solver 선택은 생성 경로와 계산비용에 영향을 준다.
  • Diffusion과 Flow Matching은 연결점이 많다. 둘을 architecture 이름이나 단순한 경쟁 관계로 보면 안 된다.

이름과 위치만 알아둘 내용

  • Continuous Normalizing Flow와 likelihood 학습
  • Conditional Flow Matching의 이론적 정당화
  • Optimal transport conditional path
  • Rectified flow와 reflow
  • Probability-flow ODE와 reverse-time SDE
  • 고차 및 adaptive ODE solver

자주 생기는 혼동과 한계

ODE이므로 정확한가

아니다. Learned vector field 자체에 근사 오차가 있고 numerical solver에도 discretization error가 있다. ODE를 사용한다는 사실만으로 data distribution과 정확히 일치하지 않는다.

Sampling step 수만 비교하면 되는가

아니다. 한 step당 network evaluation 수, model size, batch size, hardware, solver tolerance를 함께 보고 wall-clock cost를 측정해야 한다.

Flow Matching은 diffusion을 완전히 대체하는가

그렇게 단정할 수 없다. 두 계열은 겹치는 수학적 구조를 가지며 응용, 조건부 제어, architecture, solver 생태계가 다르다. 실제 선택은 목표와 검증 결과에 따라 달라진다.

좋은 생성 metric이면 과학적으로 유효한가

아니다. 일반 생성 metric과 별도로 constraint, 물성, 안정성, novelty, downstream utility를 평가해야 한다.

앞뒤 챕터와 연결하기

  • 2장과 3장: velocity network도 backpropagation과 optimizer로 학습한다.
  • 7장: Transformer는 vector field를 표현하는 backbone이 될 수 있다.
  • 10장: diffusion의 time-dependent path, score, probability-flow ODE가 직접적인 비교축이다.
  • 12장: RL은 생성모델 다음 단계가 아니라 reward를 사용하는 독립 branch다.
  • 13장: 과학 생성에서는 물리적 constraint와 OOD 평가가 추가된다.

스스로 확인하기

  • Flow Matching의 학습 단계와 생성 단계에서 각각 어떤 계산을 하는가?
  • Diffusion과 Flow Matching을 path, prediction target, sampling dynamics라는 축으로 비교하면 무엇이 같고 다른가?
  • ODE solver의 step 수만으로 두 생성모델의 실제 계산비용을 비교하기 어려운 이유는 무엇인가?

원전과 공식 읽을거리

원전은 수식과 이론을 모두 이해하기 위한 필수 과제가 아니다. 먼저 abstract, introduction, 핵심 그림을 통해 각 논문이 어떤 path와 학습 문제를 제안하는지 확인한다.