Sintesi Tecnica: Un Framework di Campo Medio per il Controllo Distributivo della Distribuzione al Tempo di Inferenza dei Modelli di Diffusione

Definizione del Problema

I modelli di diffusione vengono utilizzati sempre più spesso come campionatori controllabili, dove i processi di generazione vengono guidati durante l'inferenza per soddisfare obiettivi specifici. Mentre i metodi esistenti gestiscono efficacemente le ricompense puntuali (definite su singoli campioni, $r: \mathbb{R}^d \to \mathbb{R}$), molte applicazioni critiche richiedono ricompense distribuzionali (definite sulla misura generata, $R: \mathcal{P}(\mathbb{R}^d) \to \mathbb{R}$). Questi obiettivi distribuzionali includono la calibrazione delle popolazioni rispetto a osservazioni sperimentali, l'incoraggiamento della diversità o il bilanciamento dei modi.

Gli approcci attuali per lo steering distribuzionale incorporano spesso termini di gradiente nelle dinamiche di inferenza (ad esempio, potenziali repulsivi o verosimiglianze d'insieme). Tuttavia, similmente ai primi metodi di steering puntuale, queste euristiche mancano di garanzie teoriche riguardo alla distribuzione campionata. Nello specifico, l'aggiunta di un termine di gradiente non assicura che le particelle seguano la desiderata misura inclinata $\mu^*$, che massimizza l'obiettivo regolarizzato KL:
$$\arg \max_{\mu} \left\{ R(\mu) - \text{KL}(\mu \| p_1) \right\}$$
A differenza del caso puntuale, dove il target inclinato è un esplicito tilt esponenziale ($\mu^* \propto p_1 e^r$), il caso distribuzionale risulta in un target implicito dove il potenziale di inclinazione dipende dalla misura stessa. Attualmente non esiste un framework teoricamente fondato per correggere le dinamiche e i pesi per mirare a questa distribuzione implicita.

Metodologia

Gli autori propongono un Framework di Campo Medio per formalizzare e risolvere il controllo distribuzionale al tempo di inferenza.

1. Formulazione Teorica

Il problema è inquadrato come il puntamento di una misura inclinata $\mu^*_t$ ad ogni passo temporale $t \in [0,1]$, definita implicitamente da:
$$\mu^*_t(dx) = \frac{1}{Z_t} e^{\Psi_t(x, \mu^*_t)} p_t(x) dx$$
dove $\Psi_t(x, \mu) = \frac{\delta R_t}{\delta \mu}(x, \mu)$ è la prima variazione (derivata funzionale) della ricompensa tempo-dipendente $R_t$. Ciò crea una relazione di autoconsistenza in cui la misura target dipende dal potenziale, il quale a sua volta dipende dalla misura stessa.

2. Dinamiche di McKean-Vlasov Pesate

Per tracciare questo percorso, gli autori derivano un sistema di McKean-Vlasov pesato. A differenza dello steering di gradiente standard, il termine di steering $\nabla_x \Psi_t(X_t; \mu_t)$ induce un'interazione tra le particelle perché dipende dalla misura empirica evolutiva $\mu_t$.

Le dinamiche idealizzate per una particella $X_t$ e il suo log-peso $A_t$ sono:
$$dX_t = \left[ b_t(X_t) + \epsilon_t s_t(X_t) + \epsilon_t \nabla_x \Psi_t(X_t; \mu_t) \right] dt + \sqrt{2\epsilon_t} dB_t$$
$$dA_t = \left[ b_t(X_t) \cdot \nabla_x \Psi_t(X_t; \mu_t) + \dot{\Psi}_t(X_t) \right] dt$$
Qui, $\dot{\Psi}_t(x) = \frac{d}{dt} \Psi_t(x, \mu_t)$ è un termine di correzione cruciale. A differenza del caso puntuale dove $\dot{\Psi}_t$ è semplicemente $\partial_t r$, qui coinvolge l'evoluzione temporale della misura $\mu_t$. Gli autori dimostrano che $\dot{\Psi}_t$ soddisfa un'equazione integrale implicita che coinvolge la seconda variazione del funzionale di ricompensa.

3. Algoritmo a Particelle Finite

Per implementare questo, gli autori propongono uno schema di particelle interagenti pesate (Algoritmo 1):

• Aggiornamenti Posizionali: Le particelle sono aggiornate usando il drift di base, lo score e il gradiente della ricompensa rispetto alla misura empirica corrente $\hat{\mu}_t$.
• Aggiornamenti Log-Peso: I pesi vengono aggiornati per correggere la discrepanza tra le dinamiche guidate e il percorso inclinato target. Ciò richiede il calcolo di $\dot{\Psi}_t$.
• Risoluzione di $\dot{\Psi}_t$: Il paper fornisce due metodi per calcolare il termine della derivata temporale implicita:
  1. Sistema Lineare Implicito: Risoluzione di un sistema lineare derivato dall'equazione integrale (Proposizione 3.3).
  2. Iterazione a Punto Fisso: Uno schema di iterazione di Picard leggero che approssima gli aggiornamenti dei pesi.
• Resampling: Viene impiegato il campionamento sequenziale Monte Carlo (SMC) standard per prevene la degenerazione dei pesi.

Il framework recupera lo steering della ricompensa puntuale standard come un caso speciale in cui $R(\mu) = \int r(x) d\mu(x)$.

Contributi Chiave

1. Formalizzazione dello Steering Distribuzionale: Gli autori formalizzano lo steering a livello di distribuzione come il puntamento di una misura inclinata dove il tilt dipende dalla misura stessa, stabilendo un analogo distribuzionale del target di ricompensa inclinato puntuale.
2. Derivazione delle Dinamiche di McKean-Vlasov Pesate: Derivano un sistema dinamico fondato con un reweighting Feynman-Kac correttivo che assicura che le particelle interagenti pesate traccino il percorso inclinato desiderato $\mu^*_t$.
3. Recupero Teorico e Fondazione: Il framework recupera lo steering puntuale standard come caso speciale e fornisce una base teorica per gli esistenti metodi euristici di steering a livello di batch (ad esempio, campionamento della diversità, guida sperimentale).
4. Garanzie di Convergenza: Il paper dimostra che il sistema a particelle finite converge alle dinamiche idealizzate quando il numero di particelle $N \to \infty$.

Risultati Sperimentali

Gli autori validano il framework sia in contesti a bassa dimensione che ad alta dimensione:

• Bassa Dimensione (Miscela 1D):

  • Task: Inclinare una miscela gaussiana bimodale verso una distribuzione target con pesi dei modi differenti utilizzando una ricompensa di Maximum Mean Discrepancy (MMD) al quadrato.
  • Risultati: Lo steering di gradiente standard spinge le particelle verso il target ma fallisce nel corrispondere alla forma precisa della distribuzione inclinata $\mu^*$. Al contrario, il metodo mean-field mira accuratamente a $\mu^*$, minimizzando il gap di ottimalità $J(\mu) - J(\mu^*)$ quando la forza di steering $\lambda$ corrisponde al parametro dell'obiettivo. Ciò si verifica attraverso vari schemi di rumore.

• Alta Dimensione (Conformazione Proteica):

  • Task 1 (HIV-1 Protease): Guidare gli ensemble generati da Boltz-2 per corrispondere alle distribuzioni di distanza tra residui osservate sperimentalmente (dati DEER). Il modello base era distorto verso lo stato "chiuso"; l'obiettivo era recuperare le proporzioni dello stato "semi-aperto".
    • Risultato: Lo steering mean-field ha corrisposto accuratamente alla distribuzione empirica della distanza tra i residui, mentre lo steering basato solo sul gradiente ha prodotto una distribuzione con una forma leggermente diversa e picchi più forti.
  • Task 2 (Adenylate Kinase): Guidare verso una distribuzione 2D di angoli di legame derivata da simulazioni di Dinamica Molecolare.
    • Risultato: Lo steering mean-field ha prodotto uno spostamento verso la distribuzione target più forte rispetto ai metodi basati solo sul gradiente.
  • Task 3 (Cristallografia a raggi X): Guidare le generazioni di Protenix per adattarsi a una mappa di densità elettronica (struttura 4OLE) che presenta conformazioni eterogenee.
    • Risultato: Per la stessa forza di steering, lo steering mean-field ha ottenuto un maggiore allineamento del coseno con la densità target e ha posizionato più campioni nella conformazione alternativa (Conformazione B) rispetto allo steering basato solo sul gradiente, suggerendo che i metodi basati solo sul gradiente sottostimano lo steering rispetto al vero target del problema inverso.

Significato e Rivendicazioni

Il paper rivendica di colmare il divario tra il controllo puntuale e quello distribuzionale al tempo di inferenza fornendo una fondazione teorica rigorosa per il secondo.

• Correttezza: La rivendicazione principale è che la procedura proposta bersaglia esattamente la distribuzione inclinata dalla ricompensa $\mu^*$ nel limite di particelle infinite, correggendo le discrepanze inerenti allo steering di gradiente euristico.
• Interpretabilità: Il framework chiarisce che gli esistenti metodi di steering a livello di batch stanno effettivamente eseguendo passi di gradiente ma mancano del reweighting correttivo necessario per campionare dalla vera distribuzione posteriore/inclinata.
• Praticità: Sebbene il metodo introduca un costo computazionale aggiuntivo (risolvere per $\dot{\Psi}_t$), gli autori notano che questo è spesso trascurabile rispetto alle valutazioni della rete in contesti ad alta dimensione (ad esempio, proteine) e che il solver a punto fisso è efficiente.
• Limitazioni: Gli autori riconoscono che la dipendenza dal resampling può essere un limite per modalità come le immagini (potenzialmente portando a campioni duplicati) e che il metodo richiede funzioni di ricompensa differenziabili. Suggeriscono che l'integrazione con tecniche di riduzione della varianza (ad esempio, solver dell'equazione di Poisson) potrebbe mitigare i problemi di resampling.

Il lavoro si posiziona non come un sostituto di tutti i metodi esistenti, ma come un rigoroso fondamento che spiega perché certe euristiche funzionano e fornisce un meccanismo per garantire che mirino alla corretta distribuzione.