Variance Reduction in Hierarchical Variational Autoencoders

1 INTRODUCTION . Variational autoencoders ( VAE ) [ 10 ] are a popular latent variable model for unsupervised learning that simplifies learning by the introduction of a learned approximate posterior . Given data x and latent variables z , we specify the conditional distribution p ( x|z ) by parameterizing the distribution parameters by a neural network . Since it is difficult to learn such a model directly , another conditional distribution q ( z|x ) is introduced to approximate the posterior distribution . During learning the goal is to maximize the evidence lower bound ( ELBO ) , which lower bounds the log likelihood , log p ( x ) ≥ Eq ( z|x ) [ log p ( x|z ) +log p ( z ) − log q ( z|x ) ] . In their simplest form , the generative model p ( x|z ) and the approximate posterior q ( z|x ) are Gaussian distributions optimized in unison . A natural way to increase the modeling capacity of VAE is to incorporate a hierarchy of stochastic variables . Such models , however , turn out to be difficult to train and higher levels in the hierarchy tend to remain independent of input data – a problem termed posterior collapse . Posterior collapse in VAEs manifests itself by the latent distribution tending to fall back to the prior . With hierarchical VAEs the effect is found to be more pronounced in the top layers farther from the output . For the purpose of the paper and for clarity of exposition , we focus on the simplest extension of hierarchical variational autoencoders where stochastic layers are stacked serially on top of each other [ 2 , 21 ] , p ( x , z ) = p ( x|z1 ) p ( zL ) ∏L−1 i=1 p ( zi|zi+1 ) and q ( z|x ) = q ( z1|x ) ∏L−1 i=1 q ( zi+1|zi ) . The intermediate distributions in this model are commonly taken to be Gaussian distributions parameterized by neural network functions , so that p ( zi|zi+1 ) = N ( zi|µ ( zi+1 ) , σ ( zi+1 ) ) , where µ ( z ) , σ ( z ) are neural networks computing the mean and variance of the Gaussian distribution . We refer to them as vanilla hierarchical variational autoencoders . For each stochastic layer in this model there is a corresponding KL divergence term in the objective given by E [ KL ( q ( zi|zi−1 ) ||p ( zi|zi+1 ) ] . ( 1 ) As described later , expression 1 can be easily decomposed to show an explicit dependence on the variance of the parameterizing functions µ ( zi ) , σ ( zi ) of the intermediate Gaussian distribution . We further show the KL divergence term to be closely related to the harmonics of the parameterizing function . For complex parameterizing functions the KL divergence term has large high frequency components ( and thus high variance ) which leads to unstable training causing posterior collapse . Building on this , we suggest a method for training the simplest hierarchical extension of VAE that avoids the problem of posterior collapse without introducing further architectural complexity [ 13 , 21 ] . Given a hierarchical variational autoencoder , our training method incorporates a smoothing parameter ( we denote this by ρ ) in the neural network functions used to parameterize the intermediate latent distributions . The smoothing is done such that expected values are preserved , the higher frequencies are attenuated and the variance is reduced . Next , the gradients computed with the smooth functions are used to train the original hierarchical variational autoencoder . For the construction of the smoothing transformations for VAEs with Gaussian latent spaces we make use of ideas from the analysis of Gaussian spaces . We analyze the stochastic functions in vanilla hierarchical VAEs as Hermite expansions on Gaussian spaces [ 9 ] . The Ornstein-Uhlenbeck ( OU ) semigroup from Gaussian analysis is a set of operators that we show to smoothly interpolate between a random variable and its expectation . The OU semigroup provides the appropriate set of smoothing operators which enable us to control variance and avoid posterior collapse . We further show that by smoothing the intermediate parameterizing functions µ ( z ) , σ ( z ) in the proposed manner , the KL divergence of the top layer sees a sudden sharp drop toward zero as the amount of smoothing is decreased . This behaviour is retained when we evaluate the KL divergence on the original unsmoothed variational autoencoder model . This behaviour is reminiscent of phase transitions from statistical mechanics and we adopt the same terminology to describe the phenomenon . Our experiments suggest that the phenomenon is general across datasets and commonly used architectures . Furthermore , the critical value of the smoothing parameter ρ at which the transition occurs is fixed for a given model configuration and varies with stochastic depth and width . We make the following contributions . First , we establish a connection between higher harmonics , variance , posterior collapse and phase transitions in hierarchical VAEs . Second , we show that by using the Ornstein-Uhlenbeck semigroup of operators on the generative stochastic functions in VAEs we reduce higher frequencies and consequently variance to mitigate posterior collpase . We corroborate our findings experimentally and further obtain in CIFAR-10 likelihoods competitive with more complex architectural solutions alongside a reduction in model size . We refer to the proposed family of models as Hermite variational autoencoders ( HVAE ) . 2 HERMITE VARIATIONAL AUTOENCODERS . 2.1 ANALYSIS ON GAUSSIAN SPACES . The analysis of Gaussian spaces studies functions of Gaussian random variables . These are realvalued functions defined on Rn endowed with the Gaussian measure . Many functions employed in machine learning are instances of such functions : decoders for variational autoencoders , as is the case in this work , and generators for generative adversarial networks being two examples . By way of summary , the main facts we use from this field are that a function on a Gaussian space can be expanded in an orthonormal basis , where the basis functions are the Hermite polynomials . This orthonormal expansion is akin to a Fourier transform in this space . The second fact is that the coefficients of such an expansion can be modified in a way to reduce the variance of the expanded function by applying an operator from the Ornstein-Uhlenbeck semigroup of operators . Next , we give a brief introduction . For further details on Gaussian analysis we refer to [ 9 ] . Gaussian Spaces : Let L2 ( Rn , γ ) be the space of square integrable functions , f : Rn → R , with the Gaussian measure γ ( z ) = ∏ iN ( zi|0 , 1 ) . Given functions f , g in this space , the inner product is given by 〈f , g〉 = Eγ ( z ) [ f ( z ) g ( z ) ] . Basis functions for L2 ( R , γ ) : Taking the space of univariate functions L2 ( R , γ ) , it is known that the polynomial functions φi ( z ) = zi are a basis for this space . By a process of orthonormalization we obtain the normalized Hermite polynomial basis for this space . The first few Hermite polynomials are the following : h0 ( z ) = 1 , h1 ( z ) = z , h2 = z 2−1√ 2 , . . .. Basis functions for L2 ( Rn , γ ) : Letting α ∈ Nn be a multi-index , the basis functions for L2 ( Rn , γ ) are obtained by multiplying the univariate basis functions across dimension , hα ( z ) = ∏ i hαi ( zi ) . Hermite expansion : A function in L2 ( Rn , γ ) can be expressed as f = ∑ α∈Nn f̂ ( α ) hα , where f̂ ( α ) are the Hermite coefficients of f and are computed as f̂ ( α ) = 〈f , hα〉 = Eγ ( z ) [ f ( z ) hα ( z ) ] . Plancherel ’ s theorem is the following relation between the norm of f and f̂ which follows from orthnormality of the basis functions . 〈f , f〉 = ∑ α f̂ ( α ) 2 , ( 2 ) Ornstein-Uhlenbeck ( OU ) Semigroup : Given a parameter ρ ∈ [ 0 , 1 ] and a Gaussian variable z , we construct a correlated variable z′ as z′ = ρz + √ 1− ρ2zω , where zω ∼ N ( 0 , 1 ) is a random standard Gaussian sample . The OU semigroup is a set of operators , denoted Uρ and parameterized by ρ ∈ [ 0 , 1 ] . The action of Uρ on f at z is to average the function values on correlated z′s around z , Uρf ( z ) = Ez′|z [ f ( z′ ) ] = Ezω [ f ( ρz + √ 1− ρ2zω ) ] ( 3 ) The action of the Uρ operators on the Hermite expansion of function f ( z ) is to decay Hermite coefficients according to their degree , Uρf ( z ) = ∑ α∈Nn ρ |α|f̂ ( α ) hα . where |α| = ∑ i αi . If z is reparameterized as z = σ 1 + µ , the correlated OU sample is given by z′ = σ ( ρ 1 +√ 1− ρ2 2 ) + µ , where 1 , 2 are standard Gaussian variables . This can also be expressed in terms of z as z′ = ρz + ( 1− ρ ) µ+ σ √ 1− ρ2 2 , ( 4 ) 2.2 HERMITE EXPANSIONS FOR VAES . Our proposed method is a new training procedure for the vanilla hierarchical variational autoencoder that builds upon Hermite expansions of Gaussian functions and properties of the OU semigroup . In the context of hierarchical variational autoencoders , the Gaussian function f is the generative model µi ( zi+1 ) and σi ( zi+1 ) that receives as inputs the latent variable zi+1 to return the Gaussian latent variable of the next layer , zi ∼ N ( µi ( zi+1 ) , σi ( zi+1 ) ) . We make use of the following properties of the OU semigroup to construct Gaussian functions of lower variance . The first property we employ is that the OU semigroup of operators interpolates between a random variable ( ρ = 1 ) and its expectation ( ρ = 0 ) , where the parameter ρ controls the extent of the interpolation . Proposition 1 The operators Uρ retain the expected value of the operated function , E [ f ] = E [ Uρf ] . Proposition 2 The operators Uρ interpolate between a random variable and its expectation . In particular , as ρ→ 1 , Uρf = f . and as ρ→ 0 , Uρf = E [ f ] The second property we exploit is that the new random variable Uρf ( z ) has lower variance compared with original variable f ( z ) and is in general a smoother function than f ( z ) . The smoothing properties of the operator Uρ can be understood by examining the Hermite expansion of Uρf . First we note that we can express the expectation and variance of a function f in terms of its Hermite coefficients , specifically E [ f ] = f̂ ( 0 ) and Var ( f ) = E [ ( f −E [ f ] ) 2 ] = E [ ( f − f̂ ( 0 ) ) 2 ] =∑α : |α| > 0 f̂ ( α ) 2 , which follows from Plancherel ’ s theorem ( equation 2 ) . Replacing f with Uρf and using the Hermite expansion of Uρf from equation 3 , the mean remains the same , E [ Uρf ] = ρ0f̂ ( 0 ) = f̂ ( 0 ) , and variance reduces like Var [ Uρf ] = E [ ( Uρf − E [ f ] ) 2 ] = E [ ( f − f̂ ( 0 ) ) 2 ] = ∑ α : |α| > 0 ρ2|α|f̂ ( α ) 2 . ( 5 ) The last equation indicates that the contribution to the variance by f̂ ( α ) decays by an amount ρ2|α| when ρ ∈ ( 0 , 1 ) . This , in turn , leads to a decrease in variance . Algorithm . In essence , Hermite variational autoencoders are similar to variational autoencoders , save for applying the OU semigroup to the latent distributions p ( zi|zi+1 ) that comprise the generator to compute gradients during training only . Specifically , we apply these operators to the functions parameterizing the mean and variance of the latent Gaussian distributions . For each distribution p ( zi|zi+1 ) we substitute N ( zi|µi ( zi+1 ) , σi ( zi+1 ) ) with N ( zi|Uρµi ( zi+1 ) , Uρσi ( zi+1 ) ) . The new functions result in latent distributions with parameters that have lower variance but the same expected value relative to the conditional input latent distribution . In an alternative parameterization we apply the OU semigroup to the ratio of the mean and variance functions : Uρ µiσi ( zi+1 ) ( see next section for a justification of this ) . The OU semigroup operators can also be applied on approximate posterior functions , but we observe little benefit . In practice , we compute Uρµi ( zi+1 ) and Uρσi ( zi+1 ) by Monte Carlo averaging . As for a function f , Uρf = Ez′|z [ f ( z′ ) ] , where z′ are the correlated samples , we estimate the expectation by Monte Carlo averaging over z′ . Experiments show that 5 to 10 samples suffice . It is important to emphasize that the substitution of the lower variance functions for parameterizing the distributions is only done when computing gradients during training . All evaluations , training or test , are still done on the original hierarchical variational autoencoder model . Thus , the new training procedure has an additional computational cost only for the intermediate distributions in the generator , proportional to the number of correlated samples during training . Complexity . In Hermite VAE the OU sampling operation is only applied in the intermediate stochastic layers in the generator network . In particular , it is not applied in the inference network or in the last layer of the decoder . The fact that OU sampling is not applied in the final stochastic layer computing p ( x|z1 ) is especially important for deep VAEs for images since feature maps are upsampled to match image dimensions in this layer . Thus , for 5 OU samples , the added computational and activation memory complexity is significantly less than 5 times the total cost of the base VAE model , and is 5 times the cost in the higher decoder layers only in the base model . An empirical comparison of maximum memory usage of various models can be found in table 6 .

Summary: This paper studies the training of deep hierarchical VAEs and focuses on the problem of posterior collapse. It is argued that reducing the variance of the gradient estimate may help to overcome posterior collapse. The authors focus on reducing the variance of the functions parameterizing the variational distribution of each layer using a layer-wise smoothing operator based on the Ornstein-Uhlenbeck semigroup (parameterized by a parameter $\rho$). The operator requires additional Monte-Carlo samples. The authors provide an analytical analysis of bias and variance. Last they train multiple VAEs models, measure the posterior collapse and observe a phase transition behaviour depending on the parameter $\rho$.