Resumen Técnico: Aprendizaje de Priors Autoregresivos Discretos con Flujo de Gradiente de Wasserstein

1. Declaración del Problema

Los tokenizadores de imágenes discretos se entrenan típicamente utilizando una receta de dos etapas: primero, el tokenizador se entrena para reconstruir imágenes (centrándose en la verosimilitud condicional), y segundo, se entrena un prior autoregresivo (AR) sobre las secuencias de tokens resultantes congeladas. Este desacoplamiento crea una discrepancia fundamental: el tokenizador aprende a preservar la información de la imagen pero permanece al margen de los sesgos inductivos del prior AR. En consecuencia, los tokens aprendidos pueden ser excelentes para la reconstrucción, pero difíciles de predecir secuencialmente de izquierda a derecha por un modelo AR.

Los métodos existentes intentan mitigar esto modificando el decodificador (por ejemplo, Tail Dropout, Decodificadores de Difusión) para fomentar una estructura secuencial. Sin embargo, estos enfoques aún dependen de una supervisión centrada en el decodificador y no utilizan directamente la dificultad de modelado del prior AR para guiar el aprendizaje del tokenizador. Además, las técnicas estándar de coincidencia a nivel de distribución (por ejemplo, WAE-MMD, InfoVAE) diseñadas para espacios latentes continuos no se aplican directamente a tokenizadores discretos duros debido a la falta de caminos diferenciables a través de la cuantización.

2. Metodología

Marco Teórico: Consistencia Variacional Tripartita (TVC)

Los autores introducen la Consistencia Variacional Tripartita (TVC) para analizar el aprendizaje de variables latentes. TVC descompone el objetivo generativo en tres condiciones de consistencia derivadas de la regla de Bayes:

1. Consistencia de la Verosimilitud Condicional: Alineación entre el decodificador del modelo y la posterior anclada a los datos.
2. Consistencia del Prior: Alineación entre el prior del modelo y la posterior agregada del codificador.
3. Consistencia de la Posterior: Alineación entre la posterior del modelo y la posterior anclada a los datos.

Los autores argumentan que el entrenamiento estándar de dos etapas satisface la Consistencia de Verosimilitud en la primera etapa y la Consistencia del Prior en la segunda, pero no proporciona una señal al tokenizador para asegurar que la posterior agregada sea intrínsecamente "amigable con AR". Por el contrario, el entrenamiento estilo ELBO impone una consistencia del prior a nivel de instancia, lo que a menudo conduce al colapso de la posterior.

Solución Propuesta: WGF-DPD y wAR-Tok

Para cerrar esta brecha, el artículo propone la Divergencia de Prior Distribucional (DPD), un término que hace coincidir la distribución agregada de tokens con el prior AR objetivo sin restringir instancias individuales. Para optimizar esto para tokens discretos, emplean el Flujo de Gradiente de Wasserstein (WGF).

El mecanismo central, WGF-DPD, opera de la siguiente manera:

1. Modelo Proxy: Un modelo AR auxiliar ligero ($q_\psi$) se entrena para rastrear la distribución agregada actual de tokens producida por el tokenizador ($q_\phi$).
2. Actualización Basada en Puntuación: La dirección de actualización para el tokenizador se deriva del gradiente de Wasserstein de la divergencia KL entre la posterior agregada y el prior objetivo. Esto resulta en una actualización de gradiente a nivel de partícula:
   $$\nabla_z \mathcal{L}_{DPD} \approx -\nabla_z \log p_\theta(x|z) + (\log q_\psi(z) - \log p_\theta(z))$$
   • El primer término es el gradiente estándar de reconstrucción.
   • El segundo término es la puntuación de coincidencia del prior, que representa la diferencia entre las log-probabilidades del modelo proxy y el prior AR objetivo.
3. Optimización Eficiente: Crucialmente, esta actualización requiere solo pasadas forward a través de los modelos AR proxy y objetivo. No requiere retropropagación a través de ninguno de los modelos AR. El gradiente se propaga a los parámetros del codificador a través del Estimador Straight-Through (STE) a través del paso de cuantización no diferenciable.

El tokenizador resultante, denominado wAR-Tok, se entrena conjuntamente con el objetivo de reconstrucción y la señal WGF-DPD.

3. Contribuciones Clave

• Marco TVC: Una descomposición teórica del aprendizaje de variables latentes que identifica explícitamente la "Consistencia del Prior" como un objetivo distinto, explicando por qué el entrenamiento de dos etapas deja al tokenizador al margen de la dificultad de modelado del prior.
• Algoritmo WGF-DPD: Un método de optimización novedoso para tokenizadores discretos que alinea la distribución agregada de tokens con un prior AR utilizando flujo de gradiente de Wasserstein. Evita la intratabilidad de la retropropagación a través de cuantizadores discretos o la inestabilidad de la coincidencia adversarial.
• wAR-Tok: Una implementación práctica que integra la señal WGF-DPD en el entrenamiento del tokenizador, requiriendo solo pasadas forward a través de modelos AR auxiliares para generar señales de alineación del lado del codificador.

4. Resultados Experimentales

El método se evaluó en CIFAR-10 e ImageNet 256×256, comparando contra líneas base que incluyen Tokenización 1D, Tail Dropout, Decodificadores de Difusión y Pasos de Tiempo Atados.

• CIFAR-10:
  • Bajo calidad de reconstrucción coincidente (rFID $\approx$ 5.0), wAR-Tok logró un FID de generación (gFID) de 14.50 y una pérdida de generación AR (gLoss) de 0.42, superando significativamente a la línea base más fuerte (ImageFolder/One-D-Piece: gFID 17.60, gLoss 0.78).
  • Bajo iteraciones de entrenamiento coincidentes, wAR-Tok alcanzó un gFID de 11.30 y un gLoss de 0.34, demostrando una convergencia más rápida y una mayor calidad.
• ImageNet 256×256:
  • Utilizando la misma arquitectura de codificador-decodificador y prior AR que la línea base, wAR-Tok redujo el gFID de 6.32 a 5.42 y la pérdida de evaluación AR de 7.89 a 7.07, con un compromiso menor en la calidad de reconstrucción (rFID 2.05 $\to$ 2.35).
• Dinámicas de Entrenamiento: Las curvas de entrenamiento AR muestran que wAR-Tok mantiene una pérdida de entropía cruzada más baja y una mayor precisión de predicción de tokens durante todo el entrenamiento en comparación con las líneas base, lo que indica que los tokens aprendidos son intrínsecamente más fáciles de modelar para el prior AR.

5. Significado y Afirmaciones

El artículo afirma que wAR-Tok resuelve efectivamente la discrepancia entre tokenizadores discretos y priores autoregresivos al introducir una señal de coincidencia de prior a nivel de distribución durante el entrenamiento del tokenizador.

• Eficiencia: El método es computacionalmente eficiente, añadiendo solo dos pasadas forward (proxy y objetivo) por paso de entrenamiento y evitando la retropropagación costosa a través de los modelos AR.
• Generalidad: Aunque se centra en priores AR, el marco TVC y el enfoque WGF-DPD se presentan como agnósticos al prior, aplicables a cualquier prior discreto con log-probabilidades tratables.
• Modestia: Los autores declaran explícitamente que no afirman un rendimiento absoluto de Estado del Arte (SOTA), ya que excluyeron deliberadamente técnicas auxiliares comunes en tokenizadores recientes (por ejemplo, alineación DINO, destilación, post-entrenamiento con decodificador de difusión) para aislar el efecto de WGF-DPD. Los resultados demuestran la eficacia del objetivo de alineación propuesto en un entorno controlado.

El trabajo sugiere que optimizar explícitamente la "predecibilidad" de los tokens por parte del modelo generativo aguas abajo, en lugar de solo su reconstruibilidad, conduce a sistemas generativos discretos más eficientes y de mayor calidad.