Direct Evolutionary Optimization of Variational Autoencoders with Binary Latents

Discrete latent variables are considered important to model the generation process of real world data , which has motivated research on Variational Autoencoders ( VAEs ) with discrete latents . However , standard VAE training is not possible in this case , which has motivated different strategies to manipulate discrete distributions in order to train discrete VAEs similarly to conventional ones . Here we ask if it is also possible to keep the discrete nature of the latents fully intact by applying a direct discrete optimization for the encoding model . The studied approach is consequently strongly diverting from standard VAE training by altogether sidestepping absolute standard VAE mechanisms such as sampling approximation , reparameterization trick and amortization . Discrete optimization is realized in a variational setting using truncated posteriors in conjunction with evolutionary algorithms ( using a recently suggested approach ) . For VAEs with binary latents , we first show how such a discrete variational method ( A ) ties into gradient ascent for network weights and ( B ) uses the decoder network to select latent states for training . More conventional amortized training is , as may be expected , more efficient than direct discrete optimization , and applicable to large neural networks . However , we here find direct optimization to be efficiently scalable to hundreds of latent variables using smaller networks . More importantly , we find the effectiveness of direct optimization to be highly competitive in ‘ zero-shot ’ learning ( where high effectiveness for small networks is required ) . In contrast to large supervised neural networks , the here investigated VAEs can , e.g. , denoise a single image without previous training on clean data and/or training on large image datasets . More generally , the studied approach shows that training of VAEs is indeed possible without sampling-based approximation and reparameterization , which may be interesting for the analysis of VAE-training in general . In the regime of few data , direct optimization , furthermore , makes VAEs competitive for denoising where they have previously been outperformed by non-generative approaches . 1 INTRODUCTION AND RELATED WORK . Variational autoencoders ( Kingma & Welling , 2014 ; Rezende et al. , 2014 ) are prominent and very actively researched models for unsupervised learning . VAEs , in their many different variations , have successfully been applied to a large number of tasks including semi-supervised learning ( e.g . Maaløe et al. , 2016 ) , anomaly detection ( e.g . An & Cho , 2015 ; Kiran et al. , 2018 ) , sentence interpolation ( Bowman et al. , 2016 ) , music interpolation ( Roberts et al. , 2018 ) and drug response prediction ( Rampasek et al. , 2017 ) . The success of VAEs rests on a series of methods that enable the derivation of scalable training algorithms to optimize their model parameters ( discussed further below ) . A desired feature when applying VAEs to a given problem is that their latent variables ( i.e. , the encoder output variables ) correspond to meaningful properties of the data , ideally to those latent causes that have originally generated the data . However , many real-world datasets suggest the use of discrete latents as they often describe the data generation process more naturally . For instance , the presence or absence of objects in images is best described by binary latents ( e.g . Jojic & Frey , 2001 ) . Discrete latents are also a popular choice in modeling sounds ; for instance , describing piano sounds may naturally involve binary latents : keys are pressed or not ( e.g . Titsias & Lázaro-Gredilla , 2011 ; Goodfellow et al. , 2013 ; Sheikh et al. , 2014 ) . The success of standard forms of VAEs has consequently spurred research on novel formulations that feature discrete latents ( e.g . Rolfe , 2016 ; Khoshaman & Amin , 2018 ; Roy et al. , 2018 ; Sadeghi et al. , 2019 ; Vahdat et al. , 2019 ) . The objective of VAE training is the optimization of a generative data model which parameterizes a given data distribution . Typically we seek model parameters Θ of a VAE that maximize the data log-likelihood , L ( Θ ) = ∑ n log ( pΘ ( ~x ( n ) ) ) , where we denote by ~x ( 1 : N ) a set of N observed data points , and where pΘ ( ~x ) denotes the modeled data distribution . Like conventional autoencoders ( e.g. , Bengio et al. , 2007 ) , VAEs use a deep neural network ( DNN ) to generate ( or decode ) observables ~x from a latent code ~z . Unlike conventional autoencoders , however , the generation of data ~x is not deterministic but it takes the form of a probabilistic generative model . For VAEs with binary latent variables , as they will be of interest here , we consider the following VAE generative model : pΘ ( ~z ) = Bern ( ~z ; ~π ) = ∏ h ( πzhh ( 1− πh ) ( 1−zh ) ) , pΘ ( ~x |~z ) = N ( ~x ; ~µ ( ~z ; W ) , σ2I ) , ( 1 ) where ~z ∈ { 0 , 1 } H is a binary code and the non-linear function ~µ ( ~z ; W ) is a DNN that outputs the mean of the Gaussian distribution . pΘ ( ~x |~z ) is commonly referred to as decoder . The set of model parameters is Θ = { ~π , W , σ2 } , where W incorporates DNN weights and biases . We assume homoscedasticity of the Gaussian distribution , but note that there is no obstacle to generalizing the model by inserting a DNN non-linearity that outputs a correlation matrix . Similarly , the algorithm could easily be generalized to different noise distributions should the task at hand call for it . For the purpose of this work , however , we will focus on as elementary as possible VAEs , with the form shown in Eqn . ( 1 ) . Given standard or binary-latent VAEs , essentially all learning algorithms seek to approximately maximize the log-likelihood using the following series of methods ( we elaborate in the appendix ) : ( A ) Instead of the log-likelihood , a variational lower-bound ( a.k.a . ELBO ) is optimized . ( B ) VAE posteriors are approximated by an encoding model , that is a specific distribution ( often Gaussian ) parameterized by one or more DNNs . ( C ) The variational parameters of the encoder are optimized using gradient ascent on the lower bound , where the gradient is evaluated based on sampling and reparameterization trick to obtain sufficiently low-variance and yet efficiently computable estimates . ( D ) Using samples from the encoder , the parameters of the decoder are optimized using gradient ascent on the variational lower bound . Optimization procedures for VAEs with discrete latents follow the same steps ( Points A to D ) . However , discrete or binary latents pose substantial further obstacles in learning , mainly due to the fact that backpropagation through discrete variables is generally not possible ( Rolfe , 2016 ; Bengio et al. , 2013 ) . In order to maintain the general VAE framework for encoder optimization , different groups have therefore suggested different possible solutions : work by Rolfe ( 2016 ) , for instance , extends VAEs with discrete latents by auxiliary continuous latents such that gradients can still be computed . Work on the concrete distribution ( Maddison et al. , 2016 ) or Gumbel-softmax distribution ( Jang et al. , 2016 ) proposes newly defined continuous distributions that contain discrete distributions as limit cases . Work by Lorberbom et al . ( 2019 ) merges the Gumbel-Max reparameterization with the use of direct loss minimization for gradient estimation , enabling efficient training on structured latent spaces . Finally , work by van den Oord et al . ( 2017 ) , and Roy et al . ( 2018 ) combines VAEs with a vector quantization ( VQ ) stage in the latent layer . Latents become discrete through quantization but gradients for learning are adapted from latent values before they are processed by the VQ stage . All methods have the goal of treating discrete distributions such that standard VAE training as developed for continuous latents can still be applied . These techniques interact during training with the standard methods ( Points A-D ) already in place for VAE optimization . Furthermore , they add further types of design decisions and hyper-parameters , for example parameters for annealing from softened discrete distributions to the ( hard ) original distributions for discrete latents . For discrete VAEs , it may consequently be a desirable goal to investigate alternative , more direct optimization procedures that do not require a softening of discrete distributions or the use of other indirect solutions . Such a direct approach is challenging , however , because once DNNs are used to define the encoding model ( Point B ) standard tricks to estimate gradients ( Point C ) seem unavoidable . A direct optimization procedure , as is investigated here , consequently has to substantially change VAE training . For the data model ( 1 ) we will maintain the variational setting and a decoding model with DNNs as non-linearity ( Points A and D ) . However , we will not use an encoder model parameterized by DNNs ( Point B ) . Instead , the variational bound will be increased w.r.t . the encoder model by using a discrete optimization approach . The procedure does not require gradients to be computed for the encoder such that discrete latents are addressed without the use of reparameterization trick and sampling approximations . 2 TRUNCATED VARIATIONAL OPTIMIZATION . Let us consider the variational lower bound of the likelihood . If we denote by q ( n ) Φ ( ~z ) the variational distributions with parameters Φ ( n ) , and by 〈 h ( ~z ) 〉 q ( n ) Φ = ∑ ~z q ( n ) Φ ( ~z ) h ( ~z ) expectation values w.r.t . to q ( n ) Φ ( ~z ) , then the lower bound can be written as : F ( Φ , Θ ) = ∑ n 〈 log ( pΘ ( ~x ( n ) |~z ) pΘ ( ~z ) ) 〉 q ( n ) Φ − ∑ n 〈 log ( q ( n ) Φ ( ~z ) ) 〉 q ( n ) Φ , ( 2 ) The general challenge for the maximization of F ( Φ , Θ ) is the optimization of the encoding model q ( n ) Φ . VAEs with discrete latents add to this challenge the problem of taking gradients w.r.t . discrete latents . If we seek to avoid derivatives w.r.t . discrete variables , we have to define an alternative encoding model q ( n ) Φ but such an encoding has to remain sufficiently efficient . Considering prior work on generative models with discrete latents , variational distributions based on truncated posteriors offer themselves as such an alternative ( Lücke & Sahani , 2008 ) . Truncated posterior approximations have been shown to be functionally competitive ( e.g . Sheikh et al. , 2014 ; Hughes & Sudderth , 2016 ; Shelton et al. , 2017 ) , and they are able to efficiently train also very large scale models with hundreds or thousands of latents ( e.g . Shelton et al. , 2011 ; Sheikh & Lücke , 2016 ; Forster & Lücke , 2018 ) . However , the important question for training discrete VAEs is if or how truncated variational distributions can be used in gradient-based optimization of neural network parameters . We here , for the first time , address this question noting that all previous approaches relied on closed-form ( or pseudo-closed form ) parameter update equations in an expectation-maximization learning paradigm . Optimization of the Decoding Model . In order to optimize the parameters W of the decoder DNN ~µ ( ~z , W ) , the gradient of the variational bound ( 2 ) w.r.t . W has to be computed . We consequently need , for any VAE , a sufficiently precise and efficient approximation of the expectation value w.r.t . the encoder q ( n ) Φ ( ~z ) . Gradient estimation is of central importance for deep unsupervised learning , and approaches , e.g. , for variance reduction of estimators have played an important role and are dedicated solely to this purpose ( e.g. , Williams , 1992 ) . Reparameterization finally emerged as a key method because it allowed for sufficiently low-variance estimation of gradients based , e.g. , on Gaussian middle-layer units ( Kingma & Welling , 2014 ; Rezende et al. , 2014 ) . For discrete VAEs , however , reparameterization requires the introduction of additional manipulations of discrete distributions that we here seek to fully avoid . Instead of using reparameterization or variance reduction , we will compute gradients based on truncated posterior as variational distributions . A truncated posterior has the following form : q ( n ) Φ ( ~z ) : = pΘ ( ~z | ~x ( n ) ) ∑ ~z ′∈Φ ( n ) pΘ ( ~z ′ | ~x ( n ) ) = pΘ ( ~x ( n ) |~z ) pΘ ( ~z ) ∑ ~z ′∈Φ ( n ) pΘ ( ~x ( n ) |~z ′ ) pΘ ( ~z ′ ) if ~z ∈ Φ ( n ) , ( 3 ) where for all ~z 6∈ Φ ( n ) the probability q ( n ) Φ ( ~z ) equals zero . That is , a variational distribution q ( n ) Φ ( ~z ) is proportional to the true posteriors in a subset Φ ( n ) , which acts as its variational parameter . We can now compute the gradient of ( 2 ) w.r.t . the decoder weights W which results in : ~∇WF ( Φ , Θ ) = − 1 2σ2 ∑ n ∑ ~z∈Φ ( n ) q ( n ) Φ ( ~z ) ~∇W ‖~x ( n ) − ~µ ( ~z , W ) ‖2 . ( 4 ) The right-hand-side has salient similarities to the standard gradient ascent for VAE decoders . Especially the familiar gradient of the mean squared error ( MSE ) shows that , e.g. , standard automatic differentiation tools can be applied . However , the decisive difference are the weighting factors q ( n ) Φ ( ~z ) . Considering ( 3 ) , in order to compute the weighting factors we require all ~z ∈ Φ ( n ) to be passed through the decoder DNN . As all states of Φ ( n ) anyway have to be passed through the decoder for the MSE term of ( 4 ) , the overall computational complexity is not higher than an estimation of the gradient with samples instead of states in Φ ( n ) ( we elaborate in Appendix A ) . To complete the decoder optimization , update equations for variance σ2 and prior parameters ~π can be computed in closed-form ( compare , e.g. , Shelton et al. , 2011 ) and are given by : σ2 , new = 1DN ∑ n ∑ ~z∈Φ ( n ) q ( n ) Φ ( ~z ) ‖~x ( n ) − ~µ ( ~z , W ) ‖2 , ~πnew = 1 N ∑ n ∑ ~z∈Φ ( n ) q ( n ) Φ ( ~z ) ~z , ( 5 ) where N is the number of samples in the training dataset and D is the number of observables . Optimization of the Encoding Model . After having established that the decoder can be optimized efficiently and by using standard DNN methods , the important question is if the encoder can be trained efficiently . Encoder optimization is usually based on a reformulation of the variational bound ( 2 ) given by : F ( Φ , Θ ) = ∑ n 〈 log ( pΘ ( ~x ( n ) |~z ) ) 〉 q ( n ) Φ − ∑ nDKL ( q ( n ) Φ ( ~z ) , pΘ ( ~z ) ) . ( 6 ) Centrally for this work , truncated posteriors allow a specific alternative reformulation of the bound that enables efficient optimization . The reformulation recombines the entropy term of the original form ( 2 ) with the first expectation value into a single term , and is given by ( see Lücke , 2019 , for details ) : F ( Φ , Θ ) = ∑ n log ( ∑ ~z∈Φ ( n ) pΘ ( ~x ( n ) |~z ) pΘ ( ~z ) ) . ( 7 ) Thanks to the simplified form of the bound , the variational parameters Φ ( n ) of the encoding model can now be sought using direct discrete optimization procedures . More concretely , because of the specific form ( 7 ) , pairwise comparisons of joint probabilities are sufficient to maximize the lower bound : if we update the set Φ ( n ) for a given ~x ( n ) by replacing a state ~z old ∈ Φ ( n ) with a state ~z new 6∈ Φ ( n ) , then F ( Φ , Θ ) increases if and only if : log ( pΘ ( ~x ( n ) , ~z new ) ) > log ( pΘ ( ~x ( n ) , ~z old ) ) . ( 8 ) To obtain intuition for the pairwise comparison , consider its form when inserting the binary VAE ( 1 ) into the left- and right-hand sides . Eliding terms that do not depend on ~z we obtain : ˜log pΘ ( ~x , ~z ) = −‖~x− ~µ ( ~z , W ) ‖2 − 2σ2∑h π̃h zh ( 9 ) where π̃h = log ( ( 1 − πh ) /πh ) . The expression assumes an even more familiar form if we restrict ourselves for a moment to sparse priors π < 12 , i.e. , π̃ > 0 . Criterion ( 8 ) then becomes : ‖~x ( n ) − ~µ ( ~z new , W ) ‖2 + 2σ2π̃ |~z new| < ‖~x ( n ) − ~µ ( ~z old , W ) ‖2 + 2σ2π̃ |~z old| ( 10 ) where |~z| = ∑H h=1 zh . Such functions are routinely encountered in sparse coding or compressive sensing ( Eldar & Kutyniok , 2012 ) : for each set Φ ( n ) we seek those states ~z that are reconstructing ~x ( n ) well while being sparse ( ~z with few non-zero bits ) . For VAEs , ~µ ( ~z new , W ) is a DNN and as such much more flexible in matching the distribution of observables ~x than can be expected from linear mappings . Furthermore , criteria like ( 10 ) usually emerge for maximum a-posteriori ( MAP ) training in sparse coding ( Olshausen & Field , 1996 ) . In contrast , we here seek a population of states ~z in Φ ( n ) for each data point . It is a consequence of the reformulated lower bound ( 7 ) that it remains optimal to evaluate joint probabilities ( as for MAP ) although the constructed population of states Φ ( n ) can capture ( unlike MAP training ) a rich posterior structure . But how can new states ~z new that optimize Φ ( n ) be found efficiently in high-dimensional latent spaces ? Random search and search by sampling has recently been explored for elementary generative models ( Lücke et al. , 2018 ) . Here we will follow another recent suggestion ( Guiraud et al. , 2018 ) and make use of a search based on evolutionary algorithms ( EAs ) . In this setting we interpret sets Φ ( n ) as populations of binary genomes ~z and base the fitness function on Eqn . ( 9 ) . Concretely , using Φ ( n ) as initial parent pool , we apply the following genetic operators in sequence : firstly , parent selection picks Np states from the parent pool . In our numerical experiments we used fitness-proportional parent selection , for which we add an offset ( constant w.r.t . ~z ) to the fitness values in order to make them strictly non-negative . Each of the children undergoes mutation : one or more bits are flipped to further increase offspring diversity . In our experiments we perform random uniform selection of the bits to flip . Crossover could also be employed to increase offspring diversity . We repeat the procedure using the children generated this way as the parent pool , giving birth to multiple generations of candidate states . Finally , we update Φ ( n ) by substituting individuals with low fitness with candidates with higher fitness . The whole procedure can be seen as an evolutionary algorithm with perfect memory or very strong elitism ( individuals with higher fitness never drop out of the gene pool ) . Note that the improvement of the variational lower bound depends on generating as many as possible different children with high fitness over the course of training . We point out that the EAs optimize each Φ ( n ) independently , so this technique can be applied to large datasets in conjunction with stochastic or batch gradient descent on the model parameters Θ : it does not require to keep the full dataset ( or all sets Φ ( n ) ) in memory at a given time . Fig 1 shows how EAs produce new states that are used to update each set Φ ( n ) . The full training procedure for binary VAEs is summarized in Algorithm 1 .

Summary: the paper proposes a novel approach to training variational autoencoder models, based on non-parametric form of truncated approximate posterior. Posterior is truncated to have support on a small subset of latent space allowing for exact marginalization. The support of approximate posterior in latent space for each data point $x$ is learned via evolutionary algorithm, minimizing ELBO. Method is applied to denoising tasks for images.