Deep Attentive Variational Inference

1 INTRODUCTION . 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 . ( 2016 ) ; 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 . 2 PROPOSED MODEL . 2.1 DEEP VARIATIONAL INFERENCE . A latent variable generative model defines the joint distribution of a set of observed variables , x ∈ RD , 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 : p ( x ) = ∫ p ( x , z ) dz . ( 1 ) 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 : log p ( x ) = log ∫ p ( x , z ) dz = log ∫ q ( z | x ) q ( z | x ) p ( x , z ) dz ≥ ∫ q ( z | x ) log [ p ( x | z ) p ( z ) q ( z | x ) ] dz =⇒ log p ( x ) ≥ Eq ( z|x ) [ log p ( x | z ) ] −DKL ( q ( z | x ) ∥ p ( z ) ) , ( 2 ) 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 = { z1 , z2 , . . . , zL } , 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. , hl = T q l ( hl+1 ) , with hL+1 = x . The inference q ( z | x ) and generative p ( z ) network decompose in an identical topological ordering : q ( z | x ) = ∏ l q ( zl | x , z < l ) and p ( z ) = ∏ l p ( zl | 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.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 on all the previous in the hierarchy : p ( zl | z < l ) . Similarly , for the inference model : q ( zl | 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 ( zl | zl−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 cpl in Figure 1a ) . 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 zl 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 ) =∏ l p ( zl | z < l ) and q ( z | x ) = ∏ l q ( zl | 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 cpl 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 conditioning factors of the conditional distributions p ( zl | z < l ) and q ( zl | 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 for inference . Figure 1b is a graphical illustration of the proposed model . 2.3 DEPTH-WISE ATTENTION Figure 2 : Attend ( c < l , sl , k < l ) − a depth-wise attention block . c < l , k < l are the sequences of l − 1 contexts and corresponding keys with C and Q feature maps accordingly , while sl 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 . α < l = { αm→l } l−1m=1 are the attention scores . 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 = { cm } l−1m=1 of l − 1 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 cm are features of the same dimensionality : ĉl ∈ RH×W×C , and cm ∈ RH×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 need to handle long context sequences of large dimensions H and W . The task is characterized by a query feature sl ∈ RH×W×Q of dimensionality Q with Q≪ C. Similarly , cm is represented by a key feature km ∈ RH×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 ) ∈ RC 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 , cm ( i , j ) with m < l , and is given by : ĉl ( i , j ) = ∑ m < l αm→l ( i , j ) cm ( i , j ) , ( 3 ) αm→l ( i , j ) = exp ( sTl ( i , j ) km ( i , j ) ) ∑ m < l exp ( s T l ( i , j ) km ( i , j ) ) . ( 4 ) In words , feature sl ( i , j ) ∈ RQ queries the significance of feature cm ( i , j ) ∈ RC , represented by km ( i , j ) ∈ RQ , to form ĉl ( i , j ) ∈ RC . α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 , sl , 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 . ( 2020 ) : c < l ← c < l +Gelu ( LayerNorm ( c < l ) ) , ( 5 ) ĉl ← ĉl +Gelu ( LayerNorm ( ĉl ) ) , ( 6 ) where Gelu is the GELU non-linearity ( Hendrycks & Gimpel , 2016 ) and LayerNorm a layer normalization operation ( Ba et al. , 2016 ) .

Summary: This paper improves the architecture of deep VAEs using the attention mechanism.   The attention mechanism is used in two ways: 1. layer-wise attention (attending stochastic feature maps which are conditioned on other variables within the hierarchy, interpreted as a mixture of skip connections) 2. non-local attention (attention across the spatial dimensions, increases the size of receptive field)  The authors demonstrate the effectiveness of their architectural changes by challenging current sota deep VAEs on MNIST, OMNIGLOT and CIFAR-10. Notably, they outperform the state-of-the-art methods (in likelihood) on CIFAR-10 using fewer layers and fewer GPU hours.   The authors provide an extensive ablation study, showing the impact of each type of attention on the training performances (test likelihood on CIFAR-10).