Coevolving Representations in Joint Image-Feature Diffusion

Probleemstelling

Diffusiemodellen zijn de huidige standaard voor hoogwaardige beeldsynthese, maar ze opereren vaak voornamelijk op pixel- of VAE-latentruimtes die zich richten op lage-niveau statistieken. Recent werk heeft geprobeerd semantische priors in te bouwen door vooraf getrainde visuele encoders (zoals DINOv2) te koppelen aan diffusiemodellen.

De kernbeperking van bestaande benaderingen (zoals ReDi) is dat de semantische representatieruimte statisch is. Deze ruimte wordt vooraf gedefinieerd (bijv. via PCA) en blijft onveranderd tijdens het trainingproces, onafhankelijk van het generatieve doel. De auteurs stellen dat dit suboptimaal is: de representatieruimte die de diffusie stuurt, zou zich moeten aanpassen aan de specifieke behoeften van de beeldsynthese. Een naïeve poging om deze projectie samen met het diffusiemodel te optimaliseren, leidt echter vaak tot degenererende oplossingen (zoals feature collapse), waarbij de representaties hun betekenis verliezen.

Methodologie: CoReDi

De auteurs introduceren CoReDi (Coevolving Representation Diffusion), een raamwerk waarin de projectie van visuele features en het diffusiemodel gezamenlijk worden getraind. De representatieruimte evolueert dus tijdens het trainen.

De methode bestaat uit drie cruciale componenten om stabiele co-evolutie te garanderen:

1. Leerbaar Projectie: In plaats van een vaste PCA-projectie, wordt een learnable lineaire laag ($g_\phi$) gebruikt om de output van een bevroren visuele encoder naar een lagere dimensie te projecteren. Deze laag wordt samen met het diffusiemodel geoptimaliseerd.
2. Stop-Gradient (SG): Om te voorkomen dat het model de representatiedoelwit ($\tilde{z}_0$) triviaal aanpast om de verliesfunctie te minimaliseren, wordt een stop-gradient toegepast op de schone representatie in de loss-functie. Hierdoor leert het diffusiemodel de noise te verwijderen zonder dat de projectie de target zelf "vervalst".
3. Batch Normalization (BN): Na de learnable projectie wordt batch normalization toegepast (zonder trainbare affine parameters). Dit stabiliseert de schaal van de features, voorkomt dat de noise-schedule wordt verstoord, en fungeert als een impliciete regularisatie tegen sample collapse.
4. Expliciete Regularisatie: Zelfs met SG en BN treedt er nog steeds "feature collapse" op (waarbij kanalen redundante informatie coderen). De auteurs introduceren drie strategieën om dit te voorkomen:
   • Feature Variance Regularization: Straft kanalen af die te weinig variatie vertonen over de steekproeven.
   • Orthogonaliteitsregularisatie: Straft de projectiematrix af als deze niet orthogonaal is, om redundantie te voorkomen.
   • Covariance Regularization: Minimaliseert de correlatie tussen verschillende kanalen.

Het totale trainingsdoel combineert de image-flow matching loss, de representatie-flow matching loss, en een regularisatieterm voor de projectieparameters.

Belangrijkste Bijdragen

• CoReDi Framework: Een nieuw paradigma waarbij de semantische representatieruimte niet statisch is, maar dynamisch evolueert samen met het generatieve model.
• Identificatie van Stabiliteitsfactoren: Systematische analyse toont aan dat stop-gradient, batch normalization en expliciete regularisatie allemaal noodzakelijk zijn om degeneratie te voorkomen.
• Ruimtelijke Structuur: Het paper toont aan dat co-evoluerende representaties tijdens het trainen een sterkere ruimtelijke organisatie ontwikkelen, wat essentieel is voor kwalitatief hoogwaardige synthese.
• Extensie naar Pixelruimte: CoReDi wordt succesvol toegepast op zowel VAE-latentruimtes als directe pixelruimte-diffusie (zonder VAE-compressie), wat bewijst dat de methode robuust is voor verschillende generatieve instellingen.

Resultaten

De experimenten zijn uitgevoerd op ImageNet (256x256) met verschillende modelgroottes (SiT-B/2 en SiT-XL/2) en architecturen (VAE-latent en pixel-space/DeCo).

• Snellere Convergentie: CoReDi convergeert aanzienlijk sneller dan statische baselines.
  • In latentruimte (XL/2): CoReDi bereikt een FID van 3.3 in 2M iteraties, terwijl ReDi 4M iteraties nodig heeft voor hetzelfde resultaat. Dit is een ~13x snellere convergentie vergeleken met REPA.
  • In pixelruimte: CoReDi verbetert de convergentie met een factor 2 ten opzichte van DeCo.
• Hogere Kwaliteit: CoReDi bereikt betere FID-scores dan statische methoden (ReDi, REPA) en de standaard SiT-baselines, zowel met als zonder Classifier-Free Guidance (CFG).
  • Bijvoorbeeld: CoReDi-B/2 bereikt een FID van 16.4 na 400K iteraties, terwijl ReDi 21.4 haalt.
• Robuustheid: De methode werkt consistent goed over verschillende visuele encoders (DINOv2, MOCOv3, SigLIPv2, MAE).
• Ablatie Studies:
  • Zonder stop-gradient of batch normalization stort het trainingsproces volledig in (FID > 200).
  • Zonder expliciete regularisatie treedt feature collapse op, wat resulteert in slechtere prestaties dan statische PCA-projecties.
  • Feature Variance Regularization bleek de meest effectieve regularisatiestrategie.

Significantie

Dit paper is significant omdat het een fundamentele aanname in joint image-feature diffusion doorbreekt: dat representatieruimtes statisch moeten zijn. Door te tonen dat het gezamenlijk leren van de projectie leidt tot representaties met een beter ontwikkelde ruimtelijke structuur, biedt CoReDi een nieuwe richting voor het verbeteren van generatieve modellen. Het bewijst dat adaptieve semantische gidsen niet alleen nuttig zijn voor latentruimtes, maar ook de barrières van VAE-compressie kunnen overwinnen in pixelruimte-diffusie, wat leidt tot snellere training en hogere beeldkwaliteit.