이 기사는 Flow Matching 및 관련 생성 모델(Generative Model)에 대해 배우기 위해 작성한 개인적인 노트를 정리한 것입니다.

내용은 논문이나 공개 자료를 바탕으로 Codex를 활용하며 정리하고 있지만, 내용과 구성의 확인 및 정리는 저자 본인이 직접 수행하고 있습니다.

현재 인터넷상에 있는 많은 Flow Matching 해설은 고도의 수학적 지식을 필요로 합니다. 본 기사는 입문을 위해 이론에 치중하지 않습니다.

수학이 전혀 필요 없다는 뜻은 아니지만, 편미분 방정식(Partial Differential Equation)이나 학습 이론에 대해 깊게 파고드는 것이 아니라, 대략적으로 이해하여 코드를 읽고 쓸 수 있도록 하는 것을 목적으로 만들었습니다.

또한, 본 자료의 전문과 코드는 GitHub에 올려두었습니다. Rectified Flow에 대해서는 본 기사에서 다루지 않지만, 전문에는 포함되어 있습니다.

GitHub에 올린 전체 내용 중 제0부와 제1부를 몇 회에 걸쳐 나누어 기사로 작성할 예정입니다. 지난 기사는 다음과 같습니다.

오류나 개선점 등이 있다면 댓글로 지적해 주시면 감사하겠습니다.

본 자료는 Flow Matching 및 관련 생성 모델을 배우고 싶은 분들을 대상으로 합니다.

특히,

• 딥러닝(Deep Learning)의 기초를 한 차례 학습하신 분
• 논문을 읽거나 PyTorch로 구현을 하기 위한 발판이 되는 해설을 찾고 계신 분
• 수식뿐만 아니라 직관적인 이해도 얻고 싶은 분

을 주요 대상으로 합니다.

엄밀한 이론보다는 "왜 그 식이 되는가", "어떤 의미를 갖는가"를 중시하여 설명하고 있습니다.

단, 수학에 대해서는 다음을 전제로 합니다.

• 편미분(Partial Derivative)에 대해 이해하고 있음
• 미분 방정식(Differential Equation)을 읽을 수 있음
• 통계 기호를 읽을 수 있음

Flow Matching의 전체상을 이해하기 위해 재게시합니다. 이것만으로 전체적인 역할 분담을 알 수 있으므로, 아직 읽지 않은 분들은 꼭 읽어보시기 바랍니다.

Flow Matching에는 "화살표 지도를 만드는 학습"과 "그 지도를 사용하여 새로운 데이터를 만드는 생성"이라는 두 가지 흐름이 있습니다. 이 두 가지를 나누어 생각하면 전체상을 파악하기 쉽습니다.

학습은 간단한 난수의 집합인 source 분포 $p_0$와 배우고자 하는 데이터의 집합인 target 분포 $p_1$에서 시작합니다. source에서 노이즈 점 $x_0$를, target에서 데이터 점 $x_1$을 추출하며, 이 두 점을 어떻게 한 쌍으로 묶을지를 coupling이 결정합니다. coupling은 출발점과 목적지의 조합을 만드는 "페어 매칭(Pairing)"
이라고 생각하면 됩니다.

페어가 결정되면, time sampler가 연습할 시각 $t$를 선택하고, path (경로) 가 해당 시각의 중간 지점 $x_t$를 놓습니다. 동일한 path로부터 "그 중간 지점에서는 어느 방향으로 나아가야 하는가"라는 교사 속도 (Teacher Velocity) $u_t$도 얻을 수 있습니다. path는 연습용 길을 그리는 부품이며, 교사 속도는 그 길을 따라가는 정답 화살표입니다.

시각 $t$와 중간 지점 $x_t$를 입력받아 화살표를 예측하는 뉴럴 네트워크가 velocity model (속도장 모델) $v_ heta(t,x_t)$입니다. loss (손실 함수) 는 예측 화살표와 교사 속도의 차이를 채점하며, optimizer 는 그 점수가 작아지도록 모델을 조정합니다. 이 연습을 반복하여 얻은 "화살표 지도"를 가중치 및 설정과 함께 저장한 것이 checkpoint 입니다.

학습:
source x0 + target x1
-> coupling: 출발점과 목적지를 페어로 묶음
...

생성 시에는 학습에 사용했던 target 데이터, coupling, 교사 속도를 사용하지 않습니다. source 분포에서 새로운 노이즈 점을 취하고, checkpoint에서 복원한 velocity model에게 "지금 시각과 위치에서는 어느 방향으로 가야 하는가"라고 문의합니다.

연속적인 움직임을 컴퓨터 상에서 따라가기 위해, time grid가 문의할 시각들을 나열하고, ODE solver가 모델의 화살표를 사용하여 점을 조금씩 업데이트합니다. Euler, Heun, RK4는 이 "한 걸음의 진행 방식"이 서로 다른 solver입니다. 초기 노이즈 준비부터 solver를 통한 이동, 생성 결과 반환까지의 절차 전체를 sampler라고 부릅니다.

생성:
새로운 source noise x0
-> checkpoint에서 velocity model을 불러옴
...

모델에 화살표를 문의한 횟수를 **NFE (Number of Function Evaluations)**라고 부른다. NFE가 적으면 생성 속도가 빨라지기 쉽지만, 한 걸음(step)이 너무 거칠면 올바른 경로에서 벗어난다. evaluation (평가) 단계에서는 생성 결과가 target 분포에 가까워졌는지뿐만 아니라, 동일한 NFE에서 어떤 solver가 더 좋았는지도 조사한다.

여기서 가장 중요한 관계는 다음 세 가지이다.

• coupling, time sampler, path, 교사 속도(teacher velocity), loss, optimizer는 화살표 지도를 만드는 학습 측 부품이다.
• velocity model은 학습 시에 화살표를 연습하고, 생성 시에 화살표를 답하는 공유 부품이다.
• time grid, ODE solver, sampler, NFE는 학습된 지도를 따라가는 생성 측 부품이다.

다음 장에서는 이것들을 하나씩 수식과 코드로 옮겨갈 것이다. 지금은 용어를 암기할 필요는 없다. "쌍(pair)을 만든다 → 연습 경로와 정답 화살표를 만든다 → 화살표 지도를 배운다 → 새로운 노이즈를 지도에 따라 걷게 한다"라는 관계만 파악하면 된다.

이 장에서 다루는 개념: Conditional Flow Matching (CFM)은 끝점 쌍(endpoint pair) 등의 조건을 고정하면 계산할 수 있는 조건부 속도를 교사 신호(teacher signal)로서 회귀(regression)하여, 생성 시에 사용할 수 있는 주변 속도장(marginal velocity field)을 배우는 방법이다. 모델에는 조건 그 자체를 전달하지 않고 $(t, x_t)$만을 전달하기 때문에, 다수의 조건부 속도를 적절히 평균한 속도장이 학습된다.

주요 단계: 학습 — CFM loss를 통해 예측 속도와 교사 속도를 비교하여, 생성 시에 사용할 velocity model을 만든다.

사실은 시각 $t$의 각 지점 $x$에 대해 "데이터 분포 전체를 올바르게 움직이는 평균적인 속도"를 알고 싶다. 이를 주변 속도장(marginal velocity field)이라 부른다.

하지만 데이터 분포 전체의 올바른 속도장을 직접 계산하는 것은 어렵다. 그래서 더 간단한 문제로 분해한다.

• 먼저, source 지점 $x_0$와 target 지점 $x_1$의 쌍을 만든다.
• 그 쌍만을 보면 직선 경로와 속도를 즉시 알 수 있다.
• 많은 쌍에 대해 이를 반복한다.
• 뉴럴 네트워크는 유사한 장소와 시각에 도달하는 많은 교사 속도를 평균하도록 학습한다.

이것이 Conditional Flow Matching의 기본적인 핵심이다.

여기서 말하는 "조건"은 클래스 레이블이나 텍텍스트를 생성 모델에 입력하는 조건부 생성(conditional generation)의 조건이 아니다. 학습용으로 선택한 source 지점과 target 지점의 쌍

z=(x_0,x_1)

을 가리킨다. "$z$가 이 쌍이다"라고 고정한 상태에서 생각하기 때문에 조건부 경로(conditional path)라고 부른다. 쌍을 고정하면 어디에서 어디로 진행할지가 결정되므로, 중간 지점과 정답 속도를 쉽게 계산할 수 있다.

예를 들어, 도쿄에서 오사카로 향하는 경로와 삿포로에서 후쿠오카로 향하는 경로는 별개이다. "출발지와 목적지를 고정한 세계"가 조건부 세계에 대응한다. 학습에서는 수많은 세계를 만들지만, 뉴럴 네트워크에는 쌍 $z$ 자체를 전달하지 않는다.

일반적으로 쌍을 고정한 조건부 경로를 $\psi_t(x_0,x_1)$라고 쓰면, 중간 지점과 조건부 속도는 다음과 같다.

x_t = \psi_t(x_0,x_1)

u_t(x_t \mid z)
=
\frac{\partial}{\partial t}\psi_t(x_0,x_1)

이다. 제2식은 제1식에서 결정된 중간 지점이 시간 $t$에 대해 얼마나 변화하는지, 즉 경로의 접선 방향을 나타낸다. path.sample이 제1식을, path.velocity가 제2식을 담당한다.

Linear path에서는 경로 함수를 다음과 같이 선택한다.

x_t = \psi_t(x_0, x_1) = (1 - t)x_0 + t x_1.

이를 $t$로 미분하면 조건부 속도는 다음과 같다.

u_t(x_t \mid z) = x_1 - x_0

조건부 경로는 쌍을 고정하면 간단하다. 하지만 생성 시에는 쌍 $z$가 보이지 않는다. 생성 시에 사용할 수 있는 것은 시각 $t$와 현재 위치 $x_t$뿐이다. 따라서 뉴럴 네트워크는

v_\theta(t, x_t)

로서, $z$를 직접 받지 않는 형태로 학습한다.

다수의 조건부 경로를 섞어, 쌍(pair) $z$를 평균하여 보이지 않게 만드는 조작을 **주변화 (marginalization)**라고 한다. 주변화한 후에 남는 시각 $t$의 분포 $p_t$가 주변 분포 (marginal distribution)이며, 시각별로 나열한 $p_t$ 전체가 확률 경로 (probability path)이다.

Conditional Flow Matching의 최소 형태는 다음과 같다.

\mathcal{L}_{\mathrm{CFM}}(\theta)
=
\mathbb{E}_{t, x_0, x_1}
...

단,

x_t = (1 - t)x_0 + t x_1.

이 식은 속도(velocity)에 대한 지도 학습 회귀 (supervised regression)이다. 입력은 "현재의 시각과 중간 지점" $(t, x_t)$이고, 정답(teacher)은 "그 쌍이라면 나아가야 할 화살표" $x_1 - x_0$이다. 예측한 화살표의 방향이나 길이가 정답에서 벗어날수록 손실 (loss)이 커지며, 일치하면 0이 된다.

직관적으로는, 한 번의 학습으로 다음 과정을 수행한다.

• 소스 (source) 점과 타겟 (target) 점을 한 쌍 선택한다.
• 그 사이의 시각을 하나 선택하여 중간 지점 $x_t$를 만든다.
• 그 쌍을 잇는 정답 화살표 $u_t$를 만든다.
• 뉴럴 네트워크에 $(t, x_t)$만을 보여주고 화살표를 예측하게 한다.
• 예측 화살표와 정답 화살표의 차이를 제곱하여 미니 배치 (mini-batch)로 평균한다.

제곱을 하는 이유는 정(+)과 부(-)의 오차가 서로 상쇄되지 않게 하고, 크게 벗어난 예측에 강한 벌점 (penalty)을 주기 위해서이다. 벡터의 각 성분을 제곱하여 더하기 때문에, 손실은 "예측 화살표의 끝점과 정답 화살표의 끝점 사이 거리의 제곱"으로 읽을 수 있다.

기호 $\mathbb{E}$는 기댓값 (expectation), 즉 모든 시각과 샘플에 대한 평균을 나타낸다. 구현 시에는 무한개의 쌍을 평균할 수 없으므로, 각 스텝 (step)에서 랜덤한 $t, x_0, x_1$을 미니 배치로 가져와 그 배치 평균으로 기댓값을 근사한다. 따라서 위 수식의 기댓값은 코드 중 .mean()에 대응한다.

이 학습을 다수의 쌍에 대해 반복한다. 직관적으로 CFM은 "쌍마다의 간단한 화살표"를 대량으로 모으고, 같은 장소로 오는 화살표들을 평균하여 분포 전체를 움직이는 화살표 지도를 만들고 있는 것이다.

구현에서는 losses.py의 conditional_flow_matching_loss가 그대로 대응한다.

batch = x0.shape[0]
t = time_sampler(batch, x0.device, x0.dtype)
x_t = path.sample(t, x0, x1)
...

여기까지는 소스에서 가져온 $x_0$와 타겟에서 가져온 $x_1$을 한 쌍으로 취급했다. 하지만 어떤 $x_0$와 어떤 $x_1$을 잇는지는 유일하지 않다. 이 쌍의 결합 분포 (joint distribution)를 **커플링 (coupling)**이라고 부른다.

최소 구현에서는 소스와 타겟으로부터 독립적으로 샘플링한다.

q(x_0, x_1) = p_0(x_0)p_1(x_1).

이는 "특별한 대응 관계를 정하지 않고, 우연히 같은 미니 배치에 들어온 점들을 잇는다"는 기준이다. 구현에서는 타겟 측을 랜덤하게 섞어도 동일한 역할을 한다.

x0, x1, perm = random_coupling(x0, x1)

커플링을 바꾸어도 4.3절의 CFM loss는 변하지 않는다. 변하는 것은 loss로 전달되는 양 끝점 쌍과, 그로부터 만들어지는 $x_t, u_t$이다. 가까운 점들끼리 잇기 쉽게 만드는 최적 운송 CFM (Optimal Transport CFM, OT-CFM)은 이 쌍 만들기 과정을 개선하는 방법이며, 제5부에서는 최적화 문제와 미니 배치 근사에 집중하여 설명한다.

이 장에서 다루는 개념: 샘플러 (sampler)는 초기 노이즈 준비부터 생성 결과 반환까지를 담당하는 절차 전체이며, 솔버 (solver)는 그 내부에서 ODE의 한 단계를 계산하는 수치 적분법 (numerical integration method)이다. 본 장에서는 Euler solver를 사용하는 샘플러를 최소 예시로 하여 생성을 배운다.

주요 단계: 생성(Generation) — 체크포인트 (checkpoint)에서 학습된 속도 모델 (velocity model)을 불러와, 솔버로 소스 노이즈를 생성 샘플로 운반한다.

학습이 끝나면, $v_ heta(t, x)$는 "시각 $t$에 위치 $x$에 있는 점이 다음에 나아가야 할 방향"을 반환한다. 생성(generation)에서는 표준 정규 분포 (standard normal distribution)로부터 점을 추출하여, $t=0$부터 $t=1$까지 조금씩 진행시킨다.

이 생성 전체를 관리하는 것이 sampler이다. sampler는 출발점을 준비하고, 시간 격자 (time grid)를 나열하며, solver를 호출하여 마지막 점을 생성 결과로 반환한다.

sampler
├─ source 분포로부터 초기 노이즈 $x_0$를 생성
├─ time grid를 결정
...

반면, ODE solver가 담당하는 것은 sampler 내부의 "현재 위치에서 다음 위치로 어떻게 나아갈 것인가"뿐이다. Euler, Heun, RK4는 solver의 종류이며, sampler 전체의 종류가 아니다.

다만 실제 논문이나 라이브러리에서는 "Euler solver를 사용하는 sampler"를 짧게 Euler sampler라고 부르기도 한다. 이 교재에서도 타 자료와의 대응을 위해 이 표현이 등장하지만, 내부 관계는 항상 다음과 같다.

Euler sampler = sampler 전체 + 내부에서 사용하는 Euler solver

구분 방법은 단순하다. "초기 노이즈부터 최종 샘플까지의 생성 전체"를 묻고 있다면 sampler이고, "ODE의 다음 한 걸음을 어떻게 계산할 것인가"를 묻고 있다면 solver이다.

ODE

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

를 $N$ 스텝으로 근사한다. $\Delta t=1/N$이라고 하면,

x_{k+1}
=
x_k + \Delta t\, v_\theta(t_k, x_k),
...

이에 대응하는 구현은 solvers.py의 euler_solve이다.

euler_solve는 초기점 x0를 외부에서 받아 ODE 업데이트만을 수행하므로 solver이다. source 분포로부터 초기 노이즈를 만들고, 이 함수를 호출하며, 반환값을 생성 샘플로 취급하는 외부 프로세스까지 포함하여 sampler가 된다.

train_minimal_2d.py는 2D toy 분포에서 다음을 수행한다.

• source로서 표준 정규 분포로부터 x0를 샘플링한다.
• target로서 8개 성분의 Gaussian mixture로부터 x1을 샘플링한다. 이는 원주 위의 8개 중심 중 하나를 선택하고, 그 근처에 작은 noise를 더한 점이다.
• Linear path로 $x_t$와 $u_t$를 만든다.
• MLPVelocity를 CFM loss로 학습한다.
• 학습 후, sampler가 새로운 표준 정규 노이즈를 준비하고, 내부의 Euler solver로 구동한다.

실행 예시를 보여준다.

python .\src\train_minimal_2d.py `
--steps 2000 --batch 512 --device cpu

저장되는 minimal_2d.pt에는 학습된 가중치 (weights), 모델 구성, path, coupling, time sampler, seed, 그리고 확인용 생성 샘플이 들어있다.

이 절에서 다루는 개념: NFE (Number of Function Evaluations)는 1회의 생성에서 속도장 모델 (velocity field model)을 평가한 횟수이며, 뉴럴 생성의 주요 계산 비용 지표이다. 동일한 step 수라도 solver에 따라 NFE가 다르기 때문에, 품질과 속도를 공정하게 비교하기 위한 공통 예산으로 사용한다.

• NFE가 낮음: 모델 호출이 적어 통상적으로 생성이 빠르고 계산량도 작다. 다만, 시간 적분 (time integration)이 거칠어져 생성 품질이 떨어질 가능성이 있다.
• NFE가 높음: 모델 호출이 많아 통상적으로 생성이 느리고 계산량도 크다. ODE의 이산화 오차 (discretization error)는 작아지기 쉽지만, 학습된 속도장 자체의 오차는 교정되지 않는다.

따라서 목표는 단순히 NFE를 최소화하는 것이 아니다. "동일한 품질을 더 낮은 NFE로 달성하는 것" 또는 "동일한 NFE로 더 높은 품질을 얻는 것"이다. 이어지는 5.5절에서는 solver별 NFE를 구체화하고, 5.6절에서는 NFE를 sampler, solver, time grid의 관계 속에 위치시킨다.

ODE solver는 학습된 velocity model이 반환하는 연속 시간의 속도를 사용하여, 현재 위치에서 다음 위치를 수치적으로 결정하는 부품이다. 직관적으로 velocity model이 "그 장소의 화살표를 알려주는 지도"라면, solver는 "화살표를 보며 실제로 걷는 방법"에 대응한다.

Euler, Heun, RK4는 연속 시간 ODE (Ordinary Differential Equation)의 해를 유한한 횟수의 속도 평가로 근사하는 **명시적 수치 적분법 (Explicit Numerical Integration Method)**이다. Euler는 1 step의 시작점만, Heun은 시작점과 가상의 종료점, RK4는 구간 내 4개 지점에서 속도를 조사한다. 고차 solver일수록 1 step을 정밀하게 계산하지만, 모델을 호출하는 횟수도 늘어난다. 따라서 solver는 step 수뿐만 아니라, 모델 평가 횟수인 NFE (Number of Function Evaluations)와 이산화 오차 (Discretization Error)를 함께 비교해야 한다.

여기서, 학습 시의 path와 생성 시의 solver를 구분해야 한다.

학습: path를 통해 중간 지점과 교사 속도 (Teacher Velocity)를 생성
-> velocity model이 속도장 (Velocity Field)을 학습
생성: solver가 학습된 속도장을 유한한 횟수만큼 평가하며 궤적을 따라 이동

생성 시에는 타겟 지점 $x_1$도 교사 속도 $u_t$도 없기 때문에, solver는 model(t, x)를 반복적으로 호출하며 스스로 궤적을 따라간다. 같은 체크포인트 (Checkpoint)를 사용하더라도, solver나 step 수를 바꾸면 생성되는 종점이 달라질 수 있다.

Euler법은 현재 위치에서 속도를 1회만 평가하고, 그 방향으로 나아간다.

x_{k+1}=x_k+\Delta t\,v_\theta(t_k,x_k).

1 step당 NFE는 1이며, 단순하고 계산량이 적다. 반면, 궤적이 크게 휘어지는 구간에서는 출발점의 화살표만을 사용하기 때문에 곡선 구간을 지나쳐버리기 쉽다.

Heun법은 먼저 Euler법으로 가상의 다음 지점을 만든 뒤, 그 지점에서도 속도를 평가한다. 마지막으로 현재 지점과 가상 지점의 두 속도를 평균하여 업데이트한다.

k_1=v_\theta(t_k,x_k),\n\qquad\hat{x}_{k+1}=x_k+\Delta t\,k_1,

k_2=v_\theta(t_{k+1},\hat{x}_{k+1}),\qquad x_{k+1}=x_k+\frac{\Delta t}{2}(k_1+k_2).

"현재의 화살표만으로 결정하지 않고, 조금 앞의 화살표도 확인하여 진행 방향을 수정하는" 방식이다. 1 step당 NFE는 2가 된다.

RK4는 구간의 시작점, 중간 부근, 종료점 부근에서 총 4회 속도를 평가하고, 가중 평균을 통해 나아간다.

\begin{aligned}
k_1 &= v_\theta(t_k,x_k),\nk_2 &= v_\theta\left(t_k+\frac{\Delta t}{2},x_k+\frac{\Delta t}{2}k_1\right),\n... 
\end{aligned}

하나의 구간을 상세히 조사하기 때문에 적은 step으로도 적분 오차를 억제하기 쉽지만, 1 step당 NFE는 4이다.

교재 코드에서는 세 가지 solver가 동일한 역할을 하는 함수로 나열되어 있다.

velocity = lambda t, x: model(t, x)
samples_euler = euler_solve(velocity, x0, steps=32)
samples_heun = heun_solve(velocity, x0, steps=16)
...

이 코드 예시에서 Euler는 $32\times1$, Heun은 $16\times2$, RK4는 $8\times4$이므로, 모두 총 NFE가 32로 동일하다. 공정한 비교를 위해서는 이와 같이 모델을 호출하는 총 횟수를 맞추고, 생성 품질과 실행 시간을 측정해야 한다.

solver를 고차로 만든다고 해서 학습된 속도장 자체가 틀렸다면 해결되지 않는다. solver가 줄일 수 있는 것은 주로 "유한한 step으로 ODE를 따라갈 때 발생하는 오차"이며, 이는 모델의 속도 예측 오차와는 별개이다. 처음에는 Euler로 구현을 확인한 뒤, 이후 동일한 NFE에서 Heun이나 RK4로 교체하여 개선되는지 살펴보면 원인을 구분하기 쉽다.

5.1절의 포함 관계를 생성 시의 역할로서 정리한다.

sampler는 외부 박스이며, solver와 time grid는 박스 내부의 설정이다. NFE는 해당 설정으로 생성한 결과로서 계산되는 계산량이다. solver만으로는 초기 노이즈를 만들지 않으며, 생성 결과에 대한 평가도 수행하지 않는다.

Euler는 현재 지점의 속도를 1회 평가하여 진행한다. Heun은 Euler 방식으로 가상의 다음 지점을 만든 후, 현재 지점과 가상 지점의 속도를 평균한다. RK4는 1개 구간(step)에서 4회 속도를 평가한다. 이러한 차이로 인해, 동일한 step 수라 하더라도 계산량은 같지 않다.

이하에는 본 교재의 내용, 설명 순서, 구현 대응, 용어 선택 검토에 사용된 논문, 공식 구현(official implementation), 기술 문서, 해설 자료, 영상, 기존 교재를 게시한다.