Resumen Técnico: Alineación de la Geometría Latente para el Emparejamiento de Flujos Esféricos en la Generación de Imágenes

Declaración del Problema

El emparejamiento de flujos latentes (y los modelos de difusión) para la generación de imágenes típicamente opera transportando ruido gaussiano a latentes de un Autoencoder Variacional (VAE) a lo largo de trayectorias lineales en un espacio euclidiano. Sin embargo, este enfoque pasa por alto dos propiedades estructurales críticas de los latentes de VAE de alta dimensión:

1. Concentración de la Medida: Tanto el prior del ruido gaussiano como los latentes de datos codificados se concentran en capas esféricas delgadas en lugar de llenar el volumen del espacio. Una línea euclidiana recta (cuerda) que conecta dos puntos en estas capas atraviesa radios interiores que ninguna de las distribuciones ocupa, forzando al modelo a aprender transiciones a través de regiones del espacio latente que rara vez están pobladas.
2. Dominancia Direccional: La salida del decodificador es predominantemente sensible a la dirección (componente angular) de un token latente en lugar de su longitud (componente radial). Experimentos que intercambian el radio y la dirección de tokens de vecinos de la misma clase demuestran que preservar la dirección mientras se altera el radio resulta en un cambio perceptual mínimo, whereas alterar la dirección provoca cambios semánticos significativos.

El emparejamiento de flujos lineal estándar ignora estos hechos, asignando una porción significativa del objetivo de velocidad de entrenamiento (hasta un 50-90% dependiendo del tokenizador) al movimiento radial. Dado que el decodificador es menos sensible al radio, esta supervisión es ineficiente, lo que conduce a una convergencia más lenta y a una calidad de imagen subóptima (FID) en comparación con los métodos que respetan la geometría subyacente.

Metodología

Los autores proponen el Emparejamiento de Flujos Esféricos, un método que alinea la geometría latente con la estructura esférica de los datos sin modificar la arquitectura de difusión ni requerir codificadores auxiliares. El enfoque consta de cuatro componentes clave:

1. Proyección Esférica Token a Token:
   El método proyecta cada token latente $z$ de un VAE preentrenado sobre una hiperesfera de radio fijo $S^{d-1}(\sqrt{d})$. Esto se logra mediante un paso de normalización $L_2$ insertado entre el codificador congelado y el decodificador:
   $$z_{i,j} \leftarrow \sqrt{d} \cdot \frac{z_{i,j}}{\|z_{i,j}\|}$$
   Esto asegura que todos los tokens yacen sobre una esfera con un radio que coincide con el radio de concentración del prior gaussiano estándar.

2. Prior Esférico:
   En lugar de utilizar ruido gaussiano crudo, el punto final del ruido $z_0$ se construye muestreando ruido gaussiano isotrópico $\epsilon \sim \mathcal{N}(0, I)$ y proyectándolo sobre la esfera:
   $$z_0 = \sqrt{d} \cdot \frac{\epsilon}{\|\epsilon\|_2}$$
   Esto resulta en una distribución uniforme sobre la esfera, coincidiendo con la geometría de los datos proyectados.

3. Ajuste Fino del Decodificador:
   Para acomodar la restricción geométrica, el decodificador se ajusta fino durante pocas épocas (5 épocas en los experimentos) mientras el codificador y la capa de proyección permanecen congelados. El objetivo de reconstrucción mantiene las pérdidas estándar de píxeles (L1, LPIPS, adversarial) pero elimina el término KL, ya que el latente proyectado es determinista.

4. Transporte Slerp y Velocidad Tangente:
   La trayectoria de transporte entre los puntos finales de ruido y datos se reemplaza con Interpolación Lineal Esférica (slerp), que sigue el arco geodésico sobre la esfera.

   • Trayectoria Geodésica: La trayectoria $z_t$ permanece sobre la esfera para todo $t \in [0, 1]$.
   • Objetivo de Velocidad: La velocidad condicional $u_t$ es la derivada temporal de la geodésica, proyectada sobre el espacio tangente $T_{z_t}S^{d-1}$. Esto asegura que la velocidad sea puramente angular, con un componente radial cero por construcción.
   • Entrenamiento y Muestreo: El modelo se entrena para predecir esta velocidad tangente. Durante la inferencia, el muestreador utiliza el mapa exponencial para integrar la velocidad, garantizando que la trayectoria permanezca sobre la esfera sin requerir correcciones de proyección radial en cada paso.

Contribuciones Clave

El artículo identifica y cuantifica una "mismatch de cáscara radial" en el emparejamiento de flujos latentes estándar e introduce una intervención geométrica para resolverlo:

1. Identificación del Desajuste: Los autores demuestran que las trayectorias lineales fuerzan al modelo a predecir movimiento radial para el cual el decodificador es insensible, desperdiciando capacidad y ralentizando la convergencia.
2. Proyección Token a Token: Introducen un método para restringir los latentes de VAE preentrenados a una esfera de radio fijo utilizando únicamente el ajuste fino del decodificador, evitando la necesidad de reentrenar el codificador o utilizar codificadores de representación auxiliares.
3. Emparejamiento de Flujos Basado en Slerp: Proponen entrenar el emparejamiento de flujos a lo largo de geodésicas slerp con objetivos de velocidad tangente, asegurando que toda la trayectoria respete la variedad esférica.
4. Validación Empírica: Se demuestra que el método mejora el rendimiento en múltiples familias de tokenizadores (FLUX.2, VA-VAE, REPA-E FLUX.1) y escalas de modelos (SiT-B a SiT-XL) sin cambios arquitectónicos.

Resultados

Los experimentos en la generación condicional por clase de ImageNet-256 demuestran mejoras consistentes:

• Mejoras en FID: El método esférico-slerp logra puntuaciones FID más bajas que el emparejamiento de flujos lineal vanilla en todas las configuraciones probadas. Por ejemplo, en SiT-B/2 con FLUX.2, el FID disminuye de 26.35 (Lineal) a 20.55 (Esférico-slerp).
• Eficiencia de Entrenamiento: El método alcanza un FID objetivo de 30 en aproximadamente 2.2$\times$ menos pasos de entrenamiento en comparación con la línea base lineal vanilla.
• Generalización: Se observan ganancias en diferentes escalas de guía (CFG) y tamaños de modelos (B y XL), lo que indica que el beneficio es inherente a la geometría latente en lugar de un ajuste específico de hiperparámetros.
• Estudios de Ablación:
  • Intercambiar decodificadores entre las líneas base vanilla y esféricas degrada el rendimiento en ambas direcciones, confirmando que el flujo aprendido está estrechamente acoplado con la geometría latente específica.
  • Utilizar una trayectoria "descompuesta en cáscara" (interpolando radio y dirección por separado) en latentes vanilla no produce las mismas ganancias, confirmando que eliminar el grado de libertad radial por completo es crucial.
  • La calidad de reconstrucción (rFID) de los decodificadores esféricos ajustados finamente sigue siendo competitiva, mostrando que la restricción geométrica no compromete severamente la capacidad del VAE para reconstruir imágenes.

Significancia

El artículo establece la geometría de soporte latente como un eje de diseño crítico para el modelado generativo. Argumenta que la forma de la distribución latente define el espacio que el generador debe recorrer, e ignorar esta geometría (asumiendo linealidad euclidiana) conduce a ineficiencia.

La significancia de este trabajo radica en su capacidad para mejorar la calidad de generación de imágenes y la velocidad de entrenamiento sin:

• Cambiar la arquitectura de difusión (por ejemplo, la columna vertebral SiT).
• Añadir codificadores auxiliares (a diferencia del emparejamiento de flujos riemanniano en características DINOv2 congeladas).
• Reentrenar el codificador VAE desde cero.

Simplemente proyectando latentes preentrenados existentes sobre una esfera y alineando la trayectoria de flujo con las geodésicas resultantes, el método logra resultados de vanguardia con una sobrecarga computacional mínima. Los autores sugieren que el trabajo futuro debería explorar estas restricciones geométricas bajo regímenes de compresión más fuertes y diseñar tokenizadores que coincidan nativamente con esta estructura esférica.