D EEP A TTENTIVE V ARIATIONAL I NFERENCE
A BSTRACT
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 proximate 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 nearby and distant 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 holistic inference reduces computational footprint by alleviating the need for deep hierarchies.
I NTRODUCTION
A 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 hypotheses 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 interest while enabling tractable inference and sampling.
Latent variable probabilistic models augment the set of the observed variables with auxiliary latent variables.
They are characterized by a posterior distribution over the latent variables, one which is generally intractable and typically approximated by closed-form alternatives.
They provide an explicit parametric specification of the joint distributions 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) (Kingma & Welling, 2014) 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—specifically, it is 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.
; Chen et al.
(2017) aim at modeling more complex posterior distributions, constructed with autoregressive layers.
Theoretical research focuses on deriving tighter bounds (Burda et al., 2016; Li & Turner, 2016; Masrani et al., 2019) or 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 for semantics-preserving transformations of the data.
Taking a different approach, hierarchical VAEs
(Gulrajani et al., 2017; Sønderby 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 subsequent 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 architectures 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.
P ROPOSED M ODEL 2. D EEP V ARIATIONAL I NFERENCE
A latent variable generative model defines the joint distribution of a set of observed variables, x ∈ R D , and auxiliary latent variables, z , 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:
[Equation]
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 ( z | x ) and applying
Jensen’s inequality:
[Equation]
where θ , ϕ parameterize p ( x, z ; θ ) and q ( z | x ; ϕ ) respectively.
For ease of notation, we omit θ , ϕ in the derivations.
This objective is called the Evidence Lower BOund (ELBO) and can be optimized efficiently for continuous z via stochastic gradient descent (Kingma & Welling, 2014;
Rezende et al., 2014).
For ease of implementation, it is common to assume that both q ( z | z ) and p ( z ) are fully factorized
Gaussian distributions.
However, this assumption may be too limiting in cases of complex underlying distributions.
To compensate for this modeling constraint, many works focus on stacking and improving the stability of multiple layers of stochastic latent features which are partitioned in groups such that z
= { z 1 , z 2 , . . .
, z L } , where L is the number of such groups (Rezende et al., 2014; Gulrajani 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 a deep variational model is depicted in Figure 1.
In a bidirectional inference architecture with a deterministic bottom-up pass (left part in Figure 1a), posterior sampling is preceded by a sequence of non-linear transformations,
T ql , of the evidence, x , i.e., h l = T ql ( h l +1 ) , with h L +1 = x .
The inference q ( z | x ) and generative p ( z ) network decompose in an identical topological ordering:
q ( z | x ) = (cid:81)
l q ( z l | x , z <l ) and p ( z ) =
(cid:81) l p ( z l |
z <l ) .
The top-down pass (right part in Figure 1a) generates the posterior samples
z that feed the conditional data distribution p ( x | z ) .
2. M OTIVATION
Kingma et al.
(2016) proposed a strongly connected directed probabilistic graphical model for the generative and inference network, so that each variable depends on all the previous in the hierarchy:
[Figure]
p ( z l | z <l ) .
Similarly, for the inference model: q ( z l | x , z <l ) .
This is in contrast to other works (Sønderby et al., 2016) that consider statistical dependencies between successive layers only, i e., p ( z l | z l − 1 ) .
Maaløe et al.
(2019) also highlight the importance of this modification.
The longrange conditional dependencies are implicitly enforced via deterministic features that are mixed with the latent variables and propagated through the hierarchy (see feature c pl in Figure 1a)
[Table]
− 3.
− 6 . 3.
− 6 . 2.
− 3 . 2.
− 1 .
State-of-the-art models (Vahdat & Kautz, 2020; Child, 2020) leverage this fully-connected factorization and rely on the increased depth to improve performance and deliver results comparable 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 longrange conditional dependencies between z l and z <l − 1 as depth increases.
This means that in practice the network may no longer respect the factorization of the variational distributions p ( z ) =
(cid:81) l p ( z l |
z <l ) and q ( z | x ) = (cid:81) l q ( z l | x , z <l ) , leading to sub-optimal performance.
Table 1 reports the absolute and relative decrease in the negative log-likelihood (in bits per dimension) 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 because the effect of the latent variables of earlier layers diminishes as the context feature c pl traverses the hierarchy and is updated with latent information from subsequent layers.
In this work, we improve the flexibility of the prior p ( z ) and posterior q ( z | x ) by designing more informative representations for the conditioning factors of the conditional distributions p ( z l | z <l ) and q ( z l | x , z <l ) .
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 1b is a graphical illustration of the proposed model.
2. D EPTH -W ISE A TTENTION
[Equation]
[Equation]
We first introduce depth-wise attention.
This is the technical tool that allows our model to realize the strong couplings presented in Figure 1b and motivated in Section 2.2.
The problem can be formulated as follows:
Given a sequence c <
l
= { c m } l − 1 m =1 of l −contexts, we need to construct a feature ˆ c l that summarizes the information in c <
l that is most critical for a given task.
ˆ c l and c m are features of the same dimensionality: ˆ c l ∈ R H × W × C , and c m ∈ R H × W × C .
In our framework, this task is the construction of either prior (Section 2.4) or posterior (Section 2.5) beliefs of a variational layer in a deep
VAE.
Therefore, our architecture must be able to handle long context sequences of large dimensions H and W .
The task is characterized by a query feature s l ∈ R H × W × Q of dimensionality Q with Q ≪
C .
Similarly, c m is represented by a key feature k m ∈
R H × W × Q .
To reduce computational requirements, we treat each pixel independently from the rest.
This can equivalently be interpreted as concurrent processing of H × W independent sequences of C -dimensional features.
The feature maps ˆ c l ( i, j ) ∈ R C for pixels ( i, j ) , with 1 ≤ i ≤ H , 1 ≤ j ≤ W , are computed in parallel.
ˆ c l ( i, j ) depends only on pixel instances in c <
l at the same location ( i, j ) , i.e, c m ( i, j ) with m < l , and is given by:
[Equation]
In words, feature s l ( i, j ) ∈ R Q queries the significance of feature c m ( i, j ) ∈ R C , represented by k m ( i, j ) ∈ R Q , to form ˆ c l ( i, j ) ∈ R C .
α m → l ( i, j ) ∈ R , is the resulting relevance metric of the m − th term, with m < l , at pixel ( i, j ) .
The overall procedure is denoted as ˆ c =
Attend( c <l , s l , k <l ) , and is illustrated in Figure 2.
Finally, to improve training of models with very long sequences, we adopt the following variant of the normalization scheme proposed by Chen et al.
[Equation]
where Gelu is the GELU non-linearity (Hendrycks & Gimpel, 2016) and LayerNorm a layer normalization operation (Ba et al., 2016).
2. G ENERATIVE M ODEL
As shown in Figure 1a, the conditioning factor of the prior distribution at variational layer l is represented by context feature c pl ∈ R H × W × C .
A convolution is applied on c pl to obtain parameters θ defining the prior.
c pl is a non-linear transformation of the immediately previous latent information z l − 1 and prior context c < l − 1 containing latent information from distant layers
z <l − 1 , such that c pl = T pl ( z l − 1 ⊕ c pl − 1 ) .
T pl ( · ) is typically implemented as a cascade of ResNet cells (He et al., 2016a;b) and corresponds to the blue ResNet module in Figure 1a.
The operator ⊕ combines information from two branches in the network (e.g. by summation or concatenation).
z l − 1 , c pl − 1 are passed in from the previous layer.
Due to the locality of the architecture, the signal from c pl may be dominated by that of z l − 1 .
To prevent this, we adopt direct connections between each pair of stochastic layers, as shown in Figure 1b.
That is, variational layer l has direct access to prior context of all previous layers c p<l accompanied by keys k p<l .
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:
[Equation]
We initially let variational layer l rely on nearby dependencies captured by c pl .
During training, the prior is progressively updated with the holistic context ˆ c pl via a residual connection (Wang et al., 2018;
Bachlechner et al., 2021):
[Equation]
[Equation] 9) where γ pl ∈ R is a learnable scalar parameter initialized to zero, c p<l = { c pm } l − 1 m =1 with c pm ∈
R H × W × C , s pl ∈ R H × W × Q , k p<
l = { k pm } l − 1 m =1 with k pm ∈ R H × W × Q , and Q ≪ C .
2. I NFERENCE M ODEL
As shown in Figure 1a, the conditioning context c ql of the posterior distribution results from combining deterministic factor h l and stochastic factor c pl provided by the decoder: c ql = h l ⊕ c pl .
To improve inference, we let layer’s l encoder use both its own h l and all subsequent hidden representations h >l , 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 propagation.
As in the generative model, the bottom-up path is extended to emit low dimensional key features k ql which represent hidden features h l :
[Equation]
Prior works (Sønderby et al., 2016; Vahdat & Kautz, 2020) have sought to mitigate against exploding
D KL in Equation 2 by using parametric coordination between the prior and posterior distributions.
Motivated by this insight, we seek to establish further communication between them.
We accomplish this by allowing the generative model to choose the most explanatory features in h ≥ l by generating query feature s ql .
Finally, the holistic conditioning factor for the posterior is:
[Equation]
The detailed description is deferred to Appendix A.1.
Inference on the model is also summarized in
Appendix A.3.
2. S PATIALLY ATTENTIVE VARIATIONAL LAYERS
In Sections 2.4 and 2.5, we allowed the perceptive field of each pixel at layer l to encompass only pixels at the same locations in earlier layers during inference.
In order to fully capture distant dependencies in the latent space, we still need to coordinate portions of information in the spatial domain.
To increase the computational efficiency of our model, this occurs at a 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 intralayer, decomposition.
This observation decreases the time and memory complexity of the long-range operations for computing the attention scores from O ( L 3 × W 2 × H 2 ) to O
( L 2 + L × W 2 × H 2 ) , 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 modules in Figure 1, as shown in Figure 3 in Appendix.
The exact formulation is similar to that of Wang et al.
and is provided in Appendix A.2.
E XPERIMENTAL S TUDIES
We conduct three series of experiments.
In the first set of experiments (Section 3.1.1), we apply the 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 3.1.2, we investigate the effectiveness of the proposed techniques on large-scale latent spaces that are used for generating the
CIFAR-10 natural images.
Qualitative results are provided in Appendices E (plot of KL divergence per layer), F (visualization of attention patterns), and G (novel samples and image reconstructions).
Finally, in Section 3.2 we conduct an ablation study and report the benefits of each proposed attention 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.
3. M AIN Q UANTITATIVE R ESULTS 3.1. A PPLICATION ON BINARY IMAGES
Datasets.
We evaluate the models on two benchmark datasets: MNIST (LeCun et al., 1998), a dataset of 28 × 28 images of handwritten digits, and OMNIGLOT (Lake et al., 2013), an alphabet recognition dataset of 28 × 28 images.
For convenience, we add two zero pixels to each border of the training images.
In both cases, the observations are dynamically binarized by being resampled from the normalized real values using a Bernoulli distribution after each epoch, as suggested by Burda et 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 L = 15 variational layers.
We use a Bernoulli distribution in the image decoder.
In the generative model a non-local block is inserted at 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 at 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 is 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 × 8 while the OMNIGLOT images are downsampled only once resulting to 16 × 16 latent spaces.
Results.
Table 10 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 of normalizing flows (Kingma et al., 2016) or the data distribution by making use of autoregressive decoders (Chen et al., 2017).
The use of autoregressive layers is a model refinement orthogonal to ours and it could potentially further improve the results presented in Table 10.
The architecture closest to 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. A PPLICATION ON NATURAL IMAGES
Dataset.
CIFAR-10 is a dataset of 32 × 32 natural images.
The raw pixel values are first scaled to [ − 1 , 1] .
Set-up.
For CIFAR-10, we use a hierarchy of L = 16 variational layers.
We use a mixture of discretized 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.
Nonlocal blocks are also inserted at the pre-processing blocks, at the beginning of the bottom-up pass (see Appendix B for details).
CIFAR-10 images are downsampled only once resulting to 16 ×latent spaces.
We note that it helps optimization if we bound the log of the prior standard deviation such that log σ p ≥ − 1 . 0 , yielding less confident prior assumptions.
We also empirically find that
[Table]
adding Gaussian noise with σ noise = 0 .
001 in both the log of the prior scale log σ p and the posterior scale log σ q helps network’s generalization
[Table]
Results.
Table 3 shows the performance of the proposed model both among other VAEs and other purely generative models.
As we can see, attention-guided 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 attentiondriven skip connections mitigate posterior collapse by activating early layers in the hierarchy.
Moreover, our model closes the performance gap between VAEs and expensive generative models that rely on fully-autoregressive distributions.
Finally, it is possible our results in Table 3 could be improved by considering deeper hierarchies.
[Table]
In Table 4, we compare the computing demands for training our model with the one of the other two most competitive VAEs presented in Table 3.
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 as those proposed by Child et al.
and Choromanski et al.
(2021) could further reduce the computational requirements of our models and are left as 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.
[Table]
3. A BLATION S TUDY
To verify the effectiveness on the performance of the different attention schemes proposed in Section
[Table]
we conduct comprehensive ablation studies on architectures of a variable number of stochastic layers on the CIFAR-10 dataset.
We incrementally add different combinations of the proposed attention components of Sections 2.4, 2.5, 2.6 in the NVAE architecture and we evaluate their impact on the estimated marginal likelihood reported in bits per dimension.
Our results are shown in Table 5.
We make the following observations:
• Regardless of the number of stochastic layers in the hierarchy, we observe improvement over the baseline model (Case 1 vs Case 8 and Case 9 vs Case 16 in Table 5).
• We notice that a 8-layer architecture that employs an attention mechanism in both the generative network and the inference counterpart of the model (Case 8 in Table 5) almost reaches the performance level of the baseline 16-layer architecture (Case 9 in Table 5).
This fact indicates that our model results in a better utilization of the latent space.
This is orthogonal to other works which aim at improvement by stabilizing deeper architectures instead.
• The effect of the spatially attentive stochastic layers is larger in the shallower architecture (Case 1 vs Case 2 and Case 9 vs Case 10 in Table 5).
This may be because some
loss of the latent information is compensated by additional layers.
However, the total improvement 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 10,in Appendix F, where it is confirmed that the model learns to attend to context information that lies in distant layers of the model to improve performance.
Finally, supplementary ablations on other architectural choices of our model are provided in Appendix D.
[Table]
R ELATED W ORK 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 Kingma et al.
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 architectures which assume connectivity only between successive layers.
We improve on prior work by introducing attention-guided mechanisms to actively discover long-range statistical dependencies in the latent space.
We build on top of NVAE (Vahdat & Kautz, 2020) to demonstrate our method, but our suggestions are generic and could be applied to deeper and thinner hierarchies such as that of Child (2020), or the BIVA hierarchy which extends the deterministic bottom-up pass to latent variables (Maaløe et al., 2019).
These configurations could further improve the results reported in
Tables 10, 3.
Moreover, the proposed model is amenable to other tools for improved variational inference such as the use of an autoregressive posterior (or prior) distribution (Kingma et al., 2016), a second-stage ancenstral sampling by an auxiliary VAE that learns the aggregated posterior (Dai &
Wipf, 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 formally constitute a lower bound on the marginal log-likelihood.
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.
The authors use a fully autoregressive transformer in the latent space.
However, they manage to handle low-dimensional latent spaces, leading to log-likelihoods 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 on generative modeling leverage self-attention mechanisms (Parmar et al., 2018; Chen et al., 2018; 2020; Child et al., 2019).
However, these models are autoregressive, yielding conditional 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 a non-autoregressive data distribution.
The work of Tulsiani & Gupta (2021) uses a transformer-based encoder to learn pixel distributions conditioned on a small set of observed pixel values at arbitrary locations.
Diffusion models (Ho et al., 2019; Kingma et al., 2021) 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., 2021) can be similar to that of deep regularized VAEs (Sinha & Dieng, 2021).
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.
C ONCLUSION
We have presented deep attentive VAE, the first attention-driven 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 generative model is allowed to query its posterior counterpart based on the available latent information to identify the data-dependent representations most informative for inference.
We extensively evaluated the proposed architecture on multiple public datasets and showed that our model outperforms all existing deep VAEs, while requiring significantly less training time.
A CKNOWLEDGEMENTS
We thank anonymous reviewers for their careful reading of 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 insightful 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 us access to computing infrastructure 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-2- and by the National Science Foundation under grant number
D EEP
A
TTENTIVE V
ARIATIONAL
I
NFERENCE
A M ODEL A RCHITECTURE A. I NFERENCE N ETWORK
In this section, we provide the formal description of the attention-enhanced inference network presented in Section 2.5.
Equivalently, we present the procedure for the construction of the conditioning factor of the posterior distribution q ( z l | x , z <l ) .
The inference network is first traversed from the bottom (stochastic layer closer to the data x ) to the top, to compute successive hidden data representations h l ∈ R H × W × C and their keys:
k ql ∈
R H × W × Q , with Q ≪ C :
[Equation]
T ql ( h l +
[Equation]
k ql +
[Equation]
) , for l = L, L − 1 , . .
L is the number of stochastic layers.
We assume h L +1 ≡ x .
T ql is a non-linear transformation, typically implemented as a cascade of ResNet cells (He et al., 2016a;b).
Afterwards, samples from the posterior are drawn in top-down order following the topological ordering of the generative model:
[Equation]
The inference network receives the prior context c pl ∈ R H × W × C and a query feature s ql ∈
R H × W × Q from the generative model, see Equation 7.
c
pl and s ql contain latent information of previous layers z <l .
The observed, data-dependent feature h l and the latent feature c pl are combined via an operator ⊕ to give the context of the posterior c ql ∈ R H × W × C along with its key k p,ql ∈ R H × W × Q :
[Equation]
As already pointed out, k p,ql and c ql contain stochastic, latent information, while k ql and h l of Equation 12 contain only deterministic, observed information.
We allow variational layer l to have access to h l and all subsequent data representations h >l .
During inference, the variational layer decides parts of the deterministic features h ≥ l and the stochastic feature
c ql to pay attention to.
For pixel ( i, j ) , with 1 ≤ i ≤ H , 1 ≤ j ≤ W , the generative model queries according to s ql ( i, j ) ∈ R Q and computes an attention distribution, as follows:
[Equation]
α qp → l ( i, j ) ∈ R reflects the importance of stochastic context c ql ( i, j ) ∈ R C to pixel ( i, j ) .
Similarly, α qm → l ( i, j ) ∈ R reflects the importance of deterministic feature h m
( i, j ) ∈ R C of layer m to pixel ( i, j ) .
Z q ( i, j ) ∈ R is a normalization constant.
The posterior context for pixel ( i, j ) , ˆ c ql ( i, j ) ∈ R C , is the weighted sum:
[Equation]
To scale the deep variational model with ease, we use the following normalization:
[Equation]
19) where Gelu is the GELU non-linearity (Hendrycks & Gimpel, 2016) and LayerNorm a layer normalization operation (Ba et al., 2016).
At the beginning of the training, we let the inference model rely on nearby latent and observed information represented by c ql , by introducing the residual connection:
[Equation]
where γ ql is a learnable scalar parameter initialized to zero.
We adopt the Gaussian residual parametrization between the prior and the posterior proposed by
Vahdat & Kautz (2020).
The prior is given by:
[Equation]
[Equation]
The posterior is then given by:
[Equation]
where the summation and multiplication are pointwise, and c ql is defined in Equation 20.
µ ( · ) , σ ( · ) , ∆ µ ( · ) , and ∆ σ ( · ) are transformations implemented as convolutions.
The inference procedure is also described in detail in Algorithm 3.
A. S PATIALLY N ON -L OCAL V ARIATIONAL L AYERS
In this section, we provide the formal description of the non-local blocks (Wang et al., 2018) in the recidual cells of encoder and decoder.
As shown in Figure 3 and explained below, the output of each block is scaled and added back to the input feature.
Let c ∈ R H
× W × C a context feature.
We treat c as a sequence of N , where N = H × W , C − dimensional features, such that c ∈ R N × C .
The non-local block maps c to a feature v ∈ R N × C of the same dimension by taking into account the response of all the locations in the sequence.
The context feature c is first mapped to two feature spaces of a lower dimension Q , with Q << C , through a key α and a query β linear transformation:
[Equation]
where the linear transformations W α ∈ R C × Q , W β ∈ R C × Q are implemented as 1 × 1 convolutions, and α ( c ) ∈ R N × Q , β ( c ) ∈ R N × Q .
The energy matrix E ( c ) ∈ R N × N is computed as:
[Equation]
[Equation]
Each element E i,j ( c ) =
β i ( c ) × α Tj ( c ) , where β i is the
i − th row of β , reflects the relevance of pixel j to pixel i .
The attention distribution is computed by a scaled softmax row normalization, such that the corresponding attention score A i,j ( c ) is : [Equation] [Equation]
[Equation]
ω ∈ R is learnable scalar parameter initialized to one.
c is also linearly projected to an embedding space via a transformation η :
[Equation]
where W η ∈ R C × C and η ( c ) ∈ R N × C .
W η is implemented as a 1 × 1 convolution.
The final output v ∈ R N × C is:
[Equation]
A ( c ) ∈ R N × N is the matrix of the attention scores defined in Equation 27.
Finally, v represents a scaled, residual mapping of c such that the final output is:
[Equation]
where γ ∈ R is a learnable scalar parameter initialized to zero.
A. A LGORITHMS
Algorithms 1, 2, and 3 summarize the inference in our model.
T pl and T ql refer to a block of ResNet cells (Figure 3).
F d refers to a single convolution computing the parameters of the data distribution, while convolutions F pl and F ql give the parameters of the prior and posterior distribution of layer l respectively.
⊕ is a merging operator unless one of its operands is zero, in which case it corresponds to an identity mapping.
Algorithm 1 computes the optimization objective of Equation 2.
It consists
Algorithm 1 Infer
This algorithm computes a lower bound on the marginal data loglikelihood (Equation 2).
: Inputs 2: x : The data.
3: h , k q ← Bottom − Up Pass ( x ) (Algo. 2).
4: c p , z , p ( z ) , q ( z | x ) ← Top − Down Pass (cid:0) h , k q (cid:1) (Algo. 3).
5: c pL +1 ← T pL +
c pL ⊕ z
of three main steps:
The bottom-up pass that is described described in Algorithm 2.
The top-down pass that is described described in Algorithm 3. The formation of the conditional data distribution p ( x | z ) . algorithm
During the bottom-up pass described in Algorithm 2 (see also left part of Figures 1a, 1b), the data x are passed through a sequence of ResNet transformations T ql to give a sequence of features h and keys k q . h and k q will be used by the stochastic layers of the VAE hierarchy when forming the posterior distribution, in a reverse order.
Note that in Algorithm 2, k qL +1 is set to zero since layer L has access only to its own h L and it is not connected with another layer from the bottom.
Algorithms 1 and 2 are very similar to that of typical deep VAEs.
The only difference is that the ResNet transformations are also producing the data keys k q along with the data features h .
The second, top-down phase of Algorithm 1 infers samples of latent variables z along with the prior and posterior distributions.
It is described in Algorithm 3 and explained in detail below.
In Algorithm 1, c pL encapsulates latent information of z <L .
Along with z L , they feed the ResNet block T pL +1 succeeded by the convolution F d to compute the conditional data likelihood p ( x | z ) that is parameterized by θ .
Algorithm 3 describes the top-down pass that is also visualized in the right part of Figures 1a, 1b.
The layers are traversed from the top ( l = 0 ) to the bottom-up ( l = L ), sampling latent variables z l from the posterior distribution and computing context features c pl that contain processed latent information of previous layers.
The prior distribution is also computed in order to obtain the prior regularization term D KL ( q ( z | x ) ∥ p ( z )) of the objective function in Equation 2.
The first layer does not receive latent information or context from a previous layer, therefore it sets z 0 and c p 0 to zero.
Similarly, since there is no other layer in the generative network to connect with, it does not produce a query s p 1 to attend with.
The last layer does not pass any prior information to a subsequent layer, therefore it does not compute a key k pL .
Similarly, the last layer can use only its own data features h L .
Therefore, the generative portion of the layer does not compute a query for the inference portion setting s qL ≡ 0 .
Each layer l , unless in a marginal case described above, computes four types of features k pl , s pl , s ql and c pl (Line 6).
Key features k pl are used at the depth-wise Attend step of next layers > l (Line 10).
s pl is a query to the generative network of previous layers < l .
s ql is a query to the inference network of current and next layers ≥ l .
Finally, c pl are stochastic features that, along with s ql , will be used to form the condition of the prior distribution.
We note here that s ql is used both to form the condition of the prior ˆ c pl (Line 7) and to query the inference network when forming the condition of the posterior ˆ c ql (Line 21).
A layer in a deep VAE that is not attention-guided does not produce k pl , s ql , s pl and the computation falls into the “else” cases of Algorithm 3 (Lines 15, 26).
In Line 10, the generative network queries previous layers via s pl to identify latent context features c pl ′ represented by keys k pl ′ that are most critical when forming its prior beliefs.
The resulting feature ˆ c p l , after being normalized in Line 12, rectifies the initial context c pl by the residual connection of
Line 13.
The initial prior context c pl is being used by next layers while the attention-guided context ˆ c pl is the conditioning factor of the prior distribution formed in Lines 17, 18.
A similar procedure takes place to form the posterior distribution.
c ql is the initial posterior context that merges deterministic information h l and stochastic, latent information ˆ c pl , as seen in Lines 20, 26.
c ql is updated with an attention-driven feature computed in Line 21 by the residual connection of
This algorithm infers a sequence of stochastic, latent representations of the data .
1: Inputs 2:
4: Set z 0 ≡ 0 , c p ≡ s p ≡ s qL ≡
5: for l = 1 , 2 , . . .
, L do 6: [ s ql , c pl , s pl , k pl ] ← T pl ( z l 1 ⊕ c pl − 1 ) .
7: c pl ← c pl ⊕ s ql .
Form the prior distribution .
8: if l > 1 then 9: c pl − 1 ← LayerNorm (cid:0) c pl − 1 (cid:1) .
10: ˆ c pl ← Attend (cid:0) { c pl ′ } l − 1 l ′
=
1 , s pl , { k pl ′ } l − 1 l ′ =1 (cid:1) .
11: ˆ c pl ← ˆ c pl + Gelu(LayerNorm
( ˆ c pl )) .
12: c pl ← c pl + Gelu(LayerNorm
( c pl )) .
13: ˆ c pl ← c pl + γ pl ˆ c pl .
14: else 15: ˆ c pl ← c pl .
16: end if 17: [ µ pl , σ pl ] ← F pl ( ˆ c pl ) .
18: p ( z l | z <l ) = N ( µ pl , σ pl ) .
Form and sample from the posterior distribution .
19: if l < L then 20: [ c ql , k p,ql ] ← h l ⊕
ˆ
c
pl .
21: ˆ c ql ← Attend (cid:0) { h } Ll ′
=
l ∪ { c ql } , s ql , { k ql ′ } Ll ′
=
l ∪ { k p,ql } (cid:1) .
22: ˆ c ql ← ˆ c ql + Gelu(LayerNorm
( ˆ c ql )) .
23: c ql ← c ql + Gelu(LayerNorm
( c ql )) .
24: ˆ c ql ← c ql + γ ql ˆ c ql .
25: else 26: ˆ c q l ← h l ⊕ ˆ c p l .
27: end if 28: [∆ µ ql , ∆ σ ql ] ← F ql ( ˆ c ql ) .
29: q ( z l | x , z <l ) =
N ( µ pl + ∆ µ ql , σ pl ∆ σ ql ) .
30: z l q ( z l | x , z <l ) .
31: end for 32: Return 33:
Line 24. In Line 29, the residual parametrization of the posterior, as proposed by (Vahdat & Kautz, 2020) and also described in Equation 22, is derived.
B I MPLEMENTATION D ETAILS
For all the experiments, we use the AdaMax (Kingma & Ba, 2014) optimizer with the default parameters for training.
The learning rate is decayed according to the cosine schedule with warm-up during the first 5 epochs.
When we train for 900 epochs, we use a cosine decay schedule with restarts with T
= 300 and T mult = 2 following the notation of Loshchilov & Hutter (2017).
We apply weight normalization (Salimans & Kingma, 2016) in all our models.
We use the Swish (Ramachandran et al., 2017) activation function in the ResNet cells of the architecture.
Table 6 summarizes the main training and architectural hyperparameters for the experiments presented in Section 3.
However, it should be noted that due to the heavy computing requirements, we do not exhaustively optimize the hyperparameters.
Therefore, the results reported could potentially be further improved.
[Table]
[Figure]
[Table]
KL Annealing.
We apply KL annealing (Sønderby et al., 2016) for the first 16 training epochs starting from 0.1 to 1.
This technique introduces a scheduled regularization coefficient β for the
KL-divergence.
Formally, the regularized ELBO objective becomes:
[Equation]
(31) where β linearly increases from β 0 to 1 for a number of epochs at the beginning of the training.
Variational Distributions.
We use Gaussian prior and posterior distributions connected via a residual parametrization, see Equation 22.
To ease training stability, we bound the log-scale of the prior log σ p and the posterior log σ q as reported in Table 6.
In case of CIFAR-10, we also apply additive zero mean, Gaussian noise to the prior and posterior scale.
The mean µ and the log-scale log σ are jointly learned by using a shared convolution.
In particular, a 1 × 1 convolution with an ELU activation is used for the prior parameters.
For the posterior parameters of all but the first layer, two 3 × 3 convolutions with an ELU non-linearity are used.
For the posterior of the first layer, a single 1 × 1 convoltion with no activation is used.
Data Distributions.
For MNIST and OMNIGLOT, we assume a logit-based parametrization of a
Bernoulli distribution.
For CIFAR-10, we use a mixture of discretized Logistic distributions (Salimans et al., 2017) with five mixture components.
The parameters are jointly learned by a single 3 × 3 convolution with an ELU activation.
Pre-Processing and Post-Processing Blocks.
The data x , before being passed to the stochastic hierarchy of L layers, are first pre-processed by a cascade of ResNet cells.
Each block is responsible for downsampling the spatial dimension by a factor of 2.
For MNIST, the image is downsampled twice by two blocks, leading to latent variables with spatial dimension 8 × 8 .
For OMNIGLOT and CIFAR-10, one pre-processing block is being used leading to latent spaces of dimension 16 × 16 .
No further downsampling is performed by encoders in the main hierarchy.
Each block consists of a series of ResNet cells.
The last cell of the block is responsible for downsampling while doubling the number of the output channels.
Symmetrically, the context of the last layer L is first being post-processed by a cascade of post-processing blocks before being passed to the data distribution of the decoder.
Each such block is responsible for upsampling the context by a factor of 2 in order to match the initial spatial dimension of the image.
The upsampling operation is performed by the first ResNet cell while halving the number of the output channels.
Spatially Non-Local Blocks.
For CIFAR-10, we interleave non-local blocks as shown in Figurein the pre-processing blocks of the encoder, as well as the encoder and decoder of each layer.
We interleave non-local blocks with the convolutions of the ResNet cells as shown in Figure 3.
For OMNIGLOT and MNIST, we interleave non-local blocks as shown in Figure 3 in the post-processing blocks of the encoder.
We noticed that it suffices to add a single non-local block right before the first
ResNet cell of the encoder and decoder of each layer.
C T RAINING P LOTS
Figures 4, 5 show the minimization of the negative ELBO (NELBO) loss in Equation 2 on the training and the validation CIFAR-10 images across 900 training epochs.
The first training phase (left part) refers to the KL annealing period.
For the next two training phases, the learning rate is annealed according to the cosine schedule.
At epoch 300, the learning rate is restarted (right part).
Figures 6, 7 focus on the KL loss of the NELBO.
In this section, we perform ablation experiments to investigate the effect of different architectural choices on model performance.
For the experiments in Tables 7, 8 and 10, the models were trained on CIFAR-10 for 500 epochs.
We used shallower attentive VAE models of 8 layers for these experiments.
We considered model instances of a varying size, with 64 and 128 channels at the beginning of the bottom-up network and at the end of the top-down network.
As already noted in Section B, the number of channels is halved when a residual cell upsamples the spatial dimension and vice versa.
In all cases, model instances with more channels in the ResNet modules in encoder and decoder perform significantly better.
We can also see that both weight normalization (Salimans & Kingma, 2016), Table 7, and batch normalization layers interleaved with the convolutions of both the bottom-up and the top-down model,
Table 8, improve performance.
These results are consistent with the observations made by Vahdat & Kautz (2020).
Moreover, their effect is greater on the larger models.
[Table] [Table]
In Table 9, we report the performance of different activation functions.
The Swish (Ramachandran et al., 2017) and GELU (Hendrycks & Gimpel, 2016) activations yield the best performance.
This is in accordance with similar choices in current deep VAE architectures (Vahdat & Kautz, 2020; Child, 2020).
The rest of the hyper-parameters are identical to that of Section B.
[Table]
In Table 10, we take a closer look at the normalized depth-wise attention proposed in Section 2.4.
The gap in the performance is more noticeable for deep VAE models of 16 layers.
We train all model instances for 400 epochs.
We omit the application of the normalization scheme presented in Equations 5 and 6 in the generative and/or inference network.
We confirm that usage of depth-wise normalization facilitates training and helps scale the depth of the proposed architecture.
[Table]
Normalized
Depth
[Table]
Attention
[Table]
E E MPIRICAL EVIDENCE OF BETTER LATENT VARIABLE UTILIZATION BY ATTENTIVE VAE
In this section, we probe the capacity of attentive VAE to utilize the latent information inserted in the model.
We first examine the Kullback–Leibler divergence D KL ( q ∥ p ) observed on the test data across the hierarchy during training.
Unless it explodes, a high D KL ( q ∥ p ) loss term is often desirable since it leads to a smaller negative ELBO (NELBO) loss in Equation 2.
This is because higher D KL ( q ∥ p ) indicates more informative latent factors z , and in turn, higher conditional likelihood p ( x | z ) .
For example, Kingma et al.
(2016) use a modified objective function that lower bounds D KL ( q ∥ p ) by some free bits.
Similarly, Vahdat & Kautz (2020) use a residual parametrization between the prior and the posterior to keep more latent variables active.
In this work, we tackle posterior collapse by imposing attention-driven connections between the layers.
The proposed depth-wise attention naturally realizes long-range couplings between the layers at a pixel level.
Therefore, it selectively preserves the most informative latent information which may lie far beyond in the hierarchy.
In Figure 8, we show log D KL ( q ∥ p ) per layer z l for the baseline
NVAE during the first 16 KL annealing epochs (left part) and for the rest of the training (right part).
Similarly, in Figure 9, we show log D KL ( q ∥ p ) for the proposed model.
Given the above plots, we make the following observations:
• Although in both models, log D KL ( q ∥ p ) remains large (always positive), our model presents a more balanced activation map in the latent space.
In contrast, in Figure 8, we see that the first two layers are significantly much less active compared to the layers closer to p ( x | z ) .
This phenomenon is exacerbated in deeper hierarchies.
• In Figure 8, log D KL ( q ∥ p ) increases with depth.
In contrast, in Figure 9 of our model, the first layer converges to a much higher D KL ( q ∥ p ) compared to layer 12 although it is much farther from p ( x | z ) .
• In Figure 9, we notice that as the learning progresses, the attention scores can help activate some latent units.
For example, layer 6 is relatively dormant at the beginning of the training, while it converges at a highly informative posterior distribution at the end.
Finally, in Table 11, we show the bits per dimension as the depth varies for our model.
All models in Table 11 were trained for 500 epochs.
As we can see the predictive improvement gained when doubling the size of the model from 8 to 16 layers is on par with the improvement accomplished when transitioning from a model of 2 layers to a model of 4 layers.
In contrast, deep VAEs without attention experience half the amount of relative improvement, as seen in Table 1.
We also note that the first two smallest models have converged during the last 100 training epochs which is not the case for the last two model instances.
Therefore, the gap in the relative decrease for the models ofand 16 layers could be further closed by longer training. [Table]
F V ISUALIZATION OF A TTENTION P ATTERNS
In order to get further insights into the attention-driven inference, we visualize the attention maps proposed in this work.
In Figure 10, we show the layer connectivity in the generative network as explained in Section 2.4 for some test images that belong to different classes.
For the same images, in Figure 11 we illustrate the depth-wise attention scores in the inference network we describe in
Section 2.5, and explain in greater detail in Section A.1.
For both cases, we chose to visualize the layer which has the largest receptive field, i.e, the last layer for L = 16 which is right before the conditional likelihood p ( x | z L ) in the decoder, and the first layer which has access to all hidden, data-dependent feature maps of the bottom-up pass in the encoder.
This analysis validates the hypothesis made in this work and suggests the following conclusions:
• For all pixels, the layer relies greatly on the immediately above layer when forming the prior distribution, as shown in Figure 10.
• Different latent pixels of the same image can discover related latent clues that lie in different layers.
• Similarly, the connectivity patterns in the generative network vary significantly across images.
• Although the context c p 14 already carries latent information from all previous layers, the decoder does not depend solely on that.
Instead, it chooses to activate direct connections with layers far beyond in the hierarchy.
This reinforces our hypothesis that the effect of latent variables in early layers may be obfuscated by more recent latent information leading in practice to non fully factorized prior distributions: p ( z )
= (cid:81) l p ( z l | z <l ) .
• The inference attention scores in Figure 11 are more balanced across layers compared to the generative attention scores in Figure 10.
• Mitigation of posterior collapse as described in E is accomplished in a dual fashion.
First, all layers during the bottom-up pass of the inference can directly interact with the initial data feature maps of the last layers in the hierarchy, as shown in Figure 11.
Second, early latent variables of the top layers condition the prior distribution of the latent variables in layers with closer affinity to the conditional data distribution in the decoder, as shown in Figure
In Figures 12, 13, 14, we plot the attention scores among pixels that lie in the same layer as discussed in Sections 2.6 and A.2.
The attention scores pertain to the non-local block of the last convolution of the last residual cell, as depicted in Figure 3, of the ResNet module in the generative portion of each layer.
We derive the following observations:
• The connectivity patterns in all layers are highly structured.
This implies that the model could be benefited by sparse attention factorizations such as those proposed by Child et al.
(2019), that rely on a similar assumption.
• The attention schemes vary widely across layers.
For example, the activation map is more scattered in the first layer of Figure 12 compared to that of the last layer of Figure 14 which is sparser and gathered in periodic spatial locations.
This may align with the functionality of the latent information as discussed in current research (Child, 2020), according to which top layers capture global, high-level features while lower layers are responsible for capturing finer details.
• Different images exhibit different spatial attention maps.
The differences are more pronounced in the middle layer, as depicted in Figure 13. [Figure]
G R ECONSTRUCTIONS AND F ANTASY S AMPLES G ENERATED FROM OUR M ODEL
In this section, we qualitatively assess the Attentive VAE presented in the paper.
Figures 15b and 16b illustrate reconstructed MNIST and CIFAR-10 test images from the inferred latent variables.
Figures 15c and 16c show novel MNIST and CIFAR-10 images sampled by the generative model.
[Figure] [Figure]