Resumen Técnico: Destilación de Transformers con Deriva mediante Autoencoders de Representación

Planteamiento del Problema

Aunque los Autoencoders de Representación (RAE) han mejorado significativamente el rendimiento de los modelos generativos basados en flujo al utilizar espacios latentes con agrupamientos por etiquetas semánticamente ricos derivados de características preentrenadas de DINO, destilar estos modelos en generadores eficientes de un solo paso o de pocos pasos sigue siendo un desafío.

Los autores identifican un desajuste geomético fundamental: los espacios latentes de los RAE exhiben una anisotropía severa y grandes curvaturas en sus trayectorias de la Ecuación Diferencial Ordinaria (ODE) subyacente en comparación con los espacios de los Autoencoders Variacionales (VAE) tradicionales. Los métodos de destilación basados en trayectorias existentes dependen de la suposición de trayectorias del profesor que son suaves y casi rectas. Las trayectorias altamente curvas en los espacios RAE causan que estos métodos se vuelvan inestables o ineficientes. Además, aunque el recientemente propuesto Modelo de Deriva (Drifting Model) ofrece un enfoque a nivel de distribución que evita el emparejamiento de trayectorias curvas, su formulación original requiere un Autoencoder Enmascarado (MAE) auxiliar, entrenado por separado, para proporcionar características concentradas, lo que introduce una sobrecarga computacional significativa.

Metodología

El artículo propone Drift-RAE, un marco de destilación que adapta el paradigma del Modelo de Deriva específicamente para los espacios latentes de RAE, eliminando la necesidad de extractores de características auxiliares.

1. Análisis Teórico de Geometría y Dinámica

Los autores analizan cuantitativamente los espacios latentes, demostrando que las características de RAE son sustancialmente más anisotrópicas (menor Ratio de Participación y Entropía Espectral) que las características de SD-VAE. Demuestran teóricamente (Teorema 1) que en espacios de alta dimensión, si las muestras positivas están excesivamente dispersas, el componente atractivo del campo de deriva tiende a desvanecerse, lo que conduce a un campo degenerado.

• Perspectiva Clave: Mientras que el Modelo de Deriva original requiere un MAE para concentrar las características dispersas, los espacios latentes de RAE están naturalmente concentrados semánticamente. Esto permite que el Modelo de Deriva opere directamente sobre los latentes de RAE sin el módulo MAE auxiliar.

2. El Marco de Trabajo Drift-RAE

El método destila un modelo de flujo preentrenado $v_\theta$ en un generador de un solo paso $G_\theta(z) = z - v_\theta(z, 1)$. En lugar de emparejar trayectorias de ODE, el generador es entrenado para minimizar un objetivo de deriva que empuja las muestras generadas hacia las distribuciones de datos reales y las aleja de otras muestras generadas.

• Objetivo: La pérdida se calcula por token en las representaciones latentes de RAE, utilizando un campo de deriva normalizado $\tilde{V}$ derivado de muestras positivas (reales) y negativas (generadas).

3. Modificaciones Prácticas para la Estabilidad

Para alinear el Modelo de Deriva con el marco de trabajo Diffusion-GAN y mejorar la estabilidad del entrenamiento, los autores introducen tres modificaciones específicas:

• Dimensión Softmax: Los autores eliminan el truco de softmax bidireccional utilizado en el Modelo de Deriva original. Mantienen un único softmax sobre los índices de las muestras, lo cual se alinea mejor con la dirección del gradiente del potencial log-sum-exp subyacente.
• Perturbación de Entrada: Para prevenir el desvanecimiento del gradiente causado por variedades (manifolds) no intersectantes entre los datos reales y los generados, las muestras generadas se perturban con ruido muestreado de una distribución de Laplace antes de calcular el campo de deriva.
• Desacoplamiento Parcial de Muestras Negativas: Para aumentar el número de muestras negativas utilizadas para la aproximación de la distribución sin exceder los presupuestos de memoria de la GPU, los autores proponen utilizar un conjunto más grande de muestras negativas ($N_{total\_neg}$) donde solo un subconjunto ($N_{neg}$) porta gradientes, mientras que el resto se desacopla (detached). Esto mejora la aproximación empírica de la distribución generada.

Contribuciones Clave

1. Análisis Geométrico: Una demostración cuantitativa de que los espacios latentes de RAE poseen alta curvatura y anisotropía, lo que hace que la destilación estándar basada en trayectorias sea inestable.
2. Justificación Teórica: Una prueba que muestra que los Modelos de Deriva fallan en espacios excesivamente dispersos (lo que requiere un MAE en la formulación original) pero prosperan en los espacios semánticos naturalmente concentrados de los RAE.
3. Drift-RAE: Un nuevo método de destilación que logra la generación de un solo paso en espacios RAE sin requerir un extractor de características MAE auxiliar.
4. Mejoras de Estabilidad: Perspectivas teóricas y empíricas sobre la modificación del objetivo de Deriva (dimensión softmax, perturbación de entrada y desacoplamiento de muestras) para alinearse con los principios de Diffusion-GAN y mejorar la convergencia.

Resultados Experimentales

El método fue evaluado en el conjunto de datos ImageNet 256 $\times$ 256 utilizando backbones DiTDH-XL y DiTDH-L.

• Rendimiento: Drift-RAE logra un FID de 1.77 (con DiTDH-XL) y 2.12 (con DiTDH-L) utilizando un solo paso de generación.
• Eficiencia: El modelo alcanza estos resultados en solo 10,000 pasos de destilación (aprox. 16 épocas).
• Comparación:
  • Supera a los métodos de destilación previos en espacios RAE (por ejemplo, MF-RAE).
  • Logra un FID comparable y un FDDINOv2 mejorado en comparación con el Modelo de Deriva original, a pesar de no utilizar el costoso MAE auxiliar.
  • Se mantiene competitivo con generadores de un solo paso y de pocos pasos entrenados en otros espacios latentes (por ejemplo, Espacio de Píxeles, SD-VAE).

Significado y Reivindicaciones

El artículo afirma que la Deriva (Drifting) proporciona un marco efectivo y prometedor para destilar modelos generativos en espacios de representación. El significado principal radica en resolver la inestabilidad de la destilación de modelos basados en RAE aprovechando su concentración semántica inherente, eliminando así la dependencia de extractores de características auxiliares.

Los autores señalan modestamente las limitaciones, reconociendo que su análisis teórico depende de modelos simplificados de alta dimensión y que el entrenamiento de Modelos de Deriva desde cero (sin destilación) en espacios RAE sigue siendo un problema abierto. También destacan que el requisito de muchas muestras positivas de la misma clase puede plantear desafíos de escalabilidad para la generación de texto a imagen, sugiriendo esto como una dirección para investigaciones futuras.