Résumé Technique : FlashAR

Énoncé du Problème

Les modèles autoregressifs (AR) à grande échelle ont atteint des performances de pointe en génération d'images en traitant les données visuelles comme des séquences de tokens discrets. Cependant, les modèles AR standards reposent sur un décodage en balayage raster, générant les tokens strictement de manière séquentielle, de gauche à droite et de haut en bas. Cette dépendance séquentielle entraîne une latence d'inférence qui évolue de manière quadratique avec la résolution de l'image ($O(HW)$), rendant la génération haute résolution prohibitivement lente et empêchant une utilisation efficace du parallélisme moderne des GPU.

Les stratégies d'accélération existantes font face à des compromis significatifs :

1. Redesign du Paradigme : Les méthodes qui modifient les ordres de prédiction des tokens (par exemple, VAR, NAR) ou adoptent une génération multi-échelle nécessitent souvent un entraînement à partir de zéro, les rendant incompatibles avec l'immense dépôt de modèles AR raster pré-entraînés existants.
2. Adaptation Post-Entraînement : Des approches comme l'adaptation par diffusion discrète modifient l'objectif de prédiction original, créant un écart entre l'entraînement et l'inférence et dégradant potentiellement la cohérence structurelle fine inhérente à la modélisation causale.

Le défi central abordé par cet article est le suivant : Un modèle AR pré-entraîné en balayage raster peut-il être transformé en un générateur hautement parallèle tout en préservant son prior génératif appris et en minimisant les frais de post-entraînement ?

Méthodologie : FlashAR

FlashAR est un cadre d'adaptation post-entraînement léger qui convertit un modèle AR pré-entraîné en balayage raster en un générateur parallèle utilisant une structure de prédiction de token suivant bidirectionnelle. Le cadre introduit trois composants techniques principaux :

1. Branchement Intermédiaire pour un Décodage à Double Tête

Au lieu d'attacher une nouvelle tête de prédiction à la dernière couche (ce qui serait biaisé vers l'objectif original de balayage raster horizontal), FlashAR bifurque le réseau au niveau d'une couche intermédiaire ($m$).

• Architecture : Le backbone pré-entraîné est divisé en un tronc partagé ($F_{0:m-1}$) et deux voies parallèles.
  • Tête Horizontale : Réutilise les couches supérieures originales ($F_{m:L-1}$) pour la prédiction ligne par ligne.
  • Tête Verticale : Utilise un bloc entraînable ($\tilde{F}_{m:L-1}$), initialisé comme une copie des couches supérieures, pour la prédiction colonne par colonne.
• Raisonnement : Les caractéristiques intermédiaires conservent des abstractions sémantiques plus riches et sont moins spécialisées dans la direction spécifique de gauche à droite que les représentations de la dernière couche. Le branchement à ce niveau permet à la voie verticale d'accéder au contexte nécessaire sans étendre la profondeur du chemin critique, maintenant ainsi l'efficacité d'exécution.

2. Porte de Fusion Apprenable

Les prédictions horizontales et verticales capturent des dépendances complémentaires, mais leur importance relative varie selon les positions spatiales (par exemple, bords horizontaux vs bords verticaux). Une moyenne statique est insuffisante et peut provoquer des conflits de contexte.

• Mécanisme : FlashAR emploie une porte de fusion apprenable qui combine dynamiquement les prédictions en fonction du contexte des deux états prédécesseurs ($h^H_{p,q-1}$ et $h^V_{p-1,q}$).
• Formule : Le logit fusionné $z_{p,q}$ est calculé comme $g_{p,q} z^H_{p,q-1} + (1 - g_{p,q}) z^V_{p-1,q}$, où la porte $g_{p,q}$ est dérivée via un MLP et une fonction sigmoïde. Cela permet au modèle de pondérer de manière adaptative les indices directionnels.

3. Pipeline d'Adaptation Stabilisé en Deux Étapes

Pour assurer une convergence rapide et prévenir l'oubli catastrophique de la variété pré-entraînée, FlashAR utilise une stratégie d'entraînement en deux étapes :

• Étape 1 (Initialisation des Branches) : Le backbone et la tête horizontale sont figés. Seule la tête verticale et la porte de fusion sont entraînées. Cela permet à la nouvelle voie verticale d'acquérir rapidement une capacité de prédiction sans perturber le prior horizontal établi.
• Étape 2 (Affinement Joint) : L'ensemble du réseau (backbone, les deux têtes et la porte) est affiné conjointement pour harmoniser les représentations spatiales orthogonales pour le nouveau paradigme de décodage parallèle diagonal.

4. Inférence Consciente du Matériel

Pour traduire le parallélisme théorique en accélérations en temps réel, les auteurs implémentent :

• FlexAttention : Compile dynamiquement des masques de proximité 2D très clairsemés à la volée pour éviter la fragmentation de la mémoire due au remplissage zéro dense.
• Cache KV par Lots : Met à jour tous les tokens dans une étape diagonale simultanément, amortissant les frais de lancement de noyaux.

Contributions Clés

1. Cadre FlashAR : Une méthode d'adaptation post-entraînement qui transforme les modèles AR raster pré-entraînés en générateurs parallèles sans le coût d'un pré-entraînement à partir de zéro.
2. Branchement Intermédiaire à Double Tête : Une architecture novatrice qui résout la tension entre la préservation du prior horizontal et l'introduction d'une prédiction verticale orthogonale en branchant au niveau des couches intermédiaires.
3. Fusion Dynamique : Une porte apprenable qui résout les incohérences spatiales en combinant de manière adaptative les signaux horizontaux et verticaux.
4. Pipeline d'Entraînement Efficace : Un processus d'adaptation en deux étapes assurant une convergence rapide avec des besoins minimaux en données.

Résultats Expérimentaux

Les auteurs ont évalué FlashAR sur LlamaGen (ImageNet 256×256) et Emu3.5 (Texte vers Image 512×512).

• Accélération : FlashAR atteint jusqu'à une accélération de 22,9× en temps réel pour la génération d'images 512×512 par rapport aux bases AR standards. Sur LlamaGen, il atteint un débit de 447,2 images/s (contre 117,9 pour l'AR standard).
• Efficacité des Données : L'adaptation nécessite seulement 0,05 % des données d'entraînement originales (par exemple, 80 000 tokens pour Emu3.5 contre 150 milliards de tokens de pré-entraînement).
• Préservation de la Qualité :
  • LlamaGen : FlashAR-L atteint un FID de 3,16 et un IS de 289,0, surpassant BlockDiffusion (FID 4,55) et même dépassant NAR-L (entraîné à partir de zéro) en Score d'Inception.
  • Emu3.5 : FlashAR réduit les étapes de décodage de 1024 à 63. Le score GenEval baisse de manière négligeable (80,48 → 80,29), tout en surpassant la base sur les métriques "Couleurs" et "Position".
• Convergence : FlashAR converge significativement plus vite que BlockDiffusion, atteignant un FID supérieur en 25 époques contre 75 époques pour BlockDiffusion.

Importance et Revendications

L'article revendique que FlashAR comble avec succès le fossé entre l'efficacité du décodage parallèle et la fidélité des modèles autoregressifs pré-entraînés. En minimisant les modifications de l'objectif d'entraînement original, FlashAR préserve les puissantes capacités génératives des modèles à grande échelle existants. Les auteurs soulignent que cette approche offre une solution pratique pour la génération d'images en temps réel et haute résolution, sans les coûts prohibitifs associés à l'entraînement de nouvelles architectures parallèles à partir de zéro ni la dégradation de qualité souvent observée dans les méthodes d'adaptation modifiant l'objectif.

Le travail démontre que le décodage parallèle diagonal est un paradigme viable et efficace pour la génération autoregressive, à condition que l'architecture du modèle soit soigneusement adaptée pour respecter les biais inhérents du modèle pré-entraîné grâce au branchement intermédiaire et à la fusion dynamique.