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

1. Énoncé du problème

Les modèles de génération d'images autoregressifs (AR) (par exemple, LlamaGen) ont démontré des performances compétitives mais font face à des obstacles majeurs pour leur adoption en production par rapport aux modèles de diffusion. Le goulot d'étranglement principal est la scalabilité concernant la résolution et les rapports d'aspect :

• Explosion des tokens : Les tokeniseurs 2D conventionnels (par exemple, VQGAN) génèrent des séquences de tokens qui croissent linéairement avec la résolution de l'image. Pour une image de 1024×1024, cela se traduit par des milliers de tokens (par exemple, 4 096 tokens).
• Coût de calcul quadratique : Les modèles AR reposent sur l'attention self, qui présente une complexité $O(T^2)$ où $T$ est la longueur de la séquence. À mesure que la résolution augmente, les FLOPs d'inférence croissent de manière quadratique, rendant la génération haute résolution prohibitivement coûteuse en calcul.
• Rigidité : La plupart des modèles AR sont entraînés sur des résolutions fixes (par exemple, 256×256) et peinent à généraliser à des rapports d'aspect arbitraires ou à des résolutions plus élevées sans réentraînement ou des pipelines complexes de super-résolution.
• Domination de la diffusion : Les modèles de diffusion gèrent naturellement des résolutions arbitraires, laissant les modèles AR derrière dans des cas d'utilisation flexibles et de niveau production.

2. Méthodologie

Les auteurs proposent une solution en deux parties : VibeToken (un tokeniseur novateur) et VibeToken-Gen (un générateur AR).

A. VibeToken : Tokeniseur 1D agnostique à la résolution

VibeToken est un tokeniseur d'images basé sur un Transformer 1D conçu pour encoder des images de n'importe quelle résolution en un nombre fixe et contrôlable par l'utilisateur de tokens (par exemple, 32–256), découplant le nombre de tokens de la résolution spatiale.

Les innovations architecturales clés incluent :

1. Encodages de position de grille dynamique : Au lieu d'encodages de position absolus fixes, VibeToken utilise une grille apprise d'encodages qui est redimensionnée de manière différentielle (via interpolation bilinéaire/bicubique) pour correspondre à la grille de patchs d'entrée ($T_H \times T_W$). Cela permet au modèle de gérer des rapports d'aspect arbitraires sans dégradation de qualité.
2. Encodages de patch adaptatifs : Inspiré par FlexiViT, le modèle prend en charge des tailles de patch variables ($k$) pendant l'entraînement et l'inférence. Un opérateur de redimensionnement des poids mappe les poids de projection de base à la taille de patch actuelle, permettant au modèle de s'adapter à différentes résolutions sans reparamétrisation.
3. Résolution de décodeur adaptative : Le décodeur prédit des pixels à une haute résolution fixe (4 fois la taille du patch) et utilise une couche de sous-échantillonnage convolutive 2D apprenable pour mapper vers la résolution cible. Cela permet une super-résolution native sans module d'upscaler séparé.
4. Tokenisation à longueur dynamique : Le modèle est entraîné avec un échantillonnage uniforme des longueurs latentes $L \in [32, 256]$. L'encodeur produit exactement $L$ tokens, et le décodeur consomme exactement $L$, éliminant le besoin de remplissage (padding) ou de contraintes de longueur fixe.
5. Quantification : Utilise une quantification multi-vecteurs (MVQ) avec 8 codebooks pour atteindre une taille de vocabulaire de 32 768, équilibrant compression et qualité de reconstruction.

B. VibeToken-Gen : Générateur AR efficace

VibeToken-Gen est un modèle AR conditionné par la classe (basé sur l'architecture LlamaGen) qui opère sur les courtes séquences de tokens agnostiques à la résolution produites par VibeToken.

• Budget de calcul constant : Puisque la longueur de la séquence de tokens $L$ est fixe (par exemple, 64 tokens) indépendamment de la résolution de sortie, les FLOPs d'inférence restent constants.
• Conditionnement : Le modèle se conditionne sur la résolution cible $(H, W)$ et le rapport d'aspect, ainsi que sur l'étiquette de classe, pour prévenir les artefacts d'étirement géométrique.
• Entraînement : Entraîné sur ImageNet-1k avec des résolutions mixtes (de 256×256 à 512×512) et des rapports d'aspect, utilisant un budget de tokens comparable aux bases de référence à résolution fixe.

3. Contributions clés

1. Tokenisation agnostique à la résolution : Le premier tokeniseur Transformer 1D qui prend en charge nativement des résolutions et des rapports d'aspect arbitraires tout en maintenant un budget de tokens court et dynamique (32–256 tokens).
2. Efficacité de calcul : Découple le coût d'inférence AR de la résolution de l'image. VibeToken-Gen maintient un 179 GFLOPs constant pour n'importe quelle résolution, tandis que des bases de référence comme LlamaGen passent à ~11 TFLOPs pour des images de 1024×1024 (un gain d'efficacité de 63,4×).
3. Super-résolution native : L'architecture prend en charge la génération d'images haute résolution à partir d'entrées basse résolution (ou de cibles arbitraires) sans modules d'upscaling séparés.
4. Performances de l'état de l'art : Atteint une qualité de reconstruction (rFID) et de génération (gFID) compétitive par rapport aux tokeniseurs 2D et aux modèles de diffusion, malgré l'utilisation de significantly moins de tokens.

4. Résultats expérimentaux

Le papier évalue VibeToken et VibeToken-Gen sur les ensembles de données ImageNet-1k et FFHQ à travers diverses résolutions (de 256×256 à 1024×1024) et rapports d'aspect.

• Qualité de reconstruction :
  • VibeToken-LL atteint 0,40 rFID à 256×256 et 2,40 rFID à 1024×1024.
  • Il surpasse les autres tokeniseurs 1D (par exemple, TiTok, One-D-Piece) en généralisation de résolution et égale ou dépasse les tokeniseurs 2D (par exemple, IBQ, MAGViT) en efficacité de compression.
• Qualité de génération :
  • VibeToken-Gen (XXL) synthétise des images de 1024×1024 en utilisant seulement 64 tokens avec un gFID de 3,94.
  • Comparaison :
    • Vs. Diffusion (NiT) : VibeToken-Gen est 2,35× plus rapide (0,46s contre 1,08s) et atteint une meilleure qualité (3,94 gFID contre 5,87 gFID) à 1024×1024.
    • Vs. Base de référence AR (LlamaGen) : LlamaGen nécessite ~11 TFLOPs pour 1024×1024, tandis que VibeToken-Gen utilise un constant de 179 GFLOPs.
• Efficacité :
  • Le nombre de tokens est réduit de ~4 096 (dans les tokeniseurs 2D) à 64 tokens pour des images de 1024×1024.
  • La vitesse d'inférence est inférieure à la seconde (0,46s) pour la génération haute résolution, comparée à plus de 32 secondes pour les modèles AR de base à des résolutions similaires.

5. Signification

Ce travail comble le fossé entre les modèles autoregressifs et les modèles de diffusion en termes de flexibilité et d'efficacité.

• Viabilité en production : En permettant aux modèles AR de gérer des résolutions arbitraires avec un calcul constant, VibeToken rend la génération AR viable pour des applications réelles où les dimensions d'entrée/sortie varient (par exemple, applications mobiles, formats d'affichage divers).
• Scalabilité : Il démontre que les modèles AR peuvent passer à l'échelle vers des résolutions élevées sans l'explosion de calcul quadratique, remettant en question l'idée que la diffusion est le seul chemin vers l'IA générative haute résolution.
• Direction future : Cette approche ouvre la voie à des modèles multimodaux unifiés (texte, image, vidéo) qui peuvent fonctionner sous un budget de calcul fixe et contrôlable par l'utilisateur, indépendamment de la complexité spatiale du contenu généré.

En résumé, VibeToken résout le problème de la "mise à l'échelle de la résolution" dans la génération AR en déplaçant le fardeau de l'adaptation de la résolution du modèle génératif vers un tokeniseur flexible et agnostique à la résolution, aboutissant à un système plus rapide, moins cher et plus flexible que les alternatives actuelles de l'état de l'art.