Technische Samenvatting: Leren van Discrete Autoregressieve Priors met Wasserstein Gradient Flow

1. Probleemstelling

Discrete image tokenizers worden doorgaans getraind volgens een tweestapsrecept: eerst wordt de tokenizer getraind om afbeeldingen te reconstrueren (met focus op conditionele waarschijnlijkheid), en vervolgens wordt een autoregressief (AR) prior getraind op de resulterende bevroren tokenreeksen. Deze ontkoppeling creëert een fundamentele mismatch: de tokenizer leert om afbeeldingsinformatie te behouden, maar blijft onbewust van de inductieve bias van het AR-prior. Bijgevolg kunnen de geleerde tokens uitstekend zijn voor reconstructie, maar moeilijk voorspelbaar zijn voor een AR-model dat sequentieel van links naar rechts leest.

Bestaande methoden proberen dit te mitigeren door de decoder aan te passen (bijvoorbeeld Tail Dropout, Diffusion Decoders) om sequentiële structuur te bevorderen. Deze benaderingen vertrouwen echter nog steeds op decoder-gerichte supervisie en maken geen direct gebruik van de modelleermoeilijkheid van het AR-prior om het leren van de tokenizer te sturen. Bovendien falen standaard technieken voor distributieniveau-matching (zoals WAE-MMD, InfoVAE), die zijn ontworpen voor continue latente ruimtes, om direct toegepast te worden op harde discrete tokenizers vanwege het ontbreken van differentieerbare paden door kwantisatie heen.

2. Methodologie

Theoretisch Kader: Tripartite Variational Consistency (TVC)

De auteurs introduceren Tripartite Variational Consistency (TVC) om het leren van latent-variabelen te analyseren. TVC decomposeert het generatieve doel in drie consistentievoorwaarden die zijn afgeleid uit de regel van Bayes:

1. Conditionele Waarschijnlijkheidsconsistentie: Uitlijning tussen de decoder van het model en de data-geankerde posterior.
2. Prior-consistentie: Uitlijning tussen het prior van het model en de geaggregeerde posterior van de encoder.
3. Posterior-consistentie: Uitlijning tussen de posterior van het model en de data-geankerde posterior.

De auteurs betogen dat standaard tweestaps training in stap één voldoet aan Likelihood Consistency en in stap twee aan Prior Consistency, maar geen signaal biedt aan de tokenizer om ervoor te zorgen dat de geaggregeerde posterior van nature "AR-vriendelijk" is. Omgekeerd dwingt ELBO-stijl training instance-voor-instance prior-consistentie af, wat vaak leidt tot posterior collapse.

Voorgestelde Oplossing: WGF-DPD en wAR-Tok

Om deze kloof te overbruggen, stelt het artikel Distributional Prior Divergence (DPD) voor, een term die de geaggregeerde distributie van tokens matcht met het doel-AR-prior zonder individuele instanties te beperken. Om dit te optimaliseren voor discrete tokens, maken ze gebruik van Wasserstein Gradient Flow (WGF).

De kernmechanisme, WGF-DPD, werkt als volgt:

1. Proxy-model: Een lichtgewicht hulp-AR-model ($q_\psi$) wordt getraind om de huidige geaggregeerde token-distributie die door de tokenizer wordt geproduceerd ($q_\phi$) te volgen.
2. Score-gebaseerde update: De update-richting voor de tokenizer wordt afgeleid uit de Wasserstein-gradiënt van de KL-divergentie tussen de geaggregeerde posterior en het doel-prior. Dit resulteert in een update op deeltjesniveau:
   $$\nabla_z \mathcal{L}_{DPD} \approx -\nabla_z \log p_\theta(x|z) + (\log q_\psi(z) - \log p_\theta(z))$$
   • De eerste term is de standaard reconstructie-gradiënt.
   • De tweede term is de prior-matching score, die het verschil vertegenwoordigt tussen de log-waarschijnlijkheden van het proxy-model en het doel-AR-prior.
3. Efficiënte optimalisatie: Cruciaal is dat deze update slechts forward passes vereist door het proxy- en doel-AR-model. Er is geen backpropagation nodig door een van de AR-modellen. De gradiënt wordt via de Straight-Through Estimator (STE) over de niet-differentieerbare kwantisatiestap doorgegeven aan de encoder-parameters.

De resulterende tokenizer, genaamd wAR-Tok, wordt samen met het reconstructiedoel en het WGF-DPD-signaal getraind.

3. Belangrijkste Bijdragen

• TVC-kader: Een theoretische decompositie van latent-variabelen leren dat "Prior Consistency" expliciet identificeert als een apart doel, en verklaart waarom tweestaps training de tokenizer onbewust laat van de modelleermoeilijkheid van het prior.
• WGF-DPD-algoritme: Een nieuwe optimalisatiemethode voor discrete tokenizers die de geaggregeerde token-distributie uitlijnt met een AR-prior met behulp van Wasserstein gradient flow. Het vermijdt de onoplosbaarheid van backpropagation door discrete kwantisatoren of de instabiliteit van adversarial matching.
• wAR-Tok: Een praktische implementatie die het WGF-DPD-signaal integreert in tokenizer-training, waarbij slechts forward passes door hulp-AR-modellen nodig zijn om encoder-zijde uitlijningssignalen te genereren.

4. Experimentele Resultaten

De methode werd geëvalueerd op CIFAR-10 en ImageNet 256×256, met vergelijkingen tegen baselines waaronder 1D Tokenization, Tail Dropout, Diffusion Decoders en Tied Timesteps.

• CIFAR-10:
  • Bij overeenkomende reconstructiekwaliteit (rFID $\approx$ 5.0) behaalde wAR-Tok een generatie-FID (gFID) van 14.50 en een AR-generatieverlies (gLoss) van 0.42, wat significant beter was dan de sterkste baseline (ImageFolder/One-D-Piece: gFID 17.60, gLoss 0.78).
  • Bij overeenkomende trainingsiteraties bereikte wAR-Tok een gFID van 11.30 en een gLoss van 0.34, wat snellere convergentie en hogere kwaliteit aantoont.
• ImageNet 256×256:
  • Met dezelfde encoder-decoder en AR-prior-architectuur als de baseline, verlaagde wAR-Tok de gFID van 6.32 naar 5.42 en de evaluatie-AR-loss van 7.89 naar 7.07, met een kleine afweging in reconstructiekwaliteit (rFID 2.05 $\to$ 2.35).
• Trainingsdynamiek: De AR-trainingscurven tonen aan dat wAR-Tok gedurende de training een lagere cross-entropy-loss en een hogere token-voorspellingnauwkeurigheid behoudt in vergelijking met baselines, wat aangeeft dat de geleerde tokens van nature makkelijker te modelleren zijn door het AR-prior.

5. Betekenis en Claims

Het artikel beweert dat wAR-Tok de mismatch tussen discrete tokenizers en autoregressieve priors effectief oplost door een distributieniveau prior-matching signaal in te voeren tijdens de tokenizer-training.

• Efficiëntie: De methode is computationeel efficiënt, met slechts twee forward passes (proxy en doel) per trainingsstap en vermijdt dure backpropagation door de AR-modellen.
• Generaliteit: Hoewel gefocust op AR-priors, worden het TVC-kader en de WGF-DPD-benadering gepresenteerd als prior-agnostisch, toepasbaar op elke discrete prior met hanteerbare log-waarschijnlijkheden.
• Bescheidenheid: De auteurs stellen expliciet dat ze geen absolute State-of-the-Art (SOTA) prestaties claimen, aangezien ze opzettelijk hulptechnieken uitsloten die gebruikelijk zijn bij recente tokenizers (zoals DINO-uitlijning, distillatie, diffusion-decoder post-training) om het effect van WGF-DPD te isoleren. De resultaten tonen de effectiviteit van het voorgestelde uitlijningsdoel in een gecontroleerde setting.

Het werk suggereert dat het expliciet optimaliseren voor de "voorspelbaarheid" van tokens door het downstream generatieve model, in plaats van alleen hun reconstrueerbaarheid, leidt tot efficiëntere en hogere kwaliteit discrete generatieve systemen.