Venue: ICLR 2022 Conference

Reviews:

final_score: rating: 8: accept, good paper
review_content: summary_of_the_paper: This paper identifies a common problem in previous VAE related models: adding more stochastic layers to an already very deep model yields small predictive improvement while substantially increasing the inference and training time. Therefore, a new model that proposes to use attention mechanisms to build more expressive variational distributions in deep probabilistic models by explicitly modelling both local and global interactions in the latent space is proposed. The model is evaluated on standard dataset MNIST and OMNIGLOT, and showed superior performance against a wide range of baseline models.
main_review: Strength:
1. So far, this is the first work that I know of that successfully combined transformer type of attention to VAE models. Even though it is a combination of two existing models, I believe it still provides benefit and value to the VI research community.
2. Evaluation showed strong superior performance compared with a wide range of baseline models. 
3. An ablation study is included to justify each components of the proposed model.
4. The paper is well written and the proposed model is easy to understand.

Weakness:
This is a solid paper, no major weakness. The only concern I have is regarding the novelty of the proposed framework. However, as discussed before, I believe this still adds value to the variational inference research community.
summary_of_the_review: Overall, this is a pretty solid paper. Please refer to the previous section for detailed discussion.
review_title: null

final_score: rating: 8: accept, good paper
review_content: summary_of_the_paper: This paper proposes a novel attention-based architecture for deep VAEs that facilitates dependencies between non-neighbouring layers and reports improvements on the marginal likelihood over recent related methods on MNIST, Omniglot and CIFAR10.
main_review: **Relevance**
Deep VAEs are an active area of research and progress is commonly quantified by the marginal likelihood on standard benchmark datasets. Hence improvements in that regard make this a relevant contribution.

**Novelty**
Multi-layer structure as well as attention are popular techniques in their respective fields, but their combination is novel as far as I am aware. The specific method in this paper is a non-trivial combination of these, so I would consider this a solid, but not outstanding paper in terms of novelty.

**Clarity**
Overall clear. A fair bit of notation, but that seems unavoidable. There are a couple of figures to visualize the computational graph structure, which is helpful. Perhaps an algorithm box for inference and generative paths as well as the loss calculation would be useful as well.

Some terminology is used a bit loosely in my view, in particular “local” and “global”. For me as a reader with more of a general probabilistic ML background, these refer to per-datapoint and shared variables across all data respectively (so in VAEs the latent variables are local, if we did inference over the weights those would be global). But this might be standard in the more closely related literature.

**Empirical evaluation**
The benchmarks (MNIST, Omniglot, CIFAR10) and metrics (marginal likelihood) are standard as far as I’m aware. There is an extensive ablation study, although it is not obvious to me why the combination of a non-local generative and inference model without non-local layers is not considered. The paper suggests that the proposed method could be combined with autoregressive generative models, adding such an experiment would strengthen the empirical contribution of the paper.

**Other notes and questions**
* First sentence final paragraph of 2.1: p(x) should be p(z)?
* Is it possible to further increase the depth of the architecture when using attention? It would be useful to have an equivalent table to Tab 1 to support the claim that the attention-based scheme does not suffer from diminishing returns on increasing depth to the same degree as the baseline. If there are limitations due to the quadratic complexity in depth on current hardware, this should be mentioned explicitly (although I might have missed it).
* The formatting of the references is extremely inconsistent, it seems like the bibtex entries were copy-pasted from google scholar without change. Please be consistent in venue names, title capitalization, abbreviations of middle names etc.
summary_of_the_review: I’m not deeply familiar with the recent literature on (more sophisticated) VAE architectures, but the method, while a combination of existing techniques and ideas, appears to be novel and I would expect the empirical results to be of interest to the community. Therefore I am **leaning towards accept**.

**Post-rebuttal note**: Increased score to accept.
review_title: null

final_score: rating: 8: accept, good paper
review_content: summary_of_the_paper: This paper improves the architecture of deep VAEs using the attention mechanism. 

The attention mechanism is used in two ways:
1. layer-wise attention (attending stochastic feature maps which are conditioned on other variables within the hierarchy, interpreted as a mixture of skip connections)
2. non-local attention (attention across the spatial dimensions, increases the size of receptive field)

The authors demonstrate the effectiveness of their architectural changes by challenging current sota deep VAEs on MNIST, OMNIGLOT and CIFAR-10. Notably, they outperform the state-of-the-art methods (in likelihood) on CIFAR-10 using fewer layers and fewer GPU hours. 

The authors provide an extensive ablation study, showing the impact of each type of attention on the training performances (test likelihood on CIFAR-10).
main_review: ## 1. Strengths
- a. The introduction of the layer-wise attention mechanism for deep VAEs is novel, and its effectiveness is supported by empirical results
- b. The introduction of non-local attention within layers is effective, and supported by experiments
- c. The overall architecture (incl. normalisation and scaling tricks) challenges the current sota deep VAEs using fewer layers and shorter training. 
- d. Literature in deep VAEs seems to be well known by the authors and sufficiently cited

## 2. Weaknesses
- a.  The paper lacks structure and clarity 
- b. The paper lacks a more qualitative study of the model:
    - it would be interesting to see what layers the layer-wise attention mechanism attends to.
    - it would be great to understand how this model uses the latent variables, for instance by measuring the KL divergence at each layer, as done in the previous work (LVAE, BIVA) (connection to "posterior collapse").
- c. Experiments are limited to CIFAR-10, larger scaler experiments (i.e. ImageNet) would be beneficial to the paper. It is not guaranteed that such an architecture would translate in the same gains for larger datasets (i.e. ImageNet).

## 3. Clarification needed: 
- a. Table 5: I interpreted the column "non-local layers" as using "attention across layers", I hope I was right. The nomenclature needs to be improved.
- b. Is the layer-wise attention mechanism specific to deep VAEs, or can it be more generally applied to ResNet architectures? 
- c. Section 2.3 (paragraph cited bellow): I get the idea, but unless demonstrated, this remains a hypothesis.
    >"... in practice the network may no longer respect the factorization of the prior $p(z)=\prod_{l} p(z_{l} \mid z_{<l})$ leading to diminished performance gains as shown in Table 1":

## 4. Minor comments / suggestions
- a. The main contributions are introducing two types of attention for deep VAEs, it might help to describe them in a separate section, and only then describe the generative and inference models. Right now the description of the layer-wise attention mechanism is scattered across sections 2.3 and 2.4.
- b. tricks like normalisation or feature scaling could be referenced in a separate section. 
- c. eq8: you might want to cite ReZero [1] here
- d. Fig 1. a: the lack of arrows going from the activations $(h_l, k_l^q)$ to the attention block $(\mathcal{A}(...))$ was confusing on the first read
- e. It would be better practice to report likelihoods for multiple random seeds
- f. Typo in section 2.1: "both $q(z|x)$ and $p(x)$ are fully factorized gaussian..." -> "both $q(z|x)$ and $p(z)$ are fully factorized gaussian..."

[1] Bachlechner, T., Prasad Majumder, B., Mao, H. H., Cottrell, G. W., and McAuley, J., “ReZero is All You Need: Fast Convergence at Large Depth”, <i>arXiv e-prints</i>, 2020.

summary_of_the_review: I have enjoyed reading your paper, nice work! The introduction of the layer-wise attention mechanism is a great contribution. Using non-local blocks within the model is less novel, but this is still new for such a model. 

The results on MNIST, Omniglot and CIFAR-10 demonstrate the effectiveness of the method, and the ablation study clearly shows the positive effect of the two attention modules (layer-wise and spatial). 

Experiments remain however limited. Large-scale experiments (ImageNet) and qualitative studies (inspective learned attention) would be a great addition. They would give the reader a better understanding of the impact of the attention modules. The paper overall writing quality and structure need to be improved.

---- edit 26/11
Raised my score to 8.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: This paper presented a variational inference with attention mechanism. A deep latent variable model was proposed.
main_review: Pros:
1. This was the first work where the attention mechanism was proposed to carry out the deep hierarchical VAE. By explicitly modeling the local and global interactions in latent space, an expressive variational distribution was constructed with probabilistic justification.
2. This paper provided the experimental evidence to show that the previous method on expressive variational model via designing the deep hierarchy of interdependent latent variables would incur the problem of diminishing returns.
3. The experimental justification with the comparison over different types of VAE models was sufficient.
4. Computational cost was reduced when compared with the previous SOTA model based on NVAE.

Cons:
1. It is suggested to provide an algorithm to enhance the comprehension of the detailed training procedure of the proposed method.
2. The experiments did not show the generated samples the authors obtained. The quality of the generated images could not be evaluated.
3. Ablation study was only conducted on the effect of non-local information. The ablation studies different normalization and activation functions are required.
summary_of_the_review: An interesting work was proposed with experimental justification. The experiments can be furthered strengthened.
review_title: null

Meta-review: This paper adds an attention mechanism to deep variational autoencoders.  The authors develop a global + local attention method and achieve better log likelihoods than a variety of recent methods on MNIST and OMNIGLOT.  Overall the reviewers found this paper strong (8, 8, 8, 6), particularly after the author rebuttal.  They found the paper to be clear, the contribution sensible and novel and the experiments thorough and compelling.  In particular, the authors added additional experimental results on a larger dataset which addressed a common concern among the reviewers.  Thus the recommendation is to accept the paper.
Decision: Accept (Poster)