기술 요약: 조건 의존적 소스 플로우 매칭 (Condition-Dependent Source Flow Matching, CSFM)

문제 정의

플로우 매칭(Flow Matching, FM)은 텍스트-이미지(T2I) 생성에서 확산 모델(Diffusion Models)의 강력한 대안으로 부상했으며, 결정론적 역학(deterministic dynamics)과 소스 분포 선택의 유연성을 제공한다. 그러나 이론적으로는 어떤 소스 분포도 선택할 수 있는 자유도가 있음에도 불구하고, 대부분의 기존 FM 시스템은 단순함을 위해 확산 모델에서 물려받은 고정된 표준 가우시안 소스($N(0, I)$)를 기본값으로 사용한다. 이러한 선택은 최적화보다는 관습에 의한 것이다.

본 논문은 다음과 같은 중요한 격차를 식별한다:

1. 왜(Why) 소스 분포를 학습하는 것이 고차원 및 강한 조건부 환경에서 FM 목적 함수를 개선하는가?
2. 어떻게(How) 현대적인 T2I 아키텍처에서 훈련 붕괴나 불안정성을 일으키지 않고 안정적으로 소스 학습을 구현할 것인가?
3. 언제(When) 소스 학습이 특히 타겟 표현 공간의 기하학적 구조와 관련하여 가장 유의미한 성능 향상을 가져오는가?

최적 운송(Optimal Transport)과 같이 소스-타겟 결합(coupling)을 수정하는 기존 방식들은 고차원 조건부 설정에서 어려움을 겪으며, 다른 학습된 소스 접근법들은 훈련 역학 및 정규화에 대한 가이드가 부족하다.

방법론: 조건 의존적 소스 플로우 매칭 (CSFM)

저자들은 플로우 매칭 목적 함수 하에서 조건 의서적 소스 분포 $p_\phi(X_0 | C)$와 속도장(velocity field) $v_\theta$를 공동으로 학습하는 CSFM 프레임워크를 제안한다.

1. 적응형 소스-타겟 결합 (Adaptive Source-Target Coupling)

고정된 조건 불가지론적 결합 $\pi(X_0, X_1, C) = p_0(X_0)p_1(X_1, C)$ 대신, CSFM은 학습 가능한 결합 $\pi_\phi(X_0, X_1, C) = p_\phi(X_0 | C)p_1(X_1, C)$를 도입한다. 소스 생성기 $g_\phi$는 조건 $C$를 조건부 가우시안 소스의 파라미터로 매핑한다:
$$p_\phi(X_0 | C) = \mathcal{N}(\mu_\phi(C), \text{diag}(\sigma^2_\phi(C)))$$
이를 통해 소스가 조건에 따라 위치와 지지 집합(support)을 적응할 수 있게 하며, 이론적으로 FM 손실의 내재적 분산(intrinsic variance) 항, 즉 속도 감독의 모호성을 특징짓는 항을 줄여준다.

2. 안정화 메커니즘 (Stabilization Mechanisms)

저자들은 공동 훈련의 두 가지 주요 장애물인 소스 붕괴(분산이 0으로 수렴하는 현상)와 밀접하게 통합된 아키텍처에서의 약한 감독 신호를 식별하였다. 이를 해결하기 위해 두 가지 특정 구성 요소를 제안한다:

• 분산 전용 정규화 ($\mathcal{L}_{VarReg}$): 표준 정규 사전 확률(prior)을 향한 표준 KL 정규화는 평균과 분산을 모두 제약하여 소스 평균 $\mu_\phi(C)$를 원점으로 강제하고, 이로 인해 조건과 호환되는 타겟 영역으로 소스를 이동시키는 능력을 제한한다. CSFM은 $\sigma^2_\phi(C)$가 단위 분산에서 벗어나는 것을 벌칙을 주되 $\mu_\phi(C)$는 제약하지 않는 분산 전용 정규화를 사용하여, 소스의 이동성을 허용하면서도 붕괴를 방지한다.
• 방향성 소스-타겟 정렬 ($\mathcal{L}_{align}$): 복잡한 T2I 설정에서 조건으로 속도장을 조절하는 플로우 모델의 능력이 소스 생성기에 대한 학습 신호를 약화시킬 수 있다. 이를 상쇄하기 위해, 저자들은 학습된 소스 샘플 $X_0$와 타겟 $X_1$ 사이의 방향성 정렬을 장려하는 음의 코사인 유사도 손실을 도입한다. 이는 소스의 크기나 확률적 지지 집합을 과도하게 제한하지 않으면서 구체적인 가이드를 제공한다.

전체 훈련 목적 함수는 다음과 같다:
$$\mathcal{L}_{total} = \mathcal{L}_{FM} + \lambda_{VarReg}\mathcal{L}_{VarReg} + \lambda_{align}\mathcal{L}_{align}$$

주요 기여

1. 이론적 통찰 (Why): 본 논문은 소스 학습을 FM 목적 함수의 내재적 분산 항을 줄이는 것과 연결한다. 조건 의존적 소스를 학습함으로써 경로 교차와 속도 모호성을 줄여, 더 깨끗한 감독, 낮은 그래디언트 분산, 그리고 더 곧은 운송 경로를 이끌어낸다.
2. 실용적 프레임워크 (How): 분산 전용 정규화와 방향성 정렬을 포함하는 안정적인 훈련 레시피를 제안하여, 소스의 이동성을 제한하거나 붕괴를 일으키지 않고 엔드 투 엔드 훈련을 가능하게 한다.
3. 조건부 유효성 (When): 저자들은 소스 학습의 이점이 타겟 표현 공간의 기하학적 구조에 따라 달라짐을 입증한다. 조건이 컴팩트한 클러스터를 형성하는 구조화된 잠재 공간(예: RAE with DINOv2)에서 이득이 가장 뚜렷하며, 조건이 매우 다봉형 분포(multimodal distribution)에 대응하는 얽힌 공간(예: SD-VAE)에서는 이득이 감소한다.

실험 결과

저자들은 여러 벤치마크, 아키텍처 및 규모에 걸쳐 CSFM을 평가하였다:

• 수렴 속도: ImageNet-1K (256×256)에서 CSFM은 고정된 가우시안 소스를 사용하는 표준 FM 대비 최대 3.01배 빠른 FID 수렴과 2.48배 빠른 CLIP Score 수렴을 달성했다.
• 성능 지표: UnifiedNextDiT 백본과 RAE (DINOv2) 표현을 사용하여, CSFM은 FID 2.453, CLIP Score 0.3420을 기록하여 가우시안 베이스라인(FID 3.036) 및 다른 조건 의존적 소스/결합 방법(예: CrossFlow, C2OT, CPD)보다 우수한 성능을 보였다.
• 플로우 기하학: 소수 단계 생성(few-step generation) 실험에서 CSFM은 샘플링 단계가 감소함에 따라(예: 50단계에서 3단계로) 더 완만하게 성능이 저하되었으며, 이는 더 곧은 플로우 필드와 적은 경로 교차를 나타낸다.
• 강건성: 이 방법은 다양한 조건부 아키텍처(LightningDiT, MMDiT, UnifiedNextDiT)와 텍스트 인코더(CLIP, Qwen3) 전반에서 일관된 개선을 보여주었다.
• 확장성: 1.3B 파라미터 모델로 GenEval 및 DPG-Bench에서 확장했을 때, CSFM은 가우시안 베이스라인보다 우수한 성능을 유지하며 GenEval에서 0.80, DPG-Bench에서 81.11을 달성했다.

의의 및 주장

본 논문은 소스 분포를 학습 가능한 조건 의존적 설계 요소로 취급하는 것이 현대적 플로우 매칭의 핵심적이면서도 아직 충분히 활용되지 못한 부분이라고 주장한다. 분산 전용 정규화와 방향성 정렬을 통해 고차원 조건부 생성의 구체적인 과제를 해결함으로써, CSFM은 훈련 효율성과 생성 품질을 개선하는 원칙적인 방법을 제공한다.

저자들은 이 접근 방식의 효과가 보편적인 것이 아니라 타겟 공간의 표현 능력과 밀접하게 연관되어 있음을 강조한다. 조건 신호가 명확한 의미론적 클러스터를 생성하는 구조화된 잠재 공간에서 소스를 학습하는 것은 상당한 이득을 가져온다. 본 연구는 향 만큼이나 미래의 생성 모델 발전이 확산 모델에서 물려받은 고정된 사전 확률에 의존하기보다, 소스와 타겟 표현의 공동 최적화를 고려해야 함을 시사한다.