Unified Latents (UL): How to train your latents

Dit paper introduceert Unified Latents (UL), een framework dat latent representaties leert door een encoder te koppelen aan een diffusiemodel, wat resulteert in een efficiëntere training met state-of-the-art prestaties op zowel afbeeldingen als video.

De kern

🎨 De Kunst van het Samenvatten: Wat is "Unified Latents"?

Stel je voor dat je een enorme, gedetailleerde olieverfschildering (een foto) wilt sturen naar een vriend, maar je hebt een heel kleine postbus (je internetverbinding of opslagruimte). Je moet de foto dus eerst samenvatten tot een klein briefje.

Het probleem? Als je het briefje te klein maakt, is de foto onherkenbaar. Als je het te groot maakt, past hij niet in de postbus. En als je het briefje te "moeilijk" maakt (vol met cryptische codes), kan je vriend het misschien wel lezen, maar duurt het eeuwen om de foto weer te maken.

De auteurs van dit paper (van Google DeepMind) hebben een nieuwe manier bedacht om die "postbus" te vullen. Ze noemen het Unified Latents (Verenigde Latenten). Het is een slimme methode om beelden te comprimeren en ze later weer perfect te reconstrueren, zodat kunstmatige intelligentie (AI) er beter en sneller mee kan werken.

Hier is hoe het werkt, stap voor stap:

1. Het Drie-Deel Team

In plaats van één persoon die alles doet, laten ze drie specialisten samenwerken:

1. De Samenvatter (De Encoder): Deze kijkt naar de foto en schrijft een samenvatting op een briefje.
2. De Controleur (De Diffusion Prior): Deze kijkt naar het briefje en zegt: "Hé, dit briefje is te rommelig. Het moet netjes en gestructureerd zijn, alsof het uit een standaard doosje komt."
3. De Schilder (De Diffusion Decoder): Deze leest het briefje en probeert de oorspronkelijke foto weer te schilderen.

2. Het Grote Nieuw: De "Ruis" als Regel

Bij oude methoden (zoals bij Stable Diffusion) was het vaak een gok. De samenvatter schreef iets op, en de schilder probeerde het te raden. Soms lukte het goed, soms niet. De "samenvatting" kon te veel informatie bevatten (te zwaar) of te weinig (te vaag).

Bij Unified Latents doen ze iets heel slims:
Ze laten de Samenvatter het briefje niet helemaal "schoon" schrijven. Ze laten hem er een beetje ruis (statiek) op zetten.

• De Analogie: Stel je voor dat je een briefje schrijft, maar je doet het in een trillende trein. Je kunt niet perfect schrijven; je moet rekening houden met die trilling.
• Het Geniale: De Controleur (die ook een AI is) is getraind om precies die trilling te begrijpen. Omdat de samenvatter wéét dat de controleur die trilling kan lezen, schrijft hij zijn boodschap precies op de manier die het makkelijkst te begrijpen is voor de controleur.

Dit zorgt ervoor dat het briefje (de "latent") altijd de perfecte hoeveelheid informatie heeft: niet te veel, niet te weinig.

3. Waarom is dit beter dan de rest?

• Geen Gokwerk: Bij oude methoden moest je handmatig instellen hoeveel informatie er op het briefje mocht staan. Dat was als proberen een auto te bouwen door blindelings bouten te draaien. Bij Unified Latents is het automatisch geregeld.
• Beter Kwaliteit: Omdat het briefje zo goed is gestructureerd, kan de Schilder (de decoder) de foto veel scherper en realistischer terugmaken. Ze krijgen betere resultaten (zoals een FID-score van 1.4, wat extreem goed is) met minder rekenkracht.
• Videofoto's: Het werkt niet alleen voor plaatjes, maar ook voor video's. Ze hebben een record neergezet voor het maken van video's (Kinetics-600 dataset).

4. De "Postbus" en de "Schilder"

Een belangrijk punt in het paper is het evenwicht tussen reconstructie (hoe goed lijkt de foto op het origineel?) en generatie (hoe makkelijk is het voor de AI om nieuwe foto's te bedenken?).

• Als je het briefje te gedetailleerd maakt, is de foto perfect, maar is het voor de AI te moeilijk om nieuwe ideeën te bedenken.
• Als je het briefje te simpel maakt, is het makkelijk voor de AI, maar ziet de foto er wazig uit.

Unified Latents vinden het gouden middenpad. Ze gebruiken een slimme "loss factor" (een soort volume-knop) om precies te bepalen hoeveel detail er op het briefje moet.

🚀 Conclusie in één zin

Unified Latents is een slimme manier om AI te leren hoe je een foto samenvat tot een perfect leesbaar briefje, zodat de AI er later niet alleen sneller mee kan werken, maar ook mooiere en scherpere beelden van kan maken.

Het is alsof je niet meer probeert een hele bibliotheek in één koffer te proppen, maar een slimme samenvatting schrijft die precies past in je postbus, maar waaruit je de hele bibliotheek weer kunt reconstrueren.

Technische samenvatting

Probleemstelling

Diffusiemodellen hebben zich bewezen als uiterst succesvol voor generatie van afbeeldingen, video en audio. Een cruciaal aspect van dit succes is het gebruik van latente representaties: compacte coderingen die het mogelijk maken om diffusiemodellen efficiënter te schalen naar hogere resoluties.

Echter, de vraag hoe deze latente variabelen het beste geleerd kunnen worden, blijft onduidelijk. Bestaande methoden hebben verschillende tekortkomingen:

• Traditionele VAE's (bijv. Stable Diffusion): Gebruiken een KL-straf (Kullback-Leibler) tussen de latente verdeling en een standaard Gaussische verdeling. Omdat de decoder geen likelihood-gebaseerd verlies gebruikt, moet het gewicht van de KL-term handmatig worden ingesteld. Dit maakt het moeilijk om de informatie-inhoud van de latente variabelen te kwantificeren of te controleren.
• Semantische latente variabelen: Recent werk gebruikt vooraf getrainde netwerken (zoals DINO) of zwaar geregulariseerde auto-encoders. Hoewel deze makkelijker te leren zijn en goede FID-scores (Fréchet Inception Distance) opleveren, gaat vaak hoge-frequentie informatie verloren, wat resulteert in slechtere PSNR-waarden (Peak Signal-to-Noise Ratio) en reconstructie-artefacten.

Er bestaat een fundamenteel trade-off tussen de informatie-inhoud van de latente variabelen en de reconstructiekwaliteit. Hoe makkelijker een latente variabele te modelleren is (minder informatie), hoe beter de generatieprestaties, maar hoe slechter de reconstructie. Het doel is om dit trade-off systematisch te navigeren.

Methodologie: Unified Latents (UL)

De auteurs introduceren Unified Latents (UL), een raamwerk waarbij latente representaties gezamenlijk worden geregulariseerd door een diffusie-prior en worden gedecodeerd door een diffusie-decoder. De kern van de methode ligt in het koppelen van de encoder-output aan de precisie van het diffusiemodel.

De aanpak rust op drie pijlers:

1. Encoderen met vaste ruis:
   In plaats van een complexe verdeling te leren, voorspelt de encoder een enkele deterministische latente variabele ($z_{clean}$). Deze wordt vervolgens expliciet "verstoord" (noised) naar een tijdstap $t=0$ met een vaste hoeveelheid Gaussische ruis. Dit koppelt de encoder-ruis direct aan het minimale ruisniveau van de prior. De KL-term reduceert hierdoor tot een eenvoudige, gewogen MSE (Mean Squared Error) over de ruisniveaus.

2. Diffusie-prior:
   Een diffusiemodel (de prior) wordt getraind om de weg van pure ruis ($z_1$) naar de licht verstoord latente variabele ($z_0$) te modelleren. Omdat de prior de informatie-inhoud van $z_0$ meet en regulariseert, ontstaat er een interpreteerbare bovengrens voor het aantal bits in de latente variabelen. Het verlies voor de prior is een ongewogen ELBO (Evidence Lower Bound) om te voorkomen dat de encoder informatie op verwaarloosde ruisniveaus verbergt.

3. Diffusie-decoder met herschaling:
   De decoder is ook een diffusiemodel, maar werkt in de beeldruimte en conditieert op zowel de verstoord beeld ($x_t$) als de latente variabele ($z_0$). Om het probleem van "posterior collapse" (waarbij de decoder de latente variabelen negeert) te voorkomen, wordt het decoder-verlies opgewogen met een loss factor (bijv. 1.3 tot 1.7). Daarnaast wordt een herschaling (sigmoid weighting) gebruikt om de bijdrage van hoge frequenties (kleine details) in het verlies te beïnvloeden.

Trainingsproces:

• Fase 1: Gezamenlijke training van encoder, prior en decoder. De encoder leert een deterministische $z_{clean}$, de prior leert de verdeling van $z_0$, en de decoder leert de reconstructie.
• Fase 2 (Base Model Training): De encoder en decoder worden bevroren. De prior wordt opnieuw getraind als een "base model" met een sigmoid-gewichting (in plaats van ongewogen ELBO) om betere generatiekwaliteit te bereiken. Omdat de encoder bevroren is, kan dit model groter zijn en op grotere batches worden getraind.

Belangrijkste Bijdragen

• Interpreteerbare Bitrate: Door de link tussen encoder-ruis en de prior, biedt UL een strakke bovengrens op het aantal bits per dimensie (bpd) in de latente ruimte. Dit maakt het mogelijk om de trade-off tussen reconstructie en modellering expliciet te controleren via hyperparameters (zoals de loss factor).
• Stabiliteit: Het gebruik van een deterministische encoder met vaste ruis elimineert de instabiliteit die vaak voorkomt bij het leren van complexe encoder-verdelingen (zoals in LSGM).
• Efficiëntie: Het framework vereist minder trainings-FLOPs (floating point operations) om vergelijkbare of betere resultaten te bereiken dan modellen die op Stable Diffusion latentes zijn getraind.

Resultaten

De auteurs evalueren UL op ImageNet-512 (afbeeldingen) en Kinetics-600 (video).

• Afbeeldingen (ImageNet-512):

  • UL bereikt een FID van 1.4, wat zeer competitief is.
  • Het behoudt hoge reconstructiekwaliteit (hoge PSNR) in tegenstelling tot methoden die semantische latentes gebruiken.
  • In vergelijking met baselines (zoals Stable Diffusion latentes en UNet-architecturen) presteert UL beter op het vlak van trainingskosten versus generatiekwaliteit.
  • Er is een duidelijke correlatie gevonden: kleinere base-modellen profiteren van lagere bitrates (minder informatie in de latente variabele), terwijl grotere modellen beter presteren met hogere bitrates.

• Video (Kinetics-600):

  • UL zet een nieuwe state-of-the-art FVD (Fréchet Video Distance) van 1.3 neer.
  • Het overtreft bestaande methoden zoals MAGVIT en W.A.L.T. op het gebied van trainingskosten versus prestaties.

• Ablatie-studies:

  • Het verwijderen van de diffusie-prior of het vervangen ervan door een standaard Gaussische prior leidt tot slechtere resultaten en instabiliteit.
  • Het gebruik van een leerbare variantie in de encoder (in plaats van vaste ruis) resulteert in instabiliteit en hogere variantie in de schatting van de KL-term.

Betekenis en Conclusie

Het Unified Latents-framework biedt een principieel en stabiel alternatief voor het ontwerpen van latente ruimtes voor diffusiemodellen. Door de encoder, prior en decoder gezamenlijk te trainen met een expliciete controle over de informatie-inhoud, lost UL het probleem op van het handmatig afstellen van regularisatie.

De belangrijkste implicaties zijn:

1. Efficiëntie: Het is de meest efficiënte pre-training-aanpak voor generatieve modellen op ImageNet en Kinetics.
2. Controle: Het biedt onderzoekers en ingenieurs eenvoudige knoppen (loss factor, sigmoid bias) om de balans tussen reconstructiekwaliteit en generatiecomplexiteit te sturen.
3. Toekomstperspectief: Hoewel dit werk zich richt op afbeeldingen en video, is het raamwerk breed toepasbaar, inclusief potentieel voor discrete data zoals tekst. De auteurs wijzen er echter op dat diffusie-decoders momenteel duurder zijn om te sample dan GAN-gebaseerde decoders, wat een uitdaging blijft voor de praktische toepassing zonder distillatie.

Samenvattend bewijst dit paper dat het gezamenlijk trainen van een diffusie-prior en decoder de sleutel is tot het leren van robuuste, informatieve en goed te modelleren latente representaties.