LGQ: Learning Discretization Geometry for Scalable and Stable Image Tokenization

Ce papier présente LGQ, un tokeniseur d'images discret qui apprend end-to-end la géométrie de discrétisation via des affectations souples contrôlées par température, permettant d'obtenir une utilisation équilibrée du codebook et une meilleure fidélité de reconstruction que les méthodes existantes tout en réduisant le nombre de codes actifs nécessaires.

L'essentiel

🎨 Le Problème : Comment résumer un tableau sans le gâcher ?

Imaginez que vous êtes un artiste qui veut décrire un magnifique tableau à un ami qui ne peut pas le voir. Vous avez deux options :

1. Décrire chaque pixel (trop long, impossible à retenir).
2. Utiliser un code secret : "C'est un ciel bleu, une montagne verte, un soleil jaune". C'est plus court, mais il faut que votre ami comprenne le code.

En intelligence artificielle (IA), c'est la même chose. Pour générer des images, les ordinateurs doivent d'abord transformer les pixels en une liste de "mots" ou de jetons (tokens) discrets. C'est ce qu'on appelle la tokenisation.

Le problème actuel, c'est que les méthodes existantes sont soit trop rigides, soit trop désordonnées :

• La méthode rigide (FSQ) : C'est comme si vous aviez une grille de cases fixes. Peu importe où se trouve le nuage dans le ciel, vous devez le coller dans la case "Nuage 1", "Nuage 2", etc. C'est stable, mais si le nuage est à cheval entre deux cases, l'image devient floue.
• La méthode désordonnée (VQ classique) : C'est comme avoir un sac de 10 000 étiquettes. L'IA essaie de coller chaque nuage sur l'étiquette la plus proche. Mais souvent, l'IA devient paresseuse : elle n'utilise que 10 étiquettes sur les 10 000, laissant les 9 990 autres dans un coin poussiéreux. C'est ce qu'on appelle l'effondrement (collapse) : le système ne profite pas de sa capacité totale.

💡 La Solution : LGQ (La Géométrie Apprenante)

Les auteurs de cet article proposent une nouvelle méthode appelée LGQ. Voici l'analogie pour comprendre comment ça marche :

Imaginez que vous devez ranger des fruits dans des paniers.

• L'ancienne méthode (VQ) : Vous avez des paniers fixes. Si une pomme est à 1 cm du panier "Pommes", elle y va. Si elle est à 1,1 cm, elle va dans le panier "Poires". C'est brutal. Et souvent, les paniers "Poires" restent vides car l'IA préfère tout mettre dans "Pommes".
• La méthode LGQ : Au lieu de paniers fixes, imaginez que les paniers sont élastiques et mobiles.
  1. L'approche douce (Soft Assignment) : Au début, quand l'IA regarde un fruit, elle ne dit pas "C'est une pomme !". Elle dit : "C'est à 80% une pomme, 20% une poire". Elle distribue la responsabilité. Cela permet à tous les paniers d'apprendre et de se déplacer légèrement pour mieux attraper les fruits.
  2. Le refroidissement (Temperature) : Au fur et à mesure de l'entraînement, on "refroidit" le système. Les pourcentages deviennent plus tranchés. À la fin, on décide fermement : "C'est une pomme". Mais grâce à la phase douce, les paniers se sont déjà bien placés pour être efficaces.
  3. L'équilibre (Regularization) : LGQ ajoute une petite règle : "Si un panier est trop plein et un autre vide, on pousse un peu le panier vide vers les fruits". Cela force l'IA à utiliser tous les paniers de manière équitable, sans en laisser de côté.

🚀 Pourquoi c'est génial ? (Les Résultats)

Grâce à cette astuce, LGQ obtient des résultats incroyables :

1. Moins de gaspillage : Alors que les autres méthodes utilisent presque tous les paniers (ce qui est inefficace) ou très peu (ce qui est pauvre), LGQ trouve le juste milieu. Elle utilise environ 50% de ses paniers pour faire un travail aussi bon, voire meilleur, que les autres qui en utilisent 100%. C'est comme réussir à cuisiner un festin avec la moitié des ingrédients, mais en choisissant les meilleurs.
2. Plus stable : L'entraînement ne "casse" pas. L'IA ne perd pas ses repères.
3. Plus précis : Les images reconstruites sont plus nettes et fidèles à l'original.

🌟 En résumé

Pensez à LGQ comme à un chef cuisinier génial qui apprend à utiliser sa cuisine :

• Il ne se contente pas d'utiliser les mêmes 3 casseroles pour tout faire (comme l'IA paresseuse).
• Il ne force pas chaque ingrédient dans une case rigide (comme la grille fixe).
• Il apprend la forme exacte de ses ingrédients et ajuste la position de ses casseroles pour qu'elles s'adaptent parfaitement.

Le résultat ? Une cuisine (un modèle d'IA) qui est plus rapide, plus économe en énergie, et qui produit des plats (des images) délicieux, même avec un grand nombre d'ingrédients (vocabulaire) à gérer.

C'est une avancée majeure pour rendre les IA génératrices d'images plus intelligentes et plus efficaces !

Résumé technique

1. Problématique

La tokenisation d'images discrètes est un goulot d'étranglement central pour la génération visuelle à grande échelle. Les modèles génératifs modernes (comme les transformeurs) nécessitent des représentations latentes discrètes et compactes. Cependant, les quantificateurs existants souffrent d'un compromis persistant :

• Vector Quantization (VQ) classique : Offre une flexibilité géométrique (apprentissage d'un codebook), mais souffre de problèmes d'optimisation majeurs : biais dans les gradients (mise à jour uniquement du voisin le plus proche), sous-utilisation du codebook (collapse de représentation) et instabilité lors de l'augmentation de la taille du vocabulaire.
• Quantification Scalaire Structurée (ex: FSQ) : Garantit une utilisation complète et stable du codebook par construction, mais repose sur des géométries de discrétisation fixes (grilles scalaires). Cela suppose que les dimensions latentes sont alignées sur les axes et homogènes, ce qui est rarement le cas dans les données réelles, entraînant une allocation inefficace de la capacité de représentation.

L'objectif est de concevoir un tokeniseur capable d'apprendre sa propre géométrie de discrétisation de manière end-to-end, tout en assurant une optimisation stable et une utilisation équilibrée du codebook, sans recourir à des heuristiques complexes (comme le réensemencement de codebook).

2. Méthodologie : LGQ (Learnable Geometric Quantization)

Les auteurs proposent LGQ, un cadre unifié qui remplace l'affectation "voisin le plus proche" (hard nearest-neighbor) par des affectations douces contrôlées par la température, permettant un apprentissage différentiable.

A. Assignations Soft-to-Hard

Au lieu d'une affectation binaire rigide, LGQ calcule la probabilité d'affectation d'un vecteur latent $z$ à un codeword $c_k$ via une distribution de Gibbs (softmax) basée sur la distance euclidienne :
$$p_{t,k} \propto \exp(-\|z_{e,t} - c_k\|^2 / \tau)$$
où $\tau$ est un paramètre de température.

• Théorème de convergence : Lorsque $\tau \to 0$, cette distribution tend vers une affectation "one-hot" (le voisin le plus proche), garantissant que l'inférence reste discrète.
• Estimateur Straight-Through (STE) : Pendant l'entraînement, les gradients traversent l'opération de sélection dure en utilisant la moyenne pondérée des codewords (affectation douce) pour le passage arrière, permettant la mise à jour de tout le codebook.

B. Objectif Variational et Régularisation

L'affectation douce correspond au minimiseur d'une fonctionnelle d'énergie libre variationnelle (compromis entre distorsion attendue et entropie). Pour éviter l'effondrement du codebook et encourager une utilisation équilibrée, LGQ introduit deux régularisateurs dans la fonction de perte :

1. Régulariseur de "Peakedness" (Pointe) : Pénalise les affectations à haute entropie, encourageant des choix confiants (proches de one-hot) tout en maintenant la douceur nécessaire à l'entraînement.
2. Régulariseur d'Utilisation Globale : Minimise la norme L2 carrée de la distribution marginale d'utilisation des codes. Cela pénalise l'allocation concentrée sur quelques codes et encourage une utilisation uniforme de tout le codebook.

C. Architecture

LGQ s'intègre dans une architecture VQGAN (encodeur-décodeur). L'encodeur produit des latents continus qui sont discrétisés par LGQ avant d'être reconstruits par le décodeur. La géométrie de discrétisation (les centres du codebook) est apprise conjointement avec l'encodeur et le décodeur.

3. Contributions Clés

1. Unification VQ/FSQ : LGQ comble le fossé entre la flexibilité géométrique du VQ et la stabilité de la quantification scalaire structurée (FSQ) en apprenant la géométrie de discrétisation directement à partir des données.
2. Fondement Théorique : La méthode est ancrée dans une formulation d'énergie libre variationnelle, avec des garanties théoriques sur la convergence vers la quantification dure et la stabilité de Lipschitz des gradients.
3. Stabilité sans Heuristiques : Contrairement aux méthodes VQ classiques qui nécessitent des mécanismes de secours (reseeding, commitment loss), LGQ évite l'effondrement du codebook grâce à ses régularisateurs et à son mécanisme d'assignation douce.
4. Efficacité des Ressources : LGQ démontre qu'une haute fidélité de reconstruction peut être atteinte avec une utilisation effective du codebook bien inférieure à celle des méthodes saturées.

4. Résultats Expérimentaux

Les expériences ont été menées sur ImageNet (128x128) avec une architecture VQGAN contrôlée, en comparant LGQ à des baselines fortes (VQ, FSQ, LFQ, SimVQ).

• Qualité de Reconstruction :
  • LGQ obtient le meilleur rFID (110.64) comparé à FSQ (125.56) et VQ (121.26).
  • Il atteint également les meilleurs scores SSIM et LPIPS, indiquant une meilleure fidélité structurelle et perceptive.
• Utilisation du Codebook et Efficacité :
  • Alors que FSQ et SimVQ utilisent presque 100% du codebook (16k entrées), LGQ n'en active que ~50% (8 199 entrées) pour obtenir une meilleure qualité.
  • LGQ atteint un taux de représentation effectif (effective rate) 49,96% plus faible que FSQ pour une meilleure qualité (rFID amélioré de 11,88%).
• Évolutivité (Scaling) :
  • Même avec un vocabulaire beaucoup plus grand ($K=65,536$), LGQ maintient une optimisation stable et une meilleure qualité de reconstruction que les autres méthodes, avec une utilisation efficace des codes (14 716 actifs sur 65k).
• Dynamique d'Apprentissage :
  • Les visualisations UMAP montrent que LGQ apprend à concentrer les codes actifs sur les régions de haute densité de la distribution latente, contrairement aux méthodes qui couvrent uniformément l'espace.
  • L'analyse de la dérive des centres (center drift) confirme que la géométrie du codebook s'adapte activement et de manière structurée aux données durant l'entraînement.

5. Signification et Impact

Ce travail remet en question le paradigme selon lequel une utilisation maximale du codebook est nécessaire pour une bonne performance. Il démontre que l'efficacité de la tokenisation dépend de l'alignement entre la géométrie de discrétisation et la distribution latente empirique, plutôt que de la simple saturation du vocabulaire.

• Impact Systémique : LGQ peut être utilisé comme un remplacement "plug-and-play" dans les architectures d'autoencodeurs existantes, éliminant le besoin de réglages fins complexes pour stabiliser l'entraînement.
• Perspectives Futures : Cette approche ouvre la voie à des tokeniseurs pour des données multimodales (vidéo, audio) et des modèles génératifs à base de transformeurs, où la stabilité de l'optimisation et l'efficacité de la capacité discrète sont critiques.

En résumé, LGQ transforme la tokenisation d'un problème heuristique en un problème d'inférence géométrique principiel, offrant une solution scalable et stable pour la prochaine génération de modèles génératifs.