Technische Samenvatting: RePack then Refine

Probleemstelling

Recente vooruitgang in Latent Diffusion Models (LDM's) maakt steeds meer gebruik van Vision Foundation Models (VFMs), zoals DINOv2/v3 en CLIP, als feature-encoders om de semantische rijkdom te vergroten. Het directe gebruik van ruwe VFM-features vertoont echter een kritieke inefficiëntie. Zoals in het artikel wordt opgemerkt, zijn ruwe features van standaard ViT-encoders (bijvoorbeeld DINOv3-B/16) hoogdimensionaal (768 kanalen) en vaak redundant. Het totale aantal elementen in deze feature maps is frequent gelijk aan of groter dan het aantal pixels van de originele afbeelding, wat in strijd is met het kernprincipe van LDM's om te opereren in een compacte, gecomprimeerde latent space.

De auteurs betogen dat het uitvoeren van diffusie in zo'n hoogdimensionale, redundante ruimte de Diffusion Transformer (DiT) dwingt om verdelingen te leren die worden gedomineerd door ruis en minder kritieke informatie, wat de convergentie aanzienlijk vertraagt en de trainings-efficiëntie verlaagt. Empirische observaties tonen aan dat methoden die ruwe features gebruiken (bijvoorbeeld RAE) langzamer convergeren dan die welke gecomprimeerde ruimtes gebruiken, ondanks de rijkere semantische input van de eerste.

Methodologie: RePack then Refine

Om de afweging tussen semantische rijkdom en trainings-efficiëntie aan te pakken, stellen de auteurs RePack then Refine voor, een drie-fasen framework dat is ontworpen om VFM-features te comprimeren tot een compacte manifold terwijl hoogfrequente details worden hersteld.

1. Representation Packing (RePack)

De eerste fase introduceert een lichtgewicht, bias-vrije lineaire projectielaag ($P_\theta$) die hoogdimensionale ruwe VFM-features ($z_{raw} \in \mathbb{R}^{D \times h \times w}$) projecteert op een compacte, laagdimensionale manifold ($z \in \mathbb{R}^{d \times h \times w}$, waarbij $d \ll D$, bijvoorbeeld $d=32$).

• Mechanisme: De bevroren VFM haalt features op, die vervolgens worden geprojecteerd door $P_\theta$. Een decoder ($D_\psi$) reconstrueert de afbeelding vanuit deze gecomprimeerde latent.
• Doel: Deze fase filtert redundantie terwijl essentiële structurele informatie behouden blijft. De auteurs veronderstellen dat semantische informatie ligt op een laagdimensionale manifold, en het projecteren op deze manifold vereenvoudigt de taak van verdelingsmodellering voor de daaropvolgende DiT.
• Training: Geoptimaliseerd end-to-end met een samengestelde loss die $L_1$, LPIPS, adversarial loss, Watson perceptual loss en Focal Frequency Loss omvat.

2. Generative Modeling (DiT)

In de tweede fase wordt een standaard Diffusion Transformer (specifiek LightningDiT) getraind op de compacte latent space die is gegenereerd door de RePack-encoder.

• Proces: De DiT leert het snelheidsveld om Gaussische ruis te transformeren naar de data-verdeling binnen de gecomprimeerde 32-kanaals manifold.
• Efficiëntie: Omdat de latent space schoon en semantisch coherent is, convergeert de DiT aanzienlijk sneller dan wanneer deze wordt getraind op ruwe, hoogdimensionale features.

3. Latent-Guided Refinement

Om de hoogfrequente textuurdetails te compenseren die verloren gaan tijdens de compressie in de RePack-fase, wordt een Latent-Guided Refiner geïntroduceerd.

• Architectuur: Een op U-Net gebaseerd model dat de basisafbeelding (gedecodeerd vanuit de latent) en de geupsamplede latent features als input neemt.
• Functie: Het modelleert expliciet hoogfrequente details in de pixel-domein, geleid door de structurele semantiek van de latent space. Dit ontkoppelt semantische generatie (afgehandeld door de DiT) van textuurverfijning (afgehandeld door de Refiner).
• Training: Offline getraind in een derde fase met ground-truth afbeeldingen om de mapping van grove semantiek naar fijne details te leren.

Belangrijkste Bijdragen

1. RePack Framework: Het voorstellen van een semantisch compressieframework dat redundantie in ruwe VFM-features identificeert en een compacte manifold-bottleneck inzet om ruis te filteren terwijl kernsemantiek behouden blijft.
2. Efficiënte Generatieve Modellering: Aantonen dat generatieve modellering in deze compacte ruimte state-of-the-art convergentiesnelheden oplevert. De RePack-DiT-XL/1 bereikt een FID van 1,82 in slechts 64 trainings-epochs.
3. Latent-Guided Refiner: De introductie van een module die structuurgeneratie ontkoppelt van textuurverfijning, hoogfrequente details die verloren gaan tijdens packing herstelt en de uiteindelijke FID verbetert naar 1,65.

Experimentele Resultaten

Geëvalueerd op ImageNet-1K (256×256), demonstreert het framework superieure efficiëntie en kwaliteit:

• Convergentiesnelheid: RePack-DiT-XL/1 bereikt een FID van 1,82 in 64 epochs. Ter vergelijking: een standaard DiT vereist 1.400 epochs om een FID van 2,27 te bereiken, en RAE vereist 800 epochs om 1,41 te bereiken (hoewel met een veel groter aantal parameters en verschillende trainingsdynamiek).
• Kwaliteitsverbetering: Met de toevoeging van de Refiner verbetert de FID naar 1,65 bij 64 epochs, waarmee recente Pixel Diffusion Models (PDM's) en andere LDM's wordt overtroffen qua convergentie-efficiëntie.
• Ablatie Studies:
  • Projectie Ontwerp: Bias-vrije lineaire projectie presteert beter dan bias-varianten en statische PCA-projecties, wat de noodzaak bevestigt van een leerbare projectie voor het vastleggen van de semantische manifold.
  • Latent Dimensie: Het verminderen van de dimensies naar $d=16$ biedt de snelste convergentie, maar $d=32$ wordt als standaard gekozen om reconstructiefideliteit en eerlijke vergelijking met baselines zoals VA-VAE in evenwicht te brengen.
  • Refiner Generalisatie: De Refiner verbetert RePack aanzienlijk, maar verslechtert de prestaties op andere VAE's (zoals VA-VAE), wat suggereert dat de Refiner specifiek is afgestemd op de semantisch-pure, textuur-sparse aard van de latent space van RePack.

Betekenis en Claims

Het artikel claimt dat RePack then Refine een uiterst effectieve strategie vertegenwoordigt voor het balanceren van generatieve fideliteit met trainings-efficiëntie. Door de redundantie in ruwe VFM-features te erkennen, betogen de auteurs dat het "pakken" van deze features in een compacte manifold superieur is aan direct gebruik. Het framework bereikt de snelst convergerende resultaten onder recente methoden, en bereikt hoogwaardige generatie (FID 1,65) in een fractie van de trainings-epochs die nodig zijn voor state-of-the-art baselines.

De auteurs benadrukken dat hun aanpak geen complexe, parameter-zware encoders vereist (in tegenstelling tot FAE) of het fine-tunen van foundation models (in tegenstelling tot PS-VAE), waardoor de universaliteit van de VFM behouden blijft terwijl de computatielast van het generatieve model drastisch wordt verminderd. Het werk suggereert dat de toekomst van efficiënte diffusie ligt in het ontkoppelen van semantische compressie van detailherstel, waarbij gebruik wordt gemaakt van de complementaire sterke punten van latent en pixel ruimtes.

Opmerking: De auteurs erkennen beperkingen met betrekking tot de huidige reikwijdte van experimenten (beperkt tot 80 epochs en ImageNet-1K) en merken op dat schalen naar text-naar-afbeelding generatie en het verkennen van uitgebreide trainingsplanningen richtingen zijn voor toekomstig werk.