Resumen Técnico: Ajuste Fino Gradual para Modelos de Coincidencia de Flujo (Flow Matching)

Planteamiento del Problema

El ajuste fino (fine-tuning) de modelos de coincidencia de flujo preentrenados presenta un desafío significativo en escenarios caracterizados por datos limitados, distribuciones en evolución o restricciones computacionales. Si bien los avances recientes en el ajuste fino basado en recompensas (por ejemplo, Adjoint Matching, enfoques de aprendizaje por refuerzo) han mostrado resultados prometedores, estos suelen sufrir de limitaciones teóricas e inestabilidades prácticas. Específicamente, los métodos existentes frecuentemente fallan al no poder garantizar simultáneamente la corrección teórica, la estabilidad del entrenamiento, la eficiencia de la inferencia y la preservación de la diversidad.

Un escenario crítico, aunque poco explorado, es el ajuste fino cuando se dispone de muestras directas de la distribución objetivo (por ejemplo, en adaptación de dominio o deriva de datos), eliminando la necesidad de funciones de recompensa externas. El ajuste fino estándar en este contexto trata al modelo preentrenado meramente como una inicialización, ignorando su dinámica aprendida (como las trayectorias de probabilidad cortas), lo que conduce a una alta inestabilidad y un potencial colapso de modo. Existe la necesidad de un marco basado en principios que equilibre la adaptación a una nueva distribución objetivo mientras preserva las propiedades beneficiosas del modelo preentrenado.

Metodología: Ajuste Fino Gradual (GFT)

Los autores proponen el Ajuste Fino Gradual (GFT), un marco que utiliza una regularización con recocido (annealed regularization) para crear una secuencia de objetivos intermedios controlada por temperatura. Este enfoque interpola suavemente entre la dinámica preentrenada y la dinámica objetivo.

Formulación Teórica

GFT se deriva de un problema de optimización sobre el espacio de trayectorias inducido por flujos estocásticos. El objetivo minimiza la divergencia de Kullback-Leibler (KL) entre el proceso ajustado $P_\theta$ y el proceso objetivo $P_q$, regularizado por la divergencia KL entre $P_\theta$ y el proceso preentrenado $P_{\theta_0}$:

$$\min_{\theta} \text{KL}(P_\theta \| P_q) + \beta \text{KL}(P_\theta \| P_{\theta_0})$$

donde $\beta \in \mathbb{R}^+$ es un parámetro de temperatura.

Teorema 4.1 establece que la minimización de este objetivo produce un campo de deriva óptimo de forma cerrada $v^*_\theta$, el cual es una combinación convexa del campo vectorial objetivo $v_q$ y el campo vectorial preentrenado $v_{\theta_0}$:

$$v^*_\theta(X_t, t) = \left( \frac{1}{1 + \beta} \right) v_q(X_t, t) + \left( \frac{\beta}{1 + \beta} \right) v_{\theta_0}(X_t, t)$$

Regularización con Recocido (Annealed Regularization): Para lograr la convergencia, GFT emplea un esquema de enfriamiento dependiente del tiempo donde $\beta$ se reduce gradualmente desde un valor inicial alto hacia $0^+$.

• $\beta$ Alto: Penaliza fuertemente la desviación del modelo preentrenado, evitando el colapso prematuro y asegurando la estabilidad.
• $\beta$ Bajo ($\to 0^+$): Relaja la restricción, permitiendo que el modelo recupere asintóticamente el campo vectorial objetivo $v_q$.

Reformulación Tratable

Dado que el campo vectorial objetivo $v_q$ es desconocido y el objetivo se define sobre medidas de trayectoria, los autores aplican el Teorema de Girsanov para reformular el objetivo en términos de términos de deriva. Esto resulta en una función de pérdida tratable que involucra la distancia $L_2$ al cuadrado entre derivas.

Además, para permitir un entrenamiento e inferencia eficientes, GFT utiliza los principios de la Coincidencia de Flujo Condicional (CFM). Al condicionar en acoplamientos arbitrarios fuente-objetivo (por ejemplo, planes de Transporte Óptimo), la pérdida se convierte en:

$$\mathcal{L}_\pi(\theta) = \mathbb{E}_{Z \sim \pi, X_t \sim P(\cdot|Z), t \sim U[0,1]} \left[ \frac{1}{2} \left( \|v_\theta - v_q(\cdot|Z)\|^2 + \beta \|v_\theta - v_{\theta_0}\|^2 \right) \right]$$

Esta formulación preserva las garantías teóricas del objetivo no condicional al tiempo que permite el uso de acoplamientos eficientes (como OT) para acelerar la generación.

Extensión a la Inferencia de Pocos Pasos

El marco se extiende naturalmente a los modelos de Flujo Medio (Mean Flow - MF), que permiten la inferencia de pocos pasos. Los autores demuestran que el campo de velocidad promedio óptimo para los modelos MF también sigue una forma de combinación convexa, permitiendo que GFT se aplique a la generación de uno o pocos pasos sin sacrificar la corrección teórica.

Principales Contribuciones

El artículo describe cuatro contribuciones primarias:

1. Convergencia Teórica: Prueba de que GFT converge a la distribución objetivo bajo supuestos estándar, con una tasa de convergencia explícita que depende de la temperatura $\beta$.
2. Equivalencia de Acoplamiento: Establecimiento de que GFT produce resultados equivalentes al condicionar en acoplamientos arbitrarios fuente-objetivo, permitiendo el uso de técnicas de entrenamiento (como OT) que mejoran la eficiencia de la inferencia.
3. Extensión a Pocos Pasos: Demostración de la extensión natural de GFT a métodos de inferencia de pocos pasos (por ejemplo, Mean Flow) sin comprometer las garantías teóricas.
4. Superioridad Empírica: Validación experimental exhaustiva que muestra que GFT logra una calidad de generación comparable o superior a otros métodos, proporcionando una mayor estabilidad de entrenamiento, diversidad de generación y trayectorias de probabilidad consistentemente más cortas.

Resultados Experimentales

Los autores evaluaron GFT en múltiples configuraciones, incluyendo la adaptación de dominio cruzado y de dominio interno en los conjuntos de datos de referencia WILDS (Camelyon17, RxRx1, FMoW) y datos simulados en 2D.

• Estabilidad y Convergencia: GFT demostró una estabilidad significativamente mayor en comparación con el ajuste fino estándar de CFM y los métodos basados en recompensas (por ejemplo, Adjoint Matching). Métricas como la varianza instantánea y la correlación de Spearman entre épocas y FID indicaron una convergencia monótona y estable.
• Calidad de Generación: En la adaptación de dominio cruzado, GFT logró puntuaciones FID comparables al ajuste fino estándar pero con longitudes de trayectoria promedio menores, lo que indica una inferencia más eficiente. En la adaptación de dominio interno, GFT igualó la precisión del ajuste fino completo manteniendo longitudes de trayectoria más bajas.
• Diversidad: En objetivos gaussianos bimodales 2D, tanto GFT como el CFM estándar cubrieron ambos modos con alta diversidad. En contraste, los métodos basados en recompensas como Adjoint Matching sufrieron de colapso de modo, generando muestras densamente agrupadas dentro de un solo modo.
• Eficiencia: GFT mantuvo la eficiencia de entrenamiento libre de simulación de CFM. Mientras que Adjoint Matching requería simulaciones de ODE costosas para cada paso (resultando en tiempos de entrenamiento un orden de magnitud superiores), GFT se mantuvo eficiente.
• Inferencia de Pocos Pasos: Al aplicarse a modelos de Flujo Medio, GFT logró un FID comparable a los modelos iterativos pero con una reducción de 50-100x en las evaluaciones del modelo (NFEs) durante la inferencia.
• Control de Regularización: Los experimentos variando la temperatura mínima $\beta$ revelaron un compromiso controlable entre el FID final y la longitud de trayectoria promedio, permitiendo a los usuarios ajustar explícitamente el equilibrio entre la precisión de la generación y la velocidad de inferencia.

Significado y Reivindicaciones

El artículo posiciona a GFT como una alternativa simple, pero teóricamente fundamentada y prácticamente efectiva para la adaptación escalable de modelos de coincidencia de flujo bajo deriva de distribución.

Los autores afirman que GFT aborda brechas críticas en la literatura existente al satisfacer todos los desiderata clave para un método de ajuste fino ideal:

• Teóricamente Correcto: Converge de manera demostrable a la distribución objetivo correcta con una tasa controlable.
• Estable: Exhibe una convergencia altamente monótona, evitando las oscilaciones e inestabilidades comunes en el ajuste fino estándar.
• Eficiente: Mantiene el entrenamiento libre de simulación de CFM y soporta una inferencia eficiente mediante acoplamientos OT y métodos de pocos pasos.
• Diverso: Preserva la diversidad del modelo preentrenado, evitando el colapso de modo observado en los enfoques basados en recompensas.

Al cerrar la brecha entre el preentrenamiento de alta calidad y el ajuste fino, GFT ofrece una solución robusta para adaptar modelos generativos a dominios complejos impulsados por datos donde se requiere un control matizado sobre el resultado del ajuste fino. El trabajo establece una base para la adaptación escalable de modelos generativos basados en flujo.