Structured State-Space Regularization for Compact and Generation-Friendly Image Tokenization

Probleemstelling

Moderne vision-modellen (zoals generatieve AI, wereldmodellen en representatiemodellen) vertrouwen steeds meer op image tokenizers. Deze tokenizers projecteren beelden naar een compacte latente ruimte, waar de modellering plaatsvindt in plaats van in de hoge-dimensionale pixelruimte. Een ideale latente ruimte moet echter twee tegenstrijdige eisen vervullen:

1. Compactheid: De representatie moet zo klein mogelijk zijn maar toch alle essentiële informatie behouden voor een nauwkeurige reconstructie.
2. Generatie-vriendelijkheid: De ruimte moet makkelijk te modelleren zijn voor generatieve frameworks (zoals Diffusion Models).

Huidige tokenizers worstelen vaak met een afweging (trade-off): een tokenizer die uitstekend reconstructeert, presteert vaak slecht bij generatie, en vice versa. De auteurs stellen dat dit komt doordat de latente ruimte geen intrinsieke structuur heeft die overeenkomt met hoe generatieve modellen beelden opbouwen (vaak van laag- naar hoogfrequente componenten).

Methodologie

De auteurs introduceren een nieuwe regularisatiemethode genaamd Structured State-Space Regularization (SSR). De kern van de methode is het overbrengen van de eigenschappen van State-Space Models (SSM's), specifiek hun "frequentiebewustzijn" (frequency awareness), naar de latente features van een image tokenizer.

1. Theoretisch Kader: SSM's en Frequentie

SSM's worden traditioneel gebruikt voor sequentiemodellering en zijn bekend om hun vermogen om signalen impliciet te ontbinden in verschillende frequentiebanden. De auteurs analyseren SSM's als een combinatie van twee componenten:

• Basisprojectie ($c(\cdot)$): Het projecteren van een input-signaal op een set orthogonale basisfuncties (bijv. Fourier-basis).
• Input-transformatie ($\theta(\cdot)$): Een transformatie op de input over de tijd (bijv. het toevoegen van nieuwe tokens of, in dit geval, Gaussian blurring).

De dynamiek van de verborgen toestand in een SSM volgt de evolutie van de coëfficiënten van deze basisprojectie onder de transformatie.

2. Structured State-Space Regularization (SSR)

In plaats van SSM's te gebruiken als het model zelf, gebruiken de auteurs de dynamiek van SSM's als een regularisatie-verlies voor de encoder van de tokenizer.

• Input-transformatie: Ze kiezen voor een progressief Gaussian blurring-proces. Dit is wiskundig goed onderzocht (gelijk aan de warmtevergelijking) en verwijdert systematisch hoogfrequente informatie, waardoor alleen laagfrequente structuren overblijven.
• Basisprojectie: Ze gebruiken 2D orthogonale basisfuncties (bijv. Fourier-basis).
• De Regularisatie: De encoder $E$ wordt getraind om de dynamiek van de basiscoëfficiënten na te bootsen. Als een beeld $I_{t-1}$ wordt geblurred naar $I_t$, dan moet de latente representatie $z_t = E(I_t)$ voldoen aan de update-regel:
  $$z_t \approx \bar{A} z_{t-1}$$
  Waarbij $\bar{A}$ een matrix is die de evolutie van de basiscoëfficiënten onder blurring beschrijft.
• Verliesfunctie: Het verlies wordt gemeten als de afstand tussen de daadwerkelijke encoder-output van het geblurred beeld en de voorspelde output via de matrix $\bar{A}$ toegepast op de encoder-output van het scherpe beeld.

3. Praktische Implementatie

• Spatial Mean-Centering: Omdat Gaussian blurring de norm van de signalen doet afnemen (energieverlies), zou de latente ruimte instabiel worden (normen die naar nul gaan). De auteurs introduceren een spatial mean-centering functie $\mathcal{C}$ om dit te corrigeren, zodat de latente waarden rond hun ruimtelijke gemiddelde blijven in plaats van naar nul te vervallen.
• Training: De tokenizer wordt getraind met een combinatie van reconstructieverlies (om beeldkwaliteit te behouden) en het SSR-verlies (om de generatie-vriendelijke structuur te leren).

Belangrijkste Bijdragen

1. Nieuwe Regularisator: Een theoretisch onderbouwde regularisator die SSM-dynamiek overbrengt naar image tokenizers.
2. Eerste Theoretisch Kader: Het is de eerste studie die de principes van SSM's (basisprojectie en input-transformatie) formaliseert voor image tokenization.
3. Frequentie-geordende Latente Ruimte: De methode dwingt de tokenizer om de latente kanalen te organiseren volgens frequentiebanden (laag naar hoog), wat perfect aansluit bij hoe Diffusion-modellen beelden genereren.

Resultaten

De methode is getest op twee populaire tokenizers: Flux (8x downsampling) en Cosmos (16x downsampling) op het ImageNet-1K dataset.

• Generatiekwaliteit: Er is een aanzienlijke verbetering in de generatieprestaties van Diffusion Models (gemeten met gFID, sFID, en Inception Score). Bijvoorbeeld, bij Flux verbeterde de gFID van 16.08 naar 13.33, terwijl de Inception Score steeg van 89.32 naar 106.61.
• Reconstructiekwaliteit: De reconstructiekwaliteit (PSNR, LPIPS) bleef zeer hoog en verloor slechts minimaal aan kwaliteit vergeleken met de baseline of andere regularisatiemethoden (zoals EQVAE en FT-SE).
• Ablatie Studies:
  • Het gebruik van Fourier-basis gaf de beste resultaten vergeleken met Chebyshev, Legendre of Hermite polynomen.
  • Het gebruik van de spatial mean-centering functie was cruciaal; zonder deze functie stortte de generatiekwaliteit in door instabiele latente normen.
  • Visuele inspectie toonde aan dat de geregulariseerde tokenizer kan reconstrueren met alleen de eerste paar kanalen (laagfrequente informatie), wat bevestigt dat de frequentie-structuur effectief is geleerd.

Betekenis en Conclusie

Dit paper biedt een nieuw perspectief op het ontwerp van image tokenizers. Door de latente ruimte te structureren volgens de wiskundige dynamiek van SSM's en Gaussian blurring, creëren de auteurs een ruimte die inherent "generatie-vriendelijk" is.

De belangrijkste implicatie is dat generatieve modellen (zoals Diffusion) beter presteren wanneer de latente ruimte de spectrale hiërarchie van beelden respecteert (eerst globale vormen, dan details). Deze methode lost het traditionele compromis tussen reconstructie en generatie op, waardoor efficiëntere en hogere kwaliteit generatieve AI-systemen mogelijk worden. Het opent ook de deur voor het gebruik van SSM-principes in andere gebieden van visuele representatieleren.