Resumen Técnico: Síntesis Eficiente de Imágenes con Codificador Latente Esférico

Declaración del Problema

La generación de imágenes en pocos pasos ha avanzado significativamente mediante métodos basados en consistencia y meanflow, sin embargo, estos enfoques a menudo sufren de inestabilidad en el entrenamiento, sensibilidad a los hiperparámetros y conflictos inherentes en los objetivos de optimización. Una alternativa, el Codificador Esférico [40], logra una generación de alta calidad en pocos pasos proyectando representaciones latentes sobre una hipersfera. No obstante, el marco del Codificador Esférico enfrenta dos limitaciones críticas:

1. Ineficiencia Computacional: El proceso de generación requiere transiciones repetidas entre el espacio de píxeles y el espacio latente (decodificación a una imagen, luego re-codificación) en cada paso de muestreo. Esto introduce una sobrecarga computacional sustancial tanto durante el entrenamiento como en la inferencia.
2. Conflicto de Objetivos: El marco optimiza conjuntamente la reconstrucción y la generación dentro de una única arquitectura codificador-decodificador. Investigaciones previas sugieren que estas tareas favorecen representaciones diferentes; por ejemplo, niveles de ruido más bajos mejoran la reconstrucción pero degradan el rendimiento generativo al reducir la cobertura hipersférica. Esta tensión fuerza el uso de arquitecturas más grandes para mantener la calidad de las muestras.

Metodología

Los autores proponen el Codificador Latente Esférico, un marco que desacopla la reconstrucción de la generación para eliminar las operaciones iterativas en el espacio de píxeles. La innovación central es realizar todo el proceso generativo dentro de un espacio latente esférico fijo.

Arquitectura

El método consta de tres componentes principales:

1. Tokenizador de Imágenes Fijo: Un autoencoder de representación preentrenado (RAE) [43] se utiliza como un codificador ($E$) y un decodificador ($D$) fijos. El codificador, basado en DINOv2, mapea imágenes a representaciones latentes, mientras que el decodificador reconstruye las imágenes. Este componente se mantiene congelado durante el entrenamiento.
2. Modelo de Eliminación de Ruido Latente: Una red de eliminación de ruido separada basada en transformadores ($G$, utilizando la arquitectura SiT [25]) opera exclusivamente en el espacio latente. Toma como entrada un latente esférico ruidoso y predice el latente limpio.
3. Proyección Esférica: Las representaciones latentes se proyectan sobre una hipersfera utilizando una función de esferificación basada en RMSNorm ($F$), asegurando una estructura latente uniforme.

Objetivo de Entrenamiento

A diferencia del Codificador Esférico, que calcula pérdidas en el espacio de píxeles, el Codificador Latente Esférico traslada toda la supervisión al espacio latente para mejorar la eficiencia. El entrenamiento involucra dos latentes esféricos ruidosos, $v_{NOISY}$ y $v_{noisy}$, generados añadiendo ruido gaussiano a diferentes escalas.

• Pérdida de Reconstrucción ($L_{recon}$): Combina una pérdida $\ell_1$ y una pérdida de similitud coseno entre el latente limpio predicho $G(v_{noisy})$ y el latente limpio ground truth $z$. Esto alinea el latente desruidado con la representación limpia.
• Pérdida de Consistencia ($L_{cons}$): Hace cumplir la consistencia entre predicciones de diferentes niveles de ruido ($v_{NOISY}$ y $v_{noisy}$). La predicción del latente con menor ruido se trata como un objetivo fijo (mediante stop-gradient) para fomentar la suavidad local en la hipersfera.
• Eliminación de la Pérdida de Consistencia Latente: Los autores eliminan explícitamente la pérdida de consistencia latente utilizada en el Codificador Esférico (que re-codifica las imágenes decodificadas), encontrando que añade sobrecarga computacional sin mejorar el rendimiento en su configuración desacoplada.

Procedimiento de Muestreo

El muestreo ocurre enteramente en el espacio latente (Algoritmo 1):

1. Se muestrea un latente desde una distribución gaussiana.
2. Se proyecta sobre la hipersfera.
3. El modelo de eliminación de ruido $G$ refina el latente utilizando guía libre de clasificadores.
4. El latente refinado se vuelve a proyectar, se perturba con ruido y se utiliza como entrada para el siguiente paso.
5. Este refinamiento iterativo continúa durante $T$ pasos.
6. El decodificador se invoca una sola vez al final para mapear el latente final a píxeles.

Los autores también proponen una estrategia de muestreo de ruido modificada utilizando una distribución logit-normal para adaptarse mejor al muestreo en pocos pasos, diferenciándose del muestreo uniforme en el Codificador Esférico.

Contribuciones Clave

1. Marco Codificador Latente Esférico: Un marco generativo que opera enteramente en el espacio latente esférico, eliminando las transiciones repetidas píxeles-latente. Esto reduce el costo de inferencia en aproximadamente un 85% (aproximadamente 6.5× menos FLOPs) en comparación con el Codificador Esférico en los mismos pasos de muestreo.
2. Diseño Desacoplado: Al aprovechar un autoencoder preentrenado fijo y un modelo de eliminación de ruido separado, el marco permite que la reconstrucción y la generación se especialicen independientemente, resolviendo el conflicto de objetivos inherente a la optimización conjunta.
3. Optimización Eficiente: El método evita objetivos de aproximación de primer orden (por ejemplo, productos vector-Jacobiano) requeridos por algunos métodos de flujo en pocos pasos, aprendiendo en su lugar un proceso directo de eliminación de ruido latente en una variedad esférica estructurada.

Resultados Experimentales

El método fue evaluado en Animal-Faces, Oxford-Flowers e ImageNet-1K (256×256).

• Rendimiento vs. Codificador Esférico: En Animal-Faces y Oxford-Flowers, el método propuesto superó significativamente al Codificador Esférico tanto en calidad de generación (FID) como en costo computacional (GFLOPs) a través de 2, 4 y 6 pasos de muestreo. Por ejemplo, en Animal-Faces con 4 pasos, el FID mejoró de 18.23 (Codificador Esférico) a 6.89 (El nuestro), mientras que los GFLOPs disminuyeron de 4554 a 390.
• Resultados en ImageNet-1K: Con 4 pasos de muestreo (8 NFE con guía), el método logró un FID de 2.25, una mejora sustancial sobre el 4.02 del Codificador Esférico bajo el mismo presupuesto. También superó a varias líneas base recientes de pocos pasos (por ejemplo, MeanFlow, $\alpha$-Flow) y se acercó al rendimiento de modelos de difusión de múltiples pasos (por ejemplo, DiT-XL/2) que requieren 250+ pasos, pero con significativamente menos evaluaciones de funciones.
• Estudios de Ablación:
  • Distribución de Ruido: El muestreo log-normal superó al muestreo uniforme.
  • Pérdidas: La combinación de pérdidas de reconstrucción y consistencia produjo los mejores resultados; añadir la pérdida de consistencia latente degradó el rendimiento.
  • Autoencoder: La elección del tokenizador fue crítica. RAE [43] (con características DINOv2 y $d=768$) superó significativamente a FLUX-VAE y GAE, destacando la necesidad de representaciones semánticas fuertes y espacios latentes de alta dimensión.

Significado y Afirmaciones

El artículo afirma que el Codificador Latente Esférico ofrece una tubería más simple y eficiente para la generación de imágenes en pocos pasos. Al eliminar la dependencia del refinamiento iterativo en el espacio de píxeles, el método logra:

• Mayor Eficiencia: Reducción drástica en FLOPs y tiempo de inferencia en comparación con el Codificador Esférico.
• Mejor Calidad: Mejora en la calidad de las muestras (FID más bajo) debido al desacoplamiento de la reconstrucción y la generación, permitiendo que el modelo de eliminación de ruido se enfoque exclusivamente en la generación sin comprometer la fidelidad de la reconstrucción.
• Estabilidad: Un proceso de entrenamiento que evita la inestabilidad asociada con los productos vector-Jacobiano y los objetivos conflictivos de la optimización conjunta.

Los autores posicionan este trabajo como un paso hacia el modelado generativo escalable y eficiente que aprovecha los beneficios de los espacios latentes esféricos sin las penalizaciones computacionales de implementaciones anteriores. Reconocen limitaciones, señalando el enfoque actual en la generación condicional a clases y la dependencia de la calidad del autoencoder preentrenado, lo cual puede restringir las capacidades de muestreo en un solo paso.