La Maldición de las Condiciones: Análisis y Mejora del Transporte Óptimo para la Generación Basada en Flujos Condicionales

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

March 13, 2025

Autores: Ho Kei Cheng, Alexander Schwing

cs.AI

Resumen

El acoplamiento de transporte óptimo en minilotes endereza las trayectorias en el emparejamiento de flujo incondicional. Esto conduce a una inferencia computacionalmente menos exigente, ya que se pueden emplear menos pasos de integración y solucionadores numéricos menos complejos al resolver numéricamente una ecuación diferencial ordinaria en el momento de la prueba. Sin embargo, en el escenario condicional, el transporte óptimo en minilotes se queda corto. Esto se debe a que la asignación de transporte óptimo predeterminada ignora las condiciones, lo que resulta en una distribución previa condicionalmente sesgada durante el entrenamiento. En contraste, en el momento de la prueba, no tenemos acceso a la distribución previa sesgada, sino que muestreamos a partir de la distribución previa completa y no sesgada. Esta brecha entre el entrenamiento y la prueba conduce a un rendimiento subóptimo. Para cerrar esta brecha, proponemos el transporte óptimo condicional C^2OT, que añade un término de ponderación condicional en la matriz de costos al calcular la asignación de transporte óptimo. Los experimentos demuestran que esta solución simple funciona tanto con condiciones discretas como continuas en tareas como 8gaussians-to-moons, CIFAR-10, ImageNet-32x32 e ImageNet-256x256. Nuestro método supera en general a las líneas de base existentes en diferentes presupuestos de evaluación de funciones. El código está disponible en 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