Introduktion till diffusionsmodeller för maskininlärning

Den här artikeln publicerades ursprungligen den AssemblyAI och återpubliceras till TOPBOTS med tillstånd från författaren.

Diffusionsmodeller är generativa modeller som har vunnit betydande popularitet under de senaste åren, och av goda skäl. En handfull framstående tidningar släpptes på 2020-talet ensam har visat världen vad Diffusionsmodeller kan, som att slå GAN^[6] på bildsyntes. Senast kommer utövare att ha sett Diffusionsmodeller användas i DALL-E2, OpenAI:s bildgenereringsmodell släpptes förra månaden.

Med tanke på den senaste framgångsvågen av Diffusion Models är många maskininlärningsutövare säkerligen intresserade av deras inre funktioner. I den här artikeln kommer vi att undersöka teoretiska grunder för diffusionsmodeller, och sedan demonstrera hur man skapar bilder med en Diffusionsmodell i PyTorch. Låt oss dyka in!

Om detta fördjupade utbildningsinnehåll är användbart för dig, prenumerera på vår AI-e-postlista att bli varnade när vi släpper nytt material. 

Diffusionsmodeller – Introduktion

Diffusionsmodeller är generativ modeller, vilket innebär att de används för att generera data som liknar den data som de tränas på. I grund och botten fungerar diffusionsmodeller efter förstöra träningsdata genom successiv tillägg av Gaussiskt brus, och sedan lära sig att återhämta sig uppgifterna av backning denna brusprocess. Efter träning kan vi använda diffusionsmodellen för att generera data helt enkelt passera slumpmässigt samplade brus genom den inlärda brusreduceringsprocessen.

Mer specifikt är en diffusionsmodell en latent variabel modell som mappar till det latenta utrymmet med hjälp av en fixerad Markov-kedja. Denna kedja lägger gradvis till brus till data för att erhålla den ungefärliga bakre q(x_1:T|x₀), där x₁, ..., x_T är de latenta variablerna med samma dimensionalitet som x₀. I figuren nedan ser vi en sådan Markov-kedja manifesterad för bilddata.

I slutändan omvandlas bilden asymptotiskt till rent Gaussiskt brus. De Målet att träna en diffusionsmodell är att lära sig vända process – dvs utbildning sid_θ(x_t−1|x_t). Genom att gå bakåt längs denna kedja kan vi generera ny data.

Fördelar med diffusionsmodeller

Som nämnts ovan har forskningen om diffusionsmodeller exploderat de senaste åren. Inspirerad av termodynamik som inte är jämvikt^[1], Diffusion Models producerar för närvarande Toppmodern bildkvalitet, exempel på det kan ses nedan:

Utöver den senaste bildkvaliteten har diffusionsmodellerna en mängd andra fördelar, inklusive inte kräver motståndskraftig utbildning. Svårigheterna med kontradiktorisk träning är väldokumenterade; och i de fall där icke-motstridiga alternativ finns med jämförbar prestanda och träningseffektivitet, är det vanligtvis bäst att använda dem. När det gäller träningseffektivitet har diffusionsmodeller också de extra fördelarna med skalbarhet och parallelliserbarhet.

Även om diffusionsmodeller nästan verkar producera resultat ur tomma intet, finns det många noggranna och intressanta matematiska val och detaljer som ligger till grund för dessa resultat, och bästa praxis utvecklas fortfarande i litteraturen. Låt oss ta en titt på den matematiska teorin som ligger till grund för diffusionsmodeller mer i detalj nu.

Diffusionsmodeller – En djupdykning

Som nämnts ovan består en diffusionsmodell av en framåt process (eller diffusionsprocess), där ett datum (vanligtvis en bild) progressivt brusas, och en omvänd process (eller omvänd diffusionsprocess), där brus omvandlas tillbaka till ett sampel från målfördelningen.

Samplingskedjeövergångarna i den framåtriktade processen kan ställas in på villkorliga Gaussians när ljudnivån är tillräckligt låg. Att kombinera detta faktum med Markov-antagandet leder till en enkel parameterisering av framåtprocessen:

Matematisk anmärkning

Vi har pratat om att korrumpera data genom tillsats Gaussiskt brus, men det kan till en början vara oklart var vi utför detta tillägg. I enlighet med ovanstående ekvation, vid varje steg i kedjan samplar vi helt enkelt från en Gauss-fördelning vars medelvärde är det föregående värdet (dvs. bilden) i kedjan.

Dessa två uttalanden är likvärdiga. Det är

För att förstå varför, kommer vi att använda ett litet missbruk av notation genom att hävda

Där den slutliga implikationen härrör från den matematiska ekvivalensen mellan summan av slumpvariabler och faltningen av deras fördelningar – se denna Wikipedia-sida för mer information.

Med andra ord har vi visat att att hävda fördelningen av ett tidssteg betingat av det föregående via medelvärdet av en Gaussfördelning är ekvivalent med att hävda att fördelningen av ett givet tidssteg är den för det föregående med tillägg av Gaussiskt brus. Vi utelämnade skalärerna som introducerades av variansschemat och visade detta för en dimension för enkelhetens skull, men ett liknande bevis gäller för multivariata Gaussians.

Där β₁,…,β_T är ett variansschema (antingen inlärt eller fast) som, om det sköter sig väl, säkerställer det x_T är nästan en isotrop Gauss för tillräckligt stort T.

Som nämnts tidigare kommer diffusionsmodellernas "magi" i omvänd process. Under utbildningen lär sig modellen att vända denna spridningsprocess för att generera ny data. Börjar med det rena Gaussiska bruset p(x_T):=N(x_T,0,I) modellen lär sig den gemensamma fördelningen sid_θ(x_0:T) som

där de tidsberoende parametrarna för de Gaussiska övergångarna lärs in. Observera särskilt att Markov-formuleringen hävdar att en given omvänd diffusionsövergångsfördelning endast beror på föregående tidssteg (eller följande tidssteg, beroende på hur du ser på det):

Utbildning

En diffusionsmodell utbildas av hitta de omvända Markov-övergångarna som maximerar sannolikheten för träningsdata. I praktiken består träning likvärdigt av att minimera variationens övre gräns för den negativa logsannolikheten.

Vi försöker skriva om L_VLB i termer av Kullback-Leibler (KL) Avvikelser. KL-divergensen är ett asymmetriskt statistiskt avståndsmått på hur mycket en sannolikhetsfördelning P skiljer sig från en referensfördelning Q. Vi är intresserade av att formulera L_VLB i termer av KL-divergenser eftersom övergångsfördelningarna i vår Markov-kedja är gaussiska, och KL-divergensen mellan Gausser har en sluten form.

Vad är KL-divergensen?

Den matematiska formen av KL-divergensen för kontinuerliga fördelningar är

Nedan kan du se KL-divergensen för en varierande fördelning P (blå) från en referensfördelning Q (röd). Den gröna kurvan anger funktionen inom integralen i definitionen för KL-divergensen ovan, och den totala arean under kurvan representerar värdet av KL-divergensen på P från Q vid varje givet ögonblick, ett värde som också visas numeriskt.

Gjutning  i termer av KL-avvikelser

Som tidigare nämnts är det möjligt [1] för att skriva om L_VLB nästan helt när det gäller KL-divergenser:

var

Härledningsdetaljer

Variationsgränsen är lika med

Om vi ​​ersätter fördelningarna med deras definitioner med tanke på vårt Markov-antagande, får vi

Vi använder logregler för att omvandla uttrycket till en summa av loggar, och sedan drar vi ut den första termen

Med hjälp av Bayes sats och vårt Markov-antagande blir detta uttryck

Vi delar sedan upp mellantermen med hjälp av loggregler

Att isolera den andra mandatperioden ser vi

Att koppla tillbaka detta till vår ekvation för L_VLB, vi har

Med hjälp av loggregler arrangerar vi om

Därefter noterar vi följande ekvivalens för KL-divergensen för alla två distributioner:

Till sist, genom att tillämpa denna ekvivalens på det tidigare uttrycket, kommer vi fram till

Konditionering av den framåtriktade processen posterior på x₀ i L_t−1 resulterar i en löslig form som leder till alla KL-divergenser är jämförelser mellan Gausser. Detta innebär att avvikelserna kan beräknas exakt med slutna uttryck snarare än med Monte Carlo uppskattningar^[3].

Modellval

Med den matematiska grunden för vår målfunktion etablerad behöver vi nu göra flera val angående hur vår diffusionsmodell ska implementeras. För den framåtriktade processen är det enda valet som krävs att definiera variansschemat, vars värden i allmänhet ökar under framåtprocessen.

För den omvända processen väljer vi mycket den Gaussiska fördelningsparameteriseringen/modellarkitekturen. Notera hög grad av flexibilitet som Diffusion Models har råd med – den endast kravet på vår arkitektur är att dess input och output har samma dimensionalitet.

Vi kommer att undersöka detaljerna om dessa val mer detaljerat nedan.

Forward Process och L_T

Som nämnts ovan, när det gäller den framåtriktade processen, måste vi definiera variansschemat. Speciellt har vi satt dem att vara tidsberoende konstanter, ignorerar det faktum att de kan läras. Till exempel^[3], ett linjärt schema från β₁= 10^-4 till β_T=0.2 kan användas, eller kanske en geometrisk serie.

Oavsett vilka särskilda värden som valts, resulterar det faktum att variansschemat är fast i L_T blir en konstant med avseende på vår uppsättning inlärbara parametrar, vilket gör att vi kan ignorera det när det gäller träning.

Omvänd process och L_1:T−1

Nu diskuterar vi de val som krävs för att definiera den omvända processen. Minns från ovan definierade vi de omvända Markov-övergångarna som en Gauss:

Vi måste nu definiera de funktionella formerna av μ_θ eller Σ_θ. Även om det finns mer komplicerade sätt att parametrisera Σ_θ^[5], vi ställer helt enkelt

Det vill säga, vi antar att den multivariata Gaussian är en produkt av oberoende Gaussianer med identisk varians, ett variansvärde som kan förändras med tiden. Vi ställ in dessa avvikelser så att de motsvarar vårt avvikelseschema för framåtriktade processer.

Med tanke på denna nya formulering av Σ_θ, vi har

som gör att vi kan transformera

till

där den första termen i skillnaden är en linjär kombination av x_t och x₀ som beror på variansschemat β_t. Den exakta formen av denna funktion är inte relevant för våra syften, men den finns i [3].

Betydelsen av ovanstående andel är att den mest enkla parametriseringen av μ_θ förutsäger helt enkelt diffusionens posteriora medelvärde. Viktigt är att författarna till [3] fann faktiskt att träning μ_θ att förutsäga brus komponent vid ett givet tidssteg ger bättre resultat. I synnerhet låt

var

Detta leder till följande alternativa förlustfunktion, som författarna till [3] visat sig leda till mer stabil träning och bättre resultat:

Författarna till [3] notera också kopplingar av denna formulering av diffusionsmodeller till poängmatchande generativa modeller baserade på Langevin-dynamik. Det verkar faktiskt som att diffusionsmodeller och poängbaserade modeller kan vara två sidor av samma mynt, besläktad med den oberoende och samtidiga utvecklingen av vågbaserad kvantmekanik och matrisbaserad kvantmekanik som avslöjar två ekvivalenta formuleringar av samma fenomen^[2].

Nätverksarkitektur

Medan vår förenklade förlustfunktion försöker träna en modell ϵ_θ, vi har ännu inte definierat arkitekturen för denna modell. Observera att endast kravet på modellen är att dess ingångs- och utmatningsdimensionalitet är identiska.

Med tanke på denna begränsning är det kanske inte förvånande att bilddiffusionsmodeller vanligtvis implementeras med U-Net-liknande arkitekturer.

Reverse Process Decoder och L₀

Vägen längs den omvända processen består av många transformationer under kontinuerliga villkorade Gaussfördelningar. I slutet av den omvända processen, kom ihåg att vi försöker producera en bild, som är sammansatt av heltalspixelvärden. Därför måste vi hitta ett sätt att få diskreta (logga) sannolikheter för varje möjligt pixelvärde över alla pixlar.

Det sätt som detta görs är genom att sätta den sista övergången i den omvända diffusionskedjan till an oberoende diskret dekoder. För att bestämma sannolikheten för en given bild x₀ givet x₁, inför vi först oberoende mellan datadimensionerna:

där D är dimensionaliteten hos datan och det upphöjda i indikerar extraktion av en koordinat. Målet är nu att bestämma hur sannolikt varje heltalsvärde är för en given pixel ges fördelningen över möjliga värden för motsvarande pixel i den något brusade bilden vid tidpunkten t=1:

där pixelfördelningarna för t=1 härleds från nedanstående multivariata Gaussian vars diagonala kovariansmatris tillåter oss att dela upp fördelningen i en produkt av univariata Gaussianer, en för varje dimension av data:

Vi antar att bilderna består av heltal i 0,1,...,255 (som standard RGB-bilder gör) som har skalats linjärt till [−1,1]. Vi bryter sedan ner den verkliga linjen i små "hinkar", där, för ett givet skalat pixelvärde x, hinken för det området är [x−1/255, x+1/255]. Sannolikheten för ett pixelvärde x, givet den univariata Gaussiska fördelningen av motsvarande pixel i x₁, är area under den univariata Gauss-fördelningen inom hinken centrerad vid x.

Nedan kan du se arean för var och en av dessa hinkar med deras sannolikheter för en medelvärde-0 Gauss som i detta sammanhang motsvarar en fördelning med ett genomsnittligt pixelvärde på 255/2 (halv ljusstyrka). Den röda kurvan representerar fördelningen av en specifik pixel i t = 1 bild, och områdena ger sannolikheten för motsvarande pixelvärde i t = 0 bild.

Med tanke på t = 0 pixelvärde för varje pixel, värdet av p_θ(x₀|x₁) är helt enkelt deras produkt. Denna process är kortfattat inkapslad av följande ekvation:

var

och

Givet denna ekvation för sid_θ(x₀|x₁), kan vi beräkna den slutliga termen för L_VLB som inte är formulerad som en KL-divergens:

Slutmål

Som nämndes i det sista avsnittet, författarna till [3] fann att förutsägelse av bruskomponenten i en bild vid ett givet tidssteg gav de bästa resultaten. I slutändan använder de följande mål:

Tränings- och samplingsalgoritmerna för vår diffusionsmodell kan därför fångas kortfattat i bilden nedan:

Sammanfattning av diffusionsmodellteori

I det här avsnittet tog vi en detaljerad dykning i teorin om diffusionsmodeller. Det kan vara lätt att fastna i matematiska detaljer, så vi noterar de viktigaste punkterna i detta avsnitt nedan för att hålla oss orienterade ur ett fågelperspektiv:

1. Vår diffusionsmodell är parametriserad som en Markov kedja, vilket betyder att våra latenta variabler x₁, ..., x_T beror endast på föregående (eller följande) tidssteg.
2. Smakämnen övergångsfördelningar i Markovkedjan finns Gauss, där framåtprocessen kräver ett variansschema, och de omvända processparametrarna lärs in.
3. Diffusionsprocessen säkerställer att x_T is asymptotiskt fördelad som en isotrop Gauss för tillräckligt stort T.
4. I vårt fall är variansschemat var fixat, men det går att lära sig också. För fasta scheman kan en geometrisk progression ge bättre resultat än en linjär progression. I båda fallen ökar varianserna i allmänhet med tiden i serien (dvs. β_i<β_j för jag
5. Diffusionsmodeller är mycket flexibel och tillåta vilken som helst arkitektur vars ingångs- och utgångsdimensionalitet är samma som ska användas. Många implementeringar använder U-Net-liknande arkitekturer.
6. Smakämnen utbildningsmål är att maximera sannolikheten för träningsdata. Detta manifesteras som att man ställer in modellparametrarna till minimera den variationsmässiga övre gränsen för den negativa log sannolikheten för data.
7. Nästan alla termer i den objektiva funktionen kan gjutas som KL Avvikelser som ett resultat av vårt Markov-antagande. Dessa värden bli hållbar att beräkna med tanke på att vi använder Gaussians, därför utelämnar behovet av att utföra Monte Carlo approximation.
8. I slutändan använder du en förenklat utbildningsmål att träna en funktion som förutsäger bruskomponenten för en given latent variabel ger de bästa och mest stabila resultaten.
9. A diskret avkodare används för att erhålla loggsannolikheter över pixelvärden som det sista steget i den omvända diffusionsprocessen.

Med den här översikten på hög nivå av Diffusionsmodeller i våra sinnen, låt oss gå vidare för att se hur man använder en Diffusionsmodeller i PyTorch.

Diffusionsmodeller i PyTorch

Även om diffusionsmodeller ännu inte har demokratiserats i samma grad som andra äldre arkitekturer/tillvägagångssätt inom maskininlärning, finns det fortfarande implementeringar tillgängliga för användning. Det enklaste sättet att använda en diffusionsmodell i PyTorch är att använda denoising-diffusion-pytorch paket, som implementerar en bildspridningsmodell som den som diskuteras i den här artikeln. För att installera paketet, skriv helt enkelt följande kommando i terminalen:

pip install denoising_diffusion_pytorch

Minimalt exempel

För att träna en modell och generera bilder importerar vi först de nödvändiga paketen:

import torch
from denoising_diffusion_pytorch import Unet, GaussianDiffusion

Därefter definierar vi vår nätverksarkitektur, i detta fall ett U-nät. De dim parametern anger antalet funktionskartor före den första nedsamplingen, och dim_mults parameter ger multiplikander för detta värde och successiva nedsamplingar:

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

Nu när vår nätverksarkitektur är definierad måste vi definiera själva diffusionsmodellen. Vi skickar in U-Net-modellen som vi precis definierat tillsammans med flera parametrar – storleken på bilder som ska genereras, antalet tidssteg i diffusionsprocessen och ett val mellan L1- och L2-normerna.

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

Nu när diffusionsmodellen är definierad är det dags att träna. Vi genererar slumpmässiga data att träna på och tränar sedan diffusionsmodellen på vanligt sätt:

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

När modellen är utbildad kan vi äntligen generera bilder genom att använda sample() metod för diffusion objekt. Här genererar vi 4 bilder, som bara är brus med tanke på att vår träningsdata var slumpmässig:

sampled_images = diffusion.sample(batch_size = 4)

Utbildning om anpassade data

Smakämnen denoising-diffusion-pytorch paketet låter dig också träna en diffusionsmodell på en specifik datauppsättning. Byt bara ut 'path/to/your/images' sträng med datauppsättningskatalogsökvägen i Trainer() objekt nedan, och ändra image_size till lämpligt värde. Därefter kör du helt enkelt koden för att träna modellen och provar sedan som tidigare. Observera att PyTorch måste kompileras med CUDA aktiverat för att kunna använda Trainer klass:

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()

Nedan kan du se progressiv avbrusning från multivariat gaussiskt brus till MNIST-siffror som liknar omvänd diffusion:

slutord

Diffusionsmodeller är ett konceptuellt enkelt och elegant förhållningssätt till problemet med att generera data. Deras toppmoderna resultat i kombination med icke-motstridande träning har drivit dem till stora höjder, och ytterligare förbättringar kan förväntas under de kommande åren med tanke på deras begynnande status. Speciellt har diffusionsmodeller visat sig vara avgörande för prestanda hos banbrytande modeller som DALL-E2.

Referensprojekt

[1] Djup oövervakad inlärning med Nonequilibrium Termodynamik

[2] Generativ modellering genom att uppskatta gradienter för datadistributionen

[3] Denoising Diffusion Probabilistic Models

[4] Förbättrade tekniker för att träna resultatbaserade generativa modeller

[5] Förbättrade probabilistiska modeller för denoising-diffusion

[6] Diffusionsmodeller slår GAN på bildsyntes

[7] GLID: Mot fotorealistisk bildgenerering och redigering med textstyrda diffusionsmodeller

[8] Hierarkisk text-villkorlig bildgenerering med CLIP-latenter

Tycker du om den här artikeln? Registrera dig för fler AI-forskningsuppdateringar.

Vi meddelar dig när vi släpper fler sammanfattande artiklar som den här.

Relaterad