논문해석 - Flow Matching for Generative Modeling

들어가며

10편 DDPM은 이미지에 잡음을 조금씩 더하는 과정을 거꾸로 되돌리도록 신경망을 학습해 새 이미지를 만들었습니다. 이 방식은 학습이 안정적이지만, 잡음을 더하는 과정의 수식이 경로를 미리 정해 버립니다. DDPM 논문은 이미지 하나를 만들 때 잡음 제거 신경망을 1000번 실행합니다.

Flow Matching(흐름 정합) 은 경로를 확산 과정에서 유도하지 않고 직접 고릅니다. 잡음에서 데이터로 가는 경로를 먼저 정하고, 그 경로를 만들어 내는 속도(벡터장)를 신경망이 회귀로 맞추게 합니다. 학습 중에 미분방정식을 풀 필요가 없고, 경로를 직선으로 고르면 표본을 만들 때 신경망 실행 횟수도 줄어듭니다. VLA 계보의 π0는 이 방식으로 로봇 행동을 생성합니다.

📄 Flow Matching for Generative Modeling — Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, Matt Le / Meta AI (FAIR)·Weizmann Institute of Science, ICLR 2023

arxiv 식별자는 2210.02747이고, 이 글은 2023년 2월에 올라온 v2를 기준으로 합니다.

초록 요약

저자들은 연속 정규화 흐름(Continuous Normalizing Flow, CNF) 을 바탕으로 한 생성 모델 학습 방법을 제안합니다. 제안한 Flow Matching 목적 함수는 미리 정해 둔 조건부 확률 경로의 벡터장에 신경망 출력을 회귀시켜 CNF를 학습합니다. 학습 과정에서 미분방정식을 수치로 풀지 않으므로 큰 규모로 학습할 수 있습니다.

Flow Matching은 잡음과 데이터 사이를 잇는 가우시안 확률 경로 전체와 함께 쓸 수 있고, 기존 확산 모델의 경로도 그중 한 경우에 포함됩니다. 확산 경로에 Flow Matching을 적용하면 기존 확산 모델 학습보다 안정적이었습니다. 확산이 아닌 경로도 쓸 수 있는데, 저자들은 최적 수송(Optimal Transport, OT) 의 변위 보간으로 정의한 경로를 특히 추천합니다. 이 경로는 확산 경로보다 학습과 표본 생성이 빠르고 일반화도 좋았습니다. ImageNet에서 Flow Matching으로 학습한 CNF는 우도와 표본 품질 모두에서 확산 기반 방법보다 좋은 결과를 냈고, 기성 상미분방정식 풀이기로 빠르고 안정적으로 표본을 만들었습니다.

그림 1 — 최적 수송 경로와 Flow Matching으로 학습한 CNF가 조건 없이 생성한 ImageNet 128×128 표본입니다.

---

1. 배경 — 잡음 분포에서 데이터 분포로의 변환

생성 모델의 입력과 출력

생성 모델은 학습 데이터와 비슷한 새 데이터를 만듭니다. 32×32 컬러 이미지라면 이미지 한 장은 숫자 $32 \times 32 \times 3 = 3072$개로 이루어진 벡터이고, 학습 데이터 전체는 3072차원 공간에 흩어진 점들입니다. 이 점들이 어디에 몰려 있는지를 나타내는 분포를 데이터 분포 $q(x_1)$라고 부릅니다. 논문은 데이터 점을 $x_1$로, 잡음 점을 $x_0$로 적습니다.

데이터 분포의 식은 아무도 모릅니다. 가진 것은 그 분포에서 뽑힌 이미지들뿐입니다. 반면 표준 정규분포 $\mathcal{N}(0, I)$처럼 식을 알고 점을 마음대로 뽑을 수 있는 분포는 쉽게 다룰 수 있습니다. 그래서 생성 모델은 대개 이렇게 동작합니다.

1. 표준 정규분포에서 3072개 숫자를 무작위로 뽑습니다. 이 벡터는 화면에 띄우면 흑백 점이 흩뿌려진 잡음 이미지입니다.
2. 학습한 신경망으로 이 벡터를 조금씩 바꿔 데이터 분포 안의 점, 곧 그럴듯한 이미지로 옮깁니다.

2단계에서 벡터를 "어떤 경로로, 몇 번에 걸쳐" 옮기는지가 생성 모델마다 다릅니다.

확산 모델의 경로

DDPM 같은 확산 모델은 데이터에 잡음을 조금씩 더하는 과정을 먼저 수식으로 정하고, 그 과정을 되돌리는 방법을 학습합니다. 그러면 잡음에서 데이터로 가는 중간 분포들은 잡음을 더하는 수식이 자동으로 결정합니다. 저자들은 이 제약 때문에 쓸 수 있는 경로의 종류가 좁고, 학습 시간이 매우 길며, 표본을 빠르게 만들려면 별도의 가속 기법이 필요하다고 지적합니다.

이 논문은 경로를 확산 과정에서 유도하지 않고, 원하는 경로를 직접 정한 뒤 그 경로를 따르도록 모델을 학습하는 방법을 제안합니다. 그 틀이 연속 정규화 흐름입니다.

---

2. 연속 정규화 흐름

벡터장과 흐름

벡터장(vector field) 은 공간의 각 위치와 시각에 "지금 이 점은 어느 방향으로 얼마나 빠르게 움직여야 하는가"를 알려 주는 함수입니다. 논문은 시각 $t \in [0, 1]$에서 위치 $x$의 속도를 $v_t(x)$로 적습니다. 3072차원 이미지 공간이라면 $v_t(x)$도 숫자 3072개짜리 벡터입니다. 픽셀 값마다 "이번 순간에 얼마나 올리거나 내릴지"가 하나씩 들어 있습니다.

점 하나를 시각 0에 $x$에 두고 이 속도를 따라 움직이면, 시각 $t$의 위치 $\phi_t(x)$가 정해집니다. 이 위치를 정하는 식이 상미분방정식(ordinary differential equation, ODE) 입니다.

$$\frac{d}{dt}\phi_t(x) = v_t(\phi_t(x)), \qquad \phi_0(x) = x$$

첫 식은 "위치가 바뀌는 속도는 지금 위치에서 벡터장이 알려 주는 속도와 같다"는 뜻이고, 둘째 식은 출발점을 정합니다. $\phi_t$를 흐름(flow) 이라고 부릅니다.

컴퓨터로는 이 식을 짧은 시간 간격으로 잘라 풉니다. 가장 단순한 방법은 오일러(Euler) 방법으로, 지금 속도에 시간 간격을 곱해 위치에 더하는 계산을 반복합니다. 1차원 벡터장 $v_t(x) = 1 - x$에서 $x = 0$을 출발점으로, 간격 0.25로 계산하면 다음과 같습니다.

표의 벡터장과 수치는 계산 과정을 보여 주기 위한 임의의 예시 값입니다.

한 줄 내려갈 때마다 벡터장 $v_t$를 한 번 계산합니다. 벡터장이 신경망이라면 신경망을 한 번 실행하는 것이고, 이 횟수를 함수 평가 횟수(number of function evaluations, NFE) 라고 부릅니다. 위 표는 NFE 4로 풀었습니다. 간격을 줄이면 해가 정확해지는 대신 NFE가 늘어납니다. 뒤의 실험에서 쓰는 dopri5는 오차를 보면서 간격을 자동으로 조절하는 풀이기이고, 중점(midpoint) 방법은 한 걸음에 벡터장을 두 번 계산하는 고정 간격 풀이기입니다.

확률 경로

점 하나가 아니라 표준 정규분포에서 뽑은 점 수만 개를 같은 벡터장으로 동시에 움직이면, 점들이 모인 모양, 곧 분포가 시각에 따라 바뀝니다. 시각 $t$의 분포를 $p_t$로 적고, $p_0$부터 $p_1$까지 시간에 따라 바뀌는 분포 전체를 확률 경로(probability path) 라고 부릅니다.

$$p_t = [\phi_t]_* \, p_0$$

$[\phi_t]_*$는 밀어내기(push-forward) 연산으로, "$p_0$에서 뽑은 점들을 $\phi_t$로 옮겼을 때 점들이 이루는 분포"를 뜻합니다. 점들이 한곳으로 모이면 그 자리의 밀도가 높아지고 흩어지면 낮아지는데, 논문의 식 4는 이 밀도 변화를 흐름의 야코비안 행렬식으로 계산합니다. 벡터장 $v_t$가 이렇게 경로 $p_t$를 만들어 낼 때 "$v_t$가 $p_t$를 생성한다(generate)"고 표현합니다.

CNF는 이 벡터장을 파라미터 $\theta$를 가진 신경망 $v_t(x; \theta)$로 두는 모델입니다. 목표는 $p_0$가 표준 정규분포이고 $p_1$이 데이터 분포와 같아지도록 신경망을 학습하는 것입니다.

시뮬레이션 기반 학습의 비용

CNF는 2018년에 제안됐지만 큰 이미지로 키우기 어려웠습니다. 기존의 최대 우도 학습은 학습 데이터 이미지 한 장의 확률을 구할 때마다 상미분방정식을 끝까지 수치로 풀어야 했습니다. 역전파에도 같은 풀이가 한 번 더 필요하고, 풀이는 시간 순서대로 진행되므로 병렬화도 어렵습니다. 저자들은 기존의 시뮬레이션 없는 학습 방법도 고차원에서 계산하기 어려운 적분이 남거나(Rozen et al., 2021), 미니배치 기울기가 편향되는(Ben-Hamu et al., 2022) 문제가 있었다고 정리합니다. 확산 모델은 잡음 제거 점수 정합(denoising score matching)으로 효율적으로 학습할 수 있는 드문 경우였습니다.

Flow Matching은 CNF를 확산 경로에 묶이지 않은 채 시뮬레이션 없이 학습하는 방법입니다.

---

3. Flow Matching 목적 함수

데이터 분포에 도달하는 확률 경로 $p_t$와, 그 경로를 생성하는 벡터장 $u_t$를 누군가 알려 준다고 가정합니다. 그러면 학습은 단순한 회귀 문제가 됩니다.

$$\mathcal{L}_{\text{FM}}(\theta) = \mathbb{E}_{t,\, p_t(x)} \big\| v_t(x) - u_t(x) \big\|^2$$

시각 $t$를 0과 1 사이에서 균등하게 뽑고, 그 시각의 분포 $p_t$에서 점 $x$를 뽑은 뒤, 신경망이 낸 속도 $v_t(x)$와 정답 속도 $u_t(x)$의 제곱 오차를 줄입니다. 손실이 0이 되면 신경망이 정답 벡터장과 같아지고, 그 벡터장을 따라 풀면 $p_t$가 그대로 재현됩니다. 이 과정에서 상미분방정식을 풀 일은 없습니다.

문제는 정답이 없다는 것입니다. 데이터 분포에 도달하는 경로는 무수히 많고, 어떤 경로를 고르더라도 그 경로를 만드는 벡터장 $u_t$를 식으로 쓸 수 있는 경우가 거의 없습니다. 다음 장은 이 문제를 데이터 점 하나 단위로 나눠서 풉니다.

---

4. 조건부 확률 경로와 Conditional Flow Matching

데이터 점 하나를 기준으로 한 경로

데이터 전체가 아니라 학습 이미지 한 장 $x_1$만 생각하면 경로를 쉽게 정할 수 있습니다. 시각 0에는 표준 정규분포이고, 시각 1에는 $x_1$ 주변에 좁게 몰린 분포, 예를 들어 평균 $x_1$, 표준편차 $\sigma$가 아주 작은 정규분포 $\mathcal{N}(x \mid x_1, \sigma^2 I)$가 되도록 중간을 이으면 됩니다. 이 경로를 $x_1$에 대한 조건부 확률 경로 $p_t(x \mid x_1)$라고 부르고, 이 경로를 생성하는 벡터장을 조건부 벡터장 $u_t(x \mid x_1)$라고 부릅니다.

조건부 경로를 합친 주변 경로

모든 학습 이미지의 조건부 경로를 데이터 분포에 따라 평균하면 전체 경로가 됩니다.

$$p_t(x) = \int p_t(x \mid x_1)\, q(x_1)\, dx_1$$

시각 1에서 이 분포는 학습 이미지마다 그 이미지를 평균으로 하는 좁은 정규분포를 하나씩 두고 섞은 분포이므로 데이터 분포와 거의 같습니다. 이 전체 경로를 주변 확률 경로(marginal probability path) 라고 부릅니다.

벡터장도 같은 방식으로 합칩니다.

$$u_t(x) = \int u_t(x \mid x_1)\, \frac{p_t(x \mid x_1)\, q(x_1)}{p_t(x)}\, dx_1$$

분수 부분은 "시각 $t$에 위치 $x$에 있는 점이 학습 이미지 $x_1$의 경로에서 왔을 확률"입니다. 위치 $x$에서의 전체 속도는 각 학습 이미지가 요구하는 속도를 이 확률로 가중 평균한 값입니다.

정리 1 은 이렇게 합친 주변 벡터장 $u_t$가 실제로 주변 확률 경로 $p_t$를 생성한다는 내용입니다. 증명은 확률 경로와 벡터장이 만족해야 하는 연속 방정식(continuity equation)에 두 식을 대입해 확인합니다.

구체적인 예로, 1차원 데이터에 학습 점이 $x_1 = -2$와 $x_1 = +2$ 두 개뿐이라고 하겠습니다. 뒤의 5장에서 소개할 최적 수송 경로를 쓰고 $\sigma_{\min}$을 0으로 두면, 조건부 벡터장은 $u_t(x \mid x_1) = (x_1 - x)/(1 - t)$입니다. 시각 $t = 0.5$, 위치 $x = 0$에서 계산하면 다음과 같습니다.

두 학습 점과 계산 결과는 설명을 위한 임의의 예시 값입니다. 두 조건부 경로가 0을 기준으로 대칭이라 가중치가 같습니다.

위치 0에 있는 점은 두 학습 점 중 어느 쪽 경로에서 왔는지 알 수 없으므로 전체 속도는 0입니다. 개별 학습 점의 속도와 주변 속도가 다를 수 있다는 점이 다음 절의 핵심입니다.

CFM 목적 함수와 기울기의 일치

주변 벡터장은 모든 학습 데이터에 대한 적분이라 여전히 계산할 수 없습니다. 저자들은 대신 조건부 벡터장에 바로 회귀하는 Conditional Flow Matching(CFM) 목적 함수를 제안합니다.

$$\mathcal{L}_{\text{CFM}}(\theta) = \mathbb{E}_{t,\, q(x_1),\, p_t(x \mid x_1)} \big\| v_t(x) - u_t(x \mid x_1) \big\|^2$$

학습 한 걸음은 다음 순서로 진행됩니다.

모든 계산이 식 한 줄이고, 상미분방정식을 풀지 않습니다.

정리 2 는 이 CFM 손실이 원래의 FM 손실과 $\theta$에 무관한 상수만큼만 다르고, 따라서 두 손실의 기울기가 같다는 내용입니다. 기대값으로 보면 CFM을 최소화하는 것이 FM을 최소화하는 것과 같습니다.

앞의 두 점 예시로 보면 이유가 드러납니다. 학습 중에 위치 0, 시각 0.5라는 같은 입력에 대해 어떤 걸음에서는 정답이 $-4$로, 어떤 걸음에서는 $+4$로 주어집니다. 같은 입력에 서로 다른 정답이 섞여 있을 때 제곱 오차를 최소화하는 출력은 정답들의 평균, 곧 0입니다. 신경망은 개별 학습 이미지로 가는 속도를 외우지 않고, 저절로 주변 속도를 학습합니다. "조건부 벡터장에 회귀하면 모델이 학습 이미지 하나만 복원하게 되지 않는가"라는 의문에 대한 답이 이것입니다.

---

5. 가우시안 조건부 경로

평균과 표준편차로 정하는 경로

CFM은 어떤 조건부 경로와도 함께 쓸 수 있습니다. 저자들은 그중 다루기 쉬운 가우시안 경로를 씁니다.

$$p_t(x \mid x_1) = \mathcal{N}\big(x \mid \mu_t(x_1),\, \sigma_t(x_1)^2 I\big)$$

$\mu_t(x_1)$은 시각에 따라 바뀌는 평균, $\sigma_t(x_1)$은 시각에 따라 바뀌는 표준편차 숫자 하나입니다. 경계 조건은 다음과 같습니다.

$\sigma_{\min}$은 충분히 작은 값으로 둡니다.

같은 확률 경로를 생성하는 벡터장은 무수히 많습니다. 예를 들어 회전해도 모양이 같은 분포라면 점들을 제자리에서 빙빙 돌리는 성분을 더해도 분포는 변하지 않습니다. 이런 성분은 분포에 아무 영향 없이 계산만 늘리므로, 저자들은 가장 단순한 흐름을 고릅니다.

$$\psi_t(x) = \sigma_t(x_1)\, x + \mu_t(x_1)$$

표준 정규분포에서 뽑은 $x$에 표준편차를 곱하고 평균을 더하는 변환입니다. 결과는 평균 $\mu_t(x_1)$, 표준편차 $\sigma_t(x_1)$인 정규분포를 따릅니다. 학습에서 중간 점을 만들 때 이 식에 잡음 $x_0$를 넣으면 됩니다.

정리 3 은 이 흐름을 만드는 조건부 벡터장이 하나로 정해지고, 그 식이 다음과 같다는 내용입니다. $f'$는 시각에 대한 미분입니다.

$$u_t(x \mid x_1) = \frac{\sigma_t'(x_1)}{\sigma_t(x_1)}\big(x - \mu_t(x_1)\big) + \mu_t'(x_1)$$

첫째 항은 표준편차가 줄어드는 만큼 점을 평균 쪽으로 당기는 속도이고, 둘째 항은 평균 자체가 움직이는 속도입니다. 이 식에 넣으면 CFM 손실은 다음처럼 잡음 $x_0$만으로 적을 수 있습니다.

$$\mathcal{L}_{\text{CFM}}(\theta) = \mathbb{E}_{t,\, q(x_1),\, p(x_0)} \Big\| v_t\big(\psi_t(x_0)\big) - \frac{d}{dt}\psi_t(x_0) \Big\|^2$$

$\mu_t$와 $\sigma_t$를 무엇으로 고르느냐에 따라 경로가 달라집니다. 논문은 두 가지 경우를 보입니다.

확산 경로

확산 모델의 경로도 이 가우시안 경로의 한 경우입니다. 논문은 시간 방향을 뒤집어 잡음이 $t = 0$, 데이터가 $t = 1$이 되도록 적습니다.

분산 폭발(Variance Exploding, VE) 경로는 평균을 데이터 점에 고정하고 표준편차만 큰 값에서 0으로 줄입니다.

$$p_t(x) = \mathcal{N}(x \mid x_1,\, \sigma_{1-t}^2 I), \qquad u_t(x \mid x_1) = -\frac{\sigma_{1-t}'}{\sigma_{1-t}}(x - x_1)$$

여기서 $\sigma_t$는 $\sigma_0 = 0$에서 출발해 증가하는 함수이고 $\sigma_1 \gg 1$입니다.

분산 보존(Variance Preserving, VP) 경로는 평균을 0에서 데이터 점 쪽으로, 표준편차를 1에서 0 쪽으로 함께 바꿉니다.

$$p_t(x \mid x_1) = \mathcal{N}\big(x \mid \alpha_{1-t}\, x_1,\, (1 - \alpha_{1-t}^2) I\big), \qquad \alpha_t = e^{-\frac{1}{2}T(t)},\quad T(t) = \int_0^t \beta(s)\, ds$$

$\beta$는 잡음 크기 일정(noise schedule) 함수입니다. 정리 3에 넣은 VP 경로의 조건부 벡터장은 논문 식 19이고, 부록 D는 이 벡터장이 Song et al. (2020b)의 확률 흐름 상미분방정식에 쓰인 벡터장과 같다는 것을 유도합니다. 부록 E의 기준선 실험에서는 $\beta(s) = \beta_{\min} + s(\beta_{\max} - \beta_{\min})$, $\beta_{\min} = 0.1$, $\beta_{\max} = 20$을 씁니다.

확산 경로로 학습하더라도 기존의 점수 정합(score matching) 대신 Flow Matching 목적 함수를 쓸 수 있고, 저자들의 실험에서는 이쪽이 더 안정적이었습니다. 확산 경로에는 한계도 하나 남습니다. 확산 과정의 해로 유도한 경로라서 유한한 시간 안에 정확한 표준 정규분포에 도달하지 않고, 실제로는 $p_0$를 적당한 정규분포로 근사해 씁니다. 가우시안 경로를 직접 정하면 이런 근사 없이 $\mu_t$와 $\sigma_t$를 원하는 대로 둘 수 있습니다.

최적 수송 경로

저자들이 더 자연스러운 선택으로 제안하는 경로는 평균과 표준편차를 시간에 따라 직선으로 바꾸는 것입니다.

$$\mu_t(x) = t\, x_1, \qquad \sigma_t(x) = 1 - (1 - \sigma_{\min})\, t$$

정리 3에 넣으면 조건부 벡터장과 조건부 흐름은 다음과 같습니다.

$$u_t(x \mid x_1) = \frac{x_1 - (1 - \sigma_{\min})\, x}{1 - (1 - \sigma_{\min})\, t}, \qquad \psi_t(x) = \big(1 - (1 - \sigma_{\min})\, t\big)\, x + t\, x_1$$

흐름 $\psi_t(x_0)$는 잡음 점 $x_0$와 데이터 점 $x_1$을 잇는 선분 위를 움직입니다. 시각에 대해 미분하면 $x_1 - (1 - \sigma_{\min})\, x_0$로 시각과 무관한 값이므로, CFM 손실은 더 간단해집니다.

$$\mathcal{L}_{\text{CFM}}(\theta) = \mathbb{E}_{t,\, q(x_1),\, p(x_0)} \Big\| v_t\big(\psi_t(x_0)\big) - \big(x_1 - (1 - \sigma_{\min})\, x_0\big) \Big\|^2$$

정답 속도가 "데이터 점에서 (거의) 잡음 점을 뺀 값"이라는 뜻입니다. 1차원에서 $\sigma_{\min} = 0$, 잡음 점 $x_0 = -1$, 데이터 점 $x_1 = 3$으로 두고 간격 0.25의 오일러 방법으로 따라가면 다음과 같습니다.

$\sigma_{\min}$, $x_0$, $x_1$은 계산 과정을 보여 주기 위한 임의의 예시 값입니다.

점은 매 걸음 같은 속도 4로 직선을 따라 움직이고, 오일러 방법 네 걸음이 정확히 데이터 점에 도착합니다. 경로가 곧으면 큰 간격으로 풀어도 오차가 작다는 것이 뒤의 표본 효율 실험과 이어집니다.

이 선택이 "최적 수송"이라고 불리는 이유는 다음과 같습니다. 조건부 흐름 $\psi_t$는 두 정규분포 $p_0(x \mid x_1)$과 $p_1(x \mid x_1)$ 사이의 최적 수송 변위 사상과 같습니다. 최적 수송은 한 분포의 질량을 다른 분포로 옮길 때 옮기는 거리의 제곱 합이 가장 작은 방법이고, 그 중간 경로인 변위 보간(McCann, 1997)에서는 각 점이 일정한 속도로 직선을 따라 움직입니다. 첫 번째 분포가 표준 정규분포인 두 정규분포 사이에서는 이 변위 사상이 위의 $\psi_t$ 식이 됩니다.

다만 이 최적성은 데이터 점 하나에 대한 조건부 흐름의 성질입니다. 저자들은 조건부 흐름을 합친 주변 벡터장이 최적 수송 해가 된다는 뜻은 아니라고 명시하고, 그래도 주변 벡터장이 비교적 단순하게 남을 것으로 기대한다고 적습니다. 앞의 두 점 예시에서도 위치 0의 주변 속도는 0이었고, 이는 두 조건부 경로 어느 쪽의 직선 속도와도 다릅니다.

확산 경로와 최적 수송 경로의 비교

그림 2(a) — 확산 경로의 조건부 점수 함수를 t = 0, 1/3, 2/3, 1에서 그린 것입니다.

그림 2(b) — 최적 수송 경로의 조건부 벡터장입니다. 파란색은 크기가 큰 곳, 빨간색은 작은 곳입니다.

위 그림은 출발 분포와 도착 분포가 같은 두 경로에서, 신경망이 회귀해야 하는 목표를 시각별로 보여줍니다. 확산 모델이 보통 회귀하는 조건부 점수 함수 $\nabla \log p_t(x \mid x_1)$는 시각에 따라 크기와 모양이 크게 바뀝니다. 최적 수송 경로의 조건부 벡터장은 방향이 시각과 무관하게 일정합니다. 식으로도 $u_t(x \mid x_1) = g(t)\, h(x \mid x_1)$처럼 시각 함수와 위치 함수의 곱으로 나뉩니다. 저자들은 이 성질 때문에 신경망이 맞춰야 할 회귀 문제가 더 단순해진다고 봅니다.

그림 8(부록) — VP 확산 경로의 조건부 벡터장입니다. 그림 2와 비교하는 용도입니다.

그림 3 — 확산 조건부 벡터장(a)과 최적 수송 조건부 벡터장(b)으로 표본을 만들 때의 궤적입니다.

확산 경로의 궤적은 곡선을 그리고, 최종 표본을 지나쳤다가 되돌아오기도 합니다. 최적 수송 경로의 궤적은 항상 직선입니다.

---

6. 2차원 체커보드 실험

이미지 실험에 앞서 저자들은 체커보드 무늬로 분포한 2차원 점 데이터를 학습합니다. 모델은 은닉 뉴런 512개짜리 층 5개로 이루어진 다층 퍼셉트론입니다.

그림 4 왼쪽 — (a) 점수 정합·확산 경로, (b) Flow Matching·확산 경로, (c) Flow Matching·최적 수송 경로로 학습한 CNF의 궤적입니다.

그림 4 오른쪽 — 세 모델을 중점 방법으로 NFE 4, 8, 10, 20에 풀어 만든 표본입니다. 묶음 순서는 점수 정합·확산, Flow Matching·확산, Flow Matching·최적 수송입니다.

그림 4 왼쪽은 시각에 따라 점들이 체커보드 무늬를 만들어 가는 과정입니다. 최적 수송 경로는 체커보드 무늬가 훨씬 이른 시각에 나타나고, Flow Matching 두 모델은 점수 정합보다 학습이 안정적이었습니다. 그림 4 오른쪽은 적은 NFE로 풀었을 때의 결과로, 최적 수송 경로 모델이 가장 적은 계산으로 무늬를 복원합니다.

그림 9(부록) — 그림 4와 같은 학습률·조절값으로 ScoreFlow 손실(a)과 DDPM 손실(b)을 써 학습한 CNF의 궤적입니다.

---

7. 이미지 생성 실험

실험 설정

데이터셋은 CIFAR-10과 해상도 32, 64, 128의 ImageNet입니다. 모든 방법은 Dhariwal & Nichol (2021)의 U-Net 구조를 거의 그대로 쓰고, 같은 구조·같은 조절값·같은 학습 세대 수로 학습합니다. 기준선에는 수렴을 위해 더 많은 반복을 허용했습니다. 모두 조건 없는(unconditional) 생성 모델입니다.

비교한 학습 방법은 다섯 가지입니다.

확산 기준선도 표본은 같은 상미분방정식 풀이기로 만듭니다. 점수 정합 모델은 학습한 점수 함수를 부록 식 46의 벡터장으로 바꿔 풀고, DDPM 모델은 잡음 예측값을 점수로 바꾼 뒤 같은 식으로 풉니다. 따라서 아래 표의 DDPM NFE는 DDPM 논문의 원래 표본 생성 절차가 아니라, DDPM 손실로 학습한 모델을 상미분방정식으로 풀었을 때의 값입니다.

평가는 따로 적지 않은 한 dopri5 풀이기를 절대·상대 허용 오차 1e-5로 써서 합니다.

표 3(부록) — 각 모델의 학습 조절값입니다.

CIFAR10과 ImageNet-32는 32비트 정밀도로, ImageNet-64/128/256은 16비트 혼합 정밀도로 학습했습니다. 최적화기는 Adam($\beta_1 = 0.9$, $\beta_2 = 0.999$, 가중치 감쇠 0.0, $\epsilon = 10^{-8}$)입니다. 다항 감소 일정은 워밍업 동안 학습률을 1e-8에서 최대값까지 선형으로 올린 뒤 마지막 걸음까지 1e-8로 선형으로 내립니다.

CIFAR10 열은 세대 수와 반복 수가 서로 맞지 않습니다. 학습 이미지 50,000장을 배치 256으로 1000세대 돌리면 약 195k번 반복이지만 표에는 391k로 적혀 있어, 두 값 중 어느 쪽이 실제 설정인지 논문만으로는 알 수 없습니다.

평가 지표

세 지표를 씁니다.

• 음의 로그 가능도(negative log-likelihood, NLL) — 모델이 시험 이미지에 매기는 확률이 얼마나 높은지를 차원당 비트(bits per dimension, BPD) 로 적은 값입니다. 픽셀 값 0~255를 아무 정보 없이 균등하게 맞히려면 값 하나에 $\log_2 256 = 8$비트가 필요합니다. BPD 3.53이라면 모델이 픽셀 값 하나를 평균 3.53비트로 표현한다는 뜻이고, 낮을수록 좋습니다. CNF는 부록 C의 방법으로 상미분방정식을 거꾸로 풀어 이 확률을 정확히 계산할 수 있습니다.
• FID(Fréchet Inception Distance) — 생성 이미지 묶음과 실제 이미지 묶음을 사전학습된 Inception 신경망에 넣어 뽑은 특징의 통계가 얼마나 다른지를 잰 값입니다. 낮을수록 실제 이미지와 비슷합니다.
• NFE — 적응형 풀이기가 정한 허용 오차에 도달할 때까지 벡터장을 계산한 횟수를 표본 50k개에 대해 평균한 값입니다. 낮을수록 표본을 빨리 만듭니다.

우도와 표본 품질

표 1 왼쪽 — 같은 모델을 방법만 바꿔 학습했을 때의 우도(BPD), 표본 품질(FID), 평가 시간(NFE)입니다.

세 데이터셋과 세 지표 모두에서 FM w/ OT가 가장 좋습니다. 같은 확산 경로를 쓰는 FM w/ Diffusion도 CIFAR-10과 ImageNet 64에서는 점수 정합 두 방법보다 NFE가 낮습니다. 저자들은 CIFAR-10의 FID가 이전 연구들보다 높게 나온 이유를, 이 구조가 CIFAR-10에 맞춰 최적화되지 않았기 때문일 수 있다고 적습니다.

표 1 오른쪽 — ImageNet 128×128 조건 없는 생성 결과입니다.

ImageNet 128×128에서 FM w/ OT의 FID 20.9는 조건 없는 생성 모델 가운데 최고 기록입니다. 자기지도 학습 ResNet50 특징으로 조건을 주는 IC-GAN은 조건 설정이 달라 표에서 뺐다고 밝힙니다.

학습 속도

그림 5 — ImageNet 64×64 학습 중 세대에 따른 FID 변화입니다.

기존 확산 모델은 매우 긴 학습을 씁니다. ScoreFlow는 1.3m번, VDM은 10m번 반복을 보고했습니다. 그림 5에서 FM w/ OT는 다른 방법보다 FID를 더 빨리, 더 낮게 내립니다.

ImageNet-128에서 Dhariwal & Nichol (2021)은 배치 256으로 4.36m번 반복했습니다. Flow Matching은 25% 더 큰 모델로 배치 1.5k, 500k번 반복했고, 학습에 넣은 이미지 수는 33% 적습니다. $4.36\text{m} \times 256 \approx 1116\text{m}$장과 $500\text{k} \times 1.5\text{k} = 750\text{m}$장을 비교한 값입니다.

그림 10(부록) — CIFAR-10 학습 중 dopri5(허용 오차 1e-5)로 표본을 만들 때 필요한 NFE의 변화입니다.

점수 정합으로 학습하면 표본 생성 비용이 학습 도중에 크게 바뀌지만, Flow Matching으로 학습하면 거의 일정하게 유지됩니다.

---

8. 표본 생성 효율

표본은 표준 정규분포에서 잡음 $x_0$를 뽑고, 학습한 벡터장으로 상미분방정식을 $t = 0$부터 $t = 1$까지 풀어 $\phi_1(x_0)$를 구해 만듭니다. 확산 모델은 확률미분방정식(SDE)으로도 표본을 만들 수 있지만 매우 비효율적일 수 있고, 빠른 표본 생성 기법 상당수가 상미분방정식 관점을 씁니다. 같은 계산 비용에서 상미분방정식 풀이기의 오차가 더 작고, 풀이기의 종류도 많기 때문입니다.

표본 경로

그림 6 — ImageNet 64×64로 학습한 세 모델에 같은 잡음을 넣고 풀어 가는 중간 과정입니다. (a) 점수 정합·확산, (b) Flow Matching·확산, (c) Flow Matching·최적 수송입니다.

같은 난수 시드에서 출발해도 최적 수송 경로 모델은 잡음을 거의 일정한 비율로 줄여 이른 시각부터 이미지 윤곽이 나타납니다. 확산 경로 모델은 마지막 시각 직전까지 잡음이 이미지를 덮고 있다가 끝에서 한꺼번에 사라집니다. 세 모델이 최종적으로 만든 이미지 자체도 서로 다릅니다.

적은 NFE에서의 오차와 품질

이번에는 고정 간격 풀이기로 NFE 100 이하에서 표본을 만들어, 표 1의 ImageNet-32 모델들을 비교합니다.

그림 7 — ImageNet 32×32 모델에서 NFE에 따른 수치 오차(a, 중점 방법)와 세 모델의 FID(b~d)입니다.

그림 7(a)는 256개 잡음 시드에 대해, 적은 NFE로 구한 해와 NFE 1000으로 구한 해의 픽셀당 평균제곱오차를 비교합니다. FM w/ OT는 같은 오차에 도달하는 데 확산 모델 NFE의 약 60%만 필요합니다. 그림 7(b)~(d)에서도 FM w/ OT는 NFE가 아주 적을 때 이미 괜찮은 FID를 내어, 품질과 계산 비용의 균형이 가장 좋습니다.

그림 17(부록, 일부) — ImageNet-128로 학습한 FM w/ OT에 같은 잡음을 넣고 NFE 10, 20, 40, 100으로 만든 표본입니다. 원 그림의 여러 행 가운데 두 행을 실었습니다.

그림 18(부록, 일부) — ImageNet 256×256으로 학습한 FM w/ OT에 같은 잡음을 넣고 NFE 10, 20, 40, 60, 100으로 만든 표본입니다. 원 그림의 여러 행 가운데 한 행을 실었습니다.

NFE 10에서도 이미지의 구도와 물체가 이미 정해져 있고, NFE를 늘리면 세부가 선명해집니다.

---

9. 저해상도 이미지를 조건으로 한 생성

Flow Matching은 조건부 생성에도 쓸 수 있습니다. 64×64 이미지를 조건으로 받아 256×256 이미지를 만드는 초해상도(super-resolution) 실험을 SR3(Saharia et al., 2022)의 평가 절차대로 진행했습니다.

표 2 — ImageNet 검증 세트의 초해상도 결과입니다.

원본 이미지 행은 검증 세트 원본 이미지로 잰 FID이고, 회귀 행은 논문이 기준선으로 둔 회귀 모델입니다. FM w/ OT는 PSNR과 SSIM이 SR3와 비슷하면서 FID와 IS를 크게 개선했습니다. PSNR과 SSIM은 픽셀 값이 원본과 얼마나 같은지를 재는 지표이고, SR3 논문은 생성 품질을 판단하기에는 FID와 IS가 더 적합하다고 주장합니다.

그림 14(부록, 일부) — FM w/ OT로 검증 세트 이미지를 64×64에서 256×256으로 올린 결과입니다. 원 그림 12장 가운데 4장을 실었습니다.

그림 15(부록, 일부) — 같은 설정의 추가 초해상도 결과입니다. 원 그림 12장 가운데 2장을 실었습니다.

---

10. 부록 자료

우도 추정의 표본 수

NLL은 픽셀 값에 균등 잡음을 더하는 균등 역양자화(uniform dequantization)를 쓰고, 잡음 표본 $K$개의 중요도 가중 평균으로 추정합니다.

표 4(부록) — 서로 다른 $K$로 추정한 시험 세트의 음의 로그 가능도(BPD)입니다.

$K$를 늘릴수록 추정값이 낮아지고, 가장 큰 $K$의 값은 표 1의 NLL과 대부분 같습니다. ImageNet 32의 두 칸만 다릅니다. Score Matching은 표 1에 3.56, 표 4의 $K = 15$에 3.57로, FM w/ Diffusion은 표 1에 3.54, 표 4에 3.56으로 적혀 있습니다. 논문만으로는 어느 값이 맞는지 알 수 없어 두 표의 값을 그대로 옮겼습니다. 어느 쪽이든 ImageNet 32에서 FM w/ OT(3.53)가 가장 낮다는 결론은 같습니다.

부록 E.1은 확산 기준선의 시간 구간 끝을 "학습과 우도 계산에는 $\epsilon = 10^{-5}$, 표본 생성에는 $\epsilon = 10^{-5}$"로 적어 같은 값을 두 번 씁니다. 두 값을 달리 쓰려던 것인지는 논문에 설명이 없습니다.

조건 없는 생성 표본

그림 11 — FM w/ OT로 학습한 CNF가 조건 없이 생성한 ImageNet-32 표본입니다. 고르지 않고 그대로 실은 것입니다.

그림 12 — 같은 방식의 ImageNet-64 표본입니다.

그림 13 — 같은 방식의 ImageNet-128 표본입니다.

부록 그림 번호는 15 다음이 17이고, 그림 16은 논문에 없습니다.

---

11. 결론과 향후 과제

논문에는 한계 절이 따로 없습니다. 결론에서 저자들은 Flow Matching이 확산 모델을 보는 다른 관점을 준다고 정리합니다. 확률적인 확산 과정을 먼저 세우는 대신 확률 경로를 직접 정하면, 표본 생성이 빠르거나 생성 품질이 좋은 경로를 설계할 수 있다는 것입니다. 앞으로는 등방성이 아닌 가우시안이나 더 일반적인 커널처럼 다양한 확률 경로를 쓸 수 있을 것으로 봅니다.

사회적 책임 절에서는 이미지 생성이 해로운 목적에 쓰일 수 있고, 내용을 관리한 학습 세트와 이미지 검증·분류가 이를 줄이는 데 도움이 된다고 적습니다. 또 큰 모델 학습의 에너지 수요가 빠르게 늘고 있어, 더 적은 기울기 갱신과 이미지 처리량으로 학습하는 방법이 시간과 에너지를 크게 줄일 수 있다고 덧붙입니다.

---

12. 정리 — VLA 선행 연구에서의 위치

• Flow Matching 목적 함수 는 확률 경로를 생성하는 벡터장에 신경망을 회귀시켜, 상미분방정식을 풀지 않고 CNF를 학습합니다.
• Conditional Flow Matching 은 데이터 점 하나 기준의 조건부 벡터장에 회귀해도 기울기가 전체 목적 함수와 같다는 정리 2 덕분에, 계산할 수 없던 주변 벡터장 없이 학습합니다.
• 가우시안 조건부 경로 는 평균과 표준편차 두 함수로 경로를 정하고, 확산 모델의 VE·VP 경로를 그 특수한 경우로 포함합니다.
• 최적 수송 경로 는 평균과 표준편차를 직선으로 바꿔 잡음 점과 데이터 점을 선분으로 잇습니다. 정답 속도가 $x_1 - (1 - \sigma_{\min})x_0$로 단순하고, 표본 궤적이 곧아 적은 NFE로 풀어도 오차가 작습니다.
• ImageNet 실험에서 FM w/ OT는 같은 구조의 확산 기반 학습보다 NLL·FID·NFE가 모두 좋았고, ImageNet 128×128에서 조건 없는 생성 FID 20.9를 기록했습니다.

VLA와의 연결

VLA 계보 5편에서 다루는 π0(arxiv 2410.24164)는 이 논문의 방법으로 로봇 행동을 생성합니다. π0 논문에 적힌 설정은 다음과 같습니다.

π0는 이미지 대신 행동 50걸음을 한 번에 생성하고, 선형 경로 위에서 오일러 방법 10걸음으로 행동 묶음 하나를 만듭니다. 행동 묶음 하나에 행동 생성 신경망을 10번 실행한다는 뜻입니다. DDPM 논문은 표본 하나를 만들 때 잡음 제거를 $T = 1000$번 반복합니다. 로봇은 정해진 제어 주기 안에 다음 행동을 내야 하므로, 신경망 실행 횟수는 곧 행동 하나를 내는 데 걸리는 시간과 연결됩니다.

다음 편은 확산 모델로 로봇 행동을 생성한 Diffusion Policy에 대한 이야기입니다. 선행 연구 전체 목차는 VLA 선행 연구 허브에 있습니다.