Introduktion til diffusionsmodeller til maskinlæring

Denne artikel blev oprindeligt offentliggjort den ForsamlingAI og genudgivet til TOPBOTS med tilladelse fra forfatteren.

Diffusionsmodeller er generative modeller, som har vundet betydelig popularitet i de sidste mange år, og det er der god grund til. En håndfuld banebrydende papirer udgivet i 2020'erne alene har vist verden, hvad Diffusion-modeller er i stand til, såsom at slå GAN'er^[6] om billedsyntese. Senest vil praktiserende læger have set diffusionsmodeller brugt i DALL-E2, OpenAIs billedgenereringsmodel udgivet i sidste måned.

I betragtning af den seneste bølge af succes fra Diffusion Models, er mange maskinlæringsudøvere helt sikkert interesserede i deres indre virke. I denne artikel vil vi undersøge teoretiske grundlag for diffusionsmodeller, og demonstrer derefter, hvordan man genererer billeder med en Diffusionsmodel i PyTorch. Lad os dykke ned!

Hvis dette dybdegående undervisningsindhold er nyttigt for dig, abonner på vores AI-mailingliste for at blive advaret, når vi udgiver nyt materiale. 

Diffusionsmodeller – Introduktion

Diffusionsmodeller er generative modeller, hvilket betyder, at de bruges til at generere data svarende til de data, de er trænet på. Grundlæggende fungerer diffusionsmodeller efter ødelægge træningsdata gennem den successive tilføjelse af Gaussisk støj, og derefter lære at restituere dataene ved bakning denne støjproces. Efter træning kan vi bruge Diffusionsmodellen til at generere data ved ganske enkelt at sende tilfældigt samplet støj gennem den indlærte denoising-proces.

Mere specifikt er en diffusionsmodel en latent variabel model, som afbildes til det latente rum ved hjælp af en fast Markov-kæde. Denne kæde tilføjer gradvist støj til dataene for at opnå den omtrentlige posteriore q(x_1:T|x₀), hvor x₁,…,x_T er de latente variable med samme dimensionalitet som x₀. I figuren nedenfor ser vi sådan en Markov-kæde manifesteret for billeddata.

I sidste ende bliver billedet asymptotisk transformeret til ren Gauss-støj. Det mål at træne en diffusionsmodel er at lære vende proces – dvs. træning s_θ(x_t−1|x_t). Ved at gå baglæns langs denne kæde kan vi generere nye data.

Fordele ved diffusionsmodeller

Som nævnt ovenfor er forskningen i diffusionsmodeller eksploderet i de senere år. Inspireret af termodynamik uden ligevægt^[1], Diffusion Models producerer i øjeblikket Topmoderne billedkvalitet, hvoraf eksempler kan ses nedenfor:

Ud over banebrydende billedkvalitet kommer diffusionsmodeller med en lang række andre fordele, bl.a ikke kræver modstanderuddannelse. Vanskelighederne ved modstridende træning er veldokumenterede; og i tilfælde, hvor der findes ikke-modstridende alternativer med sammenlignelig ydeevne og træningseffektivitet, er det normalt bedst at bruge dem. Med hensyn til træningseffektivitet har Diffusionsmodeller også de ekstra fordele skalerbarhed og paralleliserbarhed.

Mens diffusionsmodeller næsten ser ud til at producere resultater ud af den blå luft, er der en masse omhyggelige og interessante matematiske valg og detaljer, der danner grundlaget for disse resultater, og bedste praksis udvikler sig stadig i litteraturen. Lad os tage et kig på den matematiske teori, der understøtter diffusionsmodeller mere detaljeret nu.

Diffusionsmodeller – et dybt dyk

Som nævnt ovenfor består en diffusionsmodel af en fremadrettet proces (eller diffusionsproces), hvor et datum (generelt et billede) gradvist støjes, og en omvendt proces (eller omvendt diffusionsproces), hvor støj omdannes tilbage til en prøve fra målfordelingen.

Prøveudtagningskædeovergangene i den fremadrettede proces kan indstilles til betingede Gaussians, når støjniveauet er tilstrækkeligt lavt. Kombination af denne kendsgerning med Markov-antagelsen fører til en simpel parameterisering af den fremadrettede proces:

Matematisk note

Vi har talt om at korrumpere dataene ved tilføje Gaussisk støj, men det kan i første omgang være uklart, hvor vi udfører denne tilføjelse. I overensstemmelse med ovenstående ligning udtager vi ved hvert trin i kæden blot en stikprøve fra en Gauss-fordeling, hvis middelværdi er den foregående værdi (dvs. billede) i kæden.

Disse to udsagn er ækvivalente. Det er

For at forstå hvorfor, vil vi bruge et lille misbrug af notation ved at hævde

Hvor den endelige implikation stammer fra den matematiske ækvivalens mellem en sum af stokastiske variable og foldningen af ​​deres fordelinger – se denne Wikipedia-side for mere information.

Med andre ord har vi vist, at at hævde fordelingen af ​​et tidstrin betinget af det foregående via middelværdien af ​​en Gauss-fordeling svarer til at hævde, at fordelingen af ​​et givet tidstrin er fordelingen af ​​det foregående med tilføjelse af Gauss-støj. Vi udelod skalarerne introduceret af variansskemaet og viste dette for én dimension for enkelhedens skyld, men et lignende bevis gælder for multivariate Gaussians.

Hvor β₁,…,β_T er en afvigelsesplan (enten indlært eller fast), som, hvis den opfører sig vel, sikrer det x_T er næsten en isotrop Gaussian til tilstrækkelig stor T.

Som tidligere nævnt kommer "magien" ved diffusionsmodeller i omvendt proces. Under træningen lærer modellen at vende denne diffusionsproces for at generere nye data. Startende med den rene Gauss-støj p(x_T):=N(x_T,0,I) modellen lærer fællesfordelingen p_θ(x_0:T) som

hvor de tidsafhængige parametre for de Gaussiske overgange læres. Bemærk især, at Markov-formuleringen hævder, at en given omvendt diffusionsovergangsfordeling kun afhænger af det foregående tidstrin (eller følgende tidstrin, afhængigt af hvordan du ser på det):

Kurser

En Diffusionsmodel trænes af finde de omvendte Markov-overgange, der maksimerer sandsynligheden for træningsdataene. I praksis består træning tilsvarende i at minimere variationens øvre grænse for den negative logsandsynlighed.

Vi søger at omskrive L_vlb i form af Kullback-Leibler (KL) Divergenser. KL-divergensen er et asymmetrisk statistisk afstandsmål for, hvor meget én sandsynlighedsfordeling P adskiller sig fra en referencefordeling Q. Vi er interesserede i at formulere L_vlb i form af KL-divergenser, fordi overgangsfordelingerne i vores Markov-kæde er gaussiske, og KL-divergensen mellem Gausserne har en lukket form.

Hvad er KL-divergensen?

Den matematiske form for KL-divergensen for kontinuerte fordelinger er

Nedenfor kan du se KL-divergensen for en varierende fordeling P (blå) fra en referencefordeling Q (rød). Den grønne kurve angiver funktionen inden for integralet i definitionen for KL-divergensen ovenfor, og det samlede areal under kurven repræsenterer værdien af ​​KL-divergensen af P fra Q på ethvert givet tidspunkt, en værdi, som også vises numerisk.

Støbning  i form af KL-divergenser

Som tidligere nævnt er det muligt [1] for at omskrive L_vlb næsten fuldstændigt med hensyn til KL-divergenser:

hvor

Afledningsdetaljer

Variationsgrænsen er lig med

Udskiftning af distributionerne med deres definitioner givet vores Markov-antagelse, får vi

Vi bruger logregler til at omdanne udtrykket til en sum af logfiler, og så trækker vi det første led ud

Ved at bruge Bayes' sætning og vores Markov-antagelse bliver dette udtryk

Vi opdeler derefter mellemleddet ved hjælp af logregler

Isolering af den anden periode, ser vi

Sætter dette tilbage i vores ligning for L_vlb, vi har

Ved hjælp af logregler omarrangerer vi

Dernæst bemærker vi følgende ækvivalens for KL-divergensen for alle to fordelinger:

Til sidst, når vi anvender denne ækvivalens på det tidligere udtryk, når vi frem til

Konditionering af den fremadrettede proces posterior på x₀ i L_t−1 resulterer i en håndterbar form, der fører til alle KL-divergenser er sammenligninger mellem gaussere. Dette betyder, at divergenserne kan beregnes nøjagtigt med udtryk i lukket form i stedet for med Monte Carlo estimater^[3].

Modelvalg

Med det matematiske grundlag for vores objektive funktion etableret, skal vi nu træffe flere valg med hensyn til, hvordan vores diffusionsmodel skal implementeres. For den fremadrettede proces er det eneste valg, der kræves, at definere variansplanen, hvis værdier generelt stiger under fremadgående proces.

Til den omvendte proces vælger vi meget den Gaussiske distributionsparameterisering/modelarkitektur(er). Bemærk høj grad af fleksibilitet at Diffusion Models har råd til – den kun krav til vores arkitektur er, at dens input og output har samme dimensionalitet.

Vi vil undersøge detaljerne om disse valg mere detaljeret nedenfor.

Forward Process og L_T

Som nævnt ovenfor, med hensyn til den fremadrettede proces, skal vi definere variansplanen. Især sætter vi dem til at være tidsafhængige konstanter, ignorerer det faktum, at de kan læres. For eksempel^[3], en lineær tidsplan fra β₁= 10^-4 til β_T=0.2 kan bruges, eller måske en geometrisk række.

Uanset de bestemte værdier, der er valgt, resulterer det faktum, at variansplanen er fast, i L_T bliver en konstant med hensyn til vores sæt af indlærelige parametre, hvilket giver os mulighed for at ignorere det, hvad træning angår.

Omvendt proces og L_1:T−1

Nu diskuterer vi de valg, der kræves for at definere den omvendte proces. Husk fra oven, at vi definerede de omvendte Markov-overgange som en Gauss:

Vi skal nu definere de funktionelle former for μ_θ eller Σ_θ. Mens der er mere komplicerede måder at parameterisere Σ på_θ^[5], vi sætter simpelthen

Det vil sige, vi antager, at den multivariate Gauss er et produkt af uafhængige Gausser med identisk varians, en variansværdi, som kan ændre sig med tiden. Vi sæt disse afvigelser til at svare til vores fremadrettede procesafvigelsesplan.

Givet denne nye formulering af Σ_θ, vi har

som giver os mulighed for at transformere

til

hvor det første led i forskellen er en lineær kombination af x_t og x₀ der afhænger af variansskemaet β_t. Den nøjagtige form for denne funktion er ikke relevant for vores formål, men den kan findes i [3].

Betydningen af ​​ovenstående andel er det den mest ligetil parameterisering af μ_θ forudsiger blot diffusionen posterior middelværdi. Det er vigtigt, at forfatterne af [3] fandt faktisk ud af, at træning μ_θ at forudsige støj komponent på et givet tidspunkt giver bedre resultater. Især lad

hvor

Dette fører til følgende alternative tabsfunktion, som forfatterne til [3] viste sig at føre til mere stabil træning og bedre resultater:

Forfatterne af [3] bemærk også forbindelserne mellem denne formulering af diffusionsmodeller til score-matchende generative modeller baseret på Langevin-dynamik. Det ser faktisk ud til, at diffusionsmodeller og scorebaserede modeller kan være to sider af samme mønt, beslægtet med den uafhængige og samtidige udvikling af bølgebaseret kvantemekanik og matrixbaseret kvantemekanik, der afslører to ækvivalente formuleringer af de samme fænomener^[2].

Netværksarkitektur

Mens vores forenklede tabsfunktion søger at træne en model ϵ_θ, har vi endnu ikke defineret arkitekturen af ​​denne model. Bemærk, at kun kravet til modellen er, at dens input- og outputdimensionalitet er identiske.

I betragtning af denne begrænsning er det måske ikke overraskende, at billeddiffusionsmodeller almindeligvis implementeres med U-Net-lignende arkitekturer.

Reverse Process Decoder og L₀

Vejen langs den omvendte proces består af mange transformationer under kontinuerlige betingede Gauss-fordelinger. I slutningen af ​​den omvendte proces skal du huske, at vi forsøger at producere en billede, som er sammensat af heltals pixelværdier. Derfor må vi udtænke en måde at opnå diskrete (log) sandsynligheder for hver mulig pixelværdi på tværs af alle pixels.

Måden dette gøres på er ved at sætte den sidste overgang i den omvendte diffusionskæde til an uafhængig diskret dekoder. For at bestemme sandsynligheden for et givet billede x₀ givet x₁, påtvinger vi først uafhængighed mellem datadimensionerne:

hvor D er dimensionaliteten af ​​dataene og det hævede skrift i angiver udtrækning af en koordinat. Målet er nu at bestemme, hvor sandsynligt hver heltalsværdi er for en given pixel given fordelingen på tværs af mulige værdier for den tilsvarende pixel i billedet med let støj på tidspunktet t=1:

hvor pixelfordelingerne for t=1 er afledt af den nedenstående multivariate Gaussian, hvis diagonale kovariansmatrix tillader os at opdele fordelingen i et produkt af univariate Gaussianer, en for hver dimension af dataene:

Vi antager, at billederne består af heltal i 0,1,...,255 (som standard RGB-billeder gør), som er blevet skaleret lineært til [−1,1]. Vi opdeler derefter den rigtige linje i små "buckets", hvor der for en given skaleret pixelværdi x, er spanden for det område [x−1/255, x+1/255]. Sandsynligheden for en pixelværdi x, givet den univariate Gauss-fordeling af den tilsvarende pixel i x₁, er areal under den univariate Gauss-fordeling inden for spanden centreret ved x.

Nedenfor kan du se arealet for hver af disse buckets med deres sandsynligheder for en middel-0 Gaussian, som i denne sammenhæng svarer til en fordeling med en gennemsnitlig pixelværdi på 255/2 (halv lysstyrke). Den røde kurve repræsenterer fordelingen af ​​en specifik pixel i t = 1 billede, og områderne giver sandsynligheden for den tilsvarende pixelværdi i t = 0 billede.

Givet en t = 0 pixelværdi for hver pixel, værdien af ​​p_θ(x₀|x₁) er simpelthen deres produkt. Denne proces er kortfattet indkapslet af følgende ligning:

hvor

,

Givet denne ligning for p_θ(x₀|x₁), kan vi beregne den endelige term af L_vlb som ikke er formuleret som en KL-divergens:

Slutmål

Som nævnt i sidste afsnit, forfatterne af [3] fandt, at forudsigelse af støjkomponenten i et billede på et givet tidspunkt gav de bedste resultater. I sidste ende bruger de følgende mål:

Trænings- og prøveudtagningsalgoritmerne for vores diffusionsmodel kan derfor kortfattet fanges i nedenstående figur:

Resumé af diffusionsmodelteori

I dette afsnit tog vi et detaljeret dyk ned i teorien om diffusionsmodeller. Det kan være nemt at blive fanget af matematiske detaljer, så vi noterer de vigtigste punkter i dette afsnit nedenfor for at holde os orienteret fra et fugleperspektiv:

1. Vores diffusionsmodel er parametriseret som en Markov kæde, hvilket betyder, at vores latente variable x₁,…,x_T afhænger kun af det foregående (eller følgende) tidstrin.
2.  overgangsfordelinger i Markov-kæden er Gauss, hvor den fremadrettede proces kræver en variansplan, og de omvendte procesparametre læres.
3. Diffusionsprocessen sikrer, at x_T is asymptotisk fordelt som en isotrop Gauss for tilstrækkeligt stort T.
4. I vores tilfælde afvigelsesplanen blev fastsat, men det kan også læres. For faste skemaer kan det give bedre resultater at følge en geometrisk progression end en lineær progression. I begge tilfælde er varianserne generelt stigende med tiden i serien (dvs. β_i<β_j for i
5. Diffusionsmodeller er meget fleksibel og give mulighed for enhver arkitektur, hvis input og output dimensionalitet er den samme, der skal bruges. Mange implementeringer bruger U-Net-lignende arkitekturer.
6.  træningsmål er at maksimere sandsynligheden for træningsdataene. Dette kommer til udtryk ved at indstille modelparametrene til minimere variationens øvre grænse for den negative logsandsynlighed for dataene.
7. Næsten alle led i den objektive funktion kan støbes som KL Afvigelser som et resultat af vores Markov-antagelse. Disse værdier blive holdbar at beregne givet, at vi bruger Gaussians, og derfor udelader behovet for at udføre Monte Carlo-tilnærmelse.
8. I sidste ende bruger en forenklet træningsmål at træne en funktion, som forudsiger støjkomponenten af ​​en given latent variabel, giver de bedste og mest stabile resultater.
9. A diskret dekoder bruges til at opnå log-sandsynligheder på tværs af pixelværdier som det sidste trin i den omvendte diffusionsproces.

Med denne oversigt på højt niveau af diffusionsmodeller i vores sind, lad os gå videre for at se, hvordan man bruger en diffusionsmodeller i PyTorch.

Diffusionsmodeller i PyTorch

Mens diffusionsmodeller endnu ikke er blevet demokratiseret i samme grad som andre ældre arkitekturer/tilgange i Machine Learning, er der stadig implementeringer tilgængelige til brug. Den nemmeste måde at bruge en diffusionsmodel i PyTorch er at bruge denoising-diffusion-pytorch pakke, som implementerer en billeddiffusionsmodel som den, der er diskuteret i denne artikel. For at installere pakken skal du blot skrive følgende kommando i terminalen:

pip install denoising_diffusion_pytorch

Minimalt eksempel

For at træne en model og generere billeder importerer vi først de nødvendige pakker:

import torch
from denoising_diffusion_pytorch import Unet, GaussianDiffusion

Dernæst definerer vi vores netværksarkitektur, i dette tilfælde et U-Net. Det dim parameter angiver antallet af feature maps før den første down-sampling, og dim_mults parameter giver multiplikander for denne værdi og efterfølgende nedsamplinger:

model = Unet(
 dim = 64,
 dim_mults = (1, 2, 4, 8)
)

Nu hvor vores netværksarkitektur er defineret, skal vi definere selve diffusionsmodellen. Vi videregiver U-Net-modellen, som vi netop har defineret sammen med flere parametre - størrelsen på billeder, der skal genereres, antallet af tidstrin i diffusionsprocessen og et valg mellem L1- og L2-normerne.

diffusion = GaussianDiffusion(
 model,
 image_size = 128,
 timesteps = 1000, # number of steps
 loss_type = 'l1' # L1 or L2
)

Nu hvor Diffusionsmodellen er defineret, er det tid til at træne. Vi genererer tilfældige data til at træne på, og træner derefter diffusionsmodellen på sædvanlig vis:

training_images = torch.randn(8, 3, 128, 128)
loss = diffusion(training_images)
loss.backward()

Når modellen er trænet, kan vi endelig generere billeder ved at bruge sample() metode af diffusion objekt. Her genererer vi 4 billeder, som kun er støj, da vores træningsdata var tilfældige:

sampled_images = diffusion.sample(batch_size = 4)

Træning i brugerdefinerede data

 denoising-diffusion-pytorch pakken giver dig også mulighed for at træne en diffusionsmodel på et specifikt datasæt. Du skal blot udskifte 'path/to/your/images' streng med datasættets biblioteksti i Trainer() objekt nedenfor, og ændre image_size til den passende værdi. Derefter skal du blot køre koden for at træne modellen og derefter prøve som før. Bemærk, at PyTorch skal være kompileret med CUDA aktiveret for at bruge Trainer klasse:

from denoising_diffusion_pytorch import Unet, GaussianDiffusion, Trainer
model = Unet(
 dim = 64,
 dim_mults = (1, 2, 4, 8)
).cuda()
diffusion = GaussianDiffusion(
 model,
 image_size = 128,
 timesteps = 1000, # number of steps
 loss_type = 'l1' # L1 or L2
).cuda()
trainer = Trainer(
 diffusion,
 'path/to/your/images',
 train_batch_size = 32,
 train_lr = 2e-5,
 train_num_steps = 700000, # total training steps
 gradient_accumulate_every = 2, # gradient accumulation steps
 ema_decay = 0.995, # exponential moving average decay
 amp = True # turn on mixed precision
)
trainer.train()

Nedenfor kan du se progressiv denoising fra multivariat Gaussisk støj til MNIST-cifre, der ligner omvendt diffusion:

Afsluttende ord

Diffusionsmodeller er en konceptuel enkel og elegant tilgang til problemet med at generere data. Deres state-of-the-art resultater kombineret med ikke-modstridende træning har drevet dem til store højder, og yderligere forbedringer kan forventes i de kommende år givet deres begyndende status. Især har diffusionsmodeller vist sig at være afgørende for ydeevnen af ​​banebrydende modeller som DALL-E2.

Referencer

[1] Dyb uovervåget læring ved hjælp af nonequilibrium termodynamik

[2] Generativ modellering ved at estimere gradienter af datafordelingen

[3] Denoising Diffusion Probabilistic Modeller

[4] Forbedrede teknikker til træning af scorebaserede generative modeller

[5] Forbedrede Denoising Diffusion Probabilistic Modeller

[6] Diffusionsmodeller slår GAN'er på billedsyntese

[7] GLIDE: Mod fotorealistisk billedgenerering og -redigering med tekststyrede diffusionsmodeller

[8] Hierarkisk tekst-betinget billedgenerering med CLIP-latenter

Nyder du denne artikel? Tilmeld dig flere AI-forskningsopdateringer.

Vi giver dig besked, når vi udgiver flere oversigtsartikler som denne.

Relaterede