🤖 machine learning

Uncertainty Estimation for Molecular Diffusion Models

Cet article propose une méthode post-hoc pour estimer l'incertitude par échantillon dans les modèles de diffusion moléculaire préentraînés en mesurant la variabilité de la prédiction du bruit via une approximation de Laplace, démontrant que le score résultant prédit efficacement la qualité des échantillons et permet des améliorations de performance grâce au filtrage au moment du test.

Auteurs originaux : Paul Seij, Christian A. Naesseth, Stephan Mandt, Metod Jazbec

Publié 2026-06-12

📖 5 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Paul Seij, Christian A. Naesseth, Stephan Mandt, Metod Jazbec

Article original sous licence CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Ceci est une explication générée par l'IA de l'article ci-dessous. Elle n'a pas été rédigée ni approuvée par les auteurs. Pour une précision technique, consultez l'article original. Lire la clause de non-responsabilité complète

Imaginez que vous avez un robot chef super intelligent capable d'inventer de nouvelles recettes (molécules) à partir de zéro. Ce robot utilise une technique spéciale appelée « modèle de diffusion ». Il commence avec un bol de bruit aléatoire (comme de la neige sur une vieille télévision) et l'affine lentement, étape par étape, jusqu'à ce qu'une molécule parfaite et stable apparaisse.

Le problème ? Parfois, le robot s'embrouille. Il pourrait vous servir une recette qui ressemble à un gâteau, mais qui est en fait faite de verre — elle ressemble à une molécule, mais elle est chimiquement impossible ou instable. Jusqu'à présent, il n'y avait aucun moyen simple pour que le robot puisse dire : « Hé, je ne suis pas sûr à 100 % que celle-ci soit bonne » avant que vous ne passiez des heures à la tester en laboratoire.

Ce document présente un nouveau « compteur de confiance » pour ces robots moléculaires. Voici comment cela fonctionne, en utilisant des analogies simples :

1. L'astuce du « deuxième avis »

D'ordinaire, le robot crée une molécule en un seul passage fluide. Pour vérifier sa confiance, les auteurs lui demandent de faire quelque chose de différent : imaginez que le robot ait un léger tremblement dans la main.

Au lieu de réentraîner le robot, ils prennent le cerveau du robot (le réseau de neurones) et font semblant qu'il possède un petit tremblement aléatoire dans ses dernières couches. Ils demandent au robot de prédire l'étape suivante de la création de la molécule plusieurs fois, chaque fois avec un « tremblement » légèrement différent.

Si le robot est confiant : Même avec une main tremblante, il prédit la même étape suivante à chaque fois. Le résultat est un chemin clair et stable.
Si le robot est confus : La main tremblante le fait deviner des étapes suivantes radicalement différentes. Les prédictions se dispersent comme une nuée d'oiseaux effrayés par un chat.

Les auteurs mesurent cette « dispersion » ou « variabilité ». Une forte dispersion signifie une grande incertitude (la molécule est probablement mauvaise). Une faible dispersion signifie une grande confiance (la molécule est probablement bonne).

2. La phase de « nettoyage » est la plus importante

Le robot crée des molécules sur 1 000 petites étapes, partant d'un bruit pur pour finir par une forme claire. Les auteurs ont découvert qu'il n'est pas nécessaire de vérifier la confiance du robot pendant tout le processus.

C'est comme regarder quelqu'un résoudre un puzzle. S'il éprouve des difficultés, il peut hésiter au début, mais vous savez vraiment qu'il est confus s'il continue de deviner de manière désordonnée alors que l'image est presque terminée. Le document montre que le « compteur de confiance » du robot est plus précis lorsqu'on observe les toutes dernières étapes du processus de création, juste avant que la molécule ne soit terminée.

3. Utiliser le compteur pour filtrer les meilleurs

Une fois que nous avons ce « score d'incertitude », nous pouvons l'utiliser comme un filtre. Imaginez que le robot doive cuisiner 20 000 recettes, mais que vous n'ayez le temps d'en tester que 10 000.

Sans le compteur : Vous pourriez choisir 10 000 recettes au hasard, et beaucoup pourraient être des « gâteaux de verre ».
Avec le compteur : Vous laissez le robot cuisiner les 20 000 recettes, calculez le score d'incertitude pour chacune, et jetez les 10 000 qui présentent le plus de « tremblement » (incertitude). Vous gardez les 10 000 qui ont le moins de tremblement.

Le Résultat : Sur un ensemble de données de petites molécules simples (appelé QM9), cette méthode a fonctionné comme par magie. En jetant les prédictions « tremblantes », la qualité des molécules restantes a considérablement augmenté. En fait, le simple fait de filtrer les résultats de cette manière a amélioré la qualité autant que si l'on avait changé pour un modèle de robot complètement différent et plus avancé.

4. Le bémol

Le document est honnête sur ses limites. Bien que ce « compteur de confiance » ait très bien fonctionné sur de petites molécules simples, il n'a pas fonctionné sur un ensemble de données plus large et plus complexe de molécules de type médicaments (GEOM-Drugs). Les auteurs suggèrent que pour ces molécules complexes, le « tremblement » du robot ne raconte pas la même histoire sur la qualité, et que des recherches supplémentaires sont nécessaires pour comprendre pourquoi.

Résumé

En bref, les auteurs ont construit un moyen de demander à une IA moléculaire : « À quel point es-tu sûr de toi ? » en vérifiant si les prédictions de l'IA restent stables lorsque son cerveau est légèrement secoué. Ils ont découvert que lorsque l'IA est instable, les molécules résultantes sont généralement mauvaises. En utilisant ce signal pour filtrer les prédictions instables, on peut obtenir de bien meilleurs résultats sans avoir besoin de réentraîner l'IA ou d'en construire une nouvelle.

Résumé Technique : Estimation de l'Incertitude pour les Modèles de Diffusion Moléculaire

Énoncé du Problème
Les modèles de diffusion sont devenus le paradigme dominant pour la génération moléculaire 3D. Cependant, un goulot d'étranglement pratique critique subsiste : ces modèles manquent d'un signal fondé sur des principes pour identifier quand une molécule générée est chimiquement invalide ou instable. Alors que les modèles prédictifs utilisent généralement l'estimation de l'incertitude pour signaler les sorties peu fiables, les méthodes existantes pour les modèles de diffusion se sont principalement concentrées sur le domaine de l'image. Dans le domaine moléculaire, où l'évaluation en aval (par exemple, le docking ou la validation en laboratoire humide) est coûteuse, s'appuyer sur des filtres grossiers basés sur des règles ou sur la propre vraisemblance du modèle — qui est souvent mal calibrée — reste la norme. Il y a un manque de travaux analogues fournissant des signaux de qualité d'échantillon pour les modèles de diffusion moléculaire.

Méthodologie
Les auteurs proposent une méthode post-hoc pour estimer l'incertitude par échantillon pour les modèles de diffusion moléculaires pré-entraînés sans réentraîner le générateur. L'approche repose sur l'intuition que les échantillons pour lesquels le modèle présente une incertitude interne présenteront des trajectoires de débruitage plus variables.

Approximation de Laplace : Pour rester compatible avec n'importe quel modèle pré-entraîné, la méthode ajuste une approximation de Laplace aux paramètres du réseau de débruitage. Plus précisément, une distribution de type $q(\theta) \approx \mathcal{N}(\theta | \hat{\theta}, \Sigma)$ est construite autour des paramètres pré-entraînés $\hat{\theta}$ . La covariance $\Sigma$ est approximée à l'aide de la matrice d'information de Fisher empirique (appliquée uniquement aux couches finales pour garantir la scalabilité) plus un terme de précision de l'a priori.
Variabilité de la Prédiction du Bruit : Pendant la trajectoire de génération, pour un sous-ensemble sélectionné de pas de temps $\mathcal{T}$ , la méthode tire $M$ échantillons de poids $\theta_m$ à partir de la distribution a posteriori $q(\theta)$ . Pour chaque échantillon, le débruiteur prédit le bruit $\epsilon^m_t = f_{\theta_m}(\hat{x}_t, t)$ .
Agrégation du Score d'Incertitude : L'incertitude intermédiaire $u_t$ à chaque pas de temps est définie comme la variance d'échantillon élément par élément de ces prédictions de bruit à travers les $M$ échantillons. Cette incertitude par étape est ensuite agrégée à travers les pas de temps sélectionnés, tous les atomes ( $N_a$ ), et toutes les dimensions de caractéristiques ( $D$ ) pour produire un score d'incertitude scalaire unique $u(\hat{x})$ pour la molécule générée :
$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}$
Empiriquement, les auteurs constatent que mesurer l'incertitude uniquement dans les derniers 10 % du processus de génération (là où $t$ est petit) est suffisant pour la performance, limitant ainsi la surcharge de calcul.

Principales Contributions

Première Application à la Diffusion Moléculaire : Ce travail étend les approches d'incertitude de diffusion de l'image au domaine de la génération moléculaire, répondant au besoin spécifique de signaux de qualité d'échantillon dans des domaines scientifiques dotés d'évaluations en aval coûteuses.
Cadre Post-hoc : La méthode ne nécessite aucun réentraînement du modèle de diffusion sous-jacent, ce qui la rend applicable aux architectures d'état de l'art existantes telles que EDM et GeoLDM.
Mise à l'échelle au Temps de Test (Test-Time Scaling) : L'article démontre comment le score d'incertitude peut être utilisé comme critère de filtrage dans un cadre de mise à l'échelle au temps de test (sur-échantillonnage et filtrage) pour améliorer la qualité des molécules générées.

Résultats
Les auteurs ont évalué leur approche sur les jeux de données QM9 et GEOM-Drugs en utilisant les modèles EDM et GeoLDM.

Corrélation avec la Qualité : Sur le jeu de données QM9, le score d'incertitude proposé présente une corrélation négative statistiquement significative avec les métriques de qualité établies (stabilité moléculaire, stabilité atomique et validité). Cela indique que des scores d'incertitude plus élevés correspondent à des échantillons de moindre qualité.
Comparaison avec la Vraisemblance (Likelihood) : Le score d'incertitude surpasse systématiquement la log-vraisemblance négative (NLL) du modèle de diffusion en tant que prédicteur de la qualité de l'échantillon. Par exemple, sur EDM/QM9, le score d'incertitude a atteint une corrélation de Spearman de $-0,284$ avec la stabilité moléculaire, contre $-0,150$ pour la NLL.
Performance de la Mise à l'échelle au Temps de Test : En générant 20 000 molécules et en filtrant pour ne garder que les 10 000 ayant l'incertitude la plus faible, les auteurs ont observé des améliorations significatives de la stabilité moléculaire ( $\sim10\%$ ), de la stabilité atomique ( $\sim1\%$ ) et de la validité ( $\sim5\%$ ) pour le modèle EDM. Ces gains étaient comparables à l'amélioration obtenue en changeant l'architecture du modèle sous-jacent de EDM vers GeoLDM.
Études d'Ablation :
- Approximation de Fisher : Le retrait de l'approximation de Fisher (remplaçant la covariance par une gaussienne isotrope) a produit des performances presque identiques, suggérant que le score agit davantage comme une mesure de sensibilité aux petites perturbations que comme une véritable incertitude épistémique bayésienne.
- Sélection des Pas de Temps : Le signal est plus fort à l'extrémité "propre" de la trajectoire de génération (petit $t$ ), validant le choix de concentrer la mesure de l'incertitude sur les étapes finales.
Limites : La méthode n'a pas produit d'améliorations significatives sur le jeu de données plus large et plus complexe GEOM-Drugs, indiquant que l'approche ne se transfère pas actuellement aux jeux de données présentant une plus grande complexité conformationnelle.

Signification et Revendications
L'article affirme faire un "premier pas" vers l'extension des approches d'incertitude de diffusion au domaine moléculaire. La principale importance réside dans la fourniture d'un outil post-hoc pratique pour filtrer les molécules générées, améliorant ainsi l'efficacité des pipelines d'évaluation en aval. Les auteurs positionnent leur score d'incertitude comme un remplacement naturel des vérificateurs basés sur la NLL dans les méthodes de mise à l'échelle adaptative au temps de test.

Les auteurs restent modestes quant à leurs affirmations, notant explicitement que l'approche échoue actuellement sur les jeux de données plus complexes et plus larges (GEOM-Drugs) et que le mécanisme sous-jacent semble être une mesure de sensibilité plutôt qu'une véritable incertitude épistémique bayésienne. Ils suggèrent que les travaux futurs devraient explorer le calcul de l'incertitude dans un espace de caractéristiques "sémantiques" en utilisant des encodeurs moléculaires pré-entraînés, de manière analogue aux succès récents dans la génération d'images de haute dimension.