LGQ: Learning Discretization Geometry for Scalable and Stable Image Tokenization

Die Arbeit stellt LGQ (Learnable Geometric Quantization) vor, einen diskreten Bild-Tokenisierer, der durch das Erlernen der Diskretisierungsgeometrie mittels temperaturgesteuerter weicher Zuordnungen und spezieller Regularisierer eine stabile Optimierung und eine effiziente, ausgewogene Codebuch-Nutzung ermöglicht, wodurch bei deutlich reduzierter Anzahl aktiver Einträge eine überlegene Rekonstruktionsqualität im Vergleich zu bestehenden Methoden erreicht wird.

Das Wesentliche

Stell dir vor, du möchtest ein riesiges, hochauflösendes Foto (wie ein Bild von ImageNet) in eine Sprache übersetzen, die ein Computer verstehen und verarbeiten kann. Aber der Computer mag keine riesigen, chaotischen Datenmengen. Er braucht eine Art „Kurzcode" oder „Stempel", um das Bild zu speichern.

Das ist das Problem, das die Forscher mit ihrer neuen Methode namens LGQ lösen wollen. Hier ist die Erklärung in einfachen Worten:

Das Problem: Der „Wahlkampf" der Stempel

Stell dir vor, du hast einen riesigen Koffer voller verschiedener Stempel (das ist die Codebook oder das Wörterbuch des Computers). Jeder Stempel steht für ein kleines Teilchen eines Bildes (z. B. „ein Stück blauer Himmel" oder „ein Stück grünes Gras").

• Der alte Weg (VQ): Wenn das Computer-System ein Bild sieht, sucht es den einen Stempel, der am besten passt, und drückt ihn auf das Bild. Das Problem dabei: Es sucht nur nach dem „perfekten" Treffer. Wenn der Koffer 10.000 Stempel hat, benutzt das System oft nur die ersten 100, weil diese einfach am besten funktionieren. Die anderen 9.900 Stempel verstauben im Koffer und werden nie benutzt. Das nennt man „Kollaps". Es ist, als würdest du in einem Supermarkt nur die ersten drei Regale kaufen und den Rest des Ladens ignorieren, obwohl du eigentlich mehr Auswahl bräuchtest.
• Der starre Weg (FSQ): Andere Methoden versuchen, das Problem zu lösen, indem sie den Koffer in ein starres Raster legen. Jeder Stempel hat einen festen Platz. Das verhindert, dass Stempel verstauben, aber es ist unflexibel. Es ist wie ein Gitter, das nicht auf die Form der Dinge passt. Wenn das Bild eine krumme Kurve hat, passt das starre Raster nicht gut.

Die Lösung: LGQ – Der lernende Kleber

Die Forscher haben LGQ (Learnable Geometric Quantization) erfunden. Stell dir LGQ nicht als starren Stempel, sondern als einen intelligenten, lernfähigen Kleber vor.

1. Weiches Anhaften statt hartes Drücken:
   Beim Training (während das System lernt) ist LGQ nicht so streng. Es sagt nicht: „Nur Stempel Nr. 5 passt!" Stattdessen sagt es: „Stempel Nr. 5 passt zu 80 %, Nr. 6 zu 15 % und Nr. 7 zu 5 %."

   • Die Analogie: Stell dir vor, du versuchst, einen Ball in einen Korb zu werfen. Beim alten System musst du genau in den Korb treffen, sonst ist es ein Fehlwurf. Bei LGQ darfst du den Ball in die Nähe werfen, und das System lernt aus der Nähe, wie man es besser macht. Alle Stempel im Korb bekommen eine Chance, „mitzureden" und zu lernen.

2. Der Temperatur-Regler:
   LGQ benutzt einen „Temperatur-Regler".

   • Heiß (am Anfang): Alles ist weich und verschwommen. Alle Stempel werden ein bisschen benutzt. Das System lernt, welche Stempel gut funktionieren.
   • Kalt (am Ende): Das System kühlt ab. Die Unsicherheit verschwindet. Am Ende wählt es wieder einen einzigen, klaren Stempel aus – aber jetzt ist dieser Stempel der beste für diese Aufgabe, weil er durch das Lernen perfekt angepasst wurde.

3. Der faire Verteiler:
   LGQ hat zwei geheime Werkzeuge, um sicherzustellen, dass niemand benachteiligt wird:

   • Der „Entschlossenheits-Check": Er sorgt dafür, dass das System am Ende doch eine klare Entscheidung trifft (nicht nur „vielleicht").
   • Der „Gerechtigkeits-Check": Er bestraft das System, wenn es nur die gleichen 10 Stempel benutzt. Er zwingt das System, auch die anderen Stempel im Koffer zu nutzen, damit der ganze Koffer voll und effizient ist.

Warum ist das so toll?

Stell dir vor, du hast einen riesigen Werkzeugkasten mit 10.000 Werkzeugen.

• Alte Methoden: Du benutzt nur 500 davon, weil die anderen zu weit weg liegen oder du nicht weißt, wie man sie benutzt.
• LGQ: Du lernst, wie man alle 10.000 Werkzeuge benutzt, aber du benutzt sie nur dort, wo sie wirklich gebraucht werden.

Das Ergebnis:
LGQ kann Bilder fast genauso gut (oder sogar besser) wiederherstellen wie die alten Methoden, aber es braucht viel weniger aktive Stempel, um das zu tun. Es ist effizienter. Es ist wie ein Künstler, der mit weniger Farben ein besseres Bild malt, weil er genau weiß, wo er welche Farbe hinsetzen muss, statt einfach alle Farben wild auf die Leinwand zu schütten.

Zusammenfassung in einem Satz

LGQ ist wie ein lernender Übersetzer, der nicht starr nach einem Wörterbuch sucht, sondern die Bedeutung der Wörter (der Bildteile) versteht und die besten Wörter (Stempel) dynamisch und fair verteilt, um das Bild mit weniger Aufwand und besserer Qualität zu speichern.

Technische Zusammenfassung

1. Problemstellung

Diskrete Bild-Tokenisierung ist ein zentraler Engpass für skalierbare visuelle Generierungsmodelle (z. B. Text-zu-Bild-Modelle). Ein Tokenizer muss kompakt genug für effiziente latente Priors sein, gleichzeitig aber semantisch wichtige Strukturen bewahren und die diskrete Kapazität effektiv nutzen.

Es besteht ein persistenter Zielkonflikt zwischen bestehenden Ansätzen:

• Vektor-Quantisierung (VQ): Lernt flexible Geometrien, leidet aber oft unter verzerrter Optimierung („biased straight-through optimization"), einer Unterauslastung des Codebuchs und einem Zusammenbruch der Repräsentation („representation collapse"), insbesondere bei großen Vokabulargrößen. Nur ausgewählte Codes erhalten Gradientenupdates, während ungenutzte Einträge eingefroren bleiben.
• Strukturierte skalare Quantisierung (z. B. FSQ): Erreicht eine stabile und nahezu vollständige Auslastung durch feste, achsenparallele Gitter. Diese Starrheit führt jedoch zu einer ineffizienten Kapazitätszuweisung, da reale latente Verteilungen oft heterogen, anisotrop und datensatzabhängig sind. Feste Gitter passen sich nicht an die natürliche Datenmannigfaltigkeit an.

Das Ziel ist ein Tokenizer, der die geometrische Flexibilität von VQ mit der Stabilität strukturierter Quantierer vereint, ohne auf heuristische Tricks (wie Codebook-Reseeding) angewiesen zu sein.

2. Methodik: Learnable Geometric Quantization (LGQ)

Die Autoren stellen LGQ vor, einen diskreten Bild-Tokenizer, der die Diskretisierungsgeometrie end-to-end lernt.

Kernmechanismen:

• Weiche Zuweisungen (Soft Assignments): LGQ ersetzt die harte „Nearest-Neighbor"-Suche durch temperaturgesteuerte weiche Zuweisungen basierend auf einer Gibbs-Verteilung (Softmax über euklidische Distanzen).
  • Formel: $p_{t,k} \propto \exp(-\|z_{e,t} - c_k\|^2 / \tau)$
  • Diese Zuweisungen entsprechen den posterior responsibilities eines isotropen Gaußschen Mischmodells und minimieren ein variationsbasiertes Freie-Energie-Objektiv (erwartete Distanz + Entropie).
• Straight-Through Estimator (STE): Während des Trainings sind die Zuweisungen differenzierbar, sodass Gradienten zu allen Codebook-Einträgen fließen. Beim Inferenzschritt werden die weichen Zuweisungen durch eine harte Auswahl (argmax) ersetzt, wobei der Gradient über den STE erhalten bleibt.
• Theoretische Konvergenz: Es wird bewiesen, dass die weichen Zuweisungen im Grenzwert einer Temperatur $\tau \to 0$ gegen eine harte Nearest-Neighbor-Quantisierung konvergieren. Zudem wird die Lipschitz-Stabilität der Zuweisungen gezeigt.

Regularisierung zur Vermeidung von Zusammenbruch:
Um eine vertrauenswürdige Diskretisierung und eine robuste Skalierbarkeit zu gewährleisten, kombiniert LGQ zwei Regularisierer:

1. Token-Level-Peakedness-Regularizer: Bestraft hohe Entropie in den Zuweisungen und fördert vertrauensvolle (nahezu One-Hot) Selektionen, ohne die Glätte während des Trainings zu verbieten.
2. Globaler Nutzungs-Regularizer: Minimiert die quadrierte $L_2$-Norm der empirischen Code-Nutzungsverteilung. Dies bestraft eine konzentrierte Zuweisung auf wenige Codes und fördert eine ausgewogene Nutzung des gesamten Codebooks.

3. Wichtige Beiträge

1. Einheitlicher Rahmen (LGQ): Ein diskreter Tokenizer, der die Diskretisierungsgeometrie über temperaturgesteuerte Gibbs-Zuweisungen lernt und so die Lücke zwischen VQ und FSQ schließt.
2. Theoretische Fundierung: Bereitstellung einer variationsbasierten Freie-Energie-Formulierung, Konvergenzbeweise zur harten Quantisierung, Lipschitz-Stabilitätsnachweise und Regularisierungstechniken gegen Collapse.
3. Empirische Validierung: Demonstration auf ImageNet, dass LGQ bei verschiedenen Vokabulargrößen wettbewerbsfähige Rate-Distortion-Leistung bei stabileren Trainingsdynamiken und besserer Codebook-Auslastung erzielt.

4. Ergebnisse

Die Experimente wurden auf ImageNet (128x128) mit einem VQGAN-ähnlichen Encoder-Decoder-Backbone durchgeführt.

• Rekonstruktionsqualität: LGQ erreicht die beste Rekonstruktionsqualität (niedrigster rFID von 110,64 im Vergleich zu FSQ: 125,56 und VQ: 121,26) sowie die besten SSIM- und LPIPS-Werte.
• Effizienz und Auslastung:
  • Im Gegensatz zu FSQ und SimVQ, die das Codebook fast vollständig ausnutzen (nahezu 100%), nutzt LGQ nur etwa 50% der aktiven Einträge (8.199 von 16.384), erzielt aber dennoch bessere oder vergleichbare Ergebnisse.
  • Dies zeigt, dass LGQ eine effiziente, datenangepasste Diskretisierungsgeometrie lernt, anstatt auf eine Sättigung des Codebooks angewiesen zu sein.
• Skalierbarkeit: Auch bei sehr großen Vokabulargrößen (K=65.536) behält LGQ eine stabile Optimierung und eine günstige Rate-Distortion-Balance bei, während andere Methoden (wie VQ) unter einem starken Zusammenbruch leiden.
• Geometrische Anpassung: Visualisierungen (UMAP) zeigen, dass LGQ-Codebook-Einträge strukturiert und heterogen wandern, um sich an die hochdichten Regionen der latenten Verteilung anzupassen, anstatt statische Gitter zu verwenden.

5. Bedeutung und Fazit

LGQ adressiert das fundamentale Problem des „Representation Collapse" bei skalierbaren Tokenizern, indem es die Diskretisierung als Lernproblem der Geometrie formuliert.

• Systemische Vorteile: LGQ kann als Drop-in-Ersatz für existierende diskrete Tokenizer dienen. Es eliminiert die Notwendigkeit für Commitment-Losses und Codebook-Reseeding und ermöglicht das effektive Training großer Vokabulare.
• Paradigmenwechsel: Die Arbeit zeigt, dass eine hohe Codebook-Auslastung allein kein Maß für Qualität ist. Effektive Tokenisierung erfordert die Zuweisung diskreter Kapazität gemäß der empirischen latenten Verteilung.
• Zukunftsaussichten: Durch die Prinzipien der geometrischen Lernbarkeit bietet LGQ eine solide Grundlage für effiziente diskrete Repräsentationen in zukünftigen multimodalen und hochdimensionalen Anwendungen (z. B. Video).

Zusammenfassend bietet LGQ einen stabilen, skalierbaren und theoretisch fundierten Ansatz, der die Flexibilität von Vektor-Quantisierung mit der Robustheit strukturierter Quantierer vereint.