LGQ: Learning Discretization Geometry for Scalable and Stable Image Tokenization

Dit paper introduceert LGQ, een nieuwe discrete beeld-tokenisatie-methode die door het leren van discretisatiegeometrie via temperatuur-gestuurde zachte toewijzingen en variational free-energy-optimalisatie, een stabielere en efficiëntere codering bereikt dan bestaande methoden met een betere balans tussen reconstructiekwaliteit en codeboekgebruik.

De kern

Stel je voor dat je een gigantische bibliotheek met miljoenen boeken (beelden) wilt opslaan, maar je hebt maar een heel klein kastje om ze in te zetten. Hoe pak je dat aan?

In de wereld van kunstmatige intelligentie (AI) proberen computers beelden te comprimeren tot een reeks "woorden" of tokens. Dit is als het vertalen van een complex schilderij naar een simpele lijst met woorden, zodat een andere AI het kan begrijpen en later weer kan "tekenen".

Het probleem is dat de huidige methoden om deze "woorden" te kiezen, vaak vastlopen. Ze zijn ofwel te star, ofwel te chaotisch.

Hier is hoe dit nieuwe papier, LGQ, dat probleem oplost, uitgelegd met een paar simpele analogieën:

1. Het Probleem: De Twee Slechte Opties

Stel je voor dat je een groep mensen (de AI) moet leren om kleuren te benoemen. Je hebt een palet met duizenden verfkleuren (de codebook).

• Optie A (De Starre Lijst): Je geeft de mensen een vaste lijst met 100 kleuren. Als ze een kleur zien die er niet bij staat, moeten ze de dichtstbijzijnde kiezen.
  • Het nadeel: Als de lijst niet past bij de echte wereld (bijvoorbeeld: er zijn geen tinten blauw, maar wel veel groen), dan wordt de lijst inefficiënt. Veel kleuren worden nooit gebruikt, en de mensen worden verward. Dit is wat methoden als FSQ doen: ze gebruiken een star rooster. Het werkt stabiel, maar het is niet slim genoeg om zich aan te passen.
• Optie B (De Chaos): Je laat de mensen zelf hun eigen lijst maken. Ze mogen elke kleur kiezen die ze willen.
  • Het nadeel: In het begin kiezen ze allemaal dezelfde paar populaire kleuren (bijvoorbeeld alleen rood en blauw). De andere duizenden kleuren op hun lijst worden nooit gebruikt en vergeten. Dit heet "collapse" (instorting). De AI wordt lui en gebruikt maar een klein deel van zijn potentieel. Dit is wat de oude VQ-methoden doen.

2. De Oplossing: LGQ (De Slimme Leraar)

De auteurs van dit papier hebben LGQ (Learnable Geometric Quantization) bedacht. Dit is als een slimme leraar die de lijst met kleuren leert terwijl de les doorgaat.

In plaats van een vaste lijst of een volledig vrije chaos, doet LGQ het volgende:

• De "Warme" Keuze: In het begin is de keuze van een kleur "zacht" en onzeker. Stel je voor dat een leerling twijfelt tussen "donkerblauw" en "hemelsblauw". In plaats van direct te zeggen "Kies blauw!", zegt de leraar: "Ik denk dat het 60% blauw is en 40% paars."
  • Waarom is dit slim? Omdat de AI nu weet dat beide kleuren belangrijk zijn, krijgt hij feedback over beide. Niemand wordt genegeerd. Dit voorkomt dat de lijst instort.
• De "Koude" Keuze: Naarmate de training vordert, wordt de leraar strenger. De twijfel verdwijnt. Uiteindelijk kiest de AI weer één duidelijke kleur (bijvoorbeeld "blauw"), maar dan wel de perfecte blauwe kleur die bij de rest van de lijst past.
• De Slimme Verdeling: De AI leert dat hij niet alle kleuren even vaak moet gebruiken. Als er veel groene bossen zijn, leert hij dat hij meer groene "woorden" nodig heeft. Als er weinig paarse bloemen zijn, gebruikt hij minder paarse woorden. Hij verdeelt zijn ruimte slim over de hele lijst, in plaats van alles vol te proppen met dezelfde dingen.

3. Het Resultaat: Meer met Minder

Het mooie aan LGQ is dat het efficiënter is dan de oude methoden.

• Oude methoden: Moesten vaak hun hele lijst van 16.000 kleuren gebruiken om een goed plaatje te maken. Het was als een overvolle koffer waar je alles in probeert te proppen, ook al past het niet.
• LGQ: Haalt hetzelfde (of zelfs betere) resultaat, maar gebruikt maar ongeveer de helft van de lijst (ongeveer 8.000 kleuren).
  • De metafoor: Het is alsof je een koffer inpakt. De oude methoden gooien alles erin, ook dubbelingen en onnodige spullen. LGQ leert precies welke spullen je echt nodig hebt en hoe je ze slim stapelt. Je hebt minder ruimte nodig, maar je kunt nog steeds alles meenemen.

Samenvattend

Dit papier introduceert een nieuwe manier om beelden voor AI te "vertalen". In plaats van te kiezen tussen een starre, onhandige lijst of een chaotische lijst die instort, leert LGQ zelf hoe de lijst eruit moet zien.

Het is als een levendige, aanpasbare kaart in plaats van een stenen muur. Door zacht te beginnen en geleidelijk harder te worden, leert de AI de perfecte manier om zijn "woorden" te verdelen. Het resultaat? Beelden die er scherper uitzien, met minder rekenkracht en minder "dode" plekken in de geheugenlijst.

Kortom: LGQ maakt de AI slimmer, efficiënter en minder lui bij het onthouden van beelden.

Technische samenvatting

1. Het Probleem

Discrete beeldtokenisatie is een cruciale schakel in schaalbare visuele generatiemodellen (zoals VQ-VAE's en diffusion-modellen). Bestaande methoden kampen echter met een fundamenteel compromis tussen flexibiliteit en stabiliteit:

• Vector Quantization (VQ): Leren flexibele geometrieën (Voronoi-partities), maar lijden vaak aan representatie-inzakking (collapse). Hierbij worden slechts een klein deel van de codebook-entries gebruikt, terwijl de rest "dood" blijft. Dit probleem verergert bij het vergroten van de vocabulairegrootte.
• Gestructureerde Scalar Quantization (zoals FSQ): Bereiken stabiele, bijna volledige gebruik van het codebook door vaste roosters te gebruiken. Echter, deze vaste geometrieën passen zich niet aan aan de heterogene, dataset-afhankelijke statistieken van latente ruimtes, wat leidt tot inefficiënte capaciteitsbenutting.

Er is behoefte aan een tokenisator die de geometrie van de discretisatie leert (aanpasbaar aan de data) maar toch stabiel blijft zonder heuristieken zoals codebook-resampling of commitment losses.

2. Methodologie: Learnable Geometric Quantization (LGQ)

De auteurs introduceren LGQ, een framework dat de discretisatiegeometrie end-to-end leert door een overgang van harde naar zachte toewijzingen te maken.

Kernmechanismen:

• Temperatuur-gestuurde zachte toewijzingen: In plaats van een harde "nearest-neighbor" zoekopdracht, berekent LGQ toewijzingskansen ($p_{t,k}$) gebaseerd op een Gibbs-verdeling (softmax) over de Euclidische afstanden tot de codebook-entries.
  $$p_{t,k} \propto \exp(-\|z_{e,t} - c_k\|^2 / \tau)$$
  Hierbij is $\tau$ een temperatuurparameter. Deze zachte toewijzingen corresponderen met posterior-verantwoordelijkheden in een isotroop Gaussisch mengselmodel.
• Straight-Through Estimator (STE): Tijdens training zijn de toewijzingen differentieerbaar, waardoor gradiënten naar alle codebook-entries stromen (vermijdt de "dead code" problemen van VQ). Tijdens inferentie wordt de zachte toewijzing omgezet in een harde discrete index via een STE, zodat het model toch discrete tokens produceert.
• Convergentie naar Harde Quantisatie: De auteurs bewijzen dat wanneer de temperatuur $\tau \to 0$, de zachte toewijzingen wiskundig convergeren naar een harde nearest-neighbor quantisatie.

Regularisatie voor Stabiliteit:
Om te voorkomen dat het model toch instort of onbalans ontwikkelt, introduceert LGQ twee regularisatoren:

1. Peakedness Regularizer ($L_{peak}$): Straft hoge entropie in de toewijzingen af, waardoor het model wordt gestimuleerd om vertrouwen te hebben in zijn keuzes (nabij one-hot), zonder de gladheid tijdens training te verliezen.
2. Global Usage Regularizer ($L_{bins}$): Minimaliseert de $L_2$-norm van de empirische verdeling van code-gebruik. Dit straft concentratie op weinig codes af en stimuleert een gebalanceerd gebruik van het hele codebook.

3. Belangrijkste Bijdragen

1. Unificatie van VQ en FSQ: LGQ overbrugt de kloof tussen de geometrische flexibiliteit van VQ en de stabiliteit van gestructureerde quantizers door de discretisatiegeometrie te laten leren in plaats van deze vast te leggen.
2. Principiële Formulering: De methode is gebaseerd op variational free-energy minimalisatie, met theoretische garanties voor convergentie naar harde quantisatie en Lipschitz-stabiliteit.
3. Geen Heuristieken: Het elimineert de noodzaak voor complexe heuristieken zoals codebook-resampling, commitment losses of reservoir sampling die vaak nodig zijn bij VQ.
4. Empirische Validatie: Uitgebreide experimenten op ImageNet tonen aan dat LGQ stabiel blijft bij schaling van de vocabulairegrootte.

4. Resultaten

De experimenten zijn uitgevoerd met een VQGAN-achtige backbone op ImageNet (128x128) met verschillende vocabulairegroottes ($K=16,384$ en $K=65,536$).

• Reconstructiekwaliteit: LGQ presteert beter dan bestaande methoden. Bij $K=16,384$ bereikt LGQ een rFID van 110.64, wat een verbetering is ten opzichte van FSQ (125.56), VQ (121.26) en SimVQ (117.77). Ook scoort LGQ het hoogst op SSIM (structuur) en LPIPS (perceptuele kwaliteit).
• Efficiëntie en Utilisatie:
  • LGQ bereikt deze superieure kwaliteit met substantieel minder actieve codebook-entries (ongeveer 50% van het totaal, oftewel ~8.199 actieve codes) vergeleken met FSQ en SimVQ die bijna 100% van het codebook gebruiken.
  • Dit resulteert in een 49,96% lagere effectieve representatie-rate vergeleken met FSQ, met een betere rFID.
• Rate-Distortion Trade-off: LGQ leert een geometrie die beter aansluit bij de empirische latente verdeling. Het activeert een compacte, maar goed uitgelijnde subset van codes die de volledige latente manifold dekt, in plaats van het hele rooster te vullen (zoals bij FSQ).
• Stabiliteit: De training vertoont gladde convergentie zonder de instabiliteit of "collapse" die vaak voorkomt bij VQ bij grote vocabulairegroottes.

5. Betekenis en Conclusie

LGQ vertegenwoordigt een paradigmaverschuiving in discrete representatieleren. Het paper toont aan dat gebruik (utilization) op zich geen voldoende maatstaf is voor kwaliteit; een tokenizer moet discrete capaciteit efficiënt toewijzen aan de gebieden waar de data zich werkelijk bevindt.

Door de discretisatiegeometrie te leren via differentieerbare zachte toewijzingen, biedt LGQ een robuustere en schaalbaarder oplossing voor de tokenisatie van beelden. Dit maakt het een ideale "drop-in" vervanging voor bestaande tokenizers in autoencoder-architecturen en vormt een sterke basis voor toekomstige generatieve modellen (zoals MaskGIT of diffusion-modellen) die vertrouwen op discrete tokens. De methode lost het probleem van representatie-inzakking op zonder de flexibiliteit van vectorquantisatie op te geven.