DriftingMol: Decoder-Coupled Drift for One-Pass Property-Conditional Molecular Generation

Riepilogo Tecnico: DriftingMol

Enunciato del Problema
Il documento affronta la sfida della generazione molecolare condizionata alle proprietà, che richiede la produzione di molecole valide e diversificate in grado di rispondere a valori continui di proprietà target (ad esempio, QED, LogP) con bassi costi di campionamento. I metodi esistenti devono affrontare un compromesso: gli approcci di campionamento iterativo (come la diffusione o l'apprendimento per rinforzo) sono computazionalmente costosi durante l'inferenza, mentre i decoder nello spazio latente o basati su grafi spesso perdono i segnali continui delle proprietà quando la decodifica discreta collassa stati latenti vicini. Inoltre, garantire la validità molecolare (correttezza sintattica e di valenza) rimane un vincolo discreto che complica il controllo continuo delle proprietà.

Metodologia
DriftingMol introduce un framework a due stadi che adatta i "modelli di deriva" a uno spazio latente molecolare SELFIES.

1. Pipeline a Due Stadi:

   • Stadio 1: Un $\beta$-VAE SELFIES viene addestrato sul dataset ZINC250K. L'encoder mappa le sequenze di token in uno spazio latente a 256 dimensioni ($z$), e un decoder Transformer one-shot prevede le posizioni SELFIES in parallelo. Questo stadio stabilisce uno spazio latente continuo regolarizzato in cui la validità è ereditata dalla rappresentazione SELFIES piuttosto che appresa dal generatore.
   • Stadio 2: Il VAE viene congelato. Un generatore in stile DiT ($f_\theta$) viene addestrato per mappare rumore gaussiano, condizioni di proprietà opzionali ($c$) e una scala di guida incorporata ($\alpha$) in un vettore latente $\hat{z}$ in un singolo passaggio in avanti. La molecola finale viene ottenuta decodificando $\hat{z}$ attraverso il decoder SELFIES congelato.

2. Deriva Accoppiata al Decoder:
   L'innovazione centrale è la definizione del campo di deriva. Invece di utilizzare un estrattore di feature esterno, DriftingMol utilizza la rappresentazione nascosta del decoder VAE congelato ($\phi$) come mappa di feature di deriva.

   • Campo di Deriva: Viene costruito un campo di deriva a kernel tra i latenti generati ($\hat{z}_i$) e i latenti di riferimento ($z_{ref, j}$) basandosi sulle loro feature di decoder: $V_{\tau, i} = \sum_j W^\tau_{ij} (\phi(z_{ref, j}) - \phi(\hat{z}_i))$.
   • Accoppiamento del Gradiente: Crucialmente, mentre i pesi del decoder sono congelati, il percorso del gradiente dalla perdita di deriva attraverso la mappa di feature del decoder $\phi(\hat{z})$ fino ai parametri del generatore $\theta$ rimane aperto. Ciò crea una "metrica di pullback" in cui gli aggiornamenti del generatore sono ponderati dalla sensibilità strutturale locale del decoder.
   • Diversità: Un termine di repulsione latente $k$-NN ($L_{zdiv}$) viene aggiunto alla funzione di perdita per prevenire il collasso delle modalità e garantire l'unicità.

3. Obiettivo di Addestramento:
   Il generatore minimizza una perdita di deriva ($L_{drift}$) che attira le feature del decoder del latente generato verso una distribuzione target definita da riferimenti abbinati alle proprietà, combinata con la regolarizzazione della diversità.

Contributi Chiave

• Pipeline di Deriva Latente SELFIES: Un framework a basso costo per la generazione condizionata alle proprietà che separa la validità (gestita da SELFIES/VAE) dalla controllabilità.
• Deriva Accoppiata al Decoder: Un meccanismo che utilizza il decoder VAE congelato come spazio di feature di deriva preservando al contempo il percorso del gradiente attraverso il decoder fino al generatore. Ciò allinea il gradiente di deriva con il processo effettivo di decodifica molecolare.
• Evidenza di Ablazione: Uno studio completo di ablazione su 15 condizioni che dimostra come l'accoppiamento con il decoder sia il principale motore del controllo delle proprietà, superando le varianti di deriva nello spazio latente, con feature casuali o con feature esterne.
• Efficienza: Il metodo raggiunge un bias sulle proprietà con una sola valutazione del generatore più un singolo passaggio attraverso il decoder congelato, evitando il denoising iterativo o l'integrazione ODE.

Risultati
Valutati su ZINC250K con condizionamento QED e multi-proprietà:

• Controllo QED: L'impostazione predefinita di DriftingMol raggiunge una correlazione di Spearman per il QED ($\rho$) di 0,493 con un'unicità del 94,7%. La variante più forte accoppiata al decoder (senza regolarizzazione della diversità) raggiunge $\rho = 0,510$ ma con un'unicità ridotta (74,1%).
• Condizionamento Multi-Proprietà: Sotto un protocollo di condizionamento a quattro proprietà abbinato (QED, SA, LogP, MolWt), la deriva accoppiata al decoder raggiunge una correlazione media di Spearman di 0,598.
• Risultati dell'Ablazione:
  • Necessità dell'Accoppiamento: I modelli che disaccoppiano il decoder o utilizzano stop-gradienti producono una correlazione QED vicina allo zero ($\rho \approx 0,01$) e un'unicità molto bassa.
  • Qualità delle Feature: Feature esterne addestrate (ad esempio, LatentMAE) o feature casuali risultano in correlazioni significativamente più basse ($\rho \approx 0,28\text{--}0,33$), indicando che le feature predittive da sole sono insufficienti senza il Jacobiano del decoder.
  • Diversità: Rimuovere il termine di diversità $z$ aumenta leggermente la correlazione ma riduce drasticamente l'unicità, confermando il ruolo del termine nel prevenire il collasso delle modalità.
• Baseline: Le baseline con lo stesso backbone (VAE, WGAN-GP, DDPM, Flow-Matching) addestrate sulla stessa cache latente hanno raggiunto correlazioni significativamente più basse ($\rho < 0,16$).

Significato e Affermazioni
Il documento afferma che la deriva accoppiata al decoder è un meccanismo efficiente per la generazione molecolare con bias sulle proprietà. Il significato risiede nell'intuizione matematica secondo cui il decoder congelato fornisce non solo uno spazio di feature, ma una specifica metrica di pullback ($J_\phi^\top J_\phi$) che definisce la geometria in cui il generatore viene ottimizzato. Questa metrica allinea gli aggiornamenti del generatore con la sensibilità strutturale del decoder molecolare.

Gli autori sottolineano che:

1. La Validità è a Livello di Rappresentazione: La validità è garantita dalla rappresentazione SELFIES e dal decoder congelato, non appresa dal generatore condizionale.
2. Efficienza in Un Solo Passaggio: Il metodo evita gli alti costi di inferenza del campionamento iterativo.
3. Bias sulle Proprietà vs. Esattezza: Il sistema è progettato per la generazione con bias sulle proprietà (arricchendo le librerie verso regioni target) piuttosto che per l'ottimizzazione esatta delle proprietà. Non garantisce la sintetizzabilità o l'attività biologica, e le molecole generate richiedono un filtraggio a valle.
4. Specificità del Meccanismo: I guadagni di prestazioni sono attribuiti specificamente all'accoppiamento del gradiente attraverso il decoder, non semplicemente all'uso di un VAE pre-addestrato o di estrattori di feature esterni.