Resumen Técnico: Exploración y Explotación de la Estabilidad en el Emparejamiento de Flujos Latentes

1. Enunciado del Problema

El entrenamiento de modelos generativos, particularmente los Modelos de Difusión y sus contrapartes de Emparejamiento de Flujos (FM), sigue siendo computacionalmente costoso. Estos modelos requieren típicamente conjuntos de datos masivos, recursos de cómputo extensos, duraciones de entrenamiento prolongadas y, en configuraciones condicionales, anotaciones exhaustivas. Aunque el Emparejamiento de Flujos ofrece una formulación determinista para la generación de muestras mediante Ecuaciones Diferenciales Ordinarias (ODEs) que a menudo reduce los pasos de muestreo en comparación con las Ecuaciones Diferenciales Estocásticas (SDEs), el costo de entrenar estos modelos en conjuntos de datos a escala completa sigue siendo un cuello de botella.

La pregunta central abordada por este trabajo es si el tamaño del conjunto de datos y la capacidad del modelo pueden reducirse significativamente sin sacrificar la calidad de generación. Específicamente, los autores investigan la estabilidad de los modelos de Emparejamiento de Flujos Latentes (LFM) bajo perturbaciones mayores, incluyendo submuestreo del conjunto de datos (poda), cambios arquitectónicos y configuraciones de entrenamiento alteradas. Plantean la hipótesis de que si las trayectorias de transporte LFM son inherentemente estables, se podría explotar esta propiedad para derivar algoritmos de entrenamiento e inferencia más eficientes.

2. Metodología

2.1. Fundamento Teórico: Estabilidad en el Emparejamiento de Flujos

El trabajo se basa en la observación de que el Flujo Rectificado admite una expresión de forma cerrada para el campo de velocidades. Para un conjunto de datos finito, el campo de velocidades óptimo es una suma ponderada de las muestras de entrenamiento, donde los pesos están determinados por una función softmax sobre las distancias cuadradas entre el estado actual y las muestras de entrenamiento escaladas. Los autores señalan que este softmax se vuelve "puntiagudo" al inicio de la fase de transporte, lo que significa que una sola muestra de entrenamiento a menudo domina el campo de velocidades para una trayectoria dada.

Esta comprensión teórica sugiere que la trayectoria de transporte está determinada en gran medida por la presencia de muestras específicas "dominantes". En consecuencia, eliminar muestras no dominantes (poda) debería dejar la mayoría de las trayectorias sin cambios. Los autores formalizan esto definiendo una métrica de asignación: para una trayectoria que comienza en el ruido $x_0$, determinan a qué muestra de entrenamiento $x_i$ termina. Plantean la hipótesis de que esta asignación permanece invariante incluso cuando se elimina una gran fracción del conjunto de datos, siempre que se conserve la muestra dominante para esa trayectoria.

2.2. Estrategias de Poda de Datos

Para aprovechar esta estabilidad, los autores proponen tres criterios de poda de datos adaptados de modelos discriminativos pero reformulados para el objetivo FM (que utiliza trayectorias de ruido compartidas y pasos de tiempo):

1. Puntuación basada en Gradientes (G): Las muestras se clasifican por la magnitud de la norma del gradiente de la pérdida.
2. Puntuación basada en Pérdida (L): Las muestras se clasifican por la magnitud de la pérdida calculada a lo largo de trayectorias de ruido compartidas.
3. Puntuación basada en Agrupamiento (C): Las muestras se agrupan en un espacio de incrustación preentrenado (por ejemplo, CLIP). Los autores proponen dos variantes específicas:
   • Proporcional ($C_p$): Selecciona muestras proporcionales al tamaño del grupo.
   • Equilibrada ($C_b$): Selecciona un número igual de muestras de cada grupo para mitigar desequilibrios distribucionales.
   • Basada en Kernel ($C_b^\kappa$): Selecciona muestras de manera codiciosa para igualar la media del kernel gaussiano del grupo en el espacio latente.

Los autores también evalúan la "inversa" de estos criterios (seleccionar las puntuaciones más bajas) para probar la robustez de la afirmación de estabilidad.

2.3. Inferencia de Grueso a Fino (C2F)

Explotando la estabilidad de las trayectorias a través de las capacidades del modelo, los autores proponen una estrategia de inferencia en dos etapas:

1. Modelo Grueso: Una arquitectura ligera (por ejemplo, DiT-S/2) se entrena en un subconjunto podado de los datos. Maneja la primera fase de la trayectoria ($t \in [0, t_0)$).
2. Modelo Fino: Una arquitectura de alta capacidad (por ejemplo, DiT-XL/2), preentrenada en el conjunto de datos completo, maneja la segunda fase ($t \in [t_0, 1]$).

Para asegurar una transición suave en el punto de "costura" $t_0$, el Modelo Grueso se ajusta finamente con una pérdida de continuidad. Esta pérdida minimiza la diferencia entre los campos de velocidades de los Modelos Grueso y Fino en $t_0$. La trayectoria del Modelo Fino se utiliza para generar el estado objetivo $x_{t_0}$ mediante inversión de ODE (integración hacia atrás), asegurando que el Modelo Grueso aprenda a alinearse con la trayectoria del Modelo Fino.

3. Contribuciones Clave

1. Demostración Empírica de la Estabilidad LFM: Los autores muestran que los modelos LFM entrenados independientemente convergen hacia trayectorias de transporte altamente consistentes bajo submuestreo de conjuntos de datos, cambios arquitectónicos y diferentes inicializaciones aleatorias. Esta estabilidad está vinculada a las propiedades teóricas del objetivo FM.
2. Algoritmos de Poda Prácticos: Introducen tres criterios de poda (Gradiente, Pérdida, Agrupamiento) adecuados para FM. Demuestran que una poda agresiva (por ejemplo, eliminar el 50-95% de los datos) no degrada el rendimiento e incluso puede mejorar la velocidad de convergencia y la calidad de generación bajo presupuestos de cómputo limitados.
3. Aceleración de Inferencia de Grueso a Fino: Proponen una estrategia de generación de dos modelos que divide la trayectoria. Este enfoque logra una aceleración significativa en la inferencia (aprox. 2.15$\times$) manteniendo la calidad de generación, siempre que la suposición de estabilidad se mantenga a través de las arquitecturas del modelo.
4. Mitigación de Sesgos mediante Poda Equilibrada: Los autores muestran que la poda de agrupamiento equilibrada ($C_b$) puede mitigar sesgos en los datos generados (por ejemplo, género, tono de piel) sin requerir un equilibrio consciente de etiquetas explícito, al igualar la representación de grupos semánticos en el conjunto de entrenamiento.

4. Resultados Experimentales

Los métodos fueron evaluados en CelebA-HQ, FFHQ e ImageNet-1K.

• Estabilidad bajo Poda: En CelebA-HQ, los modelos entrenados en subconjuntos disjuntos (podados al 50%) produjeron imágenes visualmente similares al comenzar desde la misma semilla de ruido $x_0$. La similitud coseno ArcFace entre modelos podados y sin podar osciló entre 0.79 y 0.83, significativamente más alta que pares aleatorios (0.37).
• Métricas de Rendimiento:
  • CelebA-HQ: Podar el 50% de los datos utilizando agrupamiento equilibrado ($C_b$) mejoró ligeramente la Distancia de Inception Fréchet (FID) de 24.24 (sin podar) a 22.80.
  • ImageNet: Podar el 75% de los datos ($pr=0.75$) resultó en una convergencia más rápida y mantuvo o mejoró ligeramente el FID en comparación con el conjunto de datos completo. Incluso con una poda del 95%, el modelo convergió más rápido inicialmente, aunque el rendimiento se degradó después de cierto número de iteraciones.
• Robustez ante Perturbaciones:
  • Arquitectura: Aumentar la capacidad del modelo (de DiT-S/2 a DiT-XL/2) resultó en cambios mínimos en la trayectoria (similitud $s=0.81$). Cambiar a una arquitectura U-Net redujo la similitud ($s=0.55$) pero preservó atributos gruesos.
  • Espacio Latente: La estabilidad es sensible a cambios en la parametrización del espacio latente (por ejemplo, cambiar la semilla del VAE o aplicar transformaciones invertibles), lo cual rompió la estabilidad.
  • Objetivo: Cambiar de Emparejamiento de Flujos a difusión basada en puntuación alteró completamente las trayectorias, indicando que la estabilidad es específica del objetivo FM.
• Eficiencia de Grueso a Fino: En CelebA-HQ, el enfoque C2F con $t_0=0.7$ redujo el tiempo de inferencia de 93.95 ms/imagen (Solo Fino) a 43.53 ms/imagen (aprox. 2.15$\times$ de aceleración) manteniendo el FID.
• Equidad: La poda equilibrada ($C_b$) redujo la divergencia KL hacia una distribución uniforme a través de atributos (género, edad, tono de piel) más efectivamente que los modelos sin podar, a pesar de ser agnóstica a las etiquetas.

5. Significado y Afirmaciones

Los autores afirman que la estabilidad observada es una propiedad característica del paradigma de Emparejamiento de Flujos Rectificado bajo representaciones latentes fijas, en lugar de un artefacto de algoritmos de entrenamiento específicos.

• Eficiencia: El significado principal es la capacidad de reducir el tiempo de entrenamiento y los requisitos de datos sin sacrificar la calidad. Los autores demuestran que para conjuntos de datos grandes como ImageNet, la poda puede acelerar la convergencia, y para conjuntos de datos más pequeños como CelebA-HQ, puede mejorar la calidad de generación.
• Aceleración de Inferencia: El enfoque C2F ofrece una forma principada de reducir los costos de inferencia al delegar las etapas tempranas y gruesas de la generación a modelos ligeros, confiando en la estabilidad de la trayectoria para asegurar que el modelo de alta capacidad solo necesite refinar los detalles finales.
• Control de Sesgos: El trabajo sugiere que la poda de conjuntos de datos, específicamente mediante agrupamiento equilibrado, puede servir como un mecanismo para controlar y mitigar desequilibrios distribucionales en datos generados, ofreciendo un camino hacia modelos generativos más justos sin una reponderación compleja consciente de etiquetas.

El documento concluye que, aunque la estabilidad es robusta ante cambios en los datos y la capacidad del modelo, es sensible a cambios en la estructura del espacio latente y en el objetivo generativo en sí mismo. Esta comprensión permite a los profesionales diseñar pipelines más eficientes en recursos mientras mantienen una generación de alta fidelidad.