Technische Samenvatting: MergeTok

Probleemstelling

Huidige visuele tokenizer voor beeldgeneratie zijn verdeeld in twee afzonderlijke families, die elk lijden onder complementaire structurele beperkingen:

1. Continue VAE-gebaseerde tokenizers: Hoewel ze een hoge reconstructiefailliteit bieden, vertrouwen ze op latente ruimtes die geen intrinsieke organisatie bezitten. Onder sterke compressie raken semantische factoren verstrengeld binnen dichte latente representaties, wat resulteert in zwak ontwarde codes die de semantische controleerbaarheid belemmeren.
2. Discrete VQ-gebaseerde tokenizers: Hoewel ze autoregressieve generatie mogelijk maken, kampen ze met optimalisatieproblemen. Het niet-differentieerbare karakter van kwantisatie leidt tot ijle en ongelijkmatige gradiëntupdates, waardoor veel codeboek-entries onvoldoende supervisie ontvangen. Dit vergroot het risico op codeboek-instorting (codebook collapse) en vermindert de representatie-efficiëntie.

Deze beperkingen creëren een hardnekkige spanning tussen reconstructiefailliteit, semantische structuur, efficiëntie en trainingsstabiliteit.

Methodologie: MergeTok Framework

MergeTok stelt een unified dual-branch tokenizer voor die continue (VAE) en discrete (VQ) tokenisatie gezamenlijk optimaliseert binnen één encoder-decoder architectuur. De kerninnovatie is het gebruik van Token Merging (ToMe) als een semantische brug tussen de twee takken.

Architectuur Overzicht

Gegeven een invoerbeeld $X$, produceert een gedeelde CNN-encoder $E_c$ een feature map $Z$. Deze map wordt verwerkt door twee parallelle takken:

1. VAE Tak (Continu):

   • Token Merging: Een attention-gebaseerde encoder uitgerust met ToMe-modules comprimeert de token-sequentie van $L$ tokens naar $K$ tokens ($K < L$) door vergelijkbare tokens samen te voegen. Dit proces genereert een gecondenseerde representatie $Z^{(vae)}_K$ en een source map $S$ die de toewijzing van originele tokens aan samengevoegde groepen vastlegt.
   • Semantische Uitlijning: De samengevoegde tokens worden uitgelijnd met een bevroren DINO-stijl docentmodel (eveneens uitgerust met ToMe) met behulp van een cosine distance loss ($L_{align}$). Dit reguleert de continue latente ruimte richting semantisch bewuste, ontwarde representaties.
   • Reconstructie: Een hybride decoder gebruikt zowel de gecondenseerde tokens als de source map $S$ om de groepen te "ont-mergen" en de pixel-ruimte afbeelding te reconstrueren.

2. VQ Tak (Discreet):

   • Groep-bewuste Kwantisatie: Deze tak omzeilt de ToMe-operatie tijdens de forward pass om de volledige sequentielengte ($L$ tokens) te behouden. Echter, het maakt gebruik van de source map $S$ gegenereerd door de VAE-tak als een structurele prior.
   • Regularisatie: De source map partitioneert de $L$ tokens in $K$ semantische groepen. Twee groep-bewuste losses worden toegepast op de codeboek-toewijzingen:
     • Intra-groep Diversiteit ($L_{div}$): Maximaliseert de entropie binnen elke groep om instorting naar een enkele code te voorkomen.
     • Inter-groep Exclusiviteit ($L_{cons}$): Minimaliseert overlap tussen groepen om te waarborgen dat verschillende semantische groepen onderscheidende code-subsets gebruiken.
   • Reconstructie: Een standaard VQ-decoder reconstrueert de afbeelding vanuit de gekwantiseerde tokens.

Trainingsstrategie

• Unified Objective: Het model wordt end-to-end getraind met een totale loss die reconstructie, adversarial en merge-bewuste regularisaties combineert:
  $$L_{total} = \lambda_{vae}L_{VAE} + \lambda_{vq}L_{VQ} + \lambda_{ali}L_{align} + \lambda_{div}L_{div} + \lambda_{cons}L_{cons}$$
• Granulariteit-bewuste Sampling: Tijdens de training wordt de merge-ratio (en daarmee het aantal behouden tokens $K$) dynamisch gesampled uit een discrete distributie gecentreerd rond een hyperparameter $K^*$. Dit stelt de encoder bloot aan variërende bottleneck-breedtes, wat robuustheid over verschillende semantische granulariteiten stimuleert.
• Inference: Tijdens de inference geeft de tokenizer altijd een vaste sequentie van 256 tokens uit. Downstream generatoren (zowel autoregressief als diffusie) opereren op deze vaste sequentie en zijn agnostisch tegenover de interne mergestructuur.

Belangrijkste Bijdragen

1. Unified Dual-Branch Tokenizer: MergeTok is de eerste architectuur die zowel continue als discrete latents gezamenlijk optimaliseert onder één enkele objective, waarbij token merging als een semantische interface dient om de stabiliteit van VAE-training en de structurele voordelen van VQ te versterken.
2. Merge-bewuste Trainingsbeperkingen: De auteurs introduceren twee nieuwe objectives afgeleid van het mergingproces:
   • Merged-token alignment: Verbetert de semantische structuur in de continue latente ruimte.
   • Groep-bewuste kwantisatie: Stabiliseert VQ-training en verbetert de codeboek-benutting via intra-groep diversiteit en inter-groep exclusiviteit.
3. Granulariteit-bewuste Merge Ratio Sampling: Een strategie tijdens de training die verschillende token-granulariteiten samplet, waardoor het model een verenigde tokenisatie kan leren die robuust is bij variërende compressieratio's.

Experimentele Resultaten

Experimenten werden uitgevoerd op ImageNet-256 (256×256 resolutie) met een vaste token-budget van 256 tokens.

• Reconstructiefailliteit:
  • Discreet (VQ): MergeTok-BL bereikte een rFID van 0.78, waarmee het sterke baselines zoals GigaTok-BL (0.81) en MergeVQ-GR (1.12) overtreft. Met een bevroren DINO discriminator bereikte MergeTok-BL* een rFID van 0.48.
  • Continu (VAE): MergeTok-BL bereikte een rFID van 0.47, wat marginaal beter is dan RAE (0.49) en aanzienlijk beter dan SoftVQ-L (0.61) en DC-AE-f32 (0.69).
• Representatiekwaliteit:
  • MergeTok-BL behaalde 78.3% linear probing nauwkeurigheid, de hoogste onder de vermelde discrete tokenizers, waarmee het UniTok (70.8%) en MergeVQ-GR (77.9%) overtreft. Dit duidt op een superieure semantische organisatie.
• Downstream Generatie:
  • Autoregressief (VQ Tak): In combinatie met LlamaGen-XXL bereikte MergeTok-BL een gFID van 1.93 (zonder CFG) en 2.13 (met CFG), waarmee het GigaTok-BL (2.03) overtreft.
  • Diffusie/Flow (VAE Tak): In combinatie met SiT-XL/1 bereikte MergeTok-BL een gFID van 1.79 (zonder CFG) en 1.29 (met CFG), wat overeenkomt met de beste resultaten in de tabel.

Betekenis en Claims

Het artikel claimt dat MergeTok aantoont dat een enkele architectuur visuele tokenizers kan voorzien van zowel robuuste semantische organisatie als generator-vriendelijke discretie.

• Unificatie: Het lost de trade-off op tussen de hoge reconstructiefailliteit van VAE's en de discrete structuur die vereist is voor autoregressieve generatie.
• Semantische Brug: Door token merging te gebruiken als een structurele prior, slaagt het model erin om semantische groeperingsinformatie van de continue tak over te dragen om de discrete tak te stabiliseren, en vice versa.
• Compatibiliteit: De tokenizer produceert 256-token sequenties die compatibel zijn met zowel autoregressieve (bijv. LlamaGen) als diffusie/flow (bijv. DiT, SiT) generatoren zonder dat de generatoren zelf aangepast hoeven te worden.

De auteurs positioneren MergeTok niet louter als een reconstructietool, maar als een verenigde interface die gelijktijdig discrete autoregressieve generatie, continue diffusie-generatie en semantisch gestructureerde representatie-lering ondersteunt.