73 - VAE: il peso beta, il collasso del posterior e le instabilità dell'addestramento

Un VAE addestrato male non ti dà un errore: la perdita scende, il programma termina, le immagini escono. Il difetto è dentro, e lo vedi solo se sai che cosa misurare. Nella lezione precedente il VAE ha funzionato al primo colpo, ma è un caso fortunato: tutti i suoi ingredienti erano in equilibrio. Qui rompiamo l'equilibrio, un ingrediente per volta, e guardiamo che cosa si rompe: il peso β fra ricostruzione e KL, il collasso del posterior (dimensioni del codice che smettono di portare informazione) e le instabilità numeriche. Per ognuno c'è una misura che lo rivela in pochi secondi. Tutti i numeri vengono da esecuzioni reali, ripetute due volte con uscite identiche, su cifre scritte a mano 8×8 (load_digits di scikit-learn) e non su MNIST: l'esecuzione non ha scaricato dataset, e le cifre dei libri su MNIST qui non valgono.

Le parole di questa lezione

Si usano i termini delle due lezioni precedenti (ELBO, KL, prior, riparametrizzazione, dimensione attiva) e ne servono pochi altri. Un β-VAE è un VAE in cui la KL è moltiplicata per un coefficiente: perdita = ricostruzione + β · KL. Con β = 1 è il VAE standard, cioè il vero ELBO; il nome viene da Higgins e colleghi (2017), che usano β > 1 per ottenere rappresentazioni più «ordinate» (attribuzione non verificata sul testo, vedi le fonti). Il collasso del posterior avviene quando, per una o più dimensioni latenti, l'encoder produce lo stesso N(0, 1) qualunque sia l'immagine: la KL di quella dimensione è zero e il decoder la ignora, per lui è una dimensione morta.

Il KL annealing (o warm-up) parte con β = 0 e lo porta gradualmente a 1 nelle prime epoche, per lasciare al modello il tempo di imparare a usare il codice prima di penalizzarlo. Il posterior aggregato, q(z), è la media di q(z | x) sull'intero dataset, cioè la distribuzione dei codici che l'encoder produce davvero: per generare con z ~ N(0, I) deve somigliare al prior. Infine il compromesso rate-distortion: la KL è il costo del codice in nat (rate), l'errore di ricostruzione è la distorsione, e ogni valore di β sceglie un punto su quella curva.

L'analogia: il budget del corriere

Immagina un corriere che deve trasmettere un'immagine a un collega e paga a nat: più dettagli manda, più spende. β è il prezzo al nat. Se il messaggio costa zero (β = 0) il corriere manda tutto, con precisione folle, e il collega ricostruisce benissimo; ma i messaggi sono scritti in un dialetto che solo i due conoscono, quindi nessun altro può ordinare «un messaggio a caso» e ottenere qualcosa di sensato. Se costa troppo, il corriere non manda più niente e il collega risponde sempre con la stessa immagine media. Il VAE standard sta a metà strada, e l'esperimento che segue dice dove stanno i confini.

Il punto di partenza, autonomo

Gli esperimenti riusano modello e perdita della lezione precedente, ma riscritti qui per intero, così che il blocco si copi e giri da solo. Cambiano due cose: la dimensione del codice latent è un parametro, perché la varieremo, e in fondo c'è il «giudice», un piccolo classificatore di cifre addestrato a parte. Il giudice non misura la qualità percettiva: serve a contare, su 2000 immagini generate, quali cifre compaiono e con quale frequenza, cioè la diversità dei campioni.

import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split

torch.set_num_threads(2)
X, y = load_digits(return_X_y=True)                                # 1797 digits, 8x8 pixels, values 0..16
Xtr, Xte, ytr, yte = train_test_split(X / 16.0, y, test_size=0.3, random_state=0, stratify=y)
Xtr, Xte = torch.tensor(Xtr, dtype=torch.float32), torch.tensor(Xte, dtype=torch.float32)
ytr, yte = torch.tensor(ytr), torch.tensor(yte)
D = 64

class VAE(nn.Module):
    def __init__(self, latent=10, hidden=128):
        super().__init__()
        self.enc = nn.Sequential(nn.Linear(D, hidden), nn.ReLU())
        self.mu = nn.Linear(hidden, latent)
        self.logvar = nn.Linear(hidden, latent)
        self.dec = nn.Sequential(nn.Linear(latent, hidden), nn.ReLU(), nn.Linear(hidden, D))

    def encode(self, x):
        h = self.enc(x)
        return self.mu(h), self.logvar(h)

    def reparametrize(self, mu, logvar):
        return mu + torch.exp(0.5 * logvar) * torch.randn_like(mu)

    def decode(self, z):
        return self.dec(z)                                          # logits

    def forward(self, x):
        mu, logvar = self.encode(x)
        return self.decode(self.reparametrize(mu, logvar)), mu, logvar

def vae_loss(logits, x, mu, logvar):
    """Negative ELBO per image in nats: reconstruction (BCE summed over pixels) + KL (summed over dimensions)."""
    recon = F.binary_cross_entropy_with_logits(logits, x, reduction="sum") / x.size(0)
    kl = (0.5 * (mu.pow(2) + logvar.exp() - logvar - 1)).sum() / x.size(0)
    return recon + kl, recon, kl

# The "judge": a small classifier that tells us which digit it sees in a generated image
torch.manual_seed(0)
clf = nn.Sequential(nn.Linear(D, 64), nn.ReLU(), nn.Linear(64, 10))
opt_c = torch.optim.Adam(clf.parameters(), lr=1e-3)
for _ in range(60):
    perm = torch.randperm(len(Xtr))
    for i in range(0, len(Xtr), 64):
        idx = perm[i:i + 64]
        loss = F.cross_entropy(clf(Xtr[idx]), ytr[idx]); opt_c.zero_grad(); loss.backward(); opt_c.step()
with torch.no_grad():
    print("judge accuracy on real test digits:", round((clf(Xte).argmax(-1) == yte).float().mean().item(), 4))

judge accuracy on real test digits: 0.9778

La perdita vae_loss somma la BCE sui 64 pixel e la KL sulle dimensioni, poi divide per il numero di immagini: il risultato è in nat per immagine, ed è l'opposto dell'ELBO. Il giudice raggiunge il 97.8% sul test, abbastanza per fidarsi dei conteggi che ne ricaviamo; sulle immagini sfocate dei campioni però assegnerà comunque una cifra, quindi un conteggio equilibrato dice che nessuna classe manca, non che le immagini sono nitide.

Esperimento 1: spazzare β

Alleniamo lo stesso VAE per cinque valori di β (0, 0.25, 1, 2, 4), con codice a 10 dimensioni, 60 epoche, un seme. Per ognuno misuriamo ricostruzione e KL sul test, sempre con la vera perdita (β = 1), l'MSE decodificando la media, il numero di dimensioni attive (varianza di μ sul test > 0.01) e l'entropia delle cifre che il giudice vede in 2000 campioni da N(0, I): il massimo è ln 10 = 2.30, tutte le cifre equiprobabili.

def train_vae(beta=1.0, latent=10, epochs=60, seed=0, warmup_epochs=0):
    """beta weights the KL term: loss = reconstruction + beta * KL. With warmup_epochs, beta grows linearly from 0 to its value."""
    torch.manual_seed(seed)
    model = VAE(latent=latent)
    opt = torch.optim.Adam(model.parameters(), lr=1e-3)
    n_batches = (len(Xtr) + 63) // 64
    step = 0
    for epoch in range(epochs):
        perm = torch.randperm(len(Xtr))
        for i in range(0, len(Xtr), 64):
            x = Xtr[perm[i:i + 64]]
            logits, mu, logvar = model(x)
            _, recon, kl = vae_loss(logits, x, mu, logvar)
            frac = min(1.0, step / (warmup_epochs * n_batches)) if warmup_epochs > 0 else 1.0
            loss = recon + beta * frac * kl                       # the weight of the KL, possibly ramped up
            opt.zero_grad(); loss.backward(); opt.step()
            step += 1
    return model

@torch.no_grad()
def report(model, latent):
    logits, mu, logvar = model(Xte)
    _, recon, kl = vae_loss(logits, Xte, mu, logvar)            # always the true ELBO (beta = 1)
    mse = F.mse_loss(torch.sigmoid(model.decode(mu)), Xte).item()
    active = int((mu.var(dim=0) > 0.01).sum())
    samples = torch.sigmoid(model.decode(torch.randn(2000, latent)))
    hist = torch.bincount(clf(samples).argmax(-1), minlength=10).float() / 2000
    entropy = max(0.0, -(hist[hist > 0] * hist[hist > 0].log()).sum().item())
    return recon.item(), kl.item(), mse, active, entropy

La funzione train_vae è l'addestramento della lezione precedente con tre differenze. Il coefficiente beta moltiplica la KL nella perdita. Il parametro latent varia la dimensione del codice. Il parametro warmup_epochs, se positivo, fa crescere linearmente il fattore frac da 0 a 1 (step / (warmup_epochs · n_batches)) e lo usa al posto di 1. La funzione report invece non conosce β: valuta sempre con vae_loss, cioè con il vero ELBO, così le righe della tabella si confrontano fra loro. Se valutassi ogni modello con la perdita con cui è stato addestrato, confronteresti funzioni diverse.

models = {}
print(f"{'beta':>5s} {'recon':>7s} {'KL':>6s} {'MSE':>7s} {'active':>7s} {'class entropy (max 2.30)':>26s}")
for beta in (0.0, 0.25, 1.0, 2.0, 4.0):
    models[beta] = train_vae(beta=beta)
    recon, kl, mse, active, entropy = report(models[beta], 10)
    print(f"{beta:5.2f} {recon:7.1f} {kl:6.1f} {mse:7.4f} {active:4d}/10 {entropy:14.2f}")

 beta   recon     KL     MSE  active   class entropy (max 2.30)
 0.00    16.7  137.2  0.0127   10/10           1.64
 0.25    18.2    9.7  0.0148   10/10           2.29
 1.00    21.9    2.9  0.0302   10/10           2.28
 2.00    25.3    0.9  0.0491    4/10           2.11
 4.00    27.2    0.0  0.0743    0/10           0.00

Il grafico mette in relazione le due colonne centrali: ogni punto è un valore di β, l'asse orizzontale è la KL in scala logaritmica, quello verticale la ricostruzione. Il punto β = 4 non c'è, perché ha KL zero e non si può disegnare su una scala logaritmica.

Leggiamo la tabella dall'alto. Con β = 0 nessuna penalità tocca il codice: la KL esplode a 137 nat per immagine (contro 2.9 a β = 1), la ricostruzione è la migliore (16.7 nat, MSE 0.0127) e il modello è, di fatto, un autoencoder. Ma l'entropia delle classi scende a 1.64: generare da N(0, I) non funziona, lo stesso difetto visto con l'autoencoder della lezione 70. Con β = 0.25 la KL scende a 9.7, la ricostruzione resta vicina (18.2) e l'entropia è già 2.29, la più alta della tabella: campioni diversificati con quasi tutta la qualità di ricostruzione. Con β = 1 (il vero ELBO) i valori sono KL 2.9, ricostruzione 21.9, MSE 0.0302, entropia 2.28.

Poi il comportamento cambia in modo brusco, non graduale. A β = 2 solo 4 dimensioni su 10 restano attive, la ricostruzione sale a 25.3 e l'MSE a 0.0491. A β = 4 le dimensioni attive sono zero: KL 0.0, ricostruzione 27.2, entropia 0.00. È un collasso totale del posterior: l'encoder dà a ogni immagine lo stesso N(0, I), il decoder riceve solo rumore e impara a emettere sempre la stessa immagine. Lo verifichiamo, perché la coincidenza è sospetta: se il modello a β = 4 emette l'immagine media del training, la sua ricostruzione deve coincidere con quella dell'immagine media.

mean = Xtr.mean(0).clamp(1e-6, 1 - 1e-6)                        # the average training image
bce = -(Xte * mean.log() + (1 - Xte) * (1 - mean).log()).sum(1).mean().item()
mse = F.mse_loss(mean.expand_as(Xte), Xte).item()
print(f"average image used for every digit: reconstruction {bce:.1f} nat, MSE {mse:.4f}")
with torch.no_grad():
    s = torch.sigmoid(models[4.0].decode(torch.randn(2000, 10)))
    print("digits the judge sees in 2000 samples of the beta=4 model:", torch.bincount(clf(s).argmax(-1), minlength=10).tolist())

average image used for every digit: reconstruction 27.2 nat, MSE 0.0739
digits the judge sees in 2000 samples of the beta=4 model: [0, 0, 0, 0, 0, 0, 0, 0, 2000, 0]

Coincide: 27.2 nat e MSE 0.0739 (contro 0.0743 del modello a β = 4), e il giudice vede la stessa cifra in tutti i 2000 campioni. Il codice a β = 4 non porta nulla, e il decoder è diventato una costante. Il punto della scala in cui questo succede dipende da dati, architettura e durata dell'addestramento: su MNIST con decoder più grandi i valori di β saranno altri. Qui è un seme e una configurazione: la soglia fra «4 dimensioni attive» e «nessuna» non è una costante, la forma della curva sì.

Nessun β è giusto in assoluto. Per comprimere bene si sceglie β piccolo; per uno spazio ordinato e campionabile si sta vicino a 1; oltre un certo valore si compra semplicità del codice con un modello inutile. Gli autori del β-VAE sostengono che β > 1 favorisca fattori «disaccoppiati»: qui non l'abbiamo misurato, abbiamo misurato solo la perdita di dimensioni attive.

Esperimento 2: il collasso del posterior

Il collasso non avviene solo con β grande. Diamo al VAE un codice di 40 dimensioni invece di 10 e contiamo quante ne usa, con β = 1; poi proviamo il KL warm-up (β cresce da 0 a 1 in 10 epoche) e infine β = 3.

print(f"{'setting':38s} {'recon':>7s} {'KL':>6s} {'MSE':>7s} {'active':>9s}")
for name, kwargs in [("latent 10, beta 1",                     dict(latent=10)),
                     ("latent 40, beta 1",                     dict(latent=40)),
                     ("latent 40, beta 1, KL warm-up 10 ep",   dict(latent=40, warmup_epochs=10)),
                     ("latent 40, beta 3",                     dict(latent=40, beta=3.0))]:
    m = train_vae(**kwargs)
    recon, kl, mse, active, entropy = report(m, kwargs["latent"])
    print(f"{name:38s} {recon:7.1f} {kl:6.1f} {mse:7.4f} {active:5d}/{kwargs['latent']}")

setting                                  recon     KL     MSE    active
latent 10, beta 1                         21.9    2.9  0.0302    10/10
latent 40, beta 1                         23.0    2.4  0.0367    38/40
latent 40, beta 1, KL warm-up 10 ep       22.6    2.6  0.0343    33/40
latent 40, beta 3                         27.1    0.0  0.0789     1/40

Con 40 dimensioni disponibili il VAE ne usa 38, e ricostruzione e KL restano quasi uguali a quelle con 10 (23.0 e 2.4 nat contro 21.9 e 2.9): l'informazione che serve per le cifre sta in poche dimensioni, e le altre vengono spente dalla KL perché non servono. È un collasso benigno, una regolarizzazione che funziona. Il warm-up ha lasciato 33 dimensioni attive e −ELBO 25.2 (22.6 + 2.6) contro 25.4 (23.0 + 2.4) senza warm-up: una differenza di 0.2 nat, con un seme, che non si può distinguere dal caso. Con β = 3 il collasso è quasi totale: 1 dimensione su 40, ricostruzione 27.1.

Perché dedicare attenzione a qualcosa di benigno? Perché in altri contesti il collasso è fatale. Bowman e colleghi (2016) addestrano un VAE per frasi con un decoder LSTM e osservano che il decoder impara a spiegare i dati senza guardare z: a quel punto non passa più segnale fra decoder ed encoder, e il modello si ferma in un equilibrio con KL zero, equivalente a un normale modello di linguaggio. Nel loro articolo l'addestramento standard, senza rimedi, produce infatti modelli con prestazioni pari al modello di linguaggio di base e KL zero. Per questo propongono il KL cost annealing (peso della KL a zero, poi in crescita) e il word dropout. Il nostro decoder è a strati densi e non riesce a modellare un'immagine senza z, quindi il caso grave non si presenta. La regola generale resta: conta le dimensioni attive ogni volta che addestri un VAE.

Diagnostica: il posterior aggregato somiglia al prior?

Per generare con z ~ N(0, I) i codici che l'encoder produce sull'intero dataset devono essere distribuiti più o meno come N(0, I). Il controllo rapido: per ogni dimensione, media e varianza dei z sul test, che dovrebbero essere circa 0 e circa 1. Riusiamo i modelli addestrati nell'esperimento 1.

@torch.no_grad()
def aggregate_posterior(model, X):
    mu, logvar = model.encode(X)
    z = model.reparametrize(mu, logvar)                       # one z per image, as the decoder sees it in training
    return z.mean(dim=0), z.var(dim=0)

for beta in (0.0, 1.0, 4.0):                                  # the models trained in the sweep above
    mean_z, var_z = aggregate_posterior(models[beta], Xte)
    print(f"beta={beta:3.1f}   avg |mean| of z per dim {mean_z.abs().mean():.2f}   variance of z per dim: min {var_z.min():.2f}  mean {var_z.mean():.2f}  max {var_z.max():.2f}")

beta=0.0   avg |mean| of z per dim 2.82   variance of z per dim: min 4.13  mean 8.75  max 17.75
beta=1.0   avg |mean| of z per dim 0.05   variance of z per dim: min 0.88  mean 1.00  max 1.11
beta=4.0   avg |mean| of z per dim 0.05   variance of z per dim: min 0.95  mean 1.00  max 1.10

La funzione aggregate_posterior prende le immagini di test, calcola per ciascuna un codice con la riparametrizzazione (come lo vedrebbe il decoder in addestramento) e restituisce media e varianza per dimensione. Con β = 0 la varianza media dei codici è 8.75, fino a 17.75, e la media assoluta 2.82: il decoder lavora su una scala molto diversa da N(0, 1), quindi campionare dal prior cade in una zona che non ha mai visto. Con β = 1 la varianza media è 1.00 e la media assoluta 0.05: un posterior aggregato allineato al prior. Qui c'è però una trappola, visibile nella terza riga: anche con β = 4 la varianza media è 1.00 e la media 0.05. Il modello collassato passa il controllo perfettamente, perché il suo posterior è identico al prior per costruzione. La diagnostica del posterior aggregato è una condizione necessaria, non sufficiente: va letta insieme al conteggio delle dimensioni attive e alla diversità dei campioni.

Instabilità dell'addestramento

I VAE hanno punti numericamente delicati: exp(logvar) può andare in overflow e i gradienti attraverso σ possono diventare grandi. Proviamo cinque learning rate, con e senza due contromisure: limitare logvar a [−10, 10] e il clipping della norma del gradiente a 5 (torch.nn.utils.clip_grad_norm_).

@torch.no_grad()
def evaluate_loss(model):
    logits, mu, logvar = model(Xte)
    return vae_loss(logits, Xte, mu, logvar)[0].item()

def train_with_lr(lr, stabilise=False, epochs=30, seed=0):
    """Returns the epoch at which the loss became non-finite (or None) and the final test loss."""
    torch.manual_seed(seed)
    model = VAE(latent=10)
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    for epoch in range(1, epochs + 1):
        perm = torch.randperm(len(Xtr))
        for i in range(0, len(Xtr), 64):
            x = Xtr[perm[i:i + 64]]
            logits, mu, logvar = model(x)
            if stabilise:
                logvar = logvar.clamp(-10, 10)                     # keep exp(logvar) in a safe range
            loss, _, _ = vae_loss(logits, x, mu, logvar)
            opt.zero_grad(); loss.backward()
            if stabilise:
                nn.utils.clip_grad_norm_(model.parameters(), 5.0)  # tame huge gradients
            opt.step()
            if not torch.isfinite(loss):
                return epoch, float("nan")
    return None, evaluate_loss(model)

for lr in (1e-3, 1e-2, 5e-2, 2e-1, 5e-1):
    nan_epoch, final = train_with_lr(lr)
    stab_epoch, stab_final = train_with_lr(lr, stabilise=True)
    print(f"lr={lr:g}: plain -> non-finite at epoch {nan_epoch}, final {final:.1f}   |   with clamp+clipping -> non-finite at epoch {stab_epoch}, final {stab_final:.1f}")

lr=0.001: plain -> non-finite at epoch None, final 25.4   |   with clamp+clipping -> non-finite at epoch None, final 25.3
lr=0.01: plain -> non-finite at epoch None, final 24.9   |   with clamp+clipping -> non-finite at epoch None, final 24.7
lr=0.05: plain -> non-finite at epoch None, final 26.8   |   with clamp+clipping -> non-finite at epoch None, final 27.3
lr=0.2: plain -> non-finite at epoch None, final 14224.7   |   with clamp+clipping -> non-finite at epoch None, final 27.8
lr=0.5: plain -> non-finite at epoch 1, final nan   |   with clamp+clipping -> non-finite at epoch 1, final nan

La funzione restituisce l'epoca in cui la perdita diventa non finita (controllo con torch.isfinite), oppure None, più la perdita di test finale. Fino a lr 0.05 nessuna configurazione dà NaN, e le contromisure quasi non incidono (25.4 contro 25.3 a lr 0.001; 26.8 contro 27.3 a 0.05, dove sono perfino leggermente peggiori). A lr 0.2 succede il fatto istruttivo: senza contromisure l'addestramento non produce NaN ma la perdita finale è 14224.7; con clamp e clipping è 27.8. Il fallimento è silenzioso, una perdita enorme senza alcun errore. A lr 0.5 entrambe le versioni producono valori non finiti già alla prima epoca: le contromisure non salvano un learning rate assurdo, e il motivo non l'ho indagato.

Una precisazione sul codice, perché è facile copiarlo male: nel blocco il clamp agisce sul logvar passato alla perdita, quindi protegge il termine KL, ma il campionamento dentro model(x) usa ancora il valore non limitato. Per proteggere anche exp(0.5 · logvar) del campionamento il clamp va dentro reparametrize. Su modelli più grandi (decoder convoluzionali profondi, dati non normalizzati) i NaN diventano più probabili: i due accorgimenti restano buona pratica, ma questo esperimento non lo dimostra, mostra solo il caso a lr 0.2 su un modello piccolo.

Altri fallimenti tipici e dove cercarli

• Immagini sfocate. Non è un bug: deriva dalla verosimiglianza per pixel indipendenti, che favorisce medie di futuri possibili. Per immagini nitide servono altri decoder o altre perdite.
• Ricostruzione ottima, generazione pessima. Sintomo di β troppo piccolo: posterior aggregato lontano dal prior. Controlla la varianza dei z.
• Campioni tutti uguali, ricostruzione pessima. β troppo grande: controlla le dimensioni attive e se il giudice vede sempre la stessa cifra.
• Dimensioni attive pochissime. Collasso: riduci β, riduci la dimensione latente, prova il warm-up o rendi il decoder meno potente.
• ELBO che peggiora con più dati. Probabile errore di somma o media nelle due componenti della perdita.

L'elenco è una tabella di sintomi e cause: ogni voce collega un'osservazione a un controllo, ed è per questo che l'ordine non conta. La riga più utile è la terza, perché l'esperimento ha mostrato che il sintomo «campioni tutti uguali» e il sintomo «posterior aggregato perfetto» possono comparire insieme.

Trade-off

β piccolo dà ricostruzione migliore a spese della regolarità del codice, β grande il contrario fino al collasso (tabella dell'esperimento 1). Una dimensione latente grande dà più spazio e più dimensioni da spegnere: spesso conviene sovrastimare e lasciar fare alla KL, osservando le dimensioni attive. Il warm-up può servire con decoder potenti come quello di Bowman; qui non ha dato un beneficio distinguibile e costa un iperparametro in più.

Trappole reali

1. Confrontare modelli con β diversi usando la loro stessa perdita. Una perdita con β = 4 non è confrontabile con una con β = 1; valuta sempre con β = 1.
2. Giudicare un VAE dalla sola ricostruzione. Con β = 0 la ricostruzione è la migliore (16.7 nat) e il generatore il peggiore tra quelli non collassati.
3. Non guardare le dimensioni attive. Un VAE con 40 dimensioni e 38 attive è normale; uno con 10 dimensioni e 0 attive è rotto.
4. Fidarsi del posterior aggregato da solo. Il modello collassato lo ha perfetto.
5. Pensare che l'assenza di NaN significhi addestramento sano. A lr 0.2 non c'è NaN e la perdita è 14224.7.
6. Un seme, una configurazione. Tutti i numeri qui sono singoli: le tendenze (KL scende, ricostruzione sale con β) sono nette, le differenze piccole (33 contro 38 dimensioni attive) non sono affidabili.

Le sei trappole hanno una causa comune: ciò che il programma stampa non coincide con ciò che vuoi misurare. Una singola cifra (la perdita, la varianza dei codici, l'assenza di NaN) sembra una diagnosi e non lo è. La lista di controllo che ne segue è semplice: registra separatamente ricostruzione e KL a ogni epoca, conta le dimensioni attive, controlla media e varianza del posterior aggregato, campiona qualche migliaio di immagini dal prior e misura la diversità, confronta con un autoencoder dello stesso formato, e se usi β ≠ 1 riporta comunque la perdita a β = 1.

Un esempio svolto a mano: perché β controlla la «precisione» del codice

Un caso a una dimensione in cui il risultato si ricava con la matita. L'immagine vale x = 0 e il decoder è p(x | z) = N(z, s²) con s = 0.3. L'encoder produce una gaussiana con media μ = 0 e deviazione σ da scegliere. La perdita è «ricostruzione + β·KL»; con μ = 0 la ricostruzione attesa è σ²/(2s²) (più σ è grande, più z si discosta da 0 e più il decoder sbaglia, a meno di costanti) e la KL è ½(σ² − ln σ² − 1) (più σ è piccola, più il codice è preciso e costoso). Derivando rispetto a σ² e ponendo a zero:

1/(2s²) + (β/2)·(1 − 1/σ²) = 0     →     σ² = β / ( β + 1/s² )

Con 1/s² = 11.11:

β = 0.1:  σ² = 0.0089  (σ = 0.094)   KL = 1.86 nat
β = 1:    σ² = 0.0826  (σ = 0.287)   KL = 0.79 nat
β = 10:   σ² = 0.474   (σ = 0.688)   KL = 0.11 nat

Il risultato ha la stessa forma dell'effetto misurato: un β alto rende il codice meno preciso (σ più grande) e ne abbassa il costo in informazione (KL più piccola); un β basso lo rende preciso e costoso. A β → 0, σ² → 0: il codice diventa un punto esatto, come in un autoencoder. Nel VAE vero la stessa logica agisce su ogni componente, con la differenza che la precisione necessaria dipende dall'immagine. Questi conti sono fatti a mano, non con il programma, e ricontrollati con un calcolo a parte.

Domande frequenti

Come scelgo β in pratica?

Parti da 1, che è l'ELBO vero. Se ti servono ricostruzioni più fedeli prova β < 1 (qui 0.25 ha dato la ricostruzione di 18.2 nat con campioni ancora diversificati); se vuoi uno spazio più regolare o compatto, β > 1, ma controlla le dimensioni attive: a β = 2 ne sono rimaste 4 su 10. Valuta sempre con la perdita a β = 1.

È vero che un β alto «disaccoppia» i fattori?

È l'idea degli autori del β-VAE; i risultati dipendono dai dati e non l'abbiamo verificata. Quello che abbiamo misurato è la riduzione delle dimensioni attive, non l'interpretabilità.

Posso cambiare β durante l'addestramento?

Sì: è il KL warm-up, si parte con β = 0 e si sale fino a 1 nelle prime epoche. Con decoder potenti serve a evitare il collasso descritto da Bowman e colleghi; nel nostro esperimento non ha dato benefici netti.

Se il collasso è benigno, perché preoccuparmene?

Perché dimensioni inattive possono segnalare che il modello non usa il codice come dovrebbe, per esempio con un decoder troppo potente. Contarle costa una riga.

Un posterior aggregato vicino al prior garantisce buoni campioni?

No: è necessario, non sufficiente. Il modello a β = 4 ha varianza media 1.00 e campioni tutti uguali.

Esercizi con soluzione

Prima di andare avanti, prova a risolverli con carta e matita; le soluzioni sono subito sotto ciascuno.

1. Nel caso giocattolo con s = 0.5 e β = 2, quanto vale la varianza ottima dell'encoder σ² = β/(β + 1/s²)?
   Soluzione: 1/s² = 4, quindi σ² = 2/(2 + 4) = 0.333.
2. Che cosa succede a σ² quando β → 0 e quando β → ∞?
   Soluzione: β → 0: σ² → 0 (codice esatto, come un autoencoder); β → ∞: σ² → 1 (il codice diventa il prior: nessuna informazione).
3. Una dimensione ha varianza di μ pari a 0.004 e KL media 0.002. È attiva?
   Soluzione: No (varianza < 0.01 e KL quasi nulla): è collassata.
4. Dalla tabella dell'esperimento 1, quale fra i modelli a β = 0.25 e β = 1 ha il −ELBO migliore sul test, e perché?
   Soluzione: β = 1: 21.9 + 2.9 = 24.8 nat contro 18.2 + 9.7 = 27.9. Il modello a β = 1 ottimizza proprio quella quantità; quello a β = 0.25 ne ottimizza un'altra.

I primi due esercizi usano la formula ricavata a mano, il terzo applica la regola di diagnosi, il quarto obbliga a leggere la tabella con la perdita giusta. Se sbagli il quarto, hai confrontato la ricostruzione e non il −ELBO.

Che cosa sai fare adesso

Sai usare β come manopola e sai riconoscere dove smette di funzionare: a β = 0 la KL vale 137 nat e i codici hanno varianza media 8.75, fuori scala rispetto al prior; a β = 4 il codice è vuoto e il decoder emette l'immagine media. Sai che con un codice troppo grande il VAE spegne da solo le dimensioni inutili (38 attive su 40), e che il controllo del posterior aggregato, da solo, non scopre un modello collassato. Sai che l'instabilità si manifesta come una perdita enorme (14224.7 a lr 0.2) più che come un NaN. Chiudiamo qui il modulo dei modelli generativi; nel prossimo cambiamo paradigma con l'apprendimento per rinforzo, l'agente che impara per tentativi.