One-Sided Quantile Coupling for Flow Matching

Resumen Técnico: Acoplamiento de Cuantiles Unilateral para Flow Matching

Planteamiento del Problema
El Flow Matching (FM) entrena modelos generativos de tiempo continuo mediante la regresión del campo de velocidad de una trayectoria de probabilidad que conecta una distribución de origen simple (típicamente una Gaussiana estándar) con una distribución de datos objetivo. La calidad del acoplamiento —el emparejamiento de muestras de origen y destino— influye críticamente en la geometría de las trayectorias de entrenamiento. El FM estándar utiliza un emparejamiento aleatorio independiente, lo que a menudo resulta en trayectorias largas y conflictivas que aumentan la varianza del campo de velocidad objetivo y dificultan la eficiencia del entrenamiento.

Si bien existen acoplamientos estructurados inspirados en el Transporte Óptimo (OT) que pueden enderezar estas trayectorias y mejorar la calidad de la generación, los métodos existentes (por ejemplo, OT-CFM) dependen de la resolución de un problema de asignación por mini-lote (mini-batch), lo que requiere la construcción de matrices de costo por pares y la resolución de problemas de asignación (por ejemplo, mediante el algoritmo de Hungarian o iteraciones de Sinkhorn), lo que conlleva costos computacionales que escalan al menos de forma cuadrática (o cúbica para resoluciones exactas) con el tamaño del lote. Esto limita la escalabilidad y puede introducir inestabilidad. El artículo aborda la cuestión de si los beneficios geométricos del acoplamiento estructurado pueden lograrse sin emparejar dos lotes pre-muestreados, evitando así el alto costo del transporte por lote.

Metodología: Acoplamiento de Cuantiles para Flow Matching (QC-FM)
Los autores proponen el Acoplamiento de Cuantiles para Flow Matching (QC-FM), un esquema de acoplamiento unilateral ligero. A diferencia de los métodos tradicionales que muestrean tanto el lote de ruido como el de datos y luego los re-emparejan, el QC-FM solo muestrea el lote de datos y construye directamente las muestras de origen emparejadas.

1. Construcción Unilateral:

   • Dado un lote de datos $\mathcal{B}_x$, el método proyecta los puntos de datos sobre $k$ direcciones ortogonales aleatorias (cortes o slices) definidas por una matriz $U$.
   • Para cada dirección, los puntos de datos se clasifican dentro del lote. Estos rangos empíricos se mapean a cuantiles gaussianos utilizando la función de distribución acumulativa inversa de la normal estándar.
   • El código latente de origen $e^{(i)}$ se construye para satisfacer estas restricciones de cuantiles proyectados. Específicamente, la proyección del origen sobre las direcciones de corte coincide con los cuantiles gaussianos derivados de los rangos de los datos.
   • Las dimensiones restantes (el complemento ortogonal del subespacio de corte) se completan mediante el muestreo de una distribución Gaussiana condicional. Esto asegura que el origen permanezca estocástico en las direcciones no restringidas.

2. Extensiones Híbridas de Gaussiano-Remanente:
   Dado que el acoplamiento QC completo (donde todos los puntos de origen se construyen vía rangos) se desvía de la distribución previa Gaussiana i.i.d. estándar para $k \ge 2$ debido a las dependencias entre los códigos de corte, los autores introducen dos estrategias híbridas para retener los beneficios de QC preservando la distribución previa:

   • QC-FM-Mixture (Mezcla): Aplica la construcción QC a un subconjunto "ancla" del lote (por ejemplo, el 20%) y empareja el resto de las muestras de origen "remanentes" con las muestras de datos restantes utilizando un emparejamiento aleatorio independiente (el enfoque Baseline). Esto crea una mezcla de acoplamientos estructurados y no estructurados.
   • QC-FM-Adjacency (Adyacencia): Aplica QC al subconjunto ancla pero empareja las muestras de origen remanentes basándose en la geometría local inducida por las anclas. Los puntos de datos remanentes se asignan a su ancla más cercana (en el subespacio estructurado), y las muestras de origen remanentes se emparejan dentro de estos grupos. Esto propaga el sesgo estructurado sin duplicar los objetivos.

Contribuciones Clave

• Algoritmo QC-FM: Un acoplamiento novedoso, local al lote, que impone un transporte monótono unidimensional a lo largo de cortes aleatorios. Elimina la necesidad de matrices de costo por pares y solucionadores de asignación, reduciendo el costo computacional dominante a $O(kB \log B)$ (ordenamiento) en lugar de $O(B^2)$ u $O(B^3)$.
• Análisis Teórico:
  • Reducción de Varianza: Los autores demuestran que, para un marco fijo, el acoplamiento QC elimina la varianza de regresión irreducible a lo largo de cada corte seleccionado, haciendo que el flujo ideal sea exactamente recto en esas direcciones.
  • Preservación de la Distribución Previa: Cuantifican que para $k=1$, la marginal de origen es exactamente Gaussiana. Para $k \ge 2$, la desviación de la distribución previa Gaussiana se confina a la cópula de los códigos de corte, y el costo de transporte de esta desviación está acotado.
• Esquemas de Entrenamiento Híbridos: La introducción de las variantes de Mezcla y Adyacencia que equilibran la inyección de sesgo geométrico con la preservación de la señal Gaussiana explícita del Baseline.

Resultados Experimentales
El método fue evaluado en CIFAR-10, CelebA-64, FFHQ-64 e ImageNet-64 bajo presupuestos de entrenamiento y arquitecturas equivalentes.

• Rendimiento: QC-FM-Mixture superó consistentemente al Baseline de acoplamiento independiente y al OT-CFM de Hungarian exacto en los cuatro conjuntos de datos. Las mejoras sobre el Baseline oscilaron entre el 6.3% (CIFAR-10) y el 12.9% (FFHQ) en términos de FID-50K.
• Eficiencia: En términos de costo computacional, QC-FM es significativamente más rápido que los métodos basados en OT. Con un tamaño de lote de 2048, QC-FM-Mixture fue más de 800 veces más rápido que el solucionador de Hungarian OT y 35 veces más rápido que el solucionador de Sinkhorn.
• Diagnósticos de Juguete: En conjuntos de datos sintéticos (Checkerboard y 8-Gaussians), las variantes de QC-FM produjeron trayectorias aprendidas casi rectas (Relación de Longitud de Trayectoria $\approx 1$) y una varianza de la velocidad objetivo significativamente menor en comparación con el Baseline, confirmando los beneficios geométricos del acoplamiento.

Significancia y Reivindicaciones
El artículo sostiene que preservar la estructura de rango proyectada es una forma simple, escalable y efectiva de inyectar un sesgo geométrico útil en los acoplamientos de Flow Matching sin resolver un problema de transporte por lote.

• Escalabilidad: QC-FM ofrece un punto medio atractivo entre el Baseline no estructurado y los costosos métodos de transporte por lote. Es computacionalmente ligero y controlable mediante el número de proyecciones ($k$).
• Practicidad: El método no requiere componentes aprendidos para la inferencia; la generación sigue comenzando desde la Gaussiana estándar. El acoplamiento se construye sobre la marcha durante el entrenamiento.
• Limitaciones: Los autores señalan modestamente que QC-FM es un sustituto local al lote para un acoplamiento estructurado, no una aproximación del Transporte Óptimo global. Los cuantiles se estiman dentro de los mini-lotes, y la teoría controla la varianza por corte en lugar de garantizar la rectitud de la trayectoria global o mejoras específicas en el FID. La selección del número de cortes ($k$) y la proporción de anclas ($p$) depende actualmente de heurísticas y ajustes específicos del conjunto de datos.

En conclusión, el trabajo demuestra que el acoplamiento de cuantiles unilateral puede mejorar la calidad de la muestra y la eficiencia del entrenamiento en Flow Matching al estructurar el emparejamiento entre ruido y datos mediante la preservación de rangos, ofreciendo una alternativa práctica a los métodos de transporte óptimo computacionalmente prohibitivos.