Transition Flow Matching: 기술적 요약

1. 문제 정의 (Problem)

기존의 생성 모델 (Flow Matching, Diffusion Models) 은 주로 **국소적인 속도장 (local velocity field)**을 학습합니다. 이 방식은 생성 과정에서 오일러-라그랑주 방정식 (ODE) 또는 확률 미분 방정식 (SDE) 을 수치적으로 적분해야 하므로, 고품질 샘플을 생성하기 위해 여러 단계 (multi-step integration) 의 추론이 필요합니다.

• 기존 접근법의 한계:
  • 적분 의존성: 국소 속도만 알기 때문에 다음 상태를 예측하려면 작은 시간 간격으로 반복적으로 적분해야 함.
  • 단일 단계 생성의 어려움: Consistency Models 와 같은 단일 단계 (one-step) 생성 모델은 네트워크 출력에 일관성 제약을 부과하지만, 이는 경험적 행동 특성 (behavioral property) 에 기반하며, 학습을 유도해야 할 근본적인 속도장의 특성이 명확하지 않아 학습이 불안정하거나 복잡한 커리큘럼이 필요함.
  • Distillation 의존성: 많은 Few-step 모델이 사전 학습된 Teacher 모델로부터 증류 (distillation) 를 필요로 함.

2. 방법론 (Methodology)

저자들은 **Transition Flow Matching (TFM)**이라는 새로운 패러다임을 제안합니다. 이는 국소 속도장을 회귀하는 대신, **생성 궤적 자체 (generation trajectory)**를 직접 모델링하여 **전역적인 전이 흐름 (global transition flow)**을 학습하는 방식입니다.

핵심 개념: Transition Flow Identity

TFM 은 두 시간 $t$와 $r$ ($t \le r$) 사이의 상태 전이를 직접 예측하는 함수 $X(x_t, t, r)$를 학습합니다. 이 모델은 현재 상태 $x_t$에서 임의의 미래 시간 $r$의 상태 $x_r$을 한 번의 추론으로 매핑할 수 있습니다.

저자들은 **Transition Flow Identity (전이 흐름 항등식)**을 수학적으로 유도했습니다.
$$X(x_t, t, r) = x_{t \to r} + (r-t) \frac{d}{dt} X(x_t, t, r)$$
여기서:

• $X(x_t, t, r)$: 모델이 예측하는 전이 흐름 (미래 상태).
• $x_{t \to r}$: 실제 전이된 상태.
• $\frac{d}{dt} X(x_t, t, r)$: 전이 흐름의 시간 미분 (Jacobian-vector product 를 통해 계산).

이 항등식은 국소 속도 $v$와 전이 흐름 $X$ 사이의 관계를 명확히 하며, 모델이 이 항등식을 만족하도록 학습할 수 있는 이론적 기반을 제공합니다.

학습 목표 (Training Objective)

이론적으로 이상적인 목적 함수는 전역 (marginal) 분포에 기반하지만 계산이 불가능합니다. 따라서 조건부 (conditional) 목적 함수를 사용하여 학습합니다.

1. 조건부 전이 상태: $Z=(X_0, X_1)$를 조건으로 할 때, 선형 보간 $X_t = (1-t)X_0 + tX_1$를 가정합니다.
2. 구현 가능한 목표: 조건부 속도 $v(x_t, t|Z) = X_1 - X_0$가 상수이므로, 전이 흐름의 시간 미분 $\frac{d}{dt}X^\theta$를 계산하기 쉽습니다.
3. 손실 함수 (C-TFM):
   $$\mathcal{L}_{\text{C-TFM}} = \mathbb{E} \left[ D \left( \text{sg}[X_{t \to r}^Z + (r-t)\frac{d}{dt}X^\theta], X^\theta \right) \right]$$
   • $\text{sg}$: Stop-gradient 연산자 (목표값으로 고정).
   • $X_{t \to r}^Z$: 조건부 전이 상태 (선형 보간으로 계산 가능).
   • 이 목적 함수는 Theorem 2에 의해 전역 목적 함수와 기울기가 동일함을 보장받습니다.

추론 (Inference)

학습된 모델 $X^\theta$를 사용하여 소스 분포 (노이즈) 에서 타겟 분포 (데이터) 로의 전이를 수행합니다.

• 임의의 단계 수: $x_{r} = X^\theta(x_t, t, r)$를 반복 적용하여 임의의 시간 간격과 단계 수로 생성 가능.
• 단일 단계 (One-step): $t=0, r=1$로 설정하여 노이즈에서 바로 데이터를 생성 가능.

3. 주요 기여 (Key Contributions)

1. Transition Flow Matching 프레임워크 제안: 국소 속도장이 아닌 전이 흐름 (transition flow) 을 직접 학습하여 Few-step 및 One-step 생성을 가능하게 하는 원칙적인 프레임워크를 제시.
2. Transition Flow Identity 유도 및 이론적 기반: 전이 흐름과 국소 속도 간의 관계를 수학적으로 증명하고, 이를 만족하도록 End-to-End 학습 가능한 목적 함수를 설계.
3. Mean Velocity 모델과의 통합적 관점: 제안된 방법이 기존 Mean Velocity 모델 및 Consistency 모델과 어떻게 연결되는지 명확히 하여, 생성 역학에 대한 통합된 이론적 관점을 제공.
4. 성능 입증: CIFAR-10 및 ImageNet-256 에서 다양한 Few-step 설정 (NFE=1, 2, 5, 10) 에서 기존 최첨단 모델 (Consistency Models, MeanFlow 등) 을 능가하는 성능 달성.

4. 실험 결과 (Results)

1) 2D 합성 데이터 시각화

• 알파벳 "M" 형태의 데이터 분포로 생성 궤적을 시각화.
• 제안된 방법은 단일 단계 (One-step) 에서도 명확한 "M" 형태를 생성하며, 단계 수를 늘리면 더 정교한 결과를 보여 생성 궤적 모델링의 유효성을 입증.

2) CIFAR-10 (32x32)

• FID 점수: 단일 단계 (NFE=1) 에서 가장 낮은 FID 점수 (2.77) 를 기록하여 기존 모델들 (Flow Matching, Consistency Models 등) 을 압도.
• NFE 증가에 따른 경향: NFE 가 증가함에 따라 FID 점수가 지속적으로 개선됨 (기존 Consistency 모델들은 NFE 증가 시 성능이 저하되는 경향이 있음).

3) ImageNet-256 (256x256)

• TFM-XL/2: 676M 파라미터 모델로, 1-step 생성 시 FID 3.02 를 기록 (MeanFlow-XL/2 의 3.43 보다 우세).
• 2-step 생성: FID 2.77 로 기존 최상위 모델들을 능가.
• 시각적 품질: 1-step 생성에서도 구조적 일관성과 디테일이 우수하며, 단계 수가 증가함에 따라 화질이 점진적으로 향상됨.

4) Ablation Study

• 시간 조건부 (Conditioning on t, r): $(t, r)$ 직접 조건부 또는 $(t, \Delta t)$ 조건부 모두 효과적이었으나, 상대적 시간 정보 ($\Delta t$) 가 중요한 역할을 함.
• 샘플링 분포: Logit-normal 분포를 사용한 시간 단계 샘플링이 균일 분포보다 성능이 우수함.
• 손실 함수: 적응적 가중치를 가진 $L_2$ 손실 (Pseudo-Huber loss 유사, $p=0.5$ 또는 $1.0$) 이 표준$L_2$ 손실보다 성능이 좋음.
• CFG (Classifier-Free Guidance): 단일 단계 생성에서도 CFG 적용 시 성능이 크게 향상됨.

5. 의의 및 결론 (Significance)

Transition Flow Matching 은 생성 모델링 분야에서 국소적 접근 (local approach) 에서 전역적 접근 (global approach) 으로의 전환을 의미합니다.

• 학습의 안정성: 사전 학습된 Teacher 모델 없이 처음부터 (from scratch) End-to-End 로 학습 가능하며, 복잡한 커리큘럼 없이도 안정적인 학습이 가능합니다.
• 유연성: 임의의 시간 간격과 단계 수를 자유롭게 조절할 수 있어, 추론 속도와 생성 품질 사이의 균형을 유연하게 맞출 수 있습니다.
• 이론적 통합: 기존 Flow Matching, Mean Velocity, Consistency Model 들을 하나의 통일된 관점에서 설명하며, 생성 역학에 대한 깊은 통찰을 제공합니다.

이 연구는 단일 단계 생성 모델의 성능 한계를 극복하고, 더 빠르고 효율적인 생성 모델 개발을 위한 새로운 표준을 제시한다는 점에서 의의가 큽니다.