Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow

Background & Academic Lineage

The Origin & Academic Lineage

하나의 확률 분포를 다른 분포로 변환하는 문제, 이른바 "transport mapping problem"은 머신러닝과 통계학의 근본적인 난제이다. 역사적으로 이 문제는 분포 간의 질량을 이동시키는 가장 효율적인 방법을 탐구하는 Optimal Transport (OT) 분야에서 기원하였다. OT는 엄밀한 수학적 프레임워크를 제공하지만, 현대의 이미지 생성이나 도메인 전이(domain transfer) 작업에서 마주하는 고차원 공간에서는 이를 해결하기가 극도로 어렵다는 한계가 있다.

기존의 접근 방식, 특히 Generative Adversarial Networks (GANs)이나 Variational Autoencoders (VAEs)와 같은 생성 모델들은 데이터와 Latent Space 간의 매핑을 학습함으로써 이 문제를 해결하고자 하였다. 그러나 이러한 모델들은 고질적인 문제에 직면해 있다. GAN은 수치적 불안정성과 Mode Collapse 문제로 고통받으며, VAE를 비롯한 Likelihood 기반 모델들은 복잡하고 계산 비용이 높은 추론 과정을 요구하는 경우가 많다. 최근에는 Diffusion models나 Neural Ordinary Differential Equations (ODEs)와 같은 연속 시간 모델(continuous-time models)이 주목받고 있다. 이들은 강력한 성능을 자랑하지만, 본질적으로 "infinite-step" 프로세스라는 특성을 지닌다. 즉, 복잡한 미분 방정식을 풀기 위해 고비용의 신경망을 반복적으로 호출해야 하므로, 실시간 응용이나 빠른 추론을 수행하기에는 지나치게 느리다. 본 논문의 저자들은 이러한 연속 시간 모델의 핵심적인 한계가 곡선 형태의 비직선적 궤적(curved, non-straight trajectories)에 의존한다는 점에 있으며, 이것이 정확한 시뮬레이션을 위해 수많은 Discretization Step을 요구하게 만든다는 점을 포착하였다.

Intuitive Domain Terms

• Rectified Flow: "고속도로를 직선화하는 것"으로 이해할 수 있다. 데이터 입자들이 두 분포 사이에서 구불구불하고 비효율적인 경로를 따라 이동하게 두는 대신, 가능한 최단 직선 경로를 따르도록 강제함으로써 이동을 훨씬 빠르고 계산하기 쉽게 만든다.
• Reflow: 첫날 구불구불한 경로로 배달하던 운전자가 교통 상황을 관찰한 뒤, 경로를 완벽한 직선으로 "재조정(reflow)"하는 과정에 비유할 수 있다. 이전 모델이 생성한 경로를 바탕으로 반복 학습함으로써 시스템은 스스로 궤적을 "직선화"하며, 이를 통해 훨씬 적은 단계만으로도 고품질의 결과를 도출한다.
• Coupling: 이는 "짝짓기 계획"이다. 모래 더미(분포 $\pi_0$)를 특정 형태(분포 $\pi_1$)로 옮기려 할 때, 각 모래 알갱이가 정확히 어디로 가야 하는지를 지시하는 일련의 명령 체계이다.
• Drift Force: ODE의 맥락에서 이는 모델의 "운전대" 역할을 한다. 데이터 포인트가 목적지에 도달하기 위해 특정 시점 $t$에 어느 방향으로 이동해야 하는지를 알려주는 신경망이다.
• Discretization Step: 영상의 "프레임 레이트"와 같다. 연속적인 움직임을 시뮬레이션하기 위해 이를 작은 단위로 쪼개는 것이다. 단계 수가 많을수록 부드럽지만 느린 프로세스가 되며, 저자들은 매우 적은 단계(심지어 단 한 단계)만으로도 높은 품질을 달성하는 것을 목표로 한다.

Notation Table

Problem Definition & Constraints

Core Problem Formulation & The Dilemma

본 논문은 고차원 공간에서 경험적으로 관측된 두 데이터 분포 $\pi_0$와 $\pi_1$ 사이의 Transport Map을 학습하는 근본적인 문제를 다룬다. 이는 생성 모델링(예: 가우시안 노이즈를 이미지로 매핑) 및 도메인 전이(예: 이미지 스타일 변환)를 포함한 다양한 머신러닝 응용 분야에서 핵심적인 과제이다.

Input/Current State: 출발점은 $\mathbb{R}^d$ 상의 두 분포 $\pi_0$와 $\pi_1$에서 얻은 경험적 관측치(샘플)이다. 이 문제의 결정적인 측면은 쌍을 이루는 입출력 데이터의 부재(lack of paired input/output data)이다. 즉, 각 샘플 $X_0 \sim \pi_0$에 대해 "올바른" 변환 또는 생성 타겟으로 알려진 대응 샘플 $X_1 \sim \pi_1$이 존재하지 않는다. 우리는 단지 각 분포에서 독립적으로 추출된 샘플 집합만을 가지고 있을 뿐이다.

Output/Goal State: 최종 목표는 무한 데이터 극한에서 $Z_0 \sim \pi_0$일 때 $Z_1 := T(Z_0) \sim \pi_1$이 성립하도록 하는 Transport Map $T: \mathbb{R}^d \to \mathbb{R}^d$를 학습하는 것이다. 구체적으로, 본 논문은 최대한 "직선"에 가까운 경로를 따라 $\pi_0$에서 $\pi_1$로 샘플을 이동시킬 수 있는 Neural ODE 모델 $dZ_t = v(Z_t, t)dt$를 학습하고자 한다. 이 ODE는 새로운 데이터를 생성하거나 도메인 전이를 수행하기 위해 순방향으로 시뮬레이션 가능해야 한다.

Missing Link/Mathematical Gap: 핵심적인 공백은 생성 모델링과 도메인 전이를 통합하면서도 기존 방법론의 한계를 극복할 수 있는, 쌍을 이루지 않은 데이터로부터 인과적(causal)이고 계산 효율적인 Transport Map을 구축하는 방법이다.

이 간극을 메우려는 기존의 시도들은 다음과 같은 문제에 직면했다:
1. Naive Linear Interpolation: 단순 선형 보간 $X_t = tX_1 + (1-t)X_0$은 직선 경로를 제공하지만 "비인과적(또는 예측적)"이다. 이는 $X_t$를 결정하기 위해 최종 지점 $X_1$을 미리 알아야 하므로, 새로운 데이터를 생성하기 위한 순방향 시뮬레이션이 불가능하다.
2. Optimal Transport (OT): OT는 운송 비용을 최소화하는 매핑을 찾는 이론적으로 건전한 프레임워크를 제공하지만, 고차원 연속 측정치에 대해서는 "계산적으로 매우 어렵고", 많은 머신러닝 작업의 특정 목적과는 "직접적인 관련이 없는" 경우가 많다.
3. Continuous-Time Generative Models (ODEs/SDEs): Score-based generative models 및 DDPM과 같은 최근의 발전은 인상적인 결과를 보여주었다. 그러나 이러한 모델들은 "사실상 'infinite-step'"이며, ODE/SDE를 시뮬레이션하기 위해 고비용의 신경망 힘장(force field)을 수없이 반복 호출해야 하므로 "추론 시 높은 계산 비용"이 발생한다.

본 논문은 이 문제를 단순한 비선형 최소자승 최적화(nonlinear least squares optimization)로 정식화함으로써 이 간극을 메우고자 한다. 이는 경험적으로 샘플링된 지점들 사이의 선형 보간 $X_t = tX_1 + (1-t)X_0$의 방향을 최대한 따르도록 ODE $dZ_t = v(Z_t, t)dt$를 구동하는 속도장 $v(Z_t, t)$를 학습하는 것이다. 이는 다음과 같이 표현된다:
$$ \min_v \mathbb{E} \left[ \int_0^1 \|(X_1 - X_0) - v(X_t, t)\|^2 dt \right] $$
이 정식화는 선형 보간의 직선 경로를 "인과화(causalize)"하여 시뮬레이션 가능하게 만드는 것을 목표로 한다.

Constraints & Failure Modes

분포 간의 Transport Map을 학습하는 문제는 다음과 같은 가혹하고 현실적인 제약 조건들에 얽매여 있다:

Physical, Computational, or Data-driven Constraints:
* Unpaired Data: 가장 중요한 데이터 기반 제약은 비지도 학습 환경에서 내재된 "쌍을 이루는 입출력 데이터의 부재"이다. 이는 모델이 $X_0$에서 $X_1$으로의 직접적인 회귀를 단순히 학습할 수 없음을 의미한다.
* High-Dimensionality of Data: 실제 데이터, 특히 이미지는 매우 고차원 공간($d$가 수백만 단위일 수 있는 $\mathbb{R}^d$)에 존재한다. 이는 직접적인 OT 계산을 불가능하게 만들며, 수치적 ODE/SDE 솔버의 계산 비용을 가중시킨다.
* Computational Cost of ODE/SDE Solvers: 기존의 연속 시간 모델들은 추론 과정에서 "고비용의 신경망 힘장을 수없이 반복 호출"해야 한다. 이는 수백, 수천 단계에 걸쳐 이미지를 생성하는 것이 너무 느린 많은 응용 분야에서 엄격한 실시간 지연 시간(latency) 요구사항과 충돌한다.
* Non-Crossing Property of ODEs: 잘 정의된 ODE의 경우, 그 해는 유일해야 하며 이는 서로 다른 경로가 교차할 수 없음을 의미한다. 이는 교차할 수 있는 단순 선형 보간과 달리, 학습된 모든 Flow가 반드시 만족해야 하는 근본적인 수학적 제약이다.

Why This Approach

The Inevitability of the Choice

저자들은 기존의 생성 모델들, 특히 GAN과 Diffusion 모델이 추론 속도 측면에서 근본적인 "계산적 벽"에 부딪혔음을 확인했다. GAN은 빠르지만 악명 높은 학습 불안정성과 Mode Collapse 문제를 겪는다. 반면, Diffusion 모델(및 그 ODE 기반 변형인 PF-ODE)은 수학적으로는 견고하지만, 정확도를 유지하기 위해 많은 Discretization Step을 요구하는 복잡하고 곡선인 궤적을 풀어야 하므로 계산 비용이 높다. 저자들은 이러한 궤적의 "곡선"적 성질이 주된 병목 현상임을 깨달았다. 만약 두 분포 사이의 이동 경로를 "직선"으로 만들 수 있다면, ODE를 최소한의 Discretization(심지어 단 한 단계)으로 풀 수 있을 것이다. 이러한 깨달음은 단순히 분포를 맞추는 것에서 벗어나, 분포 사이의 가장 짧고 곧은 경로를 찾는 것으로 초점을 전환시켰다.

Comparative Superiority

Rectified Flow는 Transport 문제를 단순하고 확장 가능한 비제약 최소자승 최적화로 변환하기 때문에 질적으로 우월하다. 섬세한 Minimax 균형을 요구하는 GAN이나 복잡한 SDE/ODE 솔버에 의존하는 Diffusion 모델과 달리, Rectified Flow는 "Reflow" 절차를 사용한다. 이 절차는 Flow의 궤적을 반복적으로 직선화한다. 구조적으로 이는 Discretization 오차를 크게 줄여준다. 표준 Diffusion 모델이 고품질 이미지를 생성하기 위해 수백 번의 함수 평가(NFE)를 필요로 할 수 있는 반면, Rectified Flow는(특히 Reflow 이후) 단 한 번의 Euler step만으로도 그와 대등하거나 더 뛰어난 결과를 생성할 수 있다. 이는 VAE와 같은 One-step 모델과 연속 시간 모델 사이의 간극을 효과적으로 메우며, 후자의 높은 품질과 전자의 속도를 동시에 제공한다.

Alignment with Constraints

문제의 제약 조건들은 GAN의 불안정성이나 Diffusion의 과도한 추론 비용 없이 고차원 데이터(이미지 등)를 처리할 수 있는 모델을 요구했다. Rectified Flow는 이동 경로의 "인과화"를 통해 이러한 제약 조건들을 만족시킨다. Drift Force $v$가 선형 보간 $X_t = tX_1 + (1-t)X_0$을 따르도록 학습함으로써, 모델은 근시안적이고 교차하지 않으며 결정론적인 방식으로 질량을 이동시키는 법을 배운다. ODE 프레임워크와 직선 경로 목표의 이러한 "결합"은 모델이 계산적으로 효율적(직선 경로 덕분)이면서도 이론적으로 건전(주변 분포를 보존하고 운송 비용을 감소시키므로)하도록 보장한다.

Mathematical & Logical Mechanism

The Master Equation

Rectified Flow의 핵심 메커니즘은 소스 분포 $\pi_0$를 타겟 분포 $\pi_1$로 변환하는 속도장 $v(z, t)$를 직선 경로를 따라 학습하는 것이다. 이 속도장을 학습하기 위해 사용되는 목적 함수는 다음과 같다:

$$\min_{v} \int_{0}^{1} \mathbb{E} \left[ \left\| (X_1 - X_0) - v(X_t, t) \right\|^2 \right] dt, \quad \text{with } X_t = tX_1 + (1 - t)X_0$$

Step-by-Step Flow

1. Initialization: 데이터 분포에서 쌍 $(X_0, X_1)$을 샘플링한다.
2. Interpolation: 시스템은 무작위로 샘플링된 시간 $t$에서 중간 지점 $X_t$를 계산한다.
3. Velocity Prediction: 신경망 $v$는 현재 상태 $X_t$와 시간 $t$를 입력으로 받아 예측된 속도 벡터를 출력한다.
4. Regression: 모델은 예측된 속도를 타겟 방향 $(X_1 - X_0)$과 비교한다.
5. Update: 경사 하강법을 통해 차이를 최소화하도록 네트워크 파라미터를 업데이트한다.
6. Inference: 샘플링 시, 모델은 $Z_0 \sim \pi_0$에서 시작하여 수치 솔버(Euler 방법 등)를 사용하여 ODE $dZ_t = v(Z_t, t)dt$를 풀어 $Z_1 \sim \pi_1$에 도달한다.

Optimization Dynamics

이 메커니즘은 선형 보간을 "인과화"함으로써 학습한다. 단순 경로 $X_t$는 미래($X_1$)에 대한 지식을 요구하지만, 학습된 속도장 $v(Z_t, t)$는 현재 상태와 시간에만 의존하는 함수이므로 유효하고 인과적인 ODE가 된다.

"Reflow" 절차는 결정적인 최적화 동역학이다. 초기 모델을 학습한 후, 학습된 Flow를 시뮬레이션하여 새로운 쌍 $(Z_0, Z_1)$을 생성하는 데 사용한다. 이 새로운 쌍들은 모델을 재학습하는 데 사용된다. 첫 번째 모델에 의해 생성된 Flow는 이미 원본 데이터 커플링보다 "더 직선"이기 때문에, 두 번째 반복에서는 훨씬 더 직선인 경로가 생성된다. 이 반복적인 과정은 Flow를 효과적으로 "직선화"하여 수치 솔버의 Discretization 오차를 줄인다. 결과적으로 손실 지형(loss landscape)은 점점 더 매끄러워지며, 모델은 매우 적은(심지어 단 한 번의) Euler step만으로도 고품질 샘플을 생성할 수 있는 상태로 수렴하게 된다.

Results, Limitations & Conclusion

Experimental Design & Baselines

저자들은 CIFAR-10 데이터셋과 고해상도 데이터셋(LSUN, CelebA-HQ, AFHQ)을 사용한 무조건적 이미지 생성(unconditional image generation)을 통해 Rectified Flow를 평가한다. 엄밀한 베이스라인을 구축하기 위해 DDPM++ 프레임워크(Song et al., 2020b)의 U-Net 아키텍처를 활용한다. 실험 설계는 "Reflow" 절차의 효능과 그 결과로 나타나는 학습된 ODE 궤적의 "직선성"을 테스트하도록 구성되었다.

What the Evidence Proves

제시된 증거는 특히 Reflow 절차의 "직선화" 효과와 관련하여 설득력이 있다. 저자들은 초기(1-rectified) Flow가 효과적이기는 하지만 완벽하게 직선은 아님을 증명한다. 이전 Flow에 의해 생성된 데이터로 모델을 재학습하는 Reflow 절차를 적용함으로써 궤적은 점차 선형성을 띠게 된다.

이 메커니즘에 대한 결정적인 증거는 두 가지이다:
* Quantitative: CIFAR-10에서 증류된 2-rectified flow는 FID 4.85를 달성하여, 기존에 알려진 최고의 One-step 생성 모델(TDPM, FID 8.91)을 크게 앞선다. 또한, 0.51의 Recall은 StyleGAN2+ADA(0.49)를 상회하며, 이 방법이 높은 다양성을 유지함을 입증한다.
* Visual/Geometric: Figure 4와 Figure 18은 2-rectified flow의 궤적이 거의 직선임을 시각적으로 증명한다. 외삽(extrapolation) $\hat{z}_1^t = z_t + (1-t)v(z_t, t)$는 $t$에 관계없이 거의 일정하게 유지되는데, 이는 직선 ODE의 특징이다. 이는 모델이 Transport 과정을 성공적으로 "인과화"하여 최소한의 Discretization Step으로 정확한 시뮬레이션을 가능하게 했음을 확인시켜 준다.

Limitations & Future Directions

본 연구의 향후 방향은 다음과 같다:
* Theoretical Refinement: 궤적 직선화의 이점이 수치적 오차 누적의 이점을 상쇄하기 전까지의 Reflow 단계 수에 이론적 한계가 존재하는지 탐구한다.
* Broader Applications: "직선화" 속성을 물리 시스템 모델링이나 시계열 예측과 같은 비생성적 작업에 활용할 수 있는지 조사한다.
* Optimal Transport Integration: 저자들이 언급했듯이, Rectified Flow는 특정 비용 함수 $c$에 대한 $c$-optimal transport를 엄격하게 보장하지는 않는다. 향후 연구는 속도장 $v$가 기울기 장(예: $v = \nabla f$)이 되도록 제약하여 최적성을 명시적으로 강제하는 데 집중할 수 있다.

이러한 발견은 생성 모델링의 패러다임 전환을 시사한다. "노이즈-데이터" 확산 패러다임에서 계산적으로 더 효율적이고 이론적으로 더 투명한 "직선" Transport 패러다임으로의 이동이다.

Isomorphisms with other fields

Structural Skeleton

두 확률 분포 사이의 비인과적이고 교차하는 보간 경로를 결정론적이고 교차하지 않는 직선 ODE Flow로 변환하는 메커니즘.

Distant Cousins

1. Target Field: Fluid Dynamics

   • The Connection: 운송 비용을 최소화하기 위해 궤적을 반복적으로 직선화하는 "Reflow" 절차는 파이프 내의 층류(laminar flow)를 찾는 문제의 거울상이다. Rectified Flow가 교차를 피하고 에너지 소산을 최소화하기 위해 궤적을 "재배선(rewire)"하는 것과 마찬가지로, 유체 역학은 점성 항력을 최소화하는 매끄럽고 평행한 유선을 얻기 위해 난류 소용돌이(교차)를 제거하고자 한다.

2. Target Field: Urban Traffic Engineering

   • The Connection: 비인과적 선형 보간에서 Rectified Flow로의 전환은 정적인 격자 기반 도로망에서 적응형 지능형 교통 관리 시스템으로의 전환과 유사하다. 원래의 보간에서 경로는 맹목적으로 교차한다(교차로의 정체처럼). Rectified Flow는 교통 흐름을 "재배선"하는 중앙 제어 장치 역할을 하여, 입자(차량)가 목적지에 도달하기 위해 가장 효율적이고 충돌 없는 경로를 따라 이동하도록 보장함으로써 전체 시스템의 처리량을 효과적으로 최적화한다.

What If Scenario

만약 양자장론(Quantum Field Theory) 연구자가 이 방정식을 "차용"한다면, 그들은 Rectified Flow 메커니즘을 경로 적분(Path Integral) 정식화에 적용할 수 있을 것이다. 양자 상태 간의 전이를 모든 가능한 경로의 합이 아닌 Rectified Flow로 처리함으로써, 시뮬레이션하기에 계산적으로 사소한 최소 작용의 "직선화된" 경로를 도출할 수 있을지도 모른다. 이는 고비용의 몬테카를로 샘플링 없이도 고차원 양자 시스템에서의 전이 진폭을 정확하게 계산할 수 있게 하여, 복잡하고 비선형적인 양자 상호작용을 일련의 결정론적이고 One-step인 "직선" 전이로 효과적으로 변환할 것이다.

Universal Library of Structures

본 논문은 이미지, 확률 분포, 물리적 구성 등 두 상태 사이를 매핑하는 근본적인 과제가 본질적으로 기저 공간에서 가장 효율적이고 교차하지 않는 기하학을 찾는 문제임을 보여준다. 이는 Optimal Transport와 Flow Rectification의 원리가 모든 과학 분야에 걸쳐 복잡성을 단순화하는 보편적인 도구임을 입증한다.