Résumé Technique : MergeTok

Énoncé du Problème

Les tokeniseurs visuels actuels pour la génération d'images sont divisés en deux familles distinctes, chacune souffrant de limitations structurelles complémentaires :

1. Tokeniseurs basés sur des VAE continus : Bien qu'offrant une reconstruction de haute fidélité, ils reposent sur des espaces latents dépourvus d'organisation intrinsèque. Sous une forte compression, les facteurs sémantiques deviennent entrelacés au sein de représentations latentes denses, ce qui produit des codes faiblement désenchevêtrés qui entravent le contrôle sémantique.
2. Tokeniseurs basés sur des VQ discrets : Bien qu'ils permettent la génération autorégressive, ils font face à des goulots d'étranglement d'optimisation. La nature non différentiable de la quantification entraîne des mises à jour de gradients creuses et inégales, provoquant un manque de supervision pour de nombreuses entrées du codebook. Cela augmente le risque d'effondrement du codebook et réduit l'efficacité de la représentation.

Ces limitations créent une tension persistante entre la fidélité de la reconstruction, la structure sémantique, l'efficacité et la stabilité de l'entraînement.

Méthodologie : Le Cadre MergeTok

MergeTok propose un tokeniseur à double branche unifié qui optimise conjointement la tokenisation continue (VAE) et discrète (VQ) au sein d'une architecture encodeur-décodeur unique. L'innovation centrale est l'utilisation du Fusion de Tokens (ToMe - Token Merging) comme un pont sémantique entre les deux branches.

Aperçu de l'Architecture

Étant donné une image d'entrée $X$, un encodeur CNN partagé $E_c$ produit une carte de caractéristiques $Z$. Cette carte est traitée par deux branches parallèles :

1. Branche VAE (Continue) :

   • Fusion de Tokens (Token Merging) : Un encodeur basé sur l'attention équipé de modules ToMe compresse la séquence de tokens de $L$ tokens à $K$ tokens ($K < L$) en fusionnant les tokens similaires. Ce processus génère une représentation condensée $Z^{(vae)}_K$ et une carte source $S$ qui enregistre l'assignation des tokens originaux aux groupes fusionnés.
   • Alignement Sémantique : Les tokens fusionnés sont alignés avec un modèle enseignant de type DINO gelé (équipé également de ToMe) à l'aide d'une perte de distance cosinus ($L_{align}$). Cela régularise l'espace latent continu vers des représentations sémantiquement conscientes et désenchevêtrées.
   • Reconstruction : Un décodeur hybride utilise à la fois les tokens condensés et la carte source $S$ pour "défusionner" les groupes et reconstruire l'image dans l'espace des pixels.

2. Branche VQ (Discrète) :

   • Quantification Sensible aux Groupes : Cette branche contourne l'opération ToMe lors de la passe avant afin de préserver la longueur de séquence complète ($L$ tokens). Cependant, elle utilise la carte source $S$ générée par la branche VAE comme un a priori structurel.
   • Régularisation : La carte source partitionne les $L$ tokens en $K$ groupes sémantiques. Deux pertes sensibles aux groupes sont appliquées aux assignations du codebook :
     • Diversité Intra-groupe ($L_{div}$) : Maximise l'entropie au sein de chaque groupe pour empêcher l'effondrement vers un code unique.
     • Exclusivité Inter-groupe ($L_{cons}$) : Minimise le chevauchement entre les groupes pour garantir que différents groupes sémantiques utilisent des sous-ensembles de codes distincts.
   • Reconstruction : Un décodeur VQ standard reconstruit l'image à partir des tokens quantifiés.

Stratégie d'Entraînement

• Objectif Unifié : Le modèle est entraîné de bout en bout avec une perte totale combinant la reconstruction, l'adversaire et les régulariseurs de fusion :
  $$L_{total} = \lambda_{vae}L_{VAE} + \lambda_{vq}L_{VQ} + \lambda_{ali}L_{align} + \lambda_{div}L_{div} + \lambda_{cons}L_{cons}$$
• Échantillonnage Sensible à la Granularité : Pendant l'entraînement, le ratio de fusion (et donc le nombre de tokens conservés $K$) est échantillonné dynamiquement à partir d'une distribution discrète centrée autour d'un hyperparamètre $K^*$. Cela expose l'encodeur à des largeurs de goulot d'étranglement variables, encourageant la robustesse à travers différentes granularités sémantiques.
• Inférence : À l'inférence, le tokeniseur produit toujours une séquence fixe de 256 tokens. Les générateurs en aval (tant autorégressifs que de diffusion) opèrent sur cette séquence fixe et sont agnostiques à la structure interne de fusion.

Principales Contributions

1. Tokeniseur à Double Branche Unifié : MergeTok est la première architecture à optimiser conjointement les latents continus et discrets sous un seul objectif, en utilisant la fusion de tokens comme interface sémantique pour renforcer la stabilité de l'entraînement VAE et les avantages structurels du VQ.
2. Contraintes d'Entraînement Sensibles à la Fusion : Les auteurs introduisent deux nouveaux objectifs dérivés du processus de fusion :
   • Alignement des tokens fusionnés : Améliore la structure sémantique dans l'espace latent continu.
   • Quantification sensible aux groupes : Stabilise l'entraînement VQ et améliore l'utilisation du codebook via la diversité intra-groupe et l'exclusivité inter-groupe.
3. Échantillonnage du Ratio de Fusion Sensible à la Granularité : Une stratégie d'entraînement qui échantillonne différentes granularités de tokens, permettant au modèle d'apprendre une tokenisation unifiée qui est robuste à travers différents taux de compression.

Résultats Expérimentaux

Les expériences ont été menées sur ImageNet-256 (résolution 256×256) avec un budget de tokens fixe de 256.

• Fidélité de la Reconstruction :
  • Discret (VQ) : MergeTok-BL a atteint un rFID de 0,78, surpassant de fortes bases de référence comme GigaTok-BL (0,81) et MergeVQ-GR (1,12). Avec un discriminateur DINO gelé, MergeTok-BL* a atteint 0,48 rFID.
  • Continu (VAE) : MergeTok-BL a atteint un rFID de 0,47, surpassant légèrement RAE (0,49) et battant nettement SoftVQ-L (0,61) et DC-AE-f32 (0,69).
• Qualité de la Représentation :
  • MergeTok-BL a obtenu une précision de probing linéaire de 78,3 %, la plus élevée parmi les tokeniseurs discrets listés, dépassant UniTok (70,8 %) et MergeVQ-GR (77,9 %). Cela indique une organisation sémantique supérieure.
• Génération en Aval :
  • Autorégressive (Branche VQ) : Apparié avec LlamaGen-XXL, MergeTok-BL a atteint un gFID de 1,93 (sans CFG) et 2,13 (avec CFG), surpassant GigaTok-BL (2,03).
  • Diffusion/Flow (Branche VAE) : Apparié avec SiT-XL/1, MergeTok-BL a atteint un gFID de 1,79 (sans CFG) et 1,29 (avec CFG), égalant les meilleures entrées du tableau.

Signification et Revendications

L'article affirme que MergeTok démontre qu'une architecture unique peut doter les tokeniseurs visuels à la fois d'une organisation sémantique robuste et d'une discrétisation favorable aux générateurs.

• Unification : Il résout le compromis entre la haute fidélité de reconstruction des VAE et la structure discrète requise pour la génération autorégressive.
• Pont Sémantique : En exploitant la fusion de tokens comme un a priori structurel, le modèle transfère avec succès les informations de regroupement sémantique de la branche continue pour stabiliser la branche discrète, et vice versa.
• Compatibilité : Le tokeniseur produit des séquences de 256 tokens qui sont compatibles avec les générateurs autorégressifs (ex: LlamaGen) et de diffusion/flow (ex: DiT, SiT) sans nécessiter de modifications des générateurs eux-mêmes.

Les auteurs positionnent MergeTok non pas simplement comme un outil de reconstruction, mais comme une interface unifiée qui supporte simultanément la génération autorégressive discrète, la génération de diffusion continue, et l'apprentissage de représentations sémantiquement structurées.