Technische Samenvatting: Laminerende Representatie Autoencoders voor Efficiënte Diffusie (FlatDINO)

Probleemstelling

Recente ontwikkelingen in beeldgeneratie hebben aangetoond dat diffusiemodellen direct kunnen opereren op self-supervised learning (SSL) patch-features (bijv. van DINOv2) in plaats van op pixel-ruimte latents. Deze aanpak, geïllustreerd door RAE (Zheng et al., 2025), omzeilt de reconstructie-generatie-trade-off die inherent is aan pixel-getrainde VAE's door gebruik te maken van de semantisch betekenisvolle structuur die al in de SSL-features is gecodeerd.

Er blijft echter een aanzienlijke computationele bottleneck bestaan. Standaard SSL-encoders zoals DINOv2 produceren dichte tweedimensionale grids van hoogdimensionale patch-embeddings (bijv. 256 patches van 768 dimensies voor een 256×256 afbeelding). Deze representatie is qua omvang vergelijkbaar met het oorspronkelijke signaal, wat geen computationeel voordeel biedt ten opzichte van pixel-ruimte diffusie. Het dichte grid bevat aanzienlijke ruimtelijke redundantie omdat naburige patches veel semantische inhoud delen, waardoor diffusie op deze features onnodig duur is in termen van FLOPs en geheugen.

Methodologie

De auteurs stellen FlatDINO voor, een variational autoencoder (VAE) die ontworaceerd is om de dichte, tweedimensionale DINOv2 patch-embeddings te comprimeren tot een compacte, eendimensionale sequentie van continue tokens.

1. Architectuur

• Encoder: Een Vision Transformer (ViT-B) die de 256 patch-embeddings van een bevroren DINOv2 ViT-B/14 (met registers) als input neemt. De encoder voegt $T$ leerbare register-tokens toe aan de patch-sequentie. Na verwerking worden alleen deze $T$ register-tokens behouden als de latente representatie, waarbij de oorspronkelijke 2D-ruimtelijke structuur wordt weggegooid.
• Decoder: Een ViT-L die leerbare register-tokens gebruikt als queries om aandacht te besteden aan de gecomprimeerde latente sequentie en de oorspronkelijke 256 DINOv2 patch-embeddings te reconstrueren.
• Trainingsdoel: Het model wordt getraind met een standaard VAE-doelstelling met een Mean Squared Error (MSE) reconstructieverlies en een KL-divergentie straf. De KL-gewicht ($\beta$) wordt genormaliseerd door de totale latente dimensionaliteit ($T \times d$) om een consistente regularisatiedruk te garanderen over verschillende configuraties.
• Latente Generatie: Een generatief model (specifiek een DiT-XL met flow matching) wordt getraind op de gecomprimeerde FlatDINO latente ruimte. Het model leert een snelheidveld (velocity field) om samples te transporteren van een standaard Gaussische prior naar de datadistributie langs rechte paden.

2. Belangrijke Ontwerpkeuzes

• 1D Compressie: Geïnspireerd door TiTok (Yu et al., 2024a), comprimeert de methode het 2D-grid naar een 1D-sequentie. De auteurs vonden dat 32 tokens een gunstige afweging bieden tussen reconstructiefaillibiliteit en sequentielengte.
• Reductie van Feature-dimensie: Hoewel het aantal tokens wordt verminderd van 256 naar 32 (een 8× reductie), wordt ook de per-token feature-dimensie gecomprimeerd (bijv. van 768 naar 128) om de totale dimensionaliteit verder te verlagen.
• Flow Matching: Het generatieve model maakt gebruik van flow matching in plaats van standaard diffusie, waarbij het direct opereert op de 1D latente sequentie met geleerde positionele embeddings.

Belangrijkste Bijdragen

1. FlatDINO: De eerste methode die self-supervised patch-features comprimeert naar een eendimensionale latente representatie, waarbij de oorspronkelijke 2D-ruimtelijke structuur effectief wordt weggegooid terwijl de semantische inhoud behouden blijft.
2. Efficiëntiewinst: Demonstreert dat diffusiemodellen getraind op FlatDINO-latents competitieve generatiekwaliteit bereiken terwijl ze de computationele kosten aanzienlijk verlagen. Specifiek leidt de 8× reductie in sequentielengte tot een 8× reductie in FLOPs per forward pass en tot wel een 4,5× reductie in FLOPs per trainingsstap vergeleken met diffusie op ongecomprimeerde DINOv2 features.
3. Analyse van Tokenisatie: Door middel van token-ablatiestudies laten de auteurs zien dat het model lokale, blob-achtige receptieve velden leert voor de meeste tokens, waardoor ze de afbeelding ruimtelijk partitioneren. Ze identificeren een "faseovergang" waarbij het reduceren van tokens onder een bepaalde drempel (bijv. van 32 naar 16) de methode dwingt om minder efficiënte encoding-strategieën (horizontale strepen) te adopteren, wat leidt tot een scherpe daling in kwaliteit.

Resultaten

Experimenten werden uitgevoerd op ImageNet 256×256 met een DiT-XL architectuur.

• Generatiekwaliteit: Een DiT-XL getraind op FlatDINO-latents (32×128) bereikt een gFID van 1,85 met classifier-free guidance (CFG) en 3,34 zonder CFG. Dit is competitief met RAE (die 1,41 gFID bereikt met CFG op ongecomprimeerde features) en andere state-of-the-art latente diffusiemodellen, ondanks de significante compressie.
• Computationele Efficiëntie:
  • Inference: FlatDINO vereist ongeveer 8× minder FLOPs per forward pass vergeleken met het opereren op volledige DINOv2 patches.
  • Training: Vanwege de verminderde sequentielengte en de inclusie van de encoder-overhead, vereist FlatDINO 4,5× minder FLOPs per trainingsstap voor DiT-XL modellen vergeleken met de RAE baseline.
• Reconstructie: De 32×128 configuratie bereikt een rFID van 0,77 (vergeleken met 0,62 voor de ongecomprimeerde RAE baseline), wat wijst op een hoogwaardige reconstructie van de feature-ruimte.
• Token Ablatie: De studie onthult dat 32 tokens voldoende zijn om gelokaliseerde receptieve velden te behouden. Het reduceren naar 16 tokens veroorzaakt een kwalitatieve verschuiving naar horizontale streep-encoding en een significante daling in reconstructiekwaliteit.

Betekenis en Claims

Het artikel positioneert FlatDINO als een stap naar efficiënte, semantisch gegronde beeldsynthese. De auteurs beweren dat door gebruik te maken van de ruimtelijke redundantie die inherent is aan 2D patch-grids, het mogelijk is om een 8× reductie in het aantal tokens te realiseren zonder de voordelen van convergentie te verliezen die gepaard gaan met het opereren in een semantisch georganiseerde latente ruimte.

Het werk suggereert dat diffusie op gecomprimeerde semantische representaties een praktische aanpak is voor omgevingen met beperkte middelen. De auteurs erkennen dat hun huidige resultaten preliminair zijn en dat het model nog niet volledig is geconvergeerd (getraind voor 600 epochs versus 800 voor RAE). Ze schrijven het resterende gat in generatiekwaliteit toe aan een onvoldoende lange trainingsduur en de noodzaak voor diffusie-recepten die specifiek zijn afgestemd op gecomprimeerde semantische latents. Toekomstig werk wordt voorgesteld om uitgebreidere trainingsschema's en geoptimaliseerde SSL-autoencoders te onderzoeken om dit gat te dichten.