VFM-VAE: Vision Foundation Models Can Be Good Tokenizers for Latent Diffusion Models

1. Het Probleem

Latente Diffusiemodellen (LDMs) zijn de huidige standaard voor visuele synthese, maar hun prestaties zijn kritiek afhankelijk van de kwaliteit van hun visuele tokenizers (meestal Variational Autoencoders of VAE's).

• Huidige aanpak: Recent werk probeert Vision Foundation Models (VFMs, zoals DINOv2 of SigLIP) te integreren in tokenizers via distillatie. Hierbij wordt een VAE getraind om de representaties van een VFM na te bootsen.
• De beperking: De auteurs vinden empirisch dat deze distillatie-aanpak onvermijdelijk leidt tot een verzwakking van de robuustheid van de oorspronkelijke VFM-representaties. Onderzoek met de metric CKNNA toont aan dat deze "afgeleide" representaties kwetsbaar zijn (brittle) onder semantisch behoudende transformaties (zoals ruis of rotatie), wat wijst op verlies van cruciale informatie.
• De uitdaging: VFMs zijn geoptimaliseerd voor semantisch begrip, niet voor pixel-accurate reconstructie. Het direct gebruiken van een bevroren VFM als encoder creëert een spanning tussen semantische rijkdom en de noodzaak om realistische beelden te reconstrueren.

2. Methodologie: VFM-VAE

De auteurs stellen VFM-VAE voor, een raamwerk dat de distillatie volledig omzeilt door de VFM-encoder bevroren (frozen) te houden en direct te gebruiken als het voorste deel van de VAE.

Kernarchitectuur:

1. Bevroren Encoder: In plaats van een encoder te trainen, wordt een pre-getrainde VFM (standaard SigLIP2-Large) gebruikt. Er worden features op meerdere schalen (shallow, middle, final) uit de VFM gehaald om zowel fijne details als hoge semantiek te behouden.
2. Lightweight Projectie: Een leerbaar, lichtgewicht netwerk projecteert deze multi-scale features naar een compacte latente ruimte (mean en variance voor de Gaussian posterior).
3. Geavanceerde Decoder: Omdat VFM-features ruimtelijk grof zijn, wordt een nieuwe decoder ontworpen met twee innovaties:
   • Multi-Scale Latent Fusion: De latente code wordt opgesplitst in een globale component (voor stijl en lay-out) en ruimtelijke componenten (voor details).
   • Progressive Resolution Reconstruction Blocks: De decoder bouwt het beeld op via een reeks blokken die de resolutie stapsgewijs verhogen (van 8x8 naar 256x256).
     • De globale component stuurt elk blok aan via Modulated ConvNeXt Blocks (stijlmodulatie).
     • De ruimtelijke componenten worden alleen in de vroege, lage-resolutie blokken geïnjecteerd om de basislay-out te vormen.
     • Multi-resolution Supervision: Elke blok heeft een eigen "ToRGB" kop die een tussentijdse reconstructie genereert, wat zorgt voor stabielere training en betere detailherstel.

Trainingsdoelstellingen:
De loss-functie combineert:

• Representatie Regularisatie: Behoud van de VFM-semantiek (via KL-divergentie en een VFM-verlies).
• Multi-resolutie Reconstructie: L1-verlies op elke schaal.
• Adversariaal en Perceptueel Verlies: Gebruik van een discriminator (op DINOv2 gebaseerd) en LPIPS om fotorealistische kwaliteit te garanderen.

Nieuwe Evaluatiemetric: SE-CKNNA
Om de robuustheid van tokenizers beter te meten, introduceren de auteurs Semantic-Equivariant CKNNA (SE-CKNNA). In plaats van alleen te meten op schone beelden, wordt de alignment-score gemiddeld over een distributie van semantisch behoudende transformaties (ruis, schaalveranderingen, rotaties). Dit geeft een betrouwbaarder beeld van de stabiliteit van de representatie.

3. Belangrijkste Bijdragen

1. Directe Integratie van Bevroren VFMs: Een nieuw paradigma waarbij VFMs direct als encoder dienen zonder distillatie, wat representatie-degradatie elimineert en hoge reconstructiekwaliteit behoudt via een gespecialiseerde decoder.
2. SE-CKNNA Metric: Een diagnostische tool om de robuustheid van tokenizer-VFM alignment onder perturbaties te kwantificeren.
3. Synergie tussen Tokenizer en Diffusie: Het paper toont aan dat een tokenizer met sterke VFM-alignment het representatieleren van de diffusiemodel zelf verbetert (dual-side alignment), wat leidt tot snellere convergentie.

4. Resultaten

De prestaties zijn getest op ImageNet (256x256 en 512x512) en in tekst-naar-beeld scenario's:

• Snelheid en Efficiëntie: Het systeem bereikt een gFID (zonder CFG) van 2.22 in slechts 80 epochs. Dit is een 10x versnelling ten opzichte van eerdere tokenizers (zoals VA-VAE) die 800+ epochs nodig hadden voor vergelijkbare resultaten.
• Ultieme Kwaliteit: Na 640 epochs bereikt het een gFID van 1.62, wat een state-of-the-art resultaat is voor ImageNet generatie zonder Classifier-Free Guidance (CFG).
• Robuustheid: VFM-VAE toont een SE-CKNNA score die nauwelijks afwijkt van de schone score (+1.6%), terwijl concurrenten (VA-VAE) een drastische daling vertonen (-33.2%) onder perturbaties.
• Text-to-Image: In combinatie met BLIP3-o voor tekst-naar-beeld generatie, overtreft VFM-VAE VA-VAE aanzienlijk op zowel DPG-Bench (semantische consistentie) als MJHQ-30K (visuele kwaliteit).

5. Betekenis en Conclusie

Dit paper biedt overtuigend bewijs dat Vision Foundation Models, wanneer hun native semantiek wordt behouden door ze bevroren te houden, uitstekende visuele tokenizers kunnen zijn voor Latent Diffusion Models.

• Versnelling van Training: Door de robuustheid van de VFM-representaties direct te benutten, hoeven diffusiemodellen minder tijd te besteden aan het leren van fundamentele semantische structuren, wat leidt tot een enorme versnelling in trainingsduur.
• Architecturale Innovatie: Het bewijst dat de kloof tussen semantische features (van VFMs) en pixel-accurate reconstructie overbrugd kan worden door slimme decoder-architecturen (progressieve reconstructie en multi-scale fusie) in plaats van door distillatie.
• Toekomstperspectief: VFM-VAE biedt een solide basis voor toekomstige generatieve modellen die sneller trainen, robuuster zijn en betere semantische consistentie bereiken.