Beyond Bayesian Model Averaging over Paths in Probabilistic Programs with
Stochastic Support

Tim Reichelt1 Luke Ong1,2 Tom Rainforth1

1 University of Oxford 2 Nanyang Technological University

###### Abstract

The posterior in probabilistic programs with stochastic support decomposes as
a weighted sum of the local posterior distributions associated with each
possible program path. We show that making predictions with this full
posterior implicitly performs a Bayesian model averaging (BMA) over paths.
This is potentially problematic, as model misspecification can cause the BMA
weights to prematurely collapse onto a single path, leading to sub-optimal
predictions in turn. To remedy this issue, we propose alternative mechanisms
for path weighting: one based on _stacking_ and one based on ideas from _PAC-
Bayes_. We show how both can be implemented as a cheap post-processing step on
top of existing inference engines. In our experiments, we find them to be more
robust and lead to better predictions compared to the default BMA weights.

##  1 Introduction

Universal probabilistic programming systems (PPS) (Tolpin et al., 2016;
Goodman et al., 2008; Bingham et al., 2019; Ge et al., 2018; Mansinghka et
al., 2014) provide flexible frameworks for expressing powerful probabilistic
models, along with tools to aid performing inference in them. By permitting
branching on the outcomes of sampling statements, they allow users to express
programs with _stochastic support_ , wherein the number of latent variables
varies between program executions, leading to challenging inference problems.

Such programs can be thought of as a combination of independent sub-programs,
each with static support, known as straight-line programs (SLP) (Chaganty et
al., 2013; Sankaranarayanan et al., 2013; Luo et al., 2021). The overall
posterior is then given by the weighted sum of individual SLP posteriors, with
weights corresponding to the local normalization constants of the SLPs; a
breakdown recent work has exploited to improve inference (Zhou et al., 2020;
Reichelt et al., 2022a).

We show that this decomposition also reveals that the posterior of _any_
program with stochastic support is a Bayesian Model Averaging (BMA) (Hoeting
et al., 1999) over the constituent SLPs of the program. Thus, all PPS
inference engines are implicitly estimating a BMA when the program has
stochastic support, whether they explicitly account for this or not.

However, it is widely acknowledged in the Bayesian statistics literature that
BMA can be a problematic mechanism for combining the posteriors of individual
models (Minka, 2000; Yao et al., 2018), with alternatives often preferred in
practice, especially when our aim is to make good predictions. In particular,
BMA often performs poorly under _model misspecification_ (Gelman and Yao,
2020; Oelrich et al., 2020), wherein it tends to produce _overconfident_
posterior model weights that collapse towards a single model (Huggins and
Miller, 2021; Yang and Zhu, 2018). Given that models will rarely be perfect
when working with real data (Box, 1976; Key et al., 1999; Vehtari and Ojanen,
2012), this is a serious practical concern that has been observed to cause
notable issues in many applied fields (Yang and Zhu, 2018; Smets and Wouters,
2007; Leff et al., 2008).

We argue that PPSs need to account for these shortfalls and provide access to
more robust weighting schemes. To provide such alternatives, we suggest
optimizing the SLP weights for _predictive performance_. Specifically, we
introduce weighting schemes based on _stacking of predictive distributions_
(Wolpert, 1992; Breiman, 1996; LeBlanc and Tibshirani, 1996; Yao et al., 2018)
and _PAC-Bayes objectives_ (Masegosa, 2020; Masiha et al., 2021; Morningstar
et al., 2022; Alquier, 2023). We show how to run them as a cheap post-
processing step on the outputs of any sample-based inference scheme, and
demonstrate that they provide more robust weights with better predictive
performance.

In summary, our contributions are: (a) Showing that the posterior in programs
with stochastic support can be interpreted as a BMA, thereby revealing
shortcomings of the standard mechanism to weight SLPs inherent to all
inference algorithms. (b) Providing a general scheme to adapt PPS inference
algorithms to utilize alternative weighting schemes, with an implementation in
Pyro (Bingham et al., 2019). (c) Investigating their behaviour for a variety
of different programs and showing its benefits on synthetic and real-world
data.

##  2 Background

###  2.1 Bayesian Model Averaging

An important question in Bayesian statistics is how to best combine the
inferences of different possible models. In a pure Bayesian framework, this is
done by weighting the model posteriors according to their _posterior model
probability_ , leading to a framework called Bayesian model averaging (BMA)
(Hoeting et al., 1999).

To be more precise, assume we have a countable set of Bayesian models indexed
by $k$, each with corresponding latent parameters $\theta_{k}\in\Theta_{k}$,
prior $p_{k}(\theta_{k})$, and likelihood $p_{k}(y|\theta_{k})$, where $y$ is
data we want to condition on.111Note our formulations apply equally when there
are also inputs the model is conditioned on, i.e. we have
$p_{k}(y|\theta_{k},x)$, but we negate this from our notation to avoid
clutter. In BMA, we set a prior over which model generated the data, $p(M=k)$,
from which we can derive the posterior model probability

$p(M{=}k\mid y)\propto p(y\mid M{=}k)\,p(M{=}k)$ (1)

where $p(y\mid M{=}k)=\int p_{k}(y|\theta_{k})\,p_{k}(\theta_{k})d\theta_{k}$
is the _model evidence_ , or marginal likelihood, for the $k$th model.

Predictions and expectations can now be calculated by combining those from
individual models using $p(M{=}k\mid y)$ as weights. In particular, the
posterior predictive distribution for new hypothetical data, $\tilde{y}$, is
given by

$p(\tilde{y}\mid y)\\!=\\!\sum\nolimits_{k}p(M{=}k\mid
y)\,\mathbb{E}_{p_{k}(\theta_{k}|y)}\left[p_{k}(\tilde{y}|\theta_{k})\right],$
(2)

where $p_{k}(\theta_{k}|y)\propto p_{k}(\theta_{k})p_{k}(y|\theta_{k})$ and
$p_{k}(\tilde{y}|\theta_{k})$ are the local posterior and local parameterized
predictive distribution, respectively.

Criticisms of BMA In practice, our models will never be able to capture the
full complexities of the real world as, in the words of George Box, “all
models are wrong, some are useful” (Box, 1976). It is therefore important to
investigate the behaviour of frameworks when our models are misspecified, that
is when
$\nexists~{}(\theta_{k},k):p_{k}(y|\theta_{k})=p_{\text{true}}(y)~{}\forall
y$, where $p_{\text{true}}(y)$ is the (unknown) true data generating
distribution.

Crucially, BMA implicitly assumes that the data was sampled from exactly one
of the constituent models. This is often referred to as the
$\mathcal{M}$-closed assumption (Bernardo and Smith, 2009; Clyde and Iversen,
2013; Key et al., 1999). As a result, as the amount of data increases the BMA
weights will always (except for a few special edge cases) collapse on a single
model (Clyde and Iversen, 2013); the approach reverts to just performing model
selection. Consequently, BMA predictions are often inferior compared to other
model combination techniques (Minka, 2000; Yao et al., 2018).

Viewed another way, model misspecification tends to lead to posterior model
probabilities that are _overconfident_ : both empirical and theoretical
results have shown that they too readily collapse on a single model (Huggins
and Miller, 2021; Yang and Zhu, 2018), even when there are multiple plausible
models with similar predictive performance. Moreover, the exact model onto
which the posterior collapses can change drastically when regenerating the
data from $p_{\text{true}}(y)$. Note that these are properties of the
_analytic_ BMA weights and not artifacts of using approximate inference
algorithms.

###  2.2 Programs with Stochastic Support

A probabilistic program can be interpreted as defining an _unnormalized
density function_ $\gamma:\Theta\rightarrow\mathbb{R}^{\geq 0}$, where
$\Theta$ denotes the sample space of the latent variables in the program
(Borgström et al., 2016; Staton et al., 2016). These variables are typically
defined as the outcomes of random sampling statements and each sample
statement is associated with a unique lexical address. The goal of inference
is then to find a representation of the normalized program density
$\pi(\theta)=\gamma(\theta)/\int\gamma(\theta)d\theta$, where $d\theta$ is an
implicitly defined reference measure (Gordon et al., 2014; Rainforth, 2017;
van de Meent et al., 2018). One can informally think of $\pi(\theta)$ as a
posterior distribution, $p(\theta|y)$. See App. A for a more detailed
introduction.

Universal PPS allow users to branch on the outcomes of random sampling
statements leading to programs with _stochastic support_. An important
property of such programs is that they can be decomposed into (a countable
number of) straight-line programs (SLPs), sub-programs without any control
flow (Chaganty et al., 2013; Sankaranarayanan et al., 2013; Zhou et al., 2020;
Luo et al., 2021; Reichelt et al., 2022a). These SLPs are effectively the
different possible control-flow paths that exist in the program and they are
defined by their _address path_ , i.e. the sequence of the lexical addresses
encountered during the program’s execution.

Each SLP corresponds to a disjoint sub-region, $\Theta_{k}$, of the overall
sample space (such that $\Theta=\bigcup_{k}\Theta_{k}$) and has the _local
unnormalized density_
$\gamma_{k}(\theta):=\mathbb{I}\left[\theta\in\Theta_{k}\right]\gamma(\theta)$.
The unnormalized density for the whole program can thus be written as
$\gamma(\theta)=\sum_{k}\gamma_{k}(\theta)$. Similarly, the normalized program
density can be rewritten as

$\pi(\theta)=\sum\nolimits_{k}\left(Z_{k}\big{/}\sum\nolimits_{\ell}Z_{\ell}\right)\,\pi_{k}(\theta),$
(3)

where $Z_{k}=\int\gamma_{k}(\theta)d\theta$ and
$\pi_{k}(\theta)=\gamma_{k}(\theta)/Z_{k}$ are the _local normalization
constant_ and _local posterior_ respectively. We refer to $\pi(\theta)$ as the
_full Bayes posterior_. Note that the disjoint supports of the SLPs means that
there exists exactly one $k:\pi_{k}(\theta)>0$ for any given $\theta$.

##  3 Full Inference in Programs with Stochastic Support is BMA

Examining Eq. (3), we immediately see that $\pi(\theta)$ is a weighted sum of
localized posteriors. Our novel insight is that this decomposition also
reveals that using the full Bayes posterior to calculate predictions or
expectations is implicitly performing a BMA over the individual SLPs. To see
this, consider calculating the expectation of some parameterized predictive
density $p(\tilde{y}|\theta)$:

$\displaystyle\mathbb{E}_{\pi(\theta)}\left[p(\tilde{y}|\theta)\right]=\sum_{k}\frac{Z_{k}}{\sum\nolimits_{\ell}Z_{\ell}}\mathbb{E}_{\theta_{k}\sim\pi_{k}}\left[p_{k}(\tilde{y}|\theta_{k})\right],$
(4)

where $p_{k}$ is any conditional density function such that
$p_{k}(\tilde{y}|\theta)=p(\tilde{y}|\theta)\,\forall\tilde{y},\theta\in\Theta_{k}$,
and we have defined new random variables $\theta_{k}$ drawn from the local
posterior of the $k$th SLP. We thus have that the downstream posterior
predictive on $\tilde{y}$ is a weighted sum of the posterior predictives that
would result from using the $k$th SLP instead of our full program.

There is now a clear analog between Eq. (4) and Eq. (2). To show that the
former corresponds to a BMA, all that remains is to show that the weights can
be interpreted as posterior model probabilities. At a high-level, this follows
simply from the fact that the $Z_{k}$ are analogous to (unnormalized)
posterior model probabilities (note, though, they are _not_ analogous to the
model evidences).

To be more precise, consider the factorization
$\gamma(\theta)=f(\theta)g(\theta)$ where $f(\theta)$ corresponds to all terms
from the sampling and $g(\theta)$ all terms from conditioning statements, such
that we can think of them as prior and likelihood components respectively. The
prior probability of choosing the $k$th SLP is now given by $P_{k}\\!:=\\!\int
f(\theta)\mathbb{I}\left[\theta\in\Theta_{k}\right]d\theta$, while the “model
evidence” is

$\displaystyle
E_{k}:=\int\frac{f(\theta)\mathbb{I}\left[\theta\in\Theta_{k}\right]}{P_{k}}g(\theta)d\theta.$
(5)

The posterior model probability is then equal to

$\displaystyle\frac{P_{k}E_{k}}{\sum_{\ell}P_{\ell}E_{\ell}}\\!=\\!\frac{\int
f(\theta)\mathbb{I}\left[\theta\in\Theta_{k}\right]g(\theta)d\theta}{\sum_{\ell}\int
f(\theta)\mathbb{I}\left[\theta\in\Theta_{\ell}\right]g(\theta)d\theta}\\!=\\!\frac{Z_{k}}{\sum_{\ell}Z_{\ell}}.$
(6)

Thus Eq. (4) is a BMA with model prior $p(M=k)=P_{k}$, local posteriors
$p_{k}(\theta_{k}|y)=\pi_{k}(\theta_{k})$, model evidences $p(y|M=k)=E_{k}$,
and identical local parameterized predictive distributions
$p_{k}(\tilde{y}|\theta_{k})=p(\tilde{y}|\theta=\theta_{k})$.

Having realized that using the full Bayes posterior leads to BMA, we can
instead define a generalized model averaging scheme over the SLPs that is
explicitly parameterized by a learnable set of weights, $w$:

$\tilde{\pi}(\theta;w):=\sum\nolimits_{k}w_{k}\,\pi_{k}(\theta)$ (7)

with $\sum_{k}w_{k}=1,w_{k}\geq 0$. This opens the door for considering
alternative approaches that avoid the shortfalls of BMA. We refer to
$w_{k}\propto Z_{k}$ as the _BMA weights_ ; the choice made by all current
inference engines.

At this point, a critical reader might argue that all Bayesian inference in
general, and not just the weights in a BMA, is sensitive to model
misspecification. However, averaging over a finite discrete set of models has
been highlighted as a special case in which Bayesian inference can give
counter-intuitive results and is especially susceptible to misspecification
(Yao et al., 2018; Gelman and Yao, 2020; Oelrich et al., 2020). This thus
motivates treating the SLP weights as explicit parameters that we may wish to
set in a non-Bayesian manner, while leaving the local posteriors unchanged.

##  4 Weighting Program Paths Using Predictive Objectives

There are different ways we can choose the SLP weights, $w_{k}$, in Eq. (7),
with BMA only one possible choice. A simple, albeit crude, alternative would
be to just set them equally. While we find that this can sometimes empirically
outperform the BMA weights (see Sec. 6), it is clearly not an appropriate
general-purpose solution and there are cases where it can perform very poorly.

To provide more principled alternatives, we now show how the weights can be
optimized to maximize predictive performance. For the purpose of exposition,
we will introduce the main ideas through the eyes of _stacking_ (Yao et al.,
2018; Wolpert, 1992; Breiman, 1996; LeBlanc and Tibshirani, 1996) but we will
show in Sec. 4.4 how we can also use PAC-Bayes objectives (Morningstar et al.,
2022) to fit the SLP weights.

###  4.1 Stacking Objective For PPSs

The goal of stacking is to improve predictions by optimizing the model
weights, $w$. To achieve this, we need to define a method for making
predictions for a hypothetical new observation, which we denote as
$\tilde{y}\in\mathcal{Y}$. Just as we previously defined a generalized version
of the posterior in Eq. (7), we now need to establish a generalized version of
the posterior predictive.

For simplicity, we will assume for now that an explicit predictive density,
$p(\cdot\mid\cdot):\mathcal{Y}\times\Theta\rightarrow\mathbb{R}^{\geq 0}$, has
been provided, before showing how this can instead be derived from the program
itself in Sec. 4.2. This ensures for each SLP we have a _local posterior
predictive density_

$\rho_{k}(\tilde{y}):=\mathbb{E}_{\pi_{k}(\theta)}\left[p(\tilde{y}\mid\theta)\right].$
(8)

With this, we can define the _stacked predictive density_

$\displaystyle\tilde{\rho}(\tilde{y};w):=\mathbb{E}_{\tilde{\pi}(\theta;w)}\left[p(\tilde{y}\mid\theta)\right]=\sum\nolimits_{k}w_{k}\,\rho_{k}(\tilde{y})$
(9)

In its most general form, stacking defines an objective with a user-defined
scoring rule $S(\tilde{\rho},\tilde{y})$ which takes as input a predictive
distribution and a data point (Gneiting and Raftery (2007); c.f. App. C). It
then optimizes the weights, $w$, to maximize the expected score

$R(w;S):=\mathbb{E}_{p_{\text{true}}(\tilde{y})}\left[S(\tilde{\rho}(\cdot\mid
w),\tilde{y})\right].$ (10)

We will focus on using the logarithmic score rule, as it is by far the most
popular one used in practice, yielding

$\displaystyle
R(w):=\mathbb{E}_{p_{\text{true}}(\tilde{y})}\left[\log\left(\sum\nolimits_{k}w_{k}\,\rho_{k}(\tilde{y})\right)\right].$
(11)

As $\mathbb{E}_{p_{\text{true}}(\tilde{y})}\left[\log
p_{\text{true}}(\tilde{y})\right]$ is a constant, maximizing $R(w)$ is
equivalent to minimizing
$\mathrm{KL}(p_{\text{true}}(\tilde{y})\parallel\tilde{\rho}(\tilde{y};w))$.

We now need a mechanism to estimate the expectation w.r.t.
$p_{\text{true}}(\tilde{y})$ in Eq. (11). Multiple strategies for this exist
(Vehtari and Ojanen, 2012). We will first show how to do so using an explicit
validation set, $\\{\tilde{y}_{\ell}\\}_{\ell=1}^{L}$, before describing how
to avoid the use of a validation set for a broad class of models in Sec. 4.3.

As an aside, the stacking weights should be interpreted differently from BMA
weights. In BMA, the weight of the $k$th SLP represents the posterior
probability that the data was generated from the $k$th SLP. In contrast, the
stacking weights are the result of optimizing a linear mixture for making
predictions on held-out data.

###  4.2 Stacking as Post-Processing

Given (normalised) weighted samples $\\{(v_{s},\theta_{s})\\}_{s=1}^{S}$
generated from an arbitrary inference algorithm, the posterior of the program
is approximated by the empirical measure
$\hat{\pi}(\theta)=\sum_{s=1}^{S}v_{s}\,\delta_{\theta_{s}}(\theta)$;
unweighted sampling schemes correspond to the special case $v_{s}=1/S$. The
local posteriors of the $k$th SLP are consequently approximated by all the
samples which fall into the $k$th SLP, i.e. $\hat{\pi}_{k}(\theta):=\sum_{s\in
I_{k}}(v_{s}/V_{k})\,\delta_{\theta_{s}}(\theta)$ where
$I_{k}:=\\{s\in\\{1,\dots,S\\}\mid\theta_{s}\in\Theta_{k}\\}$ are the indices
of the samples from the $k$th SLP and $V_{k}:=\sum_{s\in I_{k}}v_{s}$ is the
sum of all the associated sample weights. Recall from Sec. 2.2 that the SLP of
a sample $\theta_{s}$ is determined by its address path. Thus, we can generate
the index sets $I_{k}$ by grouping all samples with the same path.

Algorithm 1 Stacking as Post-Processing (Sec. 4.2)

1:Program $\gamma$, Weighted samples from base inference procedure
$\\{(v_{s},\theta_{s})\\}_{s=1}^{S}$

2:For each $\theta_{s}$, record address path and return values
$\\{g(\tilde{y}_{\ell}\mid\theta_{s})\\}_{\ell=1}^{L}$ (c.f. §4.2)

3:Partition sample indices $1,\dots,S$ into subsets $\\{I_{k}\\}_{k=1}^{K}$
using address paths (c.f. §4.2)

4:Compute predictive densities $\hat{\rho}_{k}(\tilde{y}_{\ell})$ (Eq. (14))

5:Compute $w^{*}=\operatorname*{argmax}\hat{R}(w)$ (Eq. (15))

6:Compute new sample weights $\omega_{s}$ (Eq. (16))

7:return $\\{(\omega_{s},\theta_{s})\\}_{s=1}^{S}$

The full Bayes posterior hence implicitly uses the approximation
$\hat{\pi}(\theta)=\sum_{k}V_{k}\,\hat{\pi}_{k}(\theta)$, assigning the weight
$V_{k}$ to each SLP. We instead replace these with the learnable SLP weights
$w_{k}$:

$\displaystyle\tilde{\pi}(\theta;w)$
$\displaystyle\approx\sum_{k}w_{k}\,\hat{\pi}_{k}(\theta)=\sum_{k}\sum_{s\in
I_{k}}\frac{w_{k}\,v_{s}}{V_{k}}\,\delta_{\theta_{s}}(\theta).$ (12)

We can then use this approximation to get an estimate of the stacked
predictive density defined in Eq. (9)

$\displaystyle\tilde{\rho}(\tilde{y}_{\ell};w)\approx$
$\displaystyle\sum\nolimits_{k}w_{k}\,\hat{\rho}_{k}(\tilde{y}_{\ell})$ (13)
$\displaystyle\text{where}\quad\hat{\rho}_{k}(\tilde{y}_{\ell}):=$
$\displaystyle\sum\nolimits_{s\in
I_{k}}(v_{s}/V_{k})\,p(\tilde{y}_{\ell}\mid\theta_{s})$ (14)

are the local posterior approximations. In our implementation, the user
implicitly defines the $p(\tilde{y}_{\ell}\mid\cdot)$ through the program
return values. We can now approximate $R(w)$ using

$\hat{R}(w):=\frac{1}{L}\sum\nolimits_{\ell=1}^{L}\log\left(\sum\nolimits_{k}w_{k}\,\hat{\rho}_{k}(\tilde{y}_{\ell})\right).$
(15)

Note that, as the $\hat{\rho}_{k}(\tilde{y}_{\ell})$ do not depend on the SLP
weights, we can precompute these estimates before optimizing $w$ in a
separate, typically cheap, procedure.

After having obtained the optimized weights,
$w^{*}=\operatorname*{argmax}_{w}\hat{R}(w)$, we want to be able to obtain
estimates w.r.t. the reweighted normalized density
$\tilde{\pi}(\theta;w^{*})$. This can be done easily by reweighting the
individual posterior samples $\theta_{s}$. Letting $k(\theta_{s})$ denote the
SLP index of sample $\theta_{s}$ we can rewrite Eq. (12) as

$\tilde{\pi}(\theta;w)\approx\sum\nolimits_{s=1}^{S}\omega_{s}\,\delta_{\theta_{s}}(\theta);~{}\omega_{s}:=\frac{w_{k(\theta_{s})}\,v_{s}}{V_{k(\theta_{s})}}.$
(16)

Alg. 1 summarizes the high-level steps of our post-processing stacking
procedure. Given an input program and corresponding (weighted) posterior
samples, we first use the program to extract the address path and return
values for each sample $\theta_{s}$ (e.g. using the Trace data structure in
Pyro). Then, we compute the index subsets $I_{1},\dots,I_{K}$ by grouping
unique address paths together. We can then compute the estimate of the local
posterior predictive densities $\hat{\rho}_{k}(\tilde{y}_{\ell})$ and store
the estimates in a $K\times L$ matrix. Finally, this matrix can be used to
evaluate our stacking objective $\hat{R}(w)$ and thus optimize the weights.
This optimization can be done cheaply relative to the cost of inference as it
is a convex optimization of a small number of parameters and does not require
further inferences; we use the L-BFGS-B algorithm (Byrd et al., 1995; Zhu et
al., 1997) for this. In App. E we provide further details on our full
implementation in Pyro.

###  4.3 Stacking Without Validation Sets

So far, we have assumed the existence of an explicit predictive density $p$
and held-out data $\\{\tilde{y}_{\ell}\\}_{\ell=1}^{L}$, but neither is often
actually needed to utilize stacking. Specifically, if we assume that the local
unnormalized SLP densities model the observed data $y_{i}\in\mathcal{Y}$ as
conditionally independent given parameters $\theta$, then for each SLP we can
write the local unnormalized density as

$\gamma_{k}(\theta,y_{1:N})=\mathbb{I}\left[\theta\in\Theta_{k}\right]f(\theta)\prod\nolimits_{i=1}^{N}g(y_{i}\mid\theta),$

where $f_{k}:\Theta\rightarrow\mathbb{R}^{\geq 0}$ represents a prior density
on $\theta$. We then use this to derive the following leave-one out (LOO)
cross-validation estimator for Eq. (11),

$R_{\text{LOO}}(w)=\,\frac{1}{N}\sum\nolimits_{i=1}^{N}\log\sum\nolimits_{k}w_{k}\,\rho_{k}(y_{i}\mid
y_{-i}).$

where $\rho_{k}(y_{i}\mid y_{-i})$ denotes the local posterior predictive
density corresponding to $\gamma_{k}(\theta,y_{1:N}\setminus y_{i})$, i.e.
that results from removing the i-th observation term from $\gamma_{k}$.

Naively computing each of the predictive distributions $\rho_{k}(y_{i}\mid
y_{-i})$ would require running inference $N$ times to evaluate the stacking
objective. However, as shown in Yao et al. (2018), this can be avoided using
Pareto smoothed importance sampling to estimate the LOO densities (PSIS-LOO)
(Vehtari et al., 2017). Computing the PSIS-LOO approximations to the densities
$\rho_{k}(y_{i}\mid y_{-i})$ requires access the individual likelihood terms
$g(y_{i}\mid\theta_{s})$ for each posterior sample $\theta_{s}$ (c.f. App. D
). Luckily, these can be extracted automatically for many common PPSs, e.g.
using the loo function of the ArviZ library (Kumar et al., 2019) for Pyro
models.

###  4.4 Regularized Stacking and PAC-Bayes

Stacking directly optimizes an estimate of the expected predictive density on
held-out data (Eq. (11)). Such estimates are fundamentally based on a finite
amount of data and optimizing them directly can, at least in principle, lead
to overfitting. Our proposed remedy for this is to add an additional KL
regularization term inspired by PAC-Bayes objectives. Namely we consider

$\displaystyle R_{\beta}(w)$
$\displaystyle:=\,\frac{1}{L}\sum\nolimits_{\ell=1}^{L}\log\left(\sum\nolimits_{k}w_{k}\rho_{k}(\tilde{y}_{\ell})\right)$
$\displaystyle-(1/\beta
L)\,\mathrm{KL}(\text{Categorical}(w_{1},\dots,w_{K})\parallel r(k)),$

where $r(k)$ is a reference weighting we want to regularize towards. Since we
want to discourage the SLP weights from collapsing towards a single SLP, we
will generally take $r(k)$ to be the uniform distribution. The hyperparameter
$\beta$ controls the amount of regularization: $\beta\to\infty$ recovers the
standard stacking objective, and $\beta\to 0$ leads to the weights following
$r(k)$.

This regularized objective corresponds to a particular instantiation of a PAC-
Bayes bound that was proposed by Morningstar et al. (2022). In the PAC-Bayes
literature, $-R(w)$ (Eq. (11)) is sometimes referred to as the _true
predictive risk_. Hence, optimizing $R_{\beta}(w)$ can be viewed as optimizing
a stochastic bound on the true predictive risk, see App. G for details.

##  5 Related Work

(a) Computed weight for the first SLP; each dot is the weight computed in one
of the datasets

(b) $\mathrm{LPPD}_{\text{Diff}}$ on held-out data. Negative values imply
stacking is performing better. Diamonds show medians.

Figure 1:  Behaviour of the BMA and stacked weights in the models as described
in Sec. 6.1.

A host of alternatives to Bayesian model averaging have been proposed in the
Bayesian statistics literature. So-called Pseudo-BMA weights (Geisser and
Eddy, 1979) use LOO predictive densities to replace marginal likelihoods but
have been shown to work less well than stacking (Yao et al., 2018). Bayesian
model combination (BMC) (Minka, 2000; Monteith et al., 2011; Kim and
Ghahramani, 2012) aims to break the BMA assumption that the data was generated
from exactly one of the candidate models. However, fitting the new extended
model is significantly more expensive as the inference task no longer breaks
down into independent sub-problems. Yao et al. (2022) introduced Bayesian
hierarchical stacking, which infers different weights for different regions in
the covariate space (similar to the frequentist mixture of experts (Gormley
and Frühwirth-Schnatter, 2019)). BayesBag (Huggins and Miller, 2021) weights
models by generating bootstrapped datasets, and averaging the normalization
constant over datasets. Future work might look to implement some of these
approaches in PPS as further alternatives to BMA.

Warrell and Gerstein (2022) extend PAC-Bayes bounds to work with hierarchical
models inspired by deep probabilistic programs (Tran et al., 2017) and use
this extension to derive bounds for multi-task settings such as transfer and
meta-learning. However, they do not consider programs with stochastic support,
the impact of model misspecification, nor integrate their method with a PPS.
Elsewhere, PAC-Bayes style arguments have also been used to reason about the
hardness of posterior inference in PPS (Freer et al., 2010).

##  6 Experiments

We will now compare different weighting schemes on a range of models and
datasets. Our quantitative measure for comparison will be the average log
posterior predictive density (LPPD) on held-out data $\tilde{y}_{1:T}$, where
$\mathrm{LPPD}:=\frac{1}{T}\sum_{t=1}^{T}\log\tilde{\rho}(\tilde{y}_{t};w)$.
In particular, we focus on the difference in LPPD from other methods to
stacking,
$\mathrm{LPPD}_{\text{Diff}}=\mathrm{LPPD}_{\text{Other}}-\mathrm{LPPD}_{\text{Stacking}}$.

Additionally, we will investigate the behaviour of the SLP weights $w_{k}$; a
major criticism of the BMA weights is that they are too sensitive to minor
changes in the data. To ensure replicable analysis, we desire the SLP weights
to be _robust_ and _consistent_ , i.e. they should be similar across different
possible generated datasets.

Except when otherwise indicated, we use a variant of the divide, conquer, and
combine (DCC) inference algorithm (Zhou et al., 2020) for the base inference
algorithm. Our DCC implementation uses HMC (Neal, 2011; Hoffman et al., 2014;
Betancourt, 2018) for the local inference algorithm of each SLP and allocates
the computational budget uniformly between SLPs. Additionally, we also
consider the reversible-jump MCMC (RJMCMC, c.f. App. A) algorithm implemented
in Gen (Cusumano-Towner et al., 2019). Our implementation is available at
https://github.com/treigerm/beyond_bma_in_probprog.

###  6.1 When is Stacking Helpful?

Table 1:  $\mathrm{LPPD}_{\text{Diff}}$ ($\uparrow$ better) for models in Sec.
6.2, 6.3, and 6.4, results computed over 10 replications. Bold indicates no
significant difference to Stacked under a Wilcoxon signed-rank test.

Model | Stacked | Stacked (Val) | BMA | BMA (Analytic) | RJMCMC | Equal
---|---|---|---|---|---|---
Subset | $\mathbf{0.0}$ | $-0.01\pm 0.01$ | $-0.11\pm 0.05$ | $-0.11\pm 0.05$ | $-0.11\pm 0.04$ | $-0.02\pm 0.01$
Fun. Induction | $\mathbf{0.0}$ | $\mathbf{-1.73\mathrm{e}{-3}\pm 2.98\mathrm{e}{-3}}$ | $\mathbf{-9.10\mathrm{e}{-4}\pm 2.42\mathrm{e}{-3}}$ | N/A | $-0.08\pm 0.03$ | $-0.31\pm 0.01$
California | $\mathbf{0.0}$ | $-\mathbf{1.55\mathrm{e}{-3}\pm 2.89\mathrm{e}{-3}}$ | $-2.10\mathrm{e}{-2}\pm 6.19\mathrm{e}{-3}$ | $-2.10\mathrm{e}{-2}\pm 6.01\mathrm{e}{-3}$ | $-2.82\mathrm{e}{-1}\pm 1.19\mathrm{e}{-1}$ | $-1.96\mathrm{e}{-1}\pm 3.11\mathrm{e}{-3}$
Diabetes | $\mathbf{0.0}$ | $-\mathbf{8.22\mathrm{e}{-3}\pm 1.44\mathrm{e}{-2}}$ | $-\mathbf{1.01\mathrm{e}{-2}\pm 1.64\mathrm{e}{-2}}$ | N/A | $-3.83\mathrm{e}{-2}\pm 2.19\mathrm{e}{-2}$ | $-3.66\mathrm{e}{-2}\pm 9.97\mathrm{e}{-3}$
Stroke | $\mathbf{0.0}$ | $-\mathbf{7.79\mathrm{e}{-4}\pm 1.57\mathrm{e}{-3}}$ | $-6.22\mathrm{e}{-3}\pm 3.81\mathrm{e}{-3}$ | N/A | $-2.25\mathrm{e}{-1}\pm 9.94\mathrm{e}{-2}$ | $-1.31\mathrm{e}{-1}\pm 5.68\mathrm{e}{-3}$

First, to develop an understanding of the scenarios in which stacking is
beneficial, we consider three simple examples in which we limit ourselves to
input programs with two SLPs. Unless otherwise stated, BMA weights are
computed analytically and the stacking weights are based on PSIS-LOO. For each
problem, we generate $10^{3}$ datasets with 200 data points each and generate
another $10^{3}$ data points to evaluate the held-out LPPD.

Distinct SLPs. For the first setting, we assume the data is generated by
$y_{i}\sim\mathcal{N}(0,1)$. We consider a program with two, misspecified SLPs
where the unnormalized density for the $k$th SLP is given by
$\gamma_{k}(\theta_{1},\theta_{2})=\mathbb{I}\left[\theta_{2}=k\right]\frac{1}{2}\prod_{i=1}^{N}\mathcal{N}(y_{i};\theta_{1},\sigma^{2}_{k})\,\mathcal{N}(\theta_{1};0,1)$
where we set $\sigma^{2}_{1}=0.62177$ and $\sigma^{2}_{2}=2$ (cf. App. A for
the corresponding Pyro program). This example is adapted from Yang and Zhu
(2018).

SLPs with overlap. Next, we generate a dataset using the relation
$y_{i}=\sum_{d=1}^{4}\beta_{d}\,x_{i,d}+\epsilon_{i}$ where
$\epsilon_{i}\sim\mathcal{N}(0,1)$, $x_{i,d}\sim\mathcal{N}(0,1)$ and
$\beta=[1.5,1.5,0.3,0.1].$ For inference, we consider a program with two SLPs
with unnormalized densities
$\gamma_{k}(\theta_{1},\theta_{2},\theta_{3})=\mathbb{I}\left[\theta_{3}=k\right]\allowbreak\frac{1}{2}\prod_{i=1}^{N}\mathcal{N}(y_{i};f_{k}(\theta_{1},\theta_{2},x_{i}),1)\prod_{j=1}^{2}\mathcal{N}(\theta_{j};0,1)$,
where for the first SLP
$f_{1}(\theta_{1},\theta_{2},x_{i})=\theta_{1}x_{i,1}+\theta_{2}x_{i,3}$ and
for the second
$f_{2}(\theta_{1},\theta_{2},x_{i})=\theta_{1}x_{i,1}+\theta_{2}x_{i,4}$.
Note, that the covariates $x_{i,d}$ are modelled as fixed by the program. Both
SLPs are misspecified because they do not have access to all the covariates.
However, the two sub-models share complexity/expressiveness, as they both have
access to the first covariate.

Dominant SLP. Lastly, we generate data from $y_{i}=f(x_{i})+\epsilon_{i}$ with
$f(x)=2\,x_{i}+\sin(5\,x_{i})$ and $\epsilon_{i},x_{i}\sim\mathcal{N}(0,1)$.
Our program for inference has two SLPs which are defined as
$\gamma_{k}(\theta_{1},\theta_{2},\theta_{3})=\mathbb{I}\left[\theta_{3}=k\right]\allowbreak\frac{1}{2}\prod_{i=1}^{N}\mathcal{N}(y_{i};f_{k}(\theta_{1},x_{i}),\theta_{2}^{2})\,\mathcal{N}(\theta_{1};0,1)\,\Gamma(\theta_{2};1,1)$
where $f_{1}(\theta,x)=\theta\,x$, $f_{2}(\theta,x)=\sin(\theta\,x)$, and
$\Gamma(\theta;\alpha,\beta)$ is a Gamma distribution parameterized by shape
$\alpha$, and rate $\beta$. The first SLP will provide a significantly better
fit to the data as it is able to recover the dominant linear trend. We use
importance sampling to estimate the BMA weights (see App. F).

Results. The results are presented in Fig. 1. For the first two models the
stacked weights lead to better predictive performance and more robust weights.
While there is some variation in the stacked weights, their overall
distribution is _unimodal_ whereas the distribution of the BMA weights are
_bimodal_ instead. The tendency of the BMA weights to collapse on either 0 or
1 is exactly the _overconfident_ behaviour we described in Sec. 3. In the
setting with one dominant SLP, both the BMA and stacking weights lead to
similar predictive performance (with BMA doing slightly better) and
consistently collapse onto the dominant SLP; here this is desirable behaviour
as the first SLP is clearly superior.

###  6.2 Subset Regression

Figure 2:  SLP weights for problem in Sec. 6.2. Each dot represents the weight
of the corresponding SLP in the model. Results are computed over 10 generated
datasets.

To further outline the issues of the default BMA weights, we consider a
regression problem with data generated from a linear model of the form
$y_{i}=\epsilon_{i}+\sum_{d=1}^{15}\beta_{d}\,x_{i,d}$ where
$\epsilon_{i}\sim\mathcal{N}(0,1)$ and all the covariates $x_{n,d}$ are drawn
independently from $\mathcal{N}(5,1)$. Note, the covariates are sampled to
generate the synthetic data set, but they are modelled as fixed by the
program. The ground-truth values for the regression coefficients $\beta_{d}$
are set following a scheme used in Breiman (1996) and Yao et al. (2018) which
ensures all covariates are relevant for the prediction of $y_{n}$ (c.f. App.
F).

We compare multiple methods: 1. Stacked, using PSIS-LOO to compute weights,
$w$ (Sec. 4.3); 2. Stacked (Val), which uses an explicit validation set
instead (Sec. 4.2); 2. BMA, using the PI-MAIS algorithm (Martino et al., 2017)
to compute local normalization constants (this is the default weighting in
DCC); 3. BMA (Analytic), using analytic solutions for the local normalization
constants; 4. RJMCMC, implemented in Gen (Cusumano-Towner et al., 2019);222We
also have conducted experiments that run stacking on top RJMCMC and found that
this also led to improvements in predictive performance (c.f. App. F). Here,
we only present the stacking results with samples generated from DCC as this
gave the best base inference procedure. 5\. Equal, weights each SLP equally.
We use warm colors to present the results which use one of our proposed
alternative objectives to set SLP weights and cooler colours for methods which
target the BMA weights. Note BMA, BMA (Analytic), and RJMCMC all target the
posterior distribution in Eq. (3); in practice, differences between these
methods will arise due to differences in the quality of the posterior
approximation.

Our input program has 15 SLPs and each SLP only gets access to one of the
covariates so our overall model is misspecified. However, since every
covariate influences the targets $y_{i}$, each SLP is relevant for making good
predictions. We generate 50 different datasets from the true data generating
distribution, with 200 data points used to run our inferences and $10^{3}$
data points to evaluate the held-out log posterior predictive density. For the
Stacked (Val) we use half of the 200 data points for inference and use the
rest to estimate the stacking objective (c.f. Eq. (15)).

Tab. 1 shows that stacking outperforms all other methods in terms of
predictive performance. The BMA weights here actually provide even worse
predictions than weighting each SLP equally. Fig. 2 shows the behaviour of the
weights for the SLPs with $k=8$ and $k=12$ over different randomly generated
datasets (see Fig. 9 in the Appendix for others). Both BMA and BMA (Analytic)
exhibit clear signs of overconfidence as described in Sec. 2.1: the weights
often collapse onto a single SLP, but the exact SLP changes between datasets,
leading to a bimodal sampling distribution for the SLP weights. As expected,
RJMCMC produces qualitatively and quantitivaly similar results to the other
Bayesian weighting mechanisms. This is in contrast with the Stacked weights
which are more evenly spread. The Stacked (Val) weights behave qualitatively
similarly to using PSIS-LOO, but with slightly worse predictive performance.
This is likely due to the corresponding reduction in training set size.

###  6.3 Function Induction

Our next example investigates how well stacking scales to a larger number of
SLPs. Following Zhou et al. (2020), we use a probabilistic context-free
grammar (PCFG) to posit a model over functions. The PCFG has production rules
$e\rightarrow\\{x\,\mid\,\sin(a*e)\,\mid\,a*e+b*e\\}$, where $x$ is a terminal
symbol denoting an input, and $a$, $b$ are coefficients to be inferred. The
program recursively samples from the PCFG using samples from categorical
distributions to select production rules. Additionally, the program defines
latent variables for all the coefficients in a given expression. We generate
observations from the relation
$y_{i}=-x_{i}+2\,\sin(2\,x_{i}^{2})+\epsilon_{i}$ with
$\epsilon_{i}\sim\mathcal{N}(0,0.1^{2})$ and the inputs $x_{i}$ uniformly
sampled between -5 and 5. We generate 400 data points used for inference and
$10^{3}$ data points for evaluation. Note, the terminal symbol $x^{2}$ is not
contained within the PCFG of our model and therefore the grammar cannot
express the data generating function.

Tab. 1 shows that Stacked provides better predictions compared to all other
methods (note analytic BMA weights cannot be calculated here). Interestingly,
we find there that the BMA weights do not tend to collapse on a single SLP
here (c.f. App. F), which shows stacking can be beneficial for predictions in
such cases.

###  6.4 Variable Selection

Figure 3:  SLP weights for Sec. 6.5. X-tick labels indicate the different
modelling choices for $\alpha$ and $\beta$; the pattern is “$\alpha$ model
choice, $\beta$ model choice” with P = pooling, NP = no pooling, H =
hierarchical, and G = group-level predictor.

Table 2: $\mathrm{LPPD}_{\text{Diff}}$ for Radon model.

Method |  $\mathrm{LPPD}_{\text{Diff}}$ ($\uparrow$)
---|---
Stacked | $\mathbf{0.0}$
BMA | $-8.24\mathrm{e}{-3}\pm 1.20\mathrm{e}{-2}$
RJMCMC | $-5.99\mathrm{e}{-2}\pm 1.86\mathrm{e}{-2}$
Equal | $-1.88\mathrm{e}{-2}\pm 1.18\mathrm{e}{-2}$

Next, we apply stacking to real-world classification and regression tasks.
Here, we have a matrix of covariates $X\in\mathbb{R}^{N\times D}$ and targets
$y_{1:N}$, and we want to do variable selection, i.e. select a subset of the
features $\mathcal{D}\subseteq\\{1,\dots,D\\}$ to make predictions. This
problem of variable selection can be encoded as a probabilistic program with
stochastic support in which each SLP corresponds to one of the potential
subsets of the features $\mathcal{D}$. We consider three different datasets:
_California_ (regression) (Pace and Barry, 1997), _Diabetes_ (classification)
(Smith et al., 1988), and _Stroke_ (classification) (Kaggle, 2020). For the
regression, our model is a linear regression with conjugate priors, permitting
an analytic solution to the BMA weights. For the classification tasks, we use
a logistic regression model which does not permit an analytic solution. As the
true data generating process in this setting is unknown, we run each method on
different train-test splits to estimate the variation in the weights and
predictions.

Tab. 1 shows the LPPD values of the weighting schemes on different datasets.
The Stacked and Stacked (Val) weighting schemes generally give better
predictive performance compared to the alternatives. Overall, these results
show that stacking can be beneficial for predictive performance even on real-
world data.

###  6.5 Radon Contamination

Our final example considers the analysis of data about Radon contamination for
houses in different US counties (Gelman and Hill, 2006). We here only give a
high-level description of the dataset and model, full details in App. F. For
each house recorded in the dataset, we have radon measurements, $y_{i}$, as
well as, a covariate, $x_{i}\in\\{0,1\\}$, which indicates whether the
measurement was made in the basement ($x_{i}=0$) or first floor of the house
($x_{i}=1$). Our program for this dataset has at its core the regression
relation $y_{i}=\alpha+\beta\,x_{i}+\epsilon_{i}$ and the different SLPs in
the program make different assumptions for how to model the coefficients
$\alpha$ and $\beta$. For both, we can either: 1) Fit the same coefficient
across all counties; 2) Fit a separate coefficient $\alpha_{c}$ for each
county $c$; or 3) Have separate coefficients for each county but assume they
come from the same underlying population distribution. For the intercept term
we also consider a fourth option: using county-wide level uranium measurements
as a group-level predictor. The program considers all combinations of
modelling choices for $\alpha$ and $\beta$, i.e. it has $4\cdot 3=12$ SLPs.

Each SLP in this program can have potentially hundreds of latent variables and
exhibit complex posterior geometries, making this a good testing ground to
compare the different weighting schemes. In Tab. 2 we find that the Stacked
weights give better performance compared to the BMA weights (we do not
consider using a validation set here because some counties contain only a
handful of observations). Fig. 3 shows that the BMA weights tend to
concentrate on SLPs with modelling choices “H,H” or “G,H” and have a bimodal
sampling distribution. The Stacked weights are more robust, giving more
consistent results between different train-test splits and more conservative
weights. Notably, RJMCMC here collapses onto different SLPs than the BMA
weights. This is due to the fact that the RJMCMC struggles with SLPs which
have a large number of latent variables. In the limit of infinite samples, the
behaviour of RJMCMC and BMA will be identical but Fig. 3 illustrates nicely
that approximate inference algorithms might collapse onto different SLPs based
on the quality of their posterior approximation.

###  6.6 Impact of Regularization: PAC-Bayes

Figure 4:  Impact of regularization parameter $\beta$ on predictive
performance in the different models (higher is better). Plotted are mean and
standard deviation.

As we have shown in Sec. 4.4, the PAC-Bayes bound $R_{\beta}(w)$ offers an
alternative objective to fit the weights and can be interpreted as the
stacking loss with an added regularization term where the hyperparameter
$\beta$ controls the amount of regularization. Smaller values of $\beta$ push
the stacking weights closer to the uniform distribution over SLPs. In Fig. 4
we plot the effect of varying $\beta$ on the predictive performance on the
different models. In our experiments, values of $\beta$ below 1 tend to lead
to worse predictive performance and the stacking objective with no
regularization ($\beta=\infty$ in Fig. 4) is not outperformed by any form of
regularization. However, depending on the application setting some level of
regularization might still be desirable.

##  7 Conclusions

Inference in probabilistic programs with stochastic support implicitly
performs BMA over the constitutent SLPs, revealing issues related to
instability of the posterior SLP weights under model misspecification. To
address these issues, we proposed to fit the SLP weights by optimizing
predictive objectives. In its most basic form, this leads to a method called
stacking of predictive distributions. By adding a regularization term to the
objective we can recover a PAC-Bayes bound. Our experiments demonstrated that
our alternative weighting schemes lead to more robust SLP weights with better
predictive performance.

#### Acknowledgements

We would like to thank Mrinank Sharma, and Andrew Campbell for feedback on
earlier drafts of this manuscript. We would also like to thank the anonymous
reviewers for providing their reviews, especially for encouraging us to
explore the connections to PAC-Bayes. Tim Reichelt is supported by the UK
EPSRC CDT in Autonomous Intelligent Machines and Systems with the grant
EP/S024050/1 and Luke Ong acknowledges support from the National Research
Foundation, Singapore, under its RSS Scheme (NRF-RSS2022-009). The authors
would like to acknowledge the use of the University of Oxford Advanced
Research Computing (ARC) facility in carrying out this work
(http://dx.doi.org/10.5281/zenodo.22558).

## References

  * Tolpin et al. (2016) David Tolpin, Jan-Willem van de Meent, Hongseok Yang, and Frank Wood.  Design and Implementation of Probabilistic Programming Language Anglican.  In _Proceedings of the 28th Symposium on the Implementation and Application of Functional Programming Languages - IFL 2016_ , pages 1–12, Leuven, Belgium, 2016. ACM Press.  ISBN 978-1-4503-4767-9.  doi: 10.1145/3064899.3064910.  URL http://dl.acm.org/citation.cfm?doid=3064899.3064910.
  * Goodman et al. (2008) Noah D. Goodman, Vikash K. Mansinghka, Daniel Roy, Keith Bonawitz, and Joshua B. Tenenbaum.  Church: A language for generative models.  UAI’08, July 2008.
  * Bingham et al. (2019) Eli Bingham, Jonathan P. Chen, Martin Jankowiak, Fritz Obermeyer, Neeraj Pradhan, Theofanis Karaletsos, Rohit Singh, Paul Szerlip, Paul Horsfall, and Noah D. Goodman.  Pyro: Deep Universal Probabilistic Programming.  _Journal of Machine Learning Research_ , 2019.
  * Ge et al. (2018) Hong Ge, Kai Xu, and Zoubin Ghahramani.  Turing: A Language for Flexible Probabilistic Inference.  In _International Conference on Artificial Intelligence and Statistics_. PMLR, March 2018.
  * Mansinghka et al. (2014) Vikash Mansinghka, Daniel Selsam, and Yura Perov.  Venture: A higher-order probabilistic programming platform with programmable inference.  _arXiv:1404.0099 [cs, stat]_ , March 2014.
  * Chaganty et al. (2013) Arun Chaganty, Aditya Nori, and Sriram Rajamani.  Efficiently Sampling Probabilistic Programs via Program Analysis.  In _Proceedings of the Sixteenth International Conference on Artificial Intelligence and Statistics_ , pages 153–160. PMLR, April 2013.  URL https://proceedings.mlr.press/v31/chaganty13a.html.
  * Sankaranarayanan et al. (2013) Sriram Sankaranarayanan, Aleksandar Chakarov, and Sumit Gulwani.  Static analysis for probabilistic programs: Inferring whole program properties from finitely many paths.  In _Proceedings of the 34th ACM SIGPLAN Conference on Programming Language Design and Implementation_ , PLDI ’13, pages 447–458, New York, NY, USA, June 2013. Association for Computing Machinery.  ISBN 978-1-4503-2014-6.  doi: 10.1145/2491956.2462179.  URL https://doi.org/10.1145/2491956.2462179.
  * Luo et al. (2021) Yicheng Luo, Antonio Filieri, and Yuan Zhou.  Symbolic parallel adaptive importance sampling for probabilistic program analysis.  In _Proceedings of the 29th ACM Joint Meeting on European Software Engineering Conference and Symposium on the Foundations of Software Engineering_ , pages 1166–1177, 2021.
  * Zhou et al. (2020) Yuan Zhou, Hongseok Yang, Yee Whye Teh, and Tom Rainforth.  Divide, Conquer, and Combine: A New Inference Strategy for Probabilistic Programs with Stochastic Support.  In _Proceedings of the 37th International Conference on Machine Learning_ , pages 11534–11545. PMLR, November 2020.  URL https://proceedings.mlr.press/v119/zhou20e.html.
  * Reichelt et al. (2022a) Tim Reichelt, Luke Ong, and Tom Rainforth.  Rethinking variational inference for probabilistic programs with stochastic support.  In _Advances in Neural Information Processing Systems_ , 2022a.
  * Hoeting et al. (1999) Jennifer A Hoeting, David Madigan, Adrian E Raftery, and Chris T Volinsky.  Bayesian model averaging: a tutorial (with comments by m. clyde, david draper and ei george, and a rejoinder by the authors.  _Statistical science_ , 14(4):382–417, 1999.
  * Minka (2000) Thomas P Minka.  Bayesian model averaging is not model combination.  2000\.
  * Yao et al. (2018) Yuling Yao, Aki Vehtari, Daniel Simpson, and Andrew Gelman.  Using Stacking to Average Bayesian Predictive Distributions (with Discussion).  _Bayesian Analysis_ , 13(3):917–1007, September 2018.  ISSN 1936-0975, 1931-6690.  doi: 10.1214/17-BA1091.
  * Gelman and Yao (2020) Andrew Gelman and Yuling Yao.  Holes in Bayesian statistics.  _Journal of Physics G: Nuclear and Particle Physics_ , 48(1), December 2020.
  * Oelrich et al. (2020) Oscar Oelrich, Shutong Ding, Måns Magnusson, Aki Vehtari, and Mattias Villani.  When are bayesian model probabilities overconfident?  _arXiv preprint arXiv:2003.04026_ , 2020.
  * Huggins and Miller (2021) Jonathan H. Huggins and Jeffrey W. Miller.  Reproducible Model Selection Using Bagged Posteriors, December 2021\.  URL http://arxiv.org/abs/2007.14845.
  * Yang and Zhu (2018) Ziheng Yang and Tianqi Zhu.  Bayesian selection of misspecified models is overconfident and may cause spurious posterior probabilities for phylogenetic trees.  _Proceedings of the National Academy of Sciences_ , 115(8):1854–1859, February 2018.  doi: 10.1073/pnas.1712673115.  URL https://www.pnas.org/doi/full/10.1073/pnas.1712673115.
  * Box (1976) George EP Box.  Science and statistics.  _Journal of the American Statistical Association_ , 71(356):791–799, 1976.
  * Key et al. (1999) Jane T Key, Luis R Pericchi, and Adrian FM Smith.  Bayesian model choice: what and why.  _Bayesian statistics_ , 6:343–370, 1999.
  * Vehtari and Ojanen (2012) Aki Vehtari and Janne Ojanen.  A survey of Bayesian predictive methods for model assessment, selection and comparison.  _Statistics Surveys_ , 6(none):142–228, January 2012.  ISSN 1935-7516.  doi: 10.1214/12-SS102.
  * Smets and Wouters (2007) Frank Smets and Rafael Wouters.  Shocks and frictions in us business cycles: A bayesian dsge approach.  _American economic review_ , 97(3):586–606, 2007\.
  * Leff et al. (2008) Alexander P Leff, Thomas M Schofield, Klass E Stephan, Jennifer T Crinion, Karl J Friston, and Cathy J Price.  The cortical dynamics of intelligible speech.  _Journal of Neuroscience_ , 28(49):13209–13215, 2008.
  * Wolpert (1992) David H. Wolpert.  Stacked generalization.  _Neural Networks_ , 5(2):241–259, 1992.  ISSN 0893-6080.  doi: https://doi.org/10.1016/S0893-6080(05)80023-1.  URL https://www.sciencedirect.com/science/article/pii/S0893608005800231.
  * Breiman (1996) Leo Breiman.  Stacked regressions.  _Machine learning_ , 24(1):49–64, 1996.
  * LeBlanc and Tibshirani (1996) Michael LeBlanc and Robert Tibshirani.  Combining estimates in regression and classification.  _Journal of the American Statistical Association_ , 1996.
  * Masegosa (2020) Andres Masegosa.  Learning under model misspecification: Applications to variational and ensemble methods.  _Advances in Neural Information Processing Systems_ , 33:5479–5491, 2020.
  * Masiha et al. (2021) Mohammad Saeed Masiha, Amin Gohari, Mohammad Hossein Yassaee, and Mohammad Reza Aref.  Learning under distribution mismatch and model misspecification.  In _2021 IEEE International Symposium on Information Theory (ISIT)_ , pages 2912–2917. IEEE, 2021.
  * Morningstar et al. (2022) Warren R. Morningstar, Alex Alemi, and Joshua V. Dillon.  Pacm-bayes: Narrowing the empirical risk gap in the misspecified bayesian regime.  In Gustau Camps-Valls, Francisco J. R. Ruiz, and Isabel Valera, editors, _Proceedings of The 25th International Conference on Artificial Intelligence and Statistics_ , volume 151 of _Proceedings of Machine Learning Research_ , pages 8270–8298. PMLR, 28–30 Mar 2022.  URL https://proceedings.mlr.press/v151/morningstar22a.html.
  * Alquier (2023) Pierre Alquier.  User-friendly introduction to pac-bayes bounds, 2023.
  * Bernardo and Smith (2009) José M Bernardo and Adrian FM Smith.  _Bayesian theory_ , volume 405.  John Wiley & Sons, 2009.
  * Clyde and Iversen (2013) Merlise Clyde and Edwin S Iversen.  Bayesian model averaging in the M-open framework.  In _Bayesian Theory and Applications_. Oxford University Press, 01 2013.
  * Borgström et al. (2016) Johannes Borgström, Ugo Dal Lago, Andrew D. Gordon, and Marcin Szymczak.  A lambda-calculus foundation for universal probabilistic programming.  In _Proceedings of the 21st ACM SIGPLAN International Conference on Functional Programming_ , ICFP 2016. Association for Computing Machinery, September 2016.
  * Staton et al. (2016) Sam Staton, Hongseok Yang, Frank Wood, Chris Heunen, and Ohad Kammar.  Semantics for probabilistic programming: Higher-order functions, continuous distributions, and soft constraints.  In _Proceedings of the 31st Annual ACM/IEEE Symposium on Logic in Computer Science_ , LICS ’16, New York, NY, USA, July 2016\. Association for Computing Machinery.
  * Gordon et al. (2014) Andrew D. Gordon, Thomas A. Henzinger, Aditya V. Nori, and Sriram K. Rajamani.  Probabilistic programming.  In _Future of Software Engineering Proceedings_ , FOSE 2014, pages 167–181, New York, NY, USA, May 2014. Association for Computing Machinery.
  * Rainforth (2017) Thomas William Gamlen Rainforth.  _Automating Inference, Learning, and Design Using Probabilistic Programming_.  http://purl.org/dc/dcmitype/Text, University of Oxford, 2017.  URL https://ora.ox.ac.uk/objects/uuid:e276f3b4-ff1d-44bf-9d67-013f68ce81f0.
  * van de Meent et al. (2018) Jan-Willem van de Meent, Brooks Paige, Hongseok Yang, and Frank Wood.  An Introduction to Probabilistic Programming.  _arXiv:1809.10756 [cs, stat]_ , September 2018.  URL http://arxiv.org/abs/1809.10756.
  * Gneiting and Raftery (2007) Tilmann Gneiting and Adrian E Raftery.  Strictly proper scoring rules, prediction, and estimation.  _Journal of the American statistical Association_ , 102(477):359–378, 2007.
  * Byrd et al. (1995) Richard H Byrd, Peihuang Lu, Jorge Nocedal, and Ciyou Zhu.  A limited memory algorithm for bound constrained optimization.  _SIAM Journal on scientific computing_ , 16(5):1190–1208, 1995.
  * Zhu et al. (1997) Ciyou Zhu, Richard H Byrd, Peihuang Lu, and Jorge Nocedal.  Algorithm 778: L-bfgs-b: Fortran subroutines for large-scale bound-constrained optimization.  _ACM Transactions on mathematical software (TOMS)_ , 23(4):550–560, 1997.
  * Vehtari et al. (2017) Aki Vehtari, Andrew Gelman, and Jonah Gabry.  Practical bayesian model evaluation using leave-one-out cross-validation and waic.  _Statistics and computing_ , 27(5):1413–1432, 2017.
  * Kumar et al. (2019) Ravin Kumar, Colin Carroll, Ari Hartikainen, and Osvaldo Martin.  Arviz a unified library for exploratory analysis of bayesian models in python.  _Journal of Open Source Software_ , 2019.  doi: 10.21105/joss.01143.
  * Geisser and Eddy (1979) Seymour Geisser and William F Eddy.  A predictive approach to model selection.  _Journal of the American Statistical Association_ , 74(365):153–160, 1979.
  * Monteith et al. (2011) Kristine Monteith, James L. Carroll, Kevin Seppi, and Tony Martinez.  Turning bayesian model averaging into bayesian model combination.  In _The 2011 International Joint Conference on Neural Networks_ , 2011\.
  * Kim and Ghahramani (2012) Hyun-Chul Kim and Zoubin Ghahramani.  Bayesian classifier combination.  In Neil D. Lawrence and Mark Girolami, editors, _Proceedings of the Fifteenth International Conference on Artificial Intelligence and Statistics_ , volume 22 of _Proceedings of Machine Learning Research_ , pages 619–627, La Palma, Canary Islands, 21–23 Apr 2012. PMLR.  URL https://proceedings.mlr.press/v22/kim12.html.
  * Yao et al. (2022) Yuling Yao, Gregor Pirš, Aki Vehtari, and Andrew Gelman.  Bayesian hierarchical stacking: Some models are (somewhere) useful.  _Bayesian Analysis_ , 17(4):1043–1071, 2022.
  * Gormley and Frühwirth-Schnatter (2019) Isobel Claire Gormley and Sylvia Frühwirth-Schnatter.  Mixture of experts models.  In _Handbook of mixture analysis_ , pages 271–307. Chapman and Hall/CRC, 2019.
  * Warrell and Gerstein (2022) Jonathan Warrell and Mark Gerstein.  Higher-order generalization bounds: Learning deep probabilistic programs via pac-bayes objectives.  _arXiv preprint arXiv:2203.15972_ , 2022.
  * Tran et al. (2017) Dustin Tran, Matthew D Hoffman, Rif A Saurous, Eugene Brevdo, Kevin Murphy, and David M Blei.  Deep probabilistic programming.  _arXiv preprint arXiv:1701.03757_ , 2017.
  * Freer et al. (2010) Cameron E Freer, Vikash K Mansinghka, and Daniel M Roy.  When are probabilistic programs probably computationally tractable.  In _NIPS Workshop on Monte Carlo Methods for Modern Applications_ , 2010.
  * Neal (2011) Radford M. Neal.  MCMC using Hamiltonian dynamics.  In _Handbook of Markov Chain Monte Carlo_ , pages 113–162. Chapman & Hall / CRC Press, 2011.
  * Hoffman et al. (2014) Matthew D Hoffman, Andrew Gelman, et al.  The no-u-turn sampler: adaptively setting path lengths in hamiltonian monte carlo.  _J. Mach. Learn. Res._ , 15(1):1593–1623, 2014\.
  * Betancourt (2018) Michael Betancourt.  A Conceptual Introduction to Hamiltonian Monte Carlo.  _arXiv:1701.02434 [stat]_ , July 2018.
  * Cusumano-Towner et al. (2019) Marco F. Cusumano-Towner, Feras A. Saad, Alexander K. Lew, and Vikash K. Mansinghka.  Gen: A general-purpose probabilistic programming system with programmable inference.  PLDI 2019, June 2019.
  * Martino et al. (2017) L. Martino, V. Elvira, D. Luengo, and J. Corander.  Layered adaptive importance sampling.  _Statistics and Computing_ , 27(3), May 2017.
  * Pace and Barry (1997) R Kelley Pace and Ronald Barry.  Sparse spatial autoregressions.  _Statistics & Probability Letters_, 33(3):291–297, 1997.
  * Smith et al. (1988) Jack W Smith, James E Everhart, WC Dickson, William C Knowler, and Robert Scott Johannes.  Using the adap learning algorithm to forecast the onset of diabetes mellitus.  In _Proceedings of the annual symposium on computer application in medical care_ , page 261. American Medical Informatics Association, 1988.
  * Kaggle (2020) Kaggle.  Stroke prediction dataset.  https://www.kaggle.com/datasets/fedesoriano/stroke-prediction-dataset, 2020\.  Accessed: 2023-10-05.
  * Gelman and Hill (2006) Andrew Gelman and Jennifer Hill.  _Data analysis using regression and multilevel/hierarchical models_.  Cambridge university press, 2006.
  * Lew et al. (2019) Alexander K. Lew, Marco F. Cusumano-Towner, Benjamin Sherman, Michael Carbin, and Vikash K. Mansinghka.  Trace types and denotational semantics for sound programmable inference in probabilistic languages.  _Proceedings of the ACM on Programming Languages_ , 4(POPL):19:1–19:32, December 2019.  doi: 10.1145/3371087.  URL https://doi.org/10.1145/3371087.
  * Wood et al. (2014) Frank Wood, Jan Willem Meent, and Vikash Mansinghka.  A New Approach to Probabilistic Programming Inference.  In _Artificial Intelligence and Statistics_ , pages 1024–1032. PMLR, April 2014.  URL http://proceedings.mlr.press/v33/wood14.html.
  * Rainforth et al. (2016) Tom Rainforth, Christian Naesseth, Fredrik Lindsten, Brooks Paige, Jan-Willem Vandemeent, Arnaud Doucet, and Frank Wood.  Interacting particle markov chain monte carlo.  In Maria Florina Balcan and Kilian Q. Weinberger, editors, _Proceedings of The 33rd International Conference on Machine Learning_ , volume 48 of _Proceedings of Machine Learning Research_ , pages 2616–2625, New York, New York, USA, 20–22 Jun 2016. PMLR.  URL https://proceedings.mlr.press/v48/rainforth16.html.
  * Murray and Schön (2020) Lawrence M. Murray and Thomas B. Schön.  Automated learning with a probabilistic programming language: Birch.  _arXiv:1810.01539 [cs, stat]_ , April 2020.  URL http://arxiv.org/abs/1810.01539.
  * Le (2016) Tuan Anh Le.  _Inference for Higher Order Probabilistic Programs_.  Master’s Thesis, University of Oxford, 2016.
  * Baydin et al. (2019) Atilim Güneş Baydin, Lei Shao, Wahid Bhimji, Lukas Heinrich, Lawrence Meadows, Jialin Liu, Andreas Munk, Saeid Naderiparizi, Bradley Gram-Hansen, Gilles Louppe, Mingfei Ma, Xiaohui Zhao, Philip Torr, Victor Lee, Kyle Cranmer, Prabhat, and Frank Wood.  Etalumis: Bringing probabilistic programming to scientific simulators at scale.  In _Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis_. Association for Computing Machinery, November 2019.
  * Harvey et al. (2019) William Harvey, Andreas Munk, Atılım Güneş Baydin, Alexander Bergholm, and Frank Wood.  Attention for inference compilation.  _arXiv preprint arXiv:1910.11961_ , 2019.
  * Wingate et al. (2011) David Wingate, Andreas Stuhlmueller, and Noah Goodman.  Lightweight Implementations of Probabilistic Programming Languages Via Transformational Compilation.  In _Proceedings of the Fourteenth International Conference on Artificial Intelligence and Statistics_ , pages 770–778. JMLR Workshop and Conference Proceedings, June 2011.  URL https://proceedings.mlr.press/v15/wingate11a.html.
  * Yang et al. (2014) Lingfeng Yang, Patrick Hanrahan, and Noah Goodman.  Generating efficient mcmc kernels from probabilistic programs.  In _Artificial Intelligence and Statistics_ , pages 1068–1076. PMLR, 2014.
  * Tolpin et al. (2015) David Tolpin, Jan-Willem van de Meent, Brooks Paige, and Frank Wood.  Output-sensitive adaptive metropolis-hastings for probabilistic programs.  In _Machine Learning and Knowledge Discovery in Databases: European Conference, ECML PKDD 2015, Porto, Portugal, September 7-11, 2015, Proceedings, Part II 15_ , pages 311–326. Springer, 2015.
  * Roberts et al. (2019) David A. Roberts, Marcus Gallagher, and Thomas Taimre.  Reversible Jump Probabilistic Programming.  In _Proceedings of the Twenty-Second International Conference on Artificial Intelligence and Statistics_ , pages 634–643. PMLR, April 2019.  URL https://proceedings.mlr.press/v89/roberts19a.html.
  * Cusumano-Towner et al. (2020) Marco Cusumano-Towner, Alexander K. Lew, and Vikash K. Mansinghka.  Automating Involutive MCMC using Probabilistic and Differentiable Programming.  _arXiv:2007.09871 [stat]_ , July 2020.  URL http://arxiv.org/abs/2007.09871.
  * Mak et al. (2021) Carol Mak, Fabian Zaiser, and Luke Ong.  Nonparametric Hamiltonian Monte Carlo.  In _Proceedings of the 38th International Conference on Machine Learning_ , pages 7336–7347. PMLR, July 2021.  URL https://proceedings.mlr.press/v139/mak21a.html.
  * Mak et al. (2022) Carol Mak, Fabian Zaiser, and Luke Ong.  Nonparametric involutive markov chain monte carlo.  In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvári, Gang Niu, and Sivan Sabato, editors, _International Conference on Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA_ , volume 162 of _Proceedings of Machine Learning Research_ , pages 14802–14859. PMLR, 2022.  URL https://proceedings.mlr.press/v162/mak22a.html.
  * Wingate and Weber (2013) David Wingate and Theophane Weber.  Automated Variational Inference in Probabilistic Programming.  _arXiv:1301.1299 [cs, stat]_ , January 2013.  URL http://arxiv.org/abs/1301.1299.
  * Ranganath et al. (2014) Rajesh Ranganath, Sean Gerrish, and David Blei.  Black Box Variational Inference.  In _Proceedings of the Seventeenth International Conference on Artificial Intelligence and Statistics_ , pages 814–822. PMLR, April 2014.  URL https://proceedings.mlr.press/v33/ranganath14.html.
  * Paige (2016) Timothy Brooks Paige.  _Automatic Inference for Higher-Order Probabilistic Programs_.  http://purl.org/dc/dcmitype/Text, University of Oxford, 2016.  URL https://ora.ox.ac.uk/objects/uuid:d912c4de-4b08-4729-aa19-766413735e2a.
  * Bernardo (1979) José M Bernardo.  Expected information as expected utility.  _the Annals of Statistics_ , pages 686–690, 1979.
  * Vehtari et al. (2015) Aki Vehtari, Daniel Simpson, Andrew Gelman, Yuling Yao, and Jonah Gabry.  Pareto smoothed importance sampling.  _arXiv preprint arXiv:1507.02646_ , 2015.
  * Zinkov and Shan (2017) Robert Zinkov and Chung-chieh Shan.  Composing inference algorithms as program transformations.  _Proceedings of Uncertainty in Artificial Intelligence (UAI)_ , page 10, 2017.
  * Reichelt et al. (2022b) Tim Reichelt, Adam Goliński, Luke Ong, and Tom Rainforth.  Expectation programming: Adapting probabilistic programming systems to estimate expectations efficiently.  In _Uncertainty in Artificial Intelligence_ , pages 1676–1685. PMLR, 2022b.
  * Lew et al. (2023) Alexander K Lew, Mathieu Huot, Sam Staton, and Vikash K Mansinghka.  Adev: Sound automatic differentiation of expected values of probabilistic programs.  _Proceedings of the ACM on Programming Languages_ , 7(POPL):121–153, 2023.
  * Virtanen et al. (2020) Pauli Virtanen, Ralf Gommers, Travis E. Oliphant, Matt Haberland, Tyler Reddy, David Cournapeau, Evgeni Burovski, Pearu Peterson, Warren Weckesser, Jonathan Bright, Stéfan J. van der Walt, Matthew Brett, Joshua Wilson, K. Jarrod Millman, Nikolay Mayorov, Andrew R. J. Nelson, Eric Jones, Robert Kern, Eric Larson, C J Carey, İlhan Polat, Yu Feng, Eric W. Moore, Jake VanderPlas, Denis Laxalde, Josef Perktold, Robert Cimrman, Ian Henriksen, E. A. Quintero, Charles R. Harris, Anne M. Archibald, Antônio H. Ribeiro, Fabian Pedregosa, Paul van Mulbregt, and SciPy 1.0 Contributors.  SciPy 1.0: Fundamental Algorithms for Scientific Computing in Python.  _Nature Methods_ , 17:261–272, 2020.  doi: 10.1038/s41592-019-0686-2.
  * Liu and Nocedal (1989) Dong C Liu and Jorge Nocedal.  On the limited memory bfgs method for large scale optimization.  _Mathematical programming_ , 45(1-3):503–528, 1989.
  * Bishop and Nasrabadi (2006) Christopher M Bishop and Nasser M Nasrabadi.  _Pattern recognition and machine learning_ , volume 4.  Springer, 2006.
  * Burda et al. (2015) Yuri Burda, Roger Grosse, and Ruslan Salakhutdinov.  Importance weighted autoencoders.  _arXiv preprint arXiv:1509.00519_ , 2015.
  * Morningstar et al. (2021) Warren Morningstar, Sharad Vikram, Cusuh Ham, Andrew Gallagher, and Joshua Dillon.  Automatic differentiation variational inference with mixtures.  In _International Conference on Artificial Intelligence and Statistics_ , pages 3250–3258. PMLR, 2021.

Beyond Bayesian Model Averaging over Paths in Probabilistic Programs with
Stochastic Support

##  Appendix A A Short Introduction to Probabilistic Programming

###  A.1 Programs as Unnormalized Densities

⬇

def model(data):

x = sample("x", Normal(0, 1))

m1 = sample("m1", Bernoulli(0.5))

if m1:

std = 0.62177

with plate("data1", data.shape[0]):

sample("y1", Normal(x, std1), obs=data)

else:

std = 2.0

with plate("data2", data.shape[0]):

sample("y2", Normal(x, std2), obs=data)

Figure 5: Example Pyro program with stochastic support.

We here give an informal description for how the probabilistic programming
language Pyro (Bingham et al., 2019) defines unnormalized densities. For more
formal description of the semantic foundations of probabilistic programming we
refer the interested reader to Borgström et al. (2016); Staton et al. (2016)
and Lew et al. (2019). Pyro provides the sample primitive function to both
define latent random variables and condition on observed data. More precisely,
calling the function sample(addr, dist) draws a sample from the distribution
object dist with the corresponding lexical address addr. Conditioning on
observed data is made possible by passing in the observed data to the sample
function as a keyword argument like so sample(addr, dist, obs=data, where
again addr and dist are assumed to be a lexical address and distribution
object. The users is responsible for ensuring that each sample statement gets
assigned a unique address, i.e. an address that has not been encountered in
the current execution.

Pyro is a universal PPS as it is embedded within the Python programming
language and users are free to use language features such as branching on the
outcomes of sample statements. This means Pyro is able to express programs
with stochastic support i.e. the number of sample statements and their
corresponding distribution type can vary from one execution to the next. Fig.
5 gives an example of a Pyro program with stochastic support.

A Pyro program defines an unnormalized density over the _raw random draws_ ,
$\theta_{1:n_{\theta}}\in\Theta$, which are defined as the sequence of
outcomes of the sample statements (without any observed data) encountered in
the program. Crucially, in models with stochastic support the number of sample
statements $n_{x}$ can be a random variable and is not fixed for a given
program. Furthermore, the raw random draws are assumed to be the only source
of randomness in the program, such that for given raw random draws
$\theta^{\prime}_{1:n_{\theta}}$ all the intermediate variables and return
values can be determined deterministically.

Each sample statement encountered during the program’s execution contributes
one term to the program’s unnormalized density
$\gamma(\theta_{1:n_{\theta}})$. If a sample does not have associated observed
data, then it contributes the term $f_{a_{i}}(\theta_{i}\mid\eta_{i})$ where
$a_{i}$ is the lexical address, $\theta_{i}$ the outcome of the sample
statement, $f_{a_{i}}$ is the parameterized density function of the associated
distribution object, and $\eta_{i}$ are the corresponding parameters.
Similarly, a sample with observed data $y_{j}$ contributes the term
$g_{b_{j}}(y_{j}\mid\phi_{j})$ with $b_{j}$ denoting the lexical address,
$g_{b_{j}}$ the parameterized density function, and $\phi_{j}$ the parameters.
Overall, as noted by (Rainforth, 2017, §4.3.2) who formalized this for the PPS
Anglican (Wood et al., 2014), this implies the program’s unnormalized density
function is given by

$\gamma(\theta_{1:n_{\theta}}):=\prod\nolimits_{i=1}^{n_{\theta}}f_{a_{i}}(\theta_{i}\mid\eta_{i})\prod\nolimits_{j=1}^{n_{y}}g_{b_{j}}(y_{j}\mid\phi_{j}).$
(17)

_Inference_ in a probabilistic program corresponds to the task of finding a
representation of the normalized density
$\pi(\theta_{1:n})=\gamma(\theta_{1:n})/Z$ with normalization constant
$Z=\int_{\Theta}\gamma(\theta_{1:n_{\theta}})d\theta_{1:n_{\theta}}$.

###  A.2 Approaches to Inference

For most problems encountered in practice we cannot exactly compute the
normalized program density $\pi(\theta_{1:n_{\theta}})$ and instead have to
resort to approximate inference algorithms. A variety of approaches for
approximate inference in probabilistic programs with stochastic support exist.
Particle-based approaches (Wood et al., 2014; Rainforth et al., 2016; Murray
and Schön, 2020) and algorithms based on importance sampling (Le, 2016; Baydin
et al., 2019; Harvey et al., 2019) generate a set of weighted samples as an
approximate posterior. Markov chain Monte Carlo (MCMC) methods with either
automatic or manual proposals try to generate a set of samples directly from
the posterior (Wingate et al., 2011; Yang et al., 2014; Tolpin et al., 2015;
Roberts et al., 2019; Cusumano-Towner et al., 2020; Mak et al., 2021, 2022).
Variational inference algorithms create a parameterized distribution
$q(\theta;\phi)$ and optimize the parameters $\phi$ such that $q(\theta;\phi)$
is “close” to the full Bayes posterior where closeness is measured with some
divergence (most commonly the KL divergence) (Wingate and Weber, 2013;
Ranganath et al., 2014; Paige, 2016).

We now give a more detailed description of involutive MCMC which is the
foundation behind Gen’s implementation of reversible-jump MCMC (RJMCMC)
(Cusumano-Towner et al., 2020). In involutive MCMC the user specifies an
auxiliary kernel $\kappa$ which acts on an auxiliary variable $v\in Y$, s.t.
$\kappa_{\theta}(\cdot):Y\to[0,\infty)$ for all $\theta\in\Theta$ with
$\gamma(\theta)>0$, and an involution $\Phi:\Theta\times Y\to\Theta\times Y$,
i.e. $\Phi^{-1}=\Phi$. Given an initial state $\theta$, involutive MCMC
proceeds by first sampling a new auxiliary variable
$v\sim\kappa_{\theta}(\cdot)$, then applying the involution to get the newly
proposed state $(\theta,v)\leftarrow\Phi(\theta,v)$, and accepting the new
state with the acceptance probability

$\alpha:=\min\Bigl{\\{}1,\frac{\gamma(\theta^{\prime})\,\kappa_{\theta}(v)}{\gamma(\theta)\,\kappa_{\theta}(v)}\lvert\det(\nabla\Phi(\theta,v))\rvert\Bigr{\\}}.$
(18)

The usage of involutive MCMC is partly automated in Gen. The user only has to
implement the auxiliary kernel and the involution using a domain-specific
language. Based on these quantities Gen is then able to automatically
construct an involutive MCMC kernel (Cusumano-Towner et al., 2020).

##  Appendix B Probabilistic Programs without Predictive Distribution

⬇

def model(y):

# Input data is a list of length 2

if y[0] > 10:

x = sample("x1", Normal(10, 1))

sample("y1", Normal(x, 10), obs=y[0])

sample("y2", Normal(y[0], sqrt(y[0])), obs=y[1])

else:

x = sample("x2", Normal(0, 1))

sample("y1", Normal(x, 1), obs=y[0])

sample("y2", Normal(y[0], 1), obs=y[1])

Figure 6: Example of a probabilistic program without a predictive
distribution.

Conventionally, in Bayesian statistics the modeller defines a prior
$p(\theta)$ and predictive distribution $p(y_{i}\mid\theta)$. Then for a
dataset $\mathbf{y}=(y_{1},\dots,y_{N})$, these two ingredients together
define the joint density
$p(\theta,\mathbf{y})=\prod_{i=1}^{N}p(y_{i}\mid\theta)\,p(\theta)$ from which
can compute the posterior $p(\theta\mid\mathbf{y})\propto
p(\theta,\mathbf{y})$. In order to predict new data $\tilde{y}$, we can then
use the posterior predictive distribution defined as

$p(\tilde{y}\mid\mathbf{y})=\int
p(\tilde{y}\mid\theta)\,p(\theta\mid\mathbf{y})dx.$ (19)

However, more generally, the joint density does not necessarily factorize in
this manner which makes it more difficult to automatically deduce a prediction
task from the model alone. Additionally, in the setting of universal
probabilistic programming languages the input data can directly influence the
model definition as well. Take for example the Pyro program in Fig. 6. Here,
the input data $y$ directly influences what sample statements we encounter
during execution. Furthermore, in both sample statenents with address "y2" the
distribution depends on the first data point y[0]. Therefore, just from this
program definition alone, it is not clear what exactly a reasonable predictive
distribution for a new data point would be.

##  Appendix C Introduction to Scoring Rules

This introduction mainly follows from Gneiting and Raftery (2007) and Yao et
al. (2018). _Scoring rules_ are functions which take as input a _probabilistic
forecast_ and a realized event. The goal of scoring rules is then to evaluate
the quality of the probabilistic forecast. The terminology is specifically
used in meteorological forecasts. In a Bayesian context, these scores are
often instead referred to as utilities and Bayesian decision theory aims to
maximize the predicted utility of a given action (Bernardo and Smith, 2009).

More formally, assume we have a random variable on the sample space
$(\Omega,\mathcal{A})$ and $\mathcal{P}$ is a convex class of probability
measures on $(\Omega,\mathcal{A})$. Then, any member $P\in\mathcal{P}$ is
referred to as a probabilistic forecast and a scoring rule is a function
$S:\mathcal{P}\times\Omega\to[-\infty,\infty]$ s.t. $S(P,\cdot)$ is
$\mathcal{P}$-quasi-integrable for all forecasts $P\in\mathcal{P}$. So for a
probabilistic forecast $P$ and observed event $y\in\Omega$, $S(P,y)$ is the
score which indicates the quality of our forecast. For notational convenience,
if $P$ and $Q$ are both probabilistic forecasts, we define $S(P,Q)=\int
S(P,y)dQ(y)$. Then a scoring rule is _proper_ if $S(Q,Q)\geq S(P,Q)$ holds for
all $P\in\mathcal{P}$, and _strictly proper_ if the equality holds only when
$P=Q$ almost surely.

Some common examples of scoring rules include: the _quadratic score_
$S(\rho,y)=2\rho(y)-\lVert\rho\rVert^{2}_{2}$, where $\rho$ is a predictive
density; the _logarithmic score_ $S(\rho,y)=\log\rho(y)$; and the _continuous-
ranked probability score_
$S(F,y)=-\int(F(y^{\prime})-\mathbb{I}[y^{\prime}\geq y])dy^{\prime}$, where
$F$ is the cumulative distribution function of the forecast. Under regularity
conditions, Bernardo (1979) showed that the logarithmic scoring rule is the
only proper _local_ scoring rule where a local scoring rule is a rule that
depends on the predictive density $\rho$ only through the actual observed
event $y$.

We refer the reader to Gneiting and Raftery (2007) for more extensive details
on scoring rules.

##  Appendix D PSIS-LOO Approximation

We here only give a brief description of the PSIS-LOO approximation as it is a
common procedure. We refer the reader to Vehtari et al. (2017) for full
details on the approximation and to Vehtari et al. (2015) for more details on
PSIS. We want to approximate the local posterior predictive

$\displaystyle\rho_{k}(y_{i}\mid y_{-i})$
$\displaystyle=\int_{\Theta_{k}}g_{k}(y_{i}\mid\theta)\,\pi_{k}(\theta\mid
y_{-i})d\theta$ (20) which can be rewritten in terms of the posterior of the
full dataset as
$\displaystyle=\int_{\Theta_{k}}g_{k}(y_{i}\mid\theta)\,\frac{\pi_{k}(\theta\mid
y_{-i})}{\pi_{k}(\theta\mid y_{1:N})}\,\pi_{k}(\theta\mid y_{1:N})d\theta.$
(21)

Importantly, the ratio in that integral is proportional to a term that only
depends on the individual predictive density

$r_{i}:=\frac{1}{g_{k}(y_{i}\mid\theta)}\propto\frac{\pi_{k}(\theta\mid
y_{-i})}{\pi_{k}(\theta\mid y_{1:N})}.$ (22)

Hence, we can use a self-normalized importance sampler to get an estimate of
Eq. (21) as follows

$\rho_{k}(y_{i}\mid y_{-i})\approx\frac{\sum_{s\in
I_{k}}r^{s}_{i}\,g_{k}(y_{i}\mid\theta_{s})}{\sum_{s\in I_{k}}r^{s}_{i}}$ (23)

where $r^{s}_{i}=1/g_{k}(y_{i}\mid\theta_{s})$ and $\theta_{s}$ are
(approximate) samples from the posterior. However, these ratios $r^{s}_{i}$
can often have high variance since $\pi_{k}(\theta\mid y_{1:N})$ will usually
have lower variance and thinner tails than $\pi_{k}(\theta\mid y_{-i})$, in
turn, leading to unstable estimates. To avoid these instabilities, the PSIS-
LOO approximation replaces the ratios $r^{s}_{i}$ with smoothed importance
weights $\nu^{s}_{i}$. The importance weights are computed by fitting a
generalized Pareto distribution to the raw ratios $r^{s}_{i}$ and replacing
them with the expected order statistics of the fitted Pareto distribution.
This then leads to the estimate

$\hat{\rho}_{k}(y_{i}\mid y_{-i})=\frac{\sum_{s\in
I_{k}}\nu^{s}_{i}\,g_{k}(y_{i}\mid\theta_{s})}{\sum_{s\in I_{k}}\nu^{s}_{i}}.$
(24)

##  Appendix E Implementation details

⬇

def pyro_subset_regression(X, y, X_val, y_val):

k = pyro.sample(

"k", dist.Categorical(torch.ones(X.shape[1]) / X.shape[1]),

infer={"branching": True},

)

X = X[:, k]

beta = pyro.sample(f"beta_{k.item()}", dist.Normal(0, np.sqrt(10)))

sigma = pyro.sample("sigma", dist.Gamma(0.1, 0.1))

mean = X * beta

with pyro.plate("data", X.shape[0]):

pyro.sample("y", dist.Normal(mean, sigma), obs=y)

X_val = X_val[:, k]

mean_val = X_val * beta

return dist.Normal(mean_val, sigma).log_prob(y_val)

Figure 7: Pyro program for the experiments in Sec. 6.2.

To be able to evaluate the stacking objective defined in Eq. (15) we need
access to $\hat{\rho}_{k}(\tilde{y}_{\ell})$, the estimates of the predictive
densities. These, in turn, depend on the evaluations of the predictive
densities. Hence, given validation data $\tilde{y}_{1},\dots,\tilde{y}_{L}$
and posterior samples $\theta_{1},\dots,\theta_{S}$, for stacking we require
the evaluations $g(\tilde{y}_{\ell}\mid\theta_{s})$. Note that the local
posterior predictive distributions $\rho_{k}(\tilde{y}_{\ell})$ are
expectations under the posterior.

This is important because previous work noted that in the context of
probabilistic programming these types of expectations can be formalized as the
_expected return values_ of a program (Gordon et al., 2014; Zinkov and Shan,
2017; Reichelt et al., 2022b; Lew et al., 2023). Then to apply stacking, we
require the user to define a program in which the return values for a given
sample $\theta_{s}$ are the predictive density evaluations
$g(\tilde{y}_{1}\mid\theta_{s}),\dots,g(\tilde{y}_{L}\mid\theta_{s})$. For our
Pyro implementation this means that we require the user to define a program
which returns an $L$-dimensional vector which correspond to the (log)
predictive density on the validation data points. Fig. 7 shows how this can be
done for the subset regression model from Sec. 6.2.

To actually compute the stacking weights we rely on Pyro’s
Trace333https://docs.pyro.ai/en/stable/poutine.html#trace data structure which
saves important metadata of each program execution such as the distribution
type of each sample statement, the corresponding sampled value, its address,
and, crucially, also the return value of the program. Algorithm 1 can then be
implemented as a simple function which takes as input a list of posterior
samples in the form of Pyro Traces. Because each Trace stores the address path
of a given run, we can easily determine the set of SLPs and associate each
sample with its SLP. From the Trace data structure we can then also extract
the return values i.e. the evaluations $\log
g(\tilde{y}_{\ell}\mid\theta_{s})$. From these evaluations we can calculate
the estimates $\hat{\rho}_{k}(\tilde{y}_{\ell})$ which are needed in our
stacking objective in Eq. (15). The optimization of the objective is done
using SciPy’s (Virtanen et al., 2020) implementation of the L-BFGS-B optimizer
(Liu and Nocedal, 1989).

Similarly, to Zhou et al. (2020) and Reichelt et al. (2022a) we also allow the
user to annotate specific discrete sampling statements to indicate that they
influence the SLP of the program. In our implementation, this is done by
passing {"branching": True} to the infer keyword argument of Pyro’s sample
statement. Using this annotation the inference backend can then control which
SLP of the program is sampled by conditioning these sample statements on
specific values. Furthermore, these annotations allow the inference backend to
deterministically enumerate all SLPs. In our implementation, we give the user
to enumerate all SLPs using breadth-first search, i.e. in the enumeration
procedure we run the program as normal but if we encounter an annotated
sampling statement we enumerate the whole support of the distribution and put
each possible sampled value onto a queue. Once enumeration of a sample site
has finished, execution resumes by popping a value from the front of the queue
and continuing executing the program from the sample statement from which the
value was sampled. This enumeration procedure is similar to how
marginalization of discrete sample sites is implemented in Pyro. Note that we
are assuming that each annotated discrete sample site has finite support.
Future implementations could deal with finite support by truncating the
sampling distribution.

###  E.1 Alternative Interface for Stacking

⬇

def model_average(models, model_args, model_kwargs):

k = pyro.sample(

"k", dist.Categorical(torch.ones(len(models)) / len(models)),

infer={"branching": True},

)

return models[k](*model_args, **model_kwargs)

Figure 8: Alternative interface to enable stacking of a list of user-specified
models.

The main aim of our paper is to highlight the shortcomings of the default BMA
weights in posteriors of probabilistic programs with stochastic support.
However, as a side-effect our Pyro implementation also provides a convenient
mechanism for users to utilize stacking. Probabilistic programs with
stochastic support allow users to encode a large class of problem settings,
including the traditional BMA setting. However, in some cases, for
convenience, users might not wish to write a single program with stochastic
support which encodes all the different models in the BMA. Instead, the user
might want to write $K$ separate programs, each with static support, and then
want to apply either BMA or stacking to that list of models. However, this
interface is easily encoded within our framework because the user only needs
to write another program which combines all the models together. This program
is shown in Fig. 8, it takes as input a list of Pyro models, and the arguments
that should be passed to these models. The program then samples an index $k$
and chooses one of the candidate models. Running standard inference in this
program would correspond to BMA. If the user instead wants to run stacking
instead of BMA, they can simply choose our stacking implementation and apply
it to this program.

##  Appendix F Experimental Details and Additional Results

To make the stacking approach widely applicable we implemented a version of
the DCC framework in Pyro (Bingham et al., 2019). For our implementation of
DCC, we assume that all the stochastic branching happens on variables with
discrete support and that the remaining latent variables in the program are
continuous. Models of this form then permit the usage of Pyro’s built-in
Hamiltonian Monte Carlo (Neal, 2011; Hoffman et al., 2014; Betancourt, 2018)
as the local inference algorithm. To improve efficiency, our implementation
also leverages Pyro’s just-in-time (JIT) compilation ability to compile each
SLP. Note that leveraging Pyro’s JIT compilation is possible due to breaking
down the original program into its SLPs because by default the JIT compilation
cannot handle programs with stochastic support.

Our experiments were conducted on an internal compute cluster which consists
of a mix of Intel Broadwell, Haswell, and Cascade Lake CPUs. In general, we
use 16 cores to parallelize our computation and running a single replication
of an experiment finishes in a matter of minutes if not seconds. An exception
is the function induction experiment in Sec. 6.3 which takes around 3 hours
for a single replication, using 32 cores.

###  F.1 When is Stacking Helpful?

For each setting we collect $10^{3}$ HMC samples from each SLP with $400$
burn-in samples. To estimate the normalization constant for the dominant SLP
setting we use a proposal of
$q(\theta_{1},\theta_{2})=\mathcal{N}(\theta_{1};0,1)\,\Gamma(\theta_{2};1,1)$
with $10^{6}$ samples.

Overall, the experiments show that stacking is particularly useful if there
are multiple SLPs which fit the data roughly equally well (with respect to the
local normalization constant). In these cases, the SLP weights are at risk of
being unstable and stacking can provide a more robust weight estimation
procedure. On the other hand, if there is one dominant SLP which clearly
performs better than all the other SLPs then we would expect there to be less
to gain from stacking, as the weights will already be fairly stable.

###  F.2 Subset Regression

Figure 9:  SLP weights for model in Sec. 6.2. Each dot represents the weight
of the corresponding SLP in the model. Results are computed over 50 randomly
generated datasets.

Figure 10: Impact of varying regularization parameter $\beta$ on the computed
SLP weights. Smaller values of $\beta$ lead to more regularization, pushing
the SLP weights more towards the uniform distribution over SLPs.

Following Breiman (1996) and Yao et al. (2018), we generate the regression
coefficients according to
$\beta_{d}=\eta\,(\zeta_{d}(4)+\zeta_{d}(8)+\zeta_{d}(12))$ with
$\zeta_{d}(a):=\mathbb{I}\left[\lvert d-a<h\rvert\right]\,(h-\lvert
d-a\rvert)^{2}$. The parameter $h$ determines the number of “strong”
coefficients. Following Yao et al. (2018) we set $h=5$ which leads to 15 weak
coefficients and set $\eta$ such that the signal-to-noise ratio
$\mathbb{V}\left[\sum_{d=1}^{15}\beta_{d}\,X_{d}\right]/(1+\mathbb{V}\left[\sum_{d=1}^{15}\beta_{d}\,X_{d}\right])=4/5$
where $X_{d}$ is the random variable for the $d$th covariate. Our input
program has 15 SLPs and each SLP has the unnormalized density
$\gamma_{k}(\theta_{1},\theta_{2},\theta_{3})~{}=~{}\mathbb{I}\left[\theta_{3}=k\right]\allowbreak\,\prod_{i=1}^{N}\mathcal{N}(y_{i};\theta_{1}x_{i,k},\theta_{2}^{2})\,\mathcal{N}(\theta_{1};0,10)\,\Gamma(\theta_{2};0.1,0.1)\allowbreak\,\text{DiscreteUniform}(\theta_{3};1,15)$.
For DCC inference, for each SLP we collect $10^{3}$ HMC samples with $400$
burn-in samples. For RJMCMC our transition kernel samples a new $\theta_{3}$,
the variable controlling the covariate that is selected, from a uniform
categorical distribution and new $theta_{1}$, the local regression
coefficient, from a standard normal distribution. The noise variable is
independently updated using a Metropolis-Hastings kernel. The individual
weights for each SLP are shown in Fig. 9.

###  F.3 Function Induction

Figure 11: SLP weights for model in Sec. 6.3. We only plot SLPs which have
achieved a weight $>0.3$ in any run for any method.

We are using the probabilistic context-free grammar
$e\rightarrow\\{x\,\mid\,\sin(a*e)\,\mid\,a*e+b*e\\}$. In our model, we use
prior production probabilities $[0.4,0.4,0.2]$ and for each of the sampled
coefficients (denoted $a$ and $b$ in the PCFG) we use a $\mathcal{N}(0,10)$
prior. Using the PCFG we can sample an expression for a function
$f:\mathbb{R}\rightarrow\mathbb{R}$, to complete our probabilistic model we
then use a likelihood $\mathcal{N}(y_{i};f(x_{i}),\sigma^{2})$ for input data
pairs $x_{i}$ and $y_{i}$. The standard deviation $\sigma$ of the likelihood
is also inferred from the data for which we use a $\Gamma(\sigma;1,1)$ prior.
For DCC inference, we run $4$ chains with $500$ HMC samples and $200$ burn-in
samples for each SLP. The sampling distribution of SLP weights are shown in
Fig. 11. For RJMCMC, our transition kernel represents each expression as a
PCFG parse tree. To propose a new expression, the transition kernel picks a
random node in the parse tree and replaces that node with a sampled expression
from the prior PCFG. The standard deviation of the likelihood, $\sigma$, is
updated independently with a Metropolis-Hastings transition kernel. For
RJMCMC, we collect the same number of total samples as for DCC to ensure a
fair comparison.

Due to recursion in the PCFG rules, the program actually defines an infinite
number of SLPs. Similar to Zhou et al. (2020), to avoid infinite recursion we
restrict the underlying inference engine to only consider a finite number of
SLPs. Our inference engine enumerates the possible PCFG expressions using
breadth-first search and only considers the first 128 expressions.

###  F.4 Variable Selection

(a) California housing dataset.

(b) Stroke dataset.

(c) Diabetes dataset.

Figure 12: SLP weights for the models in Sec. 6.4 (conventions as in Fig. 11).
The Stroke and Diabetes problems do not permit an analytic solution to the BMA
weights.

We assume we are given data $y_{1:N}$ and an associated matrix of covariates
$X\in\mathbb{R}^{N\times D}$. We consider both regression and classification
problems. The problem of variable selection is to find a subset of the
features $\mathcal{D}\subseteq\\{1,\dots,D\\}$ to make predictions given our
data. For regression task, we have $y_{i}\in\mathbb{R}$ and our model for a
specific subset $\mathcal{D}$ of the features is given by

$\displaystyle\sigma^{2}$ $\displaystyle\sim\Gamma^{-1}(2,1),$ (25)
$\displaystyle\beta_{d}$
$\displaystyle\sim\mathcal{N}(0,\sigma^{2})~{}~{}~{}~{}\text{for
}d\in\mathcal{D},$ (26) $\displaystyle y_{i}$
$\displaystyle\sim\mathcal{N}(\sum\nolimits_{d\in\mathcal{D}}\beta_{d}x_{i,d},\sigma^{2}).$
(27)

This form of the regression model allows us to analytically calculate the
marginal likelihood (see e.g. Ch. 3.5 in Bishop and Nasrabadi (2006)). For
classification tasks, we instead have $y_{i}\in\\{0,1\\}$ and use the logistic
regression model

$\beta_{d}\sim\mathcal{N}(0,1),\text{ for }d\in\mathcal{D}\text{ and
}y_{i}\sim\mathcal{B}(S(\sum\nolimits_{d\in\mathcal{D}}\beta_{d}x_{i,d})),$

where $S(x)=1/(1+\exp(-x))$ and $\mathcal{B}$ is the Bernoulli distribution.

The three different datasets we consider are: _California_ housing
($N=20,650$, $D=8$; 50 % train, 50 % test) (Pace and Barry, 1997), a
regression dataset where the goal is to predict the median house prices for
districts of California; _Diabetes_ ($N=768$, $D=8$; 80 % train, 20 % test)
(Smith et al., 1988), a classification dataset on if a person has diabetes or
not; and _Stroke_ ($N=4908$, $D=8$; 80 % train, 20 % test), a classification
dataset on if a person will have a stroke or not. For DCC inference, we
collect $10^{3}$ HMC samples with $400$ burn-in samples for each SLP. The
sampling distribution of the SLP weights are plotted in Fig. 12. For RJMCMC,
the transition kernel randomly selects a feature dimension $d$ and for that
dimension flips the inclusion, i.e. if the feature was previously included it
will be excluded and vice versa. For a previously excluded feature, a new
coefficient is sampled from the prior. For all other coefficients, a new
coefficient is proposed from a standard normal centered at the current value.
To ensure a fair comparison to DCC, we collect the same number of samples from
RJMCMC as we do for DCC.

###  F.5 Modelling Radon Contamination in US Counties

As mentioned in the main text, our program encodes different modelling choices
for both the intercept(s), $\alpha$, and the slope parameter(s), $\beta$, of
the regression relation $y_{i}=\alpha+\beta\,x_{i}$. Below we describe in
detail the four different modelling choices for the intercept term. The
modelling choices for the slope parameter are analogous, except we do not
consider using a group-level predictor for $\beta$.

Pooling. This model corresponds to the SLP denoted “P, P” in Fig. 3.

$\displaystyle\alpha\sim\mathcal{N}(0,10),~{}~{}\beta\sim\mathcal{N}(0,10),~{}~{}f_{i}=\alpha+\beta\,x_{i},~{}~{}\sigma\sim\text{Exponential}(5),~{}~{}y_{i}\sim\mathcal{N}(f_{i},\sigma^{2}).$
(28)

No pooling. Here $c[i]$ refers to the county index of the $i$th house. This
model corresponds to the SLP denoted “NP, P” in Fig. 3.

$\displaystyle\alpha_{c}\sim\mathcal{N}(0,10)~{}\text{for each county
}c,~{}~{}\beta\sim\mathcal{N}(0,10),~{}~{}\sigma\sim\text{Exponential}(5),~{}~{}y_{i}\sim\mathcal{N}(f_{i},\sigma^{2})$
(29)

with $f_{i}=\alpha_{c[i]}+\beta\,x_{i}$.

Hierarchical. This model corresponds to the SLP denoted “H, P” in Fig. 3. We
are using a non-centered parameterization to allow for better sampling
performance from the HMC sampler.

$\displaystyle\sigma_{\alpha}$ $\displaystyle\sim\text{Exponential}(1),$
$\displaystyle\mu_{\alpha}$ $\displaystyle\sim\mathcal{N}(0,10),$
$\displaystyle\epsilon_{c}$
$\displaystyle\sim\mathcal{N}(0,1)~{}~{}~{}\text{for each county }c,$
$\displaystyle\alpha_{c}$
$\displaystyle=\mu_{\alpha}+\sigma_{\alpha}\,\epsilon_{c},$ (30)
$\displaystyle\beta$ $\displaystyle\sim\mathcal{N}(0,10),$
$\displaystyle\sigma$ $\displaystyle\sim\text{Exponential}(5),$ $\displaystyle
f_{i}$ $\displaystyle=\alpha_{c[i]}+\beta\,x_{i},$ $\displaystyle y_{i}$
$\displaystyle\sim\mathcal{N}(f_{i},\sigma^{2})$ (31)

Group-level predictor. This model corresponds to the SLP denoted “G, P” in
Fig. 3.

$\displaystyle\gamma_{0}$ $\displaystyle\sim\mathcal{N}(0,10),$
$\displaystyle\gamma_{1}$ $\displaystyle\sim\mathcal{N}(0,10),$
$\displaystyle\sigma_{\alpha}$ $\displaystyle\sim\text{Exponential}(1),$
$\displaystyle\epsilon_{c}$
$\displaystyle\sim\mathcal{N}(0,1)~{}~{}~{}\text{for each county }c,$ (32)
$\displaystyle\alpha_{c}$
$\displaystyle=\mu_{\alpha,c}+\sigma_{\alpha}\,\epsilon_{c},$
$\displaystyle\beta$ $\displaystyle\sim\mathcal{N}(0,10),$
$\displaystyle\sigma$ $\displaystyle\sim\text{Exponential}(5),$ $\displaystyle
y_{i}$ $\displaystyle\sim\mathcal{N}(f_{i},\sigma^{2})$ (33)

where $\mu_{\alpha,c}=\gamma_{0}+\gamma_{1}\,u_{c}$ and
$f_{i}=\alpha_{c[i]}+\beta\,x_{i}$.

To ensure we have a balanced representation of data in each county in both the
training and the testing data we apply stratified sampling: for each county,
we hold out 20 % of the observations for evaluation. For this dataset we do
not run stacking with a validation set because of the limited amount of data
available per county. For DCC inference, we collect $2000$ HMC samples with
$2000$ samples for burn-in for each SLP. For RJMCMC, the transition kernel
picks a modelling choice of $\alpha$ and $\beta$ and then samples the local
parameters for each modelling choice from the prior. The standard deviation
$\sigma$ is updated separately with a Metropolis-Hastings transition kernel.
To ensure a fair comparison between DCC and RJMCMC we collect the same total
number of samples.

###  F.6 Stacked RJMCMC

Figure 13: Difference in LPPD between RJMCMC and other methods (higher is
better).

We conducted further experiments to determine the impact of running stacking
on top of RJMCMC, we call this _Stacked RJMCMC_. The results can be viewed in
Fig. 13 where we plot the difference in LPPD between RJMCMC and other methods,
i.e.
$\mathrm{LPPD}_{\text{Diff}}=\mathrm{LPPD}_{\text{Other}}-\mathrm{LPPD}_{\text{RJMCMC}}$.
Note, compared to the main paper we here evaluate the difference in LPPD to
RJMCMC and not to stacking on top of DCC. This is because we care about
investigating the performance improvement relative to RJMCMC. We observe in
Fig. 13 that Stacked RJMCMC generally leads to higher LPPD than RJMCMC. The
difference is positive for all problem settings. Further, the difference is
statisticaly significant under a Wilcoxon-signed rank test in all problems,
except for _California_ and _Stroke_.

##  Appendix G PAC-Bayes and Stacking

For completeness, we first give a brief introduction into PAC-Bayes which is
mostly based on Morningstar et al. (2022). Morningstar et al. (2022) assume a
setting in which data is sampled i.i.d. from $\tilde{y}_{\ell}\sim
p_{\text{true}}(\tilde{y})$, that we have a parameterized probability model
$p(\tilde{y}\mid\theta)$ and we want to find a mechanism to fit a
distribution, $q(\theta)$, over the parameters of the probability model.

The true predictive risk is given by

$\mathcal{P}(q):=-\mathbb{E}_{p_{\text{true}}(\tilde{y})}[\log\mathbb{E}_{q(\theta)}[p(\tilde{y}\mid\theta)]]$
(34)

and in most applications is the quantity we care about in the end because it
directly measures the quality of our predictions. However, in practice we
cannot evaluate the true predictive risk because we do not have access to the
true data generating distribution. The goal of PAC-Bayes methods is then to
provide a stochastic upper bound on $\mathcal{P}(q)$ which can be used to
train $q(\theta)$ (Masegosa, 2020; Morningstar et al., 2022).

The empirical predictive risk

$\overline{\mathcal{P}}(q):=-\frac{1}{L}\sum_{\ell=1}^{L}\log\mathbb{E}_{q(\theta)}[p(\tilde{y}_{\ell}\mid\theta].$
(35)

is an empirical estimate of the true predictive risk. _Ensemble methods (which
include stacking) directly minimize the empirical predictive risk._ For
example, our stacking objective in Eq. (15) is a particular instantiation of
the empirical predictive risk. We will explain the connection in more detail
below.

The true inferential risk

$\mathcal{R}(q):=-\mathbb{E}_{p_{\text{true}}(\tilde{y})}[\mathbb{E}_{q(\theta)}[\log
p(\tilde{y}\mid\theta)]]$ (36)

is an _upper bound on the true predictive risk_ (by applying Jensen’s
inequality). In the case that our model is well specified, i.e.
$\exists\theta^{\prime}$ s.t.
$p_{\text{true}}(\cdot)=p(\cdot\mid\theta^{\prime})$, then
$\operatorname*{argmin}\mathcal{P}(q)=\operatorname*{argmin}\mathcal{R}(q)$.
Hence, when our model is well specified minimizing the true inferential risk
is equivalent to minimizing the true predictive risk.

The empirical inferential risk

$\overline{\mathcal{R}}(q):=-\frac{1}{L}\mathbb{E}_{q(\theta)}[\log
p(\tilde{y}_{\ell}\mid\theta)]$ (37)

is the empirical estimate of the true inferential risk. _Minimizing this risk
directly is equivalent to maximum likelihood estimation._ By adding an extra
regularization term to the empirical inferential risk we get the PAC-
inferential risk, given by

$\widetilde{\mathcal{R}}(q;r,\beta):=\mathbb{E}_{q(\theta)}\left[-\frac{1}{L}\sum_{\ell=1}^{L}\log
p(\tilde{y}_{\ell}\mid\theta)+\frac{1}{\beta
L}\log\frac{q(\theta)}{r(\theta)}\right]$ (38)

where $r(\theta)$ is a user specified prior distribution on the parameters
$\theta$. This is a stochastic upper bound on the true predictive risk.

Morningstar et al. (2022) use results from Burda et al. (2015) to tighten the
PAC-inferential risk and introduce the PACM bound

$\widetilde{\mathcal{P}}_{M,L}(q;r,\beta):=-\frac{1}{L}\sum_{\ell=1}^{L}\mathbb{E}_{q(\theta^{M})}\left[\log\left(\frac{1}{M}\sum_{j=1}^{M}p(\tilde{y}_{\ell}\mid\theta_{j})\right)\right]+\frac{1}{\beta
L}\mathrm{KL}(q(\theta)\parallel r(\theta)).$ (39)

Notably, as the PAC-inferential risk, this bound contains a regularization
term which is meant to prevent overfitting.

######  Theorem G.1 (Morningstar et al. (2022)).

For all $q(\theta)$ absolutely continuous with respect to $r(\theta)$,
$\tilde{y}_{\ell}\sim p_{\text{true}}(\tilde{y})$ i.i.d.,
$\beta\in(0,\infty)$, $L,M\in\mathbb{N}$,
$p(\tilde{y}\mid\theta)\in(0,\infty)$ for all $\\{\theta\in\Theta\mid
p_{\text{true}}(\tilde{y})>0\\}\times\\{\theta\in\Theta\mid r(\theta)>0\\}$,
and $\xi\in(0,1)$, then with probability at least $1-\xi$,

$\mathcal{P}(q)\leq\widetilde{\mathcal{P}}_{M,L}(q;r,\beta)+\psi(p_{\text{true}},\beta,M,L,r,\xi)-\frac{1}{\beta\,M\,L}\log\xi$
(40)

and furthermore (unconditionally)

$\widetilde{\mathcal{P}}_{M+1,L}(q;r,\beta)\leq\widetilde{\mathcal{P}}_{M,L}(q;r,\beta)$
(41)

where $\widetilde{\mathcal{P}}_{M,L}(q;r,\beta)$ as in Eq. (39) and:

$\displaystyle\psi(p_{\text{true}},\beta,M,L,r,\xi)$
$\displaystyle:=\frac{1}{\beta\,M\,L}\log\mathbb{E}_{p_{\text{true}}(\tilde{y}^{L})}\mathbb{E}_{r(\theta^{M})}\left[\exp(\beta
LM\Delta(\tilde{y}^{L},\theta^{M}))\right],$ (42)
$\displaystyle\Delta(\tilde{y}^{L},\theta^{M})$
$\displaystyle:=\frac{1}{L}\sum_{\ell=1}^{L}\log\left(\frac{1}{M}\sum_{j=1}^{M}p(\tilde{y}_{\ell}\mid\theta_{j})\right)-\mathbb{E}_{p_{\text{true}}(\tilde{y})}\left[\log\left(\frac{1}{M}\sum_{j=1}^{M}p(\tilde{y}\mid\theta_{j})\right)\right].$
(43)

For a proof see Appendix C.3 of Morningstar et al. (2022).

###  G.1 From PAC-Bayes to Regularized Stacking

In order to connect the ideas from Morningstar et al. (2021) with the stacking
objective for probabilistic programs we need to define a specific form for the
parameterized probability model and the distribution $q$ which is meant to be
optimized. We choose the latent variable of our probabilitiy model to be the
random variable $k$ which indexes into the SLPs. Our probabilistic model then
turns out to be $p(\tilde{y}\mid k)=\rho_{k}(\tilde{y})$ and our approximate
posterior $q(k)$ is a categorical distribution over $\\{1,\dots,K\\}$
parameterized by the weights $w$, i.e.
$q(k)=\text{Categorical}(w_{1},\dots,w_{K})$. In this formulation the true
predictive risk is given by

$\displaystyle\mathcal{P}(q)$
$\displaystyle=-\mathbb{E}_{p_{\text{true}}(\tilde{y})}[\log\mathbb{E}_{q(k)}[\rho_{k}(\tilde{y})]]$
(44)
$\displaystyle=-\mathbb{E}_{p_{\text{true}}(\tilde{y})}\left[\log\left(\sum_{k=1}^{K}w_{k}\rho_{k}(\tilde{y})\right)\right]$
(45)

which is equivalent to Eq. (11) in the main paper. The PACM bound becomes

$\widetilde{\mathcal{P}}_{M,L}(q;r,\beta)=-\frac{1}{L}\sum_{\ell=1}^{N}\mathbb{E}_{q(k^{M})}\left[\log\left(\frac{1}{M}\sum_{j=1}^{M}\rho_{k_{j}}(\tilde{y}_{\ell})\right)\right]+\frac{1}{\beta
L}\mathrm{KL}(q(k)\parallel r(k)).$ (46)

Note that $\beta=1$ can be interpreted as the “standard Bayesian” setting as
it provides an equal weighting of the prior term $\mathrm{KL}(q(k)\parallel
r(k))$ and likelihood term
$\sum_{\ell=1}^{N}\mathbb{E}_{q(k^{M})}\left[\log\left(\frac{1}{M}\sum_{j=1}^{M}\rho_{k_{j}}(\tilde{y}_{\ell})\right)\right]$.
We can rewrite the sum inside the log as a sum over SLPs, as follows

$\widetilde{\mathcal{P}}_{M,L}(q;r,\beta)=-\frac{1}{L}\sum_{\ell=1}^{L}\mathbb{E}_{q(k^{M})}\left[\log\left(\sum_{k=1}^{K}\frac{\lvert
I_{k}\rvert}{M}\rho_{k}(\tilde{y}_{\ell})\right)\right]+\frac{1}{\beta
L}\mathrm{KL}(q(k)\parallel r(k))$ (47)

where $I_{k}:=\\{k_{j}\mid j=\\{1,\dots,M\\},k_{j}=k\\}$ is the set of all
parameter samples that are equal to $k$. Now by the law of large numbers (LLN)
as $m\to\infty$ this converges to

$\widetilde{P}_{\infty,L}(q;r,\beta)=R_{\beta}(w_{1:K})=-\frac{1}{L}\sum_{\ell=1}^{L}\log\left(\sum_{k=1}^{K}w_{k}\rho_{k}(\tilde{y}_{\ell})\right)+\frac{1}{\beta
L}\mathrm{KL}(q(k)\parallel r(k)).$ (48)

The first term in this objective is now the stacking objective and the second
term is a regularization term pushing the distribution over weights to the
prior $r(k)$. In other words, this is Eq. (15) with an added regularization
term.

There is one caveat with pushing $M\to\infty$, as shown by Morningstar et al.
(2021) the slack term $\psi$ that controls the tightness of the bound grows
linearly in $M$. Hence, the bound becomes vacuous for infinite $M$. However,
there are several reasons why in our setting it is still reasonable to work
with the objective $\widetilde{P}_{\infty,L}(q;r,\beta)$. First of all, in our
specific setting the special case $M=1$ is the degenerate setting in which we
collapse onto a single SLP. This is exactly the behaviour we are trying to
avoid in the first place. Additionally, Morningstar et al. (2022) have shown
that performance increases with using larger $M$ and show empirically that
using $\widetilde{P}_{\infty,L}(q;r,\beta)$ can be beneficial when it is
possible to do so. The only practical consideration they mention for
increasing $M$ are issues with gradient variance which are not applicable in
our setting.

To choose the prior one could be inclined to use the posterior SLP weights
$Z_{k}/\sum_{k^{\prime}}Z_{k^{\prime}}$, as they are our Bayesian beliefs
about which SLP has generated the data. However, as we have argued in the main
text these weights can be very unstable and expensive to estimate. Hence, a
reasonable default choice could be to use the discrete uniform distribution.
This will further discourage stacking from collapsing onto a single SLP. For
the special case of the prior $r(k)$ being chosen to be the uniform
distribution the objective further simplifies to

$R_{\beta}(w_{1:K})=\,\frac{1}{L}\sum_{\ell=1}^{L}\log\left(\sum_{k=1}^{K}w_{k}\rho_{k}(\tilde{y}_{\ell})\right)-\frac{1}{\beta
L}\,(\mathrm{H}[q(k;w_{1:K})]+\log K)$ (49)

where $\mathrm{H}$ denotes the Shannon entropy.