Variational autoencoders

En vanlig autoencoder kodar varje indata till en punkt i det latenta rummet. Mellan punkterna finns hål där decodern producerar nonsens — därför går den inte att sampla ur.

En VAE kodar i stället till en fördelning: encodern ger ett medelvärde och en spridning, och man drar ett slumpmässigt prov ur den.

Två saker följer av det:

1. Eftersom varje indata täcker ett litet område överlappar områdena, och rummet blir sammanhängande.
2. En straffterm drar fördelningarna mot en standardnormalfördelning, så att hela rummet fylls ut.

Då går det att sampla: dra $z \sim \mathcal{N}(0, I)$, kör decodern, få något nytt.

Priset är suddigare rekonstruktioner än en vanlig autoencoder. Det är en direkt följd av att modellen tvingas göra rummet användbart i stället för att bara minimera rekonstruktionsfelet.

Målet är att maximera $\log p(x)$, men integralen över alla $z$ är oberäknelig:

$p(x) = \int p(x \mid z)\,p(z)\, dz$

Variationsansatsen: inför en approximativ posterior $q_\phi(z\mid x)$ och skriv om. För valfri $q$ gäller

$\log p(x) = \underbrace{\mathbb{E}_{q}\left[\log p_\theta(x\mid z)\right] - \mathrm{KL}\!\left(q_\phi(z\mid x)\,\|\,p(z)\right)}_{\text{ELBO}} + \underbrace{\mathrm{KL}\!\left(q_\phi(z\mid x)\,\|\,p(z\mid x)\right)}_{\geq 0}$

Den sista termen är okänd men alltid icke-negativ. Alltså är ELBO en undre gräns för $\log p(x)$, och att maximera ELBO trycker upp golvet.

ELBO har två delar med tydlig tolkning:

Term	Betyder
$\mathbb{E}_q[\log p(x\mid z)]$	rekonstruktion — kan decodern återskapa x?
$-\mathrm{KL}(q \| p)$	regularisering — liknar den kodade fördelningen priorn?

Reparametriseringstricket. Att sampla $z \sim \mathcal{N}(\mu, \sigma^2)$ är inte deriverbart med avseende på $\mu$ och $\sigma$. Lösningen är att flytta ut slumpen:

$z = \mu + \sigma \odot \varepsilon, \qquad \varepsilon \sim \mathcal{N}(0, I)$

Nu är $z$ en deriverbar funktion av $\mu$ och $\sigma$, och $\varepsilon$ är en konstant indata. Utan det tricket går VAE:n inte att träna med gradientnedstigning alls.

KL-termen i sluten form för gaussisk posterior och standardnormal prior:

$\mathrm{KL} = -\tfrac{1}{2}\sum_j \left(1 + \log\sigma_j^2 - \mu_j^2 - \sigma_j^2\right)$

Posterior collapse är det klassiska problemet: KL-termen går till noll, encodern ignoreras, och decodern lär sig producera ett genomsnitt oavsett $z$. Det händer särskilt när decodern är kraftfull (t.ex. autoregressiv).

Motmedel	Hur
KL-uppvärmning	börja med vikt 0 på KL-termen och öka gradvis
Fri bits	tillåt en viss KL per dimension utan straff
Svagare decoder	ge den mindre kapacitet att klara sig utan z
β-VAE med β < 1	vikta ner KL-termen

β-VAE med β > 1 är den motsatta inställningen: hårdare regularisering ger mer disentanglade latenta dimensioner, till priset av sämre rekonstruktion.

import torch, torch.nn as nn, torch.nn.functional as F

class VAE(nn.Module):
    def __init__(self, in_dim=784, dolt=400, latent=20):
        super().__init__()
        self.enc = nn.Sequential(nn.Linear(in_dim, dolt), nn.ReLU())
        self.mu = nn.Linear(dolt, latent)
        self.log_var = nn.Linear(dolt, latent)
        self.dec = nn.Sequential(nn.Linear(latent, dolt), nn.ReLU(),
                                 nn.Linear(dolt, in_dim))

    def koda(self, x):
        h = self.enc(x)
        return self.mu(h), self.log_var(h)

    def reparametrisera(self, mu, log_var):
        std = torch.exp(0.5 * log_var)
        return mu + std * torch.randn_like(std)      # slumpen flyttad ut → deriverbart

    def forward(self, x):
        mu, log_var = self.koda(x)
        z = self.reparametrisera(mu, log_var)
        return self.dec(z), mu, log_var

def elbo(logits, x, mu, log_var, beta=1.0, fri_bits=0.0):
    rek = F.binary_cross_entropy_with_logits(logits, x, reduction="sum") / len(x)
    kl_per_dim = -0.5 * (1 + log_var - mu.pow(2) - log_var.exp())
    if fri_bits > 0:
        kl_per_dim = torch.clamp(kl_per_dim, min=fri_bits)   # tillåt viss KL gratis
    kl = kl_per_dim.sum(dim=1).mean()
    return rek + beta * kl, {"rekonstruktion": float(rek), "kl": float(kl)}

# KL-uppvärmning motverkar posterior collapse
def beta_schema(steg, uppvarmning=5000, max_beta=1.0):
    return min(max_beta, max_beta * steg / uppvarmning)

def trana(modell, dataloader, opt, steg_totalt=20000, fri_bits=0.05):
    steg = 0
    for x, _ in dataloader:
        x = x.flatten(1)
        logits, mu, log_var = modell(x)
        forlust, delar = elbo(logits, x, mu, log_var,
                              beta=beta_schema(steg, 5000), fri_bits=fri_bits)
        opt.zero_grad(); forlust.backward(); opt.step()
        steg += 1
        if steg % 1000 == 0:
            print(f"  steg {steg}: rek {delar['rekonstruktion']:.1f}  kl {delar['kl']:.3f}")
        if steg >= steg_totalt:
            break

# Upptäck posterior collapse: KL per dimension nära noll
def kollapsade_dimensioner(modell, dataloader, troskel=0.01):
    modell.eval()
    kl_sum = None
    n = 0
    with torch.no_grad():
        for x, _ in dataloader:
            mu, log_var = modell.koda(x.flatten(1))
            kl = -0.5 * (1 + log_var - mu.pow(2) - log_var.exp())
            kl_sum = kl.sum(0) if kl_sum is None else kl_sum + kl.sum(0)
            n += len(x)
    per_dim = kl_sum / n
    return {"kl_per_dim": [round(float(v), 4) for v in per_dim],
            "antal_kollapsade": int((per_dim < troskel).sum()),
            "av_totalt": len(per_dim)}
# {'antal_kollapsade': 14, 'av_totalt': 20}  ← bara 6 dimensioner används

# Generera: sampla ur priorn
@torch.no_grad()
def sampla(modell, n=16, latent=20):
    return torch.sigmoid(modell.dec(torch.randn(n, latent)))

kollapsade_dimensioner är diagnostiken som saknas oftast. En VAE kan se ut att träna bra medan fjorton av tjugo latenta dimensioner är oanvända — och då är den latenta representationen mycket smalare än man tror.