Technische Samenvatting: Guided End-to-End AutoRegression voor Beeldsynthese (GEAR)

1. Probleemstelling

Moderne visuele generatieve modellen, inclusief zowel autoregressieve (AR) transformers als diffusiemodellen, volgen doorgaans een twee-fasen trainingsparadigma:

1. Tokenizer Training: Een vector-gekwantiseerde (VQ) tokenizer (bijv. VQ-VAE) wordt getraind om afbeeldingen te reconstrueren en wordt vervolgens bevroren.
2. Generator Training: Een AR-generator wordt getraind op de discrete indices die door de bevroren tokenizer worden geproduceerd.

Deze ontkoppeling creëert een fundamentele mismatch: de tokenizer is puur geoptimaliseerd voor reconstructietrouw, zonder oog te hebben voor de vraag of de resulterende discrete sequentie gemakkelijk door de downstream generator gemodelleerd kan worden via causale stappen. Hoewel recent werk in diffusiemodellen (bijv. REPA-E) succesvol end-to-end gezamenlijke training heeft verkend door latente representaties af te stemmen op externe semantische encoders (zoals DINOv2), is het toepassen hiervan op discrete AR-generatie aanzienlijk moeilijker.

De kernbelemmering is niet-differentieerbaarheid. De mapping van een continue latente ruimte naar een discrete VQ-index omvat een arg max-operatie. Bijgevolg kunnen gradiënten van de AR-generator niet terugstromen naar de tokenizer. Pogingen om deze kloof te overbruggen met behulp van een Straight-Through Estimator (STE) zijn instabiel gebleken, waarbij de codebook vaak instort (waarbij de tokenizer slechts enkele dominante codes uitgeeft om voorspelbaarheid te maximaliseren ten koste van de reconstructietrouw), wat leidt tot catastrofaal falen (gFID $\approx$ 105).

2. Methodologie: GEAR

De auteurs stellen GEAR (Guided End-to-end AutoRegression) voor, een framework dat een VQ-tokenizer en een AR-generator end-to-end gezamenlijk traint zonder dat het codebook instort.

Kernmechanisme: Dual Read-Out

GEAR lost het probleem van niet-differentieerbaarheid op door een dual read-out van de codebook-toewijzing voor elke positie te introduceren:

1. Hard Branch (One-Hot):
   • Gebruikt de standaard discrete index ($q_i = \text{arg max}_k A_{ik}$).
   • Voedt discrete tokens aan het AR-model voor Next-Token Prediction (NTP) en een harde representatie-alignment loss.
   • Cruciaal: Deze tak is niet-differentieerbaar met betrekking tot de tokenizer. Gradiënten van de NTP en de harde alignment loss updaten alleen het AR-model, niet de tokenizer. Dit voorkomt dat de NTP-loss de tokenizer dwingt om te imploderen in lage-entropie codes.
2. Soft Branch (Differentiabel):
   • Gebruikt een temperatuur-gewogen softmax-interpolatie over de codebook-entries ($\pi_i = \text{softmax}(A_i / \tau)$).
   • Voedt een zachte, continue embedding aan het AR-model.
   • Cruciaal: Deze tak is volledig differentieerbaar. Het draagt een representatie-alignment loss (bijv. tegen DINOv2-features) die terugstroomt door de AR-backbone om alleen de tokenizer te updaten.

Optimalisatiestrategie

De twee modules worden bijgewerkt met gescheiden doelstellingen:

• Tokenizer ($\theta_{tok}$): Wordt bijgewerkt door de standaard VQ-reconstructieverlies ($L_{VQ}$) plus de soft alignment loss ($L^s_{align}$). De tokenizer wordt gestuurd om zijn indexverdeling te hervormen zodat deze voorspelbaarder en semantisch afgestemd is, maar niet door de predictieverlies zelf.
• AR-Generator ($\theta_{AR}$): Wordt bijgewerkt door de NTP-loss ($L_{NTP}$) plus de harde alignment loss ($L^h_{align}$).

Dit ontwerp zorgt ervoor dat de tokenizer leert om tokens uit te geven die de AR-model gemakkelijk kan voorspellen en afstemmen op semantische kenmerken, terwijl de hoge reconstructietrouw behouden blijft.

3. Belangrijkste Bijdragen

• Guided End-to-End Training: GEAR introduceert een differentieerbare soft-assignment brug die het mogelijk maakt om de representatie-alignment doelstelling van het AR-model te gebruiken om de tokenizer te sturen, waarmee de barrière van de niet-differentieerbare index wordt overwonnen die standaard STE-benaderingen fataal wordt.
• Verschoven Alignment-last: In tegenstelling tot methoden aan de diffusiezijde (bijv. REPA-E, VA-VAE) die de latente ruimte zelf meer semantisch maken (DINOv2-achtig), verschuift GEAR de alignment-last naar de AR-generator.
  • De features van de tokenizer worden minder DINOv2-achtig, waarbij de indexverdeling wordt geherorganiseerd naar een meer voorspelbaar, lager-entropie gebruik.
  • De verborgen toestanden van de AR-generator worden meer DINOv2-achtig, met name op patch-niveau, waarbij rijke lokale ruimtelijke structuren worden gevangen.
• Generaliteit: Het mechanisme is agnostisch voor de onderliggende kwantizer en vertoont verbeteringen bij VQVAE, LFQ en IBQ. Het is ook effectief overdraagbaar van class-conditional generatie naar text-to-image synthese.

4. Experimentele Resultaten

Experimenten werden uitgevoerd op ImageNet-1K (256x256) en de GPIC text-to-image benchmark.

• Convergentiesnelheid: GEAR versnelt de ImageNet gFID-convergentie met wel 10$\times$ vergeleken met de sterke LlamaGen-REPA baseline. Een naïeve end-to-end variant met behulp van STE divergeert volledig.
• Generatiekwaliteit:
  • ImageNet: Bij 300 epochs bereikt GEAR (775M parameters) een gFID van 2.52 (met CFG), waarmee het LlamaGen-REPA (2.68) verslaat en de prestaties van sterke latente diffusiemodellen (bijv. REPA-E op 1.12, hoewel REPA-E een continue VAE met hogere reconstructietrouw gebruikt) nadert.
  • Text-to-Image: Op GPIC presteert GEAR consistent beter dan LlamaGen-REPA over alle trainingsstappen (50k tot 390k), waarbij lagere Fréchet-afstanden in de DINOv2 feature-ruimte (FDD) worden behaald.
• Reconstructie: In tegenstelling tot methoden die reconstructie opofferen voor semantiek, verbetert GEAR de reconstructiekwaliteit (rFID) of behoudt deze tijdens end-to-end training.
• Ablaties:
  • STE Falen: Het vervangen van de soft bridge door STE veroorzaakt codebook-instorting (gFID $\approx$ 105).
  • Temperatuur: Een temperatuur van $\tau=0.1$ balanceert generatiekwaliteit en reconstructie.
  • Modelgrootte: GEAR verbetert de prestaties over Base, Large en XLarge schalen, waarbij de co-getrainde tokenizer een verbeterde reconstructie laat zien naarmate de AR-modelgrootte toeneemt.

5. Betekenis en Claims

Het artikel claimt dat GEAR een belangrijke stap is naar verenigde, end-to-end visuele generatie voor discrete token-pipelines.

• Overwinnen van de Discrete Bottleneck: GEAR toont aan dat end-to-end training mogelijk is voor discrete AR-modellen zonder de instabiliteit van STE, mits de predictieverlies wordt ontkoppeld van de tokenizer-update.
• Nieuw Paradigma voor Alignment: Het artikel benadrukt een contra-intuïtieve bevinding: voor discrete AR-modellen is de optimale strategie niet om de tokenizer semantisch te maken. In plaats daarvan moet de tokenizer worden geoptimaliseerd voor voorspelbaarheid (lagere entropie, gestructureerde indices), terwijl de semantische alignment ontstaat binnen de verborgen toestanden van de AR-generator. Deze "lokale, patch-niveau structuur" is wat next-token prediction efficiënt maakt.
• Fundament voor Verenigde Systemen: De auteurs stellen dat discrete next-token formuleringen uniek geschikt zijn voor verenigde begrip- en generatiesystemen (bijv. het afhandelen van lange contexten en multimodale inputs), omdat ze de fout per stap begrenzen en profiteren van volwassen alignment-technieken uit Large Language Models (RLHF, DPO). GEAR biedt een robuuste basis voor het opschalen van deze systemen.

De auteurs erkennen dat hoewel GEAR de kloof met diffusiemodellen verkleint, de reconstructie-plafond een beperking blijft vanwege de discrete aard van VQ-tokenizers vergeleken met de continue VAE's die in top-tier diffusiemodellen worden gebruikt. Toekomstig werk suggereert het verkennen van het ontkoppelen van compressie van de sequentielengte (bijv. via patchified AR) om de getrouwheid verder te verbeteren.