Flow Matching을 이해하기 위한 노트

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

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

본 기사에서는 Flow Matching의 기본에 대해, 가능한 한 직관적으로 이해할 수 있도록 설명하고 있습니다.

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

본 기사와 GitHub의 전문은 길지만, 제가 수학과 출신이 아니기도 하여 이론 편중이 되지 않도록 노력했습니다.

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

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

특히,

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

을 주요 대상으로 합니다.

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

단, 수학에 대해서는 아래를 전제로 합니다.

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

이미지 분류 모델이라면, 입력 이미지 $x$를 받아 클래스 y를 예측한다. 학습 후에도 이용 시에는 입력 이미지가 필요하다.

식별 모델: 이미지 x -> 클래s y

생성 모델은 학습 데이터와 유사한 새로운 데이터를 만드는 것을 목표로 한다. 생성 시에는 정답 이미지를 입력하지 않고, 난수 (Random Number)로부터 시작한다.

생성 모델: 난수 z -> 새로운 데이터 x_generated

여기서 "유사한 데이터"란, 훈련 이미지를 그대로 반환하는 것이 아니다. 훈련 데이터에 나타나는 형태, 색상, 배치 등의 통계적인 규칙을 배우고, 그 규칙을 따르는 새로운 표본을 만드는 것을 의미한다.

수중에 있는 dataset은 미지의 데이터 생성 규칙으로부터 얻어진 유한개의 관측값이라고 생각한다. 이 배후의 규칙을 데이터 분포 (Data Distribution) $p_{\mathrm{data}}$라고 부른다. 실제로는 $p_{\mathrm{data}}$의 식을 알고 있는 것이 아니라, dataset에서 샘플을 추출할 수 있을 뿐이다.

예를 들어 고양이 이미지 dataset이라면, 고양이 이미지 1장이 1개의 샘플이다. 2D toy라면, 평면 위의 점 1개가 1개의 샘플이다. Flow Matching은 이미지든 점이든 동일한 사고방식을 사용하므로, 처음에는 시각화하기 쉬운 2D 점으로 배운다.

미지의 데이터 분포 p_data
-> dataset으로서 유한개의 샘플을 관측
-> 모델이 분포의 규칙을 학습
...

매번 같은 입력에서 시작하면 결정론적 모델 (Deterministic Model)은 매번 같은 출력을 반환한다. 서로 다른 이미지나 점을 생성하려면 생성할 때마다 다른 시드 (Seed)가 필요하다. 그 시드로서, 간단하게 얼마든지 샘플링할 수 있는 표준 정규 분포 (Standard Normal Distribution)의 난수를 사용한다.

Flow Matching에서는 이 단순한 난수 분포를 source 분포 $p_0$, 만들고자 하는 데이터 분포를 target 분포 $p_1$이라고 부른다. 학습하는 것은 $p_0$의 난수를 $p_1$다운 데이터로 변형하는 규칙이다.

먼저, 수식을 사용하지 않고 학습부터 생성까지 한 바퀴 돌아본다.

처음에 평면 위에 랜덤한 점들을 대량으로 흩뿌려 놓은 "노이즈의 구름"을 상상한다. 한편, 학습 데이터 측에는 고양이 이미지나 8개의 산을 가진 2D 점군(Point Cloud) 같은 "데이터의 모임"이 있다. Flow Matching의 목적은 노이즈의 구름을 데이터의 모임으로 조금씩 변형하는 것이다.

학습 시에는 노이즈 점과 데이터 점을 하나씩 짝지어 준다. 그 두 점 사이에 몇 개의 중간 지점을 두고, 각 중간 지점에서 "데이터 측으로 나아가려면 어느 방향으로 얼마나 움직여야 하는가"라는 화살표를 만든다. 이 화살표가 학습용 정답이다.

뉴럴 네트워크 (Neural Network)에는 중간 지점과 현재 시각을 보여주고, 그 장소의 정답 화살표를 예측하게 한다. 비슷한 장소로 서로 다른 쌍으로부터 화살표가 모이는 경우, 모델은 그것들을 통합한 진행 방향을 배운다. 이렇게 하여 공간의 어디에 있더라도 다음 진행 방향을 답하는 "화살표 지도"가 만들어진다.

생성 시에는 학습할 때 사용했던 데이터 점을 전달하지 않는다. 새로운 노이즈 점만을 배치하고, 학습된 화살표 지도에 현재 위치를 문의하여 조금 움직이고, 다시 한번 문의한다. 이 조작을 반복하면 노이즈 점의 모임이 데이터다운 모임으로 변한다.

학습:
noise 점 + data 점
-> 중간 지점과 정답 화살표를 만듦
...

여기서 배우는 것은 특정 노이즈 점을 특정 이미지로 보내는 대응표가 아니다. 생성 시에는 목적지인 정답 데이터가 없기 때문에, "지금의 시각과 위치라면 분포 전체로서 어느 방향으로 나아가야 하는가"를 답하는 규칙을 배울 필요가 있다. 이것이 나중에 Conditional Flow Matching (CFM)을 고려하게 되는 이유가 된다.

2차원 벡터 (3, -1)

은(는) "가로로 3, 세로로 -1만큼 이동한다"라는 하나의 화살표로 읽을 수 있다. 방향뿐만 아니라 길이도 가지기 때문에, 나아가는 방향과 속도를 동시에 나타낼 수 있다.

위치 x = (2, 4)
속도 v = (3, -1)
조금 이동한 위치 = x + dt * v

여기서 $x$는 점의 현재 위치, $v$는 그 점이 지금 가진 속도이다. 하나의 $v$는 하나의 장소에서의 하나의 속도 벡터 (velocity vector)에 불과하다.

벡터장 (Vector field)은 장소 $x$를 입력하면 그 장소에 대응하는 벡터를 반환하는 함수이다.

장소 x -> 벡터장 v(x) -> 그 장소의 화살표
(0, 0) -> (1, 0)
(1, 0) -> (0, 1)
...

"화살표를 공간 중에 빽빽하게 그린 지도"라고 생각하면 좋다. 다만 구현에서는 모든 지점의 화살표를 저장하는 것이 아니다. 뉴럴 네트워크 (Neural Network)가 질의된 장소에 대해 그때마다 벡터를 계산한다.

Flow Matching의 속도장은 장소뿐만 아니라 시각 $t$도 입력으로 받는다.

v_theta(t, x)

같은 장소 $x$라도 생성의 초반과 종반에서는 나아가야 할 방향이 다를 수 있다. 따라서 $v_ heta(x)$가 아니라, 시각 의존 벡터장 (time-dependent vector field) $v_ heta(t,x)$를 사용한다. 아래첨자 $ heta$는 뉴럴 네트워크의 학습 가능한 파라미터 (learnable parameters)를 통칭한다.

2D toy의 미니 배치 (mini-batch)에서는 입출력 shape은 다음과 같다.

t = torch.rand(batch, 1) # [B, 1]: 각 점의 시각
x = torch.randn(batch, 2) # [B, 2]: 각 점의 현재 위치
v = model(t, x) # [B, 2]: 각 점의 속도 벡터

입력 $x$와 출력 $v$의 shape이 같은 이유는 2차원 점을 움직이려면 2차원 속도가 필요하기 때문이다. 이미지라면 위치 $x$는 이미지 텐서 (image tensor)이며, 속도도 이미지와 같은 shape의 텐서가 된다.

벡터장은 "각 장소에서 어느 방향으로 나아갈지를 답하는 규칙"이다. 궤도 (Trajectory)는 하나의 초기점이 그 규칙에 따라 실제로 지나간 경로이다. 같은 벡터장이라도 초기점이 다르면 궤도는 다르다.

벡터장: 공간 전체의 화살표 규칙
궤도: 하나의 점이 화살표를 차례로 따라간 결과
분포 (Distribution): 다수의 점을 움직였을 때의 모임

시각 $t$의 위치를 $x_t$라고 쓰면, "현재 위치의 화살표를 속도로 하여 움직인다"라는 규칙이 상미분 방정식 (Ordinary Differential Equation, ODE)이다.

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

좌변은 위치가 시간에 대해 얼마나 변하는지를 나타내고, 우변은 모델이 현재의 시각과 장소에서 반환한 속도이다. Flow Matching에서는 뉴럴 네트워크가 우변을 학습하며, 생성 시에 이 규칙을 조금씩 따라간다.

생성 모델에서는 학습 시에만 이용할 수 있는 정보와 생성 시에 이용할 수 있는 정보를 구분해야 한다.

단계	보유하고 있는 것	목적
학습	dataset, 난수, 모델	데이터로 향하는 규칙을 배움
생성	새로운 난수, 학습된 모델	dataset에 없는 새로운 샘플을 만듦

Flow Matching의 학습 시에는 노이즈 점 $x_0$와 데이터 점 $x_1$을 모두 사용하여 교사 속도 (teacher velocity)를 만들 수 있다. 하지만 생성 시에는 $x_1$이 존재하지 않는다. 생성 시에 사용할 수 있는 것은 현재 위치, 시각, 학습된 속도장뿐이다. 이 비대칭성이 CFM을 이해하는 핵심이 된다.

기호	최초의 의미	PyTorch에서의 예
$p_0$	source 분포. 통상적으로 간단한 noise 분포	torch.randn(...)
...	[B,1] 크기의 시각 텐서
$x_t$	시각 $t$의 중간 상태	model에 대한 입력
$u_t$	학습용으로 계산한 교사 속도	regression target
$v_\theta(t,x_t)$	뉴럴 네트워크가 예측하는 속도	model output

처음부터 모든 용어를 암기할 필요는 없다. 제1부의 전반부에서는 위의 8가지 기호가 코드의 어떤 텐서(tensor)에 대응하는지만 따라가면 된다.

이 장에서 다루는 개념: Flow Matching은 source 분포에서 target 분포로 샘플을 운반하는 시간 의존적 속도장 (velocity field)을 뉴럴 네트워크 (neural network)로 학습하는 생성 모델 (generative model)이다. 완성된 데이터를 직접 출력하는 것이 아니라, 각 시각·각 위치에서 "다음에 어느 방향으로 나아갈 것인가"를 나타내는 벡터를 예측하고, 그 속도장을 ODE (상미분 방정식)로서 적분하여 생성한다.

생성 모델을 우선 "난수를 데이터다운 점으로 변형하는 장치"라고 보자.

표준 정규 분포 (standard normal distribution)에서 샘플링한 점 $x_0$는 처음에는 단순한 노이즈 (noise)이다. 목표는 그것을 데이터 분포 (data distribution)에서 나온 것 같은 점 $x_1$으로 이동시키는 것이다. Flow Matching에서는 이 변형을 한 번의 사상 (mapping)으로 직접 맞추는 것이 아니라, 시각 $t$ 마다의 작은 이동으로 나타낸다.

예를 들어 2차원이라면, 각 점에 작은 화살표를 할당한다고 생각하면 된다.

• 위치 $x$에 있다.
• 시각은 $t$이다.
• 그 점은 다음에 어느 방향으로 움직여야 하는가.

이 "다음에 나아가야 할 방향"을 반환하는 함수가 속도장 (velocity field)이다. 뉴럴 네트워크는

v_theta(t, x)

를 학습한다. 여기서 $v_ heta(t, x)$는 시각 $t$, 위치 $x$에 있는 점이 나아가야 할 속도 벡터 (velocity vector)이다.

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

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

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

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

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

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

연속적인 움직임을 컴퓨터 상에서 따라가기 위해, **time grid (시간 그리드)**가 질문할 시각들을 나열하고, **ODE solver (ODE 솔버)**가 모델의 화살표를 사용하여 점을 조금씩 업데이트한다. Euler, Heun, RK4는 이 "한 걸음의 전진 방식"이 서로 다른 solver들이다. 초기 노이즈 준비부터 solver에 의한 이동, 생성 결과 반환까지를 아우르는 절차 전체를 **sampler (샘플러)**라고 부른다.

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

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

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

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

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

Flow Matching의 생성은 다음 상미분 방정식 (ODE)을 푸는 것으로 쓸 수 있다.

\frac{d x_t}{d t} = v_\theta(t, x_t), \quad x_0 \sim p_0.

여기서 $p_0$는 간단한 소스 분포 (source distribution)이다. 전형적으로는 표준 정규 분포를 사용한다. ODE를 $t=0$부터 $t=1$까지 풀면 최종 지점 $x_1$을 얻을 수 있다. 이 $x_1$이 데이터 분포처럼 되어 있다면 생성 모델로서 성공이다.

직관적으로 ODE는 "속도장 (velocity field)이라는 지도를 보면서 점을 조금씩 걷게 하는 규칙"이다. 1스텝만이라면

x_{t+\Delta t} \approx x_t + \Delta t\, v_\theta(t, x_t)

이 된다. 이것이 오일러 방법 (Euler method)이다.

대응하는 코드는 solvers.py의 euler_solve이다.

t = torch.full((x.shape[0], 1), t0 + i * dt, device=x.device, dtype=x.dtype)
x = x + dt * velocity(t, x)

이 두 줄은 위의 수식 그 자체이다. velocity(t, x)가 $v_\theta(t, x_t)$, dt가 $\Delta t$, 업데이트 후의 $x$가 $x_{t+\Delta t}$에 대응한다.

여기서 점의 궤적 $x_t$와 분포 $p_t$를 나누어 생각한다. 그림 1에서 보여주듯이, ODE로 직접 움직이는 것은 하나하나의 점이다. 반면 분포 $p_t$는 대량의 점을 동일한 속도장으로 움직였을 때의 집합으로서 변형된다.

그림 1 하나의 점이 그리는 궤적과, 다수의 점이 만드는 분포의 변형

그림 1에 나타낸 구별은 Flow Matching에서 상당히 중요하다. $x_t$는 하나의 샘플의 위치이고, $p_t$는 시각 $t$에서의 샘플 전체의 분포이다. 학습에서는 미니 배치 (mini-batch) 점을 사용하지만, 이론에서 말하는 "분포가 움직인다"는 것은 모든 점을 동일한 속도장으로 움직였을 때 밀도의 형태가 변하는 것을 의미한다.

생성 시에는 $v_\theta(t, x)$를 사용하여 점을 움직인다. 그렇다면 학습 시에는 정답 속도를 어떻게 만드는가.

가장 단순한 생각은 노이즈 점 $x_0$와 데이터 점 $x_1$을 쌍으로 묶어 그 사이를 직선으로 잇는 것이다. 점이 직선 위를 등속으로 움직인다면 시각 $t$의 위치와 속도는 간단히 쓸 수 있다.

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

u_t = \frac{d x_t}{d t} = x_1 - x_0.

이 $u_t$를 교사 신호 (teacher signal)로 하여, 뉴럴 네트워크의 출력 $v_\theta(t, x_t)$를 가깝게 만든다.

여기서 중요한 것은 $x_t$는 "노이즈 점과 데이터 점 사이의 중간 지점"이며, $u_t$는 "그 중간 지점에서 나아가야 할 직선 방향"이라는 점이다.

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

def sample(self, t, x0, x1):
    return (1.0 - t) * x0 + t * x1
def velocity(self, t, x0, x1):
    ...

이 시점에서 Flow Matching의 최소 스토리는 다음과 같다.

• 노이즈 점 $x_0$와 데이터 점 $x_1$을 준비한다.
• 랜덤한 시각 $t$를 선택한다.
• 중간 지점 $x_t$를 만든다.
• 직선 속도 $u_t$를 만든다.
• $v_\theta(t, x_t)$가 $u_t$에 가까워지도록 학습한다.
• 생성 시에는 v_theta를 사용하여 노이즈를 ODE로 움직인다.

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

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

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

시각 $t=0$에서는 소스 점 $x_0$에 있다. 시각 $t=1$에서는 타겟 점 $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$가 스칼라(scalar)가 아닌 벡터(vector)가 될 뿐이다.

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

수식	코드 변수	shape	의미
$x_0$	x0	[B, D]	소스 분포로부터의 점
$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를 교체했을 때 "중간 지점"과 "정답 속도"가 동시에 변하는 것을 보기 쉽다.

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

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

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

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

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의 시각(time)을 선택하는 분포로 사용하기 용이하다. 두 개의 양수 $\alpha, \beta$에 의해 확률 밀도(probability density)의 형태가 결정된다.

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

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

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

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

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

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

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

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

수식상으로도 time sampler를 $q(t)$로 바꾸면, 근사하는 목적 함수(objective function)는 다음과 같다.

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

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

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

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

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

또한, time sampler는 학습 시에 어느 시각을 볼 것인가를 결정하는 부품이다. 생성 시의 ODE 스텝 수, 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 function)는 타임 샘플러 (time sampler)를 인자로 받는다.

loss = conditional_flow_matching_loss(
model,
path,
...

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

생략. 내용을 읽고 싶으신 분은 GitHub에 올린 본문을 읽어주세요.