Unified Latents (UL): How to train your latents

El artículo presenta Unified Latents (UL), un marco de entrenamiento que combina un prior de difusión y un modelo de difusión para generar representaciones latentes eficientes, logrando resultados de vanguardia en calidad de imagen y video con menor costo computacional.

Lo esencial

¡Claro que sí! Imagina que quieres enviar una foto de alta calidad a un amigo que vive muy lejos, pero el correo postal (nuestra red de internet o la memoria de la computadora) es muy lento y caro.

El problema es: ¿Cómo comprimes la foto lo suficiente para que sea barata de enviar, pero sin que tu amigo reciba una imagen borrosa y sin sentido?

Los científicos de Google DeepMind en este paper presentan una solución llamada "Latentes Unificados" (Unified Latents). Aquí te explico cómo funciona usando una analogía sencilla:

1. El Problema: El "Sándwich" Roto

Antes de este método, existían dos formas de hacer esto, y ambas tenían defectos:

• Opción A (El compresor estricto): Usabas un compresor muy fuerte. La foto llegaba pequeña, pero tu amigo tenía que "adivinar" muchos detalles. El resultado era una imagen borrosa o con artefactos extraños (como si la foto se hubiera descompuesto).
• Opción B (El compresor suave): Usabas un compresor que dejaba pasar casi toda la información. La foto llegaba perfecta, pero el archivo era tan gigante que tardaba días en enviarse y costaba una fortuna.

Además, en el mundo de la Inteligencia Artificial, a veces el "compresor" (el codificador) y el "descompresor" (el decodificador) no se entendían bien entre ellos, y el sistema aprendía a hacer trucos para engañar al sistema, en lugar de aprender a comprimir de verdad.

2. La Solución: El Equipo de Tres (UL)

Los autores proponen un equipo de tres personas que trabajan juntas en una cadena de montaje:

1. El Codificador (El Empaquetador): Toma la foto original y la convierte en una "caja" pequeña (el latente).
2. El Prior (El Supervisor de Calidad): Es un experto que revisa qué tan pequeña es la caja. Su trabajo es decir: "Oye, esa caja es demasiado pequeña, no cabrá nada. Oye, esa es demasiado grande, estás desperdiciando espacio".
3. El Decodificador (El Desempaquetador): Toma esa caja pequeña y trata de reconstruir la foto original lo más fielmente posible.

La magia de "Latentes Unificados" es que el Codificador y el Prior (el Supervisor) entrenan juntos desde el principio.

3. La Analogía del "Ruido Controlado"

Imagina que el Codificador quiere meter una foto en una caja.

• En el pasado, el Codificador intentaba meter la foto perfectamente limpia. Pero el Prior no sabía cómo medir si la foto era "demasiado compleja" para la caja.
• En este nuevo método, el Codificador intencionalmente añade un poco de "ruido" o estática a la foto antes de meterla en la caja.
• El Prior (el Supervisor) es un experto en "limpiar ruido". Como sabe exactamente cuánto ruido se añadió, puede medir con precisión cuánta información real hay dentro de la caja.

¿Por qué es genial esto?
Es como si le dijeras al Codificador: "No intentes guardar la foto perfecta. Guarda la foto con un poco de estática, y yo (el Prior) te diré exactamente cuánta información real cabe en tu caja antes de que te pases".

Esto crea un equilibrio perfecto:

• Si la caja es muy pequeña, el Prior le dice al Codificador: "¡Eh, estás perdiendo demasiada información!".
• Si la caja es muy grande, el Prior dice: "¡Podrías haber usado menos espacio!".

4. El Resultado: La Foto Perfecta y Barata

Gracias a este entrenamiento conjunto:

• Eficiencia: Logran comprimir la imagen mucho mejor que los métodos anteriores. En pruebas con imágenes (como las de ImageNet), lograron una calidad visual increíble (casi perfecta) usando menos energía de computadora.
• Control: Tienen un "botón de volumen" (llamado loss factor) que les permite decidir: "¿Quiero una caja más pequeña y rápida, o una caja más grande y con más detalles?". Pueden ajustar esto fácilmente sin tener que reinventar todo el sistema.
• Video: También funcionó muy bien con videos (Kinetics-600), logrando el mejor resultado hasta la fecha en la relación entre calidad y costo de entrenamiento.

En Resumen

Antes, comprimir imágenes para la IA era como intentar adivinar el peso de una caja sin una báscula. A veces metías demasiada ropa y la caja se rompía, o muy poca y quedaba vacía.

Unified Latents es como poner una báscula inteligente (el Prior) que trabaja en equipo con el empaquetador (el Codificador). Juntos aprenden a llenar la caja justo hasta el límite perfecto: ni un gramo más, ni un gramo menos, para que el desempaquetador pueda sacar una foto o video espectacular.

Es una forma más inteligente, estable y eficiente de enseñar a las máquinas a "pensar" en imágenes comprimidas.

Resumen técnico

Resumen Técnico: Unified Latents (UL)

1. El Problema

Los modelos de difusión han logrado un éxito notable en la generación de imágenes y video, a menudo utilizando representaciones latentes compactas para escalar eficientemente a altas resoluciones. Sin embargo, existen desafíos fundamentales en cómo aprender estas latencias:

• Compromiso Información-Calidad: Existe una tensión inherente entre la densidad de información de la latente y la calidad de reconstrucción. Latentes con menos información son más fáciles de modelar (mejor generación), pero pierden detalles de alta frecuencia (peor reconstrucción).
• Limitaciones de los VAEs Tradicionales: Los Autoencoders Variacionales (VAE) clásicos, como los usados en Stable Diffusion, utilizan una penalización KL manual entre la distribución latente y un Gaussiano estándar. Sin una pérdida basada en verosimilitud en el decodificador, es difícil determinar el contenido de información real de la latente y ajustar el peso del término KL manualmente.
• Representaciones Semánticas: Los enfoques recientes que usan redes preentrenadas (ej. DINO) o autoencoders fuertemente regularizados obtienen buenos resultados de generación (FID bajo), pero a menudo pierden información de alta frecuencia, resultando en reconstrucciones con artefactos o bajo PSNR.

El objetivo es encontrar una forma sistemática de regularizar las latentes para que sean óptimas tanto para la reconstrucción como para ser modeladas eficientemente por un modelo de difusión posterior.

2. Metodología: Unified Latents (UL)

El marco Unified Latents propone un enfoque de entrenamiento conjunto donde las latentes son codificadas, regularizadas y modeladas simultáneamente mediante difusión. La arquitectura consta de tres componentes principales: un Encoder, un Prior de Difusión y un Decodificador de Difusión.

Principios Clave del Entrenamiento:

1. Codificación con Ruido Fijo: En lugar de aprender una distribución compleja (media y varianza) como en los VAEs tradicionales, el encoder predice una latente determinista ($z_{clean}$). Esta latente se "ruidifica" explícitamente con una cantidad fija de ruido gaussiano hasta un nivel de ruido mínimo ($t=0$ en el esquema de difusión del prior). Esto simplifica la distribución a una con ruido fijo.
2. Prior de Difusión y Límite de Verosimilitud: Se entrena un modelo de difusión (el Prior) sobre las latentes codificadas. Al alinear el nivel de ruido mínimo del prior con el ruido de entrada del encoder, el término KL de la ELBO (Evidence Lower Bound) se reduce a una pérdida MSE ponderada sobre los niveles de ruido. Esto proporciona un límite superior estricto y interpretable sobre la tasa de bits (bitrate) de la latente.
3. Decodificador de Difusión con Ponderación: El decodificador es también un modelo de difusión que reconstruye la imagen a partir de la latente ruidosa ($z_0$). Para evitar el colapso posterior (donde el decodificador ignora la latente), se utiliza una ELBO reponderada (usando una función sigmoide) y un factor de pérdida que aumenta el peso de la reconstrucción frente al prior. Esto permite controlar el intercambio entre la calidad de reconstrucción y la complejidad de modelado.

Proceso de Entrenamiento (Dos Etapas):

• Etapa 1: Entrenamiento conjunto del Encoder, Prior y Decodificador. Se optimiza la suma de la pérdida del prior (sobre las latentes) y la pérdida del decodificador (sobre la imagen).
• Etapa 2 (Modelo Base): Una vez congelados el encoder y el decodificador, se reentrena el modelo de prior (ahora llamado "modelo base") utilizando una ponderación sigmoide para mejorar la calidad de las muestras generadas, permitiendo un tamaño de modelo y batch más grandes.

3. Contribuciones Clave

• Marco Unificado: Introduce un método donde el prior de difusión y el decodificador de difusión se entrenan conjuntamente para regularizar las latentes, eliminando la necesidad de ajustar manualmente pesos de KL.
• Control Interpretable de Bits: Proporciona un límite superior matemático sobre la cantidad de información (bits) en la latente, controlado directamente por hiperparámetros simples (factor de pérdida y sesgo de la sigmoide).
• Estabilidad: Al usar una codificación determinista con ruido fijo en lugar de una distribución variacional compleja, se evita la inestabilidad de entrenamiento común en los VAEs con priores de difusión.
• Eficiencia Computacional: Logra una mejor relación entre el costo de entrenamiento (FLOPs) y la calidad de generación en comparación con modelos entrenados en latentes de Stable Diffusion.

4. Resultados Experimentales

Los autores evaluaron el enfoque en ImageNet-512 (imágenes) y Kinetics-600 (video):

• Generación de Imágenes (ImageNet-512):

  • Lograron un FID de 1.4, competitivo con los mejores modelos actuales.
  • Mantuvieron una alta calidad de reconstrucción (PSNR alto), superando a los métodos que usan latentes semánticas que suelen tener reconstrucciones borrosas.
  • Requieren menos FLOPs de entrenamiento que los modelos basados en latentes de Stable Diffusion para alcanzar un rendimiento similar o superior.
  • En modelos de texto-a-imagen a gran escala, superaron a los baselines de difusión de píxeles y latentes de Stable Diffusion tanto en calidad perceptual (FID) como en alineación con el texto (CLIP).

• Generación de Video (Kinetics-600):

  • Establecieron un nuevo estado del arte (SOTA) con un FVD de 1.3 para el modelo "medium", superando a enfoques previos como MAGVIT y Video Diffusion en la relación costo-eficacia.

• Ablaciones:

  • Se demostró que eliminar el prior de difusión o usar una regularización L2 estándar degrada significativamente el rendimiento.
  • Se confirmó que las latentes con menos bits (menor tasa de bits) son más fáciles de modelar para el prior, pero requieren un decodificador más potente; el marco UL permite ajustar este equilibrio óptimo.

5. Significado e Impacto

El trabajo Unified Latents representa un avance significativo en el diseño de arquitecturas para modelos de difusión:

• Principio de Diseño: Ofrece un enfoque principado para el diseño de latentes, moviéndose de la heurística manual (ajuste de pesos KL) a una optimización conjunta basada en verosimilitud.
• Eficiencia: Demuestra que es posible entrenar modelos de generación de alta calidad con menor costo computacional si las latentes están optimizadas específicamente para el prior de difusión.
• Escalabilidad: Proporciona una base sólida para escalar modelos de difusión a resoluciones y datasets más grandes, permitiendo un control explícito sobre el compromiso entre la fidelidad de la reconstrucción y la capacidad de generación.
• Generalización: Aunque se centra en imágenes y video, el marco es teóricamente aplicable a otros dominios de datos, incluyendo la compresión de datos discretos como texto.

En conclusión, Unified Latents resuelve el problema de la "caja negra" en el entrenamiento de autoencoders para difusión, ofreciendo un método estable, eficiente y controlable para aprender representaciones latentes óptimas.