JLT: Clean-Latent Prediction in Latent Diffusion Transformers

Resumen Técnico: JLT: Predicción de Latente Limpio en Transformadores de Difusión Latente

Planteamiento del Problema

Los modelos de difusión de eliminación de ruido invierten tradicionalmente un proceso de corrupción, pero el objetivo específico que la red neuronal está entrenada para predecir varía según la arquitectura. Los objetivos comunes incluyen el ruido ($\epsilon$), la velocidad ($v = x - \epsilon$) o los datos limpios ($x$). Aunque estos objetivos son algebraicamente convertibles para un tiempo de corrupción fijo, investigaciones recientes (por ejemplo, JiT) sugieren que, bajo capacidad de red finita y datos finitos, la elección del objetivo de predicción altera significativamente la dificultad del problema de regresión. Específicamente, predecir datos limpios permite que un modelo se centre en variaciones estructuradas de baja dimensión en el manifold, en lugar de reconstruir ruido ambiental fuera del manifold.

La pregunta central abordada por este artículo es si esta ventaja de "predicción limpia" persiste cuando el modelo opera en un espacio latente comprimido en lugar del espacio de píxeles crudos. En la difusión latente, un Autoencoder Variacional (VAE) ya ha eliminado gran parte de la variabilidad de los píxeles crudos. Los autores investigan si las propiedades geométricas del objetivo de predicción ($x$ frente a $v$) siguen siendo un factor crítico incluso cuando el espacio de representación está fijo y comprimido.

Metodología

Los autores introducen JLT (Transformador Latente Conjunto), un marco experimental controlado diseñado para aislar el efecto del objetivo de predicción dentro de un espacio latente fijo.

1. Arquitectura y Configuración:

   • Esqueleto: Un Transformador de escala Base (130M parámetros) con 12 bloques, dimensión oculta 768 y 12 cabezas de atención, reflejando la configuración de JiT-B/16.
   • Representación: El modelo opera sobre códigos latentes congelados del VAE FLUX.2. Esto asegura que el espacio de representación sea idéntico en todas las variantes comparadas.
   • Entrenamiento: Los modelos se entrenan durante 250K pasos (200 épocas) en ImageNet 256 $\times$ 256.
   • Variables Controladas: El estudio mantiene estrictamente constantes la representación (VAE FLUX.2), la escala de la arquitectura, la configuración del optimizador, el tamaño de lote y los protocolos de muestreo.

2. Comparación Experimental:
   Los autores comparan dos variantes principales bajo los mismos ajustes:

   • JLT (Predicción de Latente Limpio): La red es supervisada directamente para predecir el vector latente limpio $x$.
   • DiT (Predicción de Velocidad): Una línea base emparejada donde la red es supervisada para predecir el vector de velocidad $v = x - \epsilon$.
   • Variantes: Los experimentos se realizan en dos tamaños de parche de la cuadrícula VAE, denominados /1 y /2 (por ejemplo, JLT-B/1 frente a DiT-B/1).

3. Análisis Teórico:
   El artículo proporciona un análisis gaussiano local para explicar la brecha empírica. Asumiendo una aproximación lineal-gaussiana local de la distribución de datos latentes ($x \sim \mathcal{N}(0, \Sigma)$) y ruido independiente ($\epsilon \sim \mathcal{N}(0, I)$):

   • Estructura de Covarianza: La covarianza del objetivo de velocidad es $Cov(y_v) = \Sigma + I$, mientras que la del objetivo limpio es $Cov(y_x) = \Sigma$. Esto implica que la predicción de velocidad añade un suelo unitario isotrópico a cada dirección. En direcciones donde los datos limpios tienen baja varianza (débilmente soportados por el manifold), la predicción de velocidad fuerza al modelo a aprender un objetivo de varianza unitaria, amplificando efectivamente el ruido en esas direcciones.
   • Ambigüedad Condicional: El análisis muestra que la varianza residual de Bayes para la predicción de velocidad se escala por un factor de $1/(1-t)^2$ en relación con la predicción limpia. A medida que el valor propio $\lambda_i$ de la covarianza de los datos se acerca a cero, el coeficiente para la predicción limpia tiende a 0 (atenuando la dirección), mientras que el coeficiente de velocidad tiende a $-1/(1-t)$ (amplificando la dirección).

Resultados Clave

Los experimentos en ImageNet 256 $\times$ 256 demuestran una brecha de rendimiento consistente y significativa que favorece la predicción de latente limpio:

• Rendimiento Cuantitativo:
  • JLT-B/1 (Limpio) logra un FID-50K de 2.50 (con guía libre de clasificador), en comparación con 6.56 para el DiT-B/1 (Velocidad) emparejado.
  • JLT-B/2 (Limpio) logra un FID-50K de 14.81, en comparación con 28.71 para DiT-B/2.
  • La Puntuación de Inception (IS) también muestra una mejora sustancial (por ejemplo, 232.51 frente a 132.12 para las variantes /1).
• Dinámica de Entrenamiento: Las curvas de entrenamiento indican que los modelos de latente limpio entran en el régimen de bajo FID antes (alrededor de 100K pasos) y mantienen un margen claro sobre los modelos de velocidad durante todo el entrenamiento.
• Robustez: La ventaja se mantiene a través de diferentes granularidades de tokenización (/1 y /2), lo que sugiere que el beneficio no es un artefacto de un tamaño de parche específico.

Significado y Afirmaciones

El artículo afirma que la elección del objetivo de predicción en la difusión latente es una elección de modelado geométrico y no una mera parametrización algebraica.

1. Geometría del Objetivo sobre Compresión: Los resultados refutan la noción de que la compresión latente por sí sola explica la eficiencia de la difusión latente. La brecha de rendimiento existe dentro del mismo espacio latente fijo, demostrando que la "geometría del objetivo" (la distribución de la señal supervisada) es un factor distinto de la "geometría de la representación" (el tokenizador VAE).
2. Mecanismo de Mejora: Los autores argumentan que la predicción limpia es superior porque amortigua las direcciones latentes de baja varianza (que a menudo son ruido o características del manifold débilmente soportadas), mientras que la predicción de velocidad amplifica estas direcciones al heredar un suelo de covarianza isotrópico.
3. Alcance de las Afirmaciones: Los autores son modestos con sus afirmaciones teóricas. No afirman que la predicción limpia sea óptima globalmente para todos los tokenizadores, programas de ruido o muestreadores. En cambio, proporcionan un mecanismo gaussiano local que es consistente con la brecha empírica medida. El estudio se presenta como una ablación controlada para aislar el efecto del objetivo, en lugar de una propuesta para una nueva arquitectura base o un reemplazo universal para todos los objetivos de difusión.

En conclusión, el artículo establece que incluso en espacios latentes comprimidos, predecir directamente el estado latente limpio ofrece una ventaja geométrica sobre la predicción de velocidad, lo que conduce a una calidad de generación significativamente mayor en ImageNet.