COT-FM: Cluster-wise Optimal Transport Flow Matching

Resumen Técnico: COT-FM (Cluster-wise Optimal Transport Flow Matching)

1. El Problema

Los modelos de Flow Matching (FM) son un marco generativo eficiente que aprende un campo vectorial para transformar una distribución de origen simple (como un ruido gaussiano) en una distribución de datos compleja. Sin embargo, enfrentan dos desafíos principales que limitan su calidad y velocidad:

• Trayectorias Curvas: Los métodos estándar de FM a menudo utilizan acoplamientos aleatorios o acoplamientos óptimos por lotes (batch-wise) para construir el campo vectorial. Esto resulta en trayectorias de transporte curvas y con intersecciones frecuentes. Las trayectorias curvas aumentan el error de discretización temporal, obligando a usar muchos pasos de integración (NFE) para obtener muestras de alta calidad.
• Compromiso entre Velocidad y Calidad: Los métodos existentes para acelerar la generación (como MeanFlow o modelos de distilación) a menudo ignoran la curvatura subyacente del campo vectorial, simplemente "saltando" pasos. Esto reduce el número de pasos pero no mejora la fidelidad de la generación, ya que el campo vectorial sigue siendo inexacto.
• Ineficiencia del Transporte Óptimo Global: Calcular el mapa de Transporte Óptimo (OT) exacto entre distribuciones completas es computacionalmente prohibitivo ($O(n^3)$) para conjuntos de datos grandes, y las aproximaciones por lotes fallan en capturar la estructura local, manteniendo las trayectorias curvas.

2. Metodología: COT-FM

El artículo propone COT-FM, un marco general y "plug-and-play" que reestructura la ruta de probabilidad del FM mediante una estrategia de dividir y conquistar. La idea central es dividir la distribución de datos en clústeres y asignar a cada uno su propia distribución de origen, en lugar de mapear toda la distribución de datos desde una única fuente global.

El proceso se divide en dos etapas principales que se alternan:

A. Construcción de Distribuciones de Origen por Clúster (Stage 1):

1. Clustering: Se dividen las muestras de datos objetivo en $K$ clústeres (usando etiquetas de clase, descripciones textuales o clustering no supervisado como K-Means en características DINO).
2. Inversión de ODE: Utilizando un modelo FM preentrenado (que tiene trayectorias reversibles y no intersectantes), se invierte el proceso de integración de la EDO (Ecuación Diferencial Ordinaria) para cada muestra de datos $x_1$ de un clúster. Esto permite recuperar la muestra de origen estimada $\hat{x}_0$.
3. Estimación de Distribución: Para cada clúster $C_k$, se calculan la media ($\mu_{0,k}$) y la covarianza ($\Sigma_{0,k}$) de las muestras de origen recuperadas. Esto define una distribución de origen específica para ese clúster: $p_{0,k} \sim \mathcal{N}(\mu_{0,k}, \Sigma_{0,k})$.

B. Aproximación de Transporte Óptimo Local y Entrenamiento (Stage 2):

1. OT por Clúster: En lugar de calcular un OT global, se calcula un mapa de Transporte Óptimo aproximado (usando acoplamientos por lotes) dentro de cada clúster entre su distribución de origen estimada $p_{0,k}$ y sus datos objetivo $C_k$.
2. Reducción de Curvatura: Al restringir el transporte a subespacios locales más homogéneos, el acoplamiento óptimo es mucho más preciso, resultando en trayectorias casi rectas y sin intersecciones dentro de cada clúster.
3. Entrenamiento Alternante: El modelo FM se entrena para regredir este nuevo campo vectorial compuesto por los OT locales. El proceso alterna entre refinar las distribuciones de origen (vía inversión de ODE) y optimizar el modelo FM.

Generación: Durante la inferencia, se selecciona un índice de clúster $k$ (proporcional al tamaño del clúster), se muestrea ruido de la distribución específica $p_{0,k}$ y se integra a través del modelo FM afinado.

3. Contribuciones Clave

• Marco General y Plug-and-Play: COT-FM no requiere cambiar la arquitectura del modelo FM ni sus mecanismos de entrada/salida. Se puede aplicar a cualquier modelo FM existente (Rectified Flow, OT-CFM, MeanFlow).
• Rectificación de Trayectorias sin Colapso: A diferencia de Rectified Flow iterativo que entrena sobre muestras generadas (riesgo de colapso de modelo), COT-FM utiliza la inversión de un modelo preentrenado para obtener fuentes, manteniendo la estabilidad.
• Eficiencia Computacional: Al reducir el problema de OT global a múltiples problemas locales más pequeños, hace factible el cálculo de OT aproximado con mayor precisión y menor costo computacional efectivo por paso de entrenamiento.
• Mejora en Escenarios de Pocos Pasos: El método está diseñado específicamente para permitir una generación de alta calidad con un número extremadamente bajo de pasos de evaluación de la función (NFE), incluso en 1 paso.

4. Resultados Experimentales

Los autores evaluaron COT-FM en múltiples dominios, demostrando mejoras consistentes en velocidad y calidad:

• Datos Sintéticos 2D: En distribuciones como "5-Gaussians", "Two Moons" y "Checkerboard", COT-FM logró la menor distancia de Wasserstein y la menor curvatura de trayectorias en comparación con Rectified Flow, OT-CFM y MeanFlow.
• Generación de Imágenes (CIFAR-10):
  • Mejoró el FID (Fréchet Inception Distance) de 12.6 a 8.23 en 10 pasos y de 4.45 a 3.97 en 50 pasos sobre Rectified Flow.
  • En generación de 1 paso, redujo el FID de 378.0 a 205.0 (Rectified Flow) y de 2.92 a 2.60 (MeanFlow).
• Generación Condicional (ImageNet 256x256): Utilizando arquitecturas SiT, COT-FM mejoró el FID en todos los escenarios de pasos, con mejoras más pronunciadas en regímenes de bajo NFE (ej. 2 pasos).
• Manipulación Robótica (LIBERO): En tareas de robot controladas por texto, COT-FM alcanzó una tasa de éxito del 96.1% (Spatial) y 94.5% (Long) con solo 1 paso (1 NFE). Esto supera a la política FLOWER, que requiere 4 pasos para alcanzar un rendimiento similar (97.1% y 93.5%).

5. Significado e Impacto

COT-FM representa un avance significativo en la teoría de modelos generativos basados en flujo al demostrar que la curvatura del campo vectorial es el factor limitante principal para la generación rápida y de alta calidad.

• Paradigma de "Dividir y Conquistar": Introduce una nueva forma de abordar el transporte óptimo en aprendizaje profundo, mostrando que la estructura local de los datos puede explotarse para simplificar problemas globales complejos.
• Viabilidad de Generación en 1 Paso: Hace que la generación de imágenes y acciones robóticas de alta fidelidad en un solo paso sea una realidad práctica, eliminando la necesidad de costosas inferencias iterativas.
• Aplicabilidad Amplia: Su naturaleza agnóstica al modelo permite su integración inmediata en pipelines de investigación y producción existentes, mejorando tanto la eficiencia computacional como la calidad de los resultados en visión por computadora y robótica.

En resumen, COT-FM resuelve el dilema entre velocidad y calidad en los modelos Flow Matching al crear trayectorias de transporte más rectas y precisas mediante la segmentación inteligente de los datos y el uso de transporte óptimo local.