Curriculum Sampling: A Two-Phase Curriculum for Efficient Training of Flow Matching

1. El Problema: La Compensación Velidad-Calidad en el Muestreo de Pasos de Tiempo

En los modelos de Flow Matching (FM), la distribución de muestreo de los pasos de tiempo, denotada como $p(t)$, es un hiperparámetro crítico que determina qué partes de la trayectoria de generación (desde el ruido $t=1$ hasta los datos $t=0$) reciben más atención durante el entrenamiento.

• Práctica Actual: La tendencia predominante favorece distribuciones estáticas sesgadas hacia el medio (como la distribución Logit-Normal), que concentran la densidad de muestreo alrededor de $t \approx 0.5$.
• La Limitación Identificada: El artículo demuestra que esta elección induce una compensación (trade-off) entre velocidad de convergencia y calidad final:
  • Las distribuciones sesgadas al medio aceleran la convergencia temprana (aprendizaje de la estructura global).
  • Sin embargo, sufren de una submuestreo de los regímenes de frontera ($t \to 0$ y $t \to 1$), lo que resulta en una fidelidad asintótica inferior y errores persistentes en los detalles finos de la imagen.
• Análisis de la Dificultad: El análisis del paisaje de pérdida por paso de tiempo revela un perfil de dificultad en forma de "U". Los errores de predicción del campo vectorial son máximos cerca de los extremos (datos y ruido), mientras que la fase de transporte central ($t \approx 0.5$) es más fácil de aprender. Muestrear estáticamente el centro ignora las regiones donde el modelo necesita más entrenamiento para lograr alta fidelidad.

2. Metodología: Muestreo de Currículo (Curriculum Sampling)

Para resolver este dilema, los autores proponen una estrategia de muestreo de currículo de dos fases, donde la distribución $p(t)$ evoluciona dinámicamente durante el entrenamiento en lugar de ser fija.

Las Dos Fases:

1. Fase de Estructura (Early Training):
   • Distribución: Logit-Normal (sesgada hacia el medio).
   • Objetivo: Aprovechar la facilidad de aprendizaje en la región central para lograr una rápida reducción de la pérdida y establecer la estructura global del campo vectorial.
2. Fase de Refinamiento (Late Training):
   • Distribución: Uniforme ($t \sim U(0, 1)$).
   • Objetivo: Una vez establecida la estructura, se cambia a muestreo uniforme para garantizar una cobertura suficiente de los regímenes de frontera ($t \to 0, 1$). Esto permite al modelo corregir los errores residuales y refinar los detalles finos que se perdieron en la primera fase.

Configuración Experimental:

• Modelo: Flow Matching Condicional (CFM) con arquitectura U-Net (~55M parámetros).
• Dataset: CIFAR-10 (imágenes de 32x32).
• Comparación: Se evaluaron distribuciones estáticas (Uniforme, Muestreo de Modos con forma de U, Logit-Normal) frente a la estrategia de currículo dinámica.
• Hiperparámetro Clave: El tiempo de cambio ($T_s$) entre la fase 1 y la fase 2.

3. Contribuciones Clave

1. Identificación de la Compensación: Demostración empírica y analítica de que el muestreo estático "sesgado al medio" (Logit-Normal) acelera el entrenamiento inicial pero impone un techo de calidad inferior al muestreo uniforme debido a la falta de atención en los bordes.
2. Análisis del Paisaje de Pérdida: Caracterización de una curva de pérdida en forma de U, donde los extremos de la trayectoria son los más difíciles de aprender, justificando la necesidad de un enfoque de currículo.
3. Propuesta de Currículo Dinámico: Introducción de un protocolo de entrenamiento de dos fases que combina lo mejor de ambos mundos: la velocidad de convergencia del Logit-Normal y la fidelidad final del muestreo uniforme.
4. Optimización del Tiempo de Cambio: Identificación de que el momento de transición es un hiperparámetro sensible, encontrando un punto óptimo alrededor del 40% del presupuesto de entrenamiento.

4. Resultados

Los experimentos en CIFAR-10 muestran mejoras significativas tanto en calidad como en eficiencia:

• Mejora en Calidad (FID):
  • El mejor FID con muestreo Uniforme (línea base estática) fue 3.85.
  • El mejor FID con Curriculum Sampling fue 3.22.
  • Esto representa una mejora relativa del 16.4% en la calidad de la generación.
• Eficiencia en Entrenamiento:
  • La línea base Uniforme alcanzó su mejor rendimiento a los 150k pasos.
  • El currículo alcanzó su mejor rendimiento a los 100k pasos.
  • Esto implica una reducción del 33% en los pasos de entrenamiento necesarios para alcanzar el estado del arte.
• Comparación con Estrategias Estáticas:
  • Las estrategias estáticas Logit-Normal (ej. $\mu = -0.4$) lograron FIDs peores (ej. 4.03) y mostraron signos de sobreajuste o estancamiento en la calidad final.
  • Las estrategias que cambiaron a muestreo uniforme en la segunda fase superaron consistentemente a sus contrapartes estáticas.

5. Significado e Impacto

Este trabajo cambia la perspectiva sobre el muestreo de pasos de tiempo en modelos generativos continuos:

• De Hiperparámetro Fijo a Currículo Evolutivo: Sugiere que $p(t)$ no debe tratarse como un hiperparámetro estático, sino como un componente dinámico que debe adaptarse a la dificultad cambiante del aprendizaje a lo largo del tiempo.
• Eficiencia Computacional: Al permitir alcanzar mejores resultados con menos pasos de entrenamiento, el método reduce significativamente el costo computacional y energético del entrenamiento de modelos de Flow Matching.
• Generalización: Aunque probado en CIFAR-10, la lógica de la dificultad en forma de U y la necesidad de cubrir tanto la estructura global como los detalles de frontera es aplicable a modelos de mayor resolución y otros tipos de modelos de flujo.
• Implicaciones para la Inferencia: Los hallazgos sugieren que también se podrían diseñar esquemas de inferencia no uniformes que asignen más resolución a los regímenes de tiempo donde el campo vectorial es más difícil de aproximar.

En resumen, el artículo demuestra que una estrategia de entrenamiento híbrida y dinámica supera a las mejores prácticas estáticas actuales, logrando un equilibrio óptimo entre la velocidad de aprendizaje y la fidelidad de la generación final.