Uncertainty Estimation for Molecular Diffusion Models

Technische Samenvatting: Onzekerheidsschatting voor Moleculaire Diffusiemodellen

Probleemstelling
Diffusiemodellen zijn de dominante paradigma geworden voor 3D moleculaire generatie. Echter, er bestaat een kritieke praktische flessenhals: deze modellen missen een principieel signaal om te identificeren wanneer een gegenereerd molecuul chemisch ongeldig of instabiel is. Terwijl predictieve modellen doorgaans onzekerheidsschatting gebruiken om onbetrouwbare outputs te markeren, hebben bestaande methoden voor diffusiemodellen zich voornamelijk gericht op het domein van afbeeldingen. In het moleculaire domein, waar downstream-evaluatie (bijv. docking of validatie in het natte laboratorium) kostbaar is, blijft het vertrouwen op grove regelgebaseerde filters of de likelihood van het model zelf—die vaak slecht gekalibreerd is—de standaard. Er is een gebrek aan analoge werk die kwaliteits-signalen voor moleculaire diffusiemodellen levert.

Methodologie
De auteurs stellen een post-hoc methode voor om per-sample onzekerheid te schatten voor vooraf getrainde moleculaire diffusiemodellen zonder de generator opnieuw te trainen. De aanpak is gebaseerd op de intuïtie dat samples waarover het model intern onzeker is, meer variabele denoising-trajecten zullen vertonen.

1. Laplace-benadering: Om compatibel te blijven met elk vooraf getraind model, past de methode een Laplace-benadering toe op de parameters van het denoising-netwerk. Specifiek wordt een posterior distributie $q(\theta) \approx \mathcal{N}(\theta | \hat{\theta}, \Sigma)$ geconstrueerd rond de vooraf getrainde parameters $\hat{\theta}$. De covariantie $\Sigma$ wordt benaderd met behulp van de empirische Fisher-informatie-matrix (toegepast op alleen de laatste lagen om schaalbaarheid te garanderen) plus een prior precisieterm.
2. Variabiliteit in Ruisvoorspelling: Tijdens het generatie-traject, voor een geselecteerde subset van tijdstappen $\mathcal{T}$, trekt de methode $M$ gewicht-samples $\theta_m$ uit de posterior $q(\theta)$. Voor elke sample voorspelt de denoiser de ruis $\epsilon^m_t = f_{\theta_m}(\hat{x}_t, t)$.
3. Aggregatie van de Onzekerheidsscore: De intermediaire onzekerheid $u_t$ bij elke tijdstap wordt gedefinieerd als de elementgewijze variantie van deze ruisvoorspellingen over de $M$ samples. Deze per-stap onzekerheid wordt vervolgens geaggregeerd over de geselecteerde tijdstappen, alle atomen ($N_a$), en alle feature-dimensies ($D$) om een enkele scalaire onzekerheidsscore $u(\hat{x})$ voor het gegenereerde molecuul te produceren:
   $$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}$$
   Empirisch stellen de auteurs vast dat het meten van onzekerheid enkel in de laatste 10% van het generatieproces (waar $t$ klein is) voldoende is voor prestaties, wat de computationele overhead beperkt.

Belangrijkste Bijdragen

• Eerste Toepassing op Moleculaire Diffusie: Dit werk breidt de onzekerheidsbenaderingen van diffusie uit van het beelddomein naar moleculaire generatie, waarmee wordt ingegaan op de specifieke behoefte aan kwaliteits-signalen voor samples in wetenschappelijke domeinen met dure downstream-evaluaties.
• Post-hoc Framework: De methode vereist geen hertraining van het onderliggende diffusiemodel, waardoor het toepasbaar is op bestaande state-of-the-art architecturen zoals EDM en GeoLDM.
• Test-Time Scaling: Het artikel demonstreert hoe de onzekerheidsscore kan worden gebruikt als een filtercriterium in een test-time scaling setting (oversampling en filtering) om de kwaliteit van gegenereerde moleculen te verbeteren.

Resultaten
De auteurs evalueerden hun aanpak op de QM9 en GEOM-Drugs datasets met behulp van EDM en GeoLDM modellen.

• Correlatie met Kwaliteit: Op de QM9-dataset vertoont de voorgestelde onzekerheidsscore een statistisch significante negatieve correlatie met gevestigde kwaliteitsmetrieken (moleculaire stabiliteit, atomaire stabiliteit en validiteit). Dit duidt erop dat hogere onzekerheidsscores overeenkomen met lagere kwaliteit samples.
• Vergelijking met Likelihood: De onzekerheidsscore presteert consistent beter dan de negatieve log-likelihood (NLL) van het diffusiemodel als voorspeller van de kwaliteit van samples. Bijvoorbeeld, op EDM/QM9 bereikte de onzekerheidsscore een Spearman-correlatie van $-0,284$ met moleculaire stabiliteit, vergeleken met $-0,150$ voor NLL.
• Prestaties van Test-Time Scaling: Door 20K moleculen te genereren en te filteren om de 10K met de laagste onzekerheid te behouden, observeerden de auteurs significante verbeteringen in moleculaire stabiliteit ($\sim10\%$), atomaire stabiliteit ($\sim1\%$) en validiteit ($\sim5\%$) voor het EDM-model. Deze winsten waren vergelijkbaar met de verbetering die werd bereikt door de onderliggende modelarchitectuur te wijzigen van EDM naar GeoLDM.
• Ablatie-studies:
  • Fisher-benadering: Het verwijderen van de Fisher-benadering (waarbij de covariantie wordt vervangen door een isotrope Gaussische verdeling) leverde bijna identieke prestaties op, wat suggereert dat de score eerder fungeert als een gevoeligheidsmaat voor kleine perturbaties dan als een strikte Bayesiaanse epistemische onzekerheid.
  • Selectie van Tijdstappen: Het signaal is het sterkst aan het "schone" einde van het generatie-traject (kleine $t$), wat de keuze valideert om de onzekerheidsmeting te concentreren op de laatste stappen.
• Beperkingen: De methode leverde geen betekenisvolle verbeteringen op de grotere, complexere GEOM-Drugs dataset, wat aangeeft dat de aanpak momenteel niet overdraagbaar is naar datasets met een hogere conformationele complexiteit.

Betekenis en Claims
Het artikel claimt een "eerste stap" te zetten naar het uitbreiden van onzekerheidsbenaderingen voor diffusie naar het moleculaire domein. De primaire betekenis ligt in het bieden van een praktische, post-hoc tool voor het filteren van gegenereerde moleculen, waardoor de efficiëntie van downstream-evaluatiepipelines wordt verbeterd. De auteurs positioneren hun onzekerheidsscore als een natuurlijke vervanging voor NLL-gebaseerde verifiers in adaptieve test-time scaling methoden.

De auteurs blijven bescheiden over hun claims en merken expliciet op dat de aanpak momenteel faalt op grotere, complexere datasets (GEOM-Drugs) en dat het onderliggende mechanisme eerder een gevoeligheidsmaat lijkt te zijn dan ware Bayesiaanse epistemische onzekerheid. Zij suggereren dat toekomstig werk de onzekerheid in een "semantische" feature-ruimte zou moeten berekenen met behulp van vooraf getrainde moleculaire encoders, analoog aan recente successen in hoog-dimensionale beeldgeneratie.