Riepilogo Tecnico: Calibrazione dei Modelli Generativi a Vincoli Distribuzionali

Enunciato del Problema

I modelli generativi soffrono frequentemente di miscalibrazione, dove le statistiche della distribuzione di campionamento (ad esempio, frequenze di classe, proporzioni di strutture secondarie o rappresentazione di genere) si discostano sistematicamente dai valori desiderati. Questo problema si manifesta in diversi domini: i modelli di immagine mostrano un collasso delle modalità; i modelli linguistici rafforzano i pregiudizi sociali; e i modelli di proteine/DNA generano strutture con frequenze di sottosezioni atipiche. Questi errori derivano spesso da squilibri nei dataset, dinamiche di training subottimali o aggiustamenti post-hoc come il campionamento a bassa temperatura.

Gli autori inquadrano la calibrazione come un problema di ottimizzazione vincolata. Dato un modello generativo di base $p_{\theta_{base}}(x)$, una statistica $h(x)$ e un'aspettativa target $h^*$, l'obiettivo è trovare una nuova distribuzione $p_\theta$ che soddisfi $E_{p_\theta}[h(x)] = h^*$ minimizzando al contempo la divergenza di Kullback-Leibler (KL) rispetto al modello di base:
$$p_{\theta^*} = \arg \min_{p_\theta} D_{KL}(p_\theta \parallel p_{\theta_{base}}) \quad \text{s.t.} \quad E_{p_\theta}[h(x)] = h^*$$
Questa formulazione cerca il modello più vicino alla distribuzione di base che soddisfi il vincolo di calibrazione. Tuttavia, risolvere esattamente questo problema è intrattabile perché sia l'obiettivo che il vincolo coinvolgono aspettative difficili da calcolare per modelli generativi ad alta dimensionalità.

Metodologia

Per affrontare l'intrattabilità, il paper introduce due obiettivi surrogati per il fine-tuning, denominati CGM-relax e CGM-reward. Entrambi i metodi si basano sull'ottimizzazione stocastica utilizzando stime di gradiente non distorte derivate dai campioni.

1. CGM-relax (Perdita di Rilassamento)

Questo approccio sostituisce il vincolo rigido con un termine di penalità. La funzione obiettivo è:
$$L_{relax}(\theta) = \| E_{p_\theta}[h(x)] - h^* \|^2 + \lambda D_{KL}(p_\theta \parallel p_{\theta_{base}})$$

• Violazione del Vincolo ($L_{viol}$): La norma quadrata $\ell_2$ della differenza tra la statistica empirica e l'obiettivo. Gli autori derivano uno stimatore non distorto per questo termine utilizzando una formula di varianza campionaria corretta per il bias.
• Divergenza KL ($L_{KL}$): Misurata tramite campionamento Monte Carlo.
• Ottimizzazione: Il metodo utilizza il campionamento per importanza per calcolare gradienti non distorti. Nello specifico, campiona dal modello corrente $p_\theta$ (con stop-gradient) e ripesa i campioni per stimare i gradienti del termine KL e della violazione del vincolo. Un iperparametro $\lambda$ bilancia tra il soddisfare il vincolo e rimanere vicino al modello di base.

2. CGM-reward (Perdita di Ricompensa)

Questo approccio collega il problema di calibrazione al principio di Massima Entropia. La soluzione al problema di massima entropia (minimizzare la KL rispetto al modello di base soggetto a vincoli di momento) è una inclinazione esponenziale del modello di base: $p_{\alpha^*}(x) \propto p_{\theta_{base}}(x) \exp(\alpha^{*\top} h(x))$.

• Stima di $\alpha^*$: Poiché la vera distribuzione di massima entropia è intrattabile, il metodo stima prima i parametri di ricompensa ottimali $\alpha^*$ risolvendo un problema di ottimizzazione duale utilizzando $N$ campioni dal modello di base.
• Funzione di Perdita: L'obiettivo del fine-tuning è la divergenza KL tra il modello corrente $p_\theta$ e la distribuzione di massima entropia stimata $p_{\hat{\alpha}_N}$:
  $$L_{reward}(\theta) = D_{KL}(p_\theta \parallel p_{\hat{\alpha}_N}) = D_{KL}(p_\theta \parallel p_{\theta_{base}}) + E_{p_\theta}[-r_{\hat{\alpha}_N}(x)] + C$$
  Questo equivale a un problema di fine-tuning con ricompensa dove la "ricompensa" è derivata dai vincoli.
• Ottimizzazione: Simile a CGM-relax, utilizza stimatori di gradiente basati sulla funzione di punteggio con baseline leave-one-out per ridurre la varianza.

Stima del Gradiente

Entrambi i metodi impiegano stimatori di gradiente basati sulla funzione di punteggio (stile REINFORCE) combinati con il campionamento per importanza. Per gestire la dipendenza dell'aspettativa da $\theta$, gli autori utilizzano un trucco "stop-gradient" dove i campioni sono estratti da una versione fissa del modello corrente e i gradienti sono calcolati tramite il rapporto di densità (funzione di punteggio). La riduzione della varianza è ottenuta utilizzando baseline leave-one-out.

Contributi Chiave

1. Formulazione: Inquadrare la calibrazione dei modelli generativi come un problema di minimizzazione KL vincolata.
2. Algoritmi: Proporre due algoritmi pratici di fine-tuning (CGM-relax e CGM-reward) che approssimano la soluzione del problema vincolato utilizzando ottimizzazione stocastica.
3. Scalabilità: Dimostrare che questi metodi possono gestire centinaia di vincoli simultanei e modelli con fino a 9 miliardi di parametri.
4. Connessione Teorica: Stabilire il legame tra calibrazione, massima entropia e fine-tuning con ricompensa, mostrando che CGM-reward mira alla soluzione di massima entropia mentre CGM-relax mira a una versione regolarizzata.

Risultati

Gli autori valutano i metodi su tre applicazioni diverse:

1. Progettazione di Proteine (Genie2 e ESM3):

   • Obiettivo: Calibrare il contenuto di struttura secondaria (eliche alfa e foglietti beta) per corrispondere alle statistiche delle proteine naturali (dataset CATH).
   • Vincoli: Fino a 99 vincoli simultanei (CDF bivariati).
   • Esito: CGM-relax ha ridotto la distanza rispetto alle distribuzioni naturali delle proteine di oltre 4 volte per Genie2 e di oltre 2 volte per ESM3. CGM-reward ha ottenuto miglioramenti più modesti, in parte a causa della difficoltà di stimare $\alpha^*$ con vincoli ad alta dimensionalità. Entrambi i metodi hanno mantenuto la plausibilità biofisica, sebbene con un leggero aumento dei "fallimenti di progettazione".

2. Generazione di Immagini (TarFlow):

   • Obiettivo: Bilanciare le proporzioni di classe all'interno della categoria "fauna selvatica" del dataset AFHQ (ad esempio, equalizzare leoni, tigri, volpi, lupi).
   • Vincoli: 5 vincoli simultanei.
   • Esito: CGM-relax ha ridotto la distanza di variazione totale delle proporzioni di classe da 0,306 a 0,101. Sebbene il realismo del campione (FID) sia diminuito leggermente, il modello calibrato ha prodotto generazioni più diversificate ed equilibrate. CGM-reward non è riuscito a ridurre significativamente la violazione del vincolo.

3. Modellazione Linguistica (Gemma-2-9B-IT):

   • Obiettivo: Ridurre il pregiudizio di genere nelle storie per bambini generate per varie professioni (ad esempio, assicurando una probabilità uguale di avvocati maschi/femmine).
   • Vincoli: 8 vincoli condizionati alla professione.
   • Esito: Entrambi i metodi hanno ridotto lo squilibrio di genere. CGM-relax ha ridotto lo squilibrio di oltre tre volte rispetto alla baseline, mentre CGM-reward ha mostrato un miglioramento statisticamente significativo ma più piccolo. I modelli fine-tunati hanno mantenuto la qualità della storia (nessun degrado nella log-probabilità).

Esperimenti Sintetici:

• Eventi Rari: CGM ha aumentato con successo il peso delle modalità rare nelle miscele gaussiane, anche quando la probabilità di base era bassa quanto $10^{-3}$.
• Alta Dimensionalità: CGM-relax ha gestito efficacemente fino a $10^3$ vincoli. Al contrario, CGM-reward è diventato non fattibile per vincoli $>30$ a causa della difficoltà di risolvere il problema empirico di massima entropia (Equazione 7).

Significato e Affermazioni

Il paper afferma che CGM fornisce un framework pratico per calibrare modelli generativi ad alta dimensionalità in modo da soddisfare vincoli distribuzionali senza richiedere modifiche all'architettura del modello o un addestramento da zero.

• Robustezza: CGM-relax si è dimostrato più robusto di CGM-reward, specialmente in contesti di vincoli ad alta dimensionalità dove la stima dei parametri di massima entropia diventa intrattabile.
• Compromessi: I metodi offrono un compromesso navigabile tra accuratezza della calibrazione e fedeltà al modello di base (controllato da $\lambda$).
• Limitazioni: Gli autori riconoscono che rimangono errori residui, specialmente per eventi molto rari. Inoltre, il framework dipende attualmente da modelli con verosimiglianze trattabili (modelli di diffusione, flussi normalizzanti, modelli autoregressivi), limitandone l'applicazione diretta a modelli impliciti come i GAN.

Il lavoro posiziona la calibrazione come una direzione di ricerca distinta e fruttuosa, separata dal normale fine-tuning con ricompensa, utilizzando i vincoli stessi come segnale trainante per l'allineamento distribuzionale.