UniCompress: Token Compression for Unified Vision-Language Understanding and Generation

Probleemstelling

Unificatie-modellen (Unified Models) zijn een recente doorbraak in multimodaal leren, waarbij zowel beeldbegrip (understanding) als beeldgeneratie (generation) worden uitgevoerd binnen één autoregressief raamwerk. Deze modellen coderen afbeeldingen naar discrete visuele tokens die samen met tekst worden verwerkt door een taalmodel (LLM).

Hoewel deze architectuur voordelen biedt zoals gedeelde parametrisatie en naadloze overdracht tussen modaliteiten, introduceert het een aanzienlijk efficiëntieprobleem:

1. Hoge Token-aantallen: Bestaande tokenizers (zoals VQ-VAE, VQGAN) coderen een afbeelding van 512x512 vaak naar 1024 tokens (32x32). Dit leidt tot lange sequenties.
2. Resource-intensief: De lange sequenties veroorzaken hoge geheugenvoet, trainingskosten en inferentielatentie, wat de inzet beperkt op systemen met beperkte middelen (bijv. embodied AI).
3. Het Dilemma van compressie: Eenvoudige compressiemethoden (zoals naïeve downsampling of uniform pruning) werken redelijk voor beeldbegrip, maar degraderen de beeldgeneratiekwaliteit drastisch (vaak >15% verlies). Generatie vereist fijne, ruimtelijk consistente details die kwetsbaar zijn voor pruning.
4. Kosten van hertraining: Het trainen van een nieuwe, efficiëntere tokenizer van scratch vereist vaak het opnieuw trainen van het hele taalmodel, wat prohibitief duur is.

Er is dus behoefte aan een modulaire, plug-in oplossing die de token-aantallen verlaagt zonder de prestaties op beide taken (begrip en generatie) significant te schaden en zonder volledige hertraining van het LLM.

Methodologie: UniCompress

De auteurs stellen UniCompress voor, een plug-in token-compressieframework dat visuele tokens tot 4x reduceert (bijv. van 256 naar 64) terwijl de prestaties behouden blijven. Het framework bestaat uit drie lichtgewicht modules die rondom een bestaande discrete tokenizer worden geplaatst:

1. Global Token Extraction (Cross-Attention):

   • Een klein aantal leerbare "meta query tokens" ($Q$) gebruikt cross-attention om globale context uit de volledige reeks visuele tokens te extraheren.
   • Dit resulteert in een set van globale meta-tokens ($G$) die de scène-level semantiek en objectrelaties vastleggen, zonder de lokale details te verliezen.

2. Token Compression (Pooling):

   • De lokale visuele tokens worden samengevoegd via gemiddelde pooling (bijv. 2x2 of 4x4 blokken) binnen niet-overlappende ruimtelijke patches.
   • Dit reduceert de sequentielengte aanzienlijk (lokaal) terwijl de grove structuur behouden blijft.
   • De sequentie wordt nu samengesteld uit: [IMG_BOS], globale tokens, [IMG_SEP], en gecomprimeerde lokale tokens.

3. Autoregressive Decompression (Global Guided):

   • Tijdens de generatie voorspelt het LLM de indices voor zowel de globale als de gecomprimeerde lokale tokens.
   • Een decompressor (geïmplementeerd als een Transformer-decoder) gebruikt de globale tokens als "semantische ankers" om de gecomprimeerde representatie autoregressief uit te breiden naar een dicht raster van tokens op de originele resolutie.
   • De globale tokens zorgen ervoor dat de decompressor langeafstandsconsistentie en details kan herstellen, wat het verlies door compressie compenseert.

Trainingsproces (Twee fasen):

• Fase 1 (Tokenizer Training): Het LLM wordt bevroren. De tokenizer-stack (extractor, compressor, decompressor) wordt getraind met een reconstructieverlies om de afbeelding zo goed mogelijk te reconstrueren uit de gecomprimeerde tokens.
• Fase 2 (LLM Fine-tuning): De tokenizer wordt bevroren. Het LLM wordt lichtjes gefinetuned op de gecomprimeerde data voor zowel begrip als generatie. Dit vereist geen volledige hertraining van het model.

Belangrijkste Bijdragen

1. Identificatie van de Bottleneck: De auteurs benadrukken dat token-efficiëntie een kritieke beperking is voor unified modellen en tonen aan dat naïeve compressie generatie onevenredig schaadt.
2. UniCompress Framework: Een plug-in architectuur met een global-guided autoregressive decompressor. Dit is uniek omdat het zowel continu (voor begrip) als discreet (voor generatie) tokens ondersteunt binnen één compressiemechanisme.
3. Plug-and-Play Integratie: Het framework kan worden geïntegreerd in bestaande unified modellen (zoals UNITOK, VILA-U, VARGPT) zonder de architectuur van het taalmodel te wijzigen of volledige hertraining te vereisen.
4. Empirische Validatie: Uitgebreide experimenten tonen aan dat het mogelijk is om visuele tokens met 4x te reduceren met slechts minimale prestatieverlies (<3% op begrip, <5% op generatiekwaliteit).

Resultaten

De methode is getest op diverse state-of-the-art unified modellen (UNITOK, VILA-U, VARGPT, UNIFORK, OPENUNI, BAGEL) op standaard benchmarks.

• Beeldbegrip (Understanding):
  • Prestaties op benchmarks zoals GQA, MME, POPE en MM-Bench blijven zeer stabiel.
  • Bijvoorbeeld: UNITOK-COMPRESSED toont slechts een kleine daling op GQA (55.71 → 53.07) en behoudt prestaties op MME en MMMU.
• Beeldgeneratie (Generation):
  • Hoewel er een lichte toename is in FID (Frechet Inception Distance, wat betekent iets minder scherpe beelden), blijft de semantische alignatie (CLIP-score) hoog.
  • Modellen zoals VILA-U en VARGPT behouden een hoge kwaliteit (FID stijgt slechts marginaal, bijv. 14.80 → 16.37 voor VILA-U).
  • De globale meta-tokens zijn cruciaal: zonder deze daalt de generatiekwaliteit aanzienlijk.
• Efficiëntie:
  • Inferentie: Tot 41,8% lagere latentie bij generatie. Bijvoorbeeld, de inferentietijd voor UNITOK daalt van 32,25 minuten naar 18,96 minuten (een versnelling van >40%).
  • Training: Tot 15,4% kortere trainingstijd door de kortere sequenties.
  • Tokenreductie: Tot 4x minder tokens (bijv. van 256 naar 64) zonder volledige hertraining.

Betekenis en Conclusie

UniCompress biedt een praktische oplossing voor de schaalbaarheid van multimodale modellen. Het bewijst dat compacte visuele representaties haalbaar zijn voor zowel begrip als generatie, mits de compressie wordt begeleid door globale semantische tokens en een autoregressieve decompressiestrategie.

De belangrijkste implicaties zijn:

• Toepasbaarheid: Het maakt unified modellen inzetbaar op resource-beperkte platforms (zoals robots of edge devices) waar geheugen en rekkracht beperkt zijn.
• Architecturale Onafhankelijkheid: Omdat het een plug-in is, kan het worden toegepast op diverse bestaande en toekomstige unified modellen zonder de complexe LLM-architectuur aan te passen.
• Balans: Het lost het compromis op tussen efficiëntie en kwaliteit, wat eerder een hardnekkig probleem was in token-compressie voor generatieve taken.

Kortom, UniCompress opent de weg naar schaalbare, kostenefficiënte en hoogwaardige unified vision-language modellen voor real-world toepassingen.