¿Por qué la predicción "clean" es clave en modelos de difusión latente?

Predicción clean en el espacio latente: un estudio controlado del objetivo

Los modelos de difusión pueden entrenarse para predecir los datos limpios (clean), el ruido añadido o un campo de velocidad. Aunque estos objetivos son algebraicamente convertibles, un trabajo reciente (JiT) demostró que, en el espacio de píxeles, regresar directamente la imagen clean explota mejor la estructura de baja dimensión que predecir el ruido ambiente o la velocidad. Este artículo pregunta: ¿sigue importando la elección del objetivo de predicción cuando el modelo opera en un espacio latente comprimido de un VAE, donde ya se ha eliminado gran parte de la variabilidad de los píxeles? Los autores presentan JLT, un Transformer de difusión latente de 130M que utiliza predicción clean-latente sobre códigos congelados del VAE FLUX.2. Comparando JLT con un DiT emparejado de predicción de velocidad bajo idénticas condiciones de representación, arquitectura y entrenamiento, el estudio aísla el efecto del objetivo de regresión directa. El hallazgo central es que la predicción clean-latente supera consistentemente a la predicción de velocidad, demostrando que la parametrización del objetivo es una elección de modelado geométrico, no un mero cambio algebraico.

Objetivos de predicción y equivalencia algebraica

El proceso de corrupción hacia adelante mezcla un latente limpio $x$ y ruido gaussiano $\epsilon$ a lo largo de una trayectoria lineal:

$z_t = t x + (1-t) \epsilon, \quad t \in [0,1].$

Tres objetivos directos comunes son el latente limpio $y_x = x$, el ruido $y_\epsilon = \epsilon$ y la velocidad $y_v = x - \epsilon$. Para un $t$ fijo, cualquier objetivo determina los demás mediante una lectura afín. Por ejemplo, a partir de un valor limpio predicho $\hat{x}_\theta$, se puede recuperar $\hat{\epsilon}_\theta = (z_t - t\hat{x}_\theta)/(1-t)$ y $\hat{v}_\theta = (\hat{x}_\theta - z_t)/(1-t)$. Esta equivalencia algebraica suele llevar a los investigadores a tratar la elección del objetivo como un simple cambio de notación. Sin embargo, la red se entrena antes de aplicar esa lectura, y la lectura escala los errores de predicción de manera diferente según el nivel de ruido. La comparación controlada del artículo modifica únicamente el objetivo directo —latente limpio para JLT, velocidad para el DiT emparejado— manteniendo fijos la representación, la arquitectura principal y el entrenamiento, y revela que el problema de regresión inducido difiere sustancialmente.

Análisis de la geometría del objetivo: por qué importa la elección

Un análisis local lineal-gaussiano explica la brecha empírica. Suponiendo $x \sim \mathcal{N}(0,\Sigma)$ y $\epsilon \sim \mathcal{N}(0,I)$. Las covarianzas marginales de los objetivos son:

$\operatorname{Cov}(y_x) = \Sigma, \quad \operatorname{Cov}(y_v) = \Sigma + I.$

La predicción de velocidad añade un piso isotrópico unitario a cada dirección. Si $\Sigma$ es anisotrópica, las direcciones latentes de baja varianza adquieren varianza unitaria en $y_v$, mientras que la predicción clean mantiene su varianza objetivo pequeña. La ambigüedad condicional también difiere. Para una coordenada individual con autovalor $\lambda_i$,

$\operatorname{Var}(v_i \mid z_i) = \frac{1}{(1-t)^2} \operatorname{Var}(x_i \mid z_i).$

Los estimadores de Bayes revelan un mecanismo adicional: a medida que $\lambda_i \to 0$, el coeficiente del objetivo clean tiende a 0, atenuando las direcciones de baja varianza, mientras que el coeficiente del objetivo de velocidad tiende a $-1/(1-t)$, amplificándolas. Por tanto, aunque los objetivos sean linealmente convertibles después de la predicción, presentan distintos problemas de regresión supervisada a la red.

Arquitectura y configuración de entrenamiento

JLT es un Transformer latente de escala Base con 12 bloques, dimensión oculta 768, 12 cabezas de atención y un embedding de parche con cuello de botella de 128 dimensiones, totalizando 130M de parámetros. Sigue de cerca la configuración de JiT-B/16, pero reemplaza los parches de imagen cruda por tokens latentes fijos del VAE FLUX.2. Se evalúan dos variantes de tamaño de parche: JLT-B/1 y JLT-B/2, correspondientes a parches en la cuadrícula del VAE de tamaño 1 y 2. La línea base emparejada de velocidad, DiT-B/1 y DiT-B/2, utiliza la misma arquitectura y representación pero predice $v = x - \epsilon$. Todos los modelos se entrenan durante 250K pasos (200 épocas) con AdamW, una tasa de aprendizaje base de $5\times10^{-5}$ (escalada a $2\times10^{-4}$) y un tamaño de lote efectivo de 1024. Para aislar el efecto del objetivo, la implementación omite la concatenación repetida del token de clase en contexto y la pérdida auxiliar de clasificación usadas en JiT. El condicionamiento por clase es, por lo demás, estándar.

Ablación de objetivo emparejada: Clean vs.

Velocity

El experimento central fija la representación, la escala del Transformer y la configuración de entrenamiento, variando únicamente el objetivo de predicción directa. Los resultados en ImageNet $256\times256$ se muestran en la Tabla 1 y la Figura 2.

Model	Target	Patch	FID-50K
DiT-B/1	velocity	/1	6.56
JLT-B/1	clean	/1	2.56
DiT-B/2	velocity	/2	28.71
JLT-B/2	clean	/2	14.81
JLT-B/1 (guided)	clean	/1	2.50

Tabla 1: Ablación de objetivo latente emparejada en ImageNet $256\times256$. La predicción clean‑latente domina a la predicción de velocidad en ambos tamaños de parche.

Figura 2: Curvas de entrenamiento para la ablación de objetivo emparejada. Las variantes clean‑latente mantienen un FID más bajo y una puntuación Inception más alta durante todo el entrenamiento.

La predicción clean‑latente mejora el FID de 6.56 a 2.56 en parche /1, y de 28.71 a 14.81 en parche /2. Las curvas de entrenamiento muestran que el modelo clean‑latente entra antes en el régimen de bajo FID y mantiene una ventaja clara. Con guiado libre de clasificador, JLT‑B/1 alcanza un FID de 2.50. La ventaja no está ligada a un tamaño de parche específico, lo que confirma que es la propia geometría del objetivo la que impulsa la mejora.

Comparación con líneas base representativas

La Tabla 2 sitúa el resultado guiado de JLT-B/1 junto a modelos establecidos en ImageNet $256\times256$.
JLT es un modelo latente de 130M entrenado en solo 250K pasos y, sin embargo, consigue un FID competitivo.

Modelo	FID-50K	IS	Entrenamiento
ADM (guiado)	3.94	215.3	1000K
LDM-4	3.60	247.7	178K
DiT-XL/2	2.27	278.2	7M
SiT-XL/2	2.06	270.3	7M
JiT-B/16	2.09	282.3	800K
JiT-B/32	2.28	278.4	800K
JLT-B/1 (nuestro)	2.50	—	250K

Tabla 2: Comparación guiada en ImageNet $256\times256$.
JLT-B/1 es un modelo a menor escala entrenado en menos pasos y, sin embargo, su FID se aproxima al de sistemas mucho más grandes.

Aunque los modelos de escala XL o alineados en representación alcanzan un FID más bajo, alteran varios factores simultáneamente.
La comparación contextualiza la magnitud del beneficio de la predicción limpia sin reclamar un nuevo estado del arte.
La conclusión clave es que un simple cambio de objetivo dentro de un espacio latente fijo produce ganancias sustanciales, incluso a escala modesta.

Conclusión e implicaciones

Este estudio demuestra que la predicción de latente limpio supera consistentemente a la predicción de velocidad cuando la representación, la arquitectura y el entrenamiento se mantienen constantes.
El análisis gaussiano local proporciona un mecanismo: la predicción de velocidad añade un suelo de covarianza isotrópica y amplifica las direcciones latentes de baja varianza, mientras que la predicción limpia las atenúa.
Estos hallazgos replantean la parametrización del objetivo en difusión latente como una elección de modelado geométrico, no un detalle algebraico.
El resultado no se explica únicamente por la compresión latente: la brecha aparece dentro del mismo espacio latente.
Las limitaciones incluyen el enfoque en ImageNet $256\times256$ y una configuración de 130M; el trabajo futuro debería validar el mecanismo en distintos tokenizadores y conjuntos de datos.
Este artículo anima a los profesionales a tratar el objetivo de predicción directa como una dimensión de diseño de primera clase, incluso en espacios latentes comprimidos.