Introduksjon til diffusjonsmodeller for maskinlæring

Denne artikkelen ble opprinnelig publisert på AssemblyAI og publisert på nytt til TOPBOTS med tillatelse fra forfatteren.

Diffusjonsmodeller er generative modeller som har vunnet betydelig popularitet de siste årene, og med god grunn. En håndfull banebrytende papirer utgitt på 2020-tallet alene har vist verden hva diffusjonsmodeller er i stand til, for eksempel å slå GAN-er^[6] på bildesyntese. Senest vil utøvere ha sett diffusjonsmodeller brukt i DALL-E2, OpenAIs bildegenereringsmodell utgitt forrige måned.

Gitt den nylige suksessbølgen til Diffusion Models, er mange maskinlæringsutøvere sikkert interessert i deres indre funksjoner. I denne artikkelen vil vi undersøke teoretisk grunnlag for diffusjonsmodeller, og demonstrer deretter hvordan du genererer bilder med en Diffusjonsmodell i PyTorch. La oss dykke inn!

Hvis dette inngående pedagogiske innholdet er nyttig for deg, abonner på vår AI-e-postliste å bli varslet når vi slipper nytt materiale. 

Diffusjonsmodeller – Introduksjon

Diffusjonsmodeller er generative modeller, noe som betyr at de brukes til å generere data som ligner på dataene de er trent på. I utgangspunktet fungerer diffusjonsmodeller etter ødelegge treningsdata gjennom påfølgende tillegg av Gaussisk støy, og deretter lære å komme seg dataene av reversering denne støyprosessen. Etter trening kan vi bruke diffusjonsmodellen til å generere data ganske enkelt sende tilfeldig samplet støy gjennom den lærte denoising-prosessen.

Mer spesifikt er en diffusjonsmodell en latent variabel modell som kartlegger til det latente rommet ved hjelp av en fast Markov-kjede. Denne kjeden legger gradvis til støy til dataene for å oppnå omtrentlig bakre q(x_1:T|x₀), hvor x₁,…, X_T er de latente variablene med samme dimensjonalitet som x₀. I figuren nedenfor ser vi en slik Markov-kjede manifestert for bildedata.

Til syvende og sist blir bildet asymptotisk transformert til ren Gaussisk støy. De mål å trene en diffusjonsmodell er å lære reversere prosess – dvs. opplæring s_θ(x_t−1|x_t). Ved å gå bakover langs denne kjeden kan vi generere nye data.

Fordeler med diffusjonsmodeller

Som nevnt ovenfor har forskningen på diffusjonsmodeller eksplodert de siste årene. Inspirert av termodynamikk uten likevekt^[1], Diffusion Models produserer for tiden Toppmoderne bildekvalitet, eksempler på dette kan ses nedenfor:

Utover banebrytende bildekvalitet, kommer diffusjonsmodeller med en rekke andre fordeler, inkludert krever ikke motstanderutdanning. Vanskelighetene med motstandertrening er godt dokumentert; og i tilfeller der det eksisterer ikke-motstridende alternativer med sammenlignbar ytelse og treningseffektivitet, er det vanligvis best å bruke dem. Når det gjelder treningseffektivitet, har diffusjonsmodeller også de ekstra fordelene skalerbarhet og parallelliserbarhet.

Mens diffusjonsmodeller nesten ser ut til å produsere resultater ut av løse luften, er det mange nøye og interessante matematiske valg og detaljer som danner grunnlaget for disse resultatene, og beste praksis utvikler seg fortsatt i litteraturen. La oss ta en titt på den matematiske teorien som ligger til grunn for diffusjonsmodeller mer detaljert nå.

Diffusjonsmodeller – et dypdykk

Som nevnt ovenfor består en diffusjonsmodell av en videre prosess (eller diffusjonsprosess), der et datum (vanligvis et bilde) støyes gradvis, og en omvendt prosess (eller omvendt diffusjonsprosess), der støy transformeres tilbake til en prøve fra målfordelingen.

Samplingskjedeovergangene i foroverprosessen kan settes til betingede Gaussians når støynivået er tilstrekkelig lavt. Å kombinere dette faktum med Markov-antagelsen fører til en enkel parameterisering av fremgangsprosessen:

Matematisk notat

Vi har snakket om å korrupte dataene ved legge Gaussisk støy, men det kan først være uklart hvor vi utfører dette tillegget. I henhold til ligningen ovenfor, ved hvert trinn i kjeden prøver vi ganske enkelt fra en gaussisk fordeling hvis gjennomsnitt er den forrige verdien (dvs. bildet) i kjeden.

Disse to utsagnene er likeverdige. Det er

For å forstå hvorfor, vil vi bruke et lite misbruk av notasjon ved å hevde

Der den endelige implikasjonen stammer fra den matematiske ekvivalensen mellom en sum av tilfeldige variabler og konvolusjonen av deres fordelinger – se denne Wikipedia-siden for mer informasjon.

Med andre ord, vi har vist at å hevde fordelingen av et tidstrinn betinget av det forrige via gjennomsnittet av en Gaussisk fordeling er ekvivalent med å hevde at fordelingen av et gitt tidstrinn er den for det forrige med tillegg av Gaussisk støy. Vi utelot skalarene introdusert av variansskjemaet og viste dette for én dimensjon for enkelhets skyld, men et lignende bevis gjelder for multivariate gaussere.

Hvor β₁,…,β_T er en variasjonsplan (enten lært eller fast) som, hvis den oppfører seg vel, sikrer det x_T er nesten en isotrop gaussisk for tilstrekkelig stor T.

Som nevnt tidligere kommer "magien" til diffusjonsmodeller i omvendt prosess. Under trening lærer modellen å reversere denne diffusjonsprosessen for å generere nye data. Starter med den rene Gauss-støyen p(x_T):=N(x_T,0,I) modellen lærer fellesfordelingen s_θ(x_0:T) som

hvor de tidsavhengige parameterne til de gaussiske overgangene læres. Merk spesielt at Markov-formuleringen hevder at en gitt omvendt diffusjonsovergangsfordeling bare avhenger av forrige tidstrinn (eller følgende tidstrinn, avhengig av hvordan du ser på det):

Kurs

En diffusjonsmodell er opplært av finne de omvendte Markov-overgangene som maksimerer sannsynligheten for treningsdataene. I praksis består trening tilsvarende av å minimere den øvre variasjonsgrensen på sannsynligheten for negativ logg.

Vi prøver å omskrive L_etc i form av Kullback-Leibler (KL) Divergenser. KL-divergensen er et asymmetrisk statistisk avstandsmål for hvor mye en sannsynlighetsfordeling P skiller seg fra en referansefordeling Q. Vi er interessert i å formulere L_etc når det gjelder KL-divergenser fordi overgangsfordelingene i vår Markov-kjede er gaussiske, og KL-divergensen mellom gaussere har en lukket form.

Hva er KL-divergensen?

Den matematiske formen for KL-divergensen for kontinuerlige fordelinger er

Nedenfor kan du se KL-divergensen til en varierende fordeling P (blå) fra en referansefordeling Q (rød). Den grønne kurven angir funksjonen innenfor integralet i definisjonen for KL-divergensen ovenfor, og det totale arealet under kurven representerer verdien av KL-divergensen på P fra Q til enhver tid, en verdi som også vises numerisk.

Avstøpning  i form av KL divergenser

Som nevnt tidligere, er det mulig [1] for å omskrive L_etc nesten fullstendig når det gjelder KL-divergenser:

hvor

Avledningsdetaljer

Variasjonsgrensen er lik

Å erstatte distribusjonene med deres definisjoner gitt vår Markov-antakelse, får vi

Vi bruker logregler for å transformere uttrykket til en sum av logger, og så trekker vi ut det første leddet

Ved å bruke Bayes' teorem og vår Markov-antakelse blir dette uttrykket

Vi deler deretter opp mellomleddet ved å bruke loggregler

Å isolere den andre perioden, ser vi

Plugger dette tilbake inn i ligningen vår for L_etc, vi har

Ved å bruke loggregler omorganiserer vi

Deretter bemerker vi følgende ekvivalens for KL-divergensen for alle to fordelinger:

Til slutt, ved å bruke denne ekvivalensen på det forrige uttrykket, kommer vi til

Kondisjonere den fremre prosessen posterior på x₀ i L_t−1 resulterer i en håndterbar form som fører til alle KL-divergenser er sammenligninger mellom gaussere. Dette betyr at divergensene kan beregnes nøyaktig med lukkede uttrykk i stedet for med Monte Carlo-estimat^[3].

Modellvalg

Med det matematiske grunnlaget for vår objektive funksjon etablert, må vi nå ta flere valg angående hvordan vår diffusjonsmodell skal implementeres. For fremdriftsprosessen er det eneste valget som kreves å definere variansplanen, hvis verdier generelt øker under fremdriftsprosessen.

For den omvendte prosessen velger vi mye Gaussisk distribusjonsparameterisering/modellarkitektur(er). Legg merke til høy grad av fleksibilitet som Diffusion Models har råd til – den bare kravet til arkitekturen vår er at input og output har samme dimensjonalitet.

Vi vil utforske detaljene i disse valgene mer detaljert nedenfor.

Videresend prosess og L_T

Som nevnt ovenfor, angående fremdriftsprosessen, må vi definere avviksplanen. Spesielt setter vi dem til å være tidsavhengige konstanter, ignorerer det faktum at de kan læres. For eksempel^[3], en lineær tidsplan fra β₁= 10^-4 til β_T=0.2 kan brukes, eller kanskje en geometrisk serie.

Uavhengig av de spesielle verdiene som er valgt, resulterer det faktum at variasjonsplanen er fast i L_T blir en konstant med hensyn til vårt sett med lærbare parametere, slik at vi kan ignorere det når det gjelder trening.

Omvendt prosess og L_1:T−1

Nå diskuterer vi valgene som kreves for å definere den omvendte prosessen. Husk at vi ovenfra definerte de omvendte Markov-overgangene som en gaussisk:

Vi må nå definere de funksjonelle formene til μ_θ eller Σ_θ. Mens det er mer kompliserte måter å parameterisere Σ på_θ^[5], vi bare setter

Det vil si at vi antar at den multivariate gaussiske er et produkt av uavhengige gaussere med identisk varians, en variansverdi som kan endre seg med tiden. Vi angi disse avvikene til å være ekvivalente med vår avviksplan for fremadgående prosess.

Gitt denne nye formuleringen av Σ_θ, vi har

som lar oss transformere

til

hvor det første leddet i forskjellen er en lineær kombinasjon av x_t og x₀ som avhenger av variansskjemaet β_t. Den nøyaktige formen for denne funksjonen er ikke relevant for våre formål, men den kan finnes i [3].

Betydningen av andelen ovenfor er det den mest enkle parameteriseringen av μ_θ forutsier ganske enkelt diffusjonen bakre middelverdi. Det er viktig at forfatterne av [3] fant faktisk ut at trening μ_θ å forutsi støy komponent på et gitt tidspunkt gir bedre resultater. Spesielt la

hvor

Dette fører til følgende alternative tapsfunksjon, som forfatterne av [3] funnet å føre til mer stabil trening og bedre resultater:

Forfatterne av [3] legg også merke til sammenhenger mellom denne formuleringen av diffusjonsmodeller til score-matchende generative modeller basert på Langevin-dynamikk. Det ser faktisk ut til at diffusjonsmodeller og poengbaserte modeller kan være to sider av samme sak, beslektet med den uavhengige og samtidige utviklingen av bølgebasert kvantemekanikk og matrisebasert kvantemekanikk som avslører to ekvivalente formuleringer av de samme fenomenene^[2].

Nettverksarkitektur

Mens vår forenklede tapsfunksjon søker å trene en modell ϵ_θ, har vi ennå ikke definert arkitekturen til denne modellen. Merk at bare kravet til modellen er at dens inngangs- og utgangsdimensjonalitet er identiske.

Gitt denne begrensningen, er det kanskje ikke overraskende at bildediffusjonsmodeller ofte implementeres med U-Net-lignende arkitekturer.

Omvendt prosessdekoder og L₀

Banen langs den omvendte prosessen består av mange transformasjoner under kontinuerlige betingede Gauss-fordelinger. På slutten av den omvendte prosessen, husk at vi prøver å produsere en bilde, som er sammensatt av heltalls pikselverdier. Derfor må vi finne en måte å oppnå diskrete (logg) sannsynligheter for hver mulig pikselverdi på tvers av alle piksler.

Måten dette gjøres på er ved å sette den siste overgangen i den omvendte diffusjonskjeden til an uavhengig diskret dekoder. For å bestemme sannsynligheten for et gitt bilde x₀ gitt x₁, pålegger vi først uavhengighet mellom datadimensjonene:

hvor D er dimensjonaliteten til dataene og hevet skrift i indikerer uttrekking av en koordinat. Målet nå er å bestemme hvor sannsynlig hver heltallsverdi er for en gitt piksel gitt fordelingen over mulige verdier for den korresponderende pikselen i det svakt støyede bildet på tidspunktet t=1:

hvor pikselfordelingene for t=1 er utledet fra den under multivariate Gaussian hvis diagonale kovariansmatrise tillater oss å dele fordelingen inn i et produkt av univariate Gaussians, en for hver dimensjon av dataene:

Vi antar at bildene består av heltall i 0,1,…,255 (som standard RGB-bilder gjør) som er skalert lineært til [−1,1]. Vi bryter deretter ned den virkelige linjen i små "bøtter", hvor det for en gitt skalert pikselverdi x, er bøtten for det området [x−1/255, x+1/255]. Sannsynligheten for en pikselverdi x, gitt den univariate gaussiske fordelingen av den tilsvarende pikselen i x₁, er den område under den univariate Gauss-fordelingen innenfor bøtta sentrert ved x.

Nedenfor kan du se arealet for hver av disse bøttene med deres sannsynligheter for en gjennomsnittlig-0 Gauss som i denne sammenheng tilsvarer en fordeling med en gjennomsnittlig pikselverdi på 255/2 (halv lysstyrke). Den røde kurven representerer fordelingen av en spesifikk piksel i t = 1 bildet, og områdene gir sannsynligheten for den tilsvarende pikselverdien i t = 0 bilde.

Gitt en t = 0 pikselverdi for hver piksel, verdien av p_θ(x₀|x₁) er ganske enkelt deres produkt. Denne prosessen er kortfattet innkapslet av følgende ligning:

hvor

og

Gitt denne ligningen for s_θ(x₀|x₁), kan vi beregne sluttleddet til L_etc som ikke er formulert som en KL-divergens:

Endelig mål

Som nevnt i den siste delen, forfatterne av [3] fant at å forutsi støykomponenten til et bilde ved et gitt tidstrinn ga de beste resultatene. Til syvende og sist bruker de følgende mål:

Trenings- og prøvetakingsalgoritmene for vår diffusjonsmodell kan derfor fanges opp i figuren nedenfor:

Sammendrag av diffusjonsmodellteori

I denne delen tok vi et detaljert dykk inn i teorien om diffusjonsmodeller. Det kan være lett å bli fanget opp i matematiske detaljer, så vi noterer oss de viktigste punktene i denne delen nedenfor for å holde oss orientert fra et fugleperspektiv:

1. Vår diffusjonsmodell er parameterisert som en Markov-kjeden, som betyr at våre latente variabler x₁,…, X_T avhenger bare av forrige (eller følgende) tidstrinn.
2. De overgangsfordelinger i Markov-kjeden er Gaussian, hvor foroverprosessen krever en variasjonsplan, og de omvendte prosessparametrene læres.
3. Diffusjonsprosessen sørger for at x_T is asymptotisk fordelt som en isotrop Gaussian for tilstrekkelig stor T.
4. I vårt tilfelle er avviksplanen ble fastsatt, men det kan også læres. For faste tidsplaner kan det å følge en geometrisk progresjon gi bedre resultater enn en lineær progresjon. I begge tilfeller øker variansene generelt med tiden i serien (dvs. β_i<β_j for jeg
5. Diffusjonsmodeller er svært fleksibel og gi rom for noen arkitektur hvis inngangs- og utgangsdimensjonalitet er den samme som skal brukes. Mange implementeringer bruker U-Net-lignende arkitekturer.
6. De treningsmål er å maksimere sannsynligheten for treningsdataene. Dette manifesteres som å justere modellparametrene til minimere variasjonens øvre grense for den negative loggsannsynligheten for dataene.
7. Nesten alle ledd i objektivfunksjonen kan støpes som KL Divergenser som et resultat av vår Markov-antagelse. Disse verdiene bli holdbar å beregne gitt at vi bruker Gaussians, og derfor utelater behovet for å utføre Monte Carlo-tilnærming.
8. Til slutt bruker du en forenklet opplæringsmål å trene en funksjon som forutsier støykomponenten til en gitt latent variabel gir de beste og mest stabile resultatene.
9. A diskret dekoder brukes til å oppnå loggsannsynligheter på tvers av pikselverdier som siste trinn i den omvendte diffusjonsprosessen.

Med denne høynivåoversikten over diffusjonsmodeller i tankene våre, la oss gå videre for å se hvordan du bruker en diffusjonsmodeller i PyTorch.

Diffusjonsmodeller i PyTorch

Mens diffusjonsmodeller ennå ikke er blitt demokratisert i samme grad som andre eldre arkitekturer/tilnærminger innen maskinlæring, er det fortsatt tilgjengelige implementeringer for bruk. Den enkleste måten å bruke en diffusjonsmodell i PyTorch er å bruke denoising-diffusion-pytorch pakke, som implementerer en bildediffusjonsmodell som den som er omtalt i denne artikkelen. For å installere pakken, skriv ganske enkelt inn følgende kommando i terminalen:

pip install denoising_diffusion_pytorch

Minimalt eksempel

For å trene en modell og generere bilder, importerer vi først de nødvendige pakkene:

import torch
from denoising_diffusion_pytorch import Unet, GaussianDiffusion

Deretter definerer vi nettverksarkitekturen vår, i dette tilfellet et U-Net. De dim parameteren spesifiserer antall funksjonskart før den første nedsamplingen, og dim_mults parameter gir multiplikander for denne verdien og påfølgende ned-samplinger:

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

Nå som nettverksarkitekturen vår er definert, må vi definere selve diffusjonsmodellen. Vi sender inn U-Net-modellen som vi nettopp definerte sammen med flere parametere – størrelsen på bildene som skal genereres, antall tidstrinn i diffusjonsprosessen og et valg mellom L1- og L2-normene.

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

Nå som diffusjonsmodellen er definert, er det på tide å trene. Vi genererer tilfeldige data å trene på, og trener deretter diffusjonsmodellen på vanlig måte:

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

Når modellen er opplært, kan vi endelig generere bilder ved å bruke sample() metode av diffusion gjenstand. Her genererer vi 4 bilder, som kun er støy gitt at treningsdataene våre var tilfeldige:

sampled_images = diffusion.sample(batch_size = 4)

Opplæring i tilpassede data

De denoising-diffusion-pytorch pakken lar deg også trene en diffusjonsmodell på et spesifikt datasett. Bare bytt ut 'path/to/your/images' streng med datasettkatalogbanen i Trainer() objekt nedenfor, og endre image_size til riktig verdi. Deretter kjører du bare koden for å trene modellen, og prøver deretter som før. Merk at PyTorch må være kompilert med CUDA aktivert for å bruke 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 progressive denoising fra multivariat Gaussisk støy til MNIST-sifre som ligner på omvendt diffusjon:

Avsluttende ord

Diffusjonsmodeller er en konseptuelt enkel og elegant tilnærming til problemet med å generere data. Deres toppmoderne resultater kombinert med ikke-motstridende trening har drevet dem til store høyder, og ytterligere forbedringer kan forventes i de kommende årene gitt deres begynnende status. Spesielt har diffusjonsmodeller vist seg å være avgjørende for ytelsen til banebrytende modeller som DALL-E2.

Referanser

[1] Dyp uovervåket læring ved bruk av nonequilibrium termodynamikk

[2] Generativ modellering ved å estimere gradienter av datadistribusjonen

[3] Denoising Diffusion Probabilistic Models

[4] Forbedrede teknikker for opplæring av resultatbaserte generative modeller

[5] Forbedrede Denoising Diffusion Probabilistic Modeller

[6] Diffusjonsmodeller slår GAN-er på bildesyntese

[7] GLIDE: Mot fotorealistisk bildegenerering og redigering med tekststyrte spredningsmodeller

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

Liker du denne artikkelen? Registrer deg for flere AI-forskningsoppdateringer.

Vi gir beskjed når vi gir ut flere sammendragsartikler som denne.

I slekt