Zrozumieć modele dyfuzyjne: Dogłębne spojrzenie na generatywną sztuczną inteligencję

Modele dyfuzyjne wyłoniły się jako potężne podejście w generatywnej sztucznej inteligencji, osiągając najwyższej klasy wyniki w generowaniu obrazów, dźwięku i wideo. W tym dogłębnym artykule technicznym zbadamy, jak działają modele dyfuzyjne, ich kluczowe innowacje i dlaczego odniosły taki sukces. Omówimy matematyczne podstawy, proces uczenia, algorytmy próbkowania oraz najnowocześniejsze zastosowania tej ekscytującej nowej technologii.

Wprowadzenie do modeli dyfuzyjnych

Modele dyfuzyjne to klasa modeli generatywnych, które uczą się stopniowo usuwać szum z danych poprzez odwrócenie procesu dyfuzji. Główna idea polega na rozpoczęciu od czystego szumu i iteracyjnym udoskonalaniu go w celu uzyskania wysokiej jakości próbki z docelowego rozkładu.

To podejście zostało zainspirowane termodynamiką nierównowagową – a konkretnie procesem odwracania dyfuzji w celu odzyskania struktury. W kontekście uczenia maszynowego można o tym myśleć jako o uczeniu się odwracania stopniowego dodawania szumu do danych.

Niektóre kluczowe zalety modeli dyfuzyjnych obejmują:

• Najwyższej jakości obrazy, przewyższające w wielu przypadkach GAN-y
• Stabilne uczenie bez dynamiki przeciwstawnej
• Wysoka równoległość
• Elastyczna architektura – można użyć dowolnego modelu, który mapuje wejścia na wyjścia o tym samym wymiarze
• Solidne podstawy teoretyczne

Zanurzmy się głębiej w to, jak działają modele dyfuzyjne.

Źródło: Song et al.

Stochastyczne równania różniczkowe rządzą procesami do przodu i wstecz w modelach dyfuzyjnych. Równanie SDE do przodu dodaje szum do danych, stopniowo przekształcając je w rozkład szumu. Równanie SDE wstecz, kierowane przez wyuczoną funkcję score, stopniowo usuwa szum, prowadząc do generowania realistycznych obrazów z losowego szumu. To podejście jest kluczowe dla osiągnięcia wysokiej jakości generatywnej w ciągłych przestrzeniach stanów.

Proces dyfuzji do przodu

Proces dyfuzji do przodu zaczyna się od punktu danych x₀ pobranego z rzeczywistego rozkładu danych i stopniowo dodaje szum Gaussa przez T kroków czasowych, aby wytworzyć coraz bardziej zaszumione wersje x₁, x₂, …, xT.

W każdym kroku czasowym t dodajemy niewielką ilość szumu zgodnie z:

x_t = √(1 - β_t) * x_{t-1} + √(β_t) * ε

Gdzie:

• β_t to harmonogram wariancji kontrolujący, ile szumu jest dodawane w każdym kroku
• ε to losowy szum Gaussa

Ten proces trwa, aż xT stanie się niemal czystym szumem Gaussa.

Matematycznie możemy to opisać jako łańcuch Markowa:

q(x_t | x_{t-1}) = N(x_t; √(1 - β_t) * x_{t-1}, β_t * I)

Gdzie N oznacza rozkład Gaussa.

Harmonogram β_t jest zazwyczaj wybierany tak, aby był mały dla wczesnych kroków czasowych i wzrastał z czasem. Typowe wybory obejmują harmonogramy liniowe, cosinusowe lub sigmoidalne.

Proces dyfuzji wstecznej

Celem modelu dyfuzyjnego jest nauczenie się odwrotności tego procesu – rozpoczęcie od czystego szumu xT i stopniowe odszumianie go w celu odzyskania czystej próbki x₀.

Modelujemy ten proces odwrotny jako:

p_θ(x_{t-1} | x_t) = N(x_{t-1}; μ_θ(x_t, t), σ_θ^2(x_t, t))

Gdzie μ_θ i σ_θ^2 to wyuczone funkcje (zazwyczaj sieci neuronowe) sparametryzowane przez θ.

Kluczową innowacją jest to, że nie musimy jawnie modelować pełnego rozkładu odwrotnego. Zamiast tego możemy go sparametryzować za pomocą procesu do przodu, który znamy.

Konkretnie, możemy pokazać, że optymalna średnia procesu odwrotnego μ* to:

μ* = 1/√(1 - β_t) * (x_t - β_t/√(1 - α_t) * ε_θ(x_t, t))

Gdzie:

• α_t = 1 – β_t
• ε_θ to wyuczona sieć predykcji szumu

To daje nam prosty cel – wytrenuj sieć neuronową ε_θ, aby przewidywała szum, który został dodany w każdym kroku.

Cel uczenia

Cel uczenia dla modeli dyfuzyjnych można wyprowadzić z wnioskowania wariacyjnego. Po pewnym uproszczeniu dochodzimy do prostej straty L2:

L = E_t,x₀,ε [ ||ε - ε_θ(x_t, t)||² ]

Gdzie:

• t jest próbkowane równomiernie od 1 do T
• x₀ jest próbkowane z danych treningowych
• ε to próbkowany szum Gaussa
• x_t jest konstruowane przez dodanie szumu do x₀ zgodnie z procesem do przodu

Innymi słowy, trenujemy model, aby przewidywał szum, który został dodany w każdym kroku czasowym.

Architektura modelu

Źródło: Ronneberger et al.

Architektura U-Net jest centralna dla kroku odszumiania w modelu dyfuzyjnym. Charakteryzuje się strukturą enkoder-dekoder z połączeniami skip, które pomagają zachować drobne szczegóły podczas procesu rekonstrukcji. Enkoder stopniowo zmniejsza rozdzielczość obrazu wejściowego, przechwytując cechy wysokiego poziomu, a dekoder zwiększa rozdzielczość zakodowanych cech, aby zrekonstruować obraz. Ta architektura jest szczególnie skuteczna w zadaniach wymagających precyzyjnej lokalizacji, takich jak segmentacja obrazu.

Sieć predykcji szumu ε_θ może używać dowolnej architektury, która mapuje wejścia na wyjścia o tym samym wymiarze. Architektury w stylu U-Net są popularnym wyborem, szczególnie w zadaniach generowania obrazów.

Typowa architektura może wyglądać następująco:

 class DiffusionUNet(nn.Module): def __init__(self): super().__init__() # Downsampling self.down1 = UNetBlock(3, 64) self.down2 = UNetBlock(64, 128) self.down3 = UNetBlock(128, 256) # Bottleneck self.bottleneck = UNetBlock(256, 512) # Upsampling self.up3 = UNetBlock(512, 256) self.up2 = UNetBlock(256, 128) self.up1 = UNetBlock(128, 64) # Output self.out = nn.Conv2d(64, 3, 1) def forward(self, x, t): # Embed timestep t_emb = self.time_embedding(t) # Downsample d1 = self.down1(x, t_emb) d2 = self.down2(d1, t_emb) d3 = self.down3(d2, t_emb) # Bottleneck bottleneck = self.bottleneck(d3, t_emb) # Upsample u3 = self.up3(torch.cat([bottleneck, d3], dim=1), t_emb) u2 = self.up2(torch.cat([u3, d2], dim=1), t_emb) u1 = self.up1(torch.cat([u2, d1], dim=1), t_emb) # Output return self.out(u1)