플로 매칭

Flow matching; 플로 매칭

잡음 분포에서 데이터 분포로 표본을 옮기는 속도장을 회귀로 직접 학습하고, 생성할 때는 그 속도장을 따라 상미분 방정식을 적분하는 생성 모델 학습법

플로 매칭(flow matching, FM)은 잡음과 데이터를 잇는 확률 경로를 미리 정해 두고, 그 경로를 만들어 내는 속도장(vector field)을 신경망이 평균 제곱 오차로 맞히도록 학습하는 방법이다. 2022년 메타 AI와 바이츠만 연구소의 립먼(Yaron Lipman) 등이 연속 정규화 흐름(CNF)을 ODE 시뮬레이션 없이 학습하는 방법으로 제안했다. 비슷한 시기에 류싱차오(Xingchao Liu) 등의 정류 흐름(rectified flow)과 알베르고(Michael Albergo)·반덴에인던(Eric Vanden-Eijnden)의 확률적 보간(stochastic interpolant)이 같은 생각을 독립적으로 내놓았다.

확산 모델의 잡음 제거 학습과 수식이 매우 가깝고, 가우스 경로를 고르면 확산 모델이 특수한 경우로 들어온다. 스테이블 디퓨전 3, FLUX.1, 메타의 영상 모델 무비 젠(Movie Gen)처럼 2024년 이후 공개된 대형 이미지·영상 생성 모델 다수가 이 방식으로 학습한다. 확산 모델 문서의 "플로 매칭과 정류 흐름" 절이 개요라면, 이 문서는 목적 함수 유도와 구현, 설계 선택을 다룬다.

배경: 연속 정규화 흐름

노멀라이징 플로를 연속 시간으로 확장하면, 시간 $t\in [0,1]$에서 속도장 $v_t(x)$가 상미분 방정식(ODE)

$\frac{d}{dt}{\varphi }_t(x)=v_t({\varphi }_t(x)),{\varphi }_0(x)=x$

으로 흐름 ${\varphi }_t$를 정한다. $t=0$의 단순한 분포 $p_0$(표준 정규)가 이 흐름을 따라 움직여 $t=1$에서 데이터 분포 $p_1\approx q$가 되게 하는 것이 목표다. 시간에 따른 밀도 $p_t$와 속도장은 연속 방정식 ${\partial }_t{p}_t+\mathrm{\nabla }\cdot (p_t{v}_t)=0$으로 묶인다. 예전에는 최대 우도로 학습하려고 학습 단계마다 ODE를 적분하고 그 해를 역전파해야 해서 느렸다.

목적 함수

플로 매칭

원하는 확률 경로 $p_t$와 그것을 만드는 속도장 $u_t$를 안다면, 신경망 $v_{\theta }$를 그냥 회귀시키면 된다.

${ℒ}_{\mathrm{F}\mathrm{M}}(\theta )={𝔼}_{t\sim U[0,1],x\sim p_t}\Vert v_{\theta }(x,t)-u_t(x){\Vert }^2$

문제는 데이터 분포 전체에 대한 $p_t$와 $u_t$를 모른다는 점이다.

조건부 플로 매칭

립먼 등의 핵심 관찰은 데이터 한 점 $x_1$에 조건을 건 경로 $p_t(x\mid x_1)$와 속도 $u_t(x\mid x_1)$는 쉽게 정할 수 있고, 이것들을 데이터에 대해 섞으면 전체 경로가 된다는 것이다. 그리고 조건부 목적 함수

${ℒ}_{\mathrm{C}\mathrm{F}\mathrm{M}}(\theta )={𝔼}_{t,x_1\sim q,x\sim p_t(\cdot \mid x_1)}\Vert v_{\theta }(x,t)-u_t(x\mid x_1){\Vert }^2$

는 ${ℒ}_{\mathrm{F}\mathrm{M}}$과 상수만큼만 달라 θ에 대한 기울기가 같다(논문의 정리 2). 따라서 데이터 한 점과 잡음 한 점만 뽑아 계산할 수 있는 회귀 손실로, 알 수 없는 전체 속도장을 학습하게 된다. 확산 모델의 잡음 제거 점수 매칭이 조건부 점수로 전체 점수를 학습하는 것과 같은 구조다.

가우스 경로와 최적 수송 경로

조건부 경로를 $p_t(x\mid x_1)=𝒩(x;{\mu }_t(x_1),{\sigma }_t(x_1{)}^{2}I)$로 두면, 흐름 ${\psi }_t(x_0)={\sigma }_t{x}_0+{\mu }_t$와 속도 $u_t=\frac{{{\sigma }_t}^{\prime }}{{\sigma }_t}(x-{\mu }_t)+{{\mu }_t}^{\prime }$가 정해진다.

• 확산 경로: ${\mu }_t,{\sigma }_t$를 분산 보존(VP)·분산 폭발(VE) 확산의 일정표에 맞추면 기존 확산 모델의 확률 흐름 ODE와 같은 경로가 된다.
• 최적 수송(OT) 경로: 평균과 표준편차를 시간에 선형으로 바꾼다. ${\mu }_t=t{x}_1,{\sigma }_t=1-(1-{\sigma }_{\min })t$이면

${\psi }_t(x_0)=(1-(1-{\sigma }_{\min })t)x_0+t{x}_1,{ℒ}_{\mathrm{C}\mathrm{F}\mathrm{M}}=𝔼\Vert v_{\theta }({\psi }_t(x_0),t)-(x_1-(1-{\sigma }_{\min })x_0){\Vert }^2$

이 경로는 두 가우스 분포 사이의 최적 수송 변위 보간(displacement interpolation)과 같아서, 조건부 궤적이 직선이고 속도 방향이 시간에 따라 바뀌지 않는다. 논문은 확산 경로 궤적이 목표를 지나쳤다가 되돌아오는 반면 OT 경로는 곧게 간다고 설명했다.

${\sigma }_{\min }\to 0$이면 결국 "잡음 $x_0$와 데이터 $x_1$을 직선으로 잇고, 그 직선의 속도 $x_1-x_0$를 맞혀라"라는 아주 단순한 규칙이 된다. 정류 흐름의 1단계 학습과 같은 식이다.

시간 방향 표기

논문마다 시간 방향이 다르므로 코드를 읽을 때 주의한다. 립먼 등과 정류 흐름 원 논문은 $t=0$이 잡음, $t=1$이 데이터다. 스테이블 디퓨전 3 논문과 Hugging Face diffusers의 FlowMatchEulerDiscreteScheduler는 반대로 $z_t=(1-t)x_0+tϵ$, 곧 $t=1$(σ=1)이 잡음이다. 그래서 예측 대상 속도의 부호도 $x_1-x_0$와 $ϵ-x_0$로 갈린다.

정류 흐름과 리플로

정류 흐름은 두 분포 ${\pi }_0,{\pi }_1$에서 뽑은 점을 직선으로 잇고 그 속도를 최소 제곱으로 학습한다. 학습된 ODE는 원래 짝짓기(coupling)보다 볼록 수송 비용이 늘지 않는 새로운 결정적 짝짓기를 만든다. 학습한 모델로 잡음 $Z_0$에서 만든 결과 $Z_1$을 짝으로 삼아 다시 학습하는 리플로(reflow)를 반복하면 궤적이 점점 곧아지고, 곧은 경로는 오일러 한 걸음으로도 정확히 따라갈 수 있다. 논문은 리플로를 한 번만 해도 거의 직선 흐름이 되어 한 번의 오일러 단계로 좋은 결과가 나온다고 보고했고, 여기에 증류를 더해 1단계 생성기를 만들 수 있다.

미니배치 안에서 잡음과 데이터를 무작위로 짝짓는 대신 최적 수송 계획으로 짝지으면(OT-CFM) 서로 엇갈리는 경로가 줄어 학습이 안정되고 추론이 빨라진다는 보고도 있다. 이 방식은 출발 분포가 가우스일 필요도 없어 세포 상태 변화나 영상 간 변환처럼 두 데이터 분포 사이의 수송에도 쓴다.

확산 모델과의 관계

항목	확산 모델(DDPM 계열)	플로 매칭(OT 경로, 정류 흐름)
순방향 경로	$x_t=\sqrt{{\overline{\alpha }}_t}{x}_0+\sqrt{1-{\overline{\alpha }}_t}ϵ$, 잡음 일정표에 따름	$x_t=(1-t)x_0+t{x}_1$ 직선 보간
학습 대상	잡음 $ϵ$(또는 점수, v-예측)	속도 $x_1-x_0$
끝점	무한 시간에서만 순수 잡음, 실제로는 근사	유한 시간(t=0)에 정확히 잡음 분포
생성	SDE 또는 확률 흐름 ODE를 수십~수백 단계	ODE를 오일러·중점법 등으로 적분
출발 분포	가우스	가우스가 기본이지만 임의 분포 가능

가우스 경로를 쓰는 한 두 방법은 같은 확률 경로를 다른 매개변수화로 학습하는 것이고, 잡음 예측과 속도 예측은 선형 변환으로 서로 바꿀 수 있다. 차이는 경로(잡음 일정표)의 선택과 시간 표본 분포, 그리고 손실 가중치에서 나온다. 립먼 등은 확산 경로를 플로 매칭으로 학습하는 것만으로도 점수 매칭보다 안정적이었고, OT 경로가 CIFAR-10과 ImageNet 32·64에서 우도(BPD), FID, ODE 함수 평가 횟수 모두 가장 좋았다고 보고했다. 예를 들어 ImageNet 64×64에서 DDPM은 FID 17.36에 함수 평가 264회, FM-OT는 FID 14.45에 138회였다.

대형 모델에서의 쓰임

• 스테이블 디퓨전 3(2024): 61가지 확산·정류 흐름 정식화를 같은 조건에서 비교해, 정류 흐름에 시간 $t$를 로짓 정규 분포(logit-normal, $m=0,s=1$)로 뽑아 중간 시점에 학습을 몰아주는 방식이 가장 일관되게 좋았다고 보고했다. 고해상도에서는 같은 시점이라도 신호가 덜 지워지므로 시간 일정표를 잡음 쪽으로 미는 해상도 의존 이동(shift, 1024² 기준 $\alpha =3.0$)을 썼다. 가장 큰 모델은 80억 매개변수의 MMDiT였다. → 확산 트랜스포머
• FLUX.1(2024): 블랙 포레스트 랩스는 2024년 8월 1일 공개하며 모든 공개 FLUX.1 모델이 멀티모달·병렬 확산 트랜스포머 블록을 섞은 구조를 120억 매개변수로 키운 것이고, 확산을 특수한 경우로 포함하는 플로 매칭을 바탕으로 한다고 밝혔다. FLUX.1 [schnell]은 120억 매개변수 정류 흐름 트랜스포머를 잠재 적대 확산 증류로 1~4단계 생성이 되게 만든 모델로 Apache 2.0 라이선스다.
• 메타 무비 젠(2024): 가장 큰 영상 모델이 300억 매개변수인 트랜스포머로, 플로 매칭 목적 함수로 학습했다. 논문은 플로 매칭이 종단 신호 대 잡음비 0을 자연스럽게 보장하고 잡음 일정표 선택에 덜 민감하며 확산 손실보다 성능이 좋아 채택했다고 설명했다.
• 음성: 메타의 보이스박스(Voicebox, 2023)는 비자기회귀 플로 매칭 모델로 5만 시간 이상의 음성을 학습해, 문맥 음성을 주고 빈 구간을 채우는 방식으로 제로샷 음성 합성, 잡음 제거, 내용 편집을 한 모델로 처리했다.
• 메타는 2024년 12월 해설서 〈Flow Matching Guide and Code〉와 PyTorch 패키지 flow_matching을 공개했다. 해설서는 플로 매칭이 이미지, 영상, 오디오, 음성, 생체 분자 구조에서 최고 성능을 냈다고 소개한다.

최소 구현

2차원 두 반달 데이터에 OT 경로(σ_min = 0) 조건부 플로 매칭을 적용하고, 오일러 방법 50단계로 표본을 만든다. PyTorch 2.5에서 실행해 생성 표본의 평균·표준편차가 데이터와 거의 같아지는 것을 확인했다.

import math, torch, torch.nn as nn

def two_moons(n):
    t = torch.rand(n) * math.pi
    upper = torch.stack([torch.cos(t), torch.sin(t)], 1)
    lower = torch.stack([1 - torch.cos(t), 0.5 - torch.sin(t)], 1)
    x = torch.where(torch.rand(n, 1) < 0.5, upper, lower)
    return x + 0.05 * torch.randn_like(x)

# 속도장 v_theta(x_t, t): 입력은 좌표 2개와 시간 1개
v = nn.Sequential(nn.Linear(3, 256), nn.SiLU(),
                  nn.Linear(256, 256), nn.SiLU(),
                  nn.Linear(256, 2))
opt = torch.optim.Adam(v.parameters(), lr=1e-3)

for step in range(5000):
    x1 = two_moons(512)                     # 데이터 (t = 1)
    x0 = torch.randn_like(x1)               # 잡음 (t = 0)
    t = torch.rand(x1.shape[0], 1)
    xt = (1 - t) * x0 + t * x1              # 직선 경로 위의 점
    target = x1 - x0                        # 조건부 속도 u_t(x | x0, x1)
    loss = ((v(torch.cat([xt, t], 1)) - target) ** 2).mean()
    opt.zero_grad(); loss.backward(); opt.step()

@torch.no_grad()
def sample(n, steps=50):                    # dx/dt = v(x, t)를 오일러 방법으로 적분
    x = torch.randn(n, 2)
    for i in range(steps):
        t = torch.full((n, 1), i / steps)
        x = x + v(torch.cat([x, t], 1)) / steps
    return x

손실 값은 0으로 수렴하지 않는다. 같은 $x_t$를 지나는 직선이 여러 개라 조건부 목표가 서로 다르고, 신경망은 그 평균(주변 속도장)을 배우기 때문이다. 손실이 1 근처에 머물러도 표본은 잘 나온다.

이미지 모델에서는 위 MLP 대신 U-Net이나 확산 트랜스포머를 쓰고, 조건(텍스트 임베딩 등)을 함께 넣으며, 추론할 때 분류기 없는 가이던스를 속도에 똑같이 적용한다. diffusers에서는 FlowMatchEulerDiscreteScheduler가 시그마 일정표 이동(shift)과 오일러 갱신 $x\leftarrow x+({\sigma }_{\text{next}}-\sigma )v$를 맡는다.

설계 선택과 실무

선택	선택지	고려할 점
경로	OT(직선), VP 확산, 코사인	직선 경로가 가장 단순하고 대형 모델의 기본값
시간 표본 분포	균등, 로짓 정규, 모드 표본	중간 시점이 가장 어렵고 정보가 많아 SD3는 로짓 정규 사용
짝짓기	독립(무작위), 미니배치 OT, 리플로	경로 교차를 줄일수록 적은 단계로 생성 가능, 대신 학습 비용 증가
해상도 보정	시간 일정표 이동(shift)	해상도가 높을수록 잡음 쪽으로 이동
ODE 풀이	오일러, 중점법, 고차 적응형	경로가 곧을수록 저차 방법으로도 충분
단계 축소	리플로, 증류, 적대 증류	1~4단계 생성 모델(예: FLUX.1 [schnell])

같이 보기

• 확산 모델
• 노멀라이징 플로
• 확산 트랜스포머
• 분류기 없는 가이던스
• 생성 모델
• 스테이블 디퓨전
• 음성 합성

참고 문헌

• facebookresearch/flow_matching 공식 저장소
• 정류 흐름 공식 코드
• TorchCFM(조건부 플로 매칭) 라이브러리