CaTok: Taming Mean Flows for One-Dimensional Causal Image Tokenization

Dit paper introduceert CaTok, een nieuwe 1D-causale beeldtokenizer met een MeanFlow-decoder en REPA-A-regularisatie die state-of-the-art reconstructie op ImageNet bereikt en autoregressieve generatie voor visuele data mogelijk maakt.

De kern

Stel je voor dat je een kunstenaar bent die een schilderij moet maken, maar in plaats van met verf en kwasten werkt, gebruik je een reeks instructies. De kunst van het "generatieve AI" (AI die nieuwe beelden maakt) draait al lang om het vinden van de juiste manier om deze instructies te geven.

Deze paper introduceert CaTok, een nieuwe manier om beelden om te zetten in een reeks instructies die voor computers makkelijker te begrijpen zijn. Hier is hoe het werkt, vertaald naar alledaagse taal:

1. Het Probleem: De Verwarde Boodschappenlijst

Vroeger probeerden computers beelden te begrijpen door ze op te delen in een rooster van stukjes (zoals een pixelmat). Ze maakten hier een lange lijst van, maar de volgorde was vaak willekeurig of onlogisch.

• De analogie: Stel je voor dat je een boek leest, maar de zinnen staan door elkaar heen. Je begint met de laatste zin, dan de eerste, dan de middelste. Dat is heel lastig om een verhaal (of een beeld) te begrijpen of te herscheppen.
• Het doel: We willen dat de computer beelden leest en schrijft zoals wij boeken lezen: van links naar rechts, van begin tot eind. Dit noemen ze "causaal" (oorzaak en gevolg).

2. De Oplossing: CaTok (De Slimme Boekhouder)

CaTok is een systeem dat beelden omzet in een één-dimensionale lijst (een lange rij) van token's (instructie-blokjes). Maar het slimme deel is hoe het deze lijst maakt.

De auteurs gebruiken een slimme truc genaamd MeanFlow.

• De analogie: Stel je voor dat je een auto wilt besturen van punt A naar punt B.
  • De oude methoden keken alleen naar de snelheid op dit exacte moment (te onnauwkeurig voor een snelle rit).
  • Andere methoden keken naar de hele rit, maar dat maakte het moeilijk om halverwege te stoppen of te beginnen.
  • CaTok kijkt naar het gemiddelde van de snelheid over een stukje van de rit. Het zegt: "Als we dit stukje van de weg afleggen, wat is dan de gemiddelde richting?"
• Het resultaat: Hierdoor kan de computer het beeld in één keer "schetsen" (snel) of in stapjes verfijnen (hoogwaardig), afhankelijk van hoe snel je wilt werken.

3. De Twee Slimme Trucs

Om dit systeem te laten werken, gebruiken de auteurs twee belangrijke hulpmiddelen:

A. De "Tijdsinterval"-Truc (Causality)
In plaats van de computer alle informatie tegelijk te geven, geeft CaTok de informatie in blokken, gebaseerd op tijd.

• De analogie: Stel je voor dat je een film bekijkt. Je ziet eerst de achtergrond (het begin van de film), dan de hoofdpersonages, en dan de details. CaTok leert de computer om de beelden op deze natuurlijke volgorde te "ontleden". Als je alleen de eerste paar instructies geeft, krijg je een vaag silhouet. Geef je meer instructies, dan worden de details scherper. Dit zorgt voor een logische opbouw.

B. De "Onderwijzer"-Truc (REPA-A)
Tijdens het trainen van de AI gebruiken ze een bestaande, zeer slimme AI (een "Vision Foundation Model") als onderwijzer.

• De analogie: Stel je voor dat een student (CaTok) aan het leren is om te tekenen. In plaats van alleen maar te oefenen, kijkt hij elke dag naar de tekeningen van een meester (de onderwijzer) en zegt: "Kijk, dit is hoe een meester een oog tekent."
• Dit helpt de student om veel sneller te leren en fouten te voorkomen. In de paper noemen ze dit REPA-A. Het zorgt ervoor dat de AI de beelden sneller en beter leert begrijpen.

4. Waarom is dit belangrijk?

Vroeger moest je kiezen tussen:

1. Snelheid: Beelden snel maken, maar dan vaak wazig of onvolledig.
2. Kwaliteit: Prachtige beelden, maar dan duurt het lang om ze te genereren.

CaTok haalt het beste van beide werelden:

• Het kan beelden in één stap maken (supersnel), wat perfect is voor snelle toepassingen.
• Het kan ook in veel stappen werken om kristalheldere beelden te maken.
• Omdat het de beelden in een logische volgorde leert (zoals een verhaal), werkt het perfect in combinatie met de nieuwste generatie taalmodellen (zoals die voor tekst). Dit opent de deur voor AI die tekst en beeld naadloos met elkaar kan verbinden.

Samenvattend

CaTok is als een nieuwe manier om een schilderij te beschrijven. In plaats van een rommelige lijst met losse details, schrijft het een verhaal: "Eerst de lucht, dan de bomen, dan de mensen." Dankzij slimme wiskunde (MeanFlow) en een slimme onderwijzer (REPA-A) kan de computer dit verhaal zowel razendsnel als in hoge kwaliteit vertellen.

Technische samenvatting

Probleemstelling

Autoregressieve (AR) taalmodellen (LLMs) zijn zeer succesvol omdat ze tekst tokeniseren in causale sequenties die voldoen aan een natuurlijke leesvolgorde, waardoor "next-token prediction" mogelijk is. Het uitbreiden van dit paradigma naar visuele generatie is echter complex. Bestaande visuele tokenizers hebben belangrijke tekortkomingen:

1. 2D Tokenizers (zoals VQGAN): Deze tokeniseren afbeeldingen in 2D roosters en flatten ze naar een 1D-sequentie (bijv. raster- of willekeurige volgorde). Dit mist echte causaliteit tussen opeenvolgende tokens, wat AR-modellen belemmert.
2. Multi-schaal Tokenizers (zoals VAR): Deze garanderen causaliteit via een grof-naar-fijn benadering, maar breken het standaard "next-token prediction"-patroon van LLMs.
3. Diffusion Autoencoders (1D): Recentere 1D-tokenizers gebruiken vaak "naïeve" decoders die conditioneren op alle tokens tegelijk (geen causaliteit) of "nested dropout" (conditioneren op de eerste $k$ tokens). De laatste methode introduceert een onbalans: vroege tokens worden vaker geselecteerd dan latere tokens, wat de kwaliteit van AR-generatie negatief beïnvloedt.

Het doel is een tokenizer te ontwikkelen die causale 1D-tokens leert, balans behoudt tussen tokens, en snelle één-stap generatie ondersteunt.

Methodologie: CaTok

CaTok is een diffusion autoencoder met een causale Vision Transformer (ViT)-encoder en een MeanFlow Diffusion Transformer (DiT)-decoder.

1. Architectuur en Causaliteit

• Encoder: Een causale ViT met "registers" (leerbare parameters) die beeldfeatures comprimeert tot een reeks 1D-tokens. Een causale attention-mask zorgt ervoor dat tokens alleen naar voorgaande tokens kunnen kijken.
• Decoder (MeanFlow): In plaats van de instantane snelheid te voorspellen (zoals bij standaard rectified flows), voorspelt de decoder het gemiddelde snelheidsveld ($u$) over een tijdsinterval $[r, t]$.
• Token Selectie: Om causaliteit en balans te garanderen, selecteert de decoder niet alle tokens, maar alleen de tokens die corresponderen met het tijdsinterval $[r, t]$ (aangeduid als $V_{r:t}$).
  • Dit zorgt ervoor dat de tokens de causale relatie in het "noise-to-image" proces modelleren.
  • Het lost het onbalans-probleem op van eerdere methoden die alleen op vroege tokens conditioneren.

2. Training Objectieven

CaTok wordt geoptimaliseerd met vier doelen:

• MeanFlow Objective: Minimaliseert de fout tussen de voorspelde gemiddelde snelheid en de werkelijke gemiddelde snelheid over het interval $[r, t]$. Dit ondersteunt direct één-stap sampling.
• Rectified Flow Objective: Minimaliseert de fout voor de instantane snelheid (waarbij $r=t$) om de training te stabiliseren.
• REPA (Representation Alignment): Aligneert de decoder-features met een pre-getraind Vision Foundation Model (VFM) om de training te versnellen.
• REPA-A (Nieuw): Een geavanceerde regularisatie die de encoder-features aligneert met de VFM-representaties. Dit zorgt voor hogere kwaliteit semantische features in de encoder, waardoor de 1D-tokens informatiever en discriminatiever worden.

3. Autoregressieve Generatie

Na het trainen van CaTok worden de 1D-tokens gebruikt om een standaard AR-model (zoals LlamaGen) te trainen via "next-token prediction". Tijdens inferentie worden tokens voorspeld en vervolgens gebruikt om de MeanFlow-decoder te conditioneren voor het genereren van de afbeelding (vaak in één stap).

Belangrijkste Bijdragen

1. Nieuwe Architectuur: Een diffusion autoencoder die MeanFlow-objectieven koppelt aan causale 1D-tokensegmenten, waardoor visuele causaliteit wordt geleerd zonder de AR-patroon te schaden.
2. Eén-stap Sampling: CaTok ondersteunt native één-stap generatie (zeer snel) én hoogwaardige multi-stap sampling, wat zeldzaam is voor diffusion autoencoders.
3. REPA-A: Een nieuwe regularisatiemethode die encoder-features aligneert met VFMs, wat de training stabiliseert en versnelt.
4. Balans en Causaliteit: Door tokens te selecteren binnen een tijdsinterval $[r, t]$ in plaats van alleen de eerste $k$, wordt de onbalans in token-bijdrage opgelost.

Resultaten

De prestaties zijn getest op ImageNet-1K (256x256):

• Reconstructie (Encoder/Decoder):
  • CaTok-L-256 bereikte een rFID van 0.75, een PSNR van 22.53 en een SSIM van 0.674.
  • Dit is state-of-the-art (SOTA) voor diffusion autoencoders en bereikt vergelijkbare resultaten met veel minder trainingsepochen (160 vs. 400+ bij concurrenten).
  • Het model toont een duidelijke "grof-naar-fijn" trend wanneer het aantal tokens wordt gereduceerd, wat de causaliteit bewijst.
• Autoregressieve Generatie:
  • Met een AR-generator (LlamaGen) bereikte CaTok een gFID van 2.95, vergelijkbaar met toonaangevende methoden zoals Semanticist en VAR, maar met een veel efficiëntere tokenisatie.
  • Het presteert beter dan LlamaGen (die VQ gebruikt) in AR-generatie, ondanks een iets hogere reconstructiefout, wat aantoont dat de causale structuur cruciaal is voor AR-modellen.
• Ablatiestudies:
  • Het gebruik van token-selectie in $[r, t]$ is essentieel voor AR-prestaties (gFID daalt van 13.54 naar 4.91).
  • REPA-A vermindert trainingsspiezen en verbetert de convergentie aanzienlijk.

Betekenis en Impact

CaTok is een belangrijke stap in het overbruggen van de kloof tussen autoregressieve taalmodellen en visuele generatie.

• Paradigmaverschuiving: Het bewijst dat 1D-causale tokenisatie effectief kan zijn voor afbeeldingen zonder de complexiteit van multi-schaal benaderingen.
• Efficiëntie: De mogelijkheid tot één-stap sampling maakt visuele generatie aanzienlijk sneller, wat essentieel is voor real-time toepassingen.
• Stabiliteit: De introductie van REPA-A biedt een nieuwe route om diffusion autoencoders sneller en stabieler te trainen door gebruik te maken van bestaande foundation models.

Kortom, CaTok levert een robuust framework voor visuele tokenisatie dat de voordelen van autoregressieve modellen (schaalbaarheid, causaliteit) combineert met de hoge kwaliteit van diffusion-modellen.