Uncertainty Estimation for Molecular Diffusion Models

Riepilogo Tecnico: Stima dell'Incertezza per Modelli di Diffusione Molecolare

Problematica
I modelli di diffusione sono diventati il paradigma dominante per la generazione 3D di molecole. Tuttavia, esiste un collo di bottiglia critico dal punto di vista pratico: questi modelli mancano di un segnale fondato per identificare quando una molecola generata è chimicamente invalida o instabile. Mentre i modelli predittivi utilizzano tipicamente la stima dell'incertezza per segnalare output inaffidabili, i metodi esistenti per i modelli di diffusione si sono concentrati principalmente sul dominio delle immagini. Nel dominio molecolare, dove la valutazione a valle (ad esempio, docking o validazione in laboratorio) è costosa, affidarsi a filtri grossolani basati su regole o alla verosimiglianza (likelihood) del modello stesso — che è spesso scarsamente calibrata — rimane lo standard. Vi è una mancanza di lavori analoghi che forniscano segnali sulla qualità del campione per i modelli di diffusione molecolare.

Metodologia
Gli autori propongono un metodo post-hoc per stimare l'incertezza per singolo campione per modelli di diffusione molecolare pre-addestrati senza riaddestrare il generatore. L'approccio si basa sull'intuizione secondo cui i campioni di cui il modello è internamente incerto esibiranno traiettorie di denoising più variabili.

1. Approssimazione di Laplace: Per rimanere compatibile con qualsiasi modello pre-addestrato, il metodo adatta un'approssimazione di Laplace ai parametri della rete di denoising. Nello specifico, viene costruita una distribuzione a posteriori $q(\theta) \approx \mathcal{N}(\theta | \hat{\theta}, \Sigma)$ attorno ai parametri pre-addestrati $\hat{\theta}$. La covarianza $\Sigma$ è approssimata utilizzando la matrice di informazione di Fisher empirica (applicata solo agli ultimi strati per garantire la scalabilità) più un termine di precisione del prior.
2. Variabilità della Predizione del Rumore: Durante la traiettoria di generazione, per un sottoinsieme selezionato di timestep $\mathcal{T}$, il metodo estrae $M$ campioni di peso $\theta_m$ dalla distribuzione a posteriori $q(\theta)$. Per ogni campione, il denoiser predice il rumore $\epsilon^m_t = f_{\theta_m}(\hat{x}_t, t)$.
3. Aggregazione dello Score di Incertezza: L'incertezza intermedia $u_t$ a ogni timestep è definita come la varianza del campione elemento per elemento di queste predizioni di rumore attraverso i $M$ campioni. Questa incertezza per step viene poi aggregata attraverso i timestep selezionati, tutti gli atomi ($N_a$) e tutte le dimensioni delle caratteristiche ($D$) per produrre un singolo score di incertezza scalare $u(\hat{x})$ per la molecola generata:
   $$u(\hat{x}) := \frac{1}{|\mathcal{T}| \cdot N_a \cdot D} \sum_{t \in \mathcal{T}} \sum_{n=1}^{N_a} \sum_{d=1}^{D} [u_t]_{n,d}$$
   Empiricamente, gli autori riscontrano che misurare l'incertezza solo nell'ultimo 10% del processo di generazione (dove $t$ è piccolo) è sufficiente per le prestazioni, limitando l'overhead computazionale.

Contributi Chiave

• Prima Applicazione alla Diffusione Molecolare: Questo lavoro estende gli approcci di incertezza della diffusione dal dominio delle immagini alla generazione molecolare, affrontando la specifica necessità di segnali sulla qualità del campione in domini scientifici con valutazione a valle costosa.
• Framework Post-hoc: Il metodo non richiede il riaddestramento del modello di diffusione sottostante, rendendolo applicabile ad architetture esistenti allo stato dell'arte come EDM e GeoLDM.
• Scalabilità al Test-Time: Il documento dimostra come lo score di incertezza possa essere utilizzato come criterio di filtraggio in un contesto di scaling al test-time (sovracampionamento e filtraggio) per migliorare la qualità delle molecole generate.

Risultati
Gli autori hanno valutato il loro approccio sui dataset QM9 e GEOM-Drugs utilizzando i modelli EDM e GeoLDM.

• Correlazione con la Qualità: Sul dataset QM9, il proposto score di incertezza mostra una correlazione negativa statisticamente significativa con le metriche di qualità stabilite (stabilità molecolare, stabilità atomica e validità). Ciò indica che score di incertezza più elevati corrispondono a campioni di qualità inferiore.
• Confronto con la Verosimiglianza (Likelihood): Lo score di incertezza supera costantemente la log-verosimiglianza negativa (NLL) del modello di diffusione come predittore della qualità del campione. Ad esempio, su EDM/QM9, lo score di incertezza ha ottenuto una correlazione di Spearman di $-0.284$ con la stabilità molecolare, rispetto a $-0.150$ per la NLL.
• Performance di Test-Time Scaling: Generando 20K molecole e filtrando per mantenere le 10K con l'incertezza più bassa, gli autori hanno osservato miglioramenti significativi nella stabilità molecolare ($\sim10\%$), stabilità atomica ($\sim1\%$) e validità ($\sim5\%$) per il modello EDM. Questi guadagni sono comparabili al miglioramento ottenuto passando dall'architettura EDM alla GeoLDM.
• Studi di Ablazione:
  • Approssimazione di Fisher: Rimuovere l'approssimazione di Fisher (sostituendo la covarianza con una Gaussiana isotropa) ha prodotto prestazioni quasi identiche, suggerendo che lo score agisce più come una misura di sensibilità a piccole perturbazioni che come una vera incertezza epistemica bayesiana.
  • Selezione dei Timestep: Il segnale è più forte alla fine "pulita" della traiettoria di generazione (piccolo $t$), validando la scelta di concentrare la misurazione dell'incertezza sugli ultimi step.
• Limitazioni: Il metodo non ha prodotto miglioramenti significativi sul dataset più grande e complesso GEOM-Drugs, indicando che l'approccio non si trasferisce attualmente a dataset con maggiore complessità conformazionale.

Significatività e Rivendicazioni
Il documento sostiene di compiere un "primo passo" verso l'estensione degli approcci di incertezza della diffusione al dominio molecolare. La principale significatività risiede nel fornire uno strumento post-hoc pratico per filtrare le molecole generate, migliorando così l'efficienza delle pipeline di valutazione a valle. Gli autori posizionano il loro score di incertezza come un naturale sostituto dei verificatori basati su NLL nei metodi di scaling adattivo al test-time.

Gli autori rimangono modesti riguardo alle loro rivendicazioni, notando esplicitamente che l'approccio attualmente fallisce su dataset più grandi e complessi (GEOM-Drugs) e che il meccanismo sottostante sembra essere una misura di sensibilità piuttosto che una vera incertezza epistemica bayesiana. Suggeriscono che il lavoro futuro dovrebbe esplorare il calcolo dell'incertezza in uno spazio di caratteristiche "semantiche" utilizzando encoder molecolari pre-addestrati, analogamente ai recenti successi nella generazione di immagini ad alta dimensionalità.