Constrained Flow Optimization via Sequential Fine Tuning for Molecular Design

Riepilogo Tecnico: Ottimizzazione di Flusso Vincolata tramite Fine-Tuning Sequenziale per la Progettazione Molecolare

1. Definizione del Problema

Il documento affronta la sfida di adattare i modelli generativi di base, specificamente i modelli di flusso (flow) e di diffusione, per applicazioni di scoperta scientifica nel mondo reale, come la progettazione molecolare e l'ingegneria delle proteine. Sebbene questi modelli eccellano nel catturare distribuzioni di dati complesse, il campionamento standard è insufficiente per i compiti di scoperta, dove l'obiettivo è generare candidati che massimizzino specifici premi (ad esempio, affinità di legame, momento di dipolo) pur soddisfacendo rigorosamente vincoli specifici del dominio (ad esempio, accessibilità sintetica, limiti di tossicità o stabilità energetica).

I metodi di fine-tuning esistenti, spesso basati sull'apprendimento per rinforzo (RL) o sulla teoria del controllo, ottimizzano tipicamente le funzioni di premio sotto regolarizzazione KL per prevenire lo scostamento distributivo (distributional drift). Tuttavia, essi faticano a gestire in modo affidabile i vincoli "hard". Gli approcci ingenui, come l'incorporazione dei vincoli come termini pesati nella funzione obiettivo, sono considerati inaffidabili perché il peso ottimale tra premi e vincoli varia tra i diversi compiti e le fasi di addestramento, richiedendo un inefficiente processo di tentativi ed errori. Inoltre, tali metodi non possono garantire che i campioni generati soddisfino esplicitamente i vincoli, producendo spesso campioni ad alto premio che violano proprietà fisiche o chimiche critiche.

La domanda centrale posta è: Come può un modello di flusso o di diffusione pre-addestrato essere sottoposto a fine-tuning per bilanciare in modo affidabile e prevedibile l'ottimizzazione del premio e la soddisfazione dei vincoli?

2. Metodologia

Gli autori propongono un framework formale e una soluzione algoritmica a questo problema.

Framework di Ottimizzazione Generativa Vincolata:
Il problema è formulato come la ricerca di una policy $\pi^*$ che massimizzi il premio atteso $r(x)$ pur soddisfacendo un vincolo atteso $c(x) \leq B$, soggetto a regolarizzazione KL rispetto a un modello pre-addestrato $\pi_{pre}$. L'obiettivo è:
$$\max_{\pi} \mathbb{E}_{x \sim p^{\pi}_1}[r(x)] - \alpha D_{KL}(p^{\pi}_1 || p^{pre}_1) \quad \text{s.t.} \quad \mathbb{E}_{x \sim p^{\pi}_1}[c(x)] \leq B$$
Questa formulazione generalizza sia l'ottimizzazione guidata dal premio che la generazione vincolata (dove $r$ è costante).

Ottimizzazione di Flusso Vincolata (CFO):
Per risolvere questo problema, gli autori introducono la Constrained Flow Optimization (CFO), un algoritmo basato sullo schema Augmented Lagrangian (AL). CFO trasforma il problema vincolato in una sequenza di sottoproblemi di fine-tuning non vincolati e regolarizzati dall'entropia. L'algoritmo opera iterativamente con i seguenti passaggi:

1. Costruzione dell'Obiettivo Aumentato: In ogni iterazione $k$, viene costruita una funzione obiettivo aumentata $f_k(x)$. Questa funzione combina il premio $r(x)$ con un termine di penalità derivato dalla violazione del vincolo, regolato da un moltiplicatore di Lagrange $\lambda_k$ e un parametro di penalità $\rho_k$:
   $$f_k(x) = r(x) - \frac{\rho_k}{2} \left[ \max\left(0, c(x) - B - \frac{\lambda_k}{\rho_k}\right) \right]^2$$
2. Passaggio di Fine-Tuning: Un risolutore di fine-tuning standard (ad esempio, Adjoint Matching, DiffusionNFT) viene utilizzato per ottimizzare la policy $\pi_k$ rispetto all'obiettivo aumentato $f_k$, regolarizzato dalla divergenza KL rispetto al modello pre-addestrato.
3. Aggiornamento delle Variabili Duali:
   • Moltiplicatore di Lagrange ($\lambda$): Aggiornato in base al gap del vincolo empirico $G_k = \mathbb{E}[c(x)] - B$. Se il vincolo è violato ($G_k > 0$), $\lambda$ viene diminuito per rafforzare la penalità; se soddisfatto, viene rilassato.
   • Parametro di Penalità ($\rho$): Aggiornato in base a una statistica di contrazione $V_k$. Se il gap del vincolo non si contrae sufficientemente, $\rho$ viene aumentato per imporre una maggiore fattibilità.

Questo approccio sequenziale consente all'algoritmo di bilanciare automaticamente la massimizzazione del premio e la soddisfazione dei vincoli senza la necessità di una regolazione manuale dei pesi.

3. Contributi Principali

Il documento presenta quattro contributi primari:

• Framework Formale: Formula l'Ottimizzazione Generativa Vincolata tramite fine-tuning di flusso, catturando il compito pratico dell'adattamento guidato dal premio sotto vincoli arbitrari.
• Algoritmo CFO: Introduce CFO, un metodo basato su Augmented Lagrangian, che risolve il problema vincolato tramite fine-tuning sequenziale.
• Garanzie Teoriche: Sotto assunzioni realistiche di un risolutore approssimato e specifiche condizioni oracle, gli autori forniscono garanzie di convergenza. Dimostrano che CFO restituisce una policy che minimizza l'infeasibilità e, sotto condizioni più forti (errore di approssimazione evanescente), raggiunge l'ottimalità globale.
• Validazione Empirica: Il metodo è valutato sia in contesti illustrativi a bassa dimensionalità che in compiti di progettazione molecolare ad alta dimensionalità, dimostrando la sua capacità di scambiare efficacemente premi e vincoli.

4. Risultati Sperimentali

Gli autori valutano CFO su due fronti principali:

Ambienti Illustrativi a Bassa Dimensionalità:

• Ottimizzazione Vincolata: Utilizzando una miscela di Gaussiane, CFO ha indirizzato con successo il modello verso regioni ad alto premio (vicino a un punto target) rimanendo al contempo entro regioni di vincolo valide (triangoli). Rispetto allo standard Adjoint Matching (che ignora i vincoli), CFO ha ridotto significativamente le violazioni dei vincoli (da 2.47 a 0.12) mantenendo un alto premio.
• Generazione Vincolata: Impostando il premio a un valore costante, CFO ha recuperato le capacità di generazione vincolata. Variando il limite del vincolo $B$, è stata dimostrata la possibilità di un rilassamento controllato del vincolo, evidenziando la flessibilità del metodo.
• Agnosticismo del Solutore: CFO ha operato coerentemente sia con risolutori di primo ordine (Adjoint Matching) che con risolutori privi di gradiente (DiffusionNFT), confermando la sua modularità.

Compiti di Progettazione Molecolare:

• Massimizzazione del Dipolo con Vincolo Energetico: Utilizzando il modello FlowMol sul dataset GEOM Drugs, CFO ha ottimizzato il momento di dipolo (premio) imponendo un limite superiore all'energia totale (vincolo, proxy per la stabilità).
  • CFO ha aumentato il momento di dipolo da 6.55 D (pre-addestrato) a 8.39 D rispettando il vincolo energetico (energia attesa -82.28 Ha rispetto al limite -80 Ha).
  • Al contrario, il fine-tuning non vincolato (Adjoint Matching) ha ottenuto premi simili ma ha violato il vincolo energetico (energia attesa -78.31 Ha).
• Confronto con i Baseline: CFO ha superato i baseline a peso fisso regolati manualmente. La regolazione manuale ha richiesto una ricerca esaustiva su 18 ordini di grandezza dei pesi per trovare una singola configurazione che eguagliasse le prestazioni di CFO, evidenziando l'inefficienza della regolazione manuale.
• Validità ed Efficienza: Sebbene il fine-tuning abbia ridotto la percentuale di molecole chimicamente valide (una sfida nota nella progettazione generativa), CFO ha mantenuto una validità superiore rispetto ai baseline non vincolati a parità di livello di premio. CFO ha inoltre dimostrato efficienza computazionale, richiedendo un solo passaggio di fine-tuning rispetto agli alti costi di campionamento dei metodi di vincolo durante l'inferenza come DiffOpt.

5. Significato e Rivendicazioni

Il documento afferma che CFO fornisce un meccanismo affidabile e prevedibile per l'ottimizzazione generativa vincolata. La sua importanza risiede in:

• Scambio Automatizzato: Elimina la necessità di una regolazione manuale e inefficiente dei pesi dei vincoli, che rappresenta un collo di bottiglia principale nei moderni workflow di progettazione generativa.
• Rigore Teorico: Offre garanzie di convergenza sia per la soddisfazione dei vincoli che per l'ottimalità del premio, fondando il metodo nella consolidata teoria dell'ottimizzazione (Augmented Lagrangian).
• Utilità Pratica: Applicando con successo il metodo a compiti di progettazione molecolare ad alta dimensionalità, gli autori dimostrano che il metodo scala a problemi scientifici reali dove i vincoli (come la stabilità chimica) non sono negoziabili.
• Modularità: Il framework è agnostico rispetto al risolutore di fine-tuning sottostante, permettendo l'integrazione con varie strategie di ottimizzazione (basate su gradiente o prive di gradiente) e l'applicabilità a diverse modalità di dati (continui, discreti o misti).

Gli autori concludono che CFO consente l'adattamento di modelli di flusso pre-addestrati per massimizzare i premi attesi imponendo al contempo vincoli arbitrari, offrendo uno strumento robusto per le applicazioni di scoperta scientifica.