Venue: NeurIPS 2022 Conference

Reviews:

final_score: rating: 8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.
review_content: 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.
strengths_and_weaknesses: 
	- originality
		- 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).
	- quality
		- 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
	- clarity
		- 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.
	- significance
		- 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.
questions: 1. Look into XLNet and orderless NADE. Can decompose the variables in alternate orders and support more flexible decoding.
	2. Wonder how constrained the meta-distribution needs to be to make the problem solvable (see discusssion @ 190)
	3. Consider making the lines in figures have patterns, not just colors (figure 5)
	4. minor edits:
		1. 84: space between "programitself"
		2. Figure A1 is very hard to read on paper
		3. 685: distirbutions -> distributions
limitations: No significant negative societal impact
review_title: null

final_score: rating: 6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
review_content: summary:   This paper use meta-amortised learning to learn a neural network which can approximate not just a
  family of posterior distributions but also one that is generalised across different probabilistic
  programs. This is done through using compiled inference to generate a corpus of probabilistic programs
  which are then used to fine-tune a transformer model. The language model is trained to predict different
  posterior distributions based on different observations.

strengths_and_weaknesses:   Originality:

  The paper seems very similar to the work of Gwonsoo Che and Hongseok Yang and I'm struggling to
  distinguish them. The main difference seems to be the use of large transformer models instead
  of something that follows the structure of a particular probabilistic program.

  Technical Quality:

  The work mostly focuses on toy problems and the results are promising for them. I have a concern that
  a massive language model is overkill for handling the kinds of toy problems used in this paper. The experiments
  involving Stan and PosterioDB are very compelling though it is a shame that it doesn't out-perform NUTS. I would have
  also liked to see this work compared to the work of Che and Yang from 2021. I'm also concerned that there
  doesn't seem to be any benefit in using an augmented dataset from one toy problem to do inference in
  another toy problem. Without achieving that kind of generalisation, what makes this method better than
  just using compiled inference in each toy domain?
  
  Clarity:

  The paper is well-written and easy to follow. I had no difficulty understanding the method though
  have not tried to replicate the experiments in the paper.
  
  Significance:

  I think this work as tremendous potential as being able to do amortised inference over a set of
  related probabilistic programs would be very useful. If this can shown to generalise across several
  domains even related ones it would make the work much more compelling.

questions: How does this work differ from Che and Yang? What's the contribution of this work in particular?

What 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.
review_title: null

final_score: rating: 7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.
review_content: 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.
strengths_and_weaknesses: **Originality/Significance:** 
Good: 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.  

**Quality:**        
Good: The paper appears to be technically sound. 


**Clarity:**          
Good: 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. 

**Weaknesses:**        
One main concern of the paper is its computation burden. I will explain in detail in the following section. 
questions: - expressivity of the supported programs. 
If 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. 

- experiments. 
I 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. 

- computation cost. 
Though 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
review_title: null

final_score: rating: 7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.
review_content: 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.
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.

**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.

**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.

Overall, 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. 

My 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.

Minor: 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.

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?
2. 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? 
3. 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?)
4. 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?
5. 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?
6. What happens if you run ADVI for longer in Figure 5? 
7. What representation is ADVI using and would it do better with more parameters? 
8. How many parameters are used in each of the models in the experiments?

[Reid et al, 2022] Can Wikipedia Help Offline Reinforcement Learning? Machel Reid, Yutaro Yamada, and Shixiang Shane Gu. https://arxiv.org/abs/2201.12122 


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.
review_title: null

Meta-review: metareview: This paper proposes a novel and interesting perspective on leveraging large masked language models as ways to initialize posterior distributions across probabilistic programming language (PPL) tasks. The idea is that this distribution can later be fine tuned over different probabilistic programs.

All reviewers acknowledged that the idea is a novel application of masked language models and, despite being a natural analogy to the way these models are already used nowadays in NLP, can be potentially impactful for amortized inference in PPLs.

The paper is accepted upon the introduction of the discussions emerged during the rebuttal concerning related works and presentation. I also advise the authors to think about substituting the term "Foundation" with a more precise technical term ("Masked Language Model Posteriors"?, "Transformer Posteriors"?,...). Imprecise umbrella terms, nowadays, bring more noise than help and inflate the hype around simple concepts.  
confidence: Certain
Decision: Accept