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: In this paper, the authors use Natural Evolution Strategies (NES) to learn discrete structured VAE. They empirically show that optimizing discrete structured VAEs using NES is as effective as gradient-based approximations. And they also prove that NES converges for non-Lipschitz functions in discrete structured VAEs.