COT-FM: 클러스터별 최적 수송 흐름 매칭 (Cluster-wise Optimal Transport Flow Matching) 기술 요약

이 논문은 Flow Matching (FM) 모델의 생성 속도와 품질을 동시에 향상시키기 위한 새로운 프레임워크인 COT-FM (Cluster-wise Optimal Transport Flow Matching) 을 제안합니다. 기존 FM 모델이 겪는 비선형적이고 구부러진 흐름 (curved trajectories) 문제를 해결하여, 더 적은 샘플링 단계로 고품질 데이터를 생성할 수 있도록 합니다.

---

1. 문제 정의 (Problem)

Flow Matching 은 간단한 소스 분포 (예: 가우시안) 를 복잡한 데이터 분포로 변환하는 결정론적 벡터 필드를 학습하는 프레임워크입니다. 그러나 실제 적용 시 다음과 같은 한계가 존재합니다.

• 구부러진 경로 (Curved Trajectories): 기존 FM 모델은 무작위 커플링 (Random Coupling) 또는 배치 단위 최적 수송 (Batch-wise OT) 을 사용하여 벡터 필드를 학습합니다. 이는 서로 다른 샘플 쌍 간의 경로가 교차하거나 평균화되어 전체적으로 구부러진 흐름을 생성하게 합니다.
• 이산화 오차 (Discretization Error): 구부러진 경로는 수치적 통합 (ODE integration) 시 오차를 증가시켜, 생성 단계에서 많은 단계 (NFE) 가 필요하거나 저품질의 샘플을 생성하게 만듭니다.
• 기존 해결책의 한계:
  • Rectified Flow: 생성된 샘플과 소스 간의 커플링을 반복적으로 최적화하지만, 모델 붕괴 (Model Collapse) 를 유발하여 품질이 저하될 수 있습니다.
  • MeanFlow 등: 평균 벡터 필드를 학습하여 단계를 건너뛰지만, 근본적인 벡터 필드의 곡률을 줄이지 않아 생성 품질 향상에는 한계가 있습니다.

2. 방법론 (Methodology)

COT-FM 은 "분할 정복 (Divide-and-Conquer)" 전략을 기반으로 합니다. 전체 데이터 분포를 하나의 소스에서 매핑하는 대신, 데이터를 클러스터로 나누고 각 클러스터에 전용 소스 분포를 할당하여 국소적인 최적 수송 (Optimal Transport, OT) 을 수행합니다.

핵심 단계

1. 클러스터별 소스 분포 식별 (Reverse ODE):

   • 사전 학습된 FM 모델을 사용하여 각 데이터 클러스터의 샘플들을 역방향으로 ODE 적분 (Reverse ODE) 하여 소스 공간 (Noise space) 으로 되돌립니다.
   • 이렇게 얻은 소스 샘플들의 통계량 (평균 $\mu$, 공분산 $\Sigma$) 을 계산하여 각 클러스터 $k$에 대한 전용 가우시안 소스 분포 $p_{0,k}$를 구성합니다.
   • 이 과정은 서로 다른 클러스터 간의 경로 교차를 최소화하고, 국소적인 OT 문제를 더 효율적으로 만듭니다.

2. 클러스터 내 최적 수송 (Cluster-wise OT):

   • 각 클러스터 $C_k$와 해당 클러스터의 소스 분포 $p_{0,k}$ 사이에 배치 단위 OT 맵을 근사적으로 계산합니다.
   • 전체 데이터에 대한 전역 OT 는 계산 비용이 너무 크지만, 클러스터로 나누면 샘플 수가 줄어들어 OT 근사가 훨씬 효과적이고 정확해집니다.

3. 교차 최적화 (Alternating Optimization):

   • Stage 1: 클러스터별 OT 맵을 기반으로 타겟 벡터 필드를 구성합니다.
   • Stage 2: 구성된 벡터 필드를 FM 모델이 회귀 (Regress) 하도록 학습합니다.
   • 이 두 단계를 반복하여 모델과 소스 분포를 함께 정제합니다.

4. 샘플링 (Inference):

   • 생성 시, 먼저 클러스터 인덱스 $k$를 확률에 따라 샘플링한 후, 해당 클러스터의 소스 분포 $p_{0,k}$에서 노이즈를 추출하여 FM 모델을 통해 데이터를 생성합니다.

3. 주요 기여 (Key Contributions)

• 플러그 앤 플레이 (Plug-and-Play) 프레임워크: 모델 아키텍처나 입력/출력 메커니즘을 변경하지 않고도 기존 FM 모델 (Rectified Flow, MeanFlow 등) 에 적용 가능합니다.
• 직선화 된 흐름 (Straight Flows): 클러스터별 OT 를 통해 경로 교차를 줄이고 벡터 필드를 직선화하여, **적은 샘플링 단계 (Low NFE)**에서도 고품질 생성이 가능합니다.
• 범용성: 무조건적 생성 (Unconditional), 조건부 생성 (Conditional, 예: 텍스트, 클래스 레이블), 로봇 조작 (Robotic Manipulation) 등 다양한 도메인에서 유효함을 입증했습니다.

4. 실험 결과 (Results)

COT-FM 은 2D 합성 데이터, 이미지 생성 (CIFAR-10, ImageNet), 로봇 제어 (LIBERO) 에서 기존 최첨단 방법들을 압도적으로 능가했습니다.

• 2D 데이터 (Mixture of 5-Gaussians, Two Moons):
  • Wasserstein 거리와 곡률 (Curvature) 모두에서 가장 낮은 값을 기록했습니다.
  • 예: 5-Gaussians 에서 Wasserstein 거리를 0.5421 (Rectified Flow) 에서 0.1995로 감소.
• 이미지 생성 (CIFAR-10):
  • 1-step 생성: FID 가 378.0 (Rectified Flow) 에서 205.0으로 대폭 개선.
  • MeanFlow 기반: 1-step FID 를 2.92 에서 2.60으로 개선.
  • 50-step 생성: FID 를 4.45 에서 3.97로 개선.
• 이미지 생성 (ImageNet 256x256):
  • SiT 아키텍처 기반에서 모든 NFE 단계에서 FID 를 개선 (예: 100-step 기준 5.82 → 5.11).
• 로봇 조작 (LIBERO):
  • 1-step 샘플링으로 기존 4-step 모델 (FLOWER) 과 유사하거나 더 높은 성공률 (Spatial: 96.1%, Long: 94.5%) 을 달성했습니다.

5. 의의 및 결론 (Significance)

COT-FM 은 Flow Matching 의 근본적인 한계인 비선형적인 흐름을 데이터의 구조적 특성 (클러스터) 을 활용하여 해결했습니다.

• 효율성: 생성 속도를 획기적으로 높임 (단일 단계 또는 소수 단계로 고품질 생성 가능).
• 품질: 경로 교차를 줄여 왜곡 없는 데이터 생성을 가능하게 함.
• 확장성: 복잡한 고차원 데이터 (이미지, 로봇 행동) 에서도 효과적으로 작동하며, 조건부 생성 및 로봇 정책 학습 등 다양한 응용 분야에 적용 가능.

결론적으로, COT-FM 은 생성 모델의 속도와 품질의 트레이드오프를 극복하고, 직선적이고 구조를 보존하는 수송 경로를 학습함으로써 차세대 생성 모델의 새로운 표준을 제시합니다.