Deep Attentive Variational Inference
Authors: Ifigeneia Apostolopoulou
Ian Char
Elan Rosenfeld
Artur Dubrawski
Venue: ICLR 2022 Poster (2022)
Keywords: variational inference
approximate inference
deep probabilistic models
deep probabilistic learning
variational autoencoder
probabilistic methods for deep learning
attention

Abstract: Stochastic Variational Inference is a powerful framework for learning large-scale probabilistic latent variable models. However, typical assumptions on the factorization or independence  of the latent variables can substantially restrict its capacity for inference and generative modeling. A major line of active research aims at building more expressive variational models by designing deep hierarchies of interdependent latent variables. Although these models exhibit superior performance and enable richer latent representations, we show that they incur diminishing returns: adding more stochastic layers to an already very deep model yields small predictive improvement while substantially increasing the inference and training time. Moreover, the architecture for this class of models favors local interactions among the latent variables between neighboring layers when designing the conditioning factors of the involved distributions. This is the first work that proposes attention mechanisms to build more expressive variational distributions in deep probabilistic models by explicitly modeling both local and global interactions in the latent space. Specifically, we propose deep attentive variational autoencoder and test it on a variety of established datasets. We show it achieves state-of-the-art log-likelihoods while using fewer latent layers and requiring less  training time than existing models. The proposed non-local inference reduces computational footprint by alleviating the need for deep hierarchies. Project code:
https://github.com/ifiaposto/Deep_Attentive_VI

Published as a conference paper at ICLR 2022
DEEP ATTENTIVE VARIATIONAL INFERENCE
Ifigeneia Apostolopoulou" .2, Ian Charl.2 Elan Rosenfeld' & Artur Dubrawski Maahhnee Learning Department AAttnnLa,, Carnegie Mellon University Correspondence to: iapostol@andrew cmu cmu.edu, ifiaposto@gmall. com
ABSTRACT
Stochastic Variational Inference is a powerful framework for learning large-scale probabilistic latent variable models. However, typical assumptions on the fac- torization or independence of the latent variables can substantially restrict its ca- pacity for inference and generative modeling. A major line of active research aims at building more expressive variational models by designing deep hierar- chies of interdependent latent variables. Although these models exhibit superior performance and enable richer latent representations, we show that they incur di- minishing returns: adding more stochastic layers to an already very deep model yields small predictive improvement while substantially increasing the inference and training time. Moreover, the architecture for this class of models favors proxi- mate interactions among the latent variables between neighboring layers when de- signing the conditioning factors of the involved distributions. This is the first work that proposes attention mechanisms to build more expressive variational distribu- tions in deep probabilistic models by explicitly modeling both nearby and distant interactions in the latent space. Specifically, we propose deep attentive variational autoencoder and test on variety of established datasets. We show it achieves tiie-of-thh-art log-likelihoode while using fewer latent layers and requiring less training time than existing models. The proposed holistic inference reduces com- putational footprint by alleviating the need for deep hierarchies. Project code: https: pects://githuse com. com wwtoom/fffapooo http:///cccc..aaaaaeeeeeeeeeeeee Atteeeppttteeeiee___.
INTRODUCTION
core line of research in both supervised and unsupervised learning relies on deep probabilistic models. This class of models uses deep neural networks to model distributions that express hy- potheses about the way in which the data have been generated. Such architectures are preferred due to their capacity to express complex, non-linear relationships between the random variables of inter- est while enabling tractable inference and sampling. Latent variable probabilistic models augment the set of the observed variables with auxiliary latent variables. They characterized by posterior distribution over the latent variables, one which is generally intractable and typically approximated by closed-form alternatives. They provide explicit parametric specification of the joint distribu- tions over the expanded random variable space, while the distribution of the observed variables is computed by marginalizing over the latent variables. The Variational Autoencoder (VAE) (Kingmal & Welling, (0044) belongs to this model category. The VAE uses neural networks for learning the parametrization of both the inference network (which defines the posterior distribution of the latent variables) and the generative network (which defines the prior distribution of the latent variables and the conditional data distribution given the latent variables). Their parameters are jointly learned via stochastic variational inference (Paisley et al.) 2012 Hoffman et al. 2013).
Early VAE architectures (Rezende et al., 2014) make strong assumptions about the posterior distribution--spciifcce it standard to assume that the posterior is approximately factorial. Since then, research has progressed on learning more expressive latent variable models. For example, Rezende & Mohamed 2015);; Kingma et al. (2016); Chen et al. (2017) aim at modeling more com- plex posterior distributions, constructed with autoregressive layers. Theoretical research focuses on deriving tighter bounds (Burda et al.) 2016 Li & Turner 2016 Masrani et al.] (019) building
more informative latent variables by mitigating posterior collapse (Razavi et al. 2019a] Lucas et al.
2019). Sinha & Dieng (2021) improve generalization by enforcing regularization in the latent space

---

Published as a conference paper at ICLR 2022
for semantics-preservingg transformations of the data. Taking different approach, hierarchical VAEs (Gulrajami et al.) 2017 Sonderby et al. 2016 Maaløe et al.) 2019; Vahdat & Kautz 2020 Child 2020) leverage increasingly deep and interdependent layers of latent variables, similar to how sub- sequent layers in a discriminative network are believed to learn increasingly abstract representations of the data. These architectures exhibit superior generative and reconstructive capabilities since they allow for modeling of much richer structures in the latent space.
Previous work overlooks the effect of long-range correlations among the latent variables. In this work, we propose to restructure common hierarchical, convolutional VAE architectrres in order to increase the receptive field of the variational distributions. We first provide experimental evidence that conditional dependencies in deep probabilistic hierarchies may be implicitly disregarded by current models. Subsequently, we propose a decomposed, attention-guided scheme that allows a long-range flow of both the latent and the observed information both across different, potentially far apart, stochastic layers and within the same layer and we investigate the importance of each proposed change through extensive ablation studies. Finally, we demonstrate that our model is both computationally more economical and can attain state-of-the-art performance across a diverse set of benchmark datasets.
2 PROPOSED MODEL
2.1 DEEP VARIATIONAL INFERENCE
A latent variable generative model defines the joint distribution of set of observed variables, æ RD, and auxiliary latent variables, 2, coming from a prior distribution p(z). To perform inference, the marginal likelihood of the distribution of interest, p(x), can be computed by integrating out the latent variables:
p(x) p(xx dz.
(I)
Since this integration is generally intractable, a lower bound on the marginal likelihood is maximized instead. This is done by introducing an approximate posterior distribution q(= x)) and applying Jensen's 's inequality:
log p(x) = log / p(.r, 2) dz log / q(= x))
p(ar, 2) 2 / g(> 2:) log log
p(a z))(z)]
dz
g(> x)
q(2 x))
logg(x)) 2 (e [log =)]] D (xz(a(= x)) p(=)),
where 0, Φ parameterize p(x, 0) and q(= x;Φ) respectively. For ease of notation, omit 0, 6 in the derivations. This objective is called the Evidence Lower BOund (ELBO) and can be optimized efficiently for continuous via stochastic gradient descent (Kingma & Welling 2014
Rezende et al. 2014).
For ease of implementation, it is common to assume that both q(> 2) and p(z) are fully factorized Gaussian distributions. However, this assumption may be too limiting in cases of complex under- lying distributions. To compensate for this modeling constraint, many works focus stacking and improving the stability of multiple layers of stochastic latent features which are partitioned in groups where L is the number of such groups (Rezende al. 2014. Gul- such that z {21, (21,22..... 22..... ZL
rajani et al. 2017 Kingma et al., 2016 Sønderby et al. 2016 Maaløe et al. 2019 Vahdat & Kautz 2020 Child 2020). Our work builds on architectures of bidirectional inference with a deterministic bottom-up pass. The schematic diagram of a stochastic layer in such deep variational model is
depicted Figure I
bidirectional inference architecture with a deterministic bottom-up pass (left part in Figure Ta posterior sampling is preceded by sequence of non-linear transformations,
TV, of the evidence, æ. i..., hi = (h:+1) with hL+1 = æ. The inference q x)) and generative p(z) network decompose in identical topological ordering: g(z
II 9(2) x, 2<1) and p(z) = II p(z1 z<1). The top-down pass (right part in Figure Tab generates the posterior samples
x))
that feed the conditional data distribution p( 2).
2.2 MOTIVATION
Kingma et al. 2016) proposed a strongly connected directed probabilistic graphical model for the generative and inference network, so that each variable depends all the previous in the hierarchy:
2

---

Published as a conference paper at ICLR 2022
hu
k?
ResNet
RosNet
1.
ResNot
4
hr
k
hetz
ky
Attend
3.
Attend
ResNet 1.
cf
3.
3.
4.
p(4: =xx)
g(zi
p(zi
hat
(a) Locally connected variational layer.
(b) Strongly connected variational layer (Ours).
Figure I: Computational graph of the 1-th variational layer in a hierarchy. The numbers denote the order of the computations. The inference path forms the posterior g(z) æ, z<1) given the condition c¹- The generative path forms the prior p(z; x<1) given the condition cf- of represents context information that consists of both deterministic features and latent samples with of being a constant, h; represents hidden features of the data with h = æ. Multiple such blocks are
stacked. p( =) receives the sample y and the context cy of the last layer. The symbol denotes a module responsible for combining two streams of features. In Figure La the layer is connected only with the adjacent layers in the hierarchy. Latent information of earher layers #<1-1 carried through c- In Figure 1b] the layer is connected with all the layers below (above) at the bottom- up pass (top-down pass). The ResNet transformations are extended produce key & and query feature maps. The generative network queries the layers of the generative network (inference network) above (below) to identify the most relevant conditioning factors of the prior (posterior) according to attention scores computed by the attention modules (see Figure 2).
The
p(z: zez). Similarly, for the inference model: q(=;
221). This is in contrast to other works (Sønderby et al. 2016) that consider statistical dependencies between successive layers only, ie., p(z) 1/-1). Maaløe et al. (2019) also highlight the importance of this modification. The long- range conditional dependencies are implicitly enforced via deterministic features that are mixed with the latent variables and propagated through the hierarchy (see feature cf in Figure Tap.
State-of-the-art models (Vahdat & Kautz 2020 Child 2020) leverage this fully-connected factorization and rely on the in- creased depth to improve performance and deliver results com-
Table log ((x) for varying depth L (bits/dim).
Depth (T)
bitsidim
400%
parable to that of autoregressive models (Salimans et al., 2017). However, by construction, very deep VAE architectures favor only proximate dependencies in the latent space, limiting long- range conditional dependencies between z₁ and #<1-1 depth increases. This means that practice the network may no longer respect the factorization of the variational distributions p(*)
3.26
-6.8
16
2.96
3.2
2.91
11.
1[ Z<)) and q(z x) = II. q(2) x, Z<1), leading to sub-optimal performance.
Table D reports the absolute and relative decrease in the negative log-likelihood (in bits per dimen- sion) as one increases the number of stochastic layers in an NVAE (Vahdat & Kautz 2020 We observe that the predictive gains diminish as depth increases. We hypothesize that this may be be- cause the effect the latent variables of earlier layers diminishes as the context feature of traverses the hierarchy and is updated with latent information from subsequent layers.
3

---

Published as a conference paper at ICLR 2022
In this work, we improve the flexibility of the prior p(z) and posterior g(> x) by designing more informative representations for the conditioning factors of the conditional distributions p(z) x<1) and [ I, z<1). We do this by designing a hierarchy of densely connected stochastic layers that learn to attend to latent and observed information most critical to inference. Figure [b a graphical illustration of the proposed model.
2.3 DEPTH- WISE ATTENTION
We first introduce depth-wise attention. This the technical tool that allows our model to realize the strong couplings presented in Figure
Figure 2: Attend(c<z,, S1, k<i)- depth-wise
attention block. Cct, kei are the sequences of I - 1 contexts and corresponding keys with C and Q feature maps accordingly, while 8; is the query feature. The multiplication is applied to the inner matrix dimensions. The normalization of the softmax is applied to the last dimension, treating each pixel independently from the others. act =
lb and motivated in Section 2.2 The problem can be formulated as follows:
Given a sequence C21 = (Cm)im 1 contexts, need to construct a feature én that summarizes the information in Cea that is most critical for a given task. é and Cm are features of the same dimensionality: &g E Rx and Cm. E HxxXC.
1/-1 are the attention scores.
è: H. W, C
I=1-1
In our framework, this task is the construction of either prior (Section 2.4] posterior (Sec- tion 2.5 beliefs of a variational layer in deep VAE. Therefore, our architecture must be able to handle long context sequences of large di- mensions H and W.
0eet : [H, W.Y
softmax
[H, 1,
[1, W. 1, QQ
The task is characterized by query feature 8) E R Hxx xQ of dimensionality Q with Q < C. Similarly, Cm is represented by a key feature km E H xxXQ To reduce computational re- quirements, treat each pixel independently from the rest. This can equivalently be interpreted as concurrent processing of H x W indepen- dent sequences of C-dimensional features. The feature maps ē)(i,) E R for pixels (i, with 1 < i S H, 1 j & W, are computed in parallel. (i. depends only on pixel instances in Cet at the same location (i,j), i.e, Cm (i, with m < and is given by:
GCl: : HH W. C.1l
#1 H. W. Q]
Re H. W. Q.
(1,5)) = Σ am am-- (i,j),
(3)
m<l
exp(s] ( (1,
am.) Emer exp(s} of km
(4)
In words, feature si(j.) E Rº queries the significance of feature (i, E RC, repre- sented by ji, E R4, to form ((i,)) R am--r(i,j) E IR the resulting relevance metric of the m-th term, with m < L. at pixel (i.)). The overall procedure is denoted as = Attend(Caz 8t, ka).. and is illustrated in Figure
Finally, to improve training of models with very long sequences, we adopt the following variant of the normalization scheme proposed by Chen et al. (2020):
Cel t- Cet + Gelu(TayerNorm((zz)).
(5)
è + è + Gelu(LayerNormm(())
(6)
where Gelu is the GELU non-linearity (Hendrycks & Gimpel 2016) and LayerNorm a layer nor- malization operation (Ba et al. 2016).
2.4 GENERATIVE MODEL
As shown in Figure la the conditioning factor of the prior distribution at variational layer is represented by context feature of E RxC A convolution is applied on ef obtain parameters
4

---

Published as a conference paper at ICLR 2022
0 defining the prior. cj is a non-linear transformation of the immediately previous latent information -1-1 and prior context C2t-1 containing latent information from distant layers <<|_1, such that cf Tp ( + 2. (-) is typically implemented as a cascade of ResNet cells (He et al. 2016a b) and corresponds to the blue ResNet module in Figure Ta The operator I combines information from two branches in the network (e.g. by summation concatenation). 2t-1, G/_ passed in from the previous layer.
Due to the locality of the architecture, the signal from of may be dominated by that of 21-1. To pre- vent this, we adopt direct connections between each pair of stochastic layers, as shown in Figure That is, variational layer has direct access to prior context of all previous layers c²t accompanied by keys k This means each variational layer can actively determine the most important latent contexts when evaluating its prior beliefs. During training, the context, query, and key are jointly learned:
c sf, + 1
We initially let variational layer rely on nearby dependencies captured by ci. During training, the prior is progressively updated with the holistic context e. via residual connection (Wang et al. 2018 Bachlechner et al. 202I):
2 t- Atttond(ce sy k a
(8)
a of +
(9)
where of E R is a learnable scalar parameter initialized to zero, di2 fch 1--1 with c
E
R xxXxXC xc of E R2xxxx k2 {ke ( 11-1 with ky E # and Q < C.
2.5 INFERENCE MODEL.
As shown in Figure Ia the conditioning context of of the posterior distribution results from com- bining deterministic factor h, and stochastic factor cf provided by the decoder: of = h4 @ cf. To improve inference, we let layer's I encoder use both its own h4 and all subsequent hidden represen- tations h>1, as shown in Figure 1b This modification can also be viewed as skip connections which selectively connect each layer with representations closer to the data helping deep signal propa- gation. As in the generative model, the bottom-up path is extended to emit low dimensional key features ky which represent hidden features hy:
[h,, k:]1 + (hi41 @ k(L1)-
(10)
Prior works (Sønderby al. 2016 Vahdat Kautz, 2020 have sought to mitigate against exploding DKL in Equation by using parametric coordination between the prior and posterior distributions.
Motivated by this insight, we seek establish further communication between them. We accomplish
this by allowing the generative model to choose the most explanatory features in hz by generating query feature s/. Finally, the holistic conditioning factor for the posterior is:
c Attend(hzi, S7, 1.))-
(11)
The detailed description is deferred to Appendix A.1 Inference the model is also summarized in Appendix A.3
2.6 SPATIALLY ATTENTIVE VARIATIONAL LAYERS
In Sections 2.4 and 2.5 we allowed the perceptive field of each pixel at layer / encompass only pixels at the same locations in earlier layers during inference. In order to fully capture distant depen- dencies in the latent space, we still need coordinate portions of information in the spatial domain. To increase the computational efficiency of our model, this occurs at second stage. Conceptually, this two-step procedure resembles the factorized attention of efficient transformers (Child et al.) 2019). albeit deep variational models are inherently amenable to a different, inter-layer and intra- layer, decomposition. This observation decreases the time and memory complexity of the long-range operations for computing the attention scores from O(L3 x W H²²) to O(L + L x W x If²), while still providing inference with holistic information. To accomplish this, we interleave spatially non-local blocks (Wang et al.] 2018) with the convolutions in the residual cells of the ResNet mod- ules in Figure as shown in Figure 3 in Appendix. The exact formulation is similar to that of Wang et al. (2018) and provided in Appendix A.2
5

---

Published as a conference paper at ICLR 2022
3 EXPERIMENTAL STUDIES
We conduct three series of experiments. In the first set of experiments (Section attentive variational path proposed in this paper on VAEs that are trained on two datasets of binary images: the dynamically binarized MNIST and OMNIGLOT. In Section B.1.2] we investigate the effectiveness of the proposed techniques on large-scale latent spaces that are used for generating the 10FAR-10 natural images. Qualitative results are provided in Appendices F (plot of KI. divergence per layer). (visualization of attention patterns), and[G] (novel samples and image reconstructions) Finally, Section 3.2 we conduct an ablation study and report the benefits of each proposed at- tention module separately. All hyperparameters for experiments are available in Appendix [B] We note here that we choose to not use spectral regularization, and we find this does not compromise training stability of our model (in contrast to results reported by Vahdat & Kautz (2020). The fact that our model is able to achieve better performance without this regularization may indicate that the attention operations help stabilize training.
we apply the
3.1 MAIN QUANTITATIVE RESULTS
3.1.1 APPLICATION ON BINARY IMAGES
Datasets. We evaluate the models on two benchmark datasets: MNIST (LeCun et al.) 1998). dataset of 28 x 28 images of handwritten digits, and OMNIGLOT (Lake et al. (013)) an alphabet recognition dataset of 28 28 images. For convenience, add two pixels each border of the training images. In both cases, the observations are dynamically binarized by being resampled from the normalized real values using Bernoulli distribution after each epoch, as suggested by [urdda jet al. (2016). which prevents over-fitting We use the standard splits of MNIST into 60,000 training and 10,000 test examples, and of OMNIGLOT into 24,345 training and 8,070 test examples.
Set-up. For both datasets, we use a hierarchy of = 15 variational layers. We use a Bernoulli distribution in the image decoder. In the generative model non-local block is inserted the end of the module responsible for combining the context with the latent sample provided by the previous layer. In the inference layer, a non-local block is inserted the end of the module responsible for concatenating the deterministic, data-dependent features of the bottom-up pass and the context containing latent information that provided by the generative model. Non-local blocks are also inserted at the post-processing cells, at the end of the variational hierarchy and right before the layers that implement the data distribution (see Appendix [B] for details). MNIST images are downsampled twice before being passed to the hierarchy leading to latent spaces of spatial dimension 8 x 8 while the OMNIGLOT images are downsampled only once resulting to 16 x 16 latent spaces.
Results. Table 2 reports the estimated marginal likelihood of our model along with the performance achieved by state-of-the-art models. We observe that in both datasets, holistic variational inference consistently improves performance. More importantly, our model outperforms models that increase the flexibility of the variational posterior by making use normalizing flows (Kingma et al. 2016) or the data distribution by making use of autoregressive decoders (Chen et al., (017). The use of autoregressive layers is a model refinement orthogonal to ours and could potentially further improve the results presented in Table The architecture elosest ours is NVAE (Vahdat & Kautz, 2020), and the performance gap can be solely attributed to long-range inference operations described
in Sections 2.4 2.5 and 2.6
3.1.2 APPLICATION ON NATURAL IMAGES
Dataset. CIFAR-10 is a dataset of 32 x 32 natural images. The raw pixel values are first scaled to [-,
Set-up. For CIFAR-10, we use hierarchy of L 16 variational layers. We use mixture of dis- cretized Logistic distributions (Salimans et al., 2017) for the data distribution. We use the spatially non-local residual cells of Figure 3 in both the generative and inference network of each layer. Non- local blocks are also inserted at the pre-procsssin blocks, at the beginning of the bottom-up pass (see Appendix B for details). CIFAR-10 images downsampled only once resulting 16 x 16 latent spaces. We note that it helps optimization we bound the log of the prior standard deviation such that log op 2 11.0, yielding less confident prior assumptions. We also empirically find that

---

Published as a conference paper at ICLR 2022
Table Dynamically binarized MNIST (Burda al. 2016 and OMNIGLOT Lake 2013) performance on the test set. All models except for IWAE are trained with single importance sample. IWAE is trained with 50 importance samples. The marginal loglikelihood is estimated with 500 importance samples. Attentive VAE outperforms all state-of-the-art VAEs with without autoregressive components.
Model
log p(a) 2
MNIST OMNIGLOT
Attentive VAE (ours)
77.63
89.50
NVAE Vahdat & Kautz 2020
-78.01
99.188
MAE (Ma et al. 2019
-77.98
BIVA (Maaloc of al. 2019
-78.41
-93.54
PixeIVAE++ Sadeghi et 2019
-78.00
DVAE++ Vahdal el al. 2018
78.49
VampPrior Tomczak & Welling 2018
-8.45
89.76
SA- (Kim al. 2018
90.05
Lossy VAI
2017
-78.33
89.83
IAF- VAE (Kingma et al.
Ladder
2016
-81.74
102.11
2016
-79.10
DVAE (Rolfe|201) (011
-81.01
97.43
Conv DRAW (Grgor et al. 2016
-91.00
IWAE (Burda cl al. 2016
82.90
-00.38
adding Gaussian noise with Onosse 0.001 in both the log of the prior scale log op and the posterior scale log og helps network's 's generalization.
Results. Table shows the performance of the proposed model both among other VAEs and other purely generative models. As we can see, attention-guic variational inference is conducive to building more informative latent spaces, rendering the model the best performing in its class. This claim is further corroborated by qualitative results in Appendix E We can see there that the attention- driven skip connections mitigate posterior collapse by activating early layers in the hierarchy. More- over, our model closes the performance gap between VAEs and expensive generative models that rely on aully-autoregressive distributions. Finally, it is possible our results in Table B could be im- proved by considering deeper hierarchies.
Table 3: CIFAR-10 (Krizhessky et al., 2009 performance on the test set. The marginal log- likelihood estimated with 100 importance samples. A shallower Attentive VAE outperforms all state-of-the-art VAEs with or without autoregressive components. Attentive VAE performs on par with fully autoregressive generative models. However, it permits fast sampling that requires single network evaluation per sample as opposed to D, where D the dimension of the data distribution.
Model
VAF
Autoregressive Decollrr
log, (hlis/dim)
Attentive VAE (ours) trained for 40ll epuchs
16
2.82
Attentive VAE (ours) trained For 500 epochs
2.81
Attentive (ours)
epuchs
Very Deep VAE Child 2020
45
X
2.87
NVAE Vahdar & Kaurz 2010
2.91
BIVA Maniee al 2019
15
X
3.08
TAR
V.I: Kilggn et 2016
12
3.11
el
PixelVAPP (halegh al 2019
2.00
Lossy VAL Chea
2.95
MAF
2019
2.95
PixelCNNFF limin
2017
X
2.92
PixcISNAIL Chen
2.85
Image Transfirmer (Pamar a.
2018 X
2.00
Sparse Transformel mild 2019
2,80
In Table compare the computing demands for training our model with the one of the other two most competitive VAEs presented in Table As we see, our 16-layer VAE requires almost 3/5 of the GPU hours needed for training a 30-layer NVAE, while still outperforming the more efficient VAE of Child (2020). In this work, we implement full rank matrices for the self-attention scores of the spatially non-local residual cells. Efficient attention approximations such those proposed by Child let al. (2019) and Choromanski et al.] (2021) could further reduce the computational requirements of models and are left future work. This suggestion is also bolstered by our findings in Appendix F where we observe that the spatial attention maps are sparse and highly structured.

---

Published as a conference paper at ICLR 2022
Table Comparison of the computational requirements for training deep state-of-the-art VAE models. All models are trained on 32GB V100 GPUs. The additional cost for computing the attention scores compensated by the smaller number of stochastic layers in the hierarchy without sacrificing the generative capacity of the model, see Table 3
Model
batch size / GPU GPUs
Training Time
Total GPU hours
Attentive VAE (ours), 400 epochs
hours
272
Attentive VAF (ours), 500 epochs
32
hours
336
Attentive VAE (ours). 900 epochs
hours
608
NVAE
hours
440
Very Deep VAE
days
288
3.2 ABLATION STUDY
To verify the effectiveness on the performance of the different attention schemes proposed in Sec- 2 we conduct comprehensive ablation studies on architectures of variable number of stochastic layers on the CIFAR-10 dataset. We incrementally add different combinations of the proposed at- tention components of Sections 2.4 2.5 2.6 the NVAE architecture and we evaluate their impact on the estimated marginal likelihood reported in bits per dimension. Our results are shown in Table We make the following observations:
2.4
5
Regardless of the number of stochastic layers in the hierarchy, we observe improvement over the baseline model (Case Case and Case 9 Case 16 in Table
We notice that 8-layer architecture that employs attention mechanism in both the gener- ative network and the inference counterpart the model (Case in Table almost reaches the performance level of the baseline 16-laye 16-layer architecture (Case in Table [5) This fact indicates that our model results in a better utilization the latent space. This is orthogonal other works which aim at improvement by stabilizing deeper architectures instead.
The effect of the spatially attentive stochastic layers is larger in the shallower architec- ture (Case vs Case 2 and Case 9 vs Case 10 in Table 5). This may be because some loss of the latent information compensated by additional layers. However, the total im- provement yielded by the depth-wise attention modules is larger in the deeper architecture, highlighting the importance of the inter-layer connectivity when constructing the prior and the posterior of each layer. The inter-layer connectivity is also visualized in Figures [0] II Appendix FF where confirmed that the model learns to attend to context information that lies in distant layers of the model improve performance.
Finally, supplementary ablations other architectural choices of our model are provided in Ap- pendix D
Table 5: A comparison of the effect of the attention operations in a deep variational model on the CIFAR-10 dataset. The NVAE (Vahdat & Kautz, 2020) is used as baseline (no attention operation).
Depth (2)
Case
Depth wase generative Section 2.4
Depth wise inference Scation 2.5
Non local residual cells Section E.6
# Parameters (x (x.10F)
log p(x) (bits/dim)
3,002
44.585
2.998
44.716
3.019
46.3
3.034
X
52.151
2.985
51.925
2.388
2.053
73.200
2335
10.
89.602
2.912
90.916
2,918
16
12.
102.744
2,872
13.
X
92.965
2.92
2.005
15.
105.541
2,361
16.
118.096
2,823
8

---

Published as a conference paper at ICLR 2022
4 RELATED WORK AND DISCUSSION
Our work builds on recent deep variational autoencoders (Vahdat & Kautz 2020 Child 2020
which first perform a fully deterministic bottom-up pass followed by a top-down pass involving both stochastic and deterministic features, as first proposed by KKngmaa et al. (016). This hybrid top-down pass enables fully connected factorizations of the joint prior and posterior distributions, according to which the latent variables of each layer are conditioned on the latent variables of all the layers above in the hierarchy. However. modeling such conditional dependencies relies on local aarcitecctress which assume connectivity only between successive layers. We improve prior work by introducing attention-guide guideed mechanisms to actively discover long-range statistical dependencies in the latent space. We build on top of NVAE (Vahdat & Kautz] (020) to demonstrate our method, but our suggestions are generic and could be applied to deeper and thinner hierarchies such as that of Child 2020), the BIVA hierarchy which extends the deterministic bottom-up pass to latent variables Maalde et al.) 2019). These configurations could further improve the results reported in Tables 2 Moreover, the proposed model is amenable other tools for improved variational inference such as the use of an autoregressive posterior (or prior) distribution (Kingma et al. 2016), second-stage ancenstral sampling by an auxiliary VAE that learns the aggregated posterior (Dai & Wipi 2019). or regularization techniques that promote congruence between latent representations of the original and transformed inputs (Sinha & Dieng 2021).
Vector Quantized Variational Autoencoders (VQ-VAEs) (Oord et al., 2017 Razavi et al., 2019b)
learn discrete latent representations. However, the training objective of these models does not for- mally constitute a lower bound on the marginal log- logllikelihood. Although VQ-VAEs have exhibited high-quality generative capability on high-resolution images, they base their performance on fully autoregressive PixelCNN priors (Van Den Oord et al., 2016) which render them too slow to sample from in case of large latent spaces. Most related to ours is the work of Esser et al. (202I). The authors use fully autoregressive transformer in the latent space. However, they manage to handle low-dimensional latent spaces, leading to logllikelsooods infererior to that of Razavi et al. (2019b). Scaling this model to larger latent spaces, by leveraging the decomposed, attention-guided latent hierarchy proposed in this work could be an interesting future research direction.
Several works generative modeling leverage self-attention mechanisms (Parmar al. 2018 Chen let al. 2018 2020; Child et al. 2019). However, these models are autoregressive, yielding condi- tional per-pixel distributions, and apply attention in the pixel space. Therefore, as already discussed, they are slow to sample from. Our work has log-likelihood performance comparable to that of these works, while using aotrautoreeresvvve data distribution. The work of Tulsiani & Guptal (202T) uses wransformer-based encoder to learn pixel distributions conditioned on small set of observed pixel values at arbitrary locations. Diffusion models (Ho et al.) [019 Kingma al., (20TT) can be viewed as VAEs and achieve impressive generative quality. However, similarly to the autoregressive models, sampling from them requires multiple network evaluations. Morever, their performance when trained with data augmentation (Kingma et al., (02T) can be similar to that of deep regu- larized VAEs (Sinha & Dieng] 202I). Finally, the work by Zhang et al.] (2019) boosts GANs by incorporating non-local blocks (Wang et al. 2018) in the generator and the discriminator.
Future research includes investigating the effect of the proposed depth-wise attention mechanism on ResNet architectures for different tasks and efficient attention approximations specific to deep variational inference.
5 CONCLUSION
We have presented deep attentive VAE, the first attention-drive framework for variational inference in deep probabilistic models. We argue that the expressivity of current deep probabilistic models can be improved by selectively prioritizing statistical dependencies between latent variables that are potentially distant from each other. We introduced a scalable factorized attention scheme that first discovers intra-layer correlations, succeeded by inter-layer attention operations. Moreover, the gen- crative model is allowed to query its posterior counterpart based on the available latent information to identify the dete-deeendent representations most informative for inference. We extensively eval- uated the proposed architecture on multiple public datasets and showed that our model outperforms all existing deep VAEs, while requiring significantly less training time.

---

Published as a conference paper at ICLR 2022
To aid reproducibility of the results and methods presented in paper, made code to reproduce the main results the paper publicly available, including detailed instructions; see our github page: https: / /tggttthusss com/ wtt.:faapotoooooo////ooooaoooo..
Attentive VI
ACKNOWLEDGEMENTS
We thank anonymous reviewers for their careful reading our manuscript and their thoughtful comments and suggestions that helped us improve the quality of the paper. We also thank Christos Louizos for helpful pointers to prior works on VAEs, Katerina Fragkiadaki for helpful discussions on generative models and attention mechanisms for computer vision tasks, Andrej Risteski for in- sightful conversations on approximate inference, and Jeremy Cohen for his remarks on a late draft of this work. We are also very grateful to Radium Cloud for granting access computing infrastruc-- ture that enabled us to scale up our experiments. This material is based upon work supported by the Defense Advanced Research Projects Agency under award number FA8750-17-244 17-20-7----13, and by the National Science Foundation under grant number 2038612. The first author acknowledges support from the Alexander Onassis Foundation and from A. G. Leventis Foundation. The second author is supported by the National Science Foundation Graduate Research Fellowship Program under Grant No. DGEI 6 and 14039, Any opinions, findings, and conclusions recommendations expressed in this material are those of the author(s) and do not necessarily reflect the views of the National Science Foundation.
REFERENCES
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv: 3xXiv:607.06445 (607006550, 2016.
Thomas Bachlechner, Bodhisattwa Prasad Majumder, Huanru Henry Mao, Garrison W Cottrell, and Julian McAuley. Rezero is all you need: Fast convergence at large depth. In Conference Uncertainty in Artificial Intelliencee (UAI), 2021.
Yuri Burda, Roger B. Grosse, and Ruslan Salakhutdinov. Importance weighted autoencoders. In International Conference on Learning Representations (ICLR), 2016.
Mark Chen, Alec Radford, Rewon Child, Jeffrey Wu, Heewoo Jun, David Luan, and Hyaa Sutskever. Generative pretraining from pixels. International Conference on Machine Learning (ICML), 2020.
Xi Chen, Diederik P. Kingma, Tim Salimans, Yan Duan, Prafulla Dhariwal, John Schulman, Ilya Sutskever, and Pieter Abbeel. Variational lossy autoencoder. In International Conference on Learning Representations (ICLR), 2017.
Xi Chen, Nikhil Mishra, Mostafa Rohaninejad, and Pieter Abbeel. Pixelsnail: An improved autore- gressive generative model. International Conference on Machine Learning (ICMI), 2018.
Rewon Child. Very deep vaes generalize autoregressive models and can outperform them on images. In International Conference on Learning Representations (ICLR), 2020.
Rewon Child, Scott Gray, Alec Radford, and llya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904. 10509, 2019.
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. In International Conference on Learning Representations (ICLR), 2021.
Bin Dai and David Wipf. Diagnosing and enhancing vae models. In International Conference on Learning Representations (ICLR), 2019.
Patrick Esser, Robin Rombach, and Bjom Ommer. Taming transformers for high-resolution image synthesis. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recog- nition (CVPR), 2021.
10

---

Published as a conference paper at ICLR 2022
Karol Gregor, Frederic Besse, Danilo Jimenez Rezende, Ivo Danihelka, and Daan Wierstra. Towards conceptual compression. Advances in Neural Information Processing Systems (NeurIPS), 2016.
Ishaan Gulrajani, Kundan Kumar, Faruk Ahmed, Adrien Ali Taiga, Francesco Visin, David Vázquez, and Aaron C. Courville. PixeIVAE: A latent variable model for natural images. In International Conference Learning Representations (ICLR), 2017.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2016a.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Identity mappings in deep residual networks. In European conference on computer vision (ECCV), 2016b.
Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv preprint arXiv:1600: 084 2016.
Jonathan Ho, Ajay Jain, and Picter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems (NeurlPS), 2019.
Matthew D. Hoffman, David Blei, Chong Wang, and John Paisley. Stochastic variational inference. Journal of Machine Learning Research, 14(5), 2013.
Yoon Kim, Sam Wiseman, Andrew Miller, David Sontag, and Alexander Rush. Semi-amortized variational autoencoders. In International Conference Machine Learning (ICML), 2018.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations (ICLR), 2014.
Diederik P. Kingma and Max Welling. Auto-encoding variational Bayes. In International Confer- ence on Learning Representations (ICLR), 2014.
Diederik P. Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, Ilya Sutskever, and Max Welling. Improved variational inference with inverse autoregressive flow. Advances in Neural Informa- fion Processing Systems (NeurIPS), 2016.
Diederik P. Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. arXiv preprint arXiv:2107.00630 2021.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. In Citeseer, 2009.
Brenden M Lake, Ruslan Salakhutdinov, and Joshua B Tenenbaum. One-shot learning by inverting a compositional causal process. In Advances in Neural Information Processing Systems (NeurIPS), 2013.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86, 1998.
Yingzhen Li and Richard E Turner. R\ ''nyi divergence variational inference. In Advances Neural Information Processing Systems (NeurIPS), 2016.
Ilya Loshchilov and Frank Hutter. SGDR: Stochastic gradient descent with warm restarts. In Inter- national Conference Learning Representations (ICLR), 2017.
James Lucas, George Tucker, Roger Grosse, and Mohammad Norouzi. Don't blame the elbo!
linear vae perspective posterior collapse. In Advances in Neural Information Processing Systems (NeurIPS), 2019.
Xuezhe Ma, Chunting Zhou, and Eduard Hovy. MAE: Mutual posterior diverree regularization for variational autoencoders. In International Conference on Learning Representations (ICLR), 2019.
11

---

Published as a conference paper at ICLR 2022
Lars Maalée, Marco Fraccaro, Valentin Liévin, and Ole Winther. BIVA: A very deep hierarchy of latent variables for generative modeling. Advances in Neural Information Processing Systems (NeurIPS), 2019.
Vaden Masrani, Tuan Anh Le, and Frank Wood. The thermodynamic variational objective. In Advances in Neural Information Processing Systems (NeurIPS), 2019.
Aaron van den Oord, Oriol Vinyals, and Koray Kavukcuoglu. Neural discrete representation learn- ing. 2017.
John Paisley, David Blei, and Michael Jordan. Variational bayesian inference with stochastic search. In International Conference on Machine Learning (ICML). 2012.
Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran. Image transformer. In International Conference on Machine Learning (ICML), 2018.
Prajit Ramachandran, Barret Zoph, and Quoe Le. Searching for activation functions. arXiv preprint arXiv:1710.0594 2017.
Ali Razavi, Aäron Van Den Oord, Ben Poole, and Oriol Vinyals. Preventing posterior collapse with vaes.-vaes. In International Conference on Learning Representations (ICLR), 2019a.
Ali Razavi, Aäron Van Den Oord, and Oriol Vinyals. Generating diverse high-fidelity images with vq-vae-2. In Advances in Neural Information Processing Systems (NeurIPS), 2019b.
Danilo Rezende and Shakir Mohamed, Variational inference with normalizing flows Interna- tional conference on machine learning (ICML), 2015.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and ap- proximate inference in deep generative models. In International Conference on Machine Learning (ICML), 2014.
Jason Tyler Rolfe. Discrete variational autoencoders. International Conference on Learning Representations (ICLR), 2017.
Hossein Sadeghi, Evgeny Andriyash, Walter Vinci, Lorenzo Buffoni, and Mohammad H Amin. Pixelvac++: Improved pixelvae with discrete prior. arXiv preprint arXiv:1908. 190,:9998009888, 2019.
Tim Salimans and Diederik P. Kingma. Weight normalization: simple reparameterization to ac- celerate training of deep neural networks. In Advances in Neural Information Processing Systems (NeurIPS), 2016.
Tim Salimans, Andrej Karpathy, Xi Chen, and Diederik P. Kingma. Pixelenn++: Improving the pixelcnn with discretized logistic mixture likelihood and other modifications. In International Conference on Learning Representations (ICLR), 2017.
Samarth Sinha and Adji B Dieng. Consistency regularization for variational auto-encoders. In Advances in Neural Information Processing Systems (NeurIPS), 2021.
Casper Kaae Sønderby, Tapani Raiko, Lars Maaløe, Søren Kaae Sønderby, and Ole Winther. Lad- der variational autoencoders. Advances in Neural Information Processing Systems (NeurIPS), 2016.
Jakub Tomezak and Max Welling, Vae with a vampprior. In International Conference on Artificial Intelligence and Statistics (AISTATS), 2018.
Shubham Tulsiani and Abhinav Gupta. Pixeltransformer: Sample conditioned signal generation. In International Conference Machine Learning (ICML), 2021.
Arash Vahdat and Jan Kautz. NVAE: deep hierarchical variational autoencoder. Advances in Neural Information Processing Systems (NeurIPS), 2020.
12

---

Published as a conference paper at ICLR 2022
Arash Vahdat, William Macready, Zhengbing Bian, Amir Khoshaman, and Evgeny Andriyash. Dvae++: Discrete variational autoencoders with overlapping transformations. In International Conference on Machine Learning (ICML), 2018.
Aäron Van Den Oord, Nal Kalchbrenner, and Koray Kavukcuoglu. Pixel recurrent neural networks. In International Conference on Machine Learning (ICML), 2016.
Xiaolong Wang, Ross Girshick, Abhinav Gupta, and Kaiming He. Non-local neural networks. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2018.
Han Zhang, Ian Goodfellow, Dimitris Metaxas, and Augustus Odena. Self-attention generative adversarial networks. In International Conference on Machine Learning (ICML), 2019.
13


Reviews:
[{"review_id": "ZfKs64kIb3n", "reviewer": "Reviewer_VA9m", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "Pros:\n1. 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.\n2. 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.\n3. The experimental justification with the comparison over different types of VAE models was sufficient.\n4. Computational cost was reduced when compared with the previous SOTA model based on NVAE.\n\nCons:\n1. It is suggested to provide an algorithm to enhance the comprehension of the detailed training procedure of the proposed method.\n2. The experiments did not show the generated samples the authors obtained. The quality of the generated images could not be evaluated.\n3. 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_paper": "This paper presented a variational inference with attention mechanism. A deep latent variable model was proposed.", "summary_of_the_review": "An interesting work was proposed with experimental justification. The experiments can be furthered strengthened.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "Vm1AV6X-454", "reviewer": "Reviewer_DYsA", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "**Relevance**\nDeep 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.\n\n**Novelty**\nMulti-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.\n\n**Clarity**\nOverall 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.\n\nSome 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.\n\n**Empirical evaluation**\nThe 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.\n\n**Other notes and questions**\n* First sentence final paragraph of 2.1: p(x) should be p(z)?\n* 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).\n* 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_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.", "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**.\n\n**Post-rebuttal note**: Increased score to accept.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "JbGr1kZAjhG", "reviewer": "Reviewer_NQd9", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "## 1. Strengths\n- a. The introduction of the layer-wise attention mechanism for deep VAEs is novel, and its effectiveness is supported by empirical results\n- b. The introduction of non-local attention within layers is effective, and supported by experiments\n- c. The overall architecture (incl. normalisation and scaling tricks) challenges the current sota deep VAEs using fewer layers and shorter training. \n- d. Literature in deep VAEs seems to be well known by the authors and sufficiently cited\n\n## 2. Weaknesses\n- a.  The paper lacks structure and clarity \n- b. The paper lacks a more qualitative study of the model:\n    - it would be interesting to see what layers the layer-wise attention mechanism attends to.\n    - 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\").\n- 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).\n\n## 3. Clarification needed: \n- 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.\n- b. Is the layer-wise attention mechanism specific to deep VAEs, or can it be more generally applied to ResNet architectures? \n- c. Section 2.3 (paragraph cited bellow): I get the idea, but unless demonstrated, this remains a hypothesis.\n    >\"... 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\":\n\n## 4. Minor comments / suggestions\n- 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.\n- b. tricks like normalisation or feature scaling could be referenced in a separate section. \n- c. eq8: you might want to cite ReZero [1] here\n- 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\n- e. It would be better practice to report likelihoods for multiple random seeds\n- 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...\"\n\n[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.\n", "summary_of_the_paper": "This paper improves the architecture of deep VAEs using the attention mechanism. \n\nThe attention mechanism is used in two ways:\n1. layer-wise attention (attending stochastic feature maps which are conditioned on other variables within the hierarchy, interpreted as a mixture of skip connections)\n2. non-local attention (attention across the spatial dimensions, increases the size of receptive field)\n\nThe 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. \n\nThe authors provide an extensive ablation study, showing the impact of each type of attention on the training performances (test likelihood on CIFAR-10).", "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. \n\nThe 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). \n\nExperiments 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.\n\n---- edit 26/11\nRaised my score to 8.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "5mxuL8mGLUL", "reviewer": "Reviewer_5Yd5", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "Strength:\n1. 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.\n2. Evaluation showed strong superior performance compared with a wide range of baseline models. \n3. An ablation study is included to justify each components of the proposed model.\n4. The paper is well written and the proposed model is easy to understand.\n\nWeakness:\nThis 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_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.", "summary_of_the_review": "Overall, this is a pretty solid paper. Please refer to the previous section for detailed discussion.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

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)