Resumen Técnico: Deriva de Características del Profesor (TFD)

Enunciado del Problema

Los modelos de difusión y de ajuste de flujo preentrenados logran una generación de imágenes de alta fidelidad, pero sufren costos de inferencia elevados, que a menudo requieren decenas o cientos de pasadas hacia adelante para generar una sola imagen. Aunque existen métodos de destilación para comprimir estos profesores de múltiples pasos en generadores eficientes de uno o pocos pasos, los enfoques actuales enfrentan una complejidad significativa. Métodos como la Destilación de Correspondencia de Distribución (DMD) y sus variantes dependen de redes auxiliares (por ejemplo, redes de puntuación falsa, discriminadores adversarios) o calendarios de entrenamiento multifase. Otros enfoques, como los Modelos de Consistencia, pueden sufrir entrenamientos inestables o requerir productos complejos de Jacobiano-vector. Además, los objetivos recientes de "Modelo en Deriva", que ofrecen una pérdida directa de atracción-repulsión para la generación en un solo paso, dependen típicamente de un extractor de características preentrenado externo (por ejemplo, modelos auto-supervisados o latent-MAE) para definir el espacio de representación. Esto introduce una sobrecarga adicional de entrenamiento y complejidad arquitectónica.

El problema central abordado es cómo lograr una destilación de difusión eficiente y de alta calidad en un solo paso utilizando un único objetivo sin depender de redes auxiliares, pipelines de optimización complejos o extractores de características externos.

Metodología: Deriva de Características del Profesor (TFD)

Los autores proponen la Deriva de Características del Profesor (TFD), un marco que simplifica la destilación en un solo paso aprovechando las representaciones internas del propio modelo profesor preentrenado.

1. Espacio de Representación Inducido por el Profesor

A diferencia del Modelo en Deriva original, que requiere un codificador de características externo, TFD utiliza los estados ocultos intermedios del profesor preentrenado congelado como el espacio de características para calcular la pérdida de deriva.

• Mecanismo: Para una muestra generada $x$ y una muestra real $x^+$, se extraen características de capas específicas $\ell$ de la red del profesor $T$: $\phi^\ell_T(x, c) = h^\ell_T(x, t_f, c)$.
• Pérdida de Deriva: El campo de deriva se calcula directamente en este espacio inducido por el profesor. El generador $f_\theta$ se optimiza para minimizar la distancia entre su salida y un "objetivo transportado congelado" definido por el campo de atracción-repulsión del espacio de características del profesor. Esto alinea el objetivo de destilación con el modelo que se está destilando, eliminando la necesidad de una red de representación separada.

2. Extracción de Características Ruidosas

Un hallazgo empírico clave es que extraer características de entradas moderadamente ruidosas produce mejores resultados que el uso de muestras limpias.

• Antes de alimentar una muestra al profesor para la extracción de características, se perturba con el proceso de adición de ruido hacia adelante del profesor: $x_{tf} = x + \sigma_{tf} \xi$.
• Los autores hipotetizan que esta perturbación suaviza la geometría de las características del profesor, suprimiendo discrepancias de alta frecuencia y haciendo que el espacio de representación sea más adecuado para la coincidencia a nivel de distribución. Los experimentos indican que un nivel de ruido de $\sigma_{tf} = 0.1$ es óptimo.

3. Regularización de Margen de Anclaje

Para abordar el problema del colapso de modos (donde el generador falla al cubrir regiones diversas soportadas por el profesor), los autores introducen una ligera pérdida de margen de anclaje.

• Anclas: Se extrae un conjunto de características de anclaje $\{a_i\}$ de muestras reales o generadas por el profesor.
• Mecanismo: La pérdida alienta a las muestras generadas a ocupar regiones alrededor de estas anclas. Ajusta adaptativamente el "soporte" requerido para cada ancla basándose en la densidad local de la propia distribución de anclas.
• Objetivo: Esto proporciona una señal de cobertura complementaria al objetivo local de deriva de atracción-repulsión, asegurando que el generador estudiante cubra diversas regiones soportadas por el profesor sin introducir nuevas redes.

La función objetivo final combina la pérdida de Deriva de Características del Profesor ($L_{TFD}$) y la pérdida de margen de anclaje ($L_{anchor}$):
$$L = L_{TFD} + \lambda_{anchor} L_{anchor}$$

Contribuciones Clave

1. Marco Simplificado: TFD propone un marco de destilación de difusión en un solo paso que calcula el objetivo de deriva directamente en el espacio de representación oculta del profesor preentrenado, eliminando la necesidad de una red de representación adicional.
2. Extracción de Características Ruidosas: El trabajo demuestra que extraer características del profesor de entradas moderadamente ruidosas mejora la geometría de las características para la destilación, lo que conduce a mejores puntuaciones FID.
3. Mitigación del Colapso de Modos: Se introduce una ligera pérdida de margen de anclaje para fomentar la diversidad de muestras y la cobertura de regiones de características soportadas por el profesor, abordando un modo de fallo común en la destilación.
4. Rendimiento: El método logra calidad y diversidad de imágenes competitivas con una sola pasada hacia adelante, simplificando sustancialmente el pipeline general de destilación en comparación con las líneas base más avanzadas.

Resultados Experimentales

Los autores evaluaron TFD en la generación de ImageNet-64×64 (condicionada por clase) y SDXL (texto a imagen).

• ImageNet-64×64:

  • TFD logró un FID de 1.58 con una sola pasada hacia adelante después de 200k iteraciones de entrenamiento.
  • Este resultado es competitivo con DMD2 (FID 1.51) y supera significativamente a otras líneas base de un solo paso como DMD (2.62), EMD (2.20) y CTM (1.92).
  • Convergencia: TFD alcanzó un FID $\le$ 10 en 1.0k actualizaciones, mientras que la línea base DMD2 requirió 5.5k actualizaciones (5.5 veces más rápido). Alcanzó un FID $\le$ 3 en 12.5k actualizaciones en comparación con las 24.0k de DMD2.

• Texto a Imagen SDXL:

  • En el conjunto de datos COCO 2014, TFD logró un FID de 18.39 con una pasada hacia adelante, superando a otras variantes destiladas de SDXL (por ejemplo, DMD2 en 19.01, SDXL-Turbo en 24.57).
  • Mantuvo una puntuación CLIP competitiva de 0.332, lo que indica una fuerte alineación texto-imagen.
  • Los resultados cualitativos mostraron que TFD preserva una estructura coherente, detalles visuales y alineación con el prompt comparables a las líneas base destiladas de 4 pasos y al profesor de 50 pasos, a pesar de utilizar solo un paso.

Significado y Afirmaciones

El artículo afirma que TFD ofrece una alternativa más simple y eficiente a los marcos de destilación existentes. Al utilizar el propio modelo profesor como extractor de características, TFD evita la complejidad de entrenar redes auxiliares o diseñar pipelines multifase. La introducción de la extracción de características ruidosas y la regularización de margen de anclaje aborda desafíos específicos en la geometría de la representación y la cobertura de modos, respectivamente.

Los autores posicionan a TFD como un método que logra generación eficiente en un solo paso con calidad y diversidad de imágenes competitivas, demostrando que un único objetivo en el espacio de representación nativo del profesor es suficiente para una destilación de alta fidelidad. El trabajo sugiere que el concepto de "deriva" puede realizarse efectivamente sin codificadores de características externos, siempre que los estados internos del profesor se utilicen correctamente.

Limitaciones: Los autores señalan que el trabajo actual se centra en profesores basados en U-Net y aún no ha explorado arquitecturas de Transformadores de Difusión (DiT). Además, la selección de capas de características y niveles de ruido actualmente carece de guía teórica, y el método no ha sido probado en conjuntos de datos de mayor escala o configuraciones de mayor resolución más allá de los experimentos actuales.