조건의 저주: 조건부 흐름 기반 생성을 위한 최적 수송의 분석과 개선

The Curse of Conditions: Analyzing and Improving Optimal Transport for Conditional Flow-Based Generation

March 13, 2025

저자: Ho Kei Cheng, Alexander Schwing

cs.AI

초록

미니배치 최적 수송 커플링은 무조건적 흐름 매칭에서 경로를 직선화합니다. 이는 테스트 시점에 상미분 방정식을 수치적으로 해결할 때 더 적은 수의 적분 단계와 덜 복잡한 수치 솔버를 사용할 수 있게 함으로써 계산적으로 덜 부담스러운 추론을 가능하게 합니다. 그러나 조건부 설정에서는 미니배치 최적 수송이 한계를 보입니다. 이는 기본 최적 수송 매핑이 조건을 무시하기 때문에 훈련 중에 조건적으로 치우친 사전 분포가 생성되기 때문입니다. 반면, 테스트 시점에는 치우친 사전 분포에 접근할 수 없고, 대신 완전하고 편향되지 않은 사전 분포에서 샘플링하게 됩니다. 이러한 훈련과 테스트 간의 차이는 성능 저하로 이어집니다. 이 차이를 해소하기 위해, 우리는 최적 수송 할당을 계산할 때 비용 행렬에 조건부 가중치 항을 추가하는 조건부 최적 수송(C^2OT)을 제안합니다. 실험 결과, 이 간단한 수정이 8gaussians-to-moons, CIFAR-10, ImageNet-32x32, ImageNet-256x256에서 이산적 및 연속적 조건 모두에서 효과적으로 작동함을 보여줍니다. 우리의 방법은 다양한 함수 평가 예산에 걸쳐 기존 베이스라인보다 전반적으로 더 나은 성능을 보입니다. 코드는 https://hkchengrex.github.io/C2OT에서 확인할 수 있습니다.

English

Minibatch optimal transport coupling straightens paths in unconditional flow matching. This leads to computationally less demanding inference as fewer integration steps and less complex numerical solvers can be employed when numerically solving an ordinary differential equation at test time. However, in the conditional setting, minibatch optimal transport falls short. This is because the default optimal transport mapping disregards conditions, resulting in a conditionally skewed prior distribution during training. In contrast, at test time, we have no access to the skewed prior, and instead sample from the full, unbiased prior distribution. This gap between training and testing leads to a subpar performance. To bridge this gap, we propose conditional optimal transport C^2OT that adds a conditional weighting term in the cost matrix when computing the optimal transport assignment. Experiments demonstrate that this simple fix works with both discrete and continuous conditions in 8gaussians-to-moons, CIFAR-10, ImageNet-32x32, and ImageNet-256x256. Our method performs better overall compared to the existing baselines across different function evaluation budgets. Code is available at https://hkchengrex.github.io/C2OT