Calibrating Generative Models to Feature Distributions with MMD Finetuning

Sintesi Tecnica: Calibrazione di Modelli Generativi alle Distribuzioni di Caratteristiche con Fine-tuning MMD (kCGM)

Definizione del Problema
I modelli generativi spesso producono campioni che sono individualmente plausibili ma esibiscono un disallineamento distributivo rispetto a chiavi caratteristiche di interesse. Ad esempio, un modello preaddestrato su uno spazio chimico ampio può generare molecole che differiscono significamente nella distribuzione delle caratteristiche rispetto a una specifica classe terapeutica, come gli antibiotici. Correggere questo fenomeno tramite il fine-tuning diretto sul set target è problematico: rischia l'overfitting su piccoli dataset, richiede l'intero set di campioni target (che può essere scarso) e non riesce a specificare quali caratteristiche specifiche debbano essere corrispondute. Inoltre, i metodi esistenti spesso faticano con caratteristiche a valori discreti o "black-box" (ad esempio, i fingerprint di Morgan) che sono non differenziabili. Il documento identifica una lacuna negli algoritmi capaci di calibrare i modelli generativi per far corrispondere una distribuzione di caratteristiche target utilizzando solo la supervisione a livello di caratteristica, senza richiedere la differenziabilità della mappa delle caratteristiche o della procedura di campionamento.

Metodologia: kCGM
Gli autori introducono i kernel Calibrating Generative Models (kCGM), un metodo di fine-tuning che minimizza la Maximum Mean Discrepancy (MMD) tra la distribuzione delle caratteristiche generate e una distribuzione di caratteristiche target, regolarizzata dalla divergenza di Kullback-Leibler (KL) rispetto al modello preaddestrato.

• Funzione Obiettivo: Il metodo ottimizza un obiettivo $L_\theta$ composto da due termini:

  1. Perdita MMD: Un stimatore non distorto (unbiased) a campioni finiti del quadrato della MMD tra la distribuzione delle caratteristiche indotta dal modello $P$ (dove $y = h(x)$ e $x \sim p_\theta$) e la distribuzione empirica target $Q$ (derivata dalle caratteristiche target $y^{targ}_{1:N}$).
  2. Regolarizzazione KL: Un termine $\lambda D_{KL}(p_\theta || p_{\theta_{pre}})$ per garantire che il modello sottoposto a fine-tuning rimanga vicino alla distribuzione preaddestrata, prevenendo il dimenticanza catastrofica (catastrophic forgetting) o il collasso del modo (mode collapse).

• Selezione del Kernel: A differenza degli approcci precedenti che potrebbero fare affidamento su kernel dot-product (che corrispondono solo alle medie), kCGM utilizza kernel che definiscono una metrica sulle distribuzioni di probabilità (zero solo per distribuzioni identiche).

  • Per caratteristiche euclidee continue, viene utilizzato un kernel energy-distance.
  • Per caratteristiche a vettori di bit discreti (es. fingerprint di Morgan), viene impiegato un kernel Tanimoto/Jaccard.
  • Queste scelte evitano l'introduzione di ulteriori iperparametri di larghezza di banda (bandwidth).

• Stima del Gradiente: Poiché la mappa delle caratteristiche $h$ e il processo di campionamento possono essere non differenziabili, kCGM impiega gradienti della funzione di punteggio (stile REINFORCE).

  • Baseline Leave-One-Out (LOO): Un contributo tecnico critico è la derivazione di una nuova baseline LOO per la perdita MMD. Le baseline LOO standard mediano i coefficienti di altri campioni in un batch, ma i coefficienti MMD sono accoppiati tra i campioni. Gli autori derivano uno stimatore specifico "MMD-LOO" che approssima il coefficiente medio che si sarebbe avuto se un campione specifico non fosse apparso nel mini-batch, riducendo la varianza senza introdurre distorsione (bias).
  • Algoritmo: Il processo prevede il campionamento dal modello corrente (senza tracciamento del gradiente), il calcolo dei vettori delle caratteristiche, il calcolo dei coefficienti MMD e KL con le baseline LOO, e l'aggiornamento dei pesi del modello tramite gradienti della funzione di punteggio.

Contributi Chiave

1. Formulazione: Il documento formula il matching della distribuzione delle caratteristiche come fine-tuning MMD regolarizzato da KL utilizzando solo campioni di caratteristiche target, eliminando la necessità di campioni target completi nello spazio originale dei campioni.
2. Stimatore Non Distorto: Gli autori derivano uno stimatore del gradiente della funzione di punteggio non distorto per l'obiettivo MMD, consentendo l'uso di mappe delle caratteristiche non differenziabili e campionatori non riparametrizzabili (es. modelli autoregressivi, diffusione discreta).
3. Generalizzazione: kCGM generalizza il metodo Calibrating Generative Models (CGM). Mentre CGM fa corrispondere le medie delle caratteristiche target, kCGM fa corrispondere l'intera distribuzione tramite MMD. Il documento nota che l'uso di un kernel dot-product in kCGM recupera l'obiettivo CGM-relax.

Risultati Sperimentali
Gli autori valutano kCGM attraverso tre domini: generazione di piccole molecole, generazione di strutture proteiche e generazione di DNA regolatorio.

• Piccole Molecole Antibiotiche:

  • Task: Calibrare un modello a grafo autoregressivo G2PT per far corrispondere la distribuzione delle caratteristiche di 174 antibiotici noti.
  • Caratteristiche: Fingerprint di Morgan (sottostrutture locali), scaffold di Murcko (topologia) e descrittori fisico-chimici.
  • Risultati: kCGM raggiunge un compromesso superiore tra il matching della distribuzione delle caratteristiche target (misurato tramite MMD) e la distanza dal modello preaddestrato rispetto a CGM e al fine-tuning diretto. Crucialmente, mentre il fine-tuning diretto sul piccolo set di antibiotici ha portato a un'alta proporzione di molecole chimicamente invalide, kCGm ha migliorato il matching delle caratteristiche aumentando al contempo la validità chimica. kCGM ha anche dimostrato il trasferimento cross-feature, migliorando il matching di scaffold e descrittori anche quando addestrato solo su fingerprint.

• Diversità delle Strutture Proteiche:

  • Task: Effettuare il fine-tuning di Genie 2 (un modello di diffusione proteica) per far corrispondere la distribuzione della struttura secondaria (alfa-eliche vs beta-foglietti) del dataset CATH.
  • Risultati: Il modello preaddestrato produceva principalmente alfa-eliche. kCGM ha ridotto il disaccordo con la distribuzione target di circa 2 volte rispetto a CGM. kCGM ha mostrato una curva di tradeoff fluida con il peso KL $\lambda$, mentre le prestazioni di CGM erano meno sensibili a $\lambda$, rendendo difficile l'ottimizzazione.

• Generazione di DNA Regolatorio:

  • Task: Effettuare il fine-tuning di un modello di diffusione discreta per generare sequenze di DNA con profili di attività cellulare-specifici realistici (predetti da AlphaGenome).
  • Risultati: kCGM ha con successo fatto corrispondere l'intera distribuzione dei profili di attività per 47 diversi argomenti cis-regolatori, superando CGM (che corrisponde solo alle medie). Il modello sottoposto a fine-tuning ha generato sequenze con profili di attività simili a enhancer reali, mentre il modello preaddestrato generava sequenze con scarsa attività.

Significatività e Affermazioni
Il documento afferma che kCGM fornisce un'alternativa robusta al fine-tuning diretto, in particolare quando i dataset target sono piccoli o quando solo aspetti specifici di una distribuzione devono essere preservati. La sua principale significatività risiede nella capacità di:

• Allineare i modelli generativi alle distribizioni di caratteristiche specificate dall'utente utilizzando solo la supervisione a livello di caratteristica.
• Gestire caratteristiche black-box e non differenziabili (es. fingerprint, output di predittori) che sono comuni nei domini scientifici ma incompatibili con il tipico fine-tuning basato su gradienti.
• Mantenere o migliorare la qualità del campione (validità) pur correggendo la calibrazione distributiva, un tipico fallimento del fine-tuning diretto.
• Essere applicabile a diverse architetture, inclusi modelli autoregressivi, diffusione in spazio continuo e diffusione discreta.

Gli autori notano limitazioni, come la potenziale inefficacia per caratteristiche ad altissima dimensionalità a causa della maledizione della dimensionalità nei test di ipotesi MMD, e la necessità di log-probabilità delle traiettorie di campionamento trattabili (escludendo i modelli di flow matching). Suggeriscono che il lavoro futuro potrebbe estendere questo framework a contesti in cui esistono solo vincoli distributivi senza campioni completi nello spazio generativo originale.