Adapting Self-Supervised Representations as a Latent Space for Efficient Generation

Probleemstelling

Bestaande generatieve modellen, zoals Diffusion Models en Flow Matching-modellen, genereren vaak beelden in de hoge-dimensionale pixelruimte of in een tweedimensionale (2D) latente ruimte (zoals bij Latent Diffusion Models). Dit brengt twee grote nadelen met zich mee:

1. Hoge rekentijd: Het regresseren van vectorvelden in pixelruimte of zelfs in een 2D latente rooster is computatief zwaar.
2. Ruin in de latente ruimte: Zelfs geoptimaliseerde 2D latente ruimtes (zoals bij VAEs) lijden nog steeds onder ruimtelijke redundantie. Beelden hebben vaak veel herhalende patronen die niet efficiënt worden gecodeerd door een raster van tokens.

Recente pogingen zoals TiTok proberen dit op te lossen door 2D roosters te vervangen door 1D tokenreeksen, maar deze gebruiken nog steeds meerdere tokens (bijv. 32 tokens) per afbeelding. Het paper stelt de vraag of het mogelijk is om een afbeelding te vertegenwoordigen met slechts één enkel continu token zonder in te leveren op reconstructiekwaliteit of generatieve prestaties.

Methodologie: RepTok

De auteurs introduceren RepTok (Representation Tokenizer), een framework dat zelf-supervisie (SSL) gebruikt om een compacte latente ruimte te creëren. De kern van de methode bestaat uit drie componenten:

1. Aangepaste SSL-Encoder (Single Token):

   • Het framework maakt gebruik van een voorgeprogrammeerde Self-Supervised Learning (SSL) encoder (bijv. DINOv2, MAE, CLIP).
   • In plaats van de volledige encoder te finetunen of een nieuwe tokenizer te trainen, wordt alleen de [cls]-token embedding (de class token) van de SSL-encoder fijn afgestemd (fine-tuned). De rest van de encoder blijft bevroren.
   • Deze [cls]-token, die normaal gesproken alleen semantische informatie bevat, wordt zo aangepast dat het ook lage-niveau reconstructiedetails (ruimtelijke informatie) kan bevatten.
   • Het resultaat is dat een hele afbeelding wordt gecodeerd in één enkel continu vector (bijv. dimensie 768), in plaats van een raster van tokens.

2. Generatieve Decoder met Flow Matching:

   • Een generatieve decoder wordt getraind om dit enkele token terug te vertalen naar de pixelruimte (of een tussenliggende VAE-ruimte).
   • De training gebruikt een Flow Matching doelstelling, waarbij het model leert een vectorveld dat een Gaussische prior naar de data-distributie transformeert.
   • De decoder wordt getraind om de afbeelding te reconstrueren op basis van het enkele token $z = E(x)$.

3. Cosine-Similarity Loss (Regularisatie):

   • Een kritisch probleem is dat het vrij laten van de [cls]-token tijdens het finetunen de token te ver kan wegtrekken van de oorspronkelijke, goed gestructureerde SSL-ruimte. Dit zou de gladde geometrie verstoren die nodig is voor succesvolle generatie.
   • Om dit te voorkomen, wordt een cosine-similarity loss toegevoegd. Deze loss zorgt ervoor dat de aangepaste token $z$ dicht bij de oorspronkelijke token van de bevroren encoder $z_{frozen}$ blijft.
   • Dit creëert een balans: de token krijgt genoeg informatie voor perfecte reconstructie, maar behoudt de gladde, semantisch gestructureerde geometrie van de oorspronkelijke SSL-ruimte, wat essentieel is voor interpolatie en generatie.

4. Architectuur voor Generatie (MLP-Mixer):

   • Omdat de input slechts één token is, zijn er geen interacties tussen tokens nodig. Dit maakt het mogelijk om een attention-free architectuur, specifiek een MLP-Mixer, te gebruiken voor de generatieve stap.
   • Dit elimineert de kwadratische schaalbaarheid van attention-mechanismen, wat leidt tot extreme efficiëntie.

Belangrijkste Bijdragen

• Single-Token Latent Space: Het bewijs dat een enkele, continue token, afgeleid van een SSL-encoder, voldoende informatie bevat voor zowel hoogwaardige reconstructie als generatie, waardoor ruimtelijke redundantie volledig wordt geëlimineerd.
• Efficiënte Finetuning: Een strategie waarbij slechts de [cls]-token wordt bijgewerkt, wat de trainingstijd en -kosten drastisch verlaagt ten opzichte van het trainen van volledige encoders of decoders.
• Attention-Free Generatie: Het tonen dat generatieve modellen zonder self-attention (MLP-Mixer) competitief kunnen presteren wanneer de inputruimte extreem gecomprimeerd is (één token).
• Schaalbaarheid naar Text-to-Image: Het framework kan worden uitgebreid naar tekst-naar-beeld synthese door cross-attention toe te voegen, terwijl de latente ruimte compact blijft.

Resultaten

De auteurs evalueren RepTok op ImageNet-1K (class-conditional) en MS-COCO (text-to-image):

• Reconstructie: RepTok bereikt een rFID van 1.85 en een PSNR van 14.94 op ImageNet 256x256, wat concurrerend is met of beter is dan bestaande tokenizers zoals TiTok en FlexTok, ondanks het gebruik van slechts één token.
• Generatie: Voor class-conditional ImageNet generatie bereikt het model een gFID van 1.88 (met CFG).
• Efficiëntie:
  • De training vereist >90% minder rekentijd (GFLOPs) dan transformer-based baselines zoals DiT en SiT.
  • Het model bereikt vergelijkbare prestaties met slechts 1,7% van de FLOPs die SiT verbruikt.
  • De inferentie is extreem snel: het totale tijdsverloop voor generatie is 0.27 seconden (inclusief decoder), vergeleken met 2.43 seconden voor SiT-XL.
• Text-to-Image: Het model bereikt competitieve zero-shot prestaties op MS-COCO met een zeer beperkt trainingsbudget (< 20 uur op 4x A100 GPU's).
• Interpolatie: De latent space behoudt de gladde eigenschappen van de SSL-encoder, wat zorgt voor vloeiende overgangen in zowel semantiek als ruimtelijke configuratie bij interpolatie.

Betekenis en Impact

RepTok markeert een paradigmaverschuiving in generatieve modellering door te laten zien dat de "slimme" compressie van een voorgeprogrammeerde SSL-encoder (die al een goed begrip van beelden heeft) direct kan worden gebruikt als latente ruimte, mits er een minimale aanpassing wordt gedaan voor reconstructiedetails.

De belangrijkste implicaties zijn:

1. Democratisering van Generatieve AI: Door de rekentijd en kosten drastisch te verlagen, wordt het trainen van hoogwaardige generatieve modellen toegankelijker voor onderzoekers met beperkte middelen.
2. Efficiëntie: Het elimineren van attention-mechanismen in de generatieve stap, gecombineerd met extreme compressie, opent de deur voor snelle, real-time generatieve toepassingen.
3. Nieuwe Richting voor Latent Spaces: Het paper suggereert dat de toekomst van efficiënte generatie ligt in het benutten van bestaande, semantisch rijke representaties (SSL) in plaats van het trainen van nieuwe, zware auto-encoders van scratch.

Kortom, RepTok bewijst dat "less is more": één goed gekozen token, aangepast met slimme regularisatie, is superieur aan complexe roosters van tokens in termen van efficiëntie en prestatie.