VibeToken: Scaling 1D Image Tokenizers and Autoregressive Models for Dynamic Resolution Generations

1. Planteamiento del Problema

Los modelos de generación de imágenes autoregresivos (AR) (por ejemplo, LlamaGen) han demostrado un rendimiento competitivo, pero enfrentan barreras significativas para su adopción en producción en comparación con los modelos de difusión. El principal cuello de botella es la escalabilidad en cuanto a resolución y relaciones de aspecto:

• Explosión de Tokens: Los tokenizadores 2D convencionales (por ejemplo, VQGAN) generan secuencias de tokens que crecen linealmente con la resolución de la imagen. Para una imagen de 1024×1024, esto resulta en miles de tokens (por ejemplo, 4,096 tokens).
• Coste Computacional Cuadrático: Los modelos AR dependen de la autoatención, que tiene una complejidad de $O(T^2)$, donde $T$ es la longitud de la secuencia. A medida que aumenta la resolución, los FLOPs de inferencia crecen cuadráticamente, haciendo que la generación de alta resolución sea computacionalmente prohibitiva.
• Rigidez: La mayoría de los modelos AR se entrenan con resoluciones fijas (por ejemplo, 256×256) y luchan por generalizar a relaciones de aspecto arbitrarias o resoluciones más altas sin reentrenamiento o pipelines complejos de superresolución.
• Dominio de la Difusión: Los modelos de difusión manejan naturalmente resoluciones arbitrarias, dejando a los modelos AR atrás en casos de uso flexibles y listos para producción.

2. Metodología

Los autores proponen una solución en dos partes: VibeToken (un tokenizer novedoso) y VibeToken-Gen (un generador AR).

A. VibeToken: Tokenizador 1D Agnóstico a la Resolución

VibeToken es un tokenizer de imágenes basado en Transformer 1D diseñado para codificar imágenes de cualquier resolución en un número fijo y controlable por el usuario de tokens (por ejemplo, 32–256), desacoplando la cantidad de tokens de la resolución espacial.

Las innovaciones arquitectónicas clave incluyen:

1. Embebidos de Posición de Cuadrícula Dinámicos: En lugar de embebidos de posición absoluta fijos, VibeToken utiliza una cuadrícula aprendida de embebidos que se redimensiona diferencialmente (mediante interpolación bilineal/bicúbica) para coincidir con la cuadrícula de parches de entrada ($T_H \times T_W$). Esto permite que el modelo maneje relaciones de aspecto arbitrarias sin degradación de calidad.
2. Embebidos de Parche Adaptativos: Inspirado en FlexiViT, el modelo admite tamaños de parche variables ($k$) durante el entrenamiento y la inferencia. Un operador de redimensionamiento de pesos mapea los pesos de proyección base al tamaño de parche actual, permitiendo que el modelo se adapte a diferentes resoluciones sin reparametrización.
3. Resolución de Decodificador Adaptativa: El decodificador predice píxeles a una resolución fija alta (4 veces el tamaño del parche) y utiliza una capa de muestreo descendente convolucional 2D aprendible para mapear a la resolución objetivo. Esto habilita la superresolución nativa sin un módulo de escalado superior separado.
4. Tokenización de Longitud Dinámica: El modelo se entrena con una muestreo uniforme de longitudes latentes $L \in [32, 256]$. El codificador produce exactamente $L$ tokens, y el decodificador consume exactamente $L$, eliminando la necesidad de relleno o restricciones de longitud fija.
5. Cuantización: Utiliza Cuantización de Vectores Múltiples (MVQ) con 8 libros de códigos para lograr un tamaño de vocabulario de 32,768, equilibrando la compresión y la calidad de reconstrucción.

B. VibeToken-Gen: Generador AR Eficiente

VibeToken-Gen es un modelo AR condicionado por clase (basado en la arquitectura LlamaGen) que opera sobre las secuencias de tokens cortas y agnósticas a la resolución producidas por VibeToken.

• Presupuesto Computacional Constante: Dado que la longitud de la secuencia de tokens $L$ es fija (por ejemplo, 64 tokens) independientemente de la resolución de salida, los FLOPs de inferencia permanecen constantes.
• Condicionamiento: El modelo se condiciona en la resolución objetivo $(H, W)$ y la relación de aspecto junto con la etiqueta de clase para evitar artefactos de estiramiento geométrico.
• Entrenamiento: Entrenado en ImageNet-1k con resoluciones mixtas (de 256×256 a 512×512) y relaciones de aspecto, utilizando un presupuesto de tokens comparable a los baselines de resolución fija.

3. Contribuciones Clave

1. Tokenización Agnóstica a la Resolución: El primer tokenizer Transformer 1D que admite nativamente resoluciones y relaciones de aspecto arbitrarias mientras mantiene un presupuesto de tokens corto y dinámico (32–256 tokens).
2. Eficiencia Computacional: Desacopla el coste de inferencia AR de la resolución de la imagen. VibeToken-Gen mantiene un 179 GFLOPs constante para cualquier resolución, mientras que baselines como LlamaGen escalan a ~11 TFLOPs para imágenes de 1024×1024 (una ganancia de eficiencia de 63.4×).
3. Superresolución Nativa: La arquitectura admite generar imágenes de alta resolución a partir de entradas de baja resolución (o objetivos arbitrarios) sin módulos de escalado superior separados.
4. Rendimiento de Vanguardia: Logra una calidad de reconstrucción (rFID) y generación (gFID) competitiva en comparación tanto con tokenizadores 2D como con modelos de difusión, a pesar de utilizar significativamente menos tokens.

4. Resultados Experimentales

El artículo evalúa VibeToken y VibeToken-Gen en los conjuntos de datos ImageNet-1k y FFHQ a través de varias resoluciones (de 256×256 a 1024×1024) y relaciones de aspecto.

• Calidad de Reconstrucción:
  • VibeToken-LL logra un 0.40 rFID en 256×256 y un 2.40 rFID en 1024×1024.
  • Supera a otros tokenizadores 1D (por ejemplo, TiTok, One-D-Piece) en generalización de resolución y iguala o supera a los tokenizadores 2D (por ejemplo, IBQ, MAGViT) en eficiencia de compresión.
• Calidad de Generación:
  • VibeToken-Gen (XXL) sintetiza imágenes de 1024×1024 utilizando solo 64 tokens con un gFID de 3.94.
  • Comparación:
    • Vs. Difusión (NiT): VibeToken-Gen es 2.35× más rápido (0.46s vs 1.08s) y logra mejor calidad (3.94 gFID vs 5.87 gFID) en 1024×1024.
    • Vs. Baseline AR (LlamaGen): LlamaGen requiere ~11 TFLOPs para 1024×1024, mientras que VibeToken-Gen utiliza un constante 179 GFLOPs.
• Eficiencia:
  • La cantidad de tokens se reduce de ~4,096 (en tokenizadores 2D) a 64 tokens para imágenes de 1024×1024.
  • La velocidad de inferencia es inferior a un segundo (0.46s) para la generación de alta resolución, en comparación con más de 32 segundos para los modelos AR baselines en resoluciones similares.

5. Significado

Este trabajo cierra la brecha entre los modelos autoregresivos y los modelos de difusión en términos de flexibilidad y eficiencia.

• Viabilidad en Producción: Al permitir que los modelos AR manejen resoluciones arbitrarias con cómputo constante, VibeToken hace viable la generación AR para aplicaciones del mundo real donde las dimensiones de entrada/salida varían (por ejemplo, aplicaciones móviles, formatos de pantalla diversos).
• Escalabilidad: Demuestra que los modelos AR pueden escalar a altas resoluciones sin la explosión computacional cuadrática, desafiando la noción de que la difusión es el único camino hacia la IA generativa de alta resolución.
• Dirección Futura: El enfoque abre la puerta a modelos multimodales unificados (texto, imagen, video) que pueden operar bajo un presupuesto de cómputo fijo y controlable por el usuario, independientemente de la complejidad espacial del contenido generado.

En resumen, VibeToken resuelve el problema de la "escalabilidad de resolución" en la generación AR desplazando la carga de la adaptación de resolución del modelo generativo a un tokenizer flexible y agnóstico a la resolución, resultando en un sistema más rápido, económico y flexible que las alternativas actuales de vanguardia.