Foundation Posteriors for Approximate Probabilistic Inference
Abstract
Probabilistic programs provide an expressive representation language for generative models.
Given a probabilistic program, we are interested in the task of posterior inference: estimating a latent variable given a set of observed variables.
Existing techniques for inference in probabilistic programs often require choosing many hyper-parameters, are computationally expensive, and/or only work for restricted classes of programs.
Here we formulate inference as masked language modeling: given a program, we generate a supervised dataset of variables and assignments, and randomly mask a subset of the assignments.
We then train a neural network to unmask the random values, deﬁning an approximate posterior distribution.
By optimizing a single neural network across a range of programs we amortize the cost of training, yielding a “foundation” posterior able to do zero-shot inference for new programs.
The foundation posterior can also be ﬁne-tuned for a particular program and dataset by optimizing a variational inference objective.
We show the efﬁcacy of the approach, zero-shot and ﬁne-tuned, on a benchmark of STAN programs.
Introduction
The primary goal of probabilistic programming is to enable practitioners from any domain to easily reason about random variables of interest [29, 61].
The main challenge is to build posterior inference algorithms that are both efﬁcient for practical usage and universal – working for any program that might be written.
Many probabilistic programming languages (PPLs) have been developed [27, 45, 28, 44, 18, 58, 47, 53, 59, 12, 23, 17, 7, 57] each with inference approaches with strengths and weaknesses.
Some require users to assist in the inference process by programmatically specifying conditional independencies [7] or hand-crafting variational distributions [52, 7, 17].
Others place strong restrictions on what kinds of programs can be represented [12].
In this paper, we propose a complementary approach to scale inference for probabilistic programs.
By treating a program as a constrained form of language with some additional nomenclature, we show that probabilistic inference can be viewed as masked language modeling [62, 19], a technique popular in natural language processing.
Since little is assumed about the program, we need impose few constraints on users upfront.
Further, the generality of treating inference as a text prediction problem naturally enables more advanced features, such as plating, with little additional work.
We call this approach masked language inference , or MLI.
While MLI can be used to solve inference queries for a single program, it becomes more powerful when meta-amortized (or twice-amortized) to do inference across different queries and programs.
Success requires a dense set of programs that cover sufﬁcient diversity for the meta-algorithm to generalize.
Since this is not always obtainable, we propose program augmentations to enlarge a small dataset of programs.
The result of meta-amortized MLI on the augmented programs is a foundation posterior : a large neural network trained to do inference across many probabilistic 36th Conference on Neural Information Processing Systems (NeurIPS 2022). programs.
In experiments, we ﬁnd the foundation posterior to be capable of both zero-shot inference and variational ﬁne-tuning: given a program from the test set, we can achieve higher quality using the foundation posterior as an initial distribution.
Indeed, ﬁne-tuning the foundation posterior gives the best performance, for a wide range of runtimes, for a set of standard STAN programs.
Background Probabilistic Programs
We assume a class of probabilistic programs without loops and such that every line assigns a value to a variable.
We note that all ﬁnite loops can be “unrolled” into a program without loops.
Our constraint primarily disallows probabilistic programs with undecided runtime.
Later, we relax this assumption to support loops over conditionally independent variables where unrolling may not be practical.
Variable assignment can take many different forms including sampling x ⇠ gaussian ( z, 1) , direct assignment x = z , and function evaluation x = sqrt ( z ) .
Importantly, we do not constrain what distributions and functions are allowed.
Approximate Inference
Let p ( x, z ) be a joint distribution of latent variables z and observed variables x .
An inference query seeks to compute posterior beliefs p ( z | x ) =
⇣ p ( x,z ) p ( x ) ⌘ which is usually intractable as computing p ( x ) =
R z p ( x, z ) dz faithfully requires solving a difﬁcult integral.
So, we settle for approximate techniques.
Markov Chain Monte Carlo (MCMC) [32, 24] and variational inference (VI) [37, 64, 8] are two widely used examples.
Focusing on the latter, VI introduces a family of tractable distributions Q over the latent variables to ﬁnd the member q ⇤ 2 Q that minimizes the Kullback-Leibler (KL) divergence between itself and the exact posterior q ⇤ ( z ) =
arg min q KL ( q ( z ) || p ( z | x )) .
Once found, q ⇤ serves as a proxy for the true posterior.
VI as described above ﬁnds q ⇤ for a single observation x
but we commonly need to solve multiple inference queries of the same kind but for different values of the observed variables.
To amortize [26] the computational cost, we learn a single deterministic mapping f  : X !
Q that predicts (the parameters of the distribution)
q ⇤ as a function of x .
When scoring a value z , we often write the expression
f  ( x )( z ) as a conditional distribution q  ( z | x ) .
To ﬁnd the function f  , we maximize a lower bound of E p ( x ) log p ( x ) using Jensen’s inequality to get L elbo (  ) = E p
[Equation]
i where p ( x ) is a distribution over the observations.
For probabilistic programs, p ( x, z ) is easy to sample from (by simply executing the program) but difﬁcult to score.
In these cases, we apply a useful trick by optimizing KL ( p || q ) rather than KL ( q || p ) :
[Equation]
Equation 1 is called compiled inference [40].
Dropping the constant, the remaining expression is close to a supervised objective: choose the parameters of the distribution q ( z | x ) that maximizes the likelihood of ( x, z ) ⇠ p ( x, z ) , samples from the probabilistic program.
Masked Language Modeling
A seemingly unrelated technique is masked language modeling (MLM) [19], used widely by large language models like BERT [19].
The MLM objective is a word prediction task [56].
Given an input sentence, each token has a probability of being replaced by a special mask token.
The model then predicts the original token given the rest of the sentence.
In slightly more notation, let x =
( t 1 , . .
, t i  1 , <mask> , t i +1 , . . . ,
t n ) be a sequence of tokens that form a sentence where the i -th token was randomly masked.
The label is then just y
=
t i .
(This presentation assumes only one token is masked for simplicity whereas in general, multiple tokens can be masked in a single sentence.)
Then, the MLM objective is written as L mlm =
E p ( x ) [log p ( y | x )] where log p ( ·|· ) is cross entropy applied to softmax beliefs over the full vocabulary.
Masked Language Inference
We now connect approximate inference to MLM.
Suppose we are given a single probabilistic program; we wish to do inference, but do not know apriori which variables in the program are observed and which are latent.
During evaluation, an inquirer might hand you observations for any subset of the variables and present you with inference queries for the remaining unknown variables.
To solve this challenge we aim to ﬁnd a dataset of observed and latent assignments from which we can train a model to perform inference.
We leverage the programitself to generate such a dataset.
To do this, we edit the program slightly:
For every line (which by design must be a variable assignment), we add an annotation with the value that the variable takes for a single execution of the statement.
We do this with the syntax variable = expression !
value .
For instance, x ⇠ gaussian (0 , 1) ! 0 .
Unlike most prior PPLs, our program does not include an explicit observe statement.
Rather the values annotated on the right side of each arrow is the observation for the variable in the statement.
Given a probabilistic program, execute it and annotate each line with assignments.
Next, for each line, randomly mask the assignment with some probability.
In other words, x ⇠ gaussian (0 , 1) ! <mask> .
Save this masked program and the true assignments for all masked variables as a single data point.
To generate a dataset, we loop this procedure until we have a sufﬁciently large corpus.
Since the masking decision is made independently for each line, we can theoretically generate every possible permutation of observed and latent variables.
3.1 Objective
The dataset of programs as described is not far from a natural language corpus used by BERT.
If we could use it to teach a model to unmask assignments of latent variables, then that is tantamount to inference.
We formalize this intuition into a procedure we call masked language inference , or MLI.
See Fig.
1 for an overview.
The objective for MLI is a sum of two loss functions: one for traditional
MLM and the other to score the true assignment for a latent variable under a posterior distribution parameterized by a neural network.
In more detail, given a program x =
( t 1 , . . .
t n ) , we make two versions: for the ﬁrst x mlm , we do what MLM typically does, choosing random tokens in the program string to mask; for the second, x inf , we do as described above and randomly mask assignments (the values to the right of the ! symbol per line).
Note that the two inputs x mlm and x inf have different tokens masked.
We use a transformer network f ✓ : X !
R d that takes a raw string as input where ✓ are trainable parameters.
We compute v mlm =
f ✓ ( x mlm ) and v inf = f ✓ ( x inf ) , both of which are a sequence of contextual vector embeddings v = ( v 1 , . . .
v n ) , one for each token.
For each masked index i , we perform two actions, one for each input.
For MLM, we have a classiﬁcation head g ✓ : R d ! | V | (e.g. a linear layer) that maps a vector v mlm
[ i ] to a probability for each token in the vocabulary V .
For inference, we have an inference head h ✓ : R d !
Q that maps a vector v inf [ i ] to an approximate posterior distribution q
( t i | x inf ) in the family Q for the latent variable corresponding to token t i .
For continuous variables, a common choice for Q is the Gaussian family, in which the network h ✓ would return two vectors representing the mean and standard deviation.
Alternatively, if the latent variable is binary, we might choose the family Q to be Bernoulli where h ✓ returns a probability vector.
The choice of distribution is ﬂexible as long as scoring is differentiable (though sampling need not be). [Table]
In summary, the MLI objective is:
[Equation]
where ↵ is a weight balancing the two losses [Table] and t mlm , t inf are masked tokens.
While Equation 2 only masked a single token per loss, in practice we randomly mask 15% of tokens in x mlm
[Table]
and mask an increasing amount of assignments in x inf according to a schedule: we begin at 15% but increase this masking probability throughout training to 50%, thereby increasing the difﬁculty of inference.
The
MLI objective in Equation 2 can be maximized with stochastic gradient descent.
3.2 Orderless Auto-Regressive Decoding
After training, given a probabilistic program with m latent variables z 1 , . . .
z m , there are a number of approaches to use the approximate posterior q to perform inference.
The simplest approach is to unmask all latent variables at once using h ✓ .
However, we might believe that inferring the value of z 1 , for example, provides new information when inferring the value of z 2 .
Given some ordered queue of latent variables – e.g. z 1 , z 2 , . . .
, z m – we can loop m times, each time embedding the program using the transformer f ✓ and unmasking the next latent variable in the queue
z i using the inference head h ✓ .
Then, we remove z i from the queue and repeat for z
i +1 .
Each iteration of the loop produces a new observation, thus changing the program tokens and the resulting embedding through f ✓ .
However, our choice of ordering in the queue was arbitrary.
Because MLI trains the inference model by randomly masking random variables, any combination of observed and latent variables is within domain.
Thus, we are free to choose any random order of variables z 1 to z m .
Throughout this auto-regressive decoding, the proportion of latent to observed variables will shrink over time.
In the ﬁrst iteration, there will be the maximum number of latent variables whereas in the last iteration, there will only be one.
As such, a properly trained MLI model must handle different numbers of latent variables, not just different variables being observed.
In Section 3.1, we described how the percentage of masked assignments is varied in training.
A beneﬁt of this design is that both programs with very few and very many latent variables are in-distribution.
3.3 Toy Experiments
To build intuition and demonstrate efﬁcacy, we consider six probabilistic programs used in prior work [13].
These programs implement popular models like Gaussian mixture models, hierarchical latent variable models, astronomical models, and incorporate library imports like the Rosenbrock function whose code is not speciﬁed in the program.
We recommend the reader refer to Section A.3 for details.
While these programs are simple, they specify wide prior distributions that result in high sample diversity over repeated executions, which should make the MLI task more challenging.
Evaluation Metrics
We build a test set with 1,000 new executions of the program not used in training and randomly mask assignments.
This set is held constant across runs and ablations.
We evaluate performance with two metrics.
First, we compute log probability of unmasked assignment under the approximate posterior speciﬁed by the model, averaged over all masked tokens per execution and all test executions.
The larger this value is, the better the inference.
Second, we compute the variance of the log importance weights (IW), as in [68].
Since the approximate posterior q acts as an importance distribution estimating the true posterior, we measure the quality by the variance of p ( x,z ) q ( z | x ) , where lower variance is more desirable.
If q ( z | x ) = p ( z | x ) , the importance weight would be a ﬁxed constant, meaning zero variance.
[Figure]
Results
Table 4 reports the performance on the six programs.
We observe large probabilities (log close to 0) and small variance of log IW, which suggests good generalization of inference to new sample values.
As a baseline, we include an ablation to MLI by removing the MLM term, which amounts to training the inference head only.
We ﬁnd smaller log probabilities (every log point is a signiﬁcant difference) and much larger variance.
Together, this suggests that the MLM term is important for generalization, likely as it helps the neural network understand program structure.
For analysis on the distribution of variances, see Section A.3.
3.4 Visualizing Attention Maps
Prior works [6, 63, 1, 14] using transformer networks in natural language have re-purposed the attention weights in the later layers as an mechanism to introspect model logic.
In a similar vein, we leverage attention weights to approximate a dependency graph between random variables in the problem
. Datasets
We study simple probabilistic programs that exhibit independence and conditional independence between random variables.
First, in Figure 2(a), z 1 and y 1 are two independent
Gaussian variables.
We should expect independence between the sets { y 1 , y 2 } and { z 1 , z 2 } .
Second, in Figure 2(b), we design a graphical model with a Bernoulli random variable z , and pick a and b using if statements to be conditionally independent given z .
We add two Gaussian variables x and y that independently add noise around a and b , i.e. x ⇠ N ( a, 1) .
Finally, in Figure 2(c,d), we create a common effect model with two potential causes: choose x and y to be independent Bernoulli variables, and set z to be 1 if x or y else 0 .
We expect p ( x | z, y ) to differ meaningfully from p ( x | z ) .
[Figure]
Results
We report metrics similar to Table 4 for these programs in Section A.4.
Here, we instead visualize attention: suppose a single assignment in a program is masked.
When inferring its value, we can study the attention weights (obtained from the last transformer layer) on all other tokens in the program.
Now, divide the program into lines.
Since each line is associated with a variable declaration, we can treat the average weight for all tokens in that line as the “attention” the network is paying to the declared variable when doing inference (for the masked variable).
We hypothesize that this measure will be correlated to the graphical dependency between the two variables.
Figure 3 shows the heatmap of attention weights for the three programs.
For subplots (a) and (b), we mask each variable’s assignment one at a time.
The y-axis shows which variable is masked.
In the ‘Independent Gaussians‘ program, the model approximates the independence between { z 1 , z 2 } and { y 1 , y 2 } .
For instance, in the ﬁrst row of (a), when unmasking y 1 , we see the network upweights the tokens corresponding to y 1 and y 2 , but downweights the tokens for z 1 and z 2 .
The other rows show a consistent pattern.
While it may seem peculiar at ﬁrst for the model to pay attention to y 1 when unmasking y 1 , we note that this is important for the model to understand what kind of variable y 1 is.
Similarly, in the ‘Conditional Independence‘ program, the model approximates the independence between { a, x } and { b, y } given z .
Looking at the ﬁrst row, we see that unmasking a focuses on variables x and z and downweights b and y .
Finally, in subplot (c), the y-axis shows which variables are observed rather than which are masked, for clarity.
The bottom row plots attention weights if we observe both z and y whereas in the top row, we only observe z .
We see a stark contrast: in the bottom, the weights are roughly uniform; on top, the model gives little weight to y to unmask x .
Meta-Amortized Inference
Rather than restrict to a single program, we wish to use MLI to do meta-amortized inference , or zero-shot inference across programs.
Given a dataset of programs for training, how well can we do inference out-of-the-box for a new, unseen program?
We assume this novel program is within the same meta-distribution as the training set of programs; otherwise the problem is unsolvable.
4.1 Program Augmentations
From prior work [30, 16, 13], we know that the meta-training set must be crafted carefully for meta-amortized inference to perform well: the model must see enough programs that ‘span’ the program space to be able to generalize to new programs from the meta-distribution.
Collecting such a dataset for probabilistic programs is unfeasible.
Instead, we propose a set of program augmentations , inspired by data augmentations, to enlarge a small set of programs.
The hope is that the larger dataset provides enough coverage to generalize to a new test program.
We propose the following augmentations family, which can be recursively stacked on each other:
Fuzz Function Replaces a primitive function or distribution with a randomly sampled function or distribution.
Replacements are constrained to have the same number of arguments as the original function.
For example, the program z = rosenbrock ( a, b ) might be augmented to z ⇠ N ( a, b ) .
Fuzz Constant Constants are replaced by newly sampled constants from known prior distributions.
For example, r ⇠ N (0 , 1) becomes r ⇠ N (0 .
Line Swaps Swaps two independent program lines such that the dependency graph between random variables is unchanged.
For example, a program with three lines u ⇠ N (0 , 1); v ⇠ N (1 , 1); r = u + v could be augmented to v ⇠ N (1 , 1); u ⇠ N (0 , 1); r = u + v
but it would not be possible to swap the third line with any other given the dependency of u and v on r .
Cut and Glue Replace the usage of a variable with another variable already deﬁned in the program.
Note that this augmentation changes the dependency graph.
For example, u = a + b ; v = sqrt ( u ) could become u = a
+ b ; v = sqrt ( b ) where b is deﬁned earlier in the program.
Create and Use Creates a new random variable and uses it in lieu of another variable or constant in the right-hand side of an expression.
Note that this augmentation introduces a new random variable.
For example, u = a + b ; v = sqrt ( u ) could become u = a
+ b ; r ⇠ N (0 , 1); v = sqrt ( r ) .
Given that many of these augmentations meaningfully change the program by removing, editing or adding random variables, repeated applications of random augmentations can create novel programs in structure and content.
In practice, not all augmentations when compounded result in legitimate programs.
For instance, the augmentation rosenbrock (1 ,  1) !
N (1 ,  1) by fuzz function is ill-deﬁned.
We perform rejection sampling and discard improper programs.
4.2 Toy Experiment with Augmentations
To test program augmentations, we revisit the six programs from Section 3.3.
However, now we use ﬁve of them for meta-training, and hold out the “Rosenbrock” program for meta-test.
We apply up to ﬁve random program augmentations to each program in both sets to create the training and test splits.
Note that the test inputs are all derived from a novel program unseen by the model in training.
This is a much more difﬁcult generalization problem than prior experiments.
As an important caveat, we assume knowledge of which external functions might be used.
In the ‘Rosenbrock’ program, we use the rosenbrock external function.
We assume access to this when augmenting programs in the meta-training set such that rosenbrock may appear in training programs.
Results
We report results in Table 2.
The top row shows the log likelihood and the variance of importance weights for MLI.
The second row shows an ablation without augmentations (where we generate a dataset of equivalent size by re-executing programs as we did in Section 3.3).
We observe a 100x improvement with augmentations, suggesting better generalization to unseen programs.
Variational Finetuning
So far we have only studied “zero-shot” inference where an amortized inference model, given a new program, must do inference without any new computational expense.
In more realistic scenarios, where test programs can look quite different than the training set, obtaining high quality inference in a zero-shot manner can be challenging.
In this case, we propose to ﬁnetune our pretrained inference model using stochastic variational inference, or SVI [35, 50, 39].
Due to reparameterization requirements, in this subsection, we focus on inference of continuous random variables only.
For discrete random variables, future work can study using the concrete relaxation [42] or the family of
REINFORCE-based techniques [65, 46, 60].
More speciﬁcally, ﬁx a test program x
⇤ that we would like to do more high quality inference for.
Suppose x ⇤ has n observed data points d 1 , . . .
, d n , and m latent variables we want to estimate z 1 , . . .
z m .
Further, we are given ˆ ✓ , the parameters obtained from optimizing MLI on a pretraining dataset of programs.
Since the composed functions f · h : X !
Q map a probabilistic program to an approximate posterior, we deﬁne the shorthand q ˆ ✓ ( z 1: m | d 1: n ) =
h ( f ( x ⇤ )) .
We can optimize the evidence lower bound, which we recall is:
[Equation]
Since we are not amortizing over programs in this ﬁnetuning step, Equation 3 does not contain a second expectation over programs.
That is, z 1: m , d 1: n are the speciﬁc variables from x ⇤ .
We initialize ✓ = ˆ ✓ to leverage pretrained weights.
Note p ( . . . )
has no trainable parameters but acts as a likelihood scaling term for different values of z 1: m and thus, cannot be dropped.
Results
In Table 2, we include an ablation comparing the generalization of inference with and without ﬁne-tuning.
We ﬁne-tune the inference network (from MLI pretraining), optimizing Equationseparately for each program for 1,000 iterations.
As a baseline, we also ﬁnetune the network from a random initialization.
We observe that ﬁne-tuning MLI further improves log probability by roughlylog points with comparable IW variance.
On the other hand, ﬁne-tuning from scratch results in poor performance, highlighting the importance of pretraining.
Plating
In many applications of probabilistic inference, we have a dataset of observations that can be thousands of entries or more.
We would like to make inference queries given all entries but naive unrolling is unscalable.
Instead, we extend MLI to support ‘for’ loops over conditionally independent variables.
To do so, we will implement a form of minibatching in our inference algorithm.
Figure 4 shows such a plating example.
In the left program, we see 10 observations d 1 , . . .
, dsuch that d i ⇠ gaussian ( x, 1) where x is outside of the plate.
But suppose 10 examples are too [Table] many to unroll.
In the program on the right, we replace the ‘for’ loop with a plate(n) token where the argument n speciﬁes the number of total iterations.
Within the plate, we unroll two randomly subsampled iterations i 2 { 1 , 4 } .
For multiple executions of this program, the iterations within the plate change, just like minibatching in neural network training.
From the perspective of the transformer, this change does nothing more than add a new token to the vocabulary.
Why does this work?
Consider computing the program density p ( x, d 1 , . . .
, d n ) .
This can be rewritten as p ( x ) p ( d 1 , . . .
, d n | x ) of which the ﬁrst term is a given.
To compute the second term, observe that d 1 , . . .
, d n are conditionally independent on x .
So, p ( d 1: n | x ) =
Q n i =1 p ( d i | x ) .
Now:
[Equation]
where the minibatch is of size k ⌧ n .
We can make an unbiased (but higher variance) estimate of the full conditional distribution (and hence, the program density) using a small minibatch.
In the Appendix (Section A.8), we demonstrate plating on an item response theory (IRT) [21, 31, 51, 66, 67] model, a popular probabilistic program.
In the next section, we will leverage plating use MLI on a bank of real world probabilistic programs with thousands of observations each.
Foundation Posterior
The relationship between MLI and variational ﬁne-tuning is reminiscent of pretraining and downstream tasks, as popularized by self-supervised learning [19, 15].
Inspired by foundation models [9], we propose to frame the result of MLI as a foundation posterior that can be ﬁnetuned downstream using SVI to perform inference for a wide array of probabilistic programs.
The goal is to pay a large one-time cost in training a “general” inference model which may not be adequate for inference in all settings, but can be quickly adapted to new datasets with low cost.
To demonstrate the foundation posterior, we meta-amortize inference over a set of standard Stan [12] programs from PosteriorDB [43], a benchmark dataset for evaluating inference algorithms [4, 5, 69, 70, 20].
Setup We hold out three programs from PosterioDB for evaluation, and optimize the foundation posterior on the remaining set.
(Programs containing HMMs were removed as we are currently unable to support that graph structure.)
Plating with minibatches of size 5 is used for all programs to ﬁt observations within the transformer’s 512 token limit.
After pretraining, we optimize Equationfor each test program individually, varying the number of steps of ﬁne-tuning across 0 (zero-shot), 10, 100, and 1000.
As baselines, we use CmdStanPy to run Stan-native NUTS and ADVI.
For NUTS, we vary the number of mixing steps between 10 and 10k; for ADVI we vary the number of iterations between 100 and 1M. Each program comes with 10k pre-computed posterior samples ﬁt using long runs of expert-tuned NUTS in Stan, constituting the gold standard.
To evaluate inference quality we draw 10k samples and perform a hypothesis test between the gold and drawn samples.
If the p-value is high, we conclude the two sets of samples are likely from the same distribution.
Results Figure 5 reports results on held-out Stan programs.
We plot the cost of inference computation (wall-clock time) on the x-axis in log-seconds.
For MLI, we sum the cost of ﬁne-tuning and the cost of forward passes to sample.
For NUTS and ADVI, we ignore the cost of compilation, measuring mixing time and optimization time, respectively.
On the y-axis, we plot the p-value.
A better inference algorithm would bias towards the top left corner, achieving a higher p-value at lower [Figure] cost.
Each point in the plot represents a setting – we redo inference with different ﬁnetuning steps or mixing steps, etc.
Each line groups together a single inference algorithm.
First, observe that ADVI in Stan performs poorly (p-value < 0.5) despite being cheap.
In the KidIQ program, it is unable to surpass a p-value of 0, suggesting poor posterior samples.
Second, we see that NUTS is computationally expensive but converges to a p-value of 1 in all cases, as expected since the gold samples are derived from NUTS.
We observe foundation posterior to achieve a compromise between the two: it is cheaper than HMC but achieves higher inference quality than ADVI in Stan.
The left-most point in the MLI curve represents zero-shot inference, which we ﬁnd comparable to ADVI despite the latter training for up to 1M steps.
Finally, the gray line shows an additional baseline where we ﬁnetune a transformer from scratch.
We observe that initializing from a foundation posterior is critical to good inference.
Related Work
Amortized Inference
Traditionally, an inference query q ( z ) is solved for a single assignment of the observed variable x .
However, in many applications, we may be interested in solving the same inference query for many observations x 1 , x 2 , . . .
, x n .
Re-solving the same inference query from scratch for all n related queries seems wasteful.
Amortized inference [26, 55] was proposed as a more efﬁcient alternative by learning a function q ( z | x ) that maps a observed assignment x to a distribution over the latent variable z .
In doing so, we “amortize” the cost of inference with a large one-time cost in deﬁning q , after which an inference query can be solved with a single function application.
Without amortization, the foundation posterior would not be a very efﬁcient inference algorithm.
Meta-Inference
Meta-learning has been used to do inference amortized over both queries and a family of generative models [16, 30, 36].
These works have shown that “meta-amortized” inference can generalize well within the family of models it was trained on, even members of the family not explicitly seen in training.
The foundation posterior is a meta-amortized algorithm trained with masked language inference.
Unlike prior works, we not only evaluate zero-shot inference but also the ability to serve as an effective starting point for further ﬁne-tuning.
Most relevant to our work is a meta-learning approach [13] that builds a white-box inference algorithm by matching every step in a probabilistic program with neural network whose role is to “invert” it.
Foundation posteriors trade explainability for ﬂexibility.
Our approach is not white-box: inference is performed by a single black box neural network.
In return, our approach assumes no structure: it treats the program as nothing more than a string, making its generalization capacity greater than the approach in [13].
We view these two methods as optimizing for distinct properties.
Foundation Models
Foundation models [9] are a blanket term for large unsupervised models that map data to representations.
Examples include ResNet [33] in computer vision, BERT [19] and
GPT [10, 49] in natural language, Wav2Vec [54, 3] in speech, and CLIP [48] and Data2Vec [2] in multimodal applications.
Foundation models are treated as frozen backbones whereupon a small portion is ﬁnetuned for a downstream task.
Our motivation for the foundation posterior is precisely this: learn a good posterior “initialization” that can be quickly ﬁnetuned for downstream inference.
Discussion Limitations
First, optimizing transformer networks is computationally costly, requiring accelerated hardware.
For practitioners simply interested in a small set of inference queries, it is simpler to use
MCMC or VI.
Second, MLI as described is not sufﬁcient for all programs: in PosteriorDB, classes of graphical models like HMMs are difﬁcult to represent as text.
Third, our approach, being reliant on deep learning, is limited in its explainability.
In the case that inference fails, we are unable to reason about the root cause.
Finally, the success of foundation models [9] is largely predicated on large datasets with millions of entries.
Unfortunately, we do not know of any large collections of probabilistic programs other than PosteriorDB, which in comparison, is of modest scale.
Future Work
Future work could investigate foundation posteriors that act as proposal distributions for MCMC-based methods.
As a ﬁrst step, we run Metropolis Hastings on one of the six toy programs, and compute R-hat [25].
We achieve a score of 1 . 011 ( ± 0 .
009 over 100 test programs) when using the foundation posterior as the proposal compared to 1 . 043 ( ± 0 .
059 ) when using the prior distribution.
Best practice deems a chain sufﬁciently mixed if R-hat is less than 1 . 05 .
More analysis is needed to better understand the efﬁcacy of foundation posteriors in this context.
Summary
We proposed a meta-amortized inference algorithm for probabilistic programs using masked language modeling.
With this algorithm, we build a foundation posterior capable of fast zero-shot inference that can also be ﬁnetuned for more accuracy.
We are optimistic that the generality of the approach, despite its computational burden, can make for practical and scalable inference.