Module 3 — Autoencodeurs variationnels

Le module précédent a laissé un problème net : l'espace latent d'un autoencodeur ordinaire n'a pas de structure exploitable pour générer. L'autoencodeur variationnel (VAE) le corrige d'une manière élégante : au lieu de laisser l'encodeur produire un code quelconque, il l'oblige à produire une distribution proche d'une gaussienne standard. Cette contrainte suffit à rendre l'échantillonnage possible.

Encoder une distribution, pas un point

Dans un VAE, l'encodeur prend $x$ et renvoie deux vecteurs, une moyenne $\mu(x)$ et un écart-type $\sigma(x)$. Le code $z$ n'est plus une valeur unique : c'est un tirage aléatoire selon $\mathcal{N}(\mu(x), \sigma(x)^2)$. Le décodeur reçoit cet échantillon et reconstruit $x$.

L'objectif à maximiser est la borne inférieure de vraisemblance (ELBO), plus commode à écrire que la vraisemblance directe :

$$\mathcal{L}_{\text{ELBO}}(x) = \mathbb{E}_{z \sim q_\phi(z \mid x)} \left[ \log p_\theta(x \mid z) \right] - \mathrm{KL}\!\left(q_\phi(z \mid x)\ \|\ p(z)\right)$$

Le premier terme est la reconstruction : plus le décodeur restitue fidèlement $x$ à partir de codes tirés de $q_\phi$, plus il est grand. Le second est la divergence de Kullback-Leibler entre la distribution renvoyée par l'encodeur et une gaussienne standard $p(z) = \mathcal{N}(0, I)$. C'est ce terme qui fait tout le travail de régularisation.

Sur une gaussienne diagonale, la KL a une forme close, ce qui la rend trivialement calculable :

$$\mathrm{KL}(q_\phi \parallel p) = -\tfrac{1}{2} \sum_i \left(1 + \log \sigma_i^2 - \mu_i^2 - \sigma_i^2\right)$$

L'astuce de reparamétrisation

Un problème pratique surgit immédiatement. Le tirage $z \sim \mathcal{N}(\mu, \sigma^2)$ n'est pas différentiable par rapport à $\mu$ et $\sigma$ : on ne peut pas propager le gradient à travers un échantillonnage. Sans gradient, l'encodeur ne peut pas être entraîné par descente.

L'astuce, à retenir absolument, consiste à récrire le tirage :

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

Le hasard est déplacé vers $\varepsilon$, une variable externe qui ne dépend d'aucun paramètre du réseau. Le gradient traverse maintenant $\mu$ et $\sigma$ sans obstacle. C'est ce qui rend l'entraînement du VAE possible en pratique et, plus généralement, ce qui permet à toute méthode par variable latente de fonctionner par descente de gradient.

import torch
from torch import nn

class VAE(nn.Module):
    def __init__(self, dim_latente=20):
        super().__init__()
        self.enc = nn.Sequential(nn.Linear(784, 400), nn.ReLU())
        self.tete_mu = nn.Linear(400, dim_latente)
        self.tete_logvar = nn.Linear(400, dim_latente)
        self.dec = nn.Sequential(
            nn.Linear(dim_latente, 400), nn.ReLU(),
            nn.Linear(400, 784), nn.Sigmoid(),
        )

    def encoder(self, x):
        h = self.enc(x)
        return self.tete_mu(h), self.tete_logvar(h)

    def reparam(self, mu, logvar):
        sigma = torch.exp(0.5 * logvar)
        eps = torch.randn_like(sigma)
        return mu + sigma * eps

    def forward(self, x):
        mu, logvar = self.encoder(x)
        z = self.reparam(mu, logvar)
        return self.dec(z), mu, logvar

Prédisez log-variance, pas variance

Le réseau prédit logvar plutôt que var pour deux raisons. D'abord, une variance doit être strictement positive : passer par une exponentielle garantit ce signe sans contrainte explicite. Ensuite, l'exponentielle stabilise le gradient sur plusieurs ordres de grandeur, indispensable au début de l'entraînement où l'encodeur produit des valeurs très variables. Prédire directement $\sigma$ ou $\sigma^2$ finit presque toujours par un NaN.

La perte complète et son équilibre

La perte du VAE se compose comme la somme de deux termes que l'on cherche à minimiser :

def perte_vae(x, x_hat, mu, logvar):
    reconstruction = nn.functional.binary_cross_entropy(x_hat, x, reduction="sum")
    kl = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
    return reconstruction + kl

L'équilibre entre ces deux termes détermine tout le comportement du modèle. Un poids trop fort sur la reconstruction fait ignorer la contrainte KL : l'encodeur revient à un autoencodeur ordinaire et la génération redevient impossible. Un poids trop fort sur la KL provoque le collapse postérieur : l'encodeur renvoie $\mathcal{N}(0, I)$ pour tout $x$, le décodeur n'apprend plus rien et produit toujours la même image moyenne.

Le $\beta$-VAE ajoute un facteur $\beta$ devant la KL pour explorer ce compromis explicitement. Des valeurs de $\beta$ supérieures à 1 encouragent des dimensions latentes plus indépendantes (utile pour la démêlement), au prix d'une reconstruction dégradée.

Le flou n'est pas un bogue

En pratique, un VAE entraîné sur MNIST produit des chiffres reconnaissables mais visiblement flous, surtout comparés à un GAN entraîné dans les mêmes conditions. Cette caractéristique n'est pas due à un mauvais réglage : elle est structurelle.

Deux causes se combinent. D'une part, la reconstruction est mesurée en norme $\ell_2$ (ou en entropie croisée binaire), qui pénalise autant deux versions décalées d'un même chiffre qu'une moyenne floue des deux. Le réseau apprend donc à produire la moyenne, qui minimise la perte au prix d'un flou. D'autre part, l'échantillon $z$ est bruité par le tirage $\varepsilon$ : le décodeur voit, pour un même $x$, plusieurs $z$ légèrement différents pendant l'entraînement, et produit une image qui interpole entre eux.

Les GAN évitent le premier problème parce qu'ils ne mesurent pas la distance pixel à pixel : leur discriminateur juge la plausibilité globale. C'est le sujet du module suivant.

Ce qu'un VAE MNIST donne à voir

Après une vingtaine d'époques sur MNIST, avec $\dim(z) = 20$, un VAE atteint quelques propriétés que l'autoencodeur ordinaire n'avait pas :

• Échantillonnage : tirer $z \sim \mathcal{N}(0, I)$ et décoder produit des chiffres plausibles la plupart du temps. Certains sont ambigus, aucun n'est du bruit.
• Interpolation entre deux images encodées : la transition reste dans la variété des chiffres, sans zones aberrantes, parce que l'espace latent est plein et non troué.
• Détection d'exemples hors distribution : la reconstruction d'une image qui n'est pas un chiffre (une lettre, un motif) est mauvaise et la vraisemblance basse. C'est un usage souvent oublié du VAE.

Vérifiez la KL en cours d'entraînement

Tracez la valeur moyenne de la KL par lot en même temps que la perte de reconstruction. Une KL qui reste à zéro pendant plus de quelques époques signale un collapse postérieur en cours ; une KL qui explose signale un manque de régularisation. C'est le diagnostic le plus utile pour ajuster $\beta$ ou l'échauffement de la KL, où l'on part de $\beta = 0$ et on l'augmente linéairement.

En résumé

• Le VAE encode chaque $x$ vers une distribution gaussienne $\mathcal{N}(\mu(x), \sigma(x)^2)$ et non un point ; c'est ce qui régularise l'espace latent.
• L'astuce de reparamétrisation $z = \mu + \sigma \odot \varepsilon$ déplace le hasard hors des paramètres et rend l'entraînement possible par descente de gradient.
• La perte est la somme d'une reconstruction et d'une KL vers $\mathcal{N}(0, I)$ ; un déséquilibre produit soit un autoencodeur ordinaire, soit un collapse postérieur.
• Le flou caractéristique des reconstructions d'un VAE vient de la norme $\ell_2$ et du bruit d'échantillonnage ; il est structurel et motive les GAN.

Module suivant : le premier générateur qui produit des images nettes, au prix d'un entraînement notoirement instable.