Differentiable programming for general Bayesian decision theory

Question

It is my understanding that differentiable programming and thus libraries like TensorFlow (e.g. TFP) and JAX can be used to solve Bayesian decision theory problems where e.g. we have a probabilistic model of the loss we may incur based on a decision we need to make with information available. After all, these libraries should be able to help us optimize general and potentially large non-convex problems.

Broadly speaking, I'm referring to general problems of the form:

$\underset{d}{\operatorname{argmin}} \mathrm{E}^\pi\left[L\left(\theta,d\right)| \text{D}\right]$

where we have the standard notation:

• $L$ is the loss function
• $\pi$ is our prior over parameters of the model
• $d$ is the decision we are trying to make
• $\text{D}$ is our data or current knowledge

However, I haven't encountered yet any actual examples showing this particular connection where one flexibly defines the Loss function using probabilistic programming and optimize the decision problem via differentiable programming.

This may be because these libraries seem to be primarily geared towards learning from data and inference of model parameters. This is in contrast to e.g. optimizing the prices of items in a store, a portfolio in finance, or, say, the sample sizes for an expensive A/B/n test that are required be set a priori.

That said, there may be something else at play, perhaps the difficulty of solving the equation above for loss functions that can be much more general, and complex than e.g. NLL or 0/1 loss (MAP), gets in the way.

In other words, I'm wondering about the use of diff. programming today for optimizing general Bayesian Loss functions, i.e. using it for broader problems than just e.g. fitting a model, or running a PPL sampler per se, as in e.g. HMC for estimating model parameters (often a "standard" loss function), or obviously computing posteriors.

Can one actually fully model and solve these types of general problems today with e.g. TFP, JAX, PyTorch, etc? If not (or not so easily), why?

Perhaps motivated by the above, what type of software is otherwise available today to solve problems of the general form above, where the model and random variables can be specified flexibly via e.g. probabilistic programming?

Answer 1

TensorFlow Probability is a standalone probabilistic programming module for TensorFlow, numpyro package uses JAX, while Pyro is a PyTorch framework. All those frameworks enable you to do Variational Inference and Markov Chain Monte Carlo sampling.

The simplest way of computing the expected posterior loss, is using Monte Carlo samples from the posterior distribution (either using MCMC, or first use VI estimate posterior distribution, then sample from it). To calculate expected posterior loss, you need to do is to take samples from the posterior predictive distribution, calculate loss function over those values, and average the losses.

Frameworks that you describe (TensorFlow, JAX, PyTorch) are designed for using automatic differentiation, not for taking integrals, what is needed for Bayesian problems. You can find some packages that are focused on taking integrals, e.g. using Bayesian quadrature, but they are far less mature then the mentioned packages, so the procedure described by me above for now remains state-of-art for Bayesian problems.