Foundation Posteriors for Approximate Probabilistic Inference
Authors: Mike Wu
Noah Goodman
Venue: NeurIPS 2022 Accept (2022)
Keywords: approximate inference
masked language modeling
probabilistic program
graphical model

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, defining 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 fine-tuned for a particular program and dataset by optimizing a variational inference objective. We show the efficacy of the approach, zero-shot and fine-tuned, on a benchmark of STAN programs.

Foundation Posteriors for Approximate Probabilistic Inference
Mike Wu, Noah Goodman
Department of Computer Science
Stanford University Stanford, CA 94305
{wumike, ngoodman)@stenn (atoddmmn)@ttafoord..o
Abstract
Probabilistic programs provide an expressive representation language for generative models. Given a probabilistic program, we are interested the task of posterior inference: estimating a latent variable given set of observed variables. Existing techniques for inference in probabilistic programs often require choosing many spperpparameter, are computationally expensive, and/or only work for restricted classes of programs. Here we formulate inference as masked language modeling: given program, we generate supervised dataset of variables and assignments, and randomly mask subset the assignments. We then train a neural network unmask the random values, defining approximate posterior distribution. By optimizing a single neural network across a range of programs we amortize the cost of training, yielding "foundation" posterior able do zero-shot inference for programs. The foundation posterior can also be fine-tuned for particular program and dataset by optimizing a variational inference objective. We show the efficacy of the approach, zero-shot and fine-tuned on 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 efficient for practical usage and universal working for any program that might be written. Many probabilistic programming languages (PPLs) have been developed 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
27 45 28
IS
5% 59
12
23
strong restrictions on what kinds of programs can be represented 112. In this paper, we propose a complementary approach to scale inference for probabilistic programs. By treating program constrained form of language with some additional nomenclature, we show that probabilistic inference can be viewed as masked language modeling 62 [9]. 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 MI.I.
While MII can be used to solve inference queries for single program, it becomes powerful when meta-amortized (or twice-amortized) to do inference across different queries and programs. Success requires a dense set of programs that cover sufficient diversity for the meta-algorithm to generalize. Since this is not always obtainable, we propose program augmentations to enlarge small dataset of programs. The result of meta-amortized MLI the augmented programs is foundation posterior: large neural network trained to do inference across many probabilistic
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

vectors
Language Modeling
Masked Token Example
Shared Weights
Distribution
Score Label
Raw Example
14. 71
Inference Network
Masked Obs Examplle
Fine tuning RoDERTa
Figure 1: Masked language inference: treating a probabilistic program as a raw string, a large language model is trained unmask latent variables conditioned observed ones.
programs. experiments, we find the foundation posterior to be capable of both zero-shot inference and variational fine-tuning: given program from the test set, achieve higher quality using the foundation posterior as an initial distribution. Indeed, fine-tuning the foundation posterior gives the best performance, for wide range of runtimes, for set of standard STAN programs.
2 Background
Probabilistic Programs We assume a class of probabilistic programs without loops and such that every line assigns value to a variable. We note that all finite loops can be "unrolled" into 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 ~ gaussian(z,, 1), direct assignment x = z, and function evaluation x = sqrt(z). Importantly, we do not constrain what distributions and functions allowed.
Approximate Inference Let p(x, 2) be a joint distribution of latent variables z and observed
variables An inference query seeks compute posterior beliefs p(2(2r) p(x,2)
which is
p(z)
usually intractable computing p(x) ((,, 2)d2 faithfully requires solving difficult integral.
So, we settle for approximate techniques. Markov Chain Monte Carlo (MCMC) 12 and
variational inference (VI) 37 64 S are two widely used examples. Focusing the latter, VI
introduces family tractable distributions Q over the latent variables to find the member q" E Q
that minimizes the Bullback-Leibler (KL) divergence between itself and the exact posterior q (z)
arg min, KL (g(=))l((2(z)). Once found, qº serves as proxy for the true posterior.
VI as described above finds q" for single observation 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 So X + Q that predicts (the parameters of the distribution) q² as a function of When scoring a value 2, we often write the expression (z) as conditional distribution (2(2))
To find the function fo, we maximize lower bound of Ep(z) log p(r) using Jensen's inequality to Ep(x))
get Lato
(Ø) (Ø)
E Exs(()=))
p(x.=) where p(z) distribution over the observations. log {0(s)=)
For probabilistic programs, p(z, is easy to sample from (by simply executing the program) but difficult to score. In these cases, we apply useful trick by optimizing KL(pllg) rather than KL(gllp):
Ccomp (0) Ep(x) KL (p(z/z) (z x)) = Ep(x.s) I- (2|x) + constant
(I)

---

Dropping the constant, the remaining expression is close to a supervised objective: choose the parameters the distribution q (z) that maximizes the likelihood of (z, 2) p(x,, z), samples from the probabilistic program.
Equation is called compiled inference A0
Masked Language Modeling A seemingly unrelated technique is masked language modeling (MLM) [9. used widely by large language models like BERT [T] The MLM objective is a word prediction task [56]. Given an input sentence, each token has probability of being replaced by special mask token. The model then predicts the original token given the rest of the sentence. In slightly more notation, let æ = (f1..... .2.....44-1, <mask>, f.+1....,, t.)) be a sequence of tokens that form sentence where the i-th token was randomly masked. The label is then just t,. (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 mlm Ep(x)) [log P(y|2)) where log p(-l-) is cross entropy applied to softmax beliefs over the full vocabulary.
3 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, any subset of the variables and present you with inference queries for the remaining unknown variables. To solve this challenge we aim to find 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, 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 x ~ gaussian (0, 1) + 0.132. 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 truc assignments for all masked variables as a single data point. To generate a dataset, we loop this procedure until we have a sufficiently large corpus. Since the masking decision is made independently for each line, we can theoretically generate every possible permutation of observed and latent variables.
inquirer might hand you observations for
edit the program slightly: For every line (which by design must be variable
expression + value. For instance,
3.1 Objective
The dataset of programs described is not far from 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 procedure See Fig. M for an overview. The objective for MLI is a sum of two loss functions: one for traditional MI. and the other to score the true assignment for latent variable under posterior distribution parameterized by a neural network. In more detail, given a program x = (41,-.. .,,,, we make two versions: for the first Tmim, we do what MLM typically does, choosing random tokens in the program string to mask; for the second, Tinf, we do as described above and randomly mask assignments (the values to the right of the - symbol per line). Note that the two inputs Tmim and Tint have different tokens masked.
call masked language inference, or MLL.
We use transformer network fo : X + Rd that takes a raw string as input where f are trainable parameters. We compute Umim = (Tmmm)) and Vnnt = fo (Eant), both of which are a sequence of contextual vector embeddings = (V1,... Dn). one for each token. For each masked index i, perform two actions, one for each input. For MLM, we have classification head ge Rd [V] (c.g. a linear layer) that maps a vector vmlm [i] to a probability for each token in the vocabulary V. For inference, we have an inference head he Rd > Q that maps a vector Una[ii] [i] to an approximate posterior distribution q(t,/ in the family Q for the latent variable corresponding to token t, For continuous variables, a common choice for Q is the Gaussian family, in which the network hi would return two vectors representing the mean and standard deviation. Alternatively, if the latent variable binary, we might choose the family @ to be Bernoulli where he returns a probability vector. The choice of distribution is flexible long scoring is differentiablee (though sampling need not be).
3

---

Program
Test Set Evaluation
Ablation: No MLM
log p(=lz)
var
log p(x,s) q(aler)
logp(z/z)
var log piz,s)
Latent
-1.53840.1
1.895=1.1
-3.740+0.2
1.059e4291
Clustering
3.400620.4
1.097
-8.0373..1
5.415c325es
Hierarchical
-.2668+a.1
119.2+60
-7.088
5.162e9++2e
Multi-level
-3.359
131.3+s3
-8.363+0.7
3.219e8++2e
Milky way
2.896+0.2
66.09+44
-..619
1.14766+1es
Rosenbrock -..227+0.1
6.673
4..0
4.252e5++25
Table 1: Masked Language Inference on a suite of probabilistic programs. For each program, a test
set is built using 1 000 new executions with randomly masked assignments. We measure the average quality inferring these masked values. We show averages over runs.
In summary, the MLI objective is:
mii (0) (at [log pfmmlm [mlm)) + a logg P(tint Tiot)]
(2)
where is weight balancing the two losses, and mim, liat masked tokens. While Equation only masked a single token per loss, in practice we randomly mask 15% of tokens in Tmlm. and mask an increasing amount of assignments in Zint according to schedule: we begin at 15% but increase this masking probability throughout training to 50%, thereby increasing the difficulty of inference. The MLI objective in Equation 2 can be maximized with stochastic gradient descent.
3.2 Orderless Regr-segeesevve Decoding
After training, given probabilistic program with m latent variables 21,... 2m. there are a number of approaches use the approximate posterior q perform inference. The simplest approach is to unmask all latent variables at once using hg. However, we might believe that inferring the value of 21, for example, provides
information when inferring the value of 22-
Given some ordered queue latent variables e.g. 21, 22, 2m we can loop m times, each time embedding the program using the transformer fe and unmasking the next latent variable in the queue 2, using the inference head hg. Then, we remove 2; from the queue and repeat for 2;-1. Each iteration of the loop produces embedding through fe. 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 21 to Am-
observation, thus changing the program tokens and the resulting
Throughout this auto-regressive decoding, the proportion of latent to observed variables will shrink over time. the first iteration, there will be the maximum number of latent variables whereas the last iteration, there will only be one. As such, properly trained MLI model must handle different numbers latent variables, not just different variables being observed. Section 3.1 we described how the percentage of masked assignments is varied in training. A benefit 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 efficacy, we consider six probabilistic programs used in prior work 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 specified in the program. We recommend the reader refer to Section A.3 for details. While these programs are simple, they specifyy wide prior distributions that result in high sample diversity over repeated executions, which should make the MLI task more challenging.
13..
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 specified by the model, averaged over all masked tokens per execution and all test exccutions. The larger this value is, the better the inference.

---

Second, we compute the variance of the log importance weights (IW), as in 668 Since the approx- imate posterior q acts as an importance distribution estimating the true posterior, we measure the quality by the variance of p(z.2) importance weight would be fixed constant, meaning zero variance.
where lower variance more desirable. If ((2(x)) = p( (z), the
g(=[z))
Results Table 4 reports the performance the six programs. We observe large probabilities (log close to and small variance of log IW, which suggests good generalization of inference new sample values. As baseline, in- clude an ablation to MLI by removing the MLM term, which amounts training the inference head only. We find smaller log probabilities (every log point is significant difference) and much larger variance. Together, this suggests that the MLM term important for generaliza- tion, likely as it helps the neural network un- derstand program structure. For analysis on the distribution of variances, see Section A.3
Y1
y2
(b)
(c)
(d)
Figure 2: Graphical model representations for the 'Independent Gaussians' (a). 'Conditional Indepen- dence' (b), and 'Common Effect' (c,d) programs.
3.4 Visualizing Attention Maps
Prior works works in natural language have re-purposed the attention weights in the later layers as mech- anism to introspect model logic. In a similar vein, we leverage attention weights to approximate dependency graph between random variables in the problem.
6 63
using transformer net-
Datasets We study simple probabilistic programs that exhibit independence and conditional independence between random variables. First, in Figure 2]),, 21 and 1/1 are independent Gaussian variables. We should expect independence between the sets {y1, U2) and (=1,22) Second, in Figure bb). design graphical model with Bernoulli random variable z, and pick and b using if statements to be conditionally independent given z. We add two Gaussian variables and y that independently add noise around a and b, i.c. x N(a, 1). Finally, Figure ccdd) create common effect model with two potential causes: choose and to be independent Bernoulli variables, and be if or y else We expect p(x12 2 y) differ meaningfully from p(x122).
Morked Weathe
94
24
Observed Vonabee xxby
12
03
40
51
Input Vanable
Input Vanasto
Input Vartable
(a)
(b)
(c)
Figure 3: Heatmaps of attention norms for 'Independent Gaussians' 'Conditional Independence" and "Common Effect' programs. On the axis, we show which variable has its assignment masked. the x axis, we list all program variables. The brighter the color, the larger the weight.
Results We report metries similar to Table 4 for these programs in Section A.4 Here, instead visualize attention: suppose a single assignment in 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, treat the average weight for all tokens in that line 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 the graphical dependency between the two variables.
Figure shows the heatmap of attention weights for the three programs. For subplots (a) and (b), we mask each variable's assignment one at time. The y-axis shows which variable is masked. In the

---

"Independent Gaussians' program, the model approximates the independence between {21, 22) and {y1, V2). For instance, in the first row of (a), when unmasking y1, we see the network upweights the tokens corresponding to Y1 and y2, but downweights the tokens for 21 and 22. The other rows show a consistent pattern. While may seem peculiar at first for the model to pay attention to y1 when unmasking V1, we note that this is important for the model to understand what kind of variable is. Similarly, in the 'Conditional Independence" program, the model approximates the independence between {a, 2 and {b, y given z. Looking at the first row, see that unmasking focuses variables æ and and downweights b and y. Finally, in subplot (c), the y-axis shows which variables observed rather than which are masked, for clarity. The bottom row plots attention weights if we observe both 2 and y whereas in the top row, we only observe 2. We see a stark contrast: in the bottom, the weights roughly uniform: on top, the model gives little weight to y unmask
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 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
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 dataset for probabilistic programs is unfeasible. Instead, we propose set of program augmentations, inspired by data augmentations, 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:
From prior work 30 16
[3],
enlarge small set of programs. The hope is that the larger
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 & = rosenbrock( (a, b) might be augmented to 2 -N b).
Fuzz Constant Constants are replaced by newly sampled constants from known prior distributions. For example, r ~ (0, 1) becomes r ~ 1, 0.9).
Line Swaps Swaps two independent program lines such that the dependency graph between random variables unchanged. For example, program with three lines u (O, 1); N11,, 1);r utv could be augmented to v ~ (1, 1); u ~ (0, 1);r = u + but it would not be possible to swap the third line with any other given the dependency of and on o.
Cut and Glue Replace the usage of a variable with another variable already defined the program. Note that this augmentation changes the dependency graph. For example, u a + b; sqrt(u) could become u = a - b; u = sqrt(b) where is defined earlier in the program.
Create and Use Creates new random variable and uses it in lieu of another variable or constant in the rinh side of an expression. Note that this augmentation introduces a new random variable. For example, u + 6; sqrt(u) could become
+ b (N, 1); 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) + NN1, -1) by fuzz function is ill-defined. 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 five of them for meta-training, and hold out the "Rosenbrock" program for meta-test, We apply up to five random program augmentations each program in both sets to create the training and test splits. Note that the test inputs are all derived from novel program unseen by the model in training.

---

This is much more difficult generalization problem than prior experiments. As 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 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.
5 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 be challenging. In this case, propose finetune our pretrained inference model using stochastic variational inference, or SVI [35] 50 39 Due to reparameterization requirements, in this subsection, we focus inference of continuous random variables only. For discrete random variables. future work can study using the concrete relaxation REINFORCE- based techniques
42
or the family of
More specifically, fix a test program x' that we would like to do more high quality inference for. Suppose x² has observed data points d1,...,dy dn,, and latent variables want to estimate 21,... 2m Further, we given Ó, the parameters obtained from optimizing MLJ on pretraining dataset of programs. Since the composed functions f h X - Q map probabilistic program to an approximate posterior, we define the shorthand qá (31mm d1m)) h(f(x*))
We can optimize the evidence lower bound, which we recall is:
0" arg max Celbo (0) = arg max E Es (aromlg log p(di:n 31:m) DEG
DE
qo (Z) d1
Since we are not amortizing over programs in this finetuning step, Equation does not contain second expectation over programs. That is, Zi:m, dia are the specific variables from xt. We initialize 0 = 0 to leverage pretrained weights. Note p(...) has trainable parameters but acts a likelihood scaling term for different values of 21:m and thus, cannot be dropped.
Results In Table 2 we include an ablation comparing the generalization of inference with and with- out fine-tuning. We fine-tune the inference network (from MLI pretraining), optimizing Equation S separately for each program for 1,000 iterations. As baseline, we also finetune the network from random initialization. We observe that fine-tuning MLI further improves log probability by roughly log points with comparable IW variance. On the other hand, fine-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 plating example. In the left program, we see 10 observations di...., d10 such that d, gaussian(r; 1) where is outside of the plate. But suppose 10 examples too
Model
Test Set Evaluation
log (oga(())))
var log p(z,a) ([]l))
MLI
-60 14.5
15.10 19.6
MLI- - Augmentations
-246.6100
355.9 137
MLI Finetuning
-5.727 #4.0
16.06 +7.2
Random + Finetuning
-42.54 +89
781.9 +:42
Table 2: Meta-Amortized Masked Language Inference over suite of probabílistic programs. A test
set of 1,000 programs are constructed using random augmentations on held-out program.
7

---

many to unroll. In the program on the right, we replace the 'for' loop with a plate (n) token where the argument specifies the number of total iterations. Within the plate, we unroll two randomly subsampled iterations i E (1, 4). For multiple executions of this program, the iterations within neural network training. From the perspective of the
the plate change, just like minihatching transformer, this change does nothing more than add new token to the vocabulary.
ots IPt6 Ut.
Figure 4: Example of plating within MLI. There is a special plate token that specifies the total
number of observations, This defines scope with minibatch of k < observations.
Why does this work? Consider computing the program density p(a,, d1,-.., d.). This be rewritten as p(r)p(d1,.-. dn(z)) of which the first term is a given. To compute the second term, observe that di...., dn are conditionally independent on æ. So, p(din/<) -II 2().)2)) Now:
logp(dim/æ)) = log logp(d
Σ
logp(d,(x)
Emminibacch
where the minibatch is of size k < n. We make unbiased (but higher variance) estimate of the full conditional distribution (and hence, the program density) using small minibatch.
In the Appendix (Section A.8 we demonstrate plating item response theory (IRT) [21 3ll 51
66 67 model, popular probabilistic program. In the next section, we will leverage plating use MLI on bank of real world probabilistic programs with thousands of observations each.
Foundation Posterior
The relationship between MLI and variational fine-tuning is reminiscent of pretraining and down- stream tasks, as popularized by self-supervised learning [19 [5]. Inspired by foundation models (9) we propose frame the result of MLI as foundation posterior that can be finetuned downstream using SVI to perform inference for wide array of probabilistic programs. The goal is to pay large one-time cost in training "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
benchmark dataset for evaluating inference algorithms 4 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 is used for all programs to fit observations within the transformer's 512 token limit. After pretraining, we optimize Equation for each test program individually, varying the number of steps of fine-tuning across 0 (zero-shot), 10, 100, and 1000. As baselines, we use CudStanPy 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 I M. Each program comes with 10k pre-computed posterior samples fit 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. the p-value high, 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 com- putation (wall-clock time) on the x-axis in log-seconds. For MLI, we sum the cost of fine-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. better inference algorithm would bias towards the top left corner, achieving higher p-vulue at lower

---

13-
02-
x2-
log secends
seconds
log seccnnds
(a) Earnings
(b) KidIQ
(c) NES1976
Figure 5: Comparisons of MIJ to native inference methods in the Stan programming language, such NUTS ADVI, For any inference algorithm, draw samples from the posterior and compute similarity to ground-truth posterior samples through a Kolmogorov-Smie -omirov test.
cost. Each point in the plot represents a setting we redo inference with different finetuning steps or mixing steps, etc. Each line groups together 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 a p-value of 1 in all cases, 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 MI.I curve represents zero-shot inference, which we find comparable to ADVI despite the latter training for up to 1M steps. Finally, the gray line shows additional baseline where we finetune a transformer from scratch. We observe that initializing from a foundation posterior is critical good inference.
7 Related Work
Amortized Inference Traditionally, an inference query ((2) is solved for single assignment of the observed variable x. However, in many applications, we may be interested in solving the same inference query for many observations 21, .2,,..., In. Re-solving the same inference query from scratch for all n related queries seems wasteful. Amortized inference 2 551 was proposed as more efficient alternative by learning function (==(²r) that maps a observed assignment x to distribution over the latent variable z. In doing so, we "amortize" the cost of inference with large one-time cost in defining q, after which an inference query can solved with single function application. Without amortization, the foundation posterior would not be very efficient inference algorithm.
Meta-Inference Meta-learning has been used do inference amortized over both queries and a family of generative models [[ These works have shown that amottiamortized" 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 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 fine-tuning.
3C 36
3-1tat
Most relevant to our work is meta-learning approach [3 that builds white-box inference algorithm by matching every step in probabilistic program with neural network whose role is to "invert" it. Foundation posteriors trade explainability for flexibility. Our approach is not white-box: inference is performed by single black box neural network. In return, our approach assumes no structure: it treats the program nothing more than string, making its gener- alization capacity greater than the approach [K]] We view these two methods optimizing for distinct properties.
Figure 6: Distribution of R-hat scores over 100 executions of the 'Latent' program when us- ing the approximate posterior (red) versus the prior (orange). Averages are the dotted lines.

---

Foundation Models Foundation models [ are a blanket term for large unsupervised models that map data to representations. Examples include ResNet [33 in computer vision, BERT [9] and GPT [10, -9 in natural language, Wav2Vec [54 R in speech, and CLIP [48] and Data2 Vec [2] in multimodal applications. Foundation models are treated as frozen backbones whereupon small portion is finetuned for downstream task. Our motivation for the foundation posterior is precisely
this: learn a good posterior "initialization" that can be quickly finetuned for downstream inference.
8 Discussion
First, optimizing transformer networks is computationally costly, requiring accelerated hardware. For practitioners simply interested in small set of inference queries, it is simpler to use MCMC or VI. Second, MII as described is not sufficient for programs: PosteriorDB, classes of graphical models like HMMs are difficult to represent as text. Third, 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 O 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, of modest scale.
Limitations
Future Work Future work could investigate foundation posteriors that act as proposal distributions for MCMC-based methods. As a first step, we run Metropolis Hastings one of the six toy programs, and compute R-hat 25 We achieve a score of 1.011 (#00.099 over 100 test programs) when using the foundation posterior as the proposal compared to 1.043 (#00059) when using the prior distribution. Best practice deems a chain sufficiently mixed if R- R-hat is less than 1.05. More analysis is needed to better understand the efficacy of foundation posteriors in this context.
Summary We proposed a meta- aettt-aeottized inference algorithm for probabilistic programs using masked language modeling. With this algorithm, we build foundation posterior capable of fast zero-shot inference that can also finetuned for more accuracy. We are optimistic that the generality of the approach, despite its computational burden, can make for practical and scalable inference.
Acknowledgments and Disclosure of Funding
This work was partly supported by a NSF Expeditions Grant, Award Number (FAIN) 1918771. MW supported by the Stanford Interdisciplinary Graduate Fellowship as the Karr family fellow.
References
[1] Samira Abnar and Willem Zuidema. Quantifying attention flow in transformers. arXiv preprint arXiv:2005.00928. 2020.
[2] Alexei Baevski, Wei-Ning Hsu, Qiantong Xu, Arun Babu, Jiatao Gu, and Michael Auli. Data2vec: A general framework for self-supervised learning in speech, vision and language. arXiv preprint 1.Xi::2202.035555 2022.
[3] Alexci Baevski, Yuhao Zhou, Abdelrahman Mohamed, and Michael Auli. wav2vec 2.0:
A framework for self-supervised learning of speech representations. Advances in Neural Information Processing Systems, 33: 12449-1222, 2020.
[4] Guillaume Baudart, Javier Burroni, Martin Hirzel, Louis Mandel, and Avraham Shinnar. Compil- ing stan to generative probabilistic languages and extension to deep probabilistic programming. In Proceedings of the 42nd ACM SIGPLAN International Conference on Programming Lan- guage Design and Implementation, pages 497-510 2021.
[5] Guillaume Baudart and Louis Mandel. Automatic guide generation for stan via numpyro. arXiv preprint 2 10. 1790, 2021.
[6] Alexander Binder, Grégoire Montavon, Sebastian Lapuschkin, Klaus-Rober Müller, and Woj- ciech Samek. Layer-wise relevance propagation for neural networks with local renormalization layers. In International Conference on Artificial Neural Networks, pages 63-71. Springer, 2016.
[7] 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
10

---

universal probabilistic programming. The Journal of Machine Learning Research, 20(1):973- 978. 2019.
[8] David M Blei, Alp Kucukelbir, and Jon D McAuliffe. Variational inference: A review for statisticians. Journal of the American statistical Association, 112(518):59--87,, 2017.
[9] Rishi Bommasani, Drew A Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von
Arx. Michael S Bernstein, Jeannette Bohg, Antoine Bosselut, Emma Brunskill, et al. On the opportunities and risks of foundation models. arXiv preprint arXiv:2108.077 07258, 2021.
[10] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
[11] Yanshuai Cao and David J Fleet. Generalized product of experts for automatic and principled fusion of gaussian process predictions. arXiv preprint arXiv: /4/0///00000 7827, 2014.
[12] Bob Carpenter, Andrew Gelman, Matthew D Hoffman, Daniel Lee, Ben Goodrich, Michael Betancourt, Marcus Brubaker, Jiqiang Guo, Peter Li, and Allen Riddell. Stan: A probabilistic programming language. Journal of statistical software, 76(1), 2017.
[13] Gwonsoo Che and Hongseok Yang. Meta-learning an inference algorithm for probabilistic programs. arXiv preprint (Xxx:221030007:: 2021.
[14] Hila Chefer, Shir Gur, and Lior Wolf. Transformer interpretability beyond attention visualization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 782-791, 2021.
[15] Ting Chen, Simon Korbllith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. International conference on machine learning, pages 1597-1607. PMLR, 2020.
[16] Kristy Choi, Mike Wu, Noah Goodman, and Stefano Ermon. Meta-amortized variational inference and learning. International Conference on Learning Representation. 2019.
[17] Marco F Cusumano-
Feras A Saad, Alexander K Lew, and Vikash K Mansinghka. Gen: a general-purpose probabilistic programming system with programmable inference. In Proceed- ings of the 40th acm sigplan conference on programming language design and implementation, pages 221-236, 2019.
[18] Sarah Depaoli, James P Clifton, and Patrice R Cobb. Just another gibbs sampler (jgg)) flexible software for meme implementation. Journal of Educational and Behavioral Statistics, 41(6):628- 649, 2016.
[19] Jacob Devlin, Ming-Wai Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv: (0t::::8000444 048005 2018.
[20] Oliver Dürr, Stephan Hörling, Daniel Dold, Ivonne Kovylov, and Beate Sick. Bernstein flows for flexible posteriors in variational bayes. arXiv preprint arXiv:2202.056500 2022.
[21] Francis Ysidro Edgeworth. The statistics of examinations. Journal of the Royal Statistical Society, 51(3):599-6355 1888.
[22] Zhangyin Feng, Daya Guo, Duyu Tang, Nan Duan, Xiaocheng Feng, Ming Gong, Linjun Shou, Bing Qin, Ting Liu, Daxin Jiang, et al. Codebert: A pre-trained model for programming and natural languages. arXiv preprint XXXi:2002.088555 2020.
[23] Hong Ge, Kai Xu, and Zoubin Ghahramani. Turing: Composable inference for probabilistic programming. In AISTATS, volume 84, pages 1682-1690. PMLR, 2018.
[24] Alan E Gelfand and Adrian FM Smith. Sampling-based approaches to calculating marginal densities. Journal of the American statistical association, 85(410):398-4M)) 1990.
[25] Andrew Gelman and Donald B Rubin. Inference from iterative simulation using multiple sequences. Statistical science, 7(4):477-472, 1992.
[26] Samuel Gershman and Noah Goodman. Amortized inference in probabilistic reasoning. In Proceedings of the annual meeting of the cognitive science society, volume 36, 2014.
[27] Noah Goodman, Vikash Mansinghka, Daniel M Roy, Keith Bonawitz, and Joshua B Tenenbaum. Church: language for generative models. arXiv preprint arXiv:1206.3255 2012.
11

---

[28]
Noah D Goodman and Andreas Stuhlmüller. The Design and Implementation of Probabilistic Programming Languages. http::/ 2014. Accessed: 2022-3-19.
[29] Andrew Gordon. Thomas Henzinger, Aditya Nori, and Sriram K Rajamani. Probabilistic programming. In Future of Software Engineering Proceedings, pages 167-181. 2014.
[30] Jonathan Gordon, John Bronskill, Matthias Bauer, Sebastian Nowozin, and Richard E Turner. Meta-learning probabilistic inference for prediction. arXiv preprint arXiv: XXVv:1805.00921, 2018.
[31] Ronald K Hambleton, Hariharan Swaminathan, and Jane Rogers. Fundamentals of item response theory, volume 2. Sage, 1991.
[32] W Keith Hastings. Monte carlo sampling methods using markov chains and their applications.
1970.
[33] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770-778, 2016.
[34] Geoffrey Hinton. Training products of experts by minimizing contrastive divergence. Neural computation, 14(8):1771-1800) 2002.
[35] Matthew D Hoffman, David M Blei, Chong Wang, and John Paisley. Stochastic variational inference. Journal of Machine Learning Research, 2013.
[36] Ekaterina lakovleva, Jakob Verbeek, and Karteek Alahari. Meta-learning with shared amortized variational inference. In International Conference on Machine Learning, pages 4572-4582.
PMLR, 2020.
[37] Michael I Jordan, Zoubin Ghahramani, Tommi S Jaakkola, and Lawrence K Saul. An in- troduction to variational methods for graphical models. Machine learning, 37(2)3883-233, 1999.
[38] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint a1Xiv: 2014.
[39] Alp Kucukelbir, Dustin Tran, Rajesh Ranganath, Andrew Gelman, and David M Blei. Automatic differentiation variational inference. Journal of machine learning research, 2017.
[40] Tuan Anh Le, Atilim Gunes Baydin, and Frank Wood. Inference compilation and universal probabilistic programming. In Artificial Intelligence and Statistics, pages 1338-1348. PMLR, 2017.
[41] Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint 1907 2019.
[42] Chris J Maddison, Andriy Mnih, and Yee Whye Teh. The concrete distribution: A continuous relaxation of discrete random variables. arXiv preprint arXiv: 1,001110012222111 12, 2016.
[43] Mans Magnusson, Paul Burkner, and Aki Vehtari. posteriordb: a set of posteriors for bayesian inference and probabilistic programming. 2021.
[44] Vikash Mansinghka, Daniel Selsam, and Yura Perov. Venture: a higher-order probabilistic programming platform with programmable inference. arXiv preprint arXiv:1404.0099 2014.
[45] Tom Minka. Infer. net 2.5. http://research. microxoft. com/infernet, 2012.
[46]
Andriy Mnih and Karol Gregor. Neural variational inference and learning in belief networks. In International Conference on Machine Learning, pages 1791-1799, PMLR, 2014.
[47] Praveen Narayanan, Jacques Carette, Wren Romano, Chung-chieh Shan, and Robert Zinkov. Probabilistic inference by program transformation hakaru (system description). In Interna- lional Symposium on Functional and Logic Programming, pages 62-79. Springer, 2016.
[48] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International Conference on Machine Learning, pages 9748-8763. PMLR, 2021.
[49] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 188)99 2019.
12

---

[50] Rajesh Ranganath, Sean Gerrish, and David Blei. Black box variational inference. In Artificial intelligence and statistics, pages 814-822. PMLR, 2014.
[51] Georg Rasch. Probabilistic models for some intelligence and allainment lests. ERIC, 1993.
[52] Daniel Ritchie, Ben Mildenhall, Noah D Goodman, and Pat Hanrahan. Controlling procedural modeling programs with stochastically-ordeess sequential monte carlo. ACM Transactions on Graphics (TOG), 34(4):1-11, 2015.
[53] John Salvatier, Thomas V Wiecki, and Christopher Fonnesbeck, Probabilistic programming in python using pymc3. PeerJ Computer Science, 2:e55, 2016.
[54] Steffen Schneider, Alexei Baevski, Ronan Collobert, and Michael Auli. wav2vee: Unsupervised pre-training for speech recognition. arXiv preprini arXiv: 1:XXvv/904.00 /e01/99444..22. 2019.
[55] Andreas Stuhlmüller, Jacob Taylor, and Noah Goodman. Learning stochastic inverses. Advances in neural information processing systems, 26, 2013.
[56] Wilson L Taylor. "cloze procedure": A new tool for measuring readability. Journalism quarterly. 30(4):415-433, 1953.
[57] Nazanin Tehrani, Nimar S Arora, Yucen Lily Li, Kinjal Divesh Shah, David Noursi, Michael Tingley, Narjes Torabi, Eric Lippert, Erik Meijer, et al. Bean machine: A declarative probabilistic programming language for efficient programmable inference. In International Conference on Probabilistic Graphical Models. PMLR, 2020.
[58] David Tolpin, Mai-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,
pages 1-12, 2016.
[59] Dustin Tran, Alp Kucukelbir, Adji B Dieng, Maja Rudolph, Dawen Liang, and David M Blei. Edward: A library for probabilistic modeling, inference, and criticism. arXiv preprint arXiv:1610.09787, 2016.
[60] George Tucker, Andriy Mnih, Chris J Maddison, John Lawson, and Jascha Sohl-Dickstein. Rebar: Low-variance, unbiased gradient estimates for discrete latent variable models. Advances in Neural Information Processing Systems, 30, 2017.
[61] Jan- Willem van de Meent, Brooks Paige, Hongseok Yang, and Frank Wood. An introduction to probabilistic programming. arXiv preprint 100.iv:8009.11. 10756, 2018
[62] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[63] Elena Voita, David Talbot, Fedor Moiseev, Rico Sennrich, and Ivan Titov. Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned. arXiv preprint 2019.
[64] Martin J Wainwright, Michael Jordan, et al. Graphical models, exponential families, and variational inference. Foundations and Trends@ in Machine Learning, 1(1-2):1-305, 2008
[65] Ronald J Williams. Simple statistical following algorithms for connectionist reinforce- ment learning. Machine learning, 8(3):229-2)) 222,, 1992.
[66] Mike Wu, Richard L. Davis, Benjamin W Domingue, Chris Piech, and Noah Goodman. Varia- tional item response theory: Fast, accurate, and expressive. arXiv preprint arXiv:2002.00276, 2020.
[67] Mike Wu, Richard L Davis, Benjamin W Domingue, Chris Piech, and Noah Goodman. Modeling item response theory with stochastic variational inference. arXiv preprint 1rXiv:2108.11577 2021.
[68] Mike Wu and Noah Goodman. Multimodal generative models for scalable learning. Advances in Neural Information Processing Systems, 31, 2018.
[69] Yuling Yao, Aki Vehtari, and Andrew Gelman. Stacking for non-mixing bayesian computations: The curse and blessing of multimodal posteriors. arXiv preprint arXiv:2006. 12335, 2020.
[70] Lu Zhang, Bob Carpenter, Andrew Gelman, and Aki Vehtari. Pathfinder: Parallel quasi-newton variational inference. arXiv preprint 10X.:220080038... 2021.
13


Reviews:
[{"review_id": "rjxc-fr5Cwm", "reviewer": "Reviewer_mMgG", "summary": "This paper presents a method for learning a language model that can perform approximate inference in probabilistic programs. Specifically, given a set of probabilistic programs, a dataset of masked, annotated programs is generated and then trained on using a linear combination of a masked language modeling loss and a variational inference loss. An extra head on the language model outputs the parameters of a variational distribution for the current token. This variational distribution can then be used as an approximation of the posterior distribution, and can be further fine-tuned on individual programs to achieve even better performance. Experiments on both toy programs and a real benchmark demonstrate that the learned model provides reasonable posterior probabilities.", "questions": "1. How exactly in equation (2) do you get $p(t_{mlm}|x_{mlm})$ and $p(t_{inf}|x_{inf})$ from $g_\\theta$ and $h_\\theta$, respectively?\n2. How does the architecture deal with having multiple types of variable in the program? Are there separate heads $h_\\theta$ for each type of variable, given that the parameterization of q(t_i) changes depending on its type? \n3. In Table 1, how does the performance change as a function of \\alpha? Is the method quite sensitive to this hyperparameter? How do the scores change as \\alpha gets smaller and smaller? (i.e., what happens as the MLM loss becomes more dominant, does the method still perform well or does performance fall off as quickly as when the MLM loss is removed?)\n4. There’s evidence in the literature that pre-training these architectures on *anything* can be quite beneficial (e.g., Reid et al, 2022). How does this method perform if you pre-train with only the MLM loss on either (a) the programs or (b) some other text data, and then fine-tune on the posterior inference problem?\n5. In Section 5, is $p(d_{1:n}, z_{1:m}) difficult to compute, as alluded to in Section 2? How do you compute it here?\n6. What happens if you run ADVI for longer in Figure 5? \n7. What representation is ADVI using and would it do better with more parameters? \n8. How many parameters are used in each of the models in the experiments?\n\n[Reid et al, 2022] Can Wikipedia Help Offline Reinforcement Learning? Machel Reid, Yutaro Yamada, and Shixiang Shane Gu. https://arxiv.org/abs/2201.12122 \n\n", "limitations": "The authors wrote a very nice discussion of the limitations of the proposed approach, which I much appreciated. No potential negative societal impacts were mentioned but it’s not clear to me that there are any immediate such impacts with this work.\n", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Originality:** To my knowledge, this work is quite original. While language models have been applied to any number of tasks, it seems there are still more places where they can be applied successfully. In this case to posterior inference for probabilistic programs.\n\n**Quality and clarity:** The quality and clarity of the paper are both high. While some additional information should be provided about the inference procedure (see questions), the overall exposition is quite clear, and the contribution and empirical validation is of high quality.\n\n**Significance:** This paper presents an interesting and novel method that I think will be of some interest to the community both in regards to the ability of a language model to learn meaningful representations for posterior inference and a new method for posterior inference that incorporates ideas from modern machine learning (i.e., transformers over large amounts of data). More work is still needed to become a method used for the latter but this paper solidly demonstrates that that work may be possible.\n\nOverall, this paper was well-written and interesting to read. While the approach presented here is a nice step, I wonder if using the proposed method as a prior or proposal for a method with more guarantees would be more useful in general. I see that this is mentioned in the conclusion and Figure 6 and I’d encourage the authors to further explore such directions. \n\nMy main complaint with the paper is that I found Section 3.3 to be somewhat poorly executed and explained. Why look at average L2 norms instead of attention weights from the transformer itself? It’s hard to assign much meaning to the particular values of these norms. The writing and explanation could use some editing and rethinking as well to make the results more clear and compelling.\n\nMinor: I think the term “foundation prior” may be more appropriate, particularly if you frame  fine-tuning as an inference procedure that incorporates data into a prior to produce a posterior.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "Z0INWLLLDXl", "reviewer": "Reviewer_TwoP", "summary": "Proposes to use BERT style masked language modeling to amortize inference across assignments of variables in probabilistic programs, particularly STAN programs. This learning of an inference network (a across programs is a form of meta-amortization, implemented as compiled inference (optimizing the q(z|x) distribution that maximizes (x,z) ~ p(x,z), using KL(p || q) instead of qp). It allows efficient inference of complex programs, finding high likelihood assignments to variables with low variance in the importance weights. Combined the masked language modeling objective (of random tokens) with masking out just some entries. Attempts to study the learned attention weights to see if the model is learning meaningful causal structure/relations. Performs somewhere between NUTS (best) and ADVI (cheaper) on a dataset drawn from PosterioDB. MLI (masked language inference) does as well as ADVI with a single sample, and improves given more time.", "questions": "1. Look into XLNet and orderless NADE. Can decompose the variables in alternate orders and support more flexible decoding.\n\t2. Wonder how constrained the meta-distribution needs to be to make the problem solvable (see discusssion @ 190)\n\t3. Consider making the lines in figures have patterns, not just colors (figure 5)\n\t4. minor edits:\n\t\t1. 84: space between \"programitself\"\n\t\t2. Figure A1 is very hard to read on paper\n\t\t3. 685: distirbutions -> distributions", "limitations": "No significant negative societal impact", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "\n\t- originality\n\t\t- This is a fairly original idea, combining 2 powerful ideas to more effectively amortize and speed up inference in flexible probabilistic models (BERT for PPLs). It is the first paper I am aware of that uses BERT and transformers for inference in probabilistic programs/graphical models. Similar ideas appear in other instances of compiled inference, including applications of LSTMs to learn samplers for sites in probabilistic programs (see Tuan Anh Le thesis and ETAMULIS).\n\t- quality\n\t\t- Well executed, using an established ppl (stan) to evaluate performance. Reports both the likelihood of the sampled values and the variance in importance weights, a helpful diagnostic and indicator of the quality of the approximate posterior. Impressive results compared to standard baselines\n\t- clarity\n\t\t- Well written. Illustrates main points, provides compact background on the motivation for probabilistic programs and approximate inference using compiled inference networks. Believe the results are straightforward to reproduce. Experiments are targeted and make a convincing case for the components of the method and in comparison to standard inference methods built into STAN.\n\t- significance\n\t\t- Suggests a way to use the rapid advances in large language models to aid efficient inference in existing probabilistic programming languages. This is likely to spark significant followup work to learn meta-amortized inference models which are capable of outperforming existing optimizers across a wide range of applications.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "YVjLLcvfl9t", "reviewer": "Reviewer_qzPM", "summary": "This paper has formulated probabilistic inference as masked language modeling. In particular, it proposed a meta-amortized inference algorithm for probabilistic programs where a foundation posterior is constructed and can be ﬁnetuned downstream to perform inference over various programs. The authors demonstrated the efficacy of such an approach on a benchmark of STAN programs.", "questions": "- expressivity of the supported programs. \nIf I understand correctly, the proposed approach can be applied to any probabilistic program written in STAN. Is the if-else statement allowed in your probabilistic programs? In other words, whether or what kind of conditional statements can be supported by your approach? It is appreciated if the author can clearly define the kind of probabilistic programs they can support. \n\n- experiments. \nI think the experiment sections throughout the paper are slightly weaker than other parts of the paper. It is a little bit difficult for readers with little background knowledge of STAN and the six programs, in particular, to fully appreciate the scope of the programs the paper is targeting. I would suggest enhancing related discussion both in the main paper and in the supplement. \n\n- computation cost. \nThough mentioned by the author that their key motivation for the foundation posterior is precisely this: learn a good posterior “initialization” that can be quickly ﬁnetuned for downstream inference. However, the experiments are mostly done on not too difficult programs. I am wondering what kind of probabilistic problems that do really worth the computational cost of optimizing transformer networks and are struggling to be solved by common PPSs such as STAN, pymc3, or Pyro. ", "limitations": "N/A", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Originality/Significance:** \nGood: It is interesting to see that the authors interpret probabilistic inference from a different angle and bridge ideas from unsupervised learning (language modeling in particular). As discussed in the related work section, I believe this idea is novel and also has solid foundations, which opens up interesting research directions.  \n\n**Quality:**        \nGood: The paper appears to be technically sound. \n\n\n**Clarity:**          \nGood: The paper is well organized and clearly written. Adequate intuitions, as well as details of the proposed methods, are both introduced and it is enjoying reading the paper. \n\n**Weaknesses:**        \nOne main concern of the paper is its computation burden. I will explain in detail in the following section. ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "Mf1ElYmN9TQ", "reviewer": "Reviewer_KCop", "summary": "  This paper use meta-amortised learning to learn a neural network which can approximate not just a\n  family of posterior distributions but also one that is generalised across different probabilistic\n  programs. This is done through using compiled inference to generate a corpus of probabilistic programs\n  which are then used to fine-tune a transformer model. The language model is trained to predict different\n  posterior distributions based on different observations.\n", "questions": "How does this work differ from Che and Yang? What's the contribution of this work in particular?\n\nWhat makes this work distinct from compiled inference? What approximate inference tasks can this work do that that work struggles with?", "limitations": "I feel the authors have adequately addressed the limitations of the work and potential negative social impacts of this work.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "  Originality:\n\n  The paper seems very similar to the work of Gwonsoo Che and Hongseok Yang and I'm struggling to\n  distinguish them. The main difference seems to be the use of large transformer models instead\n  of something that follows the structure of a particular probabilistic program.\n\n  Technical Quality:\n\n  The work mostly focuses on toy problems and the results are promising for them. I have a concern that\n  a massive language model is overkill for handling the kinds of toy problems used in this paper. The experiments\n  involving Stan and PosterioDB are very compelling though it is a shame that it doesn't out-perform NUTS. I would have\n  also liked to see this work compared to the work of Che and Yang from 2021. I'm also concerned that there\n  doesn't seem to be any benefit in using an augmented dataset from one toy problem to do inference in\n  another toy problem. Without achieving that kind of generalisation, what makes this method better than\n  just using compiled inference in each toy domain?\n  \n  Clarity:\n\n  The paper is well-written and easy to follow. I had no difficulty understanding the method though\n  have not tried to replicate the experiments in the paper.\n  \n  Significance:\n\n  I think this work as tremendous potential as being able to do amortised inference over a set of\n  related probabilistic programs would be very useful. If this can shown to generalise across several\n  domains even related ones it would make the work much more compelling.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept