COT-FM: 군집별 최적 이동 플로우 매칭

기존 생성 AI는 노이즈에서 이미지를 만들 때 경로가 복잡하게 꼬여 있어 여러 번의 반복 계산이 필요했다. 이 기술은 데이터를 비슷한 그룹끼리 묶어 직선 경로를 찾아줌으로써, 계산 횟수를 획기적으로 줄이면서도 선명한 결과를 얻게 해준다.

Flow Matching은 노이즈(소스)에서 데이터(타겟)로 가는 벡터장을 학습한다. 기존 방식은 무작위로 노이즈와 데이터를 짝지어 학습하므로, 경로가 서로 엉키고 구부러지는 Curvature 문제가 발생한다. 이는 수치 해석 시 오차를 키워 많은 샘플링 단계(NFE)를 요구하게 만든다. COT-FM은 비슷한 데이터는 비슷한 노이즈에서 출발해야 한다는 직관에 기반한다. 먼저 데이터를 군집화하고, 사전 학습된 모델을 거꾸로 돌려 각 군집이 어떤 노이즈 영역에서 왔는지 파악한다. 이렇게 각 군집에 전용 노이즈 영역을 배정하면 경로가 엉키지 않는 Divide-and-Conquer가 가능해진다. 결과적으로 전체 데이터 간의 복잡한 매핑 대신 군집 내에서의 단순한 직선 매핑을 학습하게 된다. 이는 벡터장을 극도로 단순화하여, 단 1~2단계의 짧은 계산만으로도 기존의 수십 단계 계산과 대등하거나 더 나은 품질의 결과물을 만들어낸다.

전체 프로세스는 2단계로 구성된다. 1단계에서는 타겟 데이터를 K-Means 등으로 군집화하고, 사전 학습된 FM 모델의 ODE를 역방향으로 적분하여 각 데이터 포인트에 대응하는 소스 샘플 $\hat{x}_0$를 추출한다. $\hat{x}_0 = x_1 - \int_0^1 v_\theta(\hat{x}_t, t) dt$ 수식을 계산한다. [타겟 데이터 $x_1$을 입력으로] -> [사전 학습된 모델 $v_\theta$를 이용해 역방향으로 적분하는 연산을 수행해] -> [노이즈 공간의 샘플 $\hat{x}_0$를 얻고] -> [이 값이 해당 데이터가 생성되기 시작한 원래의 위치임을 의미한다.] 추출된 $\hat{x}_0$ 샘플들의 평균($\mu_{0,k}$)과 공분산($\Sigma_{0,k}$)을 계산하여 각 군집 $C_k$에 대응하는 가우시안 소스 분포 $p_{0,k}$를 정의한다. $\mu_{0,k} = \frac{1}{|X_{0,k}|} \sum \hat{x}_0$를 계산한다. [군집 내 모든 역방향 소스 샘플 $\hat{x}_0$를 입력으로] -> [산술 평균을 구하는 연산을 수행해] -> [군집의 중심점 $\mu_{0,k}$를 얻고] -> [이 값이 해당 군집이 시작되는 노이즈 분포의 평균임을 의미한다.] 2단계에서는 확정된 군집별 소스 분포와 타겟 데이터 사이의 최적 이동(Optimal Transport) 맵 $\pi_k$를 계산한다. 이후 FM 모델 $v_\theta$가 이 직선 경로를 모사하도록 학습한다. [소스-타겟 쌍 입력 -> 선형 보간 $x_t = (1-t)x_0 + tx_1$ 계산 -> 속도 벡터 $x_1 - x_0$ 도출 -> 모델 예측값과의 MSE 손실 최소화]를 통해 벡터장을 직선화한다.

2D 합성 데이터셋(Mixture of 5-Gaussians, Two Moons 등)에서 기존 Rectified Flow 대비 Wasserstein 거리를 대폭 줄였으며, 경로의 곡률(Curvature)을 최소화하여 가장 직선에 가까운 플로우를 형성함을 확인했다. CIFAR-10 이미지 생성 실험에서 10단계 샘플링 시 FID 8.23을 기록하여 Rectified Flow(12.6)보다 우수한 성능을 보였다. 특히 1단계(1-step) 생성에서 FID 205.0을 기록하며 기존 모델(378.0) 대비 압도적인 품질 향상을 증명했다. ImageNet 256x256 고해상도 조건부 생성에서도 SiT-B/2 아키텍처 기준 1단계 FID를 264.36에서 231.99로 개선하며 확장성을 입증했다. 로봇 조작 작업(LIBERO)에서는 단 1회의 모델 실행만으로도 기존 4회 실행 모델(FLOWER)과 대등한 96.1%의 성공률을 달성했다.

COT-FM은 Flow Matching의 확률 경로(Probability Path)를 군집 단위로 분해하여 재설계한다. 기존의 전역적 최적 이동(Global OT)은 계산 복잡도가 $O(n^3)$으로 대규모 데이터셋에 부적합하지만, COT-FM은 이를 군집 내 국소적 문제로 치환하여 효율성을 확보했다. 핵심은 사전 학습된 모델을 활용한 Bootstrap 전략이다. 무작위 커플링으로 학습된 모델이라도 그 플로우는 가역적(Reversible)이고 비교차적(Non-intersecting)인 특성을 가지므로, 이를 역이용해 타겟 데이터의 의미론적 구조를 소스 노이즈 공간으로 전이시킬 수 있다. 조건부 생성 환경에서는 클래스 레이블이나 텍스트 임베딩을 자연스러운 군집 식별자로 활용한다. 비고정 군집 상황을 위해 강화학습 기반의 PPO 알고리즘을 사용하여 최적의 소스 분포를 동적으로 예측하는 모듈을 추가 설계했다.

데이터셋의 크기가 커질수록 역방향 ODE 계산 단계가 비효율적일 수 있으며, 성능이 군집화의 품질에 의존적이라는 한계가 있다. 또한 군집 간 이동이 전체적인 구조를 얼마나 완벽하게 포착하는지에 대한 이론적 분석이 추가로 필요하다.