One-Sided Quantile Coupling for Flow Matching

기술 요약: Flow Matching을 위한 일방향 분위수 결합 (One-Sided Quantile Coupling)

문제 정의
Flow Matching (FM)은 단순한 소스 분포(일반적으로 표준 가우시안)에서 타겟 데이터 분포로 연결되는 확률 경로의 속도장(velocity field)을 회귀함으로써 연속 시간 생성 모델을 학습시킨다. 이 과정에서 샘플 간의 쌍을 맺는 방식인 결합(coupling)의 품질은 학습 궤적의 기하학적 구조에 결정적인 영향을 미친다. 표준 FM은 독립적인 무작위 쌍 맺기를 사용하는데, 이는 종종 길고 상충하는 경로를 생성하여 타겟 속도장의 분산을 증가시키고 학습 효율을 저해한다.

Optimal Transport (OT)에서 영감을 얻은 구조적 결합은 이러한 경로를 직선화하고 생성 품질을 개선할 수 있지만, 기존 방식(예: OT-CFM)은 미니 배치 할당 문제를 해결하는 데 의존한다. 이러한 접근법은 쌍별 비용 행렬을 구축하고 할당 문제(예: Hungarian 알고리즘 또는 Sinkhorn 반복법)를 해결해야 하므로, 배치 크기에 따라 최소 제곱(quadratic) 또는 세제곱(cubic)으로 증가하는 계산 비용이 발생한다. 이는 확장성을 제한하고 불안정성을 초래할 수 있다. 본 논문은 배치 단위의 운송 문제를 피하면서도 구조적 결합의 기하학적 이점을 얻을 수 있는지에 대한 질문을 다룬다.

방법론: 분위수 결합 Flow Matching (QC-FM)
저자들은 가볍고 일방향적인 결합 방식인 **Quantile Coupling Flow Matching (QC-FM)**을 제안한다. 두 개의 배치(노이즈와 데이터)를 모두 샘플링한 후 재결합하는 전통적인 방식과 달리, QC-FM은 데이터 배치만을 샘생하며 이를 통해 쌍이 된 소스 샘플을 직접 구축한다.

1. 일방향 구축 (One-Sided Construction):

   • 데이터 배치 $\mathcal{B}_x$가 주어지면, 이 방법은 행렬 $U$로 정의된 $k$개의 무작위 직교 방향(slices) 위로 데이터 포인트들을 투영한다.
   • 각 방향에 대해, 배치 내의 데이터 포인트들의 순위(rank)를 매긴다. 이 경험적 순위는 역표준 정규 누적 분포 함수(inverse standard normal CDF)를 사용하여 가우시안 분위수(Gaussian quantiles)로 매핑된다.
   • 소스 잠재 코드 $e^{(i)}$는 이러한 투영된 분위수 제약 조건을 만족하도록 구축된다. 구체적으로, 소스의 투영값은 데이터 순위로부터 유도된 가우시안 분위수와 일치한다.
   • 나머지 차원(슬라이스 부분 공간의 직교 보공간)은 조건부 가우시안 분포로부터 샘플링하여 채운다. 이는 제약되지 않은 방향에서 소스가 확률성을 유지하도록 보장한다.

2. 가우시안 잔여물 하이브리드 확장 (Gaussian-Remainder Hybrid Extensions):
   전체 QC 결합(모든 소스 포인트가 순위를 통해 구축되는 경우)은 슬라이스 코드 간의 의존성으로 인해 $k \ge 2$일 때 표준 i.i.d. 가우시안 사전 분포에서 벗어나기 때문에, 저자들은 QC의 이점을 유지하면서 사전 분포를 보존하기 위해 두 가지 하이브리드 전략을 도입한다.

   • QC-FM-Mixture: 배치의 "앵커(anchor)" 하위 집합(예: 20%)에 QC 구축을 적용하고, 나머지 "잔여(remainder)" 소스 샘플은 남은 데이터 샘플과 독립적인 무작위 쌍 맺기(Baseline 방식)를 사용하여 결합한다. 이는 구조적 결합과 비구조적 결합의 혼합체를 생성한다.
   • QC-FM-Adjacency: 앵커 하위 집합에 QC를 적용하되, 나머지 소스 샘플은 앵커에 의해 유도된 국소 기하 구조를 기반으로 쌍을 맺는다. 잔여 데이터 포인트는 (구조적 부분 공간 내에서) 가장 가까운 앵커에 할당되며, 잔여 소스 샘플은 해당 그룹 내에서 쌍을 맺는다. 이는 타겟을 중복 생성하지 않으면서 구조적 편향을 전파한다.

주요 기여

• QC-FM 알고리즘: 무작위 슬라이스를 따라 단조적인 1차원 운송을 강제하는 새로운 배치 로컬(batch-local) 구조적 결합이다. 이는 쌍별 비용 행렬과 할당 솔버의 필요성을 제거하여, 지배적인 계산 비용을 $O(B^2)$ 또는 $O(B^3)$에서 $O(kB \log B)$ (정렬)로 줄인다.
• 이론적 분석:
  • 분산 감소: 저자들은 고정된 프레임에 대해, QC 결합이 선택된 모든 슬라이스 방향에서 불가피한 회귀 분산을 제거하여, 이상적인 흐름(ideal flow)이 해당 방향에서 정확히 직선이 됨을 증명한다.
  • 사전 분포 보존: $k=1$일 때 소스 주변 분포가 정확히 가우시안임을 정량화한다. $k \ge 2$인 경우, 가우시안 사전 분포로부터의 편차는 슬라이스 코드의 코풀라(copula)에 국한되며, 이 편차에 의한 운송 비용은 유계(bounded)임을 보인다.
• 하이브리드 학습 체계: 기하학적 편향의 주입과 Baseline의 명시적 가우시안 신호 보존 사이의 균형을 맞추는 Mixture 및 Adjacency 변형을 도입한다.

실험 결과
본 방법은 동일한 학습 예산과 아키텍처 하에서 CIFAR-10, CelebA-64, FFHQ-64, ImageNet-64에 대해 평가되었다.

• 성능: QC-FM-Mixture는 네 가지 데이터셋 모두에서 독립 결합 Baseline 및 정확한 Hungarian OT-CFM보다 일관되게 우수한 성능을 보였다. Baseline 대비 개선도는 FID-50K 기준으로 CIFAR-10에서 6.3%, FFHQ에서 12.9%에 달했다.
• 효율성: 계산 비용 측면에서 QC-FM은 OT 기반 방식보다 훨씬 빠르다. 배치 크기 2048에서 QC-FM-Mixture는 Hungarian OT 솔버보다 800배 이상, Sinkhorn 솔버보다 35배 이상 빨랐다.
• 토이 진단 (Toy Diagnostics): 합성 데이터셋(Checkerboard 및 8-Gaussians)에서 QC-FM 변형들은 Baseline에 비해 거의 직선에 가까운 학습 궤적(Path-Length Ratio $\approx 1$)과 현저히 낮은 타겟 속도 분산을 생성하여, 결합의 기하학적 이점을 확인시켜 주었다.

의의 및 주장
본 논문은 투영된 순위 구조를 보존하는 것이 배치 단위의 운송 문제를 풀지 않고도 Flow Matching 결합에 유용한 기하학적 편향을 주입하는 간단하고 확장 가능하며 효과적인 방법임을 주장한다.

• 확장성: QC-FM은 구조 없는 Baseline과 비용이 많이 드는 배치 단위 운송 방식 사이의 매력적인 중간 지점을 제공한다. 이는 계산적으로 가볍고 투영 횟수($k$)를 통해 조절 가능하다.
• 실용성: 추론 시 별도의 학습된 구성 요소를 필요로 하지 않는다. 즉, 생성은 여전히 표준 가우시안에서 시작된다. 결합은 학습 중에 즉석에서 구축된다.
• 한계점: 저자들은 QC-FM이 구조적 결합을 위한 배치 로컬(batch-local) 대리 모델이지, 전역적 Optimal Transport의 근사치가 아님을 겸허히 밝힌다. 분위수는 미니 배치 내에서 추정되며, 이론은 전역적인 궤적의 직선성이나 특정 FID 개선을 보장하기보다는 슬라이스당 분산을 제어한다. 슬라이스 개수($k$)와 앵커 비율($p$)의 선택은 현재 휴리스틱과 데이터셋별 튜닝에 의존하고 있다.

결론적으로, 본 연구는 순위 보존을 통해 노이즈-데이터 쌍을 구조화함으로써, 한 방향의 분위수 결합이 Flow Matching의 샘플 품질과 학습 효율을 개선할 수 있음을 보여주며, 계산적으로 부담이 큰 Optimal Transport 방식에 대한 실용적인 대안을 제시한다.