Learning Discrete Structured Variational Auto-Encoder using Natural Evolution Strategies
Authors: Alon Berliner
Guy Rotman
Yossi Adi
Roi Reichart
Tamir Hazan
Venue: ICLR 2022 Poster (2022)
Keywords: structured prediction
derivative-free optimization
variational autoencoder

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.

Published as a conference paper at ICLR 2022
LEARNING DISCRETE STRUCTURED VARIATIONAL AUTO-ENCODER USING NATURAL EVOLUTION STRATEGIES
Alon Berliner
Guy Rotman Technion, IIT zotmanguy@gmail com
Yossi Adi
Technion, IIT
Meta Al Research adiyoss fb. com
aion.berlineregaill com
Roi Reichart Technion, IIT
Tamir Hazan Technion, IIT tamir www.n@tecciionn..a....
roiri@technion.ac.if
ABSTRACT
Discrete variational auto-encoders (VAEs) are able 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 gra- dients without enumerating over the space of possible structures. In this work, we use Natural Evolution Strategies (NES), a class of gradient-free black-box opti- mization algorithms, to Iearn discrete structured VAEs. The NES algorithms are computationally appealing as they estimate gradients with forward pass evalua- tions only, thus they do not require to propagate gradients through their discrete structures. We demonstrate empirically that optimizing discrete structured VAEs using NES is effective gradient-baser based approximations. Lastly, prove NES converges for non-Lipschitz functions as appear in discrete structured VAEs.
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]
[3], molecular synthesis (28 15| 48), speech
and visual understanding 36 55 3. Compared their continuous counterparts, they improve
interpretability by illustrating which terms contributed to the solution [48] 44), and they can facilitate the encoding of inductive biases in the learning process, such as images consisting of small number
of objects 12 or tasks requiring intermediate alignments [36 42 2
Learning VAEs with discrete n-dimensional latent variables 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 computation of the model value. Unfortunately, the resulting mapping remains non-differentiaale due to the presence of arg max operations. In order to propagate gradients efficiently, Jang al. 199, Maddison et al. [34 proposed the Gumbel-sofftax reformulation that uses smooth relaxation of the reparametrized objective, replacing the arg max operation with softmax operation. Following such an approach may bring back the need for enumerating over 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 To better deal with the computational complexity in the structured setting, sophisticated stochastic softmax tricks were devised to learn discrete structured VAEs 48 (c.g., perturb-and seirrb-and-aars. parse for dependency parsing by 99,
for bi-partite matching (3)

---

algorithm for learning Published as a conference paper at ICLR 2022
In this work propose use the Natural Evolution Strategy (NES) [57 58 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 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 propagate gradients through the arg max operation, as only requires to evaluate the model. Moreover, the proposed approach is highly parallelizable, hence computationally appealing.
methods
Our contributions: (1) We suggest using black-box, gradient-frree based optimization methods, specifically NES, to optimize discrete structured VAËs. (2) We experimentally demonstrate that NES, which uses the models' output in black-box manner, is as effective as gradient based approx- imations 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 generative model (z) using training set S = {.I..... }. , derived from an unknown distribution p(x)) by minimizing its nega- tive log-liielioood. VAËS rely on latent variable models of the form p0( = EeEz (=)=)oo()))) z). where: 2 is a realization of the latent variable and 2 is the discrete set of its possible assignments. We focus on discrete structures such as spanning trees in graph G (V, i... = (21, (1)) represents a spanning tree I for which = 1 if the edge e E T belongs to the spanning tree T and zero otherwise. In this case, Z is the spanning trees space which typically exponential in the size of the input, as there are |V] spanning trees for a complete graph with [V] vertices.
VAES rely on an auxiliary distribution (0(2|z)) that is used to upper bound the negative log- likelihood of the observed data points: Exes log (o S >xxss [.00, o,),, where: I((0, x)
- -Ez-ge (oog (2) [x]=2] + K Lgg (==oo)))ppp(())) This formulation is known the negative Evidence Lower Bound (ELBO) [2).. where the KL--divergence measures the similarity of two distributions 46 and pg. and is defined as: (-<(((())))l = E==gs((|x)) [log(gg htop))pPpp)))))) (=))]-
Parameter estimation is generally carried out by performing gradient descent on Exes L(0, In the discrete VAE setting, the first term admits an analytical closed form g
OE BBgrgo [log (x][=]
2 log go (2(x)) ab
log (t
(1)
An exact computation of the expectation requires enumeration over all possible latent assignments since (- (e) [logpo((s))) (x]==)] = =x (=(2xx) log po (ar]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 computationally unfavorable and suffers from high variance, leading to slow training and poor performance [46].
The Gumbel-Mat 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 : ~ (0). e.g., in the case of spanning trees = (?1, T|E)) consists of an independent random variable for each edge in the graph. Let ()(2)x)) = e (x( where has (xr) is a parametric encoder that learns edge scores and p=(2)) = ef((z,2), where fo z) denotes the log-probability
learned by parametric decoder. Then, the summation Exez 4o (2 (x) log po (x/22 can be approx-
imated by the expectation
(yret [ fo (2, arg max 2 TT)].
2E2
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 arg max can be accomplished
2

---

Published as a conference paper at ICLR 2022
efficiently even when the latent space is exponentially large. It is performed by utilizing sophisti- cated MAP solvers. For instance, finding the maximum spanning tree can be achieved in polynomial time using Kruskal's algorithm [27]. Sampling by perturbing the imput and feeding it to MAP
solver is called perturb-and-map [47]. The gradient of Eq. 2 can be estimated using REINFORCE
without needing relax the arg max operator. By reparametrizing the Gumbel distribution, we get:
((cgg(g()))) fo( arrg arg massc 2 fs m))]
(nre(ha(a)) (Z, arg max.ez 7))) 8 (o (3)
de
de
Where (z)( denotes the probability density function (PDF) of the Gumbel distribution with a location parameter of (z)r). The exponential summation in [ T 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 En [ arg max 7})] 1g [Jr arg max (= (h) + +))) one can apply the Softmax to replace the non- non-differrenial arg max function. Hence, the Gumbel-Softma
trick replaces the function arg max (h₂(x) 7)}) with the differential softmax functionn
(1g()))7)
Diege e (ho60049)) 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 opera- tor 99 36 48). For instance, in dependency parsing, Corro & Titov [9] construct the differentiable perturb-ande -nd--arsee (DPP) method that exploits differentiable surrogate of the Eisner algorithm [[T] for finding the highest-scoring dependency parsing by replacing each local arg max operation with a softmax operation. Alternatively, Paulus et al. 48 utilize the Matrix-Tres theorem 24 for prop- agating approximated gradients through the space of undirected spanning trees.
Natural Evolution Strategies (NES) a class of gradient-free optimization algorithms. NES op- timizes its objective function, by evaluating it at certain points in the parameter space. Consider a function k(p),, may be non-differentiable nor continuous, instead of optimizing k(u) using a gradient method, NES optimizes a smooth version using the expected parameters of the function:
g(p) = Ew-NN( [k(w)] = Jee (2TO2)
1 I=== 202
k( k(w)dw.
(4)
Here (M,, 0²1) is a Gaussian distribution with mean and covariance 0 The expectation with |==-ul? 2n
respect to the Gaussian ensures the function g(1) is differentiable, since of any order, although k(xz) may not be differentiable. Following the chain-rule, the score function estimator for the Gaussian distribution determines the gradient:
is differentiable
0g(n) 8p
( k(w)dw.
(5)
Rd (2TO2)
That is. NES is instance of REINFORCE, which optimizes smoothed version of k(μ( by sam- pling 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(u) 0p Jgs (2x))
-
1
1n11 ("" k(u + ow)dw = EE-NNo. k(p k(u+ + ow)]
(6)
The gradient can be estimated by sampling w repeatedly from a standard Gaussian:
0g(u)
Up
N EN'
11 We k* + ow.). The obtained estimator is biased when a > O. However, the bias ap- proaches a + O. In practice, a is assigned 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.c., one can compute the gradient Vg(H) at the time of evaluating a single k(μ + aww), in parallel for all i = I,..., N, and then average Figure depicts the parallel forward passes and the update rule
these parallel computations [22]
according to NES.
Theoretical guarantees regarding the convergence of gradient-fre free methods such as NES, were es- tablished for Lipschitz functions by Nesterov & Spokoiny [4]]. In Section B we extend the current zero-order optimization theory by presenting a convergence bound for NES over non-Lipschhiz and bounded functions.
3

---

Published as a conference paper at ICLR 2022
forward pasa
upd600
+ awi
: )ax
2. Vipi0
Ni MRLD
Hip oww)
#2 awN
Figure 1: NES perturbs the parameters of a black- -bxx model (p)) for N times via and performs parallel forward passes. Then, the model parameters are updated given the gradient estimation.
3 STRUCTURED VAE OPTIMIZATION USING NES
In this work we suggest using gradient-free based method for learning discrete structured VAEs. Specifically, we propose using NES to optimize discrete structured VAËs without the need prop- agate gradients through their latent discrete structures. For readability. we consider a concatenation of both / and & as μ [0: øl where; is the concatenation operation. For notational convenience, we refer 0 as fl and Φ as P2z
Combining the discrete VAE objective function in Eq. which is a non-continuous function due to the arg max operation and hence. non-Lipschitz, together with the NES objective in Eq. (i.e., setting k(p) [(H, 2)) we get the following smooth approximation by setting g(u) to be:
Eun- B) Bhnth (n)) - fws (/ww(2 (zz, arg m >))]. [
where [Mi; W2] is the concatenation of vectors W1. W2 that denotes the decoder and encoder parameters respectively. Our goal minimize Eq, [ Following the NES setup, we use Eq. with simplified notation to better emphasize the smoothing of Overall its gradient takes the form of:
(0,~)((011) Ey-- tooe
tong (s)))
faitow, (z, arg max ( [T}]]
(8)
=E2
We provide the pseudo code for using the NES algorithm to optimize discrete VAEs, together with
their gradient update rule on Algorithm T We estimate the NES gradient in Eq. 8 by sampling
N(0 T) and y G(ppatowy (z)), which in turn induces sampling of discrete structures arg max f2 7 This sampling procedure differs from that of the REINFORCE instances de- scribed in Eq. [ and Eq. as the samples of NES are tied to the sensitivity of the scoring function hyg (x), i.e, to a random perturbation of its parameters μ12 by ow2. In contrast, the samples of REINFORCE are proportional to the scoring function h (x) itself. In our experimental validation, we empirically demonstrate that NES has lower variance.
Theoretical Guarantees. Next, we prove NES converges for non-Lipschitz function k(-), which appears in reparameterized discrete VAEs in Eq. - 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 7 update rules to the parameters p., it generates the sequence ya(*), (,. (T)) and for a sufficiently large T, there exists & E (1, for which [|V[((())]] is arbitrarily small.
For mathematical simplicity, we prove convergence theorem the expected gradient, de- scribed in Eq. (t+1) +- 10() (t)).. To easily address [V/((()))] one usually consid- ers the difference /u (t-1) (0 using the remainder of the Taylor series: g(u((+1) - g(u(()) V)(u(() (
()) + 1 (1((+1) 10)) ( (()(()(((( (()),, where ja E [1((), (
Here we denote by V the Hessian of g(A). Applying the gradient update rule, we obtain the following equation:
n 2
(ac0))=-n\[VVooe> ((u(++)) ((((() -n/ V
A bound for (g(n the fact that the discrete VAE objective, which is given in Eq. 2 is a non-negative function that is continuous almost everywhere
is key to convergence theorem. Our bound relies

---

Published as a conference paper at ICLR 2022
Algorithm Natural Evolution Strategies for discrete VAΣs
Input: Initial parameters Ja. repeat
for all = to N Sample 1i ~ NO, (nn))
Evaluate u, BB 20 (x)) - furtou (2 end for
Update ya f N Σim Siu,
until stopping condition is met
Lemma 1. Let k : I2 + [0, M] be a non-negative function that is continuous almost everywhere and let g(u) BuwN [k(w)]. Then for any p1s f2 there holds:
dM3
Vg(m)) (V))())) Vglui) <
(10)
Proof can be found in Appendix A.1
The above lemma together with Eq. imply the following bound: (((++1))) ((x(()) < mll ()( + 172 dM
20 By summing over all algorithm steps
t= -1,..., 7 we show that the average norm of the gradient can be arbitrarily small for sufficiently large T.
Theorem 1. Under the conditions of Lemmal 7 Othere holds:
M
ndM
Vilh S
(11)
n/T
20
Moreover, when setting 17
20 TaMss
then: $ + V9(H S
2dM Tot
Therefore, there exists for which [
2d114
Taa
Proof is given in Appendix A.2 Intuitively, the above theorem proves that the NES algorithm verges 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 Lipschitz func- tions [41). Instead, we rely on the non-negativity of the discrete VAE objective. We note that a stationary point of g(n) (p. [k[)]] is not necessarily stationary point k(n). Never- theless, since lim (14) = k(μ) almost everywhere, except perhaps for con-continuous points, low value of g(n) is correlated with a low value of k(j).
EXPERIMENTS
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 we analyze how NES scales with the latent space dimension and neural network size. We additionally provide analysis for non-Lipschitz functions in Appendix B together with analyzing the results concerning the theoretical guarantees as presented in Section
In experiments, 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, each iteration, sample o Gaussian noise vectors where is the number of the VAE parameter sets utilized for estimating the NES update direction. Additionally, make NES more robust and scale-invariant, outputs of the perturbed forward passes into standard scores by subtracting their mean and dividing them by the standard deviation. All reported values measured on 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 Tesla-K80 NVIDIA GPUs.
transform the
4.1 LATENT STRUCTURE RECOVERY
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., bas- ketball players the court flock of birds during migration, can often be explained using
5

---

Published as a conference paper at ICLR 2022
Table 1:
The mean and standard deviation of the ELBO and structure recovery metries. NES outperforms
the five different baselines terms of structure recovery and presents a negligible standard deviation, which reflects its robustness.
Spanning Tree
Individual Edges
Method
ELBO 1
Edge FI-Score
ELBO
Edge FI-Score f
REINFORCE (Batch)
-2260 +
41
-2180 +
39 +
REINFORCE (EMA)
-2250 = 20
40 =7
-2170 10
42 + 1
REINFORCE (Malti-sanple))
-2230 = 20
42
-2150 = 10
40 - 0
NVIL
-1570 +: 300
83 20
-2110 10
42 ::
SST
-1080 = 110
91
-2100 = 20
41 +
NES (Ours)
-1117 = 45
92 1 0.2
-2150 + 10 44 : 0.6
simple structure. However, frequently, we only have access to individual trajectories without knowl- edge 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 distribution over the space of interaction structures given the component trajectories, and the reconstruction based on graph neural networks.
In our experiments, we utilize the dataset developed by Paulus et al. 48 where the target structure 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 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 10 possible spanning trees. In the second more challenging case, remove the tree constraint and consider all possible n - unique edge combinations as the latent space, where n denotes the number of vertices. The ability recover the underlying structure is measured as the edge Fl-score against the target spanning tree.
In each of the two cases, we compare NES with the corresponding Stochastic Softmax Trick (SST) [8). SSTs are the generalization of the Gumbel-Softmax Trick (GSM) for combinatorial discrete distributions. That is, in contrast to GSM, SSTs are designed 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-foreing every steps instead of 9 steps. We fix NES parameters to be a 0.01 and N = 600. We additionally compared our method against four in- stances of REINFORCE. Each utilizes different variance reduction technique. The first is NVIL, [39, which uses two control variates. The remaining three reduce variance by subtracting the fol- lowing 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 pecially well suited for structured prediction [25 26). uses the mean ELBO over r multiple samples per data point. Paulus al.
48 tuned on the set of (2, 4,8). Results listed in Table 1
In terms of structure recovery, the gradient-free NES outperforms all REINFORCE instances. More- over, when considering SST, NES achieves superior edge F1-score with slightly worse ELBO values. This is surprising since SST is gradient-based method that generalizes the effective GSM estima- tor 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 mini-batch size of (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 line with the results obtained by Kool et al. (25), where larger values of led to inferior results. Thus, it can be concluded that even with an equal computational cost, the REINFORCE instances are inferior to NES.
4.2 DEPENDENCY PARSING
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.]] where we focused on structures over undirected edges, here we focus dependency trees which are rooted directed spanning trees. Our model is based on a VAE architecture similar to that of differentiable perturb- and-parse (DPP) 9 The encoder is comprised of graph-based parser 222 that decomposes the

---

Published as a conference paper at ICLR 2022
Unlabeled attachment scores for unsupervised domain adaptation. The column name represents the setup where the treebank is set to be the target domain. For each the two tasks, NES achieves the hest UAS performance on 5 out of the 6 target domains.
Table 2:
Method
GL.CTG GL.TREEGAL
ID.CSUI
ID_GSD
RULGSD
RU.TAIGA
Projective
DPP
68.72
71.39
68.23
71.71
71.46
69.94
NES (Ours)
68.92
71.64
68.28
71.41
71.82
70.52
SparseMAP
68,57
70.61
68.23
70.56
70.99
69.96
NES (Ours)
68.67
70.98
68.28
70.47
71.01
69,96
score of a tree to the sum of the scores of its ares and produces a distribution over the space of depen- dency trees. Sampling from the latent space is performed using the "perturb-and-mee where each are score is perturbed independently with a noise derived from a Gumbel distribution. Then, the perturbed are scores (the resulting sample). For projective dependency parsing, we utilize the Eisner algorithm [II] as the MAP solver. Similarly, for non-projective dependency parsing, we use the Sou-Liu-Eddmode (CLE) algorithm 66] [10. The decoder is modeled as a language model, that given a latent depen- dency tree, attempts to reconstruct the input sentence.
technique,
fed into MAP solver, which outputs the highest-scoring tree
We compare NES with two strong baselines. For projective dependency parsing, we consider the DPP model. DPP optimizes the VAE by utilizing differentiable surrogate of the Eisner algorithm. In that manner, it tackles both the differentiability and exponential enumeration issues. In fact, DPP can be seen Stochastic Softmax Trick (SST). For peop-rooeecive dependency parsing, we consider SparseMAP 43 as 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 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_CSUL news, ID JSGD: general) and Russian (RU_GSD: general, RU_ RUTTA:::: social media, poetry, and fiction). We conduct 6 domain adaptation experiments, where we alternate between the source and target domain 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 (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 bbove-mentioned baselines. We set the aeper-parametert to those of the original implementation of Kiperwasser & Goldberg (22) and feed the models with the multilingual FastText word embeddings
[6]] We perform grid-search for each of the methods separately over learning rates in [5 10
10 10-3 and set the mini-bate batch size to 128. We fix NES parameters to be r = 0.1 and N = 400. Adam optimizer (21) is used optimize all methods. The models selected were those who obtained the best unlabeled attachment score (UAS) on the source domain validation set.
Table summarizes the results 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 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 bottleneck dynamic programming algorithm, NES can optimize the model parallel without the necessity gradient computations.
4.3 SCALABILITY ANALYSIS
Latent space size. In the following set of experiments, we further investigate the properties of NES and the several methods it was compared in Section 4.1 and 4.2 Specifically, examine how

---

Published as a conference paper at ICLR 2022
S6T
110
333
1.5
380
1.0
0.
20
100 200
20
220 200
250
50
100
124
330
350
Figure 2: Wall-clock time function of model input size. Experiments conducted the NRI model (Left), projective parsing model (Center), non-projective parsing model (Right), We observe that NES scales well with the latent space size in contrast to most of its competitors.
the latent space size affects the method's run-time by measuring the methods wall-clock time of forward and backward pass function of the input dimension (denoted by n). Note that the latent space size grows exponentially with the model input dimension, c.g., for 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 {10, ....... 250} create random dataset. Specificallyy, for the NRI model (Sec- tion 4..1. the input is n trajectories of 10 timesteps derived from standard Gaussian distribution. We compare the run- run-time of NES with that of SST and a REINFORCE instance that relies on Eq. where sampling is performed using Markov chain Monte Carlo (MCMC) algorithm. For the pars- ing model (Section (.2). we derive 32 random sentences of length by randomly sampling words from vocabulary of size 10000. As DPP utilizes differentiable surrogate of the Eisner algorithm, is compared to NES with the Eisner algorithm as the MAP solver. Similarly, SparseMAP is com- pared NES with the CLE algorithm. Since internal cluster consists of 4 GPUs, we utilize NES with = 4 for a fair comparison with the gradient-bars baseed methods. Finally, we run the experiments various random seeds and average the wall-clock time. Figure [ depicts the results.
As can seen, the run-times of DPP and SparseMAP heavily rely on the input dimension and grow 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, can be seen that NES scales well with the latent space size in contrast to most of its competitors.
Neural network size. Next, conduct study that examines how the enlargement of neural network affects the number of NES samples needed to optimize it. We begin by optimizing VAE of 25K parameters, then we enlarge its parameter size by 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 fixed temperature of as a baseline. For each model size, we examine how many NES samples are needed to achieve test ELBO as lower the one achieved by the SST. Results are depicted in Table A detailed description of the experimental setup can be found on Section in the Appendix.
We observe that enlarging the neural network by factor of 2 does not necessarily mean that should enlarge N in the same manner. To precise, 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.
RELATED WORK
Jang et al. 19, Maddison al. proposed the GSM estimator that replaces the nonddifferentiable arg max operation with 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 [II]. Mena al. (36] extended the Gumbel-Softmax estimator [19] and proposed the Humbel-Sinkhond method for learning latent permutations. Paulus et al. (48) took these ideas one
step further and proposed unified framework for designing structured relaxations of combinatorial
8

---

alternative optimization method for learning discrete Published as a conference paper at ICLR 2022
Table 3: ELBO function of the neural network size and N. The values in the "Growth in N column express the percent of additional NES samples (with respect to the preceding step) needed to achieve the baseline performance.
# Parameters
SST
NES
NN
Growth in N
25K
-240.48
-239.40
50K
-239.49
-231.02
+ 0.00%
100K
-233.13
-232.40
90
50.00%
200K
-233.92
-233.99
100
11.11%
400K
-239.02
-234.29
100
+ 0.00%
800K
-241.03
-234.86
100
+ 0.00%
distributions. Unlike these methods, approach is generic and such, it can be applied to general structures with no additional effort, since obviates the need for a differentiable surrogate of the linear maximization oracle. Mensch & Blondel (37] proposed a framework for turning dynamic programming algorithms differentiable.
g is framework for training deep networks with sparse structured hidden layers, solved by sequential calls to a MAP
Others have taken
generic approach. For example, SparseMAP 4
oracle. In similar sense, Itkina al. [8] suggest using evidential theory perform post hoc latent space sparsification and thus reducing the discrete latent sample space at test time. Chen et al. (S) generalized this method to 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 arg max operation, Lorberbom et al. 3 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 introduces an alternative ap- proach to learning discrete latent representation. However, VQ-VAB differs from our discrete struc- tured 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 given graph. Hence do not perform unsupervised vec- tor quantization as VQ-VAE but rather use a predetermined quantization over the set of possible structures. In this work, we rather focus latent structures.
Recently, black-box optimization methods have been applied neural networks 32 52
30
53 38)
Salimans et al. (52) showed that NES is 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. Lene et al. [3]] proposed a hybrid method that alternates between NES and SGD for training large sparse models. Finally, Zhang 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 structured VAEs.
6 CONCLUSION
We suggested using NES, class of gradient-free black-box algorithms, as an alternative for learn- ing 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 implement. Additionally, we showed that NES scales well with the latent space dimension and neural network proved that NES converges for
size. To establish the theoretical soundness of approach, non-Lipschitz functions such as the objective function of a discrete VAE.
In this study, 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.

---

Published as a conference paper at ICLR 2022
REFERENCES
[1] Quentin Berthet, Mathieu Blondel, Olivier Teboul, Marco Cuturi, Jean- Phinlppilippe Vert, and Fran- cis R Bach. Learning with differentiable pertubed optimizers. In Advances Neural Informa- fion Processing Systems, 2020.
[2] Mathieu Blondel, Olivier Teboul, Quentin Berthet, and Josip Djolonga. Fast differentiable sorting and ranking. In International Conference on Machine Learning, pp. 950-959, PMLR, 2020.
[3] Gilles Boulianne. A study of inductive biases for unsupervised speech representation learning. IEEE/ACM Transactions Audio, Speech, and Language Processing, 28:2781-2795 2020.
[4] Dimo Brockhoff, Anne Auger, Nikolaus Hansen, Dirk V Arnold, and Tim Hohm. Mirrored sampling and sequential selection for evolution strategies. In International Conference on Parallel Problem Solving from Nature, pp. 11-21. Springer, 2010.
[5] Phil Chen, Masha Itkina, Ransalu Senanayake, and Mykel J Kochenderfer. Evidential softmax for sparse multimodal distributions in deep generative models, Advances Neural Information Processing Systems, 34, 2021.
[6] Yoeng-Jin Chu. On the shortest arborescence of directed graph. Scientia Sinica, 14:1396 1400, 1965.
[7] Tarin Clanuwat, Mikel Bober-Irizar, Asanobu Kitamoto, Alex Lamb, Kazuaki Yamamoto, and David Ha. Deep learning for classical japanese literature, 2018.
[8] Gonçalo Correia, Vlad Niculae, Wilker Aziz, and André Martins. Efficient marginalizatton of discrete and structured latent variables via sparsity. Advances in Neurul Information Processing Systems, 33, 2020.
[9] Caio Corro and Ivan Titov. Differentiable perturb-and- and-parse: Semi-supervised parsing with structured variational autoencoder. In International Conference on Learning Representations, 2018.
[10] Jack Edmonds. Optimum branchings. Journal of Research of the National Bureau of Stan- dards, 11:233-240, 1967.
[11] Jason M Eisner. Three new probabilistic models for dependency parsing: an exploration. In Proceedings of the 16th conference on Computational linguistics- Volume 1, pp. 340-345, 1996.
[12] SM Ali Eslami, Nicolas Heess, Theophane Weber, Yuval Tassa, David Szepesvari, Geoffrey E Hinton, et Attend, infer, repeat: Fast scene understanding with generative models. In Advances in Neural Information Processing Systems, pp. 3225-3223333 3233, 2016.
[13] Xianghong Fang, Haoli Bai, Jian Li, Zenglin Xu, Michael Lyu, and Irwin King. Discrete ruto-regressive variational attention models for text modeling. In 2021 International Joint Conference on Neural Networks (IJCNN), pp. 1-8. IEEE, 2021.
[14] John Geweke. Antithetic acceleration of monte carlo integration in bayesian inference. Journal of Ecomometrics, 38(1-2):73-89) 1988.
[15] Alex Glushkovsky, Ai discovering a coordinate system of chemical elements: dual represen- tation by variational autoencoders. arXiv preprint 2011::0011.12000 2020.
[16] Édouard Grave, Piotr Bojanowski, Prakhar Gupta, Armand Joulin, and Tomás Mikolov, Learn- ing word vectors for 157 languages. In Proceedings of the Eleventh International Conference on Language Resources and Evaluation (LREC 2018), 2018
[17] Zhiting Hu, Zichao Yang, Xiaodan Liang, Ruslan Salakhutdinov, and Eric P Xing. Toward controlled generation of text. In International Conference on Machine Learning, pp. 1587 1596, 2017.
10

---

Published as a conference paper at ICLR 2022
[18] Masha Itkina, Boris Ivanovic, Ransalu Senanayake, Mykel J Kochenderfer, and Marco Pavone. Evidential sparsification of multimodal latent spaces in conditional variational autoencoders. Advances in Neural Information Processing Systems, 33, 2020.
[19] Eric Jang. Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax. International Conference on Learning Representations, 2017.
[20] Michael Jordan, Zoubin Ghahramani, Tommi Jaakkola, and Lawrence K Saul. An intro- duction to variational methods for graphical models. Machine learning, 37(2):183-233, 1999.
[21] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Interna- lional Conference on Learning Representations, 2015.
[22] Eliyahu Kiperwasser and Yoav Goldberg. Simple and accurate dependency parsing using bidi- rectional Istm feature representations. Transactions of the Association for Computational Lin- guistics, 4:313-327, 2016.
[23] Thomas Kipf, Ethan Fetaya, Kuan-Chieh Wang. Max Welling, and Richard Zemel. Neural relational inference for interacting systems. In International Conference on Machine Learning, pp. 2688-2697, PMLR, 2018.
[24] Terry Koo, Amir Globerson, Xavier Carreras, and Michael Collins. Structured prediction models via the mree theorem. In Proceedings of the 2007 Joint Conference on Empiri- cal Methods in Natural Language Processing and Computational Natural Language Learning (EMNLP-CONLL)) pp. 141-150, 2007.
[25] Wouter Kool, Herke van Hoof, and Max Welling. Buy 4 reinforce samples, get baseline for free! 2019.
[26] Wouter Kool, Herke van Hoof, and Max Welling. Estimating gradients for discrete random variables by sampling without replacement. In International Conference on Learning Repre- sentations, 2019.
[27] Joseph Kruskal. On the shortest spanning subtree of a graph and the traveling salesman problem. Proceedings of the American Mathematical society, (1)):48-50, 1956.
[28] Matt J Kusner, Brooks Paige, and José Miguel Hernández-Lobato. Grammar variational au- toencoder. In International Conference on Machine Learning, pp. 1955-1954. PMLR, 2017.
[29] Brenden M Lake, Ruslan Salakhutdinov, and Joshua B Tenenbaum. Human-level concept learning through probabilistic program induction. Science, 35066266) 2015.
[30] Joel Lehman, Jay Chen, Jeff Clune, and Kenneth O Stanley. Es is more than just a traditional finite-difference approximator. In Proceedings of the Genetic and Evolutionary Computation Conference, pp. 450-457, 2018.
[31] Karel Lenc, Brich Elsen, Tom Schaul, and Karen Simonyan. Non-differentiable supervised learning with evolution strategies and hybrid methods. arXiv preprint arXiv: 19X.0v9990003333399 2019.
[32] Jialin Liu, Antoine Moreau, Mike Preuss, Jeremy Rapin, Baptiste Roziere, Fabien Teytaud, and Olivier Teytaud. Versatile black-box optimization. In Proceedings of the 2020 Genetic and Evolutionary Computation Conference, pp. 620-628 2020.
[33] Guy Lorberbom, Andreea Gane, Tommi Jaakkola, and Tamir Hazan. Direct optimization through argmax for discrete variational auto-encoder. In Advances in Neural Information Pro- cessing Systems, pp. 6203-6214, 2019.
[34] Chris J. Maddison, Andriy Mnih, and Yee Whye Teh. The Concrete Distribution: A Con- tinuous Relaxation of Discrete Random Variables. In International Conference on Learning Representations, 2017.
11

---

Published as a conference paper at ICLR 2022
[35] Ryan McDonald, Joakim Nivre, Yvonne Shirmbach-Brunddzz Yoav Goldberg, Dipanjan Das, Kuzman Ganchev, Keith Hall, Slav Petrov, Hao Zhang, Oscar Täckström, Claudia Bedini, Núria Bertomeu Castelló, and Jungmee Lee. Universal dependency annotation for multilingual parsing- Proceedings of the 51st Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers), volume 2, pp. 92-97, 2013.
[36] Gonzalo Mena, Jasper Snoek, Scott Linderman, and David Belanger. Learning latent permu- tations with gumbel-sinkhorn networks. In ICLR 2018 Conference Track, volume 2018, 2018.
[37] Arthur Mensch and Mathieu Blondel. Differentiable dynamic programming for structured prediction and attention. In International Conference on Machine Learning, pp. 3462-3471. PMLR, 2018
[38] Laurent Meunier, Herilalaina Rakotoarison, Pak Kan Wong, Baptiste Roziere, Jeremy Rapin, Olivier Teytaud, Antoine Moreau, and Carola Doerr Black-box optimization revisited: Im- proving algorithm selection wizards through massive benchmarking. IEEE Transactions on Evolutionary Computation, 2021.
[39] Andriy Maih and Karol Gregor. Neural variational inference and learning belief networks.
In International Conference on Machine Learning, pp. 1791-1799. PMLR, 2014.
[40] Igor Mordatch and Pieter Abbeel. Emergence of grounded compositional language in multi- agent populations. In Thirty-Second AAAI Conference on Artificial Intelligence, 2018.
[41] Yurii Nesterov and Vladimir Spokoiny. Random gradient-free minimization of convex func- tions. Foundations of Computational Mathematics, 17(2):527-566, 2017.
[42] Vlad Niculae and Andre Martins. Lp-sparsemap: Differentiable relaxed optimization for sparse structured prediction. In International Conference on Machine Learning, pp. 7348- 7359. PMLR, 2020.
[43] Vlad Niculae, Andre Martins, Mathieu Blondel, and Claire Cardie. Sparsemap: Differentiable sparse structured inference. International Conference on Machine Learning, pp. 3799-3808. PMLR, 2018
[44] Joakim Nivre, Marie-Catherine De Marneffe, Filip Ginter, Yoav Goldberg, Jan Hajic, Christo- pher D. Manning, Ryan McDonald, Slav Petrov, Sampo Pyysalo, Natalia Silveira, Reut Tsar- faty, and Daniel Zeman. Universal dependencies vl: multilingual treebank collection. In LREC, 2016.
[45] Joakim Nivre, Mitchell Abrams, Željko Agié, Lars Ahrenberg, Lene Antonsen, Maria Jesus Aranzabe, Gashaw Arutie, Masayuki Asahara, Luma Ateyah, Mohammed Attia, et al. Univer- sal dependencies 2.2. 2018.
[46] John Paisley, David Blei, and Michael Jordan. Variational bayesian inference with stochastic search. Proceedings of the 29th International Conference Machine Learning, 2012.
[47] George Papandreou and Alan Yuille. Perturb-and-map random fields: Using discrete op- timization to learn and sample from energy models. In 2011 International Conference on Computer Vision, pp. 193-200. IEEE, 2011.
[48] Max B Paulus, Dami Choi, Daniel Tarlow, Andreas Krause, and Chris J Maddison. Gradient estimation with stochastic softmax tricks. Advances in Neural Information Processing Systems, 2020.
[49] Ali Razavi, Aaron den Oord, and Oriol Vinyals. Generating diverse high-fidelity images with vq-vae-2. In Advances neural information processing systems, pp. 14866-14876, 2019.
[50] Guy Rotman and Roi Reichart. Deep contextualized self ceff-training for low resource dependency parsing. Transactions of the Association for Computational Linguistics, 7:695-713, 2019.
[51] Ruslan Salakhutdinov and lain Murray. On the quantitative analysis of deep belief networks. In Proceedings of the 25th international conference Machine learning, pp. 872-879, 2008
12

---

Published as a conference paper at ICLR 2022
[52] Tim Salimans, Jonathan Ho, Xi Chen, Szymon Sidor, and llya Sutskever. Evolution strategies as a scalable alternative to reinforcement learning. arXiv preprint arXiv: 2017.
[53] Elad Sarafian, Mor Sinay, Yoram Louzoun, Noa Agmon, and Sarit Kraus. Explicit gradient learning for black-box optimization. In International Conference Machine Learning, pp- 8480-8490. PMLR, 2020.
[54] Dinghan Shen, Qinliang Su, Paidamoyo Chapfuwa, Wenlin Wang. Guoyin Wang, Ricardo Henao, and Lawrence Carin. Nash: Toward end-to-end neural architecture for generative mantic hashing. In Proceedings of the 56th Annual Meeting of the Association for Computa- lional Linguistics (Volume 1: Long Pupers), pp. 2041-2050, 2018.
[55] Arash Vahdat, William Macready, Zhengbing Bian, Amir Khoshaman, and Evgeny Andriyash. Dvac++: Discrete variational autoencoders with overlapping transformations. In International Conference on Machine Learning, pp. 5035-5044. PMLR, 2018.
[56] Aaron van den Oord, Oriol Vinyals, and Koray Kavukcuoglu. Neural discrete representation learning. In Proceedings of the 3 st International Conference on Neural Information Process- ing Systems, pp. 6309 6318, 2017.
[57] Daan Wierstra, Tom Schaul, Jan Peters, and Juergen Schmidhuber. Natural evolution strate- gies. In 2008 IEEE Congress Evolutionary Computation (IEEE World Congress on Com- putational Intelligence), pp. 3381-3387. IEFE, 2008.
[58] Daan Wierstra, Tom Schaul, Tobias Glasmachers, Yi Sun, Jan Peters, and Jürgen Schmidhu- ber. Natural evolution strategies. The Journal of Machine Learning Research, 15(1))949-980, 2014.
[59] Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: a novel image dataset for bench- marking machine learning algorithms, 2017.
[60] Dani Yogatama, Phil Blunsom, Chris Dyer, Edward Grefenstette, and Wang Ling. Learning to compose words into sentences with reinforcement learning. In 5th International Conference on Learning Representations (ICLR 2017). International Conference on Learning Representa- tions, 2017.
[61] Xingwen Zhang, Jeff Clune, and Kenneth O Stanley. On the relationship between the openai evolution strategy and stochastic gradient descent. arXiv preprint arxiv: 2017.
A PROOFS OF NES CONVERGENCE FOR NON-LIPSCHITZ FUNCTIONS
A.1 PROOF OF LEMMA I
Proof. The proof follows two main steps: (i) the spectral norm of V2g at most M/o² and (ii)
[[V)(=)) dM² /o²
Since the Hessian matrix is symmetric, we can use Rayleigh quotient and obtain:
Vg(un) (V'((za)))) vo(uz) S Dmaazl | [g(n))1
(12)
where Amaz is the largest eigenvalue of V²9(22). Since V V2g(u2) is symmetric it can also be shown that:
Ammaz = g(22) max Is V g(2)s))
(13)
||||==
Applying the log derivative trick on Eq. 6 obtain
Vgg(u2) Juu (2T)
1nll
2 =uw'' www k(u2 + ow)dw.
(14)
13

---

Published as a conference paper at ICLR 2022
Therefore:
1wil2
V
(12)s
max
02 s' "ww" k(p2 + ow)dwl
| =1
| =1
Rd (2T)
M
Inj
max
(s (
(15)
x²
|o] =1
JRd (2T)
M
M
||
max
Since k--) is bounded by MM and w are i.d. normal Gaussian random variables, therefore for any the random s'ww w is Gaussian with zero mean and variance ||s||?. By combining Eq. 13 with Ineq. 15 we get:
M 0
V2 -
(16)
which concludes step (i). Next, bound the squared norm of the gradient:
1
W k(m1
dwwl?
(2T)
1
1
Iull
wk([1 + ow))wll?
(2Tr)
(17)
1
1
Luje
wk(p1 + =dw
0²
(2T)
M
1w12
d.M
a²
RRd (2x)
o2
where the first inequality is obtained using the Keuchy-Schwarrd inequality and the second inequality by bounding k(.) with M. Thus, overall we have showed that:
$ M σ2
Vg(un)
( V'g((z)VV((H))
(18)
dM3
σ4
A.2
PROOF OF THBOREM 1
Proof. Rearranging Eq. 9 we obtain:
n| Vo(o")) o(w(0) g(u(t+1)) + 17 dM3
(23)
204
Summing over all algorithm steps / = 1,..., T, we have:
T
Vo((e)1 El( (((((11)) +1² 1/2 dM&
T
(24)
20
Opening the telescopic sum:
Vo s (( (((((+1)) + n2² dT
(25)
2
Since k6-) is non-negative and bounded, the difference between (( and (T+()+)) is bounded from above by M:
14

---

Published as a conference paper at ICLR 2022
Vg(u')] z M + 112 2
(26)
20²
We multiply both sides of the inequality by o/.
Vo(nt) < M 17² T t=1 T/T
1
T 20r
M
dM
+
(27)
204
Next, minimize the right-hand size of the inequality in n
20
"7" V TdM2
(28)
and plug it back to Ineq. 27
T
2dM
[g(h'))
(29)
To
Then, for arbitrarily small 8 > 0 such that:
T
/2dM4
(g(nn)11
(30)
T
Tod
there exists for which:
[ 55
(31)
after most
2dM
T
(32)
82
steps.
B ADDITIONAL RESULTS
In the following experiments, we define the encoder as input = MT.P(a) = ReLU = MLP((10)
+ argmaax and the decoder as (L) + ReLU = MLP(B) + output, where B is the input dimension. Unless otherwise stated, = 300.
B. NEURAL NETWORK SIZE
The experiments were conducted the FashionMNIST dataset [59] with fixed binarization [51] We tune a for enlarging the VAE, Particularly, a picked from the set [16, 32, 64, 128, 256 512]. All models were trained using the ADAM optimizer (1 over 300 epochs with constant learning rate 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 [ NES converges to stationary point of g(4)) for non-Lipschitz functions. To empirically explore the relation between g(n) and k(p), conduct a set of experiments which we demonstrate that a low value g(4) is correlated with low value of the objective function k(J4) by estimating the average absolute distance between them.
15

---

Published as a conference paper at ICLR 2022
25
0.01
aerrge dataro
4.5
UNSOUNDED
0.3
150
150
200
spodas
opedss
Figure 3: Left: The average absolute distance between the negative ELBO k(μ) and its Gaussian approximation g(n) function of the training epoch. The smaller is, the further the proximity between the two functions. Right: Negative ELBO as a function of the training epoch. Performance improves the ELBO upper bound M increases. Result suggests that bounding discrete VAE loss with a large enough M guarantees the convergence NES within a finite number of iterations on the hand, and the other hand, does not impair performance.
331
322
503
NES
540
453
403
452
483
480
203
280
303
553
382
opacks
Figure 4: Negative ELBO as a function of the training epoch. Experiments are conducted on the binarized FashionMNIST (Left), KMNIST (Center), and Omniglot (Right) datasets. NES does not rely on computing gradients and yet achieves comparable performance with the unbiased and GSM methods. On the KMNIST benchmark NES even performs substantially better than GSM.
First, we estimate the Gaussian approximation g(n) 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 objec- tive function, and the estimated Gaussian approximation and average experiment with three different NES configurations:
the tested samples. We 300 and a E (0.01, 0.5, 0.1] on the Fash- ionMNIST dataset [59]. The results presented on the left image in Figure 3 Iidiccate that the smaller is, the further the proximity between k((p) and g(n). It can also be seen that the average distance converges and stabilizes 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 Theorem assumption. However, can be upper bounded by bounding each log probability component with 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 model trained with an unbounded ELBO, denoted by UNBOUNDED. We train the three models with N depicted on the right image in Figure 3 It can be seen that bounding the loss has minimal effect on model performance when M is big enough. Increasing M improves the performance, while using relatively small M value may cause the model diverge. Surprisingly. bounding the loss with M = 9 leads to a slightly lower loss compared to the UNBOUNDED baseline. We hypothesize this due regularization effect.
300 and a 0.1 on the FashionMNIST dataset. Results are
16

---

Published as a conference paper at ICLR 2022
B.4 UNSTRUCTURED TASKS.
Unlike structured VAES, where the latent spaces are often exponentially large, here we explore latent space that consists of only 10 different assignments. Therefore, unbiased gradient of the objective with respect to the VAE parameters can be analytically computed by enumerating over all possible latent assignments (Eq. D]. We denote this method as UNBIASED and consider the loss of a model trained with this method as 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- Gumbel-Softmax -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. [9]
Experiments are conducted on the FashionMNIST 59 KMNIST 7 and Omniglot 22 datasets
with fixed binarization [51]. models are trained using the ADAM optimizer 21 with a constant learning rate of 10 and mini bath-athhhhee size of 128. Figure 4 depicts the negative ELBO of NES and its competitors.
Surprisingly, NES achieves competitive results compared to the UNBIASED method 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 opti- mizes the VAE parameters by only evaluating the model at certain points in parameter space.
C DERIVING EQUATION
Let y be random function that associates an independent random variable 7(2) for each put Z. When the random variables follow the Gumbel distribution law with mean 2), which we denote by ((h 2)) and whose probability density function is gs (x()) = h (tt)tte-h(ss.0)) for the Euler constant 0.57. Then for g(t) - (t)
we obtain the following identity:
bo(z,=) = Progl== = 2). where 2" A arg max (y))))
(33)
BEZ
Proof. Let Gz) = (g((e))(
be the Gumbel cumulative distribution function. Then
P Pr~gl="= = == Prywl== ==
arg max {+()) E=I.....
= (1(()) (=))))).
(34)
8
Since (1) (1to--a((t))) it holds that
/ (t()) II Ga(t)dt / (-th-ho (::))tt=--(t-hoo-a-a))
II Gz(t)dt
(35)
hole)
7
where = e (8+0) I Gt(t)dt is independent of z. Since IP Pywgl = 2²1 is a distribution then Z must equal to > ch((2,))
Next, we use the Gumbel-Max trick to rewrite the expected log-likelihood in the ELBO in the following form:
Egwas Pog
P
P = =)fo(1, 2) Ey>g(fo(x,, =*)]
(36)
ZEZ
17

---

Published as a conference paper at ICLR 2022
The equality results from the identity P Pyngl (2 2] B the linearity of expectation E3ez E =f 2) = Enng 12 ===== fo 2*]] and the fact that Σxez 1===2 1.
When 2 (21,..., )( spanning tree, or more generally, belongs the structured space, one cannot assign an i.i.d. random variable to each 2. Instead we relate random variable 7 (? MEj and set ((=)
D THE FULL OBJECTIVE FUNCTION
In Eq. 7 we didn't include KL-divergence term for the sake of simplicity. In practice, the NES algorithm optimizes both terms. Thus, for completeness provide the full NES objective. For the avoidance doubt, in our experiments, we optimized both terms.
Assuming that (-(-) is the uniform distribution over the space of structures, the Gumbel-Man repa- rameterization trick let us derive the following approximation:
(T(g((()))llom()) Eg (-10)) log (0(2/x)) po(2)
1
Eo hp(x) log
(37)
where arg (rgxaaez(=T {= The resulting NES objective is:
(p ( (w)) JNy (I,, 2²r =²)+ + 2* Mong (2) -log
(38)
[2]
And its gradient takes the form of:
w By (w)) -- ( (x,2²) (0,,2)+++ + T (x) log (
[ }}
(39)
2
where w = [wi; w2] the concatenation of the two vectors W1, W2.
18


Reviews:
[{"review_id": "v1eiNq_6CKq", "reviewer": "Reviewer_qpkv", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This is an elegant paper.  Although, NES algorithms have been used elsewhere, this is the first application to VAE training.  By demonstrating its feasibility and effectiveness the authors highlight the usefulness of the approach.  The discussion and analysis is informative and of a high scientific standard.  The use is the approach in this contexts requires the use of a number of sampling tricks which makes the approach quite subtle.\n\nOne could argue that this is just an application of NES to yet another application, but I think this would be unfair.  The REINFORCE algorithm has become a important, goto, algorithm in deep-learning.  Showing that NES provides a simple and efficient algorithm for solving non-differentiable problems is an important contribution worth making.\n\nThe paper is well written, accurate and to the point.", "summary_of_the_paper": "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.", "summary_of_the_review": "This is in my view a good paper that deserves a place in ICLR.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "n7Gb_Eo-Q4c", "reviewer": "Reviewer_ANuT", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The strengths of the paper:\n- proposing another approach to optimization a neural network than through propagating gradients,\n- proving that NES converges for non-Lipschitz functions such as the objective function of a discrete VAE,\n- experimentally showing that gradient-free methods have similarly effective as perturb-and-parse gradient-based methods\n\nThe weaknesses of the paper:\n- optimizing only the mean of the distribution of parameters, a covariance matrix of the distribution of parameters is imposed by the user. The classical model VAE optimizes both these parameters.\n- too poor experiments, the authors showed experiments that presented effective optimization of the ELBO values and how the latent space size affects the method’s run-time. It could be possible to show how the model reconstructs the images or generates them (this can be troublesome due to the lack of optimization of the VAE model variance). The results of such experiments can be compared to the results of the VQ-VAE method (https://arxiv.org/pdf/1711.00937.pdf or VQ-VAE 2 - https://arxiv.org/pdf/1906.00446.pdf).", "summary_of_the_paper": "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.", "summary_of_the_review": "The paper is well written, the authors clearly describe the problem and the proposed solution that is computationally appealing because does not require propagating gradients. It is theoretical work, the experiments only show the effectiveness of the ELBO function optimization. The authors consider only optimization of the mean distribution of parameters and do not show other experiments that confirm the effectiveness of such optimization in generative models (see the section of the weaknesses of the paper). Hense I rate the work good but not sufficient for the ICLR conference.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "VMxJTOe9eX", "reviewer": "Reviewer_Fqt8", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths.\n\n1. The paper is well-organized. \n\n2. In this work, a gradient-free black-box optimization algorithm is explored for discrete structured VAEs. \nThey also experimentally show gradient-free methods are as effective as sophisticated gradient-based methods.\n\n3. They prove that the NES algorithm converges for non-Lipschitz functions. This is different from the contemporary trend that relies on Lipschitz functions.\n\nWeaknesses\n1. In VAES, there are two terms for ELBO: one is sampled log-likelihood term, the other is the KL-divergence term that measures the similarity of the auxiliary distribution q and the unknown distribution p.  A lot of people notice the importance of the second term. In this paper, the authors use gradient-free methods for the first term. However, the second term is missing in the algorithm. Or do I miss something? \n\n2. The gradient-free black-box optimization algorithm is generally with high variance (for example, refinance force algorithm). Unfortunately, there are no clear discussions on this. And how is the effect of the proposed method? Mirrored sampling is used in the experimental setup. It is not clear whether this sampling method is used in other methods. This leads to a question: whether there is a fair comparison.\n", "summary_of_the_paper": "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.", "summary_of_the_review": "It would be great to provide more details on how to get Equation (2). Especially, why an argmax operation is added. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "Uo_pxXi5B5w", "reviewer": "Reviewer_12Bq", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**Strengths:**\n* The paper addresses an open problem in the space of discrete VAEs.\n* The paper is very well written, easy to follow, and compelling.\n* The presentation of the background material logically leads the reader into the proposed methodology, and the math is included as a meaningful and coherent addition to the story.\n* The experiments are extensive and convincing, demonstrating advantages over a wide range of baselines.\n\n**Weaknesses and suggestions for improvement:**\n* Some relevant works are missing from the literature review [1-6]. Contextualizing the proposed method in these works would make the overall paper stronger.\n* Since Eq. (2) is critical to the proposed method and the presented theoretical results, it would be helpful to highlight its importance when it is first defined. I recommend you clearly state that Eq. (2) is the objective function you will consider for the remainder of the paper, and elaborate on why it is non-negative (i.e., is it because of a ReLU at the end of the $f_{\\theta}$ network?). Additionally, please highlight in the proof of Lemma 1 where the non-negativity comes in.\n* I would appreciate the inclusion of a discussion of the bias introduced by the NES algorithm in Eq. (4). Further discussion on why $N$ forward passes in NES are more desirable than a sampling-based technique with $N$ samples would also be helpful. The experiments and the parallelization arguments contribute to this discussion, but a bit more emphasis and detail in the methods section would aid the reader to better understand the advantages of the proposed method.\n* I recommend the authors spend a bit more time explaining what $w_{1}$ and $w_{2}$ correspond to, i.e., the encoder and decoder parameters, respectively? At first read, I thought you were setting $N = 2$.\n* The statement on pages 4 and 5 that for any $T$, there exists a $t \\in \\{1, ..., T\\}$ for which the magnitude of the gradient is arbitrarily small needs to be corrected. Specifically, it should say something like for *T sufficiently large* the magnitude of the gradient is arbitrarily small.\n* In the experiments, additional discussion would be helpful for why SST might be expected to do better than NES in Table 1, and the disadvantages of SparseMap compared to NES.\n* Lastly, the paper needs to be proofread for typos, as I found a number of them while reviewing. Some examples include: 'over exponentially large latent space', missing period after Eq. (3), 'using NES algorithm' (missing 'the'), incorrect quotation marks for 'parameters' and \"Growth in N\", 'in contrast to contemporary trend' (missing 'the'), 'which must carefully designs', 'being simpler and robust' (missing 'more'), 'generic, flexible, and simpler to implement' (missing 'more'). The references need to be proofread for consistency, as well: some venues are missing, some venues are only presented as acronyms, while others are presented in full, incorrect capitalization ('Ai'), etc. If possible, avoid splitting equations across multiple lines in the text and capitalize 'Section' when referencing section numbers.\n\n[1] A. van den Oord, O. Vinyals, and K. Kavukcuoglu, \"Neural discrete representation learning,\" in NeurIPS, 2017.\n\n[2] A. Razavi, A. van den Oord, and O. Vinyals, \"Generating diverse high-fidelity images with VQ-VAE-2,\" in NeurIPS, 2019.\n\n[3] G. Correia, V. Niculae, W. Aziz, and A. Martins, \"Efficient marginalization of discrete and structured latent variables via sparsity,\" in NeurIPS, 2020.\n\n[4] M. Itkina, B. Ivanovic, R. Senanayake, M. J. Kochenderfer, and M. Pavone, “Evidential sparsification of multimodal latent spaces in conditional variational autoencoders,”in NeurIPS, 2020.\n\n[5] P. Chen, M. Itkina, R. Senanayake, and M. J. Kochenderfer, \"Evidential softmax forsparse multimodal distributions in deep generative models,\" in NeurIPS, 2021.", "summary_of_the_paper": "The proposed paper presents a method for optimizing discrete structured VAEs using Natural Evolution Strategies (NES). The authors present a theoretical conclusion regarding algorithm convergence and extensive empirical results demonstrating that NES performs comparably to gradient-based methods, while being more computationally efficient.", "summary_of_the_review": "Overall, the authors present an interesting method for optimizing discrete VAEs with compelling theoretical and experimental results. The paper is well written and the story is coherent. My recommendation is to accept.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The authors propose to use genetic algorithms to learn variational autoencoders (VAEs) with discrete latent spaces. Specifically they employ natural evolution strategies (NES) to avoid backpropagating gradients through discrete variables. Experiments show how the proposed approach is competitive with the current state-of-the-art to train discrete VAEs.

Some concerns arose from the review and discussion phases, these included confusion around the justification and derivation of NES for VAEs in the presentation and the limitation of the experiments. Authors were responsive and provided the reviewers the needed clarifications, an updated presentation in the revised paper and additional experimental results which ultimately were successful in raising the reviewers' scores towards full acceptance.
Decision: Accept (Poster)