Interpretable epistemic uncertainty decomposition in sequential generative models via polynomial chaos surrogates

Résumé technique : Décomposition interprétable de l'incertitude épistémique dans les modèles génératifs séquentiels via des substituts de chaos polynomial

Énoncé du problème

Les modèles génératifs séquentiels, tels que les réseaux de flux génératifs (GFlowNets), sont de plus en plus centraux dans la découverte scientifique pilotée par l'IA, l'optimisation combinatoire et la conception biologique. Ces modèles opèrent souvent dans des conditions d'incertitude épistémique, spécifiquement lorsque les fonctions de récompense guidant leur génération sont elles-mêmes incertaines (par exemple, des substituts de réseaux de neurones entraînés sur des données expérimentales limitées). Bien que les méthodes existantes, telles que les ensembles profonds et les réseaux de neurones bayésiens, puissent quantifier l'ampleur de l'incertitude de prédiction, elles échouent à fournir une décomposition interprétable de cette incertitude. Par conséquent, les praticiens ne peuvent pas déterminer quels composants spécifiques de l'incertitude de la récompense d'entrée pilotent quelles décisions génératives, ni distinguer les signaux chimiques robustes des artefacts statistiques résultant des biais de l'ensemble de données d'entraînement.

Méthodologie

Les auteurs proposent un cadre de modélisation par substitut qui propage l'incertitude épistémique à travers les GFlowNets en utilisant des développements en chaos polynomial (PCE). L'approche se déroule en trois étapes :

1. Entraînement et paramétrisation de l'ensemble : Un ensemble modeste de GFlowNets est entraîné dans des conditions de récompense variées. L'espace des fonctions de récompense est paramétré dans un espace latent de faible dimension $\mu \in \mathbb{R}^d$ en utilisant l'analyse en composantes principales (PCA) ou des expansions analytiques (par exemple, Karhunen–Loève). La PCA est préférée car ses composantes orthogonales approximent des entrées gaussiennes indépendantes, une condition préalable à une décomposition de Sobol analytique valide.
2. Ajustement du substitut PCE : Pour chaque action $k$ à l'étape de trajectoire $t$, un PCE est ajusté pour mapper les paramètres de récompense de faible dimension $\mu$ vers les probabilités d'action transformées par le rapport logarithme (en utilisant une transformation de rapport logarithme additif par rapport à une action de référence). Les coefficients sont estimés par régression ridge régularisée.
3. Décomposition analytique : Les coefficients PCE ajustés permettent le calcul analytique des indices de sensibilité de Sobol. Contrairement aux méthodes de Monte Carlo qui nécessitent une intégration coûteuse, ces indices sont dérivés directement des coefficients polynomiaux, fournissant une décomposition exacte basée sur la variance de la contribution de chaque composante de l'incertitude de la récompense à la variance de la distribution de la politique.

Le cadre inclut des garanties formelles de convergence pour le substitut et ses indices de Sobol. Quatre des cinq propriétés théoriques clés (convergence du PCE, convergence de l'indice de Sobol, préservation du simplexe et bornes de propagation de l'incertitude) ont été formellement vérifiées dans l'assistant de preuve Lean 4.

Contributions clés

1. Cadre général : Une méthode computationnellement efficace pour propager et décomposer l'incertitude épistémique dans les modèles génératifs séquentiels conditionnés en utilisant des substituts PCE ajustés sur de petits ensembles.
2. Vérification formelle : Garanties de convergence théorique pour le substitut et ses indices de sensibilité, avec quatre des cinq théorèmes formellement vérifiés dans Lean 4.
3. Décomposition interprétable de la sensibilité : La première méthode analytique pour décomposer l'incertitude de la politique dans les modèles génératifs, révélant exactement quels composants de la récompense pilotent des décisions génératives spécifiques. Cette capacité n'est pas disponible dans les ensembles profonds standards, les réseaux de neurones bayésiens ou le dropout Monte Carlo.
4. Validation empirique : Démonstrations sur sept tâches de complexité croissante, incluant trois applications réelles et quatre tâches de validation synthétiques.

Résultats et applications

Le cadre a été évalué sur trois domaines réels principaux :

• Optimisation des réactions chimiques (Buchwald–Hartwig) : Sur le jeu de données Doyle–Dreher, le cadre a analysé la sélection du catalyseur, de la base, de l'halogénure d'aryle et de l'additif. Il a révélé que la sélection du catalyseur est robuste (faible indice de Sobol, $D \approx 71$), indiquant un signal chimique authentique, tandis que la sélection de l'additif est fragile (indice de Sobol élevé, $D \approx 179$, soit 2,5 fois plus élevé), indiquant une susceptibilité aux biais de l'ensemble de données d'entraînement. Cela répond directement au « débat Chuang–Keiser » concernant les artefacts statistiques dans les prédicteurs de rendement.
• Découverte causale bayésienne (Signalisation des protéines Sachs) : Appliqué au réseau Sachs, l'analyse a montré que, bien que toutes les arêtes aient présenté une variance élevée en raison de données limitées (200 observations), l'incertitude se décomposait en régimes distincts. Les arêtes de la voie MAPK étaient pilotées par une composante principale de la variation de covariance, tandis que les arêtes centrales PKA/PKC étaient pilotées par une autre. Cela fournit une carte ciblée pour les expériences de perturbation.
• Conception moléculaire basée sur des fragments : Dans une tâche de conception de propriétés de type médicament, le cadre a renversé l'hypothèse conventionnelle selon laquelle les échafaudages sont robustes et les décorations fragiles. Les résultats ont montré que la position du lien était le composant le plus fragile ($D \approx 28,2$), tandis que les positions des échafaudages et des décorations étaient plus robustes. Cela suggère que les efforts expérimentaux devraient prioriser la diversité des liens.

Dans les tâches synthétiques (mondes en grille, régression symbolique et proxies de raisonnement LLM), le substitut PCE a atteint une couverture de calibration de 95 % entre 0,97 et 1,00 pour les étapes dominantes. Sur le plan computationnel, le substitut a évalué 10 000 échantillons de politique en millisecondes, offrant une accélération de $10^3$ à $10^4\times$ par rapport au réentraînement exhaustif de l'ensemble.

Importance et revendications

L'article revendique que ce cadre offre une capacité unique : la décomposition interprétable de l'incertitude. Alors que d'autres méthodes peuvent affirmer qu'une prédiction est incertaine, cette approche explique pourquoi en attribuant la variance à des composants de récompense spécifiques. Cela permet aux scientifiques de :

• Distinguer les signaux robustes des artefacts statistiques dans les données d'entraînement.
• Identifier quelles parties d'un espace de conception (par exemple, les liens par rapport aux échafaudages) nécessitent une validation expérimentale supplémentaire.
• Allouer efficacement les ressources expérimentales en ciblant les sources spécifiques d'incertitude épistémique.

Les auteurs soulignent que la méthode est agnostique à l'architecture du GFlowNet et à l'objectif d'entraînement. Ils notent des limitations, telles que la dépendance à une structure de faible rang dans le paysage de récompense pour que la PCA soit efficace et la nécessité de réajuster le substitut pour différentes trajectoires (bien que sans réentraîner l'ensemble). Ce travail se positionne comme un pont entre la quantification de l'incertitude en ingénierie computationnelle et l'IA générative moderne, offrant une voie évolutive pour comprendre la « boîte noire » de la prise de décision séquentielle dans l'incertitude.