Variační autoenkodéry (VAE): nové obrázky ze spojitého latentního prostoru

Nechť 𝒳 je množina možných dat (u nás obrázky 8 × 8 s jasy pixelů v [0, 1], tedy 𝒳 = [0, 1]⁶⁴), n ≥ 1 a x₁, …, x_n ∈ 𝒳 jsou trénovací data, o kterých předpokládáme, že jsou nezávislými vzorky z neznámého rozdělení p_data na 𝒳. Nechť Θ je množina přípustných parametrů a pro každé θ ∈ Θ je p_θ rozdělení na 𝒳. Generativní model je soustava rozdělení (p_θ)_{θ ∈ Θ} spolu s postupem, jak z rozdělení p_θ vzorkovat (vyrábět nová x ~ p_θ). Učení vybere θ tak, aby p_θ bylo co nejblíž p_data, typicky jako odhad maximální věrohodnosti

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

Model s latentní proměnnou navíc zavede latentní prostor 𝒵 = ℝ^k (k ≥ 1), na něm pevnou hustotu p(z) (prior) a pro každý kód z ∈ 𝒵 rozdělení p_θ(x | z) na 𝒳 (dekodér). Rozdělení dat je pak

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

a vzorkuje se ve dvou krocích:

1. vylosuj z ~ p(z);
2. vylosuj x ~ p_θ(x | z) — v praxi (i u nás) se často bere rovnou výstup dekodéru, tj. pravděpodobnosti pixelů.

Nová x tedy vypadají jako data, ale nejsou jejich kopií. V dalším textu index θ vynecháváme a píšeme p(x | z), p(x).

Nechť 𝒳 = [0, 1]⁶⁴ je množina obrázků, k ≥ 1 rozměr kódu, I jednotková matice k × k a β > 0. Všechny logaritmy v tomto díle jsou přirozené (log = ln). Variační autoenkodér (VAE; Diederik Kingma a Max Welling, 2013) se skládá ze tří částí:

1. kodér — neuronová síť, která obrázku x ∈ 𝒳 přiřadí vektory μ(x) ∈ ℝ^k a log σ²(x) ∈ ℝ^k (logaritmus po složkách); ty určují rozdělení kódu q(z | x) = N(μ(x), diag σ²(x)), tj. souřadnice z₁, …, z_k jsou nezávislé a z_j ~ N(μ_j(x), σ_j²(x)) (normální rozdělení);
2. dekodér — neuronová síť, která kódu z ∈ ℝ^k přiřadí rozdělení p(x | z) na 𝒳 (u nás vektorem pravděpodobností pixelů y(z) ∈ (0, 1)⁶⁴);
3. prior p(z) = N(0, I), normované normální rozdělení na ℝ^k.

Parametry obou sítí se učí minimalizací průměru přes trénovací obrázky ze ztráty

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

kde pro hustoty q, p na ℝ^k je D_KL(q ‖ p) = ∫ q(z) log q(z)/p(z) dz spojitá obdoba KL divergence z dílu o informaci (s přirozeným logaritmem, v natech). První člen ztráty je chyba rekonstrukce (u nás binární křížová entropie pixelů, viz níže), druhý trest za to, jak moc se rozdělení kódu liší od prioru. Pro q = N(μ, diag σ²) vychází

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

VAE tedy kóduje obrázek do rozdělení, ne do bodu, a trestá ho za odchylku od rozdělení, ze kterého se pak generuje. S β = 1 je −L(x) dolní mez log p(x) (ELBO, věta níže).

Nechť p(z) je hustota na ℝ^k (prior), pro každé z ∈ ℝ^k je p(x | z) rozdělení na 𝒳 (dekodér), x ∈ 𝒳 je obrázek s p(x) = ∫ p(x | z) p(z) dz > 0 a

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

je aposteriorní hustota kódu (Bayesova věta). Nechť q(z | x) je libovolná hustota na ℝ^k, která je kladná jen tam, kde je kladná p(z | x), a pro kterou níže uvedené střední hodnoty E_q = E_{z ~ q(z | x)} existují a jsou konečné. Pak

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 proto

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

s rovností právě když q(· | x) = p(· | x).

Důkaz. log p(x) nezávisí na z, takže je rovno své střední hodnotě přes q. Podle Bayesovy věty je pro 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 jsme zlomek rozšířili q(z | x). Zlogaritmujeme a vezmeme střední hodnotu přes 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 podle definice −D_KL(q(z | x) ‖ p(z)), třetí D_KL(q(z | x) ‖ p(z | x)). KL divergence je vždy nezáporná a rovná nule právě pro shodná rozdělení (Gibbsova nerovnost z dílu o informaci, pro hustoty platí stejně), takže vynecháním třetího členu dostaneme nerovnost i podmínku rovnosti. ∎

Pravá strana nerovnosti je ELBO (evidence lower bound) a s β = 1 je to přesně −L(x). Minimalizací L tedy zvedáme dolní mez log p(x), aniž bychom museli počítat integrál. Mez je tím těsnější, čím blíž je q(z | x) skutečnému aposteriornímu rozdělení p(z | x).

Nechť x ∈ [0, 1]⁶⁴ je obrázek (jasy pixelů x_j, j = 1, …, 64) a y ∈ (0, 1)⁶⁴ výstup dekodéru, kde y_j chápeme jako pravděpodobnost, že pixel j svítí. Binární křížová entropie obrázku x vůči předpovědi y je

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

Pro černobílý obrázek x ∈ {0, 1}⁶⁴ a dekodér, podle kterého pixely svítí nezávisle, každý s pravděpodobností y_j(z), je přesně BCE(x, y(z)) = −log p(x | z). U šedých jasů x_j ∈ (0, 1) se stejný výraz používá jako chyba rekonstrukce.

Příklad pro jeden pixel: svítí (x = 1) a dekodér řekne 0,9 → −ln 0,9 ≈ 0,105; řekne 0,5 → ln 2 ≈ 0,693; nesvítí (x = 0) a dekodér řekne 0,9 → −ln 0,1 ≈ 2,303. Sebejistý omyl stojí nejvíc. Protože výstup je sigmoida, y_j = 1/1 + e^−a_j, je derivace podle jejího vstupu a_j prostě ∂BCE/∂a_j = y_j − x_j — stejně jako u logistické regrese.

Nechť μ(x) ∈ ℝ^k a σ(x) ∈ (0, ∞)^k jsou výstupy kodéru pro obrázek x. Trik s převzorkováním místo přímého losování z ~ q(z | x) = N(μ(x), diag σ²(x)) postupuje takto:

1. vylosuj ε = (ε₁, …, ε_k) s nezávislými ε_j ~ N(0, 1) — šum, jehož rozdělení nezávisí na parametrech sítě;
2. polož pro j = 1, …, k z_j = μ_j(x) + σ_j(x) · ε_j.

Takto získané z má rozdělení přesně q(z | x) (lineární transformace normální veličiny je normální), a přitom je to derivovatelná funkce parametrů: ∂z_j/∂μ_j = 1 a ∂z_j/∂σ_j = ε_j. Gradient tedy projde zpětným šířením až do kodéru, což u samotného losování z ~ N(μ, σ²) nejde.

Většinu trestu tu nedělá vzdálenost od počátku, ale úzkost obláčku: člen −log σ² přispívá 2,695 a 1,931. Kodér by chtěl kód co nejpřesnější, D_KL ho drží širší. Průměrná D_KL přes testovací číslice je 3,42 natu. Chyba rekonstrukce téhle číslice (ze středu μ) je 34,6 natu na 64 pixelů.

Nechť x⁽¹⁾, …, x⁽⁶⁰⁰⁾ ∈ [0, 1]⁶⁴ jsou trénovací číslice. Vzdálenost obrázku y ∈ [0, 1]⁶⁴ k nejbližší trénovací číslici je

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

tedy střední kvadratická chyba (MSE) k nejpodobnější číslici. Pro vylosované obrázky y⁽¹⁾, …, y⁽²⁰⁰⁾ je kvalita generátoru průměr

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

Malé d̄ znamená „vzorky vypadají jako některé skutečné číslice“ — pozor, nulu by dostal i generátor, který jen kopíruje trénovací data.

Nechť c⁽¹⁾, …, c^(N) ∈ ℝ^k jsou kódy, které kodér autoenkodéru přiřadí N = 200 testovacím číslicím. Napasované rozdělení je N(m, diag s²), kde pro 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 tedy losuje z_j = m_j + s_j · ε_j s nezávislými ε_j ~ N(0, 1), tj. každou souřadnici z normálního rozdělení se stejným průměrem a směrodatnou odchylkou, jaké mají jeho skutečné kódy (pro k = 2: m = (−3,9; −0,9), s = (4,3; 6,0)). Souřadnice se losují nezávisle — případnou závislost mezi nimi napasované rozdělení nezachytí.

Nat je jednotka informace odpovídající přirozenému logaritmu: jev s pravděpodobností p ∈ (0, 1] nese −ln p natů (místo −log₂ p bitů). Protože ln p = ln 2 · log₂ p, je

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

Nechť x je náhodný obrázek s rozdělením dat p_data (u nás rovnoměrně vybraná testovací číslice), kód z má při daném x rozdělení q(z | x) kodéru a p(z) = N(0, I) je prior. Označme

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

rozdělení kódů všech obrázků dohromady a I(x; z) vzájemnou informaci obrázku a kódu (spojitá obdoba definice z dílu o informaci). Pak platí

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

tedy průměrná KL divergence přes data shora omezuje informaci, kterou kód nese o obrázku. Je to „cena“, kterou kodér platí za to, co do kódu o obrázku uloží.

Při β = 0,1 je to nejvýš 8,18 natu ≈ 11,8 bitu na obrázek, při β = 5 jen 0,05 natu ≈ 0,07 bitu — méně než jedna desetina bitu, kód tedy o obrázku skoro nic neříká a dekodéru nezbývá než kreslit pořád totéž.