Surrogate-based quantification of policy uncertainty in generative flow networks

TL;DR

This paper tackles epistemic uncertainty in Generative Flow Networks (GFNs) arising when rewards are estimated from noisy data. It introduces a surrogate modeling framework based on Polynomial Chaos Expansions (PCE) trained on a small ensemble of GFNs to map low-dimensional reward representations to per-step policy distributions, enabling inexpensive Monte Carlo estimation of policy uncertainty along trajectories. The authors validate the approach across discrete and continuous grid-worlds, symbolic regression with additive noise, and Bayesian structure learning, showing the surrogate captures complex distributions (including bimodality) and agrees with empirical ensembles. This uncertainty-quantification framework enables uncertainty-aware generation with GFNs and holds promise for applications such as molecular design and causal-graph structure learning, while providing interpretable sensitivity analyses via Sobol indices.

Abstract

Generative flow networks are able to sample, via sequential construction, high-reward, complex objects according to a reward function. However, such reward functions are often estimated approximately from noisy data, leading to epistemic uncertainty in the learnt policy. We present an approach to quantify this uncertainty by constructing a surrogate model composed of a polynomial chaos expansion, fit on a small ensemble of trained flow networks. This model learns the relationship between reward functions, parametrised in a low-dimensional space, and the probability distributions over actions at each step along a trajectory of the flow network. The surrogate model can then be used for inexpensive Monte Carlo sampling to estimate the uncertainty in the policy given uncertain rewards. We illustrate the performance of our approach on a discrete and continuous grid-world, symbolic regression, and a Bayesian structure learning task.

Figures (6)

• Figure 1: Discrete grid-world with uncertain rewards. $a)$ Ground-truth reward function defined on 10 $\times$ 10 grid. $b)$ Example trajectory terminating at a high-reward location. $c)$ Random samples from the distribution over possible reward functions. $d)$ Policy along the trajectory from $b)$ for a single model. $e)$ Distribution over policies from an ensemble of 50 models trained on sampled variations of the reward grid. $f)$ Empirical testing ensemble of models compared with outputs of PCE surrogate model. The PCE model can capture the shape of the distribution, including bimodality.
• Figure 2: Continuous grid-world with uncertain rewards.$a)$ Ground-truth reward function with an example trajectory of 5 steps. $b)$ Samples of uncertain reward functions from a distribution. $c)$ Policies along the trajectory from Panel $a)$ in $x-$ and $y-$ direction (columns) for each of 5 steps (rows). Policies are taken from a testing ensemble of 100 gfn trained on samples from the reward distribution. The black line is the mean policy. $d)$ We plot the distribution of the policies sampled from the pce surrogate model along a trajectory, where redder colours corresponds to policies with higher probability density under the surrogate model. This experiment uses 50,000 samples.
• Figure 3: Symbolic regression of a function with additive noise.$a)$ The target function $f(x)=\sin(x)+2-x$ in the range $[\pi,4\pi]$, and a noisy function with additive Wiener noise. $b)$ Empirical testing policies compared with surrogate policies from a pce along a trajectory.
• Figure 4: Bayesian structure learning a linear Gaussian network. Top left. The ground-truth Gaussian network. Panels 2-8. We show the policy, a distribution over 26 actions, for each of 7 steps along a trajectory. The policy distribution of the 250 empirical testing samples is shown in orange, whilst the 10,000 surrogate samples are in teal. The dotted line shows the action taken at each step. Outlier values beyond fences are not shown, but full scatters are presented in Fig. \ref{['fig: bsl symb mlp']}.
• Figure 5: Comparison to an mlp for grid-worlds.$a)$ Testing samples compared with surrogate samples from a pce and mlp model for the discrete-grid world example from Sec. \ref{['subsec: disc grid']}. We see that both surrogate models are able to capture the relationship between inputs and outputs, and that the surrogate model with its additional samples builds up a distribution with greater variance. $b)$ Comparison for the continuous-grid world example from Sec. \ref{['subsec: cont grid']}. The mlp can learn the input-output relationship, but suffers from regression to the mean, i.e. the surrogate model predicts values which are close to the mean, thus underestimating the variance of the empirical samples shown in Fig. \ref{['fig: contworld']}.