Il limite continuo: una sola equazione, in avanti e all’indietro#
Nel 1982 Brian Anderson, ingegnere del controllo automatico, pubblica su una rivista di processi stocastici (cioè casuali) un articolo di quattordici pagine intitolato Reverse-time diffusion equation models [And82]. Il problema che aveva per le mani era di tutt’altro genere: dato un segnale rumoroso osservato fino a un certo istante, stimare che cosa fosse successo prima. Il risultato che ottiene è generale: un processo casuale che si diffonde in avanti nel tempo si può percorrere all’indietro, e l’equazione che descrive il ritorno è quella dell’andata più un termine correttivo.
Quel termine correttivo è il gradiente del logaritmo della densità dei dati rumorosi, \(\nabla_{\mathbf{x}}\log p_t(\mathbf{x})\), detto punteggio (score): in ogni punto indica la direzione in cui la densità cresce più in fretta, cioè da che parte stanno i dati più fitti, e corrisponde al «verso della salita» della mappa di colline. È la stessa quantità che la rete di DDPM stima, a meno di un fattore, quando predice il rumore.
Per quasi quarant’anni l’articolo di Anderson resta un risultato tecnico della teoria del filtraggio stocastico; poi nel 2020 qualcuno riconosce che rovinare un’immagine con del rumore è una diffusione in avanti, che generare è percorrerla all’indietro, e che quel termine correttivo è la sola cosa che una rete debba imparare [SSDK+21].
Passare dai mille passi di DDPM al tempo continuo porta un guadagno sostanziale e non estetico. Nel continuo, infatti, le due tradizioni della diffusione diventano due casi della stessa equazione; inoltre si scopre che accanto alla strada casuale ne esiste una deterministica che produce le stesse immagini, e le quattro cose che una rete può imparare si rivelano quattro modi di scrivere lo stesso oggetto.
Che cosa resta quando i passi diventano infiniti#
Finora di ricette per rovinare un’immagine ne abbiamo vista una, quella di DDPM, che a ogni passo restringe un po” l’immagine e ci aggiunge un pizzico di rumore. Ce n’è un’altra, nata nello stesso periodo, che l’immagine non la tocca e ci versa sopra rumore sempre più forte. Sembrano diverse; scritte una sotto l’altra, però, hanno la stessa forma: il nuovo valore è il vecchio, più uno spostamento sistematico, più uno scossone estratto a caso. Cambia solo quanto valgono i due pezzi.
Applichiamo ora quella ricetta un milione di volte invece di mille, con passi mille volte più corti. Al limite si ottiene una descrizione continua: invece di dire «a ogni passo fai così», si dice «in ogni istante il valore deriva un po” in questa direzione e trema un po” con questa intensità». Il primo pezzo si chiama deriva e dice dove il valore è trascinato; il secondo si chiama diffusione e dice quanto trema. Una regola così, che dice come una cosa cambia istante per istante, in matematica si chiama equazione differenziale: la sua soluzione non è un numero, come per le equazioni di scuola, ma un intero percorso.
Il tremore si comporta in modo controintuitivo. Gli scossoni non si accumulano come uno spostamento normale, ed è lo stesso fenomeno dei mille passi a caso, che portano a una trentina di metri e non a un chilometro. Ad esempio, cento scossoni da un centimetro allontanano di una decina di centimetri, non di un metro: chi va a destra e chi a sinistra si compensano parecchio, e quello che resta cresce come la radice quadrata del numero di scossoni (la radice di cento è dieci), non come il numero.
Ecco perché nella ricetta il tremore di un passo compare con la radice quadrata della sua durata, mentre la deriva compare con la durata intera. Spezzando lo stesso tratto di tempo in passetti sempre più corti, le derive si sommano fino a dare sempre lo stesso spostamento, e gli scossoni, che si sommano «a radice», danno sempre lo stesso tremore complessivo, né nullo né infinito.
È la stessa regola che fa stabilizzare la media di tante misure: la loro somma si sparpaglia come la radice del loro numero, e la media, che è la somma divisa per quel numero, si stringe come uno diviso la radice.
Il gioco vale la candela per una ragione pratica: nel continuo la ricetta si descrive con due sole funzioni del tempo, e tutte le varianti che negli anni sono state proposte si scrivono scegliendo quelle due. Invece di una famiglia di metodi diversi, una sola equazione con due manopole.
I due schemi di iniezione del rumore si scrivono, su una griglia temporale di passo \(\Delta t\), nella forma comune
con \(\mathbf{f}\) la deriva (un vettore, funzione dello stato e del tempo) e \(g\) il coefficiente di diffusione (uno scalare, funzione del solo tempo). La radice su \(\Delta t\) viene dalla natura del processo, non da una scelta di notazione. Gli incrementi di un processo di Wiener, infatti, hanno varianza proporzionale al tempo, quindi deviazione standard proporzionale alla sua radice.
Per \(\Delta t\to 0\) la ricorrenza converge all’equazione differenziale stocastica in avanti
dove \(\mathbf{w}(t)\) è un processo di Wiener standard: parte da zero, ha incrementi indipendenti, e \(\mathbf{w}(t)-\mathbf{w}(s)\) è gaussiano di media nulla e varianza \(t-s\). È continuo ovunque e derivabile in nessun punto, per cui la scrittura \(\mathrm{d}\mathbf{w}\) non si può dividere per \(\mathrm{d}t\).
La cosa che rende trattabile tutto il resto è la scelta di una deriva affine, \(\mathbf{f}(\mathbf{x},t) = f(t)\,\mathbf{x}\) con \(f\) scalare. Sotto questa ipotesi il processo condizionato al dato iniziale resta gaussiano a ogni istante, e il nucleo di perturbazione ha forma chiusa:
Le due formule risolvono le equazioni di media e varianza del processo lineare, \(\dot m = f(t)\,m\) e \(\dot v = 2f(t)\,v + g^2(t)\), con \(m(0) = \mathbf{x}_0\) e varianza per componente \(v(0) = 0\). Qui \(\alpha_t\) è quanto resta del dato di partenza e \(\sigma_t^2\) la varianza del rumore accumulato. Sono le lettere della letteratura sul tempo continuo, e nella sezione su come funziona la diffusione avevano un altro ruolo: là \(\alpha_t = 1-\beta_t\) era il fattore di un solo passo e \(\sigma_t\) la deviazione standard del rumore fresco del passo inverso. Per il processo di DDPM (della famiglia VP, quella che conserva la varianza) il dizionario è \(\alpha_t = \sqrt{\bar{\alpha}_t}\) e \(\sigma_t = \sqrt{1-\bar{\alpha}_t}\), e da qui in avanti, nel tempo continuo, valgono le lettere nuove.
La conseguenza operativa, che spiega perché questi modelli si addestrano in tempi umani, ha un nome: il campionamento è senza simulazione (simulation-free). Per avere \(\mathbf{x}_t\) a un istante qualsiasi non serve integrare l’equazione passo per passo; basta estrarre \(\boldsymbol{\epsilon}\) e scrivere \(\mathbf{x}_t = \alpha_t\mathbf{x}_0 + \sigma_t\boldsymbol{\epsilon}\).
Le due famiglie: chi lascia esplodere la varianza e chi la conserva#
Scelte la deriva e la diffusione, il processo è deciso. Le due tradizioni della diffusione corrispondono a due scelte, e portano nomi che dicono che cosa succede alla larghezza della nuvola dei dati.
La prima famiglia non tocca l’immagine e ci versa sopra rumore sempre più forte. Non c’è nessuna deriva: il dato resta dov’è, e attorno a esso la nuvola si allarga senza limite. Alla fine il rumore è così forte che dell’immagine non si distingue più niente, ma tecnicamente c’è ancora, sepolta. Il nome dice il comportamento: la varianza (la larghezza della nuvola, misurata al quadrato) esplode.
La seconda famiglia invece restringe l’immagine mentre aggiunge rumore, e lo fa con un dosaggio preciso: quanto toglie di segnale, tanto aggiunge di rumore, così che la larghezza complessiva della nuvola resti sempre più o meno uguale a uno. È la ricetta di DDPM, e il nome dice anche questo: la varianza si conserva.
La differenza pratica sta nella taglia dei numeri con cui la rete lavora. Nella prima famiglia i valori diventano enormi verso la fine, e la rete deve maneggiare grandezze che cambiano di ordini di grandezza; nella seconda restano sempre attorno a uno, che per una rete neurale è la condizione comoda. La seconda ce li tiene senza fare niente; chi usa la prima li riscala prima di darli alla rete, e fatto quel riscalamento il vantaggio si annulla.
Il fatto notevole è che sono la stessa equazione con due manopole diverse. Prima del 2020 erano due letterature separate, con due vocabolari e due insiemi di trucchi; nel continuo la separazione sparisce.
Variance Exploding (VE), la famiglia dei modelli score-based. Nessuna deriva, e la diffusione cresce con il programma di rumore scelto:
da cui \(\alpha_t = 1\) e \(\sigma_t = \sigma(t)\): il segnale resta intatto e il rumore cresce senza limite. La marginale terminale non tende a nessuna distribuzione fissa, e come prior si usa \(\mathcal{N}(\mathbf{0},\sigma_{\max}^2\mathbf{I})\) con \(\sigma_{\max}\) dell’ordine della massima distanza fra due dati.
Variance Preserving (VP), la famiglia di DDPM.
da cui
cioè \(\alpha_t^2+\sigma_t^2 = 1\) per ogni \(t\): se i dati hanno varianza unitaria la mantengono per tutta la traiettoria, e la marginale terminale tende a \(\mathcal{N}(\mathbf{0},\mathbf{I})\). Il legame con la notazione discreta della sezione su come funziona la diffusione è \(\alpha_t = \sqrt{\bar{\alpha}_t}\).
Esiste anche una sub-VP, che usa la stessa deriva con \(g^2(t)=\beta(t)\,(1-\alpha_t^4)\), da cui \(\sigma_t^2=(1-\alpha_t^2)^2\): la varianza resta sotto quella della VP a ogni istante, e nell’articolo che la introduce dà verosimiglianze migliori della VP.
Il rapporto fra le ampiezze \(\alpha_t/\sigma_t\) dice quanto segnale resta rispetto al rumore, e la letteratura chiama rapporto segnale/rumore il suo quadrato, il rapporto fra le potenze, \(\mathrm{SNR}(t) = \alpha_t^2/\sigma_t^2\). Le tre famiglie si distinguono soltanto per come lo fanno scendere e per come riscalano lo stato lungo la strada. Questa osservazione è il punto di partenza della riformulazione di Karras e colleghi [KAAL22], che tratta la scala come una riparametrizzazione libera e riscrive tutte le varianti in un unico spazio di progetto.
Verifichiamo che il limite continuo sia corretto. Bastano venti righe: si simula l’equazione a piccoli passi e si confronta la nuvola che ne esce con quella prevista dalla forma chiusa.
import numpy as np
rng = np.random.default_rng(0)
# programma di rumore VP, lo stesso di DDPM in versione continua
B_MIN, B_MAX, T = 0.1, 20.0, 1.0
beta = lambda t: B_MIN + t * (B_MAX - B_MIN)
alpha = lambda t: np.exp(-0.5 * (B_MIN * t + 0.5 * (B_MAX - B_MIN) * t**2))
sigma = lambda t: np.sqrt(1 - alpha(t)**2)
# dati: due sole possibilita', cosi' tutto il resto e' calcolabile a mano
MODI = np.array([-1.5, 1.5])
x0 = rng.choice(MODI, size=20000)
# 1) l'equazione simulata a piccoli passi
N = 1000
dt = T / N
x = x0.copy()
fotografie = {}
for k in range(N):
t = k * dt
x = (x - 0.5 * beta(t) * x * dt
+ np.sqrt(beta(t) * dt) * rng.normal(size=x.shape))
if k + 1 in (N // 4, N // 2, N):
fotografie[(k + 1) * dt] = x.copy()
# 2) la stessa cosa in forma chiusa, senza simulare niente
print(" t simulata forma chiusa")
for t, xs in fotografie.items():
chiusa = alpha(t) * x0 + sigma(t) * rng.normal(size=x0.shape)
print(f"{t:5.2f} {xs.mean():+.3f} {xs.std():.3f} "
f"{chiusa.mean():+.3f} {chiusa.std():.3f}")
# -> 0.25 +0.002 1.282 -0.001 1.286
# -> 0.50 -0.005 1.048 -0.015 1.051
# -> 1.00 +0.005 0.999 +0.008 0.995
Media e larghezza coincidono a ogni istante, e la seconda colonna si ottiene senza fare mille passi: una moltiplicazione e una somma. Si noti che la larghezza (deviazione standard) a \(t=0{,}25\) vale \(1{,}28\) e non \(1\), perché i dati di partenza hanno già una larghezza loro, essendo due punti a distanza tre.
L’equazione all’indietro, e il suo unico ingrediente ignoto#
Riavvolgere una diffusione sembra impossibile, e per una singola traiettoria lo è: il fumo che esce da un camino non rientra, e da una nuvola di fumo non si ricostruisce la scintilla che l’ha prodotta. Il risultato di Anderson dice una cosa più sottile e più utile: la singola traiettoria non si riavvolge, ma la distribuzione sì. Non si può sapere da quale scintilla venga quel fumo lì; si può però costruire un processo che, partendo da nuvole di fumo, produce scintille distribuite come quelle vere.
L’equazione all’indietro ha la stessa forma di quella in avanti, con un termine in più: in ogni punto, oltre alla deriva e al tremore, c’è una spinta che punta verso le zone dove i dati sono più densi. Quella spinta è il punteggio, il verso della salita sulla mappa di colline, ed è l’unico elemento ignoto dell’equazione. Deriva e tremore li abbiamo scelti noi quando abbiamo deciso come rovinare le immagini: sono nostri, quindi li conosciamo con precisione. Della strada del ritorno manca un pezzo solo, che la rete impara.
Andando all’indietro si continua ad aggiungere rumore. Sembra assurdo, visto che si sta cercando di ripulire, ma i due termini lavorano insieme. La spinta verso le zone dense tira verso i dati, il rumore permette di esplorare invece di precipitare sul primo posto buono, e il tremore si spegne man mano che si procede: all’inizio del ritorno è forte e si vaga, verso la fine è quasi nullo e comanda solo la spinta.
Il procedimento è parente di uno classico per estrarre punti da una distribuzione complicata, la dinamica di Langevin: si cammina verso le zone dense tremando un po”, e alla lunga i punti visitati si distribuiscono come si voleva.
Qui però la mappa delle zone dense non resta ferma. All’inizio del ritorno è larga e sfocata, perché i dati sono ancora coperti di rumore, e si fa più nitida man mano che il rumore cala, così che la folla si raccoglie verso i dati mentre la mappa si stringe attorno a loro.
Il teorema di Anderson afferma che il processo invertito nel tempo, associato alla SDE in avanti, è a sua volta una diffusione, retta da
integrata da \(t=T\) a \(t=0\) (quindi \(\mathrm{d}t<0\)), con \(\bar{\mathbf{w}}\) un processo di Wiener nel tempo invertito. Qui \(\nabla_{\mathbf{x}}\log p_t(\mathbf{x})\) è il punteggio (score) della marginale all’istante \(t\), cioè il gradiente rispetto a \(\mathbf{x}\) del logaritmo della densità dei dati rumorosi.
Il contenuto del teorema è forte e va enunciato con precisione: se la SDE all’indietro parte da \(\bar{\mathbf{x}}(T)\sim p_T\), la marginale vera dell’istante finale, e se \(\mathbf{f}\), \(g\) e \(p_t\) sono abbastanza regolari da dare alle due equazioni una soluzione unica, le marginali del processo invertito coincidono con quelle del processo in avanti a ogni istante, come uguaglianza e non come approssimazione asintotica. In pratica si parte da \(p_{\text{prior}}\), che coincide con \(p_T\) solo al limite, e si usa uno score stimato: sono questi due scarti, e l’errore di discretizzazione, a separare i campioni dai dati. Le singole traiettorie, invece, non si corrispondono: per questo il generatore non ricostruisce l’immagine da cui il rumore era partito.
L’ingrediente ignoto è uno solo. Scelti \(\mathbf{f}\) e \(g\), il resto della dinamica è determinato dal punteggio, e imparare il punteggio è il problema che la sezione su come funziona la diffusione ha già risolto in forma discreta. Infatti il minimo della loss quadratica sul rumore è la media condizionata \(\mathbb{E}[\boldsymbol{\epsilon}\mid\mathbf{x}_t]\), e questa è \(-\sigma_t\nabla\log p_t(\mathbf{x}_t)\).
La lettura come dinamica di Langevin. Ponendo \(\mathbf{f}\equiv\mathbf{0}\) e riparametrizzando il tempo con \(s := T-t\), la SDE all’indietro diventa
dove \(\pi_s := p_{T-s}\) è la densità bersaglio a quell’istante.
L’equazione ha la forma della dinamica di Langevin con cui i modelli a energia campionano, e le differenze sono due. Il bersaglio cambia a ogni istante invece di essere sempre lo stesso: \(\pi_s\) parte larga e si stringe verso i dati man mano che il rumore scende, e questo è l’annealing, quello dei livelli di rumore di Song ed Ermon e non quello di una temperatura. La deriva, inoltre, è il doppio di quella di Langevin, che a parità di diffusione vuole \(\tau\nabla\log\pi\); la differenza è sostanziale: con il bersaglio congelato questa equazione avrebbe per stazionaria \(\pi^2\) e non \(\pi\), e la Langevin vera, dove serve, si aggiunge a parte come correttore.
Il fattore \(\tau(s)\), invece, di una temperatura ha soltanto la lettera: moltiplica insieme la deriva e la varianza del rumore, quindi ne lascia fisso il rapporto, che è ciò che in Langevin decide la temperatura, e cambia soltanto la velocità dell’orologio.
In definitiva, quello che nei modelli a energia è una catena lunghissima da far mescolare qui diventa un percorso guidato di durata prefissata.
La strada deterministica, quella che non estrae niente a caso#
Fin qui il ritorno è casuale. La domanda che chiude il quadro è se il caso serva davvero.
Una folla è sparpagliata in una piazza, e ci sono due modi di riportarla a casa. Il primo dà a ciascuno una spinta verso casa e lo lascia anche barcollare un po”; il secondo assegna a ogni punto della piazza una direzione precisa e chiede a tutti di seguirla senza sbandare. I percorsi individuali sono diversissimi, e tuttavia si può fare in modo che in ogni istante la folla sia sparpagliata allo stesso modo: stessa densità in ogni punto della piazza, a ogni ora.
Questo secondo modo esiste, e si costruisce a partire dal primo togliendo il tremore e dimezzando la spinta: metà è esattamente quanto serve a rimpiazzare il rimescolamento che si è tolto. Le conseguenze sono tre.
Si può percorrere nei due sensi. Senza caso, il percorso è una linea tracciata: dai dati al rumore e dal rumore ai dati si passa lungo la stessa strada. Ogni immagine ha quindi il suo rumore, quello da cui si ottiene, e lo si può trovare percorrendo la strada all’incontrario. È così che si modifica un’immagine esistente invece di generarne una da zero: la si porta indietro fino al suo rumore, si cambia qualcosa nella richiesta, e la si riporta avanti.
Si può dire quanto è probabile un’immagine. Lungo una strada tracciata si può tenere traccia di quanto lo spazio si allarga e si stringe, e quella quantità è precisamente il determinante, cioè quanto una trasformazione gonfia lo spazio. Serve a dare un voto al modello intero: davanti a fotografie vere che non ha mai visto, il modello migliore è quello che le giudica più probabili. Con la strada casuale quel calcolo non si può fare.
Servono meno passi. Le traiettorie senza tremore sono lisce, e una curva liscia si approssima bene con pochi segmenti. Le traiettorie casuali sono frastagliate e chiedono passi corti. Quasi tutti i generatori veloci di oggi percorrono la strada deterministica, e la sezione sui campionatori veloci racconta come.
Il limite si vede quando la spinta non è quella esatta ma quella che la rete ha imparato, cioè sempre.
Sulla strada tracciata, infatti, un errore commesso a metà percorso resta lì e si porta fino in fondo; su quella casuale il barcollamento rimescola a ogni passo, e ogni rimescolamento riporta la folla un po” più vicino alla densità che dovrebbe avere in quel momento, cancellando in parte gli errori vecchi. Per questo, con una spinta imperfetta, la strada casuale arriva spesso più vicino alla folla vera, e chi vuole il massimo della qualità usa un misto: un po” di caso in un tratto del percorso, nessuno alla fine.
Con la spinta esatta, invece, le due strade portano alla stessa folla, e il barcollamento obbliga soltanto a fare più passi: con passi non abbastanza piccoli la casuale arriva un po” più sparpagliata attorno alla meta, perché trema fino all’ultimo. La prova con i due soli valori \(\pm 1{,}5\) fa vedere questo.
Accanto alla SDE all’indietro esiste un’equazione deterministica con le stesse marginali a ogni istante, la ODE del flusso di probabilità (probability flow ODE, PF-ODE):
Il confronto con la SDE all’indietro mostra le due sole differenze: il termine di rumore è sparito, e il coefficiente del punteggio è dimezzato. La ragione sta nell’equazione che governa la densità: quella della SDE in avanti è l’equazione di Fokker-Planck,
e poiché \(\Delta p_t = \nabla\!\cdot\!\big(p_t\nabla\log p_t\big)\) la si può riscrivere come
che è l’equazione di continuità di un flusso deterministico con velocità \(\mathbf{f} - \tfrac12 g^2\nabla\log p_t\): la PF-ODE, che quindi ha le stesse marginali. Il fattore \(\tfrac12\) è esattamente quanto serve a far rientrare nel trasporto il termine diffusivo che si è tolto. La stessa riscrittura, spezzando \(\tfrac12 g^2\Delta p_t\) in \(g^2\Delta p_t - \tfrac12 g^2\Delta p_t\) e invertendo il tempo, dà la SDE all’indietro di Anderson [SSDK+21].
Ne seguono tre conseguenze:
Invertibilità. La PF-ODE definisce un flusso deterministico \(\Phi_{t\to s}\) che, se il campo è lipschitziano in \(\mathbf{x}\) (e quindi la soluzione è unica), è una biiezione fra gli stati e trasporta \(p_{\text{dati}}\) in \(p_T\). Da qui l’inversione DDIM, cioè integrare la ODE dall’immagine al suo rumore, che serve a editing e interpolazione nello spazio del rumore. Per un campo stimato da una rete la condizione vale; per il punteggio esatto può cadere vicino a \(t = 0\), dove i dati si concentrano.
Verosimiglianza esatta. Trattando la PF-ODE come un flusso normalizzante continuo si ottiene \(\log p_0(\mathbf{x}_0) = \log p_T(\mathbf{x}_T) + \int_0^T \nabla\!\cdot\!\mathbf{v}_t(\mathbf{x}_t)\,\mathrm{d}t\), con \(\mathbf{v}_t\) il campo di velocità della ODE e la divergenza stimata a caso con lo stimatore di Hutchinson [Hut89], come fa FFJORD nella stessa sezione sui flussi continui. È un calcolo che la SDE non permette.
Errore di discretizzazione. Un integratore di ordine \(p\) su una traiettoria liscia accumula errore \(O(h^p)\) con \(h\) il passo; le traiettorie della SDE hanno regolarità di Hölder di esponente inferiore a \(1/2\) e i metodi stocastici si fermano a ordini bassi. È la ragione strutturale per cui i campionatori a pochi passi lavorano sulla ODE.
Esiste però un rovescio, misurato in letteratura. Con il punteggio esatto SDE e ODE hanno le stesse marginali, quindi la stessa qualità e la stessa varietà; con un punteggio stimato la SDE dà spesso campioni migliori, perché la componente di Langevin che contiene riporta a ogni passo il campione verso la marginale giusta e corregge in parte gli errori dei passi precedenti, mentre la ODE li integra fedelmente [KAAL22].
Il vantaggio ha due limiti. Il primo è che discretizzare la parte stocastica introduce un errore suo, e troppo rumore aggiunto e tolto impoverisce i dettagli (tanto che Karras e colleghi accendono la stocasticità soltanto in una fascia di livelli di rumore e ne tarano la dose caso per caso). Il secondo è che il guadagno si riduce al migliorare del modello: su CIFAR-10, con il loro addestramento migliore, qualunque dose di rumore peggiora il risultato, mentre su ImageNet \(64\times64\) resta utile. Da qui i campionatori ibridi, che iniettano rumore in un tratto del percorso e seguono la sola ODE nel resto.
Le due strade e la loro equivalenza si controllano in poche righe, e senza addestrare niente: su dati che possono valere soltanto \(-1{,}5\) o \(+1{,}5\) (i due modi della distribuzione, cioè i due valori attorno a cui i dati si ammucchiano) il punteggio esatto si scrive a mano, quindi si può isolare la matematica dalla rete.
def punteggio(x, t):
"""Il punteggio esatto della mistura, senza nessuna rete."""
a, s = alpha(t), sigma(t)
d = x[:, None] - a * MODI[None, :]
w = np.exp(-0.5 * (d / s)**2)
w /= w.sum(axis=1, keepdims=True)
return -(w * d).sum(axis=1) / s**2
n_camp, M = 5000, 500
dt = T / M
# la strada deterministica: dal rumore ai dati, senza sorteggiare niente
y = rng.normal(size=n_camp)
for k in range(M):
t = T - k * dt
y = y + dt * 0.5 * beta(t) * (y + punteggio(y, t))
# la strada casuale: la stessa cosa con il tremore acceso
z = rng.normal(size=n_camp)
for k in range(M):
t = T - k * dt
z = (z + dt * (0.5 * beta(t) * z + beta(t) * punteggio(z, t))
+ np.sqrt(beta(t) * dt) * rng.normal(size=n_camp))
for nome, v in (("ODE", y), ("SDE", z)):
print(nome, round(float((v < 0).mean()), 4),
round(float(v[v < 0].mean()), 4), round(float(v[v > 0].mean()), 4),
round(float(np.abs(np.abs(v) - 1.5).mean()), 4))
# -> ODE 0.501 -1.5003 1.5 0.0055
# -> SDE 0.4998 -1.5001 1.5005 0.0135
Le due strade arrivano nello stesso posto: metà dei campioni su ciascun modo, e i modi a \(\pm 1{,}5\). L’ultima colonna è di quanto in media un campione manca il proprio modo, e dice la differenza fra le due: la strada casuale lascia una dispersione due volte e mezzo più larga, perché continua ad aggiungere rumore fino all’ultimo passo.
La strada deterministica è davvero una strada, cioè si percorre nei due sensi.
ts = np.linspace(1e-3, T, M + 1)
campo = lambda x, t: -0.5 * beta(t) * (x + punteggio(x, t))
partenza = rng.choice(MODI, size=6) + 0.01 * rng.normal(size=6)
x = partenza.copy()
for k in range(M): # dai dati al rumore
x = x + (ts[k + 1] - ts[k]) * campo(x, ts[k])
rumore = x.copy()
for k in range(M, 0, -1): # e ritorno, sulla stessa strada
x = x - (ts[k] - ts[k - 1]) * campo(x, ts[k])
print(np.round(rumore, 3))
# -> [-0.354 2.779 0.572 -0.04 0.562 -0.623]
print(round(float(np.abs(x - partenza).max()), 4)) # -> 0.0077
Sei immagini portate al rumore e riportate indietro tornano dove erano, a meno di otto millesimi. Quel residuo è l’errore di chi integra a passi finiti, non un difetto della strada: è la grandezza che i campionatori veloci cercano di ridurre a parità di passi.
Quattro modi di dire la stessa cosa#
Chi legge il codice di più di una libreria incontra reti che predicono cose apparentemente diverse. Predicono la stessa informazione, scritta in quattro modi: da una qualsiasi delle quattro uscite le altre tre si ricavano per conversione algebrica, e la scelta cambia il condizionamento numerico e il peso che la loss dà ai livelli di rumore, mentre ciò che c’è da imparare resta lo stesso.
Data un’immagine rovinata, ci sono quattro domande che si possono fare, e rispondere a una qualsiasi permette di rispondere a tutte le altre.
Qual era il disturbo? È la domanda di DDPM.
Qual era l’immagine pulita? È la domanda che sembra più naturale.
In che direzione salire? È il punteggio.
Con quale velocità muoversi? Se si guarda l’immagine rovinata come un punto in viaggio fra l’immagine pulita e il rumore, la velocità dice da che parte e quanto in fretta quel punto si sta spostando. Si ottiene mescolando le prime due risposte, e su di essa è costruito il flow matching.
Sono quattro forme della stessa informazione perché l’immagine rovinata è fatta di due pezzi, l’immagine pulita e il disturbo, mescolati in proporzioni note. Sapendo l’immagine rovinata e uno dei due pezzi, l’altro si ricava sottraendo.
La scelta lascia invariata la matematica e cambia la precisione dei calcoli. Chiedere il disturbo funziona male quando di disturbo ce n’è pochissimo, perché di quel disturbo nell’immagine resta una traccia minuscola e va indovinato tutto a partire da lì; chiedere l’immagine pulita funziona male dall’altra parte, dove di immagine non c’è quasi più niente. La quarta domanda, la velocità, nasce per questo: mescola le due in modo da restare affidabile a tutti e due gli estremi, e si usa quando il percorso deve essere accorciato a pochi passi.
Con \(\mathbf{x}_t = \alpha_t\mathbf{x}_0 + \sigma_t\boldsymbol{\epsilon}\), le quattro parametrizzazioni sono legate da identità esatte:
dove \(\hat{\boldsymbol{\epsilon}}\) è la predizione del rumore, \(\hat{\mathbf{x}}_0\) quella del dato pulito, \(\mathbf{s}_\theta\) il punteggio e \(\hat{\mathbf{v}}\) la velocità. Una rete addestrata su una qualsiasi delle quattro si converte nelle altre senza riaddestramento, e le librerie infatti espongono un parametro di configurazione per sceglierla.
L’equivalenza è algebrica, il condizionamento numerico no. Per \(t\to 0\) si ha \(\sigma_t\to 0\): la conversione \(\hat{\boldsymbol{\epsilon}}\mapsto \mathbf{s}\) divide per \(\sigma_t\) e amplifica l’errore, mentre \(\hat{\mathbf{x}}_0\) è ben posta. Per \(t\to T\) è il contrario, perché \(\alpha_t\to 0\) e ricostruire \(\hat{\mathbf{x}}_0\) significa dividere per qualcosa che tende a zero. La \(\mathbf{v}\)-prediction [SH22] interpola fra le due e resta ben condizionata a entrambi gli estremi; è la parametrizzazione standard nella distillazione e nei programmi di rumore che portano il rapporto segnale/rumore a zero all’ultimo passo.
Le funzioni di costo corrispondenti differiscono solo per un peso dipendente da \(t\):
e ogni scelta di parametrizzazione equivale a una scelta di \(w(t)\) nella formulazione sul rumore. Con il rapporto segnale/rumore \(\mathrm{SNR}(t) = \alpha_t^2/\sigma_t^2\), fra le potenze, il confronto sta in una riga: \(\lVert\hat{\mathbf{x}}_0 - \mathbf{x}_0\rVert^2 = \lVert\hat{\boldsymbol{\epsilon}} - \boldsymbol{\epsilon}\rVert^2/\mathrm{SNR}(t)\), \(\lVert\mathbf{s}_\theta - \nabla\log p_t(\mathbf{x}_t \mid \mathbf{x}_0)\rVert^2 = \lVert\hat{\boldsymbol{\epsilon}} - \boldsymbol{\epsilon}\rVert^2/\sigma_t^2\) e, sul percorso VP, \(\lVert\hat{\mathbf{v}} - \mathbf{v}\rVert^2 = \lVert\hat{\boldsymbol{\epsilon}} - \boldsymbol{\epsilon}\rVert^2/\alpha_t^2 = \big(1 + 1/\mathrm{SNR}(t)\big)\lVert\hat{\boldsymbol{\epsilon}} - \boldsymbol{\epsilon}\rVert^2\).
La lettura per pesi è di Kingma e colleghi, che nel 2021 scrivono il bound variazionale in funzione del solo rapporto segnale/rumore [KSPH21] e più tardi mostrano che ogni peso monotono in \(t\) equivale a un ELBO su dati aumentati con rumore [KG23]. Questa prospettiva spiega perché confronti fra articoli che dichiaravano obiettivi diversi risultassero poi difficili da interpretare: cambiava il peso, non l’obiettivo.
In pratica#
# la conversione fra le quattro parametrizzazioni, che nessuna libreria
# nasconde ma tutte chiamano in modo diverso
def converti(eps, x_t, t):
a, s = alpha(t), sigma(t)
x0 = (x_t - s * eps) / a
score = -eps / s
v = a * eps - s * x0
return x0, score, v
t = 0.5
x_t = alpha(t) * MODI[0] + sigma(t) * 0.3 # un esempio costruito a mano
eps_vero = 0.3
x0, score, v = converti(eps_vero, x_t, t)
print(round(float(x0), 6), round(float(MODI[0]), 6)) # -> -1.5 -1.5
print(round(float(v), 6),
round(float(alpha(t) * eps_vero - sigma(t) * MODI[0]), 6))
# -> 1.523836 1.523836
La prima riga è il controllo che conta: partendo dal rumore vero e dall’immagine rovinata si ricostruisce l’immagine pulita, e viene quella di partenza.
Da ricordare
Le due ricette per rovinare un’immagine hanno la stessa forma (valore vecchio, più uno spostamento sistematico, più uno scossone estratto a caso) e nel limite dei passi infinitamente corti diventano una sola equazione con due manopole: la deriva, che dice dove il valore è trascinato, e la diffusione, che dice quanto trema.
Le due famiglie storiche sono due scelte di quelle manopole: una lascia l’immagine ferma e fa esplodere il rumore, l’altra restringe l’immagine mentre aggiunge rumore così che la larghezza resti sempre attorno a uno. La seconda tiene i numeri in una scala comoda senza doverli riscalare.
La strada del ritorno esiste ed è la stessa equazione con un termine in più, la spinta verso le zone dense. Quella spinta è l’unica cosa ignota, e la rete impara solo quella. Andando all’indietro il tremore c’è ancora ma si spegne piano piano, mentre il bersaglio si stringe attorno ai dati.
Accanto alla strada casuale ce n’è una deterministica che attraversa le stesse nuvole negli stessi istanti. Si percorre nei due sensi (quindi ogni immagine ha il suo rumore, e si può modificare un’immagine esistente), permette di dire quanto è probabile un’immagine, e si percorre con meno passi. In cambio non corregge per strada gli errori della rete: con una spinta imperfetta, la strada casuale arriva spesso più vicino alla folla vera.
Le quattro cose che una rete può predire (il disturbo, l’immagine pulita, la direzione di salita, la velocità) sono la stessa informazione: sapendone una si ricavano le altre tre. Cambia solo dove i calcoli restano precisi.
Da ricordare
Il limite \(\Delta t\to 0\) dà \(\mathrm{d}\mathbf{x} = \mathbf{f}(\mathbf{x},t)\mathrm{d}t + g(t)\mathrm{d}\mathbf{w}\). Con deriva affine il nucleo è gaussiano in forma chiusa, \(p_t(\mathbf{x}_t\mid\mathbf{x}_0)=\mathcal{N}(\alpha_t\mathbf{x}_0, \sigma_t^2\mathbf{I})\), e il campionamento è simulation-free.
VE: \(\mathbf{f}=\mathbf{0}\), \(g=\sqrt{\mathrm{d}\sigma^2/\mathrm{d}t}\). VP: \(\mathbf{f}=-\tfrac12\beta\mathbf{x}\), \(g=\sqrt{\beta}\), con \(\alpha_t^2+\sigma_t^2=1\). Le due differiscono per come fanno scendere il rapporto segnale/rumore e per il riscalamento.
Anderson (1982): \(\mathrm{d}\bar{\mathbf{x}} = [\mathbf{f} - g^2\nabla\log p_t]\mathrm{d}t + g\,\mathrm{d}\bar{\mathbf{w}}\), con marginali identiche a quelle del processo in avanti. Il punteggio è l’unico termine ignoto; con \(\mathbf{f}=\mathbf{0}\) l’equazione somiglia a una dinamica di Langevin verso un bersaglio che si stringe, con deriva doppia e un orologio \(\tau(s)=\tfrac12 g^2(T-s)\) che non è una temperatura.
PF-ODE: \(\dot{\tilde{\mathbf{x}}} = \mathbf{f} - \tfrac12 g^2\nabla\log p_t\), stesse marginali, traiettorie deterministiche. Dà invertibilità (con un campo lipschitziano), verosimiglianza esatta con lo stimatore di Hutchinson e traiettorie lisce (quindi integratori di ordine alto). Con un punteggio stimato la SDE dà spesso campioni migliori, perché la stocasticità corregge in parte gli errori dei passi precedenti; il vantaggio si riduce al migliorare del modello.
Le quattro parametrizzazioni sono legate da \(\hat{\mathbf{x}}_0 = (\mathbf{x}_t-\sigma_t\hat{\boldsymbol{\epsilon}})/ \alpha_t\), \(\mathbf{s}=-\hat{\boldsymbol{\epsilon}}/\sigma_t\), \(\hat{\mathbf{v}}=\alpha_t\hat{\boldsymbol{\epsilon}}-\sigma_t \hat{\mathbf{x}}_0\), e differiscono solo per il peso \(w(t)\) nella loss. La \(\mathbf{v}\)-prediction è ben condizionata a entrambi gli estremi, dove le altre due degenerano.
La diffusione ha smesso di essere una ricetta e ha una struttura: si sceglie come rovinare, il ritorno è determinato, e l’unica cosa da imparare è una funzione.
Resta però una domanda che l’impianto lascia aperta. Il percorso è stato costruito rovinando i dati con rumore gaussiano, secondo un dosaggio che cambia istante per istante, e quella scelta fissa la forma delle traiettorie della ODE, che in generale sono curve; e una curva, per seguirla, chiede molti passi. C’è un modo di scegliere il percorso invece di ereditarlo, e magari di sceglierlo dritto?