Foundation Posteriors for Approximate Probabilistic Inference
Anonymous Author(s)
AfﬁliationAddress email
Abstract
Probabilistic programs provide an expressive representation language for generativemodels.
Given a probabilistic program, we are interested in the task of posteriorinference: estimating a latent variable given a set of observed variables.
Existingtechniques for inference in probabilistic programs often require choosing manyhyper-parameters, are computationally expensive, and/or only work for restrictedclasses 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 networkto unmask the random values, deﬁning an approximate posterior distribution.
Byoptimizing a single neural network across a range of programs we amortize the costof training, yielding a “foundation” posterior able to do zero-shot inference for newprograms.
The foundation posterior can also be ﬁne-tuned for a particular programand dataset by optimizing a variational inference objective.
We show the efﬁcacyof 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 easilyreason about random variables of interest [29, 61].
The main challenge is to build posterior inferencealgorithms that are both efﬁcient for practical usage and universal – working for any programthat 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 strengthsand weaknesses.
Some require users to assist in the inference process by programmatically specifyingconditional independencies [7] or hand-crafting variational distributions [52, 7, 17].
Others placestrong restrictions on what kinds of programs can be represented [12].
In this paper, we proposea complementary approach to scale inference for probabilistic programs.
By treating a programas a constrained form of language with some additional nomenclature, we show that probabilisticinference can be viewed as masked language modeling [62, 19], a technique popular in naturallanguage processing.
Since little is assumed about the program, we need impose few constraintson users upfront.
Further, the generality of treating inference as a text prediction problem naturallyenables more advanced features, such as plating, with little additional work.
We call this approachmasked language inference , or MLI.
While MLI can be used to solve inference queries for a single program, it becomes more powerfulwhen 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 togeneralize.
Since this is not always obtainable, we propose program augmentations to enlargea small dataset of programs.
The result of meta-amortized MLI on the augmented programs isa foundation posterior : a large neural network trained to do inference across many probabilisticprograms.
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 usingthe foundation posterior as an initial distribution.
Indeed, ﬁne-tuning the foundation posterior givesthe best performance, for a wide range of runtimes, for a set of standard STAN programs.
BackgroundProbabilistic Programs
We assume a class of probabilistic programs without loops and suchthat every line assigns a value to a variable.
We note that all ﬁnite loops can be “unrolled” into aprogram without loops.
Our constraint primarily disallows probabilistic programs with undecidedruntime.
Later, we relax this assumption to support loops over conditionally independent variableswhere unrolling may not be practical.
Variable assignment can take many different forms includingsampling 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 observedvariables x .
An inference query seeks to compute posterior beliefs p ( z | x ) =
⇣ p ( x,z ) p ( x ) ⌘ which isusually 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] andvariational inference (VI) [37, 64, 8] are two widely used examples.
Focusing on the latter, VIintroduces a family of tractable distributions Q over the latent variables to ﬁnd the member q ⇤ 2 Q 54 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 multipleinference 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 (theparameters of the distribution)
q ⇤ as a function of x .
When scoring a value z , we often write theexpression
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 toget L elbo (  ) =
[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) butdifﬁ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 isclose to a supervised objective: choose the parameters of the distribution q ( z | x ) that maximizes thelikelihood 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 wordprediction task [56].
Given an input sentence, each token has a probability of being replaced by aspecial mask token.
The model then predicts the original token given the rest of the sentence.
Inslightly more notation, let x =
( t 1 , . .
, t i  1 , <mask> , t i +1 , . . . ,
t n ) be a sequence of tokens thatform a sentence where the i -th token was randomly masked.
The label is then just y
=
t i .
(Thispresentation assumes only one token is masked for simplicity whereas in general, multiple tokens canbe 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 probabilisticprogram; we wish to do inference, but do not know apriori which variables in the program areobserved and which are latent.
During evaluation, an inquirer might hand you observations forany subset of the variables and present you with inference queries for the remaining unknownvariables.
To solve this challenge we aim to ﬁnd a dataset of observed and latent assignments fromwhich we can train a model to perform inference.
We leverage the programitself to generate such adataset.
To do this, we edit the program slightly:
For every line (which by design must be a variableassignment), we add an annotation with the value that the variable takes for a single execution ofthe 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 explicitobserve statement.
Rather the values annotated on the right side of each arrow is the observationfor the variable in the statement.
Given a probabilistic program, execute it and annotate each linewith assignments.
Next, for each line, randomly mask the assignment with some probability.
In otherwords, x ⇠ gaussian (0 , 1) ! <mask> .
Save this masked program and the true assignments for allmasked variables as a single data point.
To generate a dataset, we loop this procedure until we have asufﬁciently large corpus.
Since the masking decision is made independently for each line, we cantheoretically 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 wecould use it to teach a model to unmask assignments of latent variables, then that is tantamount toinference.
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 distributionparameterized by a neural network.
In more detail, given a program x =
( t 1 , . . .
t n ) , we make twoversions: for the ﬁrst x mlm , we do what MLM typically does, choosing random tokens in the programstring to mask; for the second, x inf , we do as described above and randomly mask assignments (thevalues to the right of the ! symbol per line).
Note that the two inputs x mlm and x inf have differenttokens masked.
We use a transformer network f ✓ : X !
R d that takes a raw string as input where ✓ are trainableparameters.
We compute v mlm =
f ✓ ( x mlm ) and v inf = f ✓ ( x inf ) , both of which are a sequence ofcontextual vector embeddings v = ( v 1 , . . .
v n ) , one for each token.
For each masked index i , weperform 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 approximateposterior distribution q
( t i | x inf ) in the family Q for the latent variable corresponding to token t i .
Forcontinuous variables, a common choice for Q is the Gaussian family, in which the network h ✓ wouldreturn two vectors representing the mean and standard deviation.
Alternatively, if the latent variableis binary, we might choose the family Q to be Bernoulli where h ✓ returns a probability vector.
Thechoice of distribution is ﬂexible as long as scoring is differentiable (though sampling need not be).
In summary, the MLI objective is:
[Equation]
where ↵ is a weight balancing the two losses, and t mlm , t inf are masked tokens.
While Equation 2 onlymasked a single token per loss, in practice we randomly mask 15% of tokens in x mlm , and mask an [Table]
increasing amount of assignments in x inf according to a schedule: we begin at 15% but increase thismasking 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 numberof approaches to use the approximate posterior q to perform inference.
The simplest approach is tounmask all latent variables at once using h ✓ . However, we might believe that inferring the value ofz 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, eachtime embedding the program using the transformer f ✓ and unmasking the next latent variable in thequeue z
i using the inference head h ✓ .
Then, we remove z i from the queue and repeat for z
i +1 .
Eachiteration of the loop produces a new observation, thus changing the program tokens and the resultingembedding through f ✓ .
However, our choice of ordering in the queue was arbitrary.
Because MLItrains the inference model by randomly masking random variables, any combination of observed andlatent 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 shrinkover time.
In the ﬁrst iteration, there will be the maximum number of latent variables whereas in thelast iteration, there will only be one.
As such, a properly trained MLI model must handle differentnumbers of latent variables, not just different variables being observed.
In Section 3.1, we describedhow the percentage of masked assignments is varied in training.
A beneﬁt of this design is that bothprograms 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 latentvariable models, astronomical models, and incorporate library imports like the Rosenbrock functionwhose 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 samplediversity 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 intraining and randomly mask assignments.
This set is held constant across runs and ablations.
Weevaluate performance with two metrics.
First, we compute log probability of unmasked assignmentunder the approximate posterior speciﬁed by the model, averaged over all masked tokens per executionand 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 thequality by the variance of p ( x,z ) q ( z | x ) , where lower variance is more desirable.
If q ( z | x ) = p ( z | x ) , theimportance weight would be a ﬁxed constant, meaning zero variance. zyzy(a) z a b x y (b) z x y (c) z x y (
[Figure]
Results
Table 4 reports the performance onthe six programs.
We observe large probabilities(log close to 0) and small variance of log IW,which suggests good generalization of inferenceto new sample values.
As a baseline, we include an ablation to MLI by removing the MLMterm, which amounts to training the inferencehead only.
We ﬁnd smaller log probabilities(every log point is a signiﬁcant difference) andmuch larger variance.
Together, this suggeststhat the MLM term is important for generalization, likely as it helps the neural network understand program structure.
For analysis on thedistribution 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 theattention 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 theproblem.
Datasets
We study simple probabilistic programs that exhibit independence and conditionalindependence 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, 180 in Figure 2(b), we design a graphical model with a Bernoulli random variable z , and pick a andb using if statements to be conditionally independent given z .
We add two Gaussian variables xand y that independently add noise around a and b , i.e. x ⇠ N ( a, 1) .
Finally, in Figure 2(c,d), wecreate a common effect model with two potential causes: choose x and y to be independent Bernoullivariables, 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 insteadvisualize attention: Suppose a single assignment in a program is masked.
We ﬁrst compute the L 2norm over contextual vectors from the last layer for all tokens in the program.
Next, we segmentthe norm sequence by program line, and compute the average norm per segment.
Since each lineis associated with a variable declaration, we can treat the average norm for a line as the “attention”the network is paying to this variable when inferring the masked variable.
We hypothesize that thismeasure will be correlated to the dependency between the two variables.
Figure 3 shows the heatmap of attention norms for the three programs.
For subplots (a) and (b), wemask 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 thetokens corresponding to y 1 and y 2 , but downweights the tokens for z 1 and z 2 .
The other rows showa consistent pattern.
While it may seem peculiar at ﬁrst for the model to pay attention to y 1 whenunmasking 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 independencebetween { a, x } and { b, y } given z .
Looking at the ﬁrst row, we see that unmasking a focuses onvariables x and z and downweights b and y .
Finally, in subplot (c), the y-axis shows which variablesare observed rather than which are masked, for clarity.
The bottom row plots attention weights ifwe observe both z and y whereas in the top row, we only observe z .
We see a stark contrast: in thebottom, 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 , orzero-shot inference across programs.
Given a dataset of programs for training, how well can we doinference out-of-the-box for a new, unseen program?
We assume this novel program is within thesame 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 formeta-amortized inference to perform well: the model must see enough programs that ‘span’ theprogram space to be able to generalize to new programs from the meta-distribution.
Collecting such adataset 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 largerdataset provides enough coverage to generalize to a new test program.
We propose the followingaugmentations 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 randomvariables is unchanged.
For example, a program with three lines u ⇠ N (0 , 1); v ⇠ N (1 , 1); r = u + vcould be augmented to v ⇠ N (1 , 1); u ⇠ N (0 , 1); r = u + v
but it would not be possible to swapthe 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 inthe 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 oradding random variables, repeated applications of random augmentations can create novel programsin structure and content.
In practice, not all augmentations when compounded result in legitimateprograms.
For instance, the augmentation rosenbrock (1 ,  1) !
N (1 ,  1) by fuzz function isill-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 upto ﬁve random program augmentations to each program in both sets to create the training and testsplits.
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, weuse the rosenbrock external function.
We assume access to this when augmenting programs in themeta-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 ofimportance weights for MLI.
The second row shows an ablation without augmentations (where wegenerate a dataset of equivalent size by re-executing programs as we did in Section 3.3).
We observea 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 newprogram, 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 inferencein a zero-shot manner can be challenging.
In this case, we propose to ﬁnetune our pretrainedinference model using stochastic variational inference, or SVI [35, 50, 39].
Due to reparameterizationrequirements, in this subsection, we focus on inference of continuous random variables only.
Fordiscrete 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 estimatez 1 , . . .
z m .
Further, we are given ˆ ✓ , the parameters obtained from optimizing MLI on a pretrainingdataset of programs.
Since the composed functions f · h : X !
Q map a probabilistic program to anapproximate 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 asecond 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 likelihoodscaling 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 Equation 3separately for each program for 1,000 iterations.
As a baseline, we also ﬁnetune the network from arandom initialization.
We observe that ﬁne-tuning MLI further improves log probability by roughly 2log points with comparable IW variance.
On the other hand, ﬁne-tuning from scratch results in poorperformance, highlighting the importance of pretraining.
Plating
In many applications of probabilistic inference, we have a dataset of observations that canbe thousands of entries or more.
We would like to make inference queries given all entries but naiveunrolling is unscalable.
Instead, we extend MLI to support ‘for’ loops over conditionally independentvariables.
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 , . . .
, d 10such that d i ⇠ gaussian ( x, 1) where x is outside of the plate.
But suppose 10 examples are toomany to unroll.
In the program on the right, we replace the ‘for’ loop with a plate(n) token wherethe argument n speciﬁes the number of total iterations.
Within the plate, we unroll two randomlysubsampled iterations i 2 { 1 , 4 } .
For multiple executions of this program, the iterations within [Table] the plate change, just like minibatching in neural network training.
From the perspective of thetransformer, 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 berewritten 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 ) =
[Equation] p ( d i | x ) .
Now:
[Equation]
where the minibatch is of size k ⌧ n .
We can make an unbiased (but higher variance) estimate of thefull 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 MLIon 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 downstreamusing SVI to perform inference for a wide array of probabilistic programs.
The goal is to pay a largeone-time cost in training a “general” inference model which may not be adequate for inference in allsettings, but can be quickly adapted to new datasets with low cost.
To demonstrate the foundationposterior, 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 foundationposterior on the remaining set.
(Programs containing HMMs were removed as we are currentlyunable 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 Equation 4for 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 iterationsbetween 100 and 1M. Each program comes with 10k
pre-computed posterior samples ﬁt using longruns of expert-tuned NUTS in Stan, constituting the gold standard.
To evaluate inference quality wedraw 10k samples and perform a hypothesis test between the gold and drawn samples.
If the p-valueis 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-tuningand 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.
Abetter inference algorithm would bias towards the top left corner, achieving a higher p-value at lowercost.
Each point in the plot represents a setting – we redo inference with different ﬁnetuning steps ormixing 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 KidIQprogram, it is unable to surpass a p-value of 0, suggesting poor posterior samples.
Second, we see [Figure] that NUTS is computationally expensive but converges to a p-value of 1 in all cases, as expected sincethe gold samples are derived from NUTS.
We observe foundation posterior to achieve a compromisebetween 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 comparableto ADVI despite the latter training for up to 1M steps.
Finally, the gray line shows an additionalbaseline where we ﬁnetune a transformer from scratch.
We observe that initializing from a foundationposterior is critical to good inference.
Related WorkAmortized Inference
Traditionally, an inference query q ( z ) is solved for a single assignment ofthe observed variable x .
However, in many applications, we may be interested in solving the sameinference query for many observations x 1 , x 2 , . . .
, x n .
Re-solving the same inference query fromscratch for all n related queries seems wasteful.
Amortized inference [26, 55] was proposed as a moreefﬁcient alternative by learning a function q ( z | x ) that maps a observed assignment x to a distributionover the latent variable z .
In doing so, we “amortize” the cost of inference with a large one-timecost 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 afamily of generative models [16, 30, 36].
These works have shown that “meta-amortized” inferencecan generalize well within the family of models it was trained on, even members of the family notexplicitly seen in training.
The foundation posterior is a meta-amortized algorithm trained withmasked language inference.
Unlike prior works, we not only evaluate zero-shot inference but also theability to serve as an effective starting point for further ﬁne-tuning.
Most relevant to our work is a meta-learning approach [13] thatbuilds a white-box inference algorithm by matching every step in aprobabilistic program with neural network whose role is to “invert”it.
Foundation posteriors trade explainability for ﬂexibility.
Ourapproach is not white-box: inference is performed by a single blackbox 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 thesetwo methods as optimizing for distinct properties.
Foundation Models
Foundation models [9] are a blanket termfor 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.
Foundationmodels 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” thatcan be quickly ﬁnetuned for downstream inference.
DiscussionLimitations
First, optimizing transformer networks is computationally costly, requiring acceleratedhardware.
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, classesof graphical models like HMMs are difﬁcult to represent as text.
Third, our approach, being relianton deep learning, is limited in its explainability.
In the case that inference fails, we are unable toreason about the root cause.
Finally, the success of foundation models [9] is largely predicated onlarge datasets with millions of entries.
Unfortunately, we do not know of any large collections ofprobabilistic programs other than PosteriorDB, which in comparison, is of modest scale.
Future Work
Future work could investigate foundation posteriors that act as proposal distributionsfor 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 thefoundation 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 tobetter understand the efﬁcacy of foundation posteriors in this context.
Summary
We proposed a meta-amortized inference algorithm for probabilistic programs usingmasked language modeling.
With this algorithm, we build a foundation posterior capable of fastzero-shot inference that can also be ﬁnetuned for more accuracy.
We are optimistic that the generalityof the approach, despite its computational burden, can make for practical and scalable inference.