Venue: NeurIPS 2022 Conference

Reviews:

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 new way to automatically construct and train variational families for models expressed as general probabilistic programs. 

Given a program $p$, it first generates many (prior) samples from $p$ to identify a collection of control flow paths that collectively have high mass under the prior. Then, a separate variational family is trained to target the restriction of the model to each control flow path. During training, control-flow paths with low ELBOs are periodically pruned, so that computation is focused on paths that appear to have relatively high posterior probability. Finally, the learned variational families are combined into a mixture, with mixture weights based on estimated ELBOs for each path. 

Experiments show competitive performance against other fully automated inference strategies for probabilistic programs, on three inference problems.
strengths_and_weaknesses: Strengths:

1. Convincing empirical results, showing that for a similar computational budget, the new method outperforms a few existing fully automatic inference methods.
2. The paper grows the body of evidence that extracting SLPs and tackling inference separately in each (as done e.g. in DCC) is a useful technique for fully automated inference in universal PPLs.
3. In models with discrete latent variables that influence control flow, gradient estimators are often based on REINFORCE, and a variant of the "explore/exploit" problem from reinforcement learning arises: without any gradient signal helping it to explore other branches of the control flow graph, the variational family may settle into 'exploiting' one suboptimal branch. This work presents a nice way of introducing some kind of 'exploration', by allocating some training time up front to many different branches, and pruning only if they do not look promising. This is a neat way of getting around the high variance of the REINFORCE estimator. (It's a bit similar to the strategy of exactly marginalizing the discrete variable, except maybe a bit less expensive in the long run, because if one path is not promising, it is pruned early.) 
4. An implementation in Pyro is provided, which increases the likelihood that this work will have real impact for practitioners.

Weaknesses:

1. In Section 4.5, rejection sampling the guide changes the variational family's density, as the authors write, introducing a "Z" term that measures the probability that the guide generates a sample that lies within its assigned SLP. But this Z term, which depends on the guide's parameters and so must be considered during optimization, appears to be intractable, and the authors do not explain how it is computed when estimating the ELBO and the ELBO's gradient. If it is ignored, I am worried that the gradient estimates would be biased. Furthermore, it doesn't seem this feature is 'stress-tested' by the experiments, the latter two of which use guides that ensure no rejection is necessary.

2. I did not see any clear discussion of the method's limitations. At times, I thought the authors made misleading claims that downplayed limitations. For example, L176 says of the SLP discovery process that "experiments show it can reliably identify all SLPs with non-negligible posterior mass." But my understanding is that you only discover SLPs with non-negligible _prior_ probability mass; in all the experiments, the posterior was concentrated in SLPs likely under the prior. Consider fitting the Gaussian mixture model to a dataset with many more components than the prior expects. The posterior will concentrate far from the prior, in an SLP that may not even have been discovered by the algorithm. Of course, this is a difficult case for many inference algorithms, but the limitation should be clearly explained. 

3. I find Figure 1 and its explanation on L99-110 misleading. The explanation makes it sound as though the problem in the example is the expressiveness of the default variational family (it cannot represent the posterior on x, due to stochastic control flow). But the striking behavior in the figure, where Autoguide fails, arises just because you are using the reparameterization trick in a non-reparameterizable model. This seems to be a totally separate question from the expressiveness of variational families: it's about whether the Pyro software happens to detect when to switch to the score function gradient estimator automatically. Unless your method *does* automatically detect when the reparameterization trick can or cannot be soundly applied, without any user intervention, this feels irrelevant to the paper. It reads as though a (predictable) failure of the gradient estimator is being passed off as a limitation of the variational family. This also applies to experiments that report results using biased gradient estimators—to me that feels like an orthogonal issue.

4. I don't think the introduction (e.g. L24-42) gives an accurate summary of the state of the art. It is true that many PPLs, including those listed here, aim to support variational inference, but the aim has *not* in general been to automate the construction of guides (although of course, that is an interesting goal). Although Pyro has an Autoguide feature, many of the other PPLs in the list do not: the *intended* use case is that a user constructs a model *and* a guide as probabilistic programs, and the system automates gradient estimation and optimization of the ELBO. (Thus, L92-98 are also a bit misleading.) Even in Pyro, which does support automatic guide construction, many of the tutorials show how to construct custom guides, and the flexibility to do so is touted as a key feature. Arguably, this is what makes languages like Pyro useful for a much broader variety of problems than earlier languages that attempt to fully automate inference. I like that this paper tackles the hard problem of fully automated inference, and I think it makes progress in that direction, but the scope of the problem within the broader field of probabilistic programming should be clearly described: many PPLs already support accurate variational inference (with correct gradient estimates) in models with stochastic control flow if the user provides a sensible variational family. This paper is attempting to further automate the inference process by automatically constructing guides for a broader class of models than existing autoguide techniques can handle.


Small nitpick: on line 63, I don't think "higher-order" is relevant -- a first-order language with branching and recursion yields the same class of models.
questions: 1. In Section 4.5, you describe two methods for resolving the issue that local guides may not 'stay within' the boundaries of their SLPs. The way I read it, you are applying both solutions (mixing the target with a constant density, and rejection sampling the guide)---but isn't only one necessary? If you are rejection sampling, why introduce unnecessary discontinuities into the target?

2. I am confused about how the rejection sampling in Section 4.5 is supposed to work. As you write, the density of the variational family now contains a "Z" term, equal to the probability that q produces a sample that is accepted. This Z term must be computed when estimating the ELBO and its gradient. (Though the notation does not make this clear, Z depends on the parameters phi, so Z cannot simply be dropped when training.) How do you compute Z and its gradient, which in general will be an intractable integral over the SLP's region (which could be strangely shaped)?

(Note: my score is only "borderline accept" because I am not sure I've understood the method / if I have I am not sure it's sound, and also because I have concerns about certain misleading aspects of the current presentation, described in "Strengths and Weaknesses." If these are addressed during the author response, I hope to revise my score to a more confident recommendation.)

**EDIT AFTER DISCUSSION:** The author response has addressed my concerns and I am increasing my score to a 7.
limitations: As stated in "Weaknesses," I think the limitations are not adequately discussed. The class of probabilistic programs in universal PPLs is huge: when should we expect this algorithm to work, and when should we expect it to fail? i.e., when might users still resort to custom guides in Pyro? What assumptions does the new method make (e.g., non-negligible posterior mass on an SLP implies non-negligible prior mass)?
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: The paper proposes a new variant of variational inference for probabilistic programs with stochastic support. To address the challenges stochastic support poses, this paper proposes a novel way of constructing a variational guide by breaking down the problem into sub-programs. The paper demonstrates that this approach results in improvements in inference performance.
strengths_and_weaknesses: 
Strengths: 
1) The paper addresses a very important and useful problem and the proposed method seems to be a nice and useful addition to probabilistic programming languages. I especially appreciate that the code is made available.
2) The paper is very well written and easy to understand
3) The paper contains comprehensive experiments to demonstrate the effectiveness of the method.

Weakness: 
1) The work seems to be incremental over [33-35] especially [35]. The authors need to justify how their work differs from the aforementioned references.

2) The paper needs to include more related works and explain the differences and similarities between the proposed method and related works in more detail. There has been a lot going on in this field and it begs a more comprehensive literature review. 
questions: - Could authors justify the main differences between their work and that of for instance [35]?
- How does it compare to MCMC VI?
limitations: - This work seems a little incremental thus some justification is required to distinguish between this work and others. 
- Choosing a variational family is always a crucial task when doing VI. I am not sure how it is compared to the MCMC VI models.
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: The paper proposes a variational inference method named Support Decomposition Variational Inference (SDVI) that targets probabilistic programs with stochastic support. The main idea is to perform variational inferences separately on sub-programs with static support. The variational approximation has a mixture form where each variational component is optimized independently for each sub-program. Finally, the approach is compared to baseline algorithms on a Gaussian model with stochastic control flow, on an infinite Gaussian mixture model (IGMM) with a stochastic number of clusters, and on inferring the kernel structure of a Gaussian Process.
strengths_and_weaknesses: The paper is clearly written and the structure is easy to follow. I don’t have problems getting the main ideas.

The paper tackles the universality perspective in probabilistic programming systems. This seems to be a difficult problem. Though the main idea follows the “Divide, Conquer, and Combine” approach proposed by Zhou et al. [35], using variational inference instead of MCMC poses different challenges, which are reasonably addressed in the section 4:
+ Using the Successive Halving algorithm to avoid wasting computational resources on less interesting sub-programs.
+ Truncating the target density by a small positive value to ensure obtaining a finite ELBO (when a variational component proposes samples outside of the support of its corresponding sub-program).
+ Variational components need to be initialized smartly so that they have sufficiently large probability mass on the support of the corresponding sub-programs.

Despite that the authors already provide an implementation in Pyro, I feel that it is a bit tricky for practitioners to adopt the approach due to the complication of the above challenges.

The paper lacks discussions on how expensive the approach is (though it seems that the method is scalable through the IGMM experiment). This is important for users who want to apply the method in practice.

In addition, though the experimental results seem to be correct, it would be more interesting for readers to know in which situations we need to construct a probabilistic program with stochastic support. This might deserve a paragraph in the introduction section.
questions: The “rethinking” word in the title seems to create confusion because the main algorithm is only compared to the simple Gaussian approximation.

The mention of Pyro’s AutoGuide in Figure 1 may cause confusion. It might be better to use an explicit algorithm name, e.g. Gaussian variational approximation or the like.

It is not clear to me how we can differentiate between different paths of a program. Based on the program in Figure 1, apparently it requires users to use different variable names (z1, z2) for different paths. Is this a requirement for users?

In the discussion of Table 1, “MAP estimate of K=6, the local guide q_k corresponding to the SLP with 5 components” - should this be SLP with 6 components?
limitations: Nothing to report.
review_title: null

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: The authors address the problem of designing variational proposals for universal probabilistic programs with stochastic support.  Their method breaks down the stochastic support of the target program into a mixture model over control-flow paths through a space of straight-line programs, and the authors explain that their variational proposal method contrasts with Divide-Conquer-Combine (its closest cousin, used for MCMC sampling).  They demonstrate their method's performance on a toy program, an infinite mixture model, and learning Gaussian process kernels in a PCFG.  Their results appear to clearly set a new state of the art in log-predictive density and ELBO.
strengths_and_weaknesses: Strengths:
* Variational inference in deep universal PPLs is widely applicable and significant to a broad variety of downstream problems
* Clearly superior experimental results, with 100*K dimensionalities
* Latter two experiments are clearly nontrivial for variational inference settings
* Color-coding provides clarity
* Explanations and mathematics are clear to PPL readers

Weaknesses:
* Explanations and mathematics will be unclear to readers without probabilistic ML background
questions: When comparing against sampling-based inference methods, why not include one based on importance sampling?  Resample-move SMC inside DCC, for instance, could provide an estimate of the normalizing constant, and therefore (in log space) an ELBO to which to compare.

I don't understand what's meant by "marginal likelihood" for an SLP within a larger program.  Not every straight-line subprogram is going to make observations before reaching a site of stochastic control-flow, and so your importance weight could just be a prior/proposal density ratio.  Is that ratio's expectation used as a normalizing constant to allocate compute power?

The authors have now clarified and explained away my previous concerns.
limitations: The authors, peculiarly for a Pyro paper, make no actual use of neural networks to represent conditional densities.
review_title: null

Meta-review: metareview: The reviewers have reached consensus after processing the authors' feedback. They all agree that this manuscript presents an interesting approach to applying variational inference in a setting of probabilistic programming that is of interest to the community. The reviewers raise tangible points that the authors have incorporated into their revision. I recommend that the authors continue to polish their manuscript to clearly address these points in the final version of their manuscript.
confidence: Certain
Decision: Accept