Flow Matching 입문 (4) ― PyTorch를 이용한 Flow Matching 학습 흐름

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

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

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

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

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

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

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

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

특히,

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

을 주요 대상으로 합니다.

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

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

편미분에 대해 이해하고 있음
미분 방정식(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는 화살표 지도를 만드는 학습 측 (training side) 부품이다.
velocity model은 학습 시에는 화살표를 연습하고, 생성 시에는 화살표를 답하는 **공유 부품 (shared component)**이다.
time grid, ODE solver, sampler, NFE는 학습된 지도를 따라가는 생성 측 (generation side) 부품이다.

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

이 장에서 다루는 개념:
2D toy 실험은 저차원의 기지(known) 분포를 사용하여 경로, 속도장 (velocity field), 생성 궤도 (generation trajectory), mode coverage를 직접 관찰하는 검증 환경이다. 고품질의 최종 모델을 만들기 위함이 아니라, 수식과 구현의 대응 및 실패 원인을 눈으로 확인하기 위해 사용한다.

주요 단계: 학습 → 생성 → 평가 — 속도장을 학습하고, Euler sampler로 생성하며, 분포와 궤도를 시각화하여 확인한다.

최소 구현(minimal implementation)을 실행했을 때, 가장 먼저 봐야 할 것은 loss만이 아니다. Flow Matching은 "속도장을 배우는" 방법이므로, 다음 세 가지를 동시에 확인해야 한다.

source 분포에서 나온 점이 생성 후에 target 분포로 가까워지고 있는가.
입자의 궤도가 도중에 너무 크게 휘거나 발산하지 않는가.
학습한 속도장이 직선 경로의 교사 신호 (teacher signal)와 모순되지 않는 방향을 향하고 있는가.

2D toy 실험에서는 target을 8개의 작은 Gaussian으로 이루어진 Gaussian 혼합 분포 (Gaussian mixture)로 설정한다. 이는 이미지 생성보다 훨씬 작지만, Flow Matching이 "노이즈의 구름을 데이터의 형태로 운반하는" 모습을 관찰하기에는 충분하다.

source, target, 쌍(pair)별 직선 path의 대응 관계를 그림 2에 나타낸다.

그림 2 2D toy에서의 source, target, 쌍별 Linear path

그림 2는 학습된 모델의 출력이 아니라, 학습에 사용하는 교사 신호를 만드는 방법을 나타낸다. 파란색 점은 source 샘플 $x_0$, 빨간색 점은 target 샘플 $x_1$, 회색 선은 쌍별 linear path, 녹색 점은 중간 시점의 $x_t$이다. 실제 학습된 모델이 만드는 궤도는 이 교사 path의 평균적인 속도장을 ODE로서 적분한 것이 된다.

train_minimal_2d.py는 다음 처리를 수행한다.

source x0 ~ N(0, I)
target x1 ~ eight Gaussians
t ~ Uniform(0, 1)
...

위의 처리는 기본값인 linear path, independent coupling, uniform time sampler의 경우이다. 이 스크립트는 학습 후에 다음을 저장한다.

model: 학습된 MLPVelocity의 가중치.
model_config: MLP의 너비, 깊이, 시간 임베딩 (time embedding), 활성화 함수.
samples: Euler sampler로 생성한 점.
path, coupling, time_sampler, seed: 실험 조건.

실행 예시:

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

plot_minimal_2d.py는 저장된 모델을 불러와 source, target, generated samples, trajectory를 한 장의 그림으로 만든다.

PyTorch 환경에서 학습과 시각화를 실행하면 그림 3과 같은 완전한 확인도를 얻을 수 있다.

그림 3 2D toy의 source 분포, target 분포, 생성 결과, 입자 궤도

그림 3의 각 패널은 다음과 같이 읽는다.

좌상단: source 분포. 처음에는 표준 정규 분포(Standard Normal)의 둥근 구름 형태이다.
우상단: target 분포. 8개의 모드(mode)가 원주 위에 나열되어 있다.
좌하단: 생성 결과. Euler sampler를 통해 운반된 점들이 target의 8개 모드로 가까워지는지 확인한다.
우하단: 궤도도(Trajectory plot). 몇몇 입자들이 $t=0$에서 $t=1$까지 어떻게 움직이는지 확인한다.

그림 3은 loss 수치만으로는 보이지 않는 실패를 찾아내기 위해 중요하다. 예를 들어, loss가 낮아지고 있더라도 궤도가 크게 요동치고(serpentine) 있다면 적은 NFE(Number of Function Evaluations)로 생성하는 것이 어려워진다. 또한, 궤도는 학습된 주변 속도장(marginal velocity field)을 ODE로서 적분한 결과이므로, 쌍(pair)별 교사 경로(teacher path) 그 자체와는 일치하지 않는다.

그림 4는 저장된 체크포인트(checkpoint) minimal_2d.pt 내의 samples를 묘사한 생성 샘플이다. 학습 과정이나 궤도를 나타내는 그림이 아니라, 체크포인트에 저장된 생성 종점(endpoint)만을 확인하기 위한 그림이다.

그림 4 minimal_2d.pt 에 저장된 생성 종점과 target의 8 mode 중심

그림 4의 파란색 점은 저장된 모델로부터 생성된 샘플이며, 빨간색 점은 target의 8개 mode 중심이다. 이는 간이 확인용이며, 그림 3의 좌하단 패널을 단독으로 보고 있는 것과 유사하다. 그림 4만으로는 궤도의 굴곡이나 NFE 의존성을 알 수 없다.

실행 예시:

python .\src\plot_minimal_2d.py `
--checkpoint .\_outputs\minimal_2d.pt `
--out .\figures\minimal_2d_result.png

관찰 대상	Flow Matching에서의 의미	다음에 의심할 점
generated samples가 target에 올라타지 않음	속도장이 분포를 올바르게 운반하지 못함	학습 부족, 모델 용량 부족, 시간 임베딩 (time embedding)
...

이 시점에서는 완벽한 생성 품질이 목표가 아니다. 목표는 "Flow Matching의 식이 실제 점군(point cloud)의 움직임으로서 무엇을 일으키는가"를 눈으로 확인하는 것이다.

Meta flow_matching으로 넘어갈 때는 먼저 2D continuous example을 본다. 교재와의 대응 관계는 다음과 같다.

구체적인 URL:

Meta flow_matching: https://github.com/facebookresearch/flow_matching
2D continuous example: https://github.com/facebookresearch/flow_matching/tree/main/examples

교재에서 본 것	공식 구현에서 볼 곳	역할
`LinearPath`	`flow_matching/path`	$x_t$와 교사 속도(teacher velocity)를 만드는 부품
`conditional_flow_matching_loss`	`flow_matching/loss` 또는 example 내의 loss 계산	모델 출력과 교사 신호(teacher signal)의 비교
`euler_solve`	`flow_matching/solver`	학습된 속도장으로 점을 이동시킴
2D toy 시각화	`examples/2d_flow_matching.ipynb`	path, loss, solver를 작게 확인하는 입구

공식 코드를 읽을 때도 처음에는 모델 구조가 아니라 path -> loss -> solver 순으로 찾는다.

여기까지 최소 구현을 한 바퀴 돌았다. 제7장에서는 제1장부터 제6장까지 사용한 CFM을 구현 부품으로서 정리한다. 제8장에서는 점을 ODE로 움직이는 것이 분포의 변형으로 어떻게 이어지는지를 연속 방정식(continuity equation)과 pushforward로 보완한다. 그 후, 제9장과 제10장에서 "왜 주변 속도장을 직접 교사로 사용할 수 없는가", "CFM은 무엇을 조건부(conditional)로 하고 있는가"를 이론적으로 깊이 파고든다. 먼저 동작하는 코드의 역할을 고정하고, 분포가 움직이는 직관을 얻은 다음, CFM의 이유를 수식으로 확인하는 순서이다.

이 장에서 다루는 개념: CFM 구현은 시간을 선택하고, 경로 위의 점을 만들고, 교사 속도를 만들고, 모델로 예측하고, 제곱 오차(squared error)를 구하는, 서로 교체 가능한 부품들로 구성된다. 구현 단위로 읽는다는 것은 수식상의 선택이 어떤 클래스나 함수에 갇혀 있는지 확인하는 것이다.

주요 단계: 학습 구현 — time sampler, path, velocity model, loss를 하나의 학습 스텝(learning step)으로 연결한다.

CFM의 구현은 모델 본체보다 "훈련 데이터를 어떻게 만드는가"가 중요하다. 모델이 MLP든 U-Net이든 DiT든, 최소한의 훈련 스텝은 다음과 같은 형태가 된다.

source를 가져온다
target을 가져온다
t를 가져온다
...

이 순서대로 읽으면 Flow Matching 계열의 많은 코드는 상당히 명확하게 파악된다.

이 교재의 최소 구현에서는 역할을 다음과 같이 나누고 있다.

파일	역할	대응하는 개념
paths.py	$x_t$ 와 $u_t$ 를 생성	조건부 경로 (Conditional Path)
...

연구 코드를 읽을 때도 이러한 분해를 찾아보는 것이 좋다. 파일명은 다르더라도 대개 동일한 역할이 존재한다.

현재의 LinearPath는

x_t = (1-t)x_0 + tx_1,
\quad
u_t = x_1 - x_0

만을 구현하고 있다.

만약 Gaussian path나 variance preserving path로 변경한다면, 주로 바뀌는 것은 paths.py의 sample과 velocity이다. loss의 외형은 크게 변하지 않는다.

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

이 구조가 보인다면, "Flow Matching의 변종"은 완전히 별개의 것이 아니라 어떤 부품을 교체한 것인지로서 읽을 수 있다.

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

Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, Matt Le, "Flow Matching for Generative Modeling", ICLR 2023. Flow Matching의 기본 정식화, Gaussian path, CFM의 기초가 되는 논문.
Alexander Tong, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector-Brooks, Kilian Fatras, Guy Wolf, Yoshua Bengio, "Conditional Flow Matching: Simulation-Free Dynamic Optimal Transport", 2023. CFM, OT-CFM, simulation-free한 dynamic OT 근사를 다루는 논문.
Xingchao Liu, Chengyue Gong, Qiang Liu, "Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow", 2022. Rectified Flow, Reflow, 적은 스텝 생성의 실용적인 입구가 되는 논문.
Qiang Liu, "Rectified Flow: A Marginal Preserving Approach to Optimal Transport", 2022. Rectified Flow의 marginal preserving property와 OT와의 관계를 설명하는 이론 중심의 논문.
Sangyun Lee, Zinan Lin, Giulia Fanti, "Improving the Training of Rectified Flows", NeurIPS 2024. Reflow를 반복하기 전에, path의 교차, 시간 분포, loss 설계를 개선하는 아이디어를 다룸.
Michael S. Albergo, Eric Vanden-Eijnden, "Building Normalizing Flows with Stochastic Interpolants", 2023. 보간 과정(interpolation process)을 설계하고, 그 확률 흐름(probability flow)으로부터 속도장(velocity field)을 학습하는 Stochastic Interpolants의 기초 논문.
Michael S. Albergo, Nicholas M. Boffi, Eric Vanden-Eijnden, "Stochastic Interpolants: A Unifying Framework for Flows and Diffusions", 2023. Flow와 Diffusion을 보간(interpolation), 속도(velocity), score의 프레임워크로 통합하여 바라보는 논문.
Patrick Esser et al., "Scaling Rectified Flow Transformers for High-Resolution Image Synthesis", 2024. 고해상도 이미지 생성(High-Resolution Image Synthesis)을 위해 Rectified Flow Transformer를 스케일링하는 실무 모델 중심의 논문.
Yaron Lipman et al., "Flow Matching Guide and Code", 2024. Flow Matching의 용어, 구현, 공식 코드 대응을 정리한 포괄적인 가이드.
Peter Holderrieth, Ezra Erives, "An Introduction to Flow Matching and Diffusion Models", MIT 6.S184 lecture notes. Diffusion과 Flow Matching을 ODE/SDE, score, velocity 관점에서 연결하는 강의 노트.
Will Grathwohl et al., "FFJORD: Free-form Continuous Dynamics for Scalable Reversible Generative Models", ICLR 2019. Continuous Normalizing Flow의 instantaneous change-of-variables, 발산(divergence) 추정, NLL 계산에 관한 기초 자료.
Cheng Lu et al., "DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps", NeurIPS 2022. NFE를 맞춘 고속 샘플링(sampling) 평가와 적은 스텝(step)에서의 품질 비교를 고려하기 위한 대표적인 솔버(solver) 연구.
Tero Karras et al., "Elucidating the Design Space of Diffusion-Based Generative Models", NeurIPS 2022. 샘플러(sampler) 설정, NFE, FID, 계산량을 분리하여 평가하는 실험 설계의 참고 자료.
Jonathan Ho, Ajay Jain, Pieter Abbeel, "Denoising Diffusion Probabilistic Models", 2020. DDPM의 기본이 되는 순방향 노이즈 추가(forward noise addition), epsilon prediction, 역과정(reverse process)에 관한 논문.
Martin Heusel et al., "GANs Trained by a Two Time-Scale Update Rule Converge to a Local Nash Equilibrium", NeurIPS 2017. Fréchet Inception Distance (FID)를 도입한 논문.
Mikołaj Bińkowski et al., "Demystifying MMD GANs", ICLR 2018. Kernel Inception Distance (KID)와 그 불편 추정(unbiased estimation)을 도입한 논문.
Catherine F. Higham, Desmond J. Higham, Peter Grindrod, "Diffusion Models for Generative Artificial Intelligence: An Introduction for Applied Mathematicians", 2023. 응용 수학 관점에서 Diffusion 모델을 소개하는 입문 논문.
TorchCFM documentation, atong01/conditional-flow-matching. CFM 변형들을 FlowMatcher 클래스로 읽기 위한 구현 자료.
Diff2Flow documentation, "Training Flow Matching Models via Diffusion Model Alignment". Diffusion 구현을 Flow Matching으로 변환 및 정렬(alignment)하는 관점의 구현 자료.
Black Forest Labs, FLUX official repository. FLUX.1의 공식 추론(inference) 코드, 모델 목록, 이용 조건으로 연결되는 입구.
Black Forest Labs, FLUX.1-schnell model card. 공식 체크포인트(checkpoint), 파이프라인(pipeline) 구성, Apache-2.0 라이선스 표기 확인처.
Black Forest Labs, FLUX.1-dev model card. 게이트(gated) 체크포인트(checkpoint) 및 독자 라이선스로의 입구.
Black Forest Labs, FLUX.1 [dev] Non-Commercial License. 가중치(weight), 파생물(derivative), 출력물(output)을 구분하여 이용 조건을 확인하는 1차 자료.
Stability AI, Stable Diffusion 3.5 official repository. SD3.5 모델, 추론 예시, 모델 카드(model card)로의 공식 입구.
Stability AI, Stable Diffusion 3.5 Medium and Large model cards. MMDiT 구성, 용도, Community License의 확인처.
Stability AI, Community License. SD3.5를 포함한 대상 모델의 이용 및 재배포 조건을 확인하는 1차 자료.
Tencent, HunyuanVideo official repository. 공개 비디오 생성(video generation) 모델, 체크포인트(checkpoint), 추론 조건, 라이선스의 1차 자료.
William Peebles et al., SiT: Exploring Flow and Diffusion-based Generative Models with Scalable Interpolant Transformers. Stochastic Interpolant Transformer의 공식 구현 및 ImageNet 체크포인트(checkpoint).
Hugging Face Diffusers, DreamBooth training for FLUX. Full/LoRA 학습, 대상 모듈(module), 텍스트 인코더(text encoder) 학습, VRAM 관련 주의사항.
Hugging Face Diffusers, DreamBooth training for Stable Diffusion 3. SD3 계열 MMDiT의 LoRA 학습 예시.
Hugging Face Diffusers, LoRA training and loading adapters. 어댑터(adapter)의 학습, 로드(load), 다중 어댑터(multiple adapter), fuse/unfuse 관련 공식 문서.
Hugging Face PEFT, LoRA developer guide. rank/alpha, 대상 모듈(module), merge/unmerge, 발전된 LoRA 설정에 관한 공식 문서.
Hugging Face, safetensors documentation. 텐서(tensor)와 메타데이터(metadata)를 저장하는 형식의 사양 및 이용법.
Edward J. Hu et al., LoRA: Low-Rank Adaptation of Large Language Models. 저차원 업데이트(low-rank update)에 관한 원 논문.

Flow Matching 입문 (4) ― PyTorch를 이용한 Flow Matching 학습 흐름

요약

핵심 포인트

댓글