Flow Matching 입문 (2) 확률 경로(probability path)를 이해하기

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

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

현재 인터넷상에 있는 많은 Flow Matching 해설은 고도의 수학적 지식이 필요합니다. 본 기사는 입문용이며, 실습을 목적으로 하기에 그렇지 않은 기사를 만들었습니다.

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

또한, 본 자료의 전문과 코드는 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은 출발점과 목적지의 조합을 만드는 "페어 매칭(pair matching)"이라고 생각하면 됩니다.

쌍이 결정되면, 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)을 만든다 → 연습 경로와 정답 화살표를 만든다 → 화살표 지도를 배운다 → 새로운 노이즈를 지도에 따라 걷게 한다"라는 관계만 파악하면 된다.

이 장에서 다루는 개념: Linear Path는 source 점 $x_0$와 target 점 $x_1$을 선분으로 잇고, 시각 $t$에 따른 중간 지점을 선형 보간 (linear interpolation)으로 결정하는 조건부 경로 (conditional path)이다. 분포 전체를 하나의 선으로 만드는 것이 아니라, 샘플의 각 쌍에 대해 하나씩 직선 경로를 정한다.

주요 단계: 학습 — path, time sampler, 중간 지점, 교사 속도는 velocity model을 위한 연습 문제를 만들기 위해 사용한다.

Linear path는 두 점을 곧게 잇는 가장 간단한 길이다.

시각 $t=0$에서는 source 점 $x_0$에 있다. 시각 $t=1$에서는 target 점 $x_1$에 있다. 중간인 $t=0.25$라면 4분의 1만큼 진행한 위치, $t=0.5$라면 중점, $t=0.75$라면 4분의 3만큼 진행한 위치에 있다.

이 직관을 그대로 식으로 나타내면,

x_t = (1 - t)x_0 + t x_1

이다.

1차원에서 생각해보자. $x_0 = -2$, $x_1 = 6$이라고 하자.

t = 0.00 -> x_t = -2
t = 0.25 -> x_t = 0
t = 0.50 -> x_t = 2
...

속도는 항상

u_t = x_1 - x_0 = 8

이다. 즉, 이 점은 단위 시간당 8만큼 진행하는 등속 운동을 하고 있다.

2차원에서도 마찬가지이며, $x_0$, $x_1$, $x_t$, $u_t$가 스칼라가 아닌 벡터가 될 뿐이다.

미니배치 (mini-batch)로 구현하면 전형적인 shape은 다음과 같다.

수식	코드 변수	shape	의미
$x_0$	x0	[B, D]	source 분포로부터의 점
$x_1$	x1	[B, D]	데이터 분포로부터의 점
$t$	$t$	[B, 1]	각 샘플의 시각
$x_t$	$x_t$	[B, D]	중간 지점
$u_t$	$u_t$	[B, D]	정답 속도

$t$를 [B]가 아니라 [B, 1]로 만드는 이유는 각 샘플에 시각을 하나씩 대응시키기 위해서이다. 2D의 [B, D]에는 이대로 broadcast할 수 있다. 이미지의 [B, C, H, W]에서는 path 내부에서 [B, 1, 1, 1]로 확장할 필요가 있으며, 교재 코드는 데이터의 차원 수에 맞춰 이 확장을 자동으로 수행한다.

losses.py에서는 Linear path에 관한 부분을 다음 두 줄로 읽을 수 있다.

x_t = path.sample(t, x0, x1)
u_t = path.velocity(t, x0, x1)

여기서는 쌍(pair)과 시각으로부터 "모델에게 보여줄 중간 지점"과 "그 지점에서의 정답 화살표"를 만들고 있다. 손실 함수 (loss function) 자체는 '조건부 (conditional)'라는 말을 정의한 후의 4.3절에서 수식, 직관, 코드를 묶어서 설명한다.

코드	수식	설명
path.sample(t, x0, x1)	$x_t = (1-t)x_0 + tx_1$	중간 지점을 만듦
path.velocity(t, x0, x1)	$u_t = x_1 - x_0$	정답 속도를 만듦

이 두 가지를 분리해 두면, path를 교체했을 때 '중간 지점'과 '교사 속도 (teacher velocity)'가 동시에 변하는 것을 더 쉽게 확인할 수 있다.

time sampler는 학습용 시각 $t$를 어떤 확률 분포에서 선택할지를 담당하는 부품이다. 생성 시에 ODE를 진행하는 sampler와는 별개의 것이며, 학습 시에만 사용한다.

CFM loss에는 시각에 관한 기댓값이 포함된다.

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

구현에서는 연속적인 모든 시각에서 loss를 계산할 수 없기 때문에, 각 샘플에 하나씩 $t$를 할당하여 기댓값을 근사한다. $t$의 shape을 [B,1]로 만드는 이유는 2.3절에서 설명한 브로드캐스트 (broadcast) 때문이다.

time_samplers.py에는 비교를 위해 세 가지 분포가 있다.

이름	분포	주로 선택되는 구간	용도
uniform	일양 분포 (uniform distribution)	전 구간을 동일한 빈도로	처음에 사용하는 베이스라인 (baseline)
center	대칭적인 Beta(2,2) 분포	$t=0.5$ 부근	중간 상태를 중점적으로 학습하는 비교
endpoint	대칭적인 Beta(0.5,0.5) 분포	$t=0$ 및 $t=1$ 부근	양 끝단을 중점적으로 학습하는 비교

Beta 분포는 $[0,1]$ 위에서만 정의되기 때문에, Flow Matching의 시각을 선택하는 분포로 사용하기 쉽다. 두 개의 양수 $\alpha, \beta$에 의해 확률 밀도의 형태가 결정된다.

p(t)
=
\frac{1}{B(\alpha,\beta)}
...

$B(\alpha,\beta)$는 밀도 전체의 면적을 1로 만들기 위한 정규화 상수 (normalization constant, 베타 함수)이다. time sampler를 이해하는 단계에서는 이 상수를 직접 계산할 필요는 없다. 중요한 것은 $t$와 1-t의 지수가 분포의 형태를 변화시킨다는 점이다.

여기서의 $\alpha, \beta$는 Beta 분포의 형태를 결정하는 값이며, 제3장에서 사용하는 path 계수 $\alpha(t), \sigma(t)$와는 별개이다.

이 교재 코드에서는 $\alpha=\beta$로 설정한 좌우 대칭 Beta 분포를 사용한다.

• $\alpha=\beta=1$: 식의 $t$에 의존하는 부분이 1이 되어, 일양 분포가 된다.
• $\alpha=\beta>1$: $t=0$과 $t=1$ 부근의 밀도가 낮고, 중앙이 산 모양이 된다. 값을 크게 할수록 $t=0.5$ 부근으로 집중된다.
• $0<\alpha=\beta<1$: 중앙의 밀도가 낮고, 양 끝이 높은 U자형이 된다. 값을 0에 가깝게 할수록 끝점 부근으로 강하게 집중된다.

따라서, $\mathrm{Beta}(2,2)$는 중앙 중시, $\mathrm{Beta}(0.5,0.5)$는 양 끝단 중시가 된다. 만약 $\alpha$와 $\beta$를 서로 다른 값으로 설정하면 source 측과 target 측 중 어느 한쪽으로 편향시킬 수도 있지만, 초기 교재 코드에서는 비교를 단순화하기 위해 좌우 대칭만을 다룬다.

동일한 batch size와 학습 step 수라면 계산할 수 있는 loss의 개수는 변하지 않는다. time sampler를 바꾸는 조작은, 그 한정된 학습 횟수를 시간 구간의 어느 곳에 더 많이 배분할지를 바꾸는 조작이다.

예를 들어 center를 사용하면, 모델은 source와 target이 강하게 섞인 중앙 부근의 상태를 반복해서 보게 된다. 그 구간의 속도 예측은 개선되기 쉬운 반면, 끝점 부근을 보는 횟수는 줄어든다. endpoint에서는 반대로, source에서 움직이기 시작하는 구간과 target에 도착하는 구간을 많이 학습하지만, 중앙 부근의 학습 횟수는 줄어든다.

수식상으로도 time sampler를 $q(t)$로 바꾸면, 근사하는 목적 함수는

\mathcal{L}_{q}
=
\mathbb{E}_{t\sim q(t),x_0,x_1}
...

가 된다. 즉, 단순히 난수 생성 방법을 바꾸는 것뿐만 아니라, 시각별 오차의 가중치 (weighting)를 바꾸고 있는 것이다. 많이 선택되는 시각의 오차는 optimizer에 반복해서 전달되므로, 모델은 그 구간에 더 많은 용량 (capacity)을 사용하게 된다.

이 변경에는 다음과 같은 효과와 부작용이 있다.

sampler	기대 효과	발생 가능한 부작용
uniform	전 구간을 편향 없이 학습	특히 어려운 구간에 충분한 업데이트를 할당하지 못할 수 있음
center	중앙 부근의 복잡한 속도장 (velocity field)을 중점적으로 학습	source 직후·target 직전의 오차가 남을 수 있음
endpoint	궤도의 출발·도착 부근을 중점적으로 학습	중앙 부근의 속도장이 거칠어질 수 있음

어느 구간이 어려운지는 path, coupling, data 분포, 모델에 따라 달라진다. 따라서 center나 endpoint가 항상 uniform보다 우월한 것은 아니다. 시간을 몇 개의 구간으로 나누어 loss를 기록하고, 오차가 집중되어 있는 곳과 생성 궤도가 무너지는 곳을 확인하여 선택한다.

또한, time sampler는 학습 시에 어느 시점(時刻)을 볼 것인가를 바꾸는 부품이다. 생성 시의 ODE step 수, solver, 평가 시점의 나열 방식은 바꾸지 않는다. 학습용 time sampler와 생성용 schedule을 혼동하지 않는 것이 중요하다.

from fm_minimal import get_time_sampler
time_sampler = get_time_sampler("uniform")
t = time_sampler(batch_size, x0.device, x0.dtype) # [B, 1]

loss 함수는 time sampler를 인자로 받는다.

loss = conditional_flow_matching_loss(
model,
path,
...

첫 학습에서는 uniform을 사용한다. time sampler를 바꾸는 실험에서는 data, coupling, path, model, optimizer를 동일하게 유지하고, 변경점을 시간 분포(time distribution)로만 한정한다. 또한, 학습 loss만으로 우열을 가리지 말고, 동일한 NFE에서 생성 품질도 비교한다. 시간 분포가 다르면 "어느 시점의 오차를 무겁게 본 loss인가"도 달라지기 때문에, loss 값의 단순 비교만으로는 판단할 수 없다.

Flow Matching의 기본 이해에는 불필요하지만, 비교 대상인 다른 Path로서 작성하였습니다.

이 장에서 다루는 개념: Gaussian path는 데이터 성분과 Gaussian noise 성분을 시간 의존 계수 $\alpha(t)$, $\sigma(t)$로 혼합하여 중간 분포를 만드는 확률 경로(probability path)이다. 그 속도는 혼합 계수를 시간 미분하여 얻을 수 있으며, Linear Path와는 중간 분포와 교사 속도(teacher velocity)가 모두 다르다.

주요 단계: 학습 — Linear Path와 마찬가지로, 중간 지점 $x_t$와 교사 속도 $u_t$를 만드는 path의 선택지로서 사용한다.

Linear path에서는 중간 지점이

x_t = (1-t)x_0 + t x_1

이며, 속도는 일정했다.

Gaussian path에서는 노이즈와 데이터를 섞는 계수가 $t$에 따라 비선형적으로 변한다.

x_t = \alpha(t)x_1 + \sigma(t)x_0.

이때 속도는 계수를 미분하면 된다.

u_t
=
\frac{d x_t}{dt}
...

직관적으로 $\alpha(t)$는 데이터 성분의 볼륨, $\sigma(t)$는 노이즈 성분의 볼륨이다. 시간이 흐름에 따라 데이터 성분의 볼륨을 높이고, 노이즈 성분의 볼륨을 낮춘다. 그 변화율이 속도가 된다.

이 교재의 최소 코드에서는 예시로 다음 스케줄을 사용한다.

\alpha(t) = \sin\left(\frac{\pi t}{2}\right),
\quad
\sigma(t) = \cos\left(\frac{\pi t}{2}\right).

이때

$$
\alpha(t)^2 + \sigma(t)^2 = 1
$$

이므로, 분산을 유지하는 형태의 Gaussian path로서 직관적으로 다루기 쉽다.

미분은

\dot{\alpha}(t)
=
\frac{\pi}{2}\cos\left(\frac{\pi t}{2}\right),
...

따라서 속도는

u_t
=
\frac{\pi}{2}
...

이에 대응하는 코드는 paths.py의 TrigGaussianPath이다.

def sample(self, t, x0, x1):
    return self.alpha(t) * x1 + self.sigma(t) * x0
def velocity(self, t, x0, x1):
    ...

이와 같이, 경로(path)를 바꾸더라도 CFM loss의 외부 구조는 변하지 않는다. 변하는 것은 $x_t$와 $u_t$를 생성하는 방식이다.

항목	Linear path	Gaussian path
중간점	$(1-t)x_0+tx_1$	$\alpha(t)x_1+\sigma(t)x_0$
속도 (velocity)	x1 - x0	$\dot{\alpha}(t)x_1 + \dot{\sigma}(t)x_0$
계수의 변화	일정	스케줄 (schedule) 의존
Diffusion과의 대응	직관적인 직선 보간	노이즈 혼합식과 대응하기 쉬움

Linear path는 Flow Matching을 처음 이해할 때 가장 좋다.

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