Venue: ICLR 2023 Conference

Reviews:

final_score: rating: 5: marginally below the acceptance threshold
review_content: summary_of_the_paper: This paper presents an approach for variational inference on large-scale hierarchical models which takes advantage of conditional independence in plate structures to factorise its variational approximations and better scale with problem dimensionality. The proposed framework is based on normalising flows taking shared parameters and data-encoding vectors. The resulting algorithm is equipped a stochastic optimisation scheme for training and is shown to converge faster than other algorithmic baselines in experiments.
strength_and_weaknesses: ### Strengths
* The proposed method achieves significant (orders of magnitude) performance speed-ups when compared to baselines in experiments.
* Although the idea of using factorised variational posteriors for hierarchical models is not new, this paper implements it in a more computationally efficient way, which is also amenable for per high-performance computing with hardware accelerators.
* Diagrams help understand the concepts in the approach.

### Weaknesses

* Related work: There are other flexible variational inference approaches for hierarchical models which also factorise the variational posterior. For example, Tran et al. (2017) introduced hierarchical implicit models for this task, and other relatively simple approaches for random effects models (as in the example experiment) are also available (e.q., Dao et al., 2022).

* Experiments: There was only one benchmarking model (Eq. 6) and variations thereof for the comparisons. The other experiment (Fig. 5) did not provide comparisons against baselines due to its complexity. Another representative set of experiments would help assessing the improvements over baselines in corner cases where other methods can be applied but PAVI should in principle perform better.

#### References
Tran, D., Ranganath, R., & Blei, D. M. (2017). Hierarchical Implicit Models and Likelihood-Free Variational Inference. 31st Conference on Neural Information Processing Systems (NIPS).

Dao, V. H., Gunawan, D., Tran, M., Kohn, R., Hawkins, G. E., & Brown, S. D. (2022). Efficient Selection Between Hierarchical Cognitive Models: Cross-Validation With Variational Bayes. Psychological Methods, (April).
clarity,_quality,_novelty_and_reproducibility: * **Clarity:** The paper is quite dense to read. For instance, sections 3.1 and 3.2 seem to be repeating what is in Sec. 2.3 and 2.4 at times perhaps a bit too much. Something more interesting would be to expand on the unbiasedness of the approximation, which has been left for the appendix (A.2).

* **Novelty:** I've found some of the methodology somewhat similar to that of ADAVI (Rouillard & Wassermann, 2022). The difference lies on the use of conditional independence structures between parameters to further factorise the posterior.

* **Reproducibility:** Code is provided and the dataset is open access.

summary_of_the_review: The paper's contributions are somewhat novel, but it lacks in experimental comparisons against the baselines in other challenging scenarios.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: This paper proposes Plate-Amortized Variational Inference (PAVI), an inference framework for hierarchical Bayesian models with amortization over the population (dataset) level implemented with normalizing flows with shared parameters. The main idea is to posit an expressive variational family parameterized with normalizing flows, and yet reduce the number of parameters to be trained by sharing part of the variational parameters across populations. Another important concept is stochastic training, where for each step, instead of training all the variables present in the graphical model, subsample a small number of variables and compute stochastic gradient with them. Thanks to amortization, PAVI could scale to large-scale population studies involving high-dimensional parameters.
strength_and_weaknesses: Strengths
- The paper is well-written and easy to follow.
- The description of amortization using the language of plate diagram is interesting.
- The neuroimaging experiment seems interesting.

Weaknesses
- In my opinion, the notion of plate amortization is not entirely new, so I think the paper can be strengthened by discussing relevant existing works.
- The experiments include a synthetic Gaussian model and neuroimaging experiments. The synthetic experiment looks promising, but it is of a small scale that does not particularly highlight the benefit of the proposed method. For the neuroimaging experiment, although it seems promising that only the proposed method was manageable to train, it is not clear whether the prediction given by the model is reasonable (at least for the non-experts like me).

Detailed comments
- The idea of plate amortization is indeed interesting. I guess the proposed framework virtually extends to hierarchical Bayesian models with an arbitrary level of hierarchies, but the main application would be the amortization across populations (datasets). Regarding amortization over datasets, there have been similar ideas (although motivated differently), especially in the context of meta-learning. Neural processes families, for instance, introduce an amortized inference network that can readily generalize to a new task (corresponding to the population) which is parameterized by the dataset encoding networks (deep sets or transformers); this would roughly correspond to PAVI-F, except that PAVI-F uses normalizing flows for the parameterization. Amortized Bayesian meta-learning (Ravi and Beatson, 2019) is also worth discussing in a similar sense. Depending on the application, neural processes can be considered as a baseline to be compared to the proposed method, since it can also scale to large population study due to its use of the amortized inference network.
- The experiment section could be enhanced. I understand that the neuroimaging experiment can be compelling, but in the current form, I'm not entirely sure how well the proposed method is doing; the only result I could see from the paper is Figure 5, which I fail to interpret. It would be great to see the numbers, for instance, metrics could tell us how accurate PAVI is, or at least how well the predictions given by PAVI match the domain knowledge given by experts. Also, if the baselines are not trainable due to their limited scalability, one could consider comparing them on a smaller subsets of manageable sizes to see if PAVI would do better than baselines for real-world tasks.
clarity,_quality,_novelty_and_reproducibility: Clarity and quality: I think the paper is well-written in general, although there are some complex notations to digest which would be hard to fully grasp at a first glance.
Novelty: as mentioned above, it is worth discussing the relevant works in meta-learning literature, and clarifying the novelty of the proposed method compared to them.
Reproducibility: I think the paper provides enough details to reproduce the experiments.
summary_of_the_review: Overall, the proposed method is an interesting contribution, but the paper could be strengthened by adding more discussions and hopefully more experimental results/interpretations.
review_title: null

final_score: rating: 3: reject, not good enough
review_content: summary_of_the_paper: The work proposes an amortized VI scheme that leverages the hierarchical structure of a Bayesian model, using the latter's plate representation. Plates are used to instantiate subsets of the variables of interest, and share weights in the encoding of sets of variables used on a conditional flow architecture. The result is a plate-amortized proposal distributionI. This idea seeks to enable VI in a large scale setup. The methodology is validated empirically against the state-of-the-art, with one application with real data in neuroimaging context.

strength_and_weaknesses: ## Strengths:
* The methodology is generic and intuitive
* Potential for large scale applications
* The idea of embedding structure into the approximate inference scheme is potentially fruitful line of further research
* The authors provide detailed discussion in the appendix, including some limitations regarding the approximation gap.
* The empirical validation in a realistic large scale scenario.

## Weakness:
* Heavy notation makes the text hard to follow;
* Limited novelty (e.g., the work is incremental when compared with ADAVI). It seems that the major difference between PAVI and ADAVI is the stochastic updates of variables;
* I believe the gradient estimates for PAVI are biased;
* Limited experimental campaing, limited to a small set of methods and models. Validation against multiple methods in smaller dataset could be beneficial.

## Questions
1. In the demonstration on Appendix A.2 it indicates that PAVI-F is unbiased by showing in A.12 that the ELBO term is equal the term with all the variables, however Section A.4 implies that there is a gap between vanilla VI, PAVI-F and PAVI-E: Gvanilla VI ≤ GPAVI-F ≤ GPAVI-E . This seems unintuitive and indicative that indeed PAVI-F might be biased, given that the term that differ in those approximation schemes is indeed the ELBO term. Are both of those results consistent?

2. How does the proposed methodology would compared with MCMC/HMC? It could be beneficial given that a new approximation scheme is being proposed to compare it as well sampling methods. 

3. When it comes to the expressivity of the variational approximation. How does the predictive performance of the method compare with others variational inference proposals (empirically)? For example another family of expressive VI approximations are found in the works such as UIVI [1], SIVI [2] or VCD [3], which does not use flows and yet are generic and could be a better test bench and give the reader more confidence about the proposed method.

4. In Eq. A.11, when two variables share the same descendants, hence leading to correlated E_{i,n} in the PAVI-E scheme, wouldn’t in this case lead also to dependency between the indicator random variables and log q_{i,n} (given the dependency of log q on the encoding)?

clarity,_quality,_novelty_and_reproducibility: Clarity: Heavy notation makes the text hard to follow.

Quality: I have doubts about the theoretical soundness, the experimental campaign is limited.

Novelty: seems incremental compared to ADAVI.

Reproducibility: I have no concerns regarding reproducibility.
summary_of_the_review: Given 1) the limited novelty, 2) short experimental campaign, 3) and possible issues with theory, I recommend rejection.
review_title: null

final_score: rating: 5: marginally below the acceptance threshold
review_content: summary_of_the_paper: The authors consider approximate inference in Bayesian networks, specifically in the subset of Bayesian networks that can be represented as plate diagrams, which are common for large population studies, the population is composed of subgroups and each group is composed of many individuals. Given global the population parameters, each of the group parameters are i.i.d. and given group parameters, the individuals are i.i.d., a common hierarchical structure.

Assuming a fully specified Bayesian generative model of the data, and only leaf nodes are observed, the authors propose a variational inference method to learn the full approximate posterior distribution to all of the unobserved nodes. The method combines normalising flows with the generative model priors to approximate the posterior. The method is applied to illustrative toy datasets and shows results to be expected and a high dimensional neurological time series dataset (but results are not presented in the main paper).

strength_and_weaknesses: Overall, I struggled with the paper and I have given a low confidence score due to possible misunderstandings.

**Strengths**
- the Supplementary material seems very comprehensive and answered many of the questions I had after reading the main paper.
- the method seems well designed and thought out
- the encoding size experiment provided good confirmation that the degrees of freedom must match the sufficient statistics.

**Weaknesses**

I feel this paper is very hard to recommend for acceptance due to the writing and the somewhat basic numerical evaluation.

I felt the writing and the structure of the paper was very poor, it took me a great deal of effort and multiple rounds of re-reading and re-evaluating my own understanding of the paper before finally appreciating the paper and its contributions (and I have experience in graphical models).
  - I would make section 2.1 into it's own section that really makes the mathematical problem clear. E.g. a plate diagram, the generative model in  equation (1), a real world example with cardinalities and (non-linear) link functions and latent variables. Describe the practicalities of inferring parameters in such a model and why it is impractical scaling with cardinality.
- I don't know if this is standard or not but I was personally thrown off by the terminology "RV template $\theta_i$" and "ground RVs $\theta_{i,n}$" and "grounding a template with cardinalities". Is is possible to reduce notation? E.g. $\mathcal{T}$ is introduced but never used, plates $\mathcal{P}$ is simply indicated by $i$, Card($\mathcal{P}$) is just $N_i^{full}$. Perhaps $N^{full}$ and $N^{redu}$ could be replaced by $N$ and $\tilde{N}$, i.e. use tilde "~" to denote reduced version of various quantities?
- what does $\pi(\theta_{i,n})$ explicitly represent in equations (1) and (2)? I believe they should be precise, known, values if they are conditioning a distribution, but this is seems vaguely defined in the paper. In Ambrogioni 2021b they say "$\theta_j$ is the array of values of the parents of the parents of the $j^{th}$ variable, and $\pi_j(theta_i)$ and $\pi_j()$ is the link function".
- what is the prior for an RV, $p(\theta_{i,n}| \pi(\theta_{i,n}))$? My intuition is that prior is the should be the generative model up to $\theta_{i, n}$ with all of the parent variables marginalized out? Or is it the generative model term with conditioned on given parent values? In the latter case, having such a distribution as the input to a normalizing flow doesn't feel like it would lead to a good posterior approximation, parent values change, prior changes, normalizing flow function is constant but the input changes?
- Section 2.3: please give dimensions of all **X**, and $\theta$ and I assume one row/col of  **X** is a subgroup indicator variable?
- (nitpick) I am not entirely sure I agree with the use of the word "amortization" in the PAVI-F, the fact that $\phi_i$ is shared across all nodes in a layer does not makle inference faster at test time for new nodes? I still have to learn $E_{new}$ from scratch? (aAlthough amortization means "spreading out", in this community, I feel it refers to predicting an approximate posterior in one shot given observed variables.
- the experiments are toy examples and the only non toy example is rather brief with no quantitative results in the main paper
- Lack of Technical Novelty
  - randomly subsampling nodes in the DAG seems like a very simple straightforward idea and I think the unbiasedness is significant and shows that this isn't "just a hack" but a real valid method, however it is rather incremental.
  - claiming to have reduced memory footprint by "using a low dimensional encoding" seems ambitious claim of novelty. While it may be novel, forcing the method to use a low dimensional parameterization of a normalizing flow seems incremental.

- I was surprised that the discussion about "faithful posteriors" (Webb et. el. 2019) was in the SM, to me this felt like very closely related work and provides a method for building approximate posterior Bayesian networks,
clarity,_quality,_novelty_and_reproducibility: **Clarity**
- as above, I felt the arguments in the main paper can be made more concisely, I felt this was unnecessarily difficult for me to understand. The parts of the SM that I did read I felt were clear and thorough, and I wish the main paper were made more concise and explicit and many of the SM details may be moved to the main paper. (pseudocode, results and baselines for neuroimaging application)

**Novelty**
- I feel the main contributions seem to be the splitting of normalizing flow parameters into $\phi_i$ shared across multuiple RVs and  RV specific parameters $E_{i,n}$, as well as the unbiased mini batching, which I feel are novel, though their significance may not be sufficient for acceptance.

**Reproducibility**
- the SM is extensive and contains many details and pseudocode, I feel the authors have been very thorough and, apart from a few misunderstanding I have, the method and toy experiments appear reproducible.

summary_of_the_review: - clarity of the main paper is very poor, verbose and possibly non-standard terminology for what appears to be standard Bayesian networks.
- the technical contribution, parameterization of normalizing flows and mini batch training, seem novel but incremental.
- the experiments are mostly toy, one non toy is only briefly considered.

UPDATE: I have updated my score from 3 to 5 after reading through the responses.
review_title: null

Meta-review: metareview:_summary,_strengths_and_weaknesses: The paper aims to speed up and reduce the memory footprint variational inference in hierarchical Bayesian models. The technique proposed is to use plate amortization to share a single distribution across multiple random variables. The shared distribution takes an encoding vector to allow for the distribution to not be the same for different random variables. Multiple encoding schemes are considered. The results look good to a relatively simple baseline of raw stochastic variational inference and several flow based approaches.  

The reviewers were negative on the paper initially. This changed to a mixed leaning negative after the author response. With authors moving from reject to borderline reject. One concern was clarity in the notation which was raised by several reviewers. Another concern that was partially addressed in the revision was on the breadth of the experiments, though I still find them still a bit lacking. There was a discussion about bias that is important and brought up by the reviewers, but is buried in appendix a.2.3. In total, the experimental breadth, the writing, and more detailed discussion of bias in the main paper makes me think the paper needs more work.

Beyond those concerns, there are still missing related works such as 

@inproceedings{hoffman2015structured,
  title={Structured stochastic variational inference},
  author={Hoffman, Matthew D and Blei, David M},
  booktitle={Artificial Intelligence and Statistics},
  pages={361--369},
  year={2015}
}

and some minor bits on the definitional correctness on what stochastic variational inference is. The cited Hoffman paper (figure 4 or equation 25) propose maximizing over local variational parameters and never storing them, while the stochastic variational inference described is one that explicitly stored.
justification_for_why_not_higher_score: In total, the experimental breadth, the writing, and more detailed discussion of bias in the main paper makes me think the paper needs more work.
Decision: Reject