Purrception: Variational Flow Matching for Vector-Quantized Image Generation

1. El Problema: La Tensión Dual en Latentes Vector-Cuantizados (VQ)

El trabajo aborda un desafío fundamental en la generación de imágenes de alta resolución utilizando espacios latentes vector-cuantizados (VQ). Estos espacios representan imágenes como una cuadrícula de índices discretos (de un código o codebook) que tienen asociados vectores de incrustación (embeddings) continuos.

Los métodos existentes se ven forzados a elegir entre dos enfoques degenerados que rompen parte de esta estructura dual:

• Métodos Continuos (ej. Flow Matching Continuo - CFM, Difusión Latente): Operan en el espacio de incrustaciones continuas ($\mathbb{R}^D$). Preservan la geometría y permiten transporte suave, pero ignoran la estructura categórica. No aprenden a elegir índices específicos, no pueden cuantificar la incertidumbre sobre múltiples códigos plausibles y carecen de logits para controlar la generación mediante temperatura.
• Métodos Discretos (ej. Flow Matching Discreto - DFM, VQ-Diffusion): Predicen directamente los índices discretos. Aunque respetan la estructura cuantizada, colapsan la geometría: tratan códigos semánticamente relacionados como tokens no relacionados, convirtiendo el proceso en "saltos" discretos sin capacidad de interpolación ni modelado de incertidumbre suave.

El objetivo es crear un modelo que aproveche la geometría suave de los embeddings continuos mientras recibe supervisión categórica directa sobre los índices.

2. Metodología: Purrception (VQ-VFM)

Los autores proponen Purrception, una adaptación del Flow Matching Variacional (VFM) específicamente diseñada para latentes VQ. La metodología combina lo mejor de ambos mundos:

• Posterior Categórica Variacional: En lugar de predecir directamente el campo de velocidad o el índice final, el modelo aprende una distribución posterior categórica sobre los índices del codebook dado un punto de interpolación actual en el espacio continuo.
  $$q_\theta^t(c | z_t) = \text{Cat}(c | \pi_\theta^t(z_t))$$
  Donde $\pi_\theta^t$ es una distribución de probabilidad sobre los $K$ vectores del código.
• Transporte Continuo: Aunque la supervisión es categórica, el transporte de probabilidad ocurre en el espacio de embeddings continuos. El campo de velocidad se calcula como la esperanza de la distribución categórica sobre los vectores del código:
  $$v_\theta^t(z_t) = \sum_{k=1}^K \pi_{\theta,k}^t(z_t) \left( \frac{e_k - z_t}{1-t} \right)$$
  Esto permite que la incertidumbre sobre múltiples códigos plausibles se traduzca en un movimiento suave y consciente de la geometría, en lugar de saltos discretos.
• Objetivo de Entrenamiento: El entrenamiento se reduce a minimizar la pérdida de entropía cruzada entre la distribución categórica predicha y el índice real del código (ground-truth). Esto hace que el entrenamiento sea eficiente y estable.
• Control por Temperatura ($\tau$): Dado que el modelo opera sobre logits antes de la función softmax, se introduce un parámetro de temperatura en la inferencia.
  • $\tau$ bajo: La distribución colapsa hacia el índice más probable, generando imágenes más nítidas y deterministas (aunque posiblemente menos detalladas).
  • $\tau$ alto: La distribución se aplana, asignando peso a códigos vecinos, lo que añade variabilidad y detalles finos, aunque puede reducir la fidelidad global si es demasiado alto.

3. Contribuciones Clave

1. Enfoque Híbrido: Purrception es el primer modelo que integra exitosamente la supervisión categórica (necesaria para VQ) con la dinámica de transporte continuo (necesaria para la eficiencia y suavidad geométrica) mediante VFM.
2. Eficiencia de Entrenamiento: Al recibir una señal de aprendizaje categórica directa, el modelo converge significativamente más rápido que los métodos puramente continuos o discretos.
3. Control de Inferencia: Introduce un mecanismo de control de temperatura en la fase de inferencia, una característica ausente en los métodos continuos (que no tienen logits) y poco efectiva en los discretos (que toman decisiones duras inmediatas).
4. Cuantificación de Incertidumbre: Permite al modelo expresar incertidumbre sobre múltiples códigos plausibles y traducirla en transporte geométrico suave.

4. Resultados Experimentales

Los experimentos se realizaron en ImageNet-1k a resolución 256x256, utilizando arquitecturas DiT (Diffusion Transformer) y diferentes tokenizadores VQ (vq-f8 y vq-ds8-c2i).

• Velocidad de Convergencia:
  • Purrception converge mucho más rápido que el Flow Matching Continuo (CFM) y el Flow Matching Discreto (DFM).
  • Con el backbone DiT-L/2, alcanza el rendimiento final de DFM en un 3.0x menos de iteraciones y el de CFM en un 1.65x menos.
  • Con DiT-XL/2, la ventaja es aún mayor: 3.5x más rápido que DFM y 2.3x que CFM.
• Calidad de Muestra (FID):
  • Logra un FID competitivo de 3.88, superando a todos los modelos de difusión discreta y generativos enmascarados (Masked Generative Models).
  • Supera a muchos métodos autoregresivos con menos parámetros y decodificación más rápida.
  • Aunque no supera a los modelos de difusión continua de vanguardia (como DiT-XL/2 o SiT-XL/2) que usan autoencoders VAE de alta calidad, cierra la brecha significativamente, demostrando que el enfoque híbrido es viable.
• Efecto de la Temperatura:
  • Se observa una curva en forma de "U" en el FID en función de la temperatura ($\tau$). El rendimiento óptimo se encuentra alrededor de $\tau \approx 0.8 - 0.9$.
  • Valores bajos producen imágenes simples; valores altos añaden ruido. Esto confirma que el ajuste de temperatura es una herramienta efectiva para equilibrar fidelidad y diversidad sin reentrenar.

5. Significado e Impacto

Purrception representa un avance significativo en la generación de imágenes latentes al resolver la dicotomía entre métodos continuos y discretos.

• Eficiencia Computacional: Al reducir drásticamente el número de iteraciones necesarias para converger, disminuye el costo computacional y energético del entrenamiento de modelos generativos de gran escala.
• Nueva Dirección de Investigación: Demuestra que es posible mantener la eficiencia y suavidad de los flujos normales continuos mientras se aprovechan las ventajas de la cuantización vectorial (estabilidad, compresión).
• Controlabilidad: Proporciona un "botón" de temperatura principista para la inferencia, algo que los modelos puramente continuos no pueden ofrecer, mejorando la utilidad práctica del modelo en escenarios donde se requiere ajustar el equilibrio entre detalle y coherencia.

En conclusión, Purrception establece un nuevo estado del arte para la generación basada en latentes VQ, ofreciendo un marco robusto que combina la mejor geometría de los métodos continuos con la supervisión categórica necesaria para la representación discreta eficiente.