Learning Discrete Structured Variational Auto-Encoder using Natural Evolution Strategies

Discrete variational auto-encoders ( VAEs ) are able to represent semantic latent spaces in generative learning . In many real-life settings , the discrete latent space consists of high-dimensional structures , and propagating gradients through the relevant structures often requires enumerating over an exponentially large latent space . Recently , various approaches were devised to propagate approximated gradients without enumerating over the space of possible structures . In this work , we use Natural Evolution Strategies ( NES ) , a class of gradient-free black-box optimization algorithms , to learn discrete structured VAEs . The NES algorithms are computationally appealing as they estimate gradients with forward pass evaluations only , thus they do not require to propagate gradients through their discrete structures . We demonstrate empirically that optimizing discrete structured VAEs using NES is as effective as gradient-based approximations . Lastly , we prove NES converges for non-Lipschitz functions as appear in discrete structured VAEs . 1 INTRODUCTION . Discrete variational auto-encoders ( VAEs ) are able to represent structured latent spaces in generative learning . Consequently VAEs drive extensive research in machine learning applications , including language classification and generation [ 60 , 17 , 54 , 9 , 13 ] , molecular synthesis [ 28 , 15 , 48 ] , speech and visual understanding [ 36 , 55 , 3 ] . Compared to their continuous counterparts , they can improve interpretability by illustrating which terms contributed to the solution [ 48 , 40 ] , and they can facilitate the encoding of inductive biases in the learning process , such as images consisting of a small number of objects [ 12 ] or tasks requiring intermediate alignments [ 36 , 42 , 1 , 2 ] . Learning VAEs with discrete n-dimensional latent variables is computationally challenging since the size of the support of the posterior distribution may be exponential in n. This is particularly common under the structured settings , when the latent variables represent complex structures such as trees or graphs . The Gumbel-max reparametrization trick trades enumeration with optimization using efficient dynamic programming algorithms and enables a computation of the model value . Unfortunately , the resulting mapping remains non-differentiable due to the presence of argmax operations . In order to propagate gradients efficiently , Jang et al . [ 19 ] , Maddison et al . [ 34 ] proposed the Gumbel-softmax reformulation that uses a smooth relaxation of the reparametrized objective , replacing the argmax operation with a softmax operation . Following such an approach may bring back the need for enumerating over a large search space . This is due to the partition function of the softmax operator , which relies on a summation over all possible latent assignments , which may be exponential in n. To better deal with the computational complexity in the structured setting , sophisticated stochastic softmax tricks were devised to learn discrete structured VAEs [ 48 ] ( e.g. , perturb-and-parse for dependency parsing by [ 9 ] , Gumbel-Sinkhorn for bi-partite matching [ 36 ] ) . In this work we propose to use the Natural Evolution Strategy ( NES ) [ 57 , 58 ] algorithm for learning discrete structured VAEs . The NES algorithm is a gradient-free black-box optimization method that does not need to propagate gradients through discrete structures . Instead , the NES algorithm estimates gradients by forward-pass evaluations only . We experimentally show that gradient-free methods are as effective as sophisticated gradient based methods , such as perturb-and-parse . NES is conceptually appealing when considering discrete structured VAEs since NES does not require to construct complex solutions to propagate gradients through the argmax operation , as it only requires to evaluate the model . Moreover , the proposed approach is highly parallelizable , hence computationally appealing . Our contributions : ( 1 ) We suggest using black-box , gradient-free based optimization methods , specifically NES , to optimize discrete structured VAEs . ( 2 ) We experimentally demonstrate that NES , which uses the models ‘ output in a black-box manner , is as effective as gradient based approximations although being more general and simpler to use . ( 3 ) We rigorously describe the connection between NES and previous gradient based optimization methods ( i.e , REINFORCE ) as well as prove that the NES algorithm converges for non-Lipschitz functions . 2 BACKGROUND . Discrete Structured Variational Auto-Encoders ( VAEs ) learn a generative model pθ ( x ) using a training set S = { x1 , . . . , xm } , derived from an unknown distribution p ( x ) by minimizing its negative log-likelihood . VAEs rely on latent variable models of the form pθ ( x ) = ∑ z∈Z pθ ( z ) pθ ( x|z ) , where z is a realization of the latent variable andZ is the discrete set of its possible assignments . We focus on discrete structures such as spanning trees in a graph G = ( V , E ) , i.e. , z = ( z1 , .... , z|E| ) represents a spanning tree T for which ze = 1 if the edge e ∈ T belongs to the spanning tree T and zero otherwise . In this case , Z is the spanning trees space which is typically exponential in the size of the input , as there are |V ||V |−2 spanning trees for a complete graph with |V | vertices . VAEs rely on an auxiliary distribution qφ ( z|x ) that is used to upper bound the negative loglikelihood of the observed data points : ∑ x∈S − log pθ ( x ) ≤ ∑ x∈S L ( θ , φ , x ) , where : L ( θ , φ , x ) = −Ez∼qφ ( ·|x ) [ log pθ ( x|z ) ] +KL ( qφ ( ·|x ) ||pθ ( · ) ) . This formulation is known as the negative Evidence Lower Bound ( ELBO ) [ 20 ] , where the KL-divergence measures the similarity of two distributions qφ and pθ , and is defined as : KL ( qφ ( ·|x ) ||pθ ( · ) ) = Ez∼qφ ( ·|x ) [ log ( qφ ( z|x ) /pθ ( z ) ) ] . Parameter estimation is generally carried out by performing gradient descent on ∑ x∈S L ( θ , φ , x ) . In the discrete VAE setting , the first term admits an analytical closed form gradient : ∂Ez∼qφ ( ·|x ) [ log pθ ( x|z ) ] ∂φ = Ez∼qφ ( ·|x ) [ log pθ ( x|z ) ∂ log qφ ( z|x ) ∂φ ] . ( 1 ) An exact computation of the expectation requires enumeration over all possible latent assignments since Ez∼qφ ( ·|x ) [ log pθ ( x|z ) ] = ∑ z∈Z qφ ( z|x ) log pθ ( x|z ) . Unfortunately , in the structured setting , the number of possible latent assignments is exponential . Instead , one can rely on the score function estimator ( REINFORCE ) to generate an unbiased estimate of the gradient by sampling from the distribution over the latent space . In many cases of interest , such as sampling spanning trees , the sampling algorithm is computationally unfavorable and suffers from high variance , leading to slow training and poor performance [ 46 ] . The Gumbel-Max reparametrization trick can trade summation with optimization . This approach is computationally appealing when considering spanning trees , since finding a maximal spanning tree is more efficient than sampling a spanning tree . Consider i.i.d . zero-location Gumbel random variables γ ∼ G ( 0 ) , e.g. , in the case of spanning trees γ = ( γ1 , ... , γ|E| ) consists of an independent random variable for each edge in the graph . Let qφ ( z|x ) = ez > hφ ( x ) where hφ ( x ) is a parametric encoder that learns edge scores and pθ ( x|z ) = efθ ( x , z ) , where fθ ( x , z ) denotes the log-probability learned by a parametric decoder . Then , the summation ∑ z∈Z −qφ ( z|x ) log pθ ( x|z ) can be approximated by the expectation Eγ∼G ( hφ ( x ) ) [ − fθ ( x , argmax z∈Z { z > γ } ) ] . ( 2 ) We provide the derivation in Appendix C. This formulation is a key to our proposed approach , as in some cases , estimating the above equation is easier . Computing the argmax can be accomplished efficiently even when the latent space is exponentially large . It is performed by utilizing sophisticated MAP solvers . For instance , finding the maximum spanning tree can be achieved in polynomial run time using Kruskal ’ s algorithm [ 27 ] . Sampling by perturbing the input and feeding it to a MAP solver is called perturb-and-map [ 47 ] . The gradient of Eq . 2 can be estimated using REINFORCE without needing to relax the argmax operator . By reparametrizing the Gumbel distribution , we get : ∂Eγ∼G ( hφ ( x ) ) [ fθ ( x , argmaxz∈Z { z > γ } ) ] ∂φ = Eγ∼G ( hφ ( x ) ) [ fθ ( x , argmaxz∈Z { z > γ } ) ∂ log G ( hφ ( x ) ) ( γ ) ∂φ ] . ( 3 ) Where G ( hφ ( x ) ) denotes the probability density function ( PDF ) of the Gumbel distribution with a location parameter of hφ ( x ) . The exponential summation in Eq . 1 is replaced with optimization , and now samples can be derived efficiently by applying the perturb-and-map technique . On the other hand , the disadvantages of REINFORCE remain as they were . The Gumbel-Softmax trick is a popular approach to reparameterize and optimize Eq . 2 . Since Eγ∼G ( hφ ( x ) ) [ fθ ( x , argmaxz∈Z { z > γ } ) ] = Eγ∼G ( 0 ) [ fθ ( x , argmaxz∈Z { z > ( hφ ( x ) + γ ) } ) ] , one can apply the Softmax to replace the non-differential argmax function . Hence , the Gumbel-Softmax trick replaces the function argmaxz∈Z { z > ( hφ ( x ) + γ ) } ) with the differential softmax function ez > ( hφ ( x ) +γ ) ∑ ẑ∈Z e ẑ > ( hφ ( x ) +γ ) , cf . Jang et al . [ 19 ] , Maddison et al . [ 34 ] . Under the structured setting , sophis- ticated extensions avoid the exponential summation in the partition function of the softmax operator [ 9 , 36 , 48 ] . For instance , in dependency parsing , Corro & Titov [ 9 ] construct the differentiable perturb-and-parse ( DPP ) method that exploits a differentiable surrogate of the Eisner algorithm [ 11 ] for finding the highest-scoring dependency parsing by replacing each local argmax operation with a softmax operation . Alternatively , Paulus et al . [ 48 ] utilize the Matrix-Tree theorem [ 24 ] for propagating approximated gradients through the space of undirected spanning trees . Natural Evolution Strategies ( NES ) is a class of gradient-free optimization algorithms . NES optimizes its objective function , by evaluating it at certain points in the parameter space . Consider a function k ( µ ) , may it be non-differentiable nor continuous , instead of optimizing k ( µ ) using a gradient method , NES optimizes a smooth version using the expected parameters of the function : g ( µ ) = Ew∼N ( µ , σ2I ) [ k ( w ) ] = ∫ Rd 1 ( 2πσ2 ) d 2 e− ‖w−µ‖2 2σ2 k ( w ) dw . ( 4 ) Here N ( µ , σ2I ) is a Gaussian distribution with mean µ and covariance σ2I . The expectation with respect to the Gaussian ensures the function g ( µ ) is differentiable , since e− ‖w−µ‖2 2σ2 is differentiable of any order , although k ( µ ) may not be differentiable . Following the chain-rule , the score function estimator for the Gaussian distribution determines the gradient : ∂g ( µ ) ∂µ = ∫ Rd 1 ( 2πσ2 ) d 2 e− ‖w−µ‖2 2σ2 ( w − µ σ2 ) k ( w ) dw . ( 5 ) That is , NES is an instance of REINFORCE , which optimizes a smoothed version of k ( µ ) by sampling from a distribution over parameter space rather than latent space . The reparameterization trick allows to further simplify the gradient estimator , with respect to a standard normal distribution : ∂g ( µ ) ∂µ = ∫ Rd 1 ( 2π ) d 2 e− ‖w‖2 2 ( w σ ) k ( µ+ σw ) dw = Ew∼N ( 0 , I ) [ w σ k ( µ+ σw ) ] . ( 6 ) The gradient can be estimated by sampling w repeatedly from a standard Gaussian : ∂g ( µ ) ∂µ ≈ 1 N ∑N i=1 wi σ k ( µ + σwi ) . The obtained estimator is biased when σ > 0 . However , the bias approaches 0 as σ → 0 . In practice , σ is assigned a small value , treated as a hyper-parameter . This algorithm is computationally appealing as it only uses the evaluations of k ( · ) to compute the gradient of g ( · ) . Moreover , NES is highly parallelizable , i.e. , one can compute the gradient ∇g ( µ ) at the time of evaluating a single k ( µ + σwi ) , in parallel for all i = 1 , ... , N , and then average these parallel computations [ 52 ] . Figure 1 depicts the parallel forward passes and the update rule according to NES . Theoretical guarantees regarding the convergence of gradient-free methods such as NES , were established for Lipschitz functions by Nesterov & Spokoiny [ 41 ] . In Section 3 , we extend the current zero-order optimization theory by presenting a convergence bound for NES over non-Lipschitz and bounded functions .

Summary: The paper proposes and analyses the use of Natural Evolutionary Strategies (NES) to train a VAE with a discrete latent space representation (Spanning Trees).  NES provides a smoothed approximation of the discrete function for which the gradient can be easily computed.  This approach can be seen as an example of the REINFORCE algorithm, although it clearly shows superior performance to traditional REINFORCE implementations.  The authors show that not only is it competitive and often marginally better than state-of-the-art techniques for this problem, it is also much more straightforward and easier to generalise.