Flow Matching 입문 (5) 분포는 어떻게 움직이는가? 연속 방정식·Pushforward·CFM 이해하기

이 기사는 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)을 읽을 수 있음
• 통계 기호 (Statistical Notation)를 읽을 수 있음

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는 화살표 지도를 만드는 학습 측 (training side) 부품이다.
• velocity model은 학습 시에는 화살표를 연습하고, 생성 시에는 화살표를 답하는 공유 부품이다.
• time grid, ODE solver, sampler, NFE는 학습된 지도를 따라가는 생성 측 (inference side) 부품이다.

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

이 장에서 다루는 개념: ODE가 직접 움직이는 것은 하나하나의 점이지만, 다수의 점을 동일한 속도장 (velocity field)으로 움직이면 그 집합체인 분포도 변형된다. 연속 방정식 (continuity equation)은 분포의 국소적인 유입·유출을 나타내며, pushforward는 초기 분포 전체를 flow map으로 옮긴 결과를 나타낸다.

주요 단계: 생성을 뒷받침하는 수학적 보충— solver가 개별 점을 움직인 결과, source 분포 전체가 target 측으로 운반되는 관계를 설명한다.

먼저, 세 가지 대상을 구분한다.

대상	의미	직관
점 $x_t$	시각 $t$에서의 1개 샘플의 위치	군중 속의 한 사람
...	...	...

ODE

\frac{d x_t}{dt}=v_t(x_t)

가 직접 기술하는 것은 1개 점의 궤적이다. 하지만 source 분포 $p_0$에서 수천 개의 점을 추출하여 모두를 동일한 속도장 $v_t$로 움직이면, 점군의 형태도 변한다. 그 시각마다의 점군의 밀도가 $p_t$이다.

즉, "분포를 움직인다"라는 별도의 조작을 추가로 실행하는 것이 아니다. 하나하나의 점을 움직인 결과를 전체로서 바라보면 분포가 움직이는 것처럼 보이는 것이다.

역 통로에 사람들이 서 있는 장면을 생각해보자. 어떤 구간의 사람이 늘어나는 이유는 외부에서 사람이 들어왔거나, 내부의 사람이 나가지 않았기 때문이다. 반대로 사람이 줄어드는 것은 들어오는 인원보다 나가는 인원이 더 많기 때문이다.

확률 분포에서도 마찬가지다. 장소 $x$의 밀도를 $p_t(x)$, 그 장소에서의 속도를 $v_t(x)$라고 하면,

p_t(x)v_t(x)

은 "그 장소를 어느 정도의 확률 질량이 어느 방향으로 통과하고 있는가"를 나타내는 유량 (flux)이다. 유량의 출입과 밀도 변화의 수지를 맞추는 식이 바로 **연속 방정식 (continuity equation)**이다.

\frac{\partial p_t(x)}{\partial t}
+
\nabla\cdot\left(p_t(x)v_t(x)\right)
...

왼쪽의 제1항은 그 장소의 밀도가 시간에 따라 증감하는 속도이다. 제2항의 발산 (divergence) $\nabla\cdot$은 그 장소에서 유량이 차감되어 얼마나 빠져나가는지를 측정한다.

• 유출이 유입보다 많은 장소에서는 밀도 $p_t(x)$가 감소한다.
• 유입이 유출보다 많은 장소에서는 밀도 $p_t(x)$가 증가한다.
• 전체적으로는 확률 질량을 중간에 만들거나 없애지 않고 장소를 옮기고 있는 것이다.

Flow Matching의 최소 구현에서 이 편미분 방정식을 직접 푸는 것은 아니다. 모델이 속도장을 학습하고 solver가 다수의 점을 움직일 때, 그 분포의 변화를 배후에서 기술하고 있는 것이 연속 방정식이다.

시각 0의 점 $x_0$를 속도장을 따라 시각 $t$까지 움직여 도착한 점을 반환하는 사상을 flow map이라고 부른다.

\phi_t(x_0)=x_t.

$\phi_t$는 특정 1개 점만의 궤적이 아니라, "어떤 초기점을 넣어도 해당 시각의 도착점을 반환하는 규칙"이다. 지도상의 각 출발 지점에 대해 시각 $t$의 현재 위치를 적어둔 대응표라고 생각하면 좋다.

코드에서는 학습된 모델을 solver로 $0$부터 $t$까지 적분하는 처리 전체가 $\phi_t$를 수치적으로 계산하고 있는 것이다.

x_t = euler_solve(
lambda time, x: model(time, x),
x0,
...

소스 분포 $p_0$에서 다수의 초기점을 추출하고, 각각에 동일한 flow map $\phi_t$를 적용한다. 얻어진 도착점 전체의 분포가 $p_t$가 된다. 이 관계를, $p_0$가 $\phi_t$에 의해 pushforward (밀어내기) 되었다고 한다.

p_t=(\phi_t)_\#p_0.

기호 $#$는 "분포의 각 점을 사상(mapping)으로 옮긴 결과"를 나타낸다. 식을 암기하기보다 다음의 흐름을 파악하는 것이 중요하다.

x0를 1개 선택 -> flow map으로 xt로 이동
이를 다수의 x0에 대해 수행 -> 도착점들의 집합이 분포 pt가 됨

1차원의 간단한 예로, $x_0$를 평균 0, 분산 1인 표준 정규 분포에서 추출하고, $\phi(x)=2x$를 통해 모든 점을 2배로 만든다고 하자. 0에 가까운 점도 먼 점도 동일한 규칙으로 바깥쪽으로 퍼지기 때문에, 도착점의 분포는 평균 0, 분산 4인 정규 분포가 된다. 점마다의 변환 $x\mapsto2x$를 집단 전체에 적용한 결과, 분포의 폭이 변했다. 이것이 pushforward의 최소 단위 예시다.

지금까지의 관계를 교재의 부품으로 되돌려보자.

학습:
velocity model $v_\theta(t,x)$가 화살표 지도를 근사함
생성:
...

따라서 Flow Matching은 "분포 그 자체를 뉴럴 네트워크(Neural Network)에 입력하여 움직이는" 방법이 아니다. 뉴럴 네트워크는 각 점의 속도를 반환하고, solver가 점을 이동시키며, 그 집단의 집합으로서 분포가 pushforward 된다. 연속 방정식(Continuity Equation)과 pushforward는 이 일련의 생성 과정을 분포 측면에서 다시 읽어내기 위한 언어이다.

이 장에서 다룰 개념: 이상적인 Flow Matching 목적 함수는, 시점마다의 주변 분포(marginal distribution)를 올바르게 운반하는 주변 속도장(marginal velocity field)을 직접 회귀(regression)하는 아이디어이다. 하지만 그 정답 속도장은 통상 데이터로부터 직접 얻을 수 없기 때문에, 계산 가능한 조건부 문제(conditional problem)로 대체해야 한다.

주요 단계: 학습 이론 — 실제로 배우고자 하는 주변 속도장과, 실제로 만들 수 있는 교사(teacher) 속도장의 차이를 정리한다.

제1부의 전반부에서는 노이즈 점 $x_0$와 데이터 점 $x_1$을 쌍(pair)으로 묶어, 직선 경로

x_t = (1 - t)x_0 + t x_1

를 만들었다. 그리고 그 속도

u_t = x_1 - x_0

를 교사 신호(teacher signal)로 삼았다.

이 설명만 보면 "쌍(pair)별 속도를 배우면 끝"인 것처럼 보인다. 하지만 생성 시에는 쌍 $x_0, x_1$을 가지고 있지 않다. 생성 시에 주어진 것은 현재의 점 $x_t$와 시점 $t$뿐이다.

이 지점이 중요하다. 학습 시에는 $x_0$와 $x_1$을 알고 있지만, 생성 시에는 모른다. 따라서 최종적으로 원하는 속도장은

쌍을 알고 있을 때의 속도

가 아니라,

시점 t와 위치 x만을 보았을 때의 평균적인 속도

여야 한다.

시점 $t$에서의 분포를 $p_t$라고 하자. 이 분포를 올바르게 움직이는 속도장을 $u_t(x)$라고 쓴다. 이는 위치 $x$에 있는 확률 질량(probability mass)을 어느 방향으로 운반해야 분포 전체가 올바르게 시간 발전(time evolution)할지를 나타낸다.

이상적으로는 뉴럴 네트워크 $v_\theta(t, x)$를 이 $u_t(x)$에 근사시키고 싶다.

\mathcal{L}_{\mathrm{FM}}(\theta)
=
\mathbb{E}_{t, x \sim p_t}
...

이 식은 아름답지만 구현에서는 곤란하다. 이유는 $u_t(x)$를 직접 알 수 없기 때문이다.

쌍 $x_0, x_1$을 고정하면 직선 속도 $x_1 - x_0$는 즉시 계산할 수 있다. 하지만 특정 위치 $x$에서의 주변 속도 $u_t(x)$는 그 위치에 도달할 수 있는 모든 쌍의 기여를 평균한 것이다. 데이터 분포 전체를 알지 못하면 이를 직접 기술하기 어렵다.

직관적으로, 도로 위의 어느 교차로 $x$를 생각해보자. 그곳에는 다양한 출발지와 목적지를 가진 사람들이 지나간다. 같은 교차로에 있더라도 각자 가고자 하는 방향은 다르다. 교차로만을 보고 결정하는 속도장은 그 장소를 지나는 사람들의 진행 방향을 평균 낸 것과 같은 형태가 된다.

이상적인 FM loss를 그대로 사용하려면 각 $t$와 $x$에 대해 정답 속도 $u_t(x)$가 필요하다. 하지만 실제로 우리가 가진 것은 데이터 샘플과 노이즈 샘플뿐이다.

원하는 것	실제로 만들기 쉬운 것
주변 속도 $u_t(x)$	쌍(pair)별 조건부 속도 $u_t(x\mid x_0,x_1)$
...
Conditional Flow Matching은 이 간극을 메우기 위한 방법이다.

이 장에서 다루는 개념: CFM에서의 "조건부 (conditional)"란, 양 끝점 쌍(endpoint pair)이나 타겟 샘플(target sample)과 같은 잠재적 조건 $z$를 고정했을 때의 경로와 속도를 고려하는 것을 의미한다. 이는 클래스 레이블을 모델에 입력하는 의미에서의 조건부 생성(conditional generation)과는 별개이며, 교사 속도(teacher velocity)를 계산 가능하게 만들기 위한 조건부 설정이다.

주요 단계: 학습 이론 — 쌍별 교사 속도로부터 생성 시에 사용할 수 있는 위치 의존적 속도장(velocity field)을 배우는 개념을 설명한다.

Conditional Flow Matching의 개념은 구현 측면에서 말하자면 다음과 같다.

어려운 "분포 전체의 속도"를 직접 만들지 않는다.
대신, 샘플 쌍별로 간단한 속도를 대량으로 만든다.
뉴럴 네트워크(Neural Network)에는 쌍 정보를 전달하지 않고, 시각(time)과 위치(location)만으로부터 속도를 예측하게 한다.
...

이 "쌍별로 구성된 간단한 세계"가 바로 조건부 경로(conditional path)이다.

쌍을 묶어서

z = (x_0, x_1)

라고 쓴다. 이 $z$를 고정하면, 직선 경로(straight path)는

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

이며, 조건부 속도(conditional velocity)는

u_t(x_t \mid z) = x_1 - x_0

이다.

여기서 "조건부"라고 말하는 것은, $z$를 고정했을 때의 세계에서는 경로와 속도가 간단하게 결정된다는 의미이다.

CFM loss는 다음과 같이 쓸 수 있다.

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

직선 경로의 경우, 이는 다음과 같다.

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

이 식은 4.3절에서 다룬 conditional_flow_matching_loss에 대응한다.

대응하는 코드:

t = time_sampler(batch, x0.device, x0.dtype)
x_t = path.sample(t, x0, x1)
u_t = path.velocity(t, x0, x1)
...

이 코드에서 모델은 x0나 x1을 직접 받지 않는다. 입력은 $t$와 $x_t$뿐이다. 이 제약이 중요하다. 만약 모델에 x0와 x1을 전달해 버리면, 모델은 쌍별 속도를 암기할 수 있게 되어, 생성 시에 사용하고자 하는 "위치만으로 결정되는 속도장"이 되지 않는다.

CFM에서는 뉴럴 네트워크에 "수많은 쌍의 화살표"를 보여주는 한편, 그것이 어떤 쌍으로부터 온 화살표인지는 알려주지 않는다. 모델이 받는 것은 시각 $t$와 위치 $x_t$뿐이다.

따라서 동일한 시각·동일한 위치에 여러 방향의 교사 화살표가 나타날 때, 제곱 오차(squared error)를 최소화하는 예측은 이들의 중간값이 된다. 이것이 "같은 위치로 오는 화살표를 평균한다"는 의미이며, 그림 5에 그 관계를 나타낸다.

그림 5: 동일한 입력으로 모이는 조건부 속도와, 그 평균으로서 얻어지는 예측 속도

그림 5의 검은 점은 동일한 입력 $(t,x)$를 나타낸다. 색이 칠해진 화살표는 서로 다른 조건 $z$로부터 온 교사 속도이다. 모델에 $z$를 전달하지 않을 경우, 모델은 어떤 화살표만을 선택해야 할지 구별할 수 없다. 따라서 제곱 오차의 최적해는 빨간색 화살표, 즉 조건부 평균(conditional mean)이 된다.

실제 학습에서는 자주 나타나는 화살표일수록 미니 배치(mini-batch)에도 많이 나타나므로, 모델의 예측에 강하게 반영된다. 엄밀한 증명으로 나아가지 않더라도, "쌍별 화살표를 모아 위치별 대표 화살표를 배운다"라고 이해하면 이 교재의 구현을 따라가기에 충분하다.

1차원에서, 시각 $t=0.5$, 위치 $x=0$에 동일하게 도달하는 두 개의 쌍을 생각해보자.

pair A: x0 = -1, x1 = 1 -> x_t = 0, velocity = 2
pair B: x0 = 1, x1 = -1 -> x_t = 0, velocity = -2

이때, 동일한 $(t,x)=(0.5,0)$에 대해 교사 속도는 $2$와 $-2$가 모두 나타난다. 모델은 쌍 정보를 보지 않으므로, 제곱 오차를 최소화하는 예측은 평균인 $0$이 된다.

이 예시는 제곱 오차 (squared error)가 경합하는 교사 화살표(teacher arrows)의 중간을 선택하는 것만을 보여주기 위한 작은 회귀 예시이다.

이 예시는 CFM의 중요한 성질을 보여준다.

• 조건부 속도 (conditional velocity)는 쌍(pair) 단위로는 단순하고 날카롭다.
• 주변 속도 (marginal velocity)는 같은 위치로 모이는 조건부 속도들의 평균이 된다.
• 모델에 무엇을 입력으로 전달하느냐가 학습되는 속도장 (velocity field)을 결정한다.

CFM loss 구현 시 특히 실수하기 쉬운 점은 다음과 같다.

주의사항	이유
$t$는 [B, 1]로 설정한다	각 샘플에 하나의 시각을 대응시키고, path 측에서 데이터 차원으로 안전하게 확장하기 위해
모델 입력은 $t, x_t$로 한다	생성 시에도 사용할 수 있는 속도장으로 만들기 위해
$u_t$와 $x_t$의 끝점 컨벤션 (convention)을 맞춘다	$t=0$이 source인지 target인지에 따라 부호가 반대가 되기 때문에
loss는 차원 방향으로 sum한 뒤 배치 평균 (batch average)을 낸다	각 샘플의 벡터 오차를 측정하기 위해

본 교재의 구현에서는 $t=0$이 source, $t=1$이 target이다. 따라서,

u_t = x_1 - x_0

이다. 다른 구현이나 논문에서는 $t=0$을 데이터 측, $t=1$을 노이즈 측에 두기도 한다. 그 경우, 같은 형태의 식으로 보이더라도 생성 방향이나 속도의 부호가 달라진다.

이것으로 이 입문 기사는 종료된다. 지금까지의 내용을 아래에 정리한다.

제1부에서 파악해야 할 대응 관계는 다음과 같다.

• 주변 속도 (Marginal Velocity)
  분포 전체의 평균적인 속도가 필요하다.

  • 수식: $u_t(x)$
  • 코드: 직접 구현하지 않음.

• 조건부 속도 (Conditional Velocity)
  쌍(pair) 단위의 간단한 속도를 만든다.

  • 수식: $u_t(x_t ext{ | } z)$
  • 코드: path.velocity(t, x0, x1)

• 조건부 경로 (Conditional Path)
  쌍(pair) 단위의 중간 지점을 만든다.

  • 수식: $x_t=\psi_t(z)$
  • 코드: path.sample(t, x0, x1)

• 속도장 예측 (Velocity Field Prediction)
  위치만으로 속도를 예측한다.

  • 수식: $v_\theta(t, x_t)$
  • 코드: model(t, x_t)

• CFM Loss
  제곱 오차로 평균 속도를 학습한다.

  • 수식: $\mathbb{E}[\lVert v_\theta-u_t\rVert^2]$
  • 코드: conditional_flow_matching_loss

여기까지 이해한다면, CFM은 '마법의 loss'가 아니라, 조건부로 생성할 수 있는 교사 속도를 사용하여 주변 속도장을 회귀 (regression)로 얻는 방법으로 읽힐 것이다.

제2부 이후의 내용은 GitHub에 올린 전문에 있다. 무료로 읽을 수 있으니 관심 있는 분들은 꼭 읽어보길 바란다.

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