Variačné autoenkodéry (VAE): nové obrázky zo spojitého latentného priestoru

Nech 𝒳 je množina možných dát (u nás obrázky 8 × 8 s jasmi pixelov v [0, 1], teda 𝒳 = [0, 1]⁶⁴), n ≥ 1 a x₁, …, x_n ∈ 𝒳 sú trénovacie dáta, o ktorých predpokladáme, že sú nezávislými vzorkami z neznámeho rozdelenia p_data na 𝒳. Nech Θ je množina prípustných parametrov a pre každé θ ∈ Θ je p_θ rozdelenie na 𝒳. Generatívny model je sústava rozdelení (p_θ)_{θ ∈ Θ} spolu s postupom, ako z rozdelenia p_θ vzorkovať (vyrábať nové x ~ p_θ). Učenie vyberie θ tak, aby p_θ bolo čo najbližšie k p_data, typicky ako odhad maximálnej vierohodnosti

θ̂ = argmax_{θ ∈ Θ} ∑_i=1ⁿ log p_θ(x_i).

Model s latentnou premennou navyše zavedie latentný priestor 𝒵 = ℝ^k (k ≥ 1), na ňom pevnú hustotu p(z) (prior) a pre každý kód z ∈ 𝒵 rozdelenie p_θ(x | z) na 𝒳 (dekodér). Rozdelenie dát je potom

p_θ(x) = ∫_𝒵 p_θ(x | z) p(z) dz,

a vzorkuje sa v dvoch krokoch:

1. vylosuj z ~ p(z);
2. vylosuj x ~ p_θ(x | z) — v praxi (aj u nás) sa často berie rovno výstup dekodéra, t. j. pravdepodobnosti pixelov.

Nové x teda vyzerajú ako dáta, ale nie sú ich kópiou. V ďalšom texte index θ vynechávame a píšeme p(x | z), p(x).

Nech 𝒳 = [0, 1]⁶⁴ je množina obrázkov, k ≥ 1 rozmer kódu, I jednotková matica k × k a β > 0. Všetky logaritmy v tomto diele sú prirodzené (log = ln). Variačný autoenkodér (VAE; Diederik Kingma a Max Welling, 2013) sa skladá z troch častí:

1. enkodér — neurónová sieť, ktorá obrázku x ∈ 𝒳 priradí vektory μ(x) ∈ ℝ^k a log σ²(x) ∈ ℝ^k (logaritmus po zložkách); tie určujú rozdelenie kódu q(z | x) = N(μ(x), diag σ²(x)), t. j. súradnice z₁, …, z_k sú nezávislé a z_j ~ N(μ_j(x), σ_j²(x)) (normálne rozdelenie);
2. dekodér — neurónová sieť, ktorá kódu z ∈ ℝ^k priradí rozdelenie p(x | z) na 𝒳 (u nás vektorom pravdepodobností pixelov y(z) ∈ (0, 1)⁶⁴);
3. prior p(z) = N(0, I), normované normálne rozdelenie na ℝ^k.

Parametre oboch sietí sa učia minimalizáciou priemeru cez trénovacie obrázky zo straty

L(x) = −E_{z ~ q(z | x)}[log p(x | z)] + β · D_KL(q(z | x) ‖ N(0, I)),

kde pre hustoty q, p na ℝ^k je D_KL(q ‖ p) = ∫ q(z) log q(z)/p(z) dz spojitá obdoba KL divergencie z dielu o teórii informácie (s prirodzeným logaritmom, v natoch). Prvý člen straty je chyba rekonštrukcie (u nás binárna krížová entropia pixelov, pozri nižšie), druhý trest za to, ako veľmi sa rozdelenie kódu líši od prioru. Pre q = N(μ, diag σ²) vychádza

D_KL(q ‖ N(0, I)) = ½ ∑_j=1^k (μ_j² + σ_j² − log σ_j² − 1).

VAE teda kóduje obrázok do rozdelenia, nie do bodu, a trestá ho za odchýlku od rozdelenia, z ktorého sa potom generuje. S β = 1 je −L(x) dolná hranica log p(x) (ELBO, veta nižšie).

Nech p(z) je hustota na ℝ^k (prior), pre každé z ∈ ℝ^k je p(x | z) rozdelenie na 𝒳 (dekodér), x ∈ 𝒳 je obrázok s p(x) = ∫ p(x | z) p(z) dz > 0 a

p(z | x) = p(x | z) p(z)/p(x)

je aposteriórna hustota kódu (Bayesova veta). Nech q(z | x) je ľubovoľná hustota na ℝ^k, ktorá je kladná len tam, kde je kladná p(z | x), a pre ktorú nižšie uvedené stredné hodnoty E_q = E_{z ~ q(z | x)} existujú a sú konečné. Potom

log p(x) = E_q[log p(x | z)] − D_KL(q(z | x) ‖ p(z)) + D_KL(q(z | x) ‖ p(z | x)),

a preto

log p(x) ≥ E_q[log p(x | z)] − D_KL(q(z | x) ‖ p(z)),

s rovnosťou práve vtedy, keď q(· | x) = p(· | x).

Dôkaz. log p(x) nezávisí od z, takže sa rovná svojej strednej hodnote cez q. Podľa Bayesovej vety je pre každé z s q(z | x) > 0

p(x) = p(x | z) p(z)/p(z | x) = p(x | z) · p(z)/q(z | x) · q(z | x)/p(z | x),

kde sme zlomok rozšírili q(z | x). Zlogaritmujeme a vezmeme strednú hodnotu cez z ~ q(z | x):

log p(x) = E_q[log p(x | z)] + E_q[log p(z)/q(z | x)] + E_q[log q(z | x)/p(z | x)].

Druhý člen je podľa definície −D_KL(q(z | x) ‖ p(z)), tretí D_KL(q(z | x) ‖ p(z | x)). KL divergencia je vždy nezáporná a rovná nule práve pre zhodné rozdelenia (Gibbsova nerovnosť z dielu o teórii informácie, pre hustoty platí rovnako), takže vynechaním tretieho člena dostaneme nerovnosť aj podmienku rovnosti. ∎

Pravá strana nerovnosti je ELBO (evidence lower bound) a s β = 1 je to presne −L(x). Minimalizáciou L teda zdvíhame dolnú hranicu log p(x) bez toho, aby sme museli počítať integrál. Hranica je tým tesnejšia, čím bližšie je q(z | x) k skutočnému aposteriórnemu rozdeleniu p(z | x).

Nech x ∈ [0, 1]⁶⁴ je obrázok (jasy pixelov x_j, j = 1, …, 64) a y ∈ (0, 1)⁶⁴ výstup dekodéra, kde y_j chápeme ako pravdepodobnosť, že pixel j svieti. Binárna krížová entropia obrázka x voči predpovedi y je

BCE(x, y) = −∑_j=1⁶⁴ [ x_j log y_j + (1 − x_j) log(1 − y_j) ].

Pre čiernobiely obrázok x ∈ {0, 1}⁶⁴ a dekodér, podľa ktorého pixely svietia nezávisle, každý s pravdepodobnosťou y_j(z), je presne BCE(x, y(z)) = −log p(x | z). Pri sivých jasoch x_j ∈ (0, 1) sa rovnaký výraz používa ako chyba rekonštrukcie.

Príklad pre jeden pixel: svieti (x = 1) a dekodér povie 0,9 → −ln 0,9 ≈ 0,105; povie 0,5 → ln 2 ≈ 0,693; nesvieti (x = 0) a dekodér povie 0,9 → −ln 0,1 ≈ 2,303. Sebaistý omyl stojí najviac. Keďže výstup je sigmoida, y_j = 1/1 + e^−a_j, je derivácia podľa jej vstupu a_j jednoducho ∂BCE/∂a_j = y_j − x_j — rovnako ako pri logistickej regresii.

Nech μ(x) ∈ ℝ^k a σ(x) ∈ (0, ∞)^k sú výstupy enkodéra pre obrázok x. Reparametrizačný trik namiesto priameho losovania z ~ q(z | x) = N(μ(x), diag σ²(x)) postupuje takto:

1. vylosuj ε = (ε₁, …, ε_k) s nezávislými ε_j ~ N(0, 1) — šum, ktorého rozdelenie nezávisí od parametrov siete;
2. polož pre j = 1, …, k z_j = μ_j(x) + σ_j(x) · ε_j.

Takto získané z má rozdelenie presne q(z | x) (lineárna transformácia normálnej veličiny je normálna), a pritom je to diferencovateľná funkcia parametrov: ∂z_j/∂μ_j = 1 a ∂z_j/∂σ_j = ε_j. Gradient teda prejde spätným šírením až do enkodéra, čo pri samotnom losovaní z ~ N(μ, σ²) nejde.

Väčšinu trestu tu nerobí vzdialenosť od začiatku, ale úzkosť obláčika: člen −log σ² prispieva 2,695 a 1,931. Enkodér by chcel kód čo najpresnejší, D_KL ho drží širší. Priemerná D_KL cez testovacie číslice je 3,42 natu. Chyba rekonštrukcie tejto číslice (zo stredu μ) je 34,6 natu na 64 pixelov.

Nech x⁽¹⁾, …, x⁽⁶⁰⁰⁾ ∈ [0, 1]⁶⁴ sú trénovacie číslice. Vzdialenosť obrázka y ∈ [0, 1]⁶⁴ k najbližšej trénovacej číslici je

d(y) = min_{i ∈ {1, …, 600}} 1/64 ∑_j=1⁶⁴ (y_j − x⁽ⁱ⁾_j)²,

teda stredná kvadratická chyba (MSE) k najpodobnejšej číslici. Pre vylosované obrázky y⁽¹⁾, …, y⁽²⁰⁰⁾ je kvalita generátora priemer

d̄ = 1/200 ∑_s=1²⁰⁰ d(y^(s)).

Malé d̄ znamená „vzorky vyzerajú ako niektoré skutočné číslice“ — pozor, nulu by dostal aj generátor, ktorý len kopíruje trénovacie dáta.

Nech c⁽¹⁾, …, c^(N) ∈ ℝ^k sú kódy, ktoré enkodér autoenkodéra priradí N = 200 testovacím číslicam. Napasované rozdelenie je N(m, diag s²), kde pre j = 1, …, k

m_j = 1/N ∑_i=1^N c⁽ⁱ⁾_j,   s_j = √1/N ∑_i=1^N (c⁽ⁱ⁾_j − m_j)².

Autoenkodér teda losuje z_j = m_j + s_j · ε_j s nezávislými ε_j ~ N(0, 1), t. j. každú súradnicu z normálneho rozdelenia s rovnakým priemerom a smerodajnou odchýlkou, aké majú jeho skutočné kódy (pre k = 2: m = (−3,9; −0,9), s = (4,3; 6,0)). Súradnice sa losujú nezávisle — prípadnú závislosť medzi nimi napasované rozdelenie nezachytí.

Nat je jednotka informácie zodpovedajúca prirodzenému logaritmu: jav s pravdepodobnosťou p ∈ (0, 1] nesie −ln p natov (namiesto −log₂ p bitov). Keďže ln p = ln 2 · log₂ p, je

1 nat = 1/ln 2 bitu ≈ 1,443 bitu.

Nech x je náhodný obrázok s rozdelením dát p_data (u nás rovnomerne vybraná testovacia číslica), kód z má pri danom x rozdelenie q(z | x) enkodéra a p(z) = N(0, I) je prior. Označme

q(z) = E_{x ~ pdata}[q(z | x)]

rozdelenie kódov všetkých obrázkov dohromady a I(x; z) vzájomnú informáciu obrázka a kódu (spojitá obdoba definície z dielu o teórii informácie). Potom platí

E_{x ~ pdata}[D_KL(q(z | x) ‖ p(z))] = I(x; z) + D_KL(q(z) ‖ p(z)) ≥ I(x; z),

teda priemerná KL divergencia cez dáta zhora ohraničuje informáciu, ktorú kód nesie o obrázku. Je to „cena“, ktorú enkodér platí za to, čo do kódu o obrázku uloží.

Pri β = 0,1 je to najviac 8,18 natu ≈ 11,8 bitu na obrázok, pri β = 5 len 0,05 natu ≈ 0,07 bitu — menej než jedna desatina bitu, kód teda o obrázku takmer nič nehovorí a dekodéru nezostáva nič iné, než kresliť stále to isté.