LEARNING DISCRETE STRUCTURED VARIATIONAL AUTO-ENCODER USING NATURAL EVOLUTION STRATEGIES
Year: 2022

Abstract:
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.

text: 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.

heading: 1 INTRODUCTION
text: 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.

heading: 2 BACKGROUND
text: 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.

heading: 3 STRUCTURED VAE OPTIMIZATION USING NES
text: In this work we suggest using gradient-free based method for learning discrete structured VAEs. Specifically, we propose using NES to optimize discrete structured VAEs without the need to propagate gradients through their latent discrete structures. For readability, we consider a concatenation of both θ and φ as µ = [θ;φ] where ; is the concatenation operation. For notational convenience, we refer θ as µ1 and φ as µ2.
Combining the discrete VAE objective function in Eq. 2, which is a non-continuous function due to the argmax operation and hence, non-Lipschitz, together with the NES objective in Eq. 4 (i.e., setting k(µ) = L(µ, x)) we get the following smooth approximation by setting g(µ) to be:
Ew∼N (µ,σ2I)Eγ∼G(hw2 (x)) [ − fw1(x, argmax z∈Z {z>γ}) ] , (7)
where w = [w1;w2] is the concatenation of vectors w1, w2 that denotes the decoder and encoder parameters respectively. Our goal is to minimize Eq. 7. Following the NES setup, we use Eq. 6 with a simplified notation to better emphasize the smoothing of µ. Overall its gradient takes the form of:
Ew∼N (0,I)Eγ∼G(hµ2+σw2 (x)) [ − w σ fµ1+σw1(x, argmax z∈Z {z>γ}) ] . (8)
We provide the pseudo code for using the NES algorithm to optimize discrete VAEs, together with their gradient update rule on Algorithm 1. We estimate the NES gradient in Eq. 8 by sampling w ∼ N(0, I) and γ ∼ G(hµ2+σw2(x)), which in turn induces a sampling of discrete structures argmaxz∈Z{z>γ}. This sampling procedure differs from that of the REINFORCE instances described in Eq. 1 and Eq. 3, as the samples of NES are tied to the sensitivity of the scoring function hµ2+σw2(x), i.e, to a random perturbation of its parameters µ2 by σw2. In contrast, the samples of REINFORCE are proportional to the scoring function hµ2(x) itself. In our experimental validation, we empirically demonstrate that NES has a lower variance.
Theoretical Guarantees. Next, we prove NES converges for a non-Lipschitz function k(·), which appears in reparameterized discrete VAEs in Eq. 2. In particular, we show that the norm of the parameters’ gradient can be arbitrarily small as training progresses. More formally, when the NES algorithm performs T update rules to the parameters µ, it generates the sequence µ(1), µ(2), ..., µ(T ), and for a sufficiently large T , there exists t ∈ {1, ..., T} for which ‖∇g(µ(t))‖ is arbitrarily small. For mathematical simplicity, we prove our convergence theorem on the expected gradient, as described in Eq. 8: µ(t+1) ← µ(t) − η∇g(µ(t)). To easily address ‖∇g(µ(t))‖ one usually considers the difference µ(t+1) − µ(t) using the remainder of the Taylor series: g(µ(t+1)) − g(µ(t)) = ∇g(µ(t))>(µ(t+1) − µ(t)) + 12 (µ
(t+1) − µ(t))>∇2g(µ̂)(µ(t+1) − µ(t)), where µ̂ ∈ [µ(t), µ(t+1)]. Here we denote by ∇2g(µ̂) the Hessian of g(µ̂). Applying the gradient update rule, we obtain the following equation:
g(µ(t+1))− g(µ(t)) = −η‖∇g(µ(t))‖2 + η 2
2 ∇g(µ(t))>∇2g(µ̂)∇g(µ(t)). (9)
A bound for ∇g(µ(t))>∇2g(µ̂)∇g(µ(t)) is a key to our convergence theorem. Our bound relies on the fact that the discrete VAE objective, which is given in Eq. 2, is a non-negative function that is continuous almost everywhere.
Algorithm 1 Natural Evolution Strategies for discrete VAEs Input: Initial parameters µ. repeat
for all i = 1 to N do Sample w̃i ∼ N (0, I|µ|) Evaluate ui = Eγ∼G(hµ2+σw̃2,i (x)) [ − fµ1+σw̃1,i(x, argmaxz∈Z{z>γ}) ] end for Update µ← µ− η · 1
N ∑N i=1 w̃i σ ui
until a stopping condition is met
Lemma 1. Let k : Rd → [0,M ] be a non-negative function that is continuous almost everywhere and let g(µ) = Ew∼N (µ,σ2I)[k(w)]. Then for any µ1, µ2 there holds:
∇g(µ1)> ( ∇2g(µ2) ) ∇g(µ1) ≤ dM3
σ4 . (10)
Proof can be found in Appendix A.1. The above lemma together with Eq. 9 imply the following bound: g(µ(t+1)) − g(µ(t)) ≤ −η‖∇g(µ(t))‖2 + η2 dM 3
2σ4 . By summing over all algorithm steps t = 1, ..., T we show that the average norm of the gradient can be arbitrarily small for a sufficiently large T . Theorem 1. Under the conditions of Lemma 1 there holds:
1
T T∑ t=1 ‖∇g(µ(t))‖2 ≤ M ηT + ηdM3 2σ4 . (11)
Moreover, when setting η = √ 2σ4
TdM2 then: 1 T ∑T t=1 ‖∇g(µ(t))‖2 ≤ √ 2dM4 Tσ4 .
Therefore, there exists t for which ‖∇g(µ(t))‖2 ≤ √ 2dM4
Tσ4 .
Proof is given in Appendix A.2. Intuitively, the above theorem proves that the NES algorithm converges on discrete VAE to a stationary point, even when the original function k(·) is non-continuous and hence non-Lipschitz. This is in contrast to the contemporary trend that relies on Lipschitz functions [41]. Instead, we rely on the non-negativity of the discrete VAE objective. We note that a stationary point of g(µ) = Ew∼N (µ,σ2I)[k(w)] is not necessarily a stationary point of k(µ). Nevertheless, since limσ→0 g(µ) = k(µ) almost everywhere, except perhaps for non-continuous points, a low value of g(µ) is correlated with a low value of k(µ).

heading: 4 EXPERIMENTS
text: We start by experimentally validating our approach by learning discrete structured VAEs for latent structure recovery in Section 4.1 and dependency parsing in Section 4.2. Next, in Section 4.3 we analyze how NES scales with the latent space dimension and neural network size. We additionally provide an analysis for non-Lipschitz functions in Appendix B, together with analyzing the results concerning the theoretical guarantees as presented in Section 3.
In our experiments, we use a variance reduction technique called mirrored sampling [14, 4]: that is, on each NES iteration, we use a single Gaussian noise vector w to create two parameter sets, one by adding and the other by subtracting the noise vector. Thus, on each iteration, we sample N2 Gaussian noise vectors where N is the number of the VAE parameter sets utilized for estimating the NES update direction. Additionally, to make NES more robust and scale-invariant, we transform the outputs of the perturbed forward passes into standard scores by subtracting their mean and dividing them by the standard deviation. All reported values are measured on a test set, and the models were selected using early stopping on the validation set. All the following experiments were conducted using an internal cluster with 4 Tesla-K80 NVIDIA GPUs.

heading: 4.1 LATENT STRUCTURE RECOVERY
text: We begin by demonstrating the capability of NES to learn the internal structure of an interacting system based on graphs in an unsupervised fashion. The interplay of group components, e.g., basketball players on the court or a flock of birds during migration, can often be explained using a
simple structure. However, frequently, we only have access to individual trajectories without knowledge of the underlying interactions. The Neural Relational Inference (NRI) model [23], which we base on, is designed to infer these interactions purely from observational data. NRI takes the form of a VAE, where the encoder produces a distribution over the space of interaction structures given the component trajectories, and the reconstruction is based on graph neural networks.
In our experiments, we utilize the dataset developed by Paulus et al. [48], where the target structure is a spanning tree over 10 vertices (each vertex represents an individual component). The model attempts to learn the true tree structure that defines the interplay among the group by only observing the locations of the 10 components during several timesteps. We focus on two cases, as suggested by Paulus et al. [48]. In the first case, we use our prior knowledge regarding the true structures and define the latent space as the space of spanning trees over a 10-nodes undirected graph, which consists of 108 possible spanning trees. In the second more challenging case, we remove the tree constraint and consider all possible n − 1 unique edge combinations as the latent space, where n denotes the number of vertices. The ability to recover the underlying structure is measured as the edge F1-score against the target spanning tree.
In each of the two cases, we compare NES with the corresponding Stochastic Softmax Trick (SST) [48]. SSTs are the generalization of the Gumbel-Softmax Trick (GSM) for combinatorial discrete distributions. That is, in contrast to GSM, SSTs are designed to optimize over exponentially large discrete spaces. We run our experiments with the same set of parameters as in Paulus et al. [48], except that during decoding we use teacher-forcing every 3 steps instead of 9 steps. We fix NES parameters to be σ = 0.01 and N = 600. We additionally compared our method against four instances of REINFORCE. Each utilizes a different variance reduction technique. The first is NVIL [39], which uses two control variates. The remaining three reduce variance by subtracting the following control variate from the learning signal: EMA uses the exponential moving average of the ELBO, Batch uses the mean ELBO of the current mini-batch, and lastly, Multi-sample, which is especially well suited for structured prediction [25, 26], uses the mean ELBO over r multiple samples per data point. Paulus et al. [48] tuned r on the set of {2, 4, 8}. Results are listed in Table 1. In terms of structure recovery, the gradient-free NES outperforms all REINFORCE instances. Moreover, when considering SST, NES achieves superior edge F1-score with slightly worse ELBO values. This is surprising since SST is a gradient-based method that generalizes the effective GSM estimator to combinatorial spaces. Also, unlike SST, which requires a carefully tailored solution for each structure, NES is simple and generic. For a fair comparison with NES, we also ran the Multi-sample method with 400 Monte Carlo samples per data point and a mini-batch size of 4 (as using larger r or larger mini-batch size has exceeded the GPU memory limit). However, the results were inferior to those achieved by tuning r on {2, 4, 8}. This is in line with the results obtained by Kool et al. [25], where larger values of r led to inferior results. Thus, it can be concluded that even with an equal computational cost, the REINFORCE instances are inferior to NES.

heading: 4.2 DEPENDENCY PARSING
text: Next, we evaluate the capability of NES in learning latent projective and non-projective dependency parse trees as part of an unsupervised domain adaptation task. Unlike Section 4.1, where we focused on structures over undirected edges, here we focus on dependency trees which are rooted directed spanning trees. Our model is based on a VAE architecture similar to that of differentiable perturband-parse (DPP) [9]. The encoder is comprised of a graph-based parser [22] that decomposes the
score of a tree to the sum of the scores of its arcs and produces a distribution over the space of dependency trees. Sampling from the latent space is performed using the “perturb-and-map” technique, where each arc score is perturbed independently with a noise derived from a Gumbel distribution. Then, the perturbed arc scores are fed into a MAP solver, which outputs the highest-scoring tree (the resulting sample). For projective dependency parsing, we utilize the Eisner algorithm [11] as the MAP solver. Similarly, for non-projective dependency parsing, we use the Chu-Liu-Edmonds (CLE) algorithm [6, 10]. The decoder is modeled as a language model, that given a latent dependency tree, attempts to reconstruct the input sentence.
We compare NES with two strong baselines. For projective dependency parsing, we consider the DPP model. DPP optimizes the VAE by utilizing a differentiable surrogate of the Eisner algorithm. In that manner, it tackles both the differentiability and exponential enumeration issues. In fact, DPP can be seen as a Stochastic Softmax Trick (SST). For non-projective dependency parsing, we consider SparseMAP [43] as a baseline by replacing the CLE algorithm with a SparseMAP layer. SparseMAP uses the active set method that performs sequential calls to a MAP solver and returns a sparse linear combination of several high-scoring structures. This procedure is differentiable almost everywhere but computationally inefficient due to its sequential nature. Unlike these methods, NES does not require any modification to the architecture.
We perform extensive experiments on the task of unsupervised domain adaptation for dependency parsing. We consider the Universal Dependencies (UD) dataset [35, 44, 45]. UD is a multilingual corpus annotated with dependency trees in more than 180 treebanks of over 100 languages. We follow the setup of Rotman & Reichart [50] and choose 3 distinct languages, considering 2 distinct treebanks from different domains for each: Galician (GL CTG: science and legal, GL TREEGAL: news), Indonesian (ID CSUI: news, ID GSD: general) and Russian (RU GSD: general, RU TAIGA: social media, poetry, and fiction). We conduct 6 domain adaptation experiments, where we alternate between the source and target domain in each language. We consider the training set of our source domain as our labeled dataset and the training set of the target domain as the unlabeled dataset.
At first, we train the VAE components separately on the labeled set (source domain) for 30 epochs. Then, we optimize the pretrained VAE on the unlabeled set (target domain) for 10 additional epochs using the NES algorithm. For a fair comparison, we perform the same training procedure for the above-mentioned baselines. We set the hyper-parameters to those of the original implementation of Kiperwasser & Goldberg [22] and feed the models with the multilingual FastText word embeddings [16]. We perform a grid-search for each of the methods separately over learning rates in [5 ·10−4, 1 · 10−5] and set the mini-batch size to 128. We fix NES parameters to be σ = 0.1 andN = 400. Adam optimizer [21] is used to optimize all methods. The models we selected were those who obtained the best unlabeled attachment score (UAS) on the source domain validation set.
Table 2 summarizes the results on the UD treebanks in terms of unlabeled attachment score (UAS). The scores under each treebank name reflect performances on the setup where the treebank is set to be the target domain. Results suggest that NES reaches comparable performance (with a minor improvement) to SparseMAP and DPP while being simpler and more flexible to use. Note that unlike SparseMAP and DPP, which use sequentially complex methods to either infer the highest-scoring tree structure or to propagate gradients through a bottleneck dynamic programming algorithm, NES can optimize the model in parallel without the necessity of gradient computations.

heading: 4.3 SCALABILITY ANALYSIS
text: Latent space size. In the following set of experiments, we further investigate the properties of NES and the several methods it was compared to in Section 4.1 and 4.2. Specifically, we examine how
the latent space size affects the method’s run-time by measuring the methods wall-clock time of a forward and backward pass as a function of the input dimension (denoted by n). Note that the latent space size grows exponentially with the model input dimension, e.g., for a sentence of length n, the latent space of the VAE architecture presented in Section 4.2 is the space of all possible dependency trees over an n-nodes directed graph.
For each n ∈ {10, 20, . . . , 250} we create a random dataset. Specifically, for the NRI model (Section 4.1), the input is n trajectories of 10 timesteps derived from a standard Gaussian distribution. We compare the run-time of NES with that of SST and a REINFORCE instance that relies on Eq. 1, where sampling is performed using a Markov chain Monte Carlo (MCMC) algorithm. For the parsing model (Section 4.2), we derive 32 random sentences of length n by randomly sampling words from a vocabulary of size 10000. As DPP utilizes a differentiable surrogate of the Eisner algorithm, it is compared to NES with the Eisner algorithm as the MAP solver. Similarly, SparseMAP is compared to NES with the CLE algorithm. Since our internal cluster consists of 4 GPUs, we utilize NES with N = 4 for a fair comparison with the gradient-based methods. Finally, we run the experiments over various random seeds and average the wall-clock time. Figure 2 depicts the results.
As can be seen, the run-times of DPP and SparseMAP heavily rely on the input dimension and grow at a much higher rate than the run-time of NES. NES also scales better than SST and REINFORCE on the NRI model. However, in this case, the gap is smaller as enlarging the input dimension of the latent structure also enlarges the model size which NES updates depend on. Overall, it can be seen that NES scales well with the latent space size in contrast to most of its competitors.
Neural network size. Next, we conduct a study that examines how the enlargement of a neural network affects the number of NES samples needed to optimize it. We begin by optimizing a VAE of 25K parameters, then we enlarge its parameter size by a factor of 2 and optimize the resulting model. We repeat this process several times up to a model of 800K parameters. We utilize SST with a fixed temperature of 1 as a baseline. For each model size, we examine how many NES samples are needed to achieve test ELBO as lower as the one achieved by the SST. Results are depicted in Table 3. A detailed description of the experimental setup can be found on Section B in the Appendix.
We observe that enlarging the neural network by a factor of 2 does not necessarily mean that we should enlarge N in the same manner. To be precise, in all our experiments, we do not need to enlarge N with more than 50% samples when optimizing the two times larger network. These observations are positive and suggest that NES can scale well with the network size.

heading: 5 RELATED WORK
text: Jang et al. [19], Maddison et al. [34] proposed the GSM estimator that replaces the non-differentiable argmax operation with a differentiable softmax operation. However, structured latent spaces can be exponentially large and the softmax opertation becomes computationally intractable. Other works proposed tailor-maid solutions for specific structures. For instance, Corro & Titov [9] focused on latent projective dependency trees and propagated gradients through a differentiable surrogate of Eisner algorithm [11]. Mena et al. [36] extended the Gumbel-Softmax estimator [19] and proposed the Gumbel-Sinkhorn method for learning latent permutations. Paulus et al. [48] took these ideas one step further and proposed a unified framework for designing structured relaxations of combinatorial
distributions. Unlike these methods, our approach is generic and as such, it can be applied to general structures with no additional effort, since it obviates the need for a differentiable surrogate of the linear maximization oracle. Mensch & Blondel [37] proposed a framework for turning dynamic programming algorithms differentiable.
Others have taken a more generic approach. For example, SparseMAP [43, 42, 8] is a framework for training deep networks with sparse structured hidden layers, solved by sequential calls to a MAP oracle. In a similar sense, Itkina et al. [18] suggest using evidential theory to perform post hoc latent space sparsification and thus reducing the discrete latent sample space at test time. Chen et al. [5] generalized this method to a sparse normalization function which can be applied during both training and test time. Moreover, a recent line of works propagates gradients through the non-differentiable argmax operation, Lorberbom et al. [33] use the difference of two maximization operations, and the method of Berthet et al. [1] is based on integration by parts. Contrarily, our approach does not require constructing sophisticated solutions to propagate gradients through discrete operations, which makes it both simple and flexible.
The Vector Quantized Variational Auto-Encoder (VQ-VAE) [56, 49] introduces an alternative approach to learning discrete latent representation. However, VQ-VAE differs from our discrete structured VAEs in an important aspect. In our setting, we know the structure of the latent space, e.g., the space of all possible spanning trees in a given graph. Hence we do not perform unsupervised vector quantization as in VQ-VAE but rather use a predetermined quantization over the set of possible structures. In this work, we rather focus on an alternative optimization method for learning discrete latent structures.
Recently, black-box optimization methods have been applied to neural networks [32, 52, 31, 61, 30, 53, 38]. Salimans et al. [52] showed that NES is a competitive alternative to popular RL techniques. Moreover, they utilized the fact that NES is highly parallelizable and proposed a generic distributed version of NES that scales well with the number of CPUs. Lenc et al. [31] proposed a hybrid method that alternates between NES and SGD for training large sparse models. Finally, Zhang et al. [61], Lehman et al. [30] compare the relation between the SGD gradients and NES updates. To our knowledge, we are the first to apply NES to structured VAEs.

heading: 6 CONCLUSION
text: We suggested using NES, a class of gradient-free black-box algorithms, as an alternative for learning discrete structured VAEs. We have demonstrated empirically that NES performs substantially better than various REINFORCE instances and even better than SST on the structure recovery task while being simpler and more robust. Moreover, NES achieves better or comparable performance to DPP and sparseMAP when considering dependency tree latent structure. However, as opposed to the aforementioned methods, NES does not require complex solutions for propagating gradients through the discrete structures, which makes it more generic, flexible, and simple to implement. Additionally, we showed that NES scales well with the latent space dimension and neural network size. To establish the theoretical soundness of our approach, we proved that NES converges for non-Lipschitz functions such as the objective function of a discrete VAE.
In this study, we have limited the expressive power of the NES method by fixing the covariance matrix of the Gaussian search distribution. For future work, we would like to explore the effect of jointly optimizing the covariance and mean of the distribution of parameters.

heading: A PROOFS OF NES CONVERGENCE FOR NON-LIPSCHITZ FUNCTIONS
text: A.1 PROOF OF LEMMA 1
Proof. The proof follows two main steps: (i) the spectral norm of ∇2g is at most M/σ2 and (ii) ‖∇g(µ1)‖2 ≤ dM2/σ2. Since the Hessian matrix is symmetric, we can use Rayleigh quotient and obtain:
∇g(µ1)> ( ∇2g(µ2) ) ∇g(µ1) ≤ |λmax| · ‖∇g(µ1)‖2, (12)
where λmax is the largest eigenvalue of∇2g(µ2). Since∇2g(µ2) is symmetric it can also be shown that:
|λmax| = ‖∇2g(µ2)‖ = max ‖s‖=1 |s>∇2g(µ2)s|. (13)
Applying the log derivative trick on Eq. 6, we obtain ∇2g(µ2) = ∫ Rd 1 (2π) d 2 e− ‖w‖2 2 ( 1 σ2 ww> ) k(µ2 + σw)dw. (14)
Therefore:
max ‖s‖=1 |s>∇2g(µ2)s| = max ‖s‖=1 | ∫ Rd 1 (2π) d 2 e− ‖w‖2 2 ( 1 σ2 s>ww>s ) k(µ2 + σw)dw|
≤ M σ2 max ‖s‖=1 | ∫ Rd 1 (2π) d 2 e− ‖w‖2 2 (s>w)2dw|
= M
σ2 max ‖s‖=1 ‖s‖2 = M σ2 .
(15)
Since k(·) is bounded by M and w are i.i.d. normal Gaussian random variables, therefore for any s the random s>w is a Gaussian with zero mean and variance ‖s‖2. By combining Eq. 13 with Ineq. 15 we get:
‖∇2g(µ2)‖ ≤ M
σ2 , (16)
which concludes step (i). Next, we bound the squared norm of the gradient: ‖∇g(µ1)‖2 = ‖ ∫ Rd 1 (2π) d 2 e− ‖w‖2 2 (w σ ) k(µ1 + σw)dw‖2
= 1 σ2 ‖ ∫ Rd 1 (2π) d 2 e− ‖w‖2 2 wk(µ1 + σw)dw‖2
≤ 1 σ2 ∫ Rd ‖ 1 (2π) d 2 e− ‖w‖2 2 wk(µ1 + σw)‖2dw
≤ M 2
σ2 ∫ Rd 1 (2π) d 2 e− ‖w‖2 2 ‖w‖2dw = dM 2 σ2 ,
(17)
where the first inequality is obtained using the Cauchy-Schwarz inequality and the second inequality by bounding k(·) with M . Thus, overall we have showed that:
∇g(µ1)>∇2g(µ2)∇g(µ1) ≤ M
σ2 · ‖∇g(µ1)‖2
≤ dM 3
σ4 .
(18)
A.2 PROOF OF THEOREM 1
Proof. Rearranging Eq. 9, we obtain:
η‖∇g(µt)‖2 ≤ g(µ(t))− g(µ(t+1)) + η2 dM 3
2σ4 . (23)
Summing over all algorithm steps t = 1, . . . , T , we have:
η T∑ t=1 ‖∇g(µt)‖2 ≤ T∑ t=1 [g(µ(t))− g(µ(t+1))] + η2 dM 3T 2σ4 . (24)
Opening the telescopic sum:
η T∑ t=1 ‖∇g(µt)‖2 ≤ g(µ(1))− g(µ(T+1)) + η2 dM 3T 2σ4 . (25)
Since k(·) is non-negative and bounded, the difference between g(µ(1)) and g(µ(T+1)) is bounded from above by M :
η T∑ t=1 ‖∇g(µt)‖2 ≤M + η2 dM 3T 2σ4 . (26)
We multiply both sides of the inequality by 1ηT :
1
T T∑ t=1 ‖∇g(µt)‖2 ≤ 1 ηT [ M + η2 dM3T 2σ4 ] ≤ M ηT + η dM3 2σ4 . (27)
Next, we minimize the right-hand size of the inequality in η:
η∗ =
√ 2σ4
TdM2 , (28)
and plug it back to Ineq. 27:
1
T T∑ t=1 ‖∇g(µt)‖2 ≤ √ 2dM4 Tσ4 , (29)
Then, for arbitrarily small δ > 0 such that:
1
T T∑ t=1 ‖∇g(µt)‖2 ≤ δ ≤ √ 2dM4 Tσ4 , (30)
there exists t for which:
‖∇g(µt)‖2 ≤ δ, (31)
after at most
T ≤ 2dM 4
δ2σ4 , (32)
steps.

heading: B ADDITIONAL RESULTS
text: In the following experiments, we define the encoder as input⇒MLP (α)⇒ ReLU ⇒MLP (10) ⇒ argmax, and the decoder as MLP (α)⇒ ReLU ⇒MLP (β)⇒ output, where β is the input dimension. Unless otherwise stated, α = 300.
B.1 NEURAL NETWORK SIZE
The experiments were conducted on the FashionMNIST dataset [59] with fixed binarization [51]. We tune α for enlarging the VAE. Particularly, α is picked from the set [16, 32, 64, 128, 256, 512]. All models were trained using the ADAM optimizer [21] over 300 epochs with a constant learning rate of 10−3 and a batch size of 128.
B.2 RELATION BETWEEN g(·) AND k(·)
In section 3, we prove that under the conditions of Lemma 1, NES converges to a stationary point of g(µ) for non-Lipschitz functions. To empirically explore the relation between g(µ) and k(µ), we conduct a set of experiments in which we demonstrate that a low value of g(µ) is correlated with a low value of the objective function k(µ) by estimating the average absolute distance between them.
First, we estimate the Gaussian approximation g(µ) for each sample in the test set by perturbing the current model parameters 1000 times, computing the ELBO for each perturbed parameter vector and average. Then, we calculate the absolute difference between the ELBO, serving as the objective function, and the estimated Gaussian approximation and average over the tested samples. We experiment with three different NES configurations: N = 300 and σ ∈ {0.01, 0.5, 0.1} on the FashionMNIST dataset [59]. The results presented on the left image in Figure 3 indicate that the smaller σ is, the further the proximity between k(µ) and g(µ). It can also be seen that the average distance converges and stabilizes as the learning progresses towards saturation.
B.3 BOUNDNESS ASSUMPTION
In the general case, the objective function of discrete VAEs is not bounded from above in contrast to Theorem 1 assumption. However, it can be upper bounded by bounding each log probability component with a constant. For ease of explanation, we scale the ELBO by dividing it with the VAE output dimension. Then, we upper bound it with M = {1, 3, 9} during training and compare the test ELBO with that of a model trained with an unbounded ELBO, denoted by UNBOUNDED. We train the three models with N = 300 and σ = 0.1 on the FashionMNIST dataset. Results are depicted on the right image in Figure 3. It can be seen that bounding the loss has a minimal effect on model performance when M is big enough. Increasing M improves the performance, while using a relatively small M value may cause the model to diverge. Surprisingly, bounding the loss with M = 9 leads to a slightly lower loss compared to the UNBOUNDED baseline. We hypothesize this is due to a regularization effect.
B.4 UNSTRUCTURED TASKS.
Unlike structured VAEs, where the latent spaces are often exponentially large, here we explore a latent space that consists of only 10 different assignments. Therefore, an unbiased gradient of the objective with respect to the VAE parameters can be analytically computed by enumerating over all possible latent assignments (Eq. 1). We denote this method as UNBIASED and consider the loss of a model trained with this method as a lower bound for the loss of the same model trained with NES. Due to the relatively small latent space, we can also compare with a Gumbel-Softmax (GSM) biased estimator. For NES, the VAE is trained with σ = 0.1 and N = 300. For GSM, we use the annealing schedule of Jang et al. [19].
Experiments are conducted on the FashionMNIST [59], KMNIST [7], and Omniglot [29] datasets with fixed binarization [51]. All models are trained using the ADAM optimizer [21] with a constant learning rate of 10−3 and a mini-batch size of 128. Figure 4 depicts the negative ELBO of NES and its competitors.
Surprisingly, NES achieves competitive results compared to the UNBIASED method on all of the three benchmarks. On KMNIST, NES significantly outperforms GSM, and on the FashionMNIST and Omniglot benchmarks, it achieves comparable results. This is despite the fact that NES optimizes the VAE parameters by only evaluating the model at certain points in parameter space.

heading: C DERIVING EQUATION 2
text: Let γ be a random function that associates an independent random variable γ(z) for each input z ∈ Z . When the random variables follow the Gumbel distribution law with mean hφ(x, z), which we denote by G(hφ(x, z)) and whose probability density function is gz(γ) = e−(γ(z)+c−hφ(x,z)+e −(γ(z)+c−hφ(x,z))) for the Euler constant c ≈ 0.57. Then for g(t) = ∏k z=1 gz(t) we obtain the following identity:
ehφ(x,z) = Pγ∼g[z∗ = z], where z∗ , argmax ẑ∈Z {γ(ẑ)}. (33)
Proof. Let Gz(t) = e−e −(t+c−hφ(x,z)) be the Gumbel cumulative distribution function. Then
Pγ∼g[z∗ = z] = Pγ∼g[z = arg max ẑ=1,...,k {γ(ẑ)}]
= ∫ gz(t) ∏ ẑ 6=z Gẑ(t))dt. (34)
Since gz(t) = e−(t+c−hφ(x,z))Gz(t) it holds that
∫ gz(t) ∏ ẑ 6=z Gẑ(t)dt = ∫ e−(t−hφ(x,z)+c)Gz(t) ∏ ẑ 6=z Gẑ(t)dt
= ehφ(x,z)
Z ,
(35)
where 1Z = ∫ e−(t+c) ∏k ẑ=1Gẑ(t)dt is independent of z. Since Pγ∼g[z = z∗] is a distribution then
Z must equal to ∑k ẑ=1 e hφ(x,ẑ).
Next, we use the Gumbel-Max trick to rewrite the expected log-likelihood in the ELBO in the following form:
Ez∼qφ log pθ(x|z) = ∑ z∈Z Pγ∼g[z∗ = z]fφ(x, z) = Eγ∼g[fθ(x, z∗)]. (36)
The equality results from the identity Pγ∼g[z∗ = z] = Eγ∼g[1z∗=z], the linearity of expectation∑ z∈Z Eγ∼g[1z∗=z]fφ(x, z) = Eγ∼g[ ∑ z∈Z 1z∗=zfφ(x, z ∗)] and the fact that ∑ z∈Z 1z∗=z = 1.
When z = (z1, ..., z|E|) is a spanning tree, or more generally, belongs to the a structured space, one cannot assign an i.i.d. random variable to each z ∈ Z . Instead we relate a random variable γ = (γ1, ..., γ|E|) and set γ(z) = z>γ.

heading: D THE FULL OBJECTIVE FUNCTION
text: In Eq. 7 we didn’t include the KL-divergence term for the sake of simplicity. In practice, the NES algorithm optimizes both terms. Thus, for completeness we provide the full NES objective. For the avoidance of doubt, in our experiments, we optimized both terms.
Assuming that pθ(·) is the uniform distribution over the space of structures, the Gumbel-Max reparameterization trick let us derive the following approximation:
KL(qφ(·|x)||pθ(·)) = Ez∼qφ(·|x) [ log
qφ(z|x) pθ(z)
] ≈ Eγ∼G(hφ(x))[z ∗>hφ(x)− log 1
|Z| ], (37)
where z∗ = argmaxz∈Z{z>γ}. The resulting NES objective is:
Ew∼N(µ,σ2I)Eγ∼G(hw2 (x)) [ − fw1(x, z∗) + z∗>hw2(x)− log 1 |Z| ] . (38)
And its gradient takes the form of:
Ew∼N(0,I)Eγ∼G(hµ2+σw2 (x)) [w σ ( − fµ1+σw1(x, z∗) + z∗>hµ2+σw2(x)− log 1 |Z| )] , (39)
where w = [w1;w2] is the concatenation of the two vectors w1, w2.

Review summaries:
In this paper, the authors proposed to use the Natural Evolution Strategy (NES) algorithm for learning discrete structured VAEs. This algorithm estimate gradients with forwarding pass evaluations only and do not require propagating gradients through their discrete structures. Authors showed empirically that optimizing discrete structured VAEs using NES is as effective as gradient-based approximations.

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.

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.