흐름 매칭과 최적수송: 확률 경로를 직선으로 펴기

이미지 생성 모델의 릴리스 노트를 보면 요즘 이런 줄이 자주 있습니다. 「flow matching 기반으로 바꿔 4스텝 샘플링에서도 품질이 유지됩니다.」 지난 글까지 다룬 확산 모델은 스텝을 20~50번 밟아야 했는데, 흐름 매칭이라는 이름이 붙은 모델들은 4번이나 8번에서 쓸 만한 그림을 냅니다.

스텝 수를 줄이는 방법은 여럿입니다. 고차 solver를 쓰는 것도 그중 하나였고, 그건 같은 곡선을 더 정확히 따라가는 접근이었습니다. 흐름 매칭이 하는 것은 다릅니다 — 따라가야 할 곡선 자체를 직선에 가깝게 만듭니다. 곡선이 직선이면 오일러 한 스텝으로도 끝까지 갑니다.

이 글은 그 「직선으로 편다」가 정확히 무슨 뜻인지, 무엇을 학습해야 그게 되는지, 그리고 왜 그 자리에서 KL이 아니라 수송 비용이 필요해지는지를 정리합니다.

확률 경로와 속도장

먼저 낱말 셋을 정의합니다.

확률 경로는 시각 $t \in [0,1]$ 마다 하나씩 놓인 분포의 열 $\{p_t\}$ 입니다. 양끝을 $p_0$ (뽑기 쉬운 것, 보통 표준정규분포)와 $p_1$ (데이터 분포)로 고정합니다.

속도장 $u_t(x)$ 는 그 경로를 만드는 벡터장입니다. 입자 하나를 $p_0$ 에서 뽑아 상미분방정식

$$\frac{dx}{dt} = u_t(x), \qquad x(0) \sim p_0$$

를 따라 흘려보냈을 때, 시각 $t$ 에서 그 입자의 분포가 정확히 $p_t$ 가 되면 「$u_t$ 가 $p_t$ 를 만든다」고 합니다. 둘의 관계는 연속방정식이 적어 줍니다 — 밀도가 사라지지도 생기지도 않고 흐름을 따라 옮겨진다는 보존 법칙입니다.

$$\frac{\partial p_t}{\partial t} + \nabla \cdot \bigl( p_t \, u_t \bigr) = 0$$

여기까지는 지난 글의 확률 흐름 ODE와 같은 그림입니다. 실제로 확산 모델의 확률 흐름 ODE도 이 틀의 한 사례입니다 — 잡음을 점점 섞는 경로를 $p_t$ 로 골랐을 때의 속도장이 그것입니다. 흐름 매칭은 경로를 우리가 고르겠다는 입장입니다.

가장 단순한 선택은 두 끝점을 직선으로 잇는 것입니다.

$$x_t = (1-t)\,x_0 + t\,x_1, \qquad x_0 \sim p_0,\ x_1 \sim p_1$$

이 경로에서 한 짝 $(x_0, x_1)$ 을 고정하면 속도는 시간에 무관한 상수입니다.

$$u_t(x_t \mid x_0, x_1) = \frac{d}{dt}\bigl[(1-t)x_0 + t x_1\bigr] = x_1 - x_0$$

이것을 조건부 속도장이라고 부릅니다. 짝을 알고 있을 때의 속도라는 뜻입니다.

알 수 없는 것을 알 수 있는 것으로 바꾸기

문제는 신경망이 짝을 모른다는 것입니다. 샘플링할 때 가진 것은 지금 위치 $x$ 와 시각 $t$ 뿐이고, 그 위치를 지나는 짝은 여럿입니다. 그 자리에서 옳은 속도는 그 여럿의 평균입니다.

$$u_t(x) = \mathbb{E}\bigl[\, x_1 - x_0 \;\big|\; x_t = x \,\bigr]$$

이것이 주변 속도장이고, 우리가 학습하고 싶은 대상입니다. 그런데 이 기댓값은 계산할 수 없습니다 — 안쪽에 $p_1$ 이 들어 있고 그게 바로 모르는 것이니까요.

흐름 매칭의 핵심은 계산할 수 없는 목적식과 계산할 수 있는 목적식이 같은 최적해를 가진다는 것입니다. 두 손실을 나란히 적습니다.

$$\begin{aligned} \mathcal{L}_{\text{FM}}(\theta) &= \mathbb{E}_{t,\; x \sim p_t} \bigl\| v_\theta(x,t) - u_t(x) \bigr\|^2 \\ \mathcal{L}_{\text{CFM}}(\theta) &= \mathbb{E}_{t,\; x_0 \sim p_0,\; x_1 \sim p_1} \bigl\| v_\theta(x_t,t) - (x_1 - x_0) \bigr\|^2 \end{aligned}$$

위쪽은 못 구하는 $u_t(x)$ 가 들어 있어 계산이 안 되고, 아래쪽은 데이터 한 개와 노이즈 한 개만 뽑으면 바로 계산됩니다. 아래쪽을 조건부 흐름 매칭 목적식이라고 부릅니다.

둘이 같은 최적해를 가진다는 것을 확인해 보겠습니다. 제곱을 펼칩니다.

$$\bigl\| v_\theta - (x_1-x_0) \bigr\|^2 = \|v_\theta\|^2 - 2 \bigl\langle v_\theta,\ x_1-x_0 \bigr\rangle + \|x_1-x_0\|^2$$

$x_t = x$ 로 고정하고 짝에 대한 기댓값을 취합니다. $v_\theta(x,t)$ 는 짝에 무관하므로 첫 항은 그대로고, 둘째 항의 기댓값은 내적 안으로 들어갑니다.

$$\mathbb{E}\bigl[\langle v_\theta,\ x_1-x_0\rangle \bigm| x_t = x\bigr] = \bigl\langle v_\theta(x,t),\ \mathbb{E}[x_1-x_0 \mid x_t=x] \bigr\rangle = \bigl\langle v_\theta(x,t),\ u_t(x) \bigr\rangle$$

셋째 항 $\|x_1-x_0\|^2$ 는 $\theta$ 를 포함하지 않으니 상수입니다. 따라서

$$\mathcal{L}_{\text{CFM}}(\theta) = \mathcal{L}_{\text{FM}}(\theta) + \text{(θ와 무관한 상수)}$$

이고 두 목적식의 기울기가 같습니다. 짝마다 다른 목표로 회귀를 시켜도 최소제곱의 최적해는 목표의 조건부 기댓값이므로, 신경망이 도달하는 곳은 평균 하나입니다. 이것은 확산 모델에서 「노이즈를 예측하라」고 시키면 모델이 스코어를 배우게 되던 것과 정확히 같은 구조입니다.

학습 반복문이 이렇게 짧아집니다.

x1 = next(data_loader)                    # 데이터
x0 = torch.randn_like(x1)                 # 노이즈
t = torch.rand(x1.shape[0], 1)            # [0, 1] 균등
xt = (1 - t) * x0 + t * x1                # 직선 경로 위의 한 점
loss = ((model(xt, t) - (x1 - x0)) ** 2).mean()

시각별 가중치도, 잡음 스케줄도, $\bar{\alpha}_t$ 도 없습니다. 목표가 그냥 $x_1 - x_0$ 입니다.

직선으로 폈는데 왜 곡선이 나오나

조건부 경로는 직선입니다. 그런데 실제로 샘플러가 따라가는 것은 주변 속도장이고, 그것은 직선이 아닙니다. 한 점을 지나는 짝이 여럿이고 그들의 목표가 서로 다르면, 평균 속도는 각 짝의 속도와 다릅니다.

1차원 장난감으로 재 보겠습니다. $p_0 = N(0,1)$, $p_1 = 0.5\,N(-4, 0.5^2) + 0.5\,N(4, 0.5^2)$ 로 두면 주변 속도장을 닫힌 꼴로 적을 수 있습니다. $x_t \mid x_1 \sim N(t x_1, (1-t)^2)$ 이므로 $x_1 - x_0 = (x_1 - x_t)/(1-t)$ 이고, 따라서

$$u_t(x) = \frac{\mathbb{E}[x_1 \mid x_t = x] - x}{1-t}$$

입니다. 오른쪽의 조건부 기댓값은 정규분포 혼합의 사후확률로 정확히 계산됩니다. 이걸 써서 「경로가 얼마나 직선에서 벗어나는가」를 두 가지로 쟀습니다.

지표	값
$\mathbb{E}\bigl[\|u_t(x_t) - (x_1-x_0)\|^2\bigr]$	4.79
ODE 경로가 양끝을 잇는 직선에서 벗어난 최대 거리	평균 0.56 (이동 거리 평균 3.20의 17%)

조건부 경로를 직선으로 골랐다고 해서 샘플링 경로가 직선이 되지는 않습니다. 짝을 무작위로 지었기 때문입니다. 왼쪽 봉우리로 갈 입자와 오른쪽 봉우리로 갈 입자가 서로를 스쳐 지나가면, 그 교차점에서 속도장은 양쪽의 평균이라 어느 쪽으로도 안 가는 방향을 가리킵니다. 그러다 갈라져야 하니 경로가 휩니다.

여기서 짝을 어떻게 짓느냐가 문제로 올라옵니다.

최적수송: 어떻게 짝지을 것인가

최적수송은 한 분포를 다른 분포로 옮기는 방법 중 이동 비용의 총합이 가장 작은 것을 찾는 문제입니다. 짝짓기를 형식화한 것이 커플링입니다 — 주변분포가 각각 $p$ 와 $q$ 인 결합분포 $\pi(x,y)$ 를 말합니다. 그런 $\pi$ 전체를 $\Pi(p,q)$ 라고 쓰면 바서슈타인 거리는 이렇게 정의됩니다.

$$W_2(p, q)^2 = \inf_{\pi \in \Pi(p,q)} \ \mathbb{E}_{(x,y) \sim \pi} \bigl[\, \|x - y\|^2 \,\bigr]$$

「양끝의 분포를 지키면서 짝을 지을 때, 이동 거리 제곱의 평균을 가장 작게 만들 수 있는 값」입니다. 독립 짝짓기 $\pi = p \otimes q$ 도 $\Pi(p,q)$ 의 원소이므로 후보 중 하나일 뿐이고, 보통은 최적이 아닙니다.

1차원에서는 최적 커플링이 놀랄 만큼 간단합니다 — 양쪽 표본을 정렬해 순서대로 이으면 됩니다. 순서를 지키는 짝짓기는 경로가 교차하지 않고, 교차하지 않으면 어느 시각에도 두 입자가 같은 자리에 있지 않습니다.

왜 KL이 아니라 수송 비용인가

KL 발산을 이미 다뤘는데, 왜 여기서 새 거리가 필요할까요. KL은 두 분포의 밀도 비를 재기 때문입니다.

$$D_{\mathrm{KL}}(p \,\|\, q) = \int p(x) \log \frac{p(x)}{q(x)} \, dx$$

$p$ 가 양수인데 $q$ 가 0인 자리가 있으면 로그가 발산해 값이 $\infty$ 가 됩니다. 그리고 생성 모델이 다루는 데이터는 대개 그렇습니다 — 이미지는 픽셀 공간 전체에 퍼져 있지 않고 훨씬 낮은 차원의 얇은 집합 위에 있습니다. 모델이 만든 분포와 데이터 분포가 겹치는 곳이 없으면 KL은 두 분포가 얼마나 떨어져 있든 늘 $\infty$ 입니다.

값이 같으면 기울기가 없습니다. 「지금보다 조금 가까워지는 방향」을 알려 줄 수 없다는 뜻입니다. 반면 $W_2$ 는 겹침이 아니라 이동 거리를 재므로 떨어진 만큼의 값을 냅니다. 평균이 $a$ 만큼 떨어진 두 정규분포로 확인해 보면 차이가 분명합니다.

분포의 폭 $s$	1.0	0.5	0.2	0.1	0.05
$D_{\mathrm{KL}}$ ($a=3$)	4.5	18.0	112.5	450	1800
$W_2$	3.0	3.0	3.0	3.0	3.0

폭을 좁힐수록 KL은 걷잡을 수 없이 커지는데 실제로 두 분포가 더 멀어진 것은 아닙니다. KL은 「겹치는가」를 재고 $W_2$ 는 「얼마나 옮겨야 하는가」를 잽니다. 짝짓기를 고르는 문제에서는 뒤쪽이 옳은 질문입니다.

최적수송으로 짝지으면 얼마나 나아지나

이제 실제로 재 봅니다. 같은 1차원 장난감에서 표본 2만 개를 뽑아 두 방식으로 짝짓고, 오일러로 스텝 수를 바꿔 가며 도착한 표본과 참값 표본 사이의 $W_2$ 를 계산했습니다. 1차원이라 $W_2$ 도 정렬해서 빼면 나옵니다.

import numpy as np
rng = np.random.default_rng(0)

x0 = rng.standard_normal(20000)
x1 = np.array([-4.0, 4.0])[rng.integers(0, 2, 20000)] + 0.5*rng.standard_normal(20000)
W2 = lambda a, b: np.sqrt(np.mean((np.sort(a) - np.sort(b))**2))

# 최적수송 짝짓기 = 양쪽 정렬
o0, o1 = np.sort(x0), np.sort(x1)
def field_ot(x, t):
    return np.interp(x, (1-t)*o0 + t*o1, o1 - o0)

for N in [1, 2, 4, 8]:
    x, h = x0.copy(), 1.0/N
    for i in range(N):
        x = x + h * field_ot(x, i*h)
    print(f"OT 짝짓기 N={N}  W2={W2(x, target):.4f}")

무작위 짝짓기 N= 1  W2=4.0314
무작위 짝짓기 N= 2  W2=1.2772
무작위 짝짓기 N= 4  W2=0.5968
무작위 짝짓기 N= 8  W2=0.3650
무작위 짝짓기 N=16  W2=0.3013
무작위 짝짓기 N=32  W2=0.2898

OT 짝짓기 N=1  W2=0.2519
OT 짝짓기 N=2  W2=0.2519
OT 짝짓기 N=4  W2=0.2519
OT 짝짓기 N=8  W2=0.2519

표본 오차 바닥      W2=0.2519

최적수송으로 짝지으면 오일러 한 스텝이 곧 바닥입니다. 0.2519는 같은 분포에서 2만 개를 두 번 뽑았을 때의 표본 오차라 더 내려갈 수 없는 값입니다. 무작위 짝짓기는 같은 자리에 오는 데 32스텝이 필요했고, 그마저도 0.2898로 바닥에 못 닿았습니다.

이유는 앞의 직선성 지표에 있습니다. 최적수송 짝짓기에서는 경로가 교차하지 않으므로 한 점을 지나는 짝이 하나뿐이고, 따라서 주변 속도장이 조건부 속도와 같습니다. 앞의 지표가 4.79에서 정확히 0이 됩니다. 속도가 경로를 따라 상수이면 오일러의 이산화 오차도 0입니다 — 지난 글에서 오일러의 국소 오차가 테일러 이차항, 즉 속도의 변화율에서 나온다고 했는데, 변화율이 0이니 오차도 0입니다.

실제 모델에서는 어디까지 되나

1차원의 정렬은 마법 같지만 고차원에서는 그대로 쓸 수 없습니다. 전체 데이터셋에 대한 최적수송을 푸는 것은 표본 수의 세제곱에 가까운 비용이 듭니다. 실무에서 쓰는 것은 두 가지 우회입니다.

미니배치 최적수송은 배치 안에서만 최적 짝짓기를 풉니다. 배치가 256이면 256×256 할당 문제라 풀 만하고, 전역 최적은 아니지만 무작위보다는 훨씬 덜 교차합니다.

재정류(rectification)는 한 번 학습한 모델로 $(x_0, x_1)$ 짝을 생성한 뒤, 그 짝으로 다시 학습합니다. 모델이 만든 짝은 이미 ODE 경로로 이어져 있어 교차가 적고, 이 과정을 반복하면 경로가 점점 곧아집니다. 「1스텝 생성」을 내세우는 모델들이 이 절차를 씁니다.

둘 다 완전한 직선을 주지는 않습니다. 그래서 4스텝이지 1스텝이 아닙니다. 그리고 짝짓기를 아무리 잘해도 속도장을 학습한 신경망 자체의 오차는 스텝으로 줄지 않습니다 — 지난 글에서 확산 모델에 대해 적었던 것과 같은 한계가 여기에도 그대로 있습니다.

그래서 4스텝은 어디서 오나

「flow matching 기반으로 바꿔 4스텝에서도 품질이 유지됩니다」라는 릴리스 노트로 돌아가면, 그 4가 어디서 왔는지 말할 수 있습니다.

고차 solver를 붙여 얻은 것이 아닙니다. 경로를 다시 고른 것입니다. 조건부 경로를 직선으로 두고, 짝짓기를 최적수송 쪽으로 밀어 주변 속도장이 그 직선을 그대로 물려받게 만든 결과입니다. 속도가 경로 위에서 변하지 않으면 스텝을 잘게 쪼갤 이유가 없어집니다. 스텝 수는 곡률에 대한 값이지 품질에 대한 값이 아니었던 것입니다.

정리

• 확률 경로 $\{p_t\}$ 와 그것을 만드는 속도장 $u_t$ 는 연속방정식 $\partial_t p_t + \nabla \cdot (p_t u_t) = 0$ 으로 묶인다. 확산의 확률 흐름 ODE도 이 틀의 한 사례다.
• 직선 경로 $x_t = (1-t)x_0 + t x_1$ 의 조건부 속도는 상수 $x_1 - x_0$ 이고, 주변 속도장은 그 조건부 기댓값 $\mathbb{E}[x_1 - x_0 \mid x_t = x]$ 다.
• 조건부 흐름 매칭 목적식은 못 구하는 목적식과 상수만큼만 다르다. 제곱을 펼치면 교차항의 기댓값이 조건부 기댓값과의 내적이 되고, 나머지 항은 $\theta$ 에 무관하다. 그래서 짝마다 다른 목표로 회귀시켜도 평균에 도달한다.
• 조건부 경로가 직선이어도 짝을 무작위로 지으면 샘플링 경로는 휜다. 1차원 장난감에서 직선성 지표가 4.79였고, 경로는 양끝을 잇는 직선에서 이동 거리의 17%만큼 벗어났다.
• 바서슈타인 거리는 커플링 위의 최소 이동 비용이다. 겹치는 곳이 없으면 KL은 거리와 무관하게 $\infty$ 라 기울기가 죽는데, $W_2$ 는 떨어진 만큼의 값을 낸다. 폭을 1.0에서 0.05로 줄이자 KL은 4.5에서 1800으로 갔지만 $W_2$ 는 3.0 그대로였다.
• 1차원 최적 커플링은 양쪽을 정렬해 잇는 것이다. 그렇게 짝지으니 직선성 지표가 0이 되고 오일러 1스텝이 표본 오차 바닥 0.2519에 그대로 닿았다. 무작위 짝짓기는 32스텝에서도 0.2898이었다.
• 고차원에서는 미니배치 최적수송과 재정류로 근사한다. 완전한 직선이 아니므로 1스텝이 아니라 4스텝이고, 속도장 자체의 오차는 여전히 스텝으로 줄지 않는다.

---

읽어주셔서 감사합니다. 😊