Guiding a Diffusion Transformer with the Internal Dynamics of Itself

Resumen Técnico: Internal Guidance (IG) para Transformadores de Difusión

1. El Problema

Los modelos de difusión, especialmente los basados en arquitecturas Transformer (DiT, SiT), han demostrado una capacidad excepcional para capturar distribuciones de datos complejas. Sin embargo, enfrentan dos desafíos principales durante la generación:

• Regiones de baja probabilidad: Los modelos a menudo fallan al generar imágenes de alta calidad en áreas de baja probabilidad de la distribución de datos, ya que carecen de suficientes ejemplos de entrenamiento para aprenderlas.
• Limitaciones de las estrategias de guía actuales:
  • Classifier-Free Guidance (CFG): Aunque mejora la alineación con el prompt, un coeficiente de guía excesivo tiende a simplificar en exceso las imágenes o distorsionarlas, reduciendo la diversidad.
  • Autoguidance (Guía con una versión "mala" del modelo): Métodos anteriores que utilizan una versión degradada del propio modelo para guiar la generación requieren estrategias de degradación cuidadosamente diseñadas, entrenamiento adicional o pasos de muestreo extra, lo que limita su escalabilidad y eficiencia.

2. Metodología: Internal Guidance (IG)

Los autores proponen una estrategia simple pero efectiva llamada Internal Guidance (IG). Esta metodología aprovecha las dinámicas internas del propio Transformador de Difusión sin necesidad de modelos externos o pasos de muestreo adicionales.

Fase de Entrenamiento:

• Se introduce una supervisión auxiliar en una capa intermedia de la red durante el entrenamiento.
• Se define una pérdida adicional ($L_{inter}$) en esta capa intermedia, además de la pérdida final ($L_{final}$).
• La función de pérdida total es: $L = L_{final} + \lambda L_{inter}$.
• Esto obliga a la red a producir una salida generativa "más débil" o menos refinada desde las representaciones de nivel medio, creando efectivamente una versión interna "mala" del modelo.

Fase de Muestreo (Inferencia):

• Durante la generación, se extraen tanto la salida de la capa intermedia ($D_i$) como la salida de la capa final ($D_f$).
• Se aplica una extrapolación similar a la CFG, pero utilizando la diferencia entre la salida final y la intermedia para guiar el proceso hacia regiones de alta probabilidad, evitando las distribuciones pobres de la capa intermedia.
• La fórmula de guía es:
  $$D_w(x; c) = D_i(x; c) + w(D_f(x; c) - D_i(x; c))$$
  Donde $w$ es el coeficiente de guía.
• Ventaja clave: No requiere entrenamiento adicional para crear un modelo "malo" separado, ni pasos de muestreo extra. Es un enfoque "plug-and-play".

Aceleración de Entrenamiento:

• Los autores también exploran incorporar la señal de guía de IG directamente en la función de pérdida durante el entrenamiento (modificando el objetivo de la capa final para que aprenda la dirección de la capa intermedia), lo que acelera la convergencia y reduce los gradientes que desaparecen.

3. Contribuciones Clave

1. Estrategia de Guía Interna (IG): Un método novedoso que utiliza las salidas de capas intermedias del propio modelo para guiar la generación, logrando un efecto similar al Autoguidance pero sin los costos computacionales de entrenamiento o muestreo adicionales.
2. Compatibilidad con CFG: Se demuestra que IG es altamente compatible con la Classifier-Free Guidance (CFG). La combinación de ambas estrategias mitiga las desventajas de cada una (IG mantiene la diversidad mientras CFG mejora la fidelidad), logrando resultados superiores.
3. Mejora en la Eficiencia de Entrenamiento: La supervisión auxiliar en capas intermedias no solo sirve para la guía, sino que también ayuda a mitigar el problema de los gradientes que desaparecen, logrando una convergencia comparable o superior a métodos de regularización auto-supervisada complejos.
4. Nuevos Récords de Estado del Arte (SOTA): Lograr un FID (Fréchet Inception Distance) extremadamente bajo en ImageNet 256x256 con menos épocas de entrenamiento que los métodos anteriores.

4. Resultados Experimentales

Los experimentos se realizaron principalmente en ImageNet 256×256 utilizando modelos como SiT-XL/2 y LightningDiT-XL/1.

• SiT-XL/2 + IG:
  • Logró un FID de 5.31 en solo 80 épocas (sin CFG).
  • Alcanzó un FID de 1.75 en 800 épocas (sin CFG), superando a SiT vanilla entrenado 1400 épocas.
  • Con CFG y guía de intervalo, alcanzó un FID de 1.46.
• LightningDiT-XL/1 + IG:
  • Logró un FID de 2.42 en solo 60 épocas.
  • Alcanzó un FID de 1.34 en 680 épocas (sin CFG), superando significativamente a todos los métodos previos.
  • Estado del Arte (SOTA): Combinado con CFG y guía de intervalo, el modelo alcanzó un FID de 1.19 en 680 épocas, estableciendo un nuevo récord en la generación de imágenes condicionales.
• Eficiencia: La mejora en la calidad de generación se logra con un aumento marginal en los parámetros y FLOPs, manteniendo la latencia casi idéntica a los modelos base.

5. Significado e Impacto

Este trabajo es significativo por varias razones:

• Simplicidad y Efectividad: Demuestra que una supervisión auxiliar simple en capas intermedias puede superar a estrategias de regularización auto-supervisada mucho más complejas.
• Escalabilidad: La mejora proporcionada por IG se vuelve más significativa a medida que aumenta el tamaño del modelo, lo que lo hace ideal para los grandes modelos de difusión actuales.
• Reducción de Costos: Al eliminar la necesidad de entrenar modelos "mala" separados o añadir pasos de muestreo, IG reduce drásticamente los costos computacionales para lograr alta calidad.
• Nuevas Direcciones de Investigación: La idea de utilizar las discrepancias entre capas internas para guiar la generación abre la puerta a nuevas estrategias de aceleración de entrenamiento y diseño de funciones de pérdida que modelan la dinámica interna de la red.

En conclusión, Internal Guidance representa un avance importante en la optimización de modelos de difusión, ofreciendo una vía eficiente para alcanzar calidad de generación de nivel SOTA con menos recursos y menor tiempo de entrenamiento.