Surrogate-Based Quantification of Policy Uncertainty in Generative Flow Networks

Ramón Dineth Nartallo-Kaluarachchi, Robert Manson-Sawko, Shashanka Ubaru, Dongsung Huh, Malgorzata J. Zimon, Lior Horesh, Yoshua Bengio

Keywords: uncertainty quantification, GFlowNets, generative modelling, polynomial chaos expansions

TL;DR: Quantifying policy uncertainty in generative flow networks with uncertain reward via a PCE-surrogate model

Abstract: Generative flow networks are able to sample, via sequential construction, highreward, 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.

Primary Area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)

Submission Number: 24892