📊 statistics

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

This paper introduces a framework that uses polynomial chaos surrogates to analytically decompose and interpret epistemic uncertainty in sequential generative models, enabling the identification of specific reward components driving generative decisions with significantly higher computational efficiency and formal verification than existing deep learning methods.

Original authors: Ramón Nartallo-Kaluarachchi, Shashanka Ubaru, Małgorzata J Zimoń, Dongsung Huh, Robert Manson-Sawko, Lior Horesh, Yoshua Bengio

Published 2026-05-19

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Ramón Nartallo-Kaluarachchi, Shashanka Ubaru, Małgorzata J Zimoń, Dongsung Huh, Robert Manson-Sawko, Lior Horesh, Yoshua Bengio

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a chef trying to invent a new recipe. You have a "taste-tester" (an AI) that tells you how good a dish might be. But here's the catch: your taste-tester is a bit shaky. Sometimes it loves a dish because of the salt; other times, it loves it because of the pepper. You don't know why the taste-tester is giving you a specific score, and you don't know if the recipe it suggests is actually good or if it just got lucky with the ingredients you happened to show it.

This paper introduces a new tool to solve that exact problem for AI scientists. It helps them figure out why an AI is making a specific choice and how much they should trust that choice, especially when the "rules" the AI is learning from are imperfect.

Here is a simple breakdown of how it works, using the analogies from the paper:

1. The Problem: The "Black Box" Ensemble

Usually, to handle uncertainty, scientists train many versions of the same AI (like asking 50 different chefs to make the same dish). If all 50 chefs agree on the recipe, you feel confident. If they all suggest different things, you know you're in trouble.

But there's a flaw: You know that they disagree, but you don't know why.

Is the disagreement because the salt is uncertain?
Is it because the heat is uncertain?
Or is it just random noise?

Standard methods can tell you "the recipe is shaky," but they can't point to the specific ingredient causing the shake.

2. The Solution: The "Crystal Prism" (Polynomial Chaos Surrogates)

The authors built a special "prism" (called a Polynomial Chaos Expansion or PCE) that sits between the AI and the data.

Think of the AI's decision-making process as a beam of white light (total uncertainty) entering a dark room.

Old Method: You just see the light is bright or dim. You know there's uncertainty, but you can't see the colors inside.
New Method (The Prism): The prism splits that white light into a rainbow. Suddenly, you can see exactly which "color" (which specific part of the data) is causing the most trouble.

This prism doesn't just guess; it uses math to analytically calculate exactly how much each part of the input (like the salt, the pepper, or the heat) contributes to the AI's final decision.

3. Real-World Examples from the Paper

The authors tested this prism on three very different "kitchens" to see if it worked:

A. The Chemistry Kitchen (Making Drugs)

The Task: An AI is trying to figure out the best mix of chemicals to create a new drug.
The Surprise: Everyone thought the "base structure" of the molecule was the most important and stable part.
What the Prism Found: The prism revealed the opposite! The "linker" (the piece connecting parts of the molecule) was actually the most fragile part. If the data about the linker was slightly off, the whole recipe changed. The base structure was actually quite robust.
The Takeaway: Scientists now know they need to get more data specifically about the linker, not the base, to make the AI more reliable.

B. The Biology Kitchen (Mapping Protein Signals)

The Task: An AI is trying to draw a map of how proteins talk to each other in a cell.
The Surprise: Some parts of the map were clear, while others were a mess.
What the Prism Found: It separated the map into two distinct zones. One zone (the "MAPK" pathway) was sensitive to one type of data error, while another zone (the "PKA/PKC" hubs) was sensitive to a different type.
The Takeaway: Instead of guessing where to do more experiments, scientists can now look at the prism's rainbow and see exactly which part of the map needs more data to become clear.

C. The Language Kitchen (Teaching AI to Reason)

The Task: An AI is learning to solve math problems step-by-step.
The Surprise: The AI sometimes stops early or keeps going too long.
What the Prism Found: It showed that early steps in the reasoning were stable, but later steps (where the AI has to choose between very similar-looking answers) were highly sensitive to uncertainty.
The Takeaway: The AI's "confidence" drops at specific moments, and the prism tells us exactly when and why.

4. Why This is a Big Deal

Speed: The prism is incredibly fast. It can simulate 10,000 different scenarios in a fraction of a second. Doing this by retraining the AI 10,000 times would take days. The prism does it in milliseconds.
Clarity: It doesn't just say "I'm unsure." It says, "I'm unsure because of this specific ingredient."
Trust: It helps scientists decide which AI suggestions are solid and which ones are just lucky guesses based on incomplete data.

Summary

This paper gives scientists a magnifying glass for AI uncertainty. Instead of just seeing a blurry picture of "maybe this, maybe that," the new tool splits the blur apart to show exactly which piece of information is causing the confusion. This allows researchers to fix the specific weak spots in their data, rather than guessing blindly.

Technical Summary: Interpretable Epistemic Uncertainty Decomposition in Sequential Generative Models via Polynomial Chaos Surrogates

Problem Statement

Sequential generative models, such as Generative Flow Networks (GFlowNets), are increasingly central to AI-driven scientific discovery, combinatorial optimization, and biological design. These models often operate under conditions of epistemic uncertainty, specifically when the reward functions guiding their generation are themselves uncertain (e.g., neural network surrogates trained on limited experimental data). While existing methods like deep ensembles and Bayesian neural networks can quantify the magnitude of prediction uncertainty, they fail to provide an interpretable decomposition of that uncertainty. Consequently, practitioners cannot determine which specific components of the input reward uncertainty drive which generative decisions, nor can they distinguish between robust chemical signals and statistical artifacts arising from training set biases.

Methodology

The authors propose a surrogate modeling framework that propagates epistemic uncertainty through GFlowNets using Polynomial Chaos Expansions (PCE). The approach proceeds in three stages:

Ensemble Training and Parameterization: A modest ensemble of GFlowNets is trained under varied reward conditions. The reward function space is parameterized into a low-dimensional latent space $\mu \in \mathbb{R}^d$ using Principal Component Analysis (PCA) or analytical expansions (e.g., Karhunen–Loève). PCA is preferred because its orthogonal components approximate independent Gaussian inputs, a prerequisite for valid analytical Sobol decomposition.
PCE Surrogate Fitting: For each action $k$ at trajectory step $t$ , a PCE is fitted to map the low-dimensional reward parameters $\mu$ to the log-ratio-transformed action probabilities (using an Additive Log-Ratio transform relative to a reference action). The coefficients are estimated via regularized ridge regression.
Analytical Decomposition: The fitted PCE coefficients enable the analytical computation of Sobol sensitivity indices. Unlike Monte Carlo methods which require expensive integration, these indices are derived directly from the polynomial coefficients, providing an exact variance-based decomposition of how much each component of reward uncertainty contributes to the variance in the policy distribution.

The framework includes formal convergence guarantees for the surrogate and its Sobol indices. Four of five key theoretical properties (PCE convergence, Sobol index convergence, simplex preservation, and uncertainty propagation bounds) have been formally verified in the Lean 4 proof assistant.

Key Contributions

General Framework: A computationally efficient method for propagating and decomposing epistemic uncertainty in conditioned sequential generative models using PCE surrogates fitted to small ensembles.
Formal Verification: Theoretical convergence guarantees for the surrogate and its sensitivity indices, with four of five theorems formally verified in Lean 4.
Interpretable Sensitivity Decomposition: The first analytical method to decompose policy uncertainty in generative models, revealing exactly which reward components drive specific generative decisions. This capability is unavailable from standard deep ensembles, Bayesian neural networks, or Monte Carlo dropout.
Empirical Validation: Demonstrations across seven tasks of increasing complexity, including three real-world applications and four synthetic validation tasks.

Results and Applications

The framework was evaluated on three primary real-world domains:

Chemical Reaction Optimization (Buchwald–Hartwig): On the Doyle–Dreher dataset, the framework analyzed catalyst, base, aryl halide, and additive selection. It revealed that catalyst selection is robust (low Sobol index, $D \approx 71$ ), indicating genuine chemical signal, while additive selection is fragile (high Sobol index, $D \approx 179$ , 2.5 $\times$ higher), indicating susceptibility to training-set bias. This directly addresses the "Chuang–Keiser debate" regarding statistical artifacts in yield predictors.
Bayesian Causal Discovery (Sachs Protein Signaling): Applied to the Sachs network, the analysis showed that while all edges exhibited high variance due to limited data (200 observations), the uncertainty decomposed into distinct regimes. MAPK-pathway edges were driven by one principal component of covariance variation, while PKA/PKC hub edges were driven by another. This provides a targeted map for perturbation experiments.
Fragment-Based Molecular Design: In a drug-likeness design task, the framework overturned the conventional assumption that scaffolds are robust and decorations are fragile. The results showed the linker position was the most fragile component ( $D \approx 28.2$ ), while scaffold and decoration positions were more robust. This suggests experimental efforts should prioritize linker diversity.

In synthetic tasks (grid-worlds, symbolic regression, and LLM reasoning proxies), the PCE surrogate achieved 95% calibration coverage between 0.97 and 1.00 for dominant steps. Computationally, the surrogate evaluated 10,000 policy samples in milliseconds, offering a speedup of $10^3$ to $10^4\times$ compared to exhaustive retraining of the ensemble.

Significance and Claims

The paper claims that this framework offers a unique capability: interpretable uncertainty decomposition. While other methods can state that a prediction is uncertain, this approach explains why by attributing variance to specific reward components. This allows scientists to:

Distinguish between robust signals and statistical artifacts in training data.
Identify which parts of a design space (e.g., linkers vs. scaffolds) require additional experimental validation.
Allocate experimental resources efficiently by targeting the specific sources of epistemic uncertainty.

The authors emphasize that the method is agnostic to the GFlowNet architecture and training objective. They note limitations, such as the reliance on low-rank structure in the reward landscape for PCA to be effective and the need for re-fitting the surrogate for different trajectories (though not retraining the ensemble). The work positions itself as a bridge between computational engineering uncertainty quantification and modern generative AI, providing a scalable path to understanding the "black box" of sequential decision-making under uncertainty.