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: This paper proposes an evolutionary optimization framework for training vartional autoencoders (VAEs) with discrete latents. In contrast to the standard VAE paradigm, the proposed TVAE approach does not require an encoder for amortized inference given the input. The method instead relies on a pool of latent variable samples for each data point to activate the decoder network. The latent variable pools are maintained and iteratively updated to increase the average lower-bound of the marginal log-likelihood of the input data. Experimental results show that non-linear decoders optimized by the TVAE framework outperform their linear counterparts on a denoising task.  Further results demonstrate method's competitiveness on zero-shot denoising, where a TVAE decoder is only trained on the noisy input image to reconstruct a smoothed version of the input image.