Wetenschappers hebben een nieuwe methode ontwikkeld voor het genereren van beelden die de traditionele trainingsprocessen van AI-modellen fundamenteel verandert. Het systeem, genaamd Scalable Pixel-anchored End-to-end Diffusion (SpeeDiff), maakt het mogelijk om de trainingstijd van latent diffusion-modellen drastisch te reduceren door verschillende componenten gelijktijdig te optimaliseren.
In de gangbare praktijk worden latent diffusion-modellen in twee strikt gescheiden fasen getraind: eerst wordt een Variational Autoencoder (VAE) ontwikkeld en vastgezet, waarna het diffusiemodel wordt getraind op de resulterende latente ruimte. Volgens een publicatie op arxiv.org breekt SpeeDiff met deze conventie door beide onderdelen vanaf het begin gezamenlijk te trainen. Hierdoor kunnen gradiënten van het diffusieverlies direct invloed uitoefenen op de VAE-encoder, wat in theorie leidt tot een efficiëntere convergentie en een latente ruimte die beter is afgestemd op generatieve taken.
Bestrijding van latent space collapse
Hoewel een gezamenlijke training voordelen biedt, stuiten onderzoekers op een technisch probleem: ongebreidelde backpropagation kan leiden tot een zogenaamde 'latent space collapse'. Dit fenomeen zorgt ervoor dat de prestaties van het model ernstig achteruitgaan. Om dit te voorkomen, hebben de auteurs een innovatie geïntroduceerd genaamd de Tweedie Pixel Reconstruction (TPR) loss.
Zoals beschreven in het volledige thecvf.com, biedt deze TPR-loss extra feedback op pixelniveau. Door een voorspelde schone latente staat uit een tussenliggende ruisstatus te decoderen via de formule van Tweedie, wordt de stabiliteit van het trainingsproces gewaarborgd en de instorting van de latente ruimte voorkomen.
Prestaties en schaalbaarheid
De praktische resultaten van het model, specifiek de SpeeDiff-XL variant, tonen een enorme sprong in efficiëntie. In vergelijking met Vanilla SiT is de training meer dan 140 keer sneller, terwijl de winst ten opzichte van REPA ongeveer 61 keer bedraagt. Ondanks deze versnelling blijft de beeldkwaliteit hoog. Bij het genereren van ImageNet-beelden met een resolutie van 256x256 pixels behaalde het model een FID-score van 1,50, zonder dat er gebruik werd gemaakt van guidance.
Het systeem blijkt bovendien schaalbaar naar hogere resoluties. Door gebruik te maken van een efficiëntere VAE met een compressiefactor van 32x, kon een FID van 1,53 worden bereikt voor beelden van 512x512 pixels. De volledige architectuur is gebaseerd op transformers, wat een betere uitlijning van representaties binnen het end-to-end raamwerk faciliteert.
Academische context
Het project is het resultaat van een samenwerking tussen diverse wetenschappers, waaronder onderzoekers van Caltech. Wenda Chu werkte aan dit project onder supervisie van professor Yisong Yue en in samenwerking met Yang Song. De focus van professor Yue ligt op het verbeteren van de efficiëntie en het schalen van leerprocessen, zoals vermeld op zijn yisongyue.com.
Het werk is officieel gepresenteerd op de paperdigest.org, een van de meest vooraanstaande conferenties voor computer vision. De uitgebreide documentatie, bestaande uit 29 pagina's met talloze figuren en tabellen, onderbouwt de effectiviteit van de TPR-loss en de schaalbaarheid van de gekozen transformer-architectuur.