F OURIER S TOCHASTIC B ACKPROPAGATION
Anonymous authors
Paper under double-blind review
A BSTRACT
Backpropagating gradients through random variables is at the heart of numerous machine learning applications.
In this paper, we present a general framework for deriving stochastic backpropagation rules for any distribution, discrete or continuous.
Our approach exploits the link between the characteristic function and the
Fourier transform, to transport the derivatives from the parameters of the distribution to the random variable.
Our method generalizes previously known estimators, and results in new estimators for the gamma, beta, Dirichlet and Laplace distributions.
Furthermore, we show that the classical deterministic backproapagation rule and the discrete random variable case, can also be interpreted through stochastic backpropagation.
I NTRODUCTION
Deep neural networks with stochastic hidden layers have become crucial in multiple domains, such as generative modeling (Kingma & Welling, 2013; Rezende et al., 2014; Mnih & Gregor, 2014), deep reinforcement learning (Sutton et al., 2000), and attention mechanisms (Mnih et al., 2014).
The difﬁculty encountered in training such models arises in the computation of gradients for functions of the form L ( θ ) :
=
E z ∼ p θ [ f ( z )] with respect to the parameters θ , thus needing to backpropagate the gradient through the random variable z .
One of the ﬁrst and most used methods is the score function or reinforce method (Glynn, 1989; Williams, 1992), that requires the computation and estimation of the derivative of the log probability function.
For high dimensional applications however, it has been noted that reinforce gradients have high variance, making the training process unstable (Rezende et al., 2014).
Recently, signiﬁcant progress has been made in tackling the variance problem.
The ﬁrst class of approaches dealing with continuous random variables are reparameterization tricks.
In that case a standardization function is introduced, that separates the stochasticity from the dependency on the parameters θ .
Thus being able to transport the derivative inside the expectation and sample from a ﬁxed distribution, resulting in low variance gradient (Kingma & Welling, 2013; Rezende et al., 2014;
Titsias & L ´ azaro-Gredilla, 2014;
Ruiz et al., 2016; Naesseth et al., 2017; Figurnov et al., 2018).
The second class of approaches concerns discrete random variables, for which a direct reparameterization is not known.
The ﬁrst solution uses the score function gradient with control variate methods to reduce its variance (Mnih & Gregor, 2014; Gu et al., 2016).
The second consists in introducing a continuous relaxation admitting a reparameterization trick of the discrete random variable, thus being able to backpropagate low-variance reparameterized gradients by sampling from the concrete distribution (Jang et al., 2016;
Maddison et al., 2016; Tucker et al., 2017; Grathwohl et al., 2018).
Although recent developments have advanced the state-of-the-art in terms of variance reduction and performance, stochastic backpropagation (i.e computing gradients through random variables) still lacks theoretical foundation.
In particular, the following questions remain open: How to develop stochastic backpropagation rules, where the derivative is transferred explicitly to the function f for a broader range of distributions?
And can the discrete and deterministic cases be interpreted in the sense of stochastic backpropagation?
In this paper, we provide a new method to address these questions, and our main contributions are the following:
• We present a theoretical framework based on the link between the multivariate Fourier transform and the characteristic function, that provides a standard method for deriving stochastic backpropagation rules, for any distribution discrete or continuous.
• We show that deterministic backpropagation can be interpreted as a special case of stochastic backpropagation, where the probability distribution p θ is a Dirac delta distribution, and that the discrete case can also be interpreted as backpropagating a discrete derivative.
• We generalize previously known estimators, and provide new stochastic backpropagation rules for the special cases of the Laplace, gamma, beta, and Dirichlet distributions.
• We demonstrate experimentally that the resulting new estimators are competitive with state-of-the art methods on simple tasks.
B ACKGROUND & P RELIMINARIES
Let ( E, λ ) be a d -dimensional measure space equipped with the standard inner product, and f be a square summable positive real valued function on E , that is, f : E → R + , with (cid:82) E | f ( z ) | 2 λ ( dz ) < ∞ .
Let p θ be an arbitrary parameterized probability density on the space E .
We denote by ϕ θ its characteristic function, deﬁned as: ϕ θ ( ω ) :
=
E z ∼ p θ [ e iω T z ] .
We denote by ˆ f the Fourier transform of the function f deﬁned as:
[Equation]
The inverse Fourier transform is given in this case by:
[Equation]
(2) where µ ( dω ) represents the measure in the Fourier domain.
In this paper we treat the cases where
E = R d for which µ ( dω ) =
dω (2 π ) d , and the case where E is a discrete set, for which the measure µ is deﬁned as: µ ( dω )
= 1 [ ω ∈ [ − π, π ] d ] dω (2 π ) d .
Throughout the paper, we reserve the letter i to denote the imaginary unit: i 2 = − 1 .
To denote higher order derivatives of the function f , we use the multi-index notation (Saint Raymond, 2018).
For a multi-index n = ( n 1 , ..., n d ) ∈ N d , we deﬁne:
[Equation]
where | n | = d (cid:88)
j = n j and ω n := d (cid:89)
j = ω n j j .
To clarify the multi-index notation, let us consider the example where d = 3 , and n = (1 , 0 , 2) , in this case:
[Equation]
The objective is to derive stochastic backpropagation rules, similar to that of (Rezende et al., 2014), for functions of the form: L ( θ ) :
=
E z ∼ p θ [ f ( z )] , for any arbitrary distribution p θ , discrete or continuous.
F OURIER S TOCHASTIC B ACKPROPAGATION
Stochastic backpropagation rules similar to that of (Rezende et al., 2014) can in fact be derived for any continuous distribution, under certain conditions on the characteristic function.
In the following theorem we present the main result of our paper concerning the derivation of Fourier stochastic backpropagation rules.
Theorem 1.
(Continuous Stochastic Backpropagation) Let f ∈ C ∞ ( R d , R + ) , under the condition that ∇ θ log ϕ
θ is a holomorphic function of iω , then there exists a unique { a n ( θ ) } n ∈ N d ∈ R such that:
[Equation]
Where { a n ( θ ) }
n ∈ N d are the Taylor expansion coefﬁcients of ∇ θ
log ϕ θ ( ω ) :
[Equation]
Proof.
Let us rewrite L in terms of ˆ f :
[Equation]
Fubini’s theorem
[Equation]
By introducing the derivative under the integral sign
[Equation]
and using the reinforce trick (Williams, 1992) applied to ϕ θ , where ∇ θ ϕ
[Equation]
ω
[Equation]
ϕ
[Equation]
∇ θ log ϕ θ ( ω ) , equation 5 becomes:
[Equation]
Under analyticity conditions of the gradient of the log characteristic function
[Equation]
we can expand the gradient term ∇ θ log ϕ θ ( ω ) , in terms of Taylor series around zero as:
[Equation]
Putting everything together, and replacing the characteristic function by its expression, the gradient of L becomes:
[Equation]
By rearranging the sums using Fubini’s theorem a second time
[Equation]
we obtain the following expression for the gradient:
[Equation]
Identically, we can follow the same procedure for discrete random variables.
We suppose that p θ factorizes over disjoint cliques of the dependency graph, where each dimension z j takes values in a discrete space V al ( z j ) .
In theorem 2 we derive the result concerning the discrete case.
Theorem 2.
(Discrete Stochastic Backpropagation) Let E be a discrete space: E = (cid:81) dj
=1 V al ( z j ) , and C the set of disjoint cliques of the dependency graph over z , that is,
[Equation]
then,
[Equation]
Where:
• z ∗ c : represents the normalizing assignment p θ ( z ∗ c )
= 1 − (cid:80)
z c (cid:54) =
z ∗ c p θ
[Equation]
z c
[Equation] [Equation]
Proof.
The characteristic function for the factored distribution is given by:
[Equation]
Thus the gradient of the log characteristic function becomes:
[Equation]
By plugging this expression to equation
[Equation]
we obtain:
[Equation]
The estimator of equation 10 has been derived in the literature through Rao–Blackwellization of the score function gradient, and it has been known under different names (Titsias & L´azaro-Gredilla, 2015; Asadi et al., 2017; Cong et al., 2019).
Theorem 2 shows that the discrete case can also be seen as backpropagating a derivative of the function f , in this case a discrete derivative given by equation 11.
A PPLICATIONS OF F OURIER S TOCHASTIC B ACKPROPAGATION
Following from the previous section, we derive the stochastic backpropagation estimators for certain commonly used distributions.
The multivariate Gaussian distribution: In this case p θ ( z ) =
N ( z ; µ θ , Σ θ ) .
The log characteristic function is given by: log ϕ θ
[Equation]
=
[Equation]
+
[Equation]
Tr (cid:2) Σ θ i 2 ωω T (cid:3) .
Thus by applying theorem 1, we recover the stochastic backpropagation rule of (Rezende et al., 2014):
[Equation] [Equation]
where, ∇ z and ∇ 2 z , represent the gradient and hessian operators.
The multivariate Dirac distribution: p θ ( z ) =
δ a θ ( z ) , the log characteristic function of the Dirac distribution is given by: log ϕ θ ( ω )
=
iω T a θ .
Thus the stochastic backpropagation rule of the Dirac is given by:
[Equation] [Equation]
resulting in the classical backpropagation rule.
In other words, the deterministic backpropagation rule is a special case of stochastic backpropagation where the distribution is a Dirac delta distribution.
This result provides a link between probabilistic graphical models and classical neural networks.
We investigate this link further in Appendix A.
The multivariate Bernoulli: p θ ( z ) = (cid:81) dj
=1 B ( z j ; π ( j ) θ ) , where π ( j ) θ = P [ z j = 1] .
By applying theorem 2, we obtain the local expectation gradient of (Titsias & L´azaro-Gredilla, 2015):
[Equation]
The multivariate categorical: p θ ( z ) =
(cid:81)
dj
=1 cat ( z j ; π ( j ) θ ) , where the dimensions are independent and take values in the set { 1 , ..., K } .
Similarly to the Bernoulli case, we obtain the following stochastic backpropagation rule:
[Equation]
The Laplace distribution: p θ ( z ) =
L ( z ; µ θ , b θ ) , in this case the log characteristic function is the following: log ϕ θ ( ω ) =
iµ θ ω − log(1 + b 2 θ ω 2 ) , using the Taylor series expansion for the function x (cid:55)→ 11 − x , we get the following stochastic backpropagation rule for the Laplace distribution:
[Equation]
The gamma distribution: p θ ( z ) =
Γ( z ; k θ , µ θ ) , the log characteristic function of the Gamma distribution is given by: log ϕ θ ( ω ) = − k θ log(1 − iµ θ ω ) .
By expanding it using Taylor series of the logarithm function, we obtain the following stochastic backpropagation rule:
[Equation]
The estimator of equation 20 gives a stochastic backpropagation rule for the gamma distribution and, hence also applies by extension to the special cases of the exponential, Erlang, and chi-squared distributions.
The beta distribution: p θ ( z ) =
Beta ( z ; α θ , β θ ) , in this case the characteristic function is the conﬂuent hypergeometric function:
ϕ θ ( ω ) =
1 F 1 ( α θ ; α θ + β θ ; iω ) .
A series expansion of the gradient of the log of this function is not trivial to derive.
However, we can use the parameterization linking the gamma and beta distributions to derive a stochastic backpropagation rule.
Indeed, if ζ 1 ∼ Γ( α θ , 1) and ζ 2 ∼ Γ( β θ , 1) , then z = g ( ζ 1 , ζ 2 ) =
[Equation]
2 ∼ Beta ( α θ , β θ ) .
By substituting in the gamma stochastic backpropagation rule, we obtain:
[Equation]
The Dirichlet distribution: p θ ( z ) =
Dir ( z ; K, α θ ) , following the same procedure, as for the beta distribution and using the following parameterization: z k =
[Equation]
Kj
[Equation] [Equation]
with, ζ k ∼ Γ
[Equation]
α
[Equation]
k
[Equation]
θ , 1) , we obtain:
[Equation]
T RACTABLE CASES & A PPROXIMATIONS OF F OURIER S TOCHASTIC B ACKPROPAGATION
The Fourier stochastic backpropagation gradient as presented in previous sections presents two major computational bottlenecks for non-trivial distributions.
The ﬁrst is the computation of inﬁnite series, and the second is evaluating higher order derivatives of the function f .
Depending on the application, the function f could be chosen in order to bypass the computational bottlenecks.
A trivial example, is if the higher order derivatives of the function f vanish at a certain order: ∂ nz f = 0 .
Another example, is the exponential function f ( z ) =
exp( (cid:15) T z ) .
From the fact that it obeys the following partial differential equation ∂f ∂z j ( z ) =
(cid:15)
j f ( z ) , one can deduce that the stochastic backpropagation rule reduces in this case to:
[Equation]
In most real world applications however, the inﬁnite sum will not often reduce to a tractable expression such as that of the exponential.
An example of this case is the evidence lower bound of a generative
[Figure]
v a r i a n c
[Figure]
i
[Figure]
model with Bernoulli observations.
In this case, the natural solution is to truncate the sum up to a ﬁnite order.
The assumption (although it might be wrong), is that the components associated to higher frequencies of the spectrum of the gradient of the log characteristic function, do not contribute as much.
And by analogy to the signal processing ﬁeld, we apply a Low-pass ﬁlter to eliminate them.
In this case the gradient of the log characteristic function of equation 7 becomes:
[Equation]
E XPERIMENTS
In our experimental evaluations, we test the stochastic backpropagation estimators of equationsand 20 for the gamma and Laplace distributions.
In the case of the gamma estimator, we use toy examples where we can derive exact stochastic backpropagation rules without truncating the inﬁnite sum.
As for the Laplace stochastic backpropagation rule, we test the estimator in the case of Bayesian logistic regression with Laplacian priors and variational posteriors on the weights.
We compare our estimators with the pathwise (Jankowiak & Karaletsos, 2019; Jankowiak & Obermeyer, 2018), and score function estimators, in addition to the weak reparameterization estimator in the gamma case (Ruiz et al., 2016).
We do not use control variates in our setup, the goal is to verify the exactness of the proposed inﬁnite series estimators and how they compare to current state-of-the-art methods in simple settings.
In all our experiments, we use the Adam optimizer to update the weights (Kingma &
Ba, 2014), with a standard learning rate of 10 − 3 .
In all the curves, we report the mean and standard deviation for all the metrics considered over 5 iterations.
6. T OY PROBLEMS
In the toy problem setting, we test the gamma stochastic backpropagation rule following the same procedure as (Mohamed et al., 2019).
we consider the following cases:
Toy problem 1 : L ( θ ) =
E z ∼ p θ (cid:2) ||
z − (cid:15) ||
2 (cid:3) , where p θ ( z ) =
(cid:81)
dj
=1 Γ( z j ; k j , µ j ) , θ = { k, µ } , and (cid:15) = .
In this case, we only need to compute the ﬁrst and second order derivatives of the function f .
Toy problem 2 : L ( θ ) =
E z ∼ p θ (cid:104)(cid:80) dj
=1 exp( − (cid:15)z j ) (cid:105) , in this case, the inﬁnite sum transfers to (cid:15) , which results in the following estimator: ∇ θ L = ∇ θ log ϕ θ (cid:0) (cid:15)i (cid:1)
E z ∼ p θ [ f ( z )] .
In ﬁgures 1 and 2 we report the training loss and log variance of the gradient across iterations of gradient descent for different values of the dimension d ∈ { 1 , 10 , 100 } .
The stochastic backpropagation estimator converges to the minimal value in all cases faster than the other estimators and the variance of the gradient is competitive with the pathwise gradient.
[Figure]
[Figure]
6. B AYESIAN LOGISTIC REGRESSION WITH L APLACIAN P RIORS
We evaluate the Laplace stochastic backpropagation estimator using a Bayesian
logistic
regression model (Jaakkola & Jordan, 1997), similarly to (Mohamed et al., 2019).
In our case, we substitute the normal prior and posterior on the weights with Laplace priors and posteriors.
We adopt the same notations of (Murphy, 2012), where the data, target and weight variables are respectively: x n ∈ R d , y n ∈ {− 1 , 1 } , and w .
The probabilistic model in our case is the following:
[Equation]
where σ represents the sigmoid function.
We consider Laplacian variational posteriors of the form p θ ( w ) =
(cid:81)
d j =1 L ( w j , µ j , b j ) , with θ = { µ, b } .
The evidence lower bound of a single sample is given by:
[Equation]
where the Kullback-Leibler divergence between the two Laplace distributions is the following:
[Equation]
We test the model on the UCI women’s breast cancer dataset (Dua & Graff, 2017), with a batch size of 64 and 50 samples from the posterior to evaluate the expectation.
In the case of the stochastic backpropagation estimator we truncate the inﬁnite series for the scale parameter b of equationto N = 4 and N = 8 .
In ﬁgure 3, we report the training evidence lower bound, the log variance of the gradient, and the accuracy computed on the entire dataset for the different estimators.
The stochastic backpropagation estimator converges faster than the considered estimators and the variance is signiﬁcantly lower.
We also notice that the truncation level of the inﬁnite series for the scale parameter has little effect on the outcome.
In ﬁgure 4, we report the bias and variance of the estimator at different values of the truncation level, for a ﬁxed parameter value during the training phase (epoch=100).
The bias and variance do not vary much, with the truncation level in this case, this
[Figure]
result conﬁrms the intuition of neglecting higher frequencies presented in section 5.
In addition, we compare the mean squared error between the Laplace stochastic backpropagation estimator and the score function and reparameterization estimators.
As shown in ﬁgure 4 the mean squared error is small, thus the values of the gradients across iterations are close.
However, the reparameterization gradient is closer to our estimator than the score function gradient, probably due to the fact that the reparameterization gradient is more stable and has lower variance.
R ELATED WORK & D ISCUSSION
Computing gradients through stochastic computation graphs has received considerable attention from the community, due to its application in many ﬁelds.
The ﬁrst general approach that provides a closed form solution for any probability distribution is the score function method (Glynn, 1989; Williams, 1992; Sutton et al., 2000; Schulman et al., 2015).
The main inconvenience of this approach, is that it results in high variance gradients when the dimension of the random variable becomes high.
In order to bypass this issue, the second approach consisted of designing control variates to reduce the variance of the score function estimator (Paisley et al., 2012; Weaver & Tao, 2013; Mnih & Gregor, 2014; Ranganath et al., 2014; Tokui & Sato, 2017).
In addition to the score function gradient, it was proposed to use an importance weighted estimator instead of the classical score function with a multi-sample objective (Mnih & Rezende, 2016; Burda et al., 2015).
The second class of approaches is that concerning reparameterization tricks (Kingma & Welling, 2013;
Rezende et al., 2014).
Through the decoupling of the computation of the gradient from the expectancy, reparameterization tricks have shown that they provide low-variance gradients using often a single sample.
The issue for these methods is the necessity to ﬁnd a reparameterization for each probability distribution.
Certain distributions such as the Gaussian are easy to reparametrize but others like the gamma are not.
In addition, discrete random variables do not admit an easy reparameterization as well.
Recently, these issues has been partially solved through implicit reparameterization, the generalized reparameterization gradient, and the pathwise gradient (Ruiz et al., 2016;
Figurnov et al., 2018; Jankowiak & Obermeyer, 2018).
For the discrete case, continuous relaxations that are reparameterizable have been proposed and combined with control variate methods (Gu et al., 2016;
Maddison et al., 2016; Jang et al., 2016; Tucker et al., 2017; Grathwohl et al., 2018).
Our approach, in contrast provides a new broad family of stochastic backpropagation rules derived using the Fourier transform.
One interesting aspect of our approach is the fact that the weighting a n is separated from the expectation of the higher order derivatives of the function f .
Thus the sampled variable does not intervene in the weighting in contrast to other methods such as reparameterization and pathwise gradients.
In addition, if the function f contains weakly correlated terms, by applying the derivative, some random variables would be eliminated.
Thus the estimators of the higher order derivatives would be sampled with respect to the derivation variable (and other correlated variables), which would result in lower variance.
It is worth noting that deriving stochastic backpropagation rules using the Fourier transform has been proposed for the Gaussian case (Fellows et al., 2018).
In our work, we extend it to non Gaussian distributions by way of the characteristic function, and exploiting the invariance of the functional inner product under Fourier transformation (Parseval’s theorem).
C ONCLUSION
In conclusion, in this paper we presented a new method to compute gradients through random variables for any probability distribution, by explicitly transferring the derivative to the random variable using the Fourier transform.
Our approach, gives a framework to be applied for any distribution, where the gradient of the log characteristic function is analytic, resulting in a new broad family of stochastic backpropagation rules, that are unique for each distribution
T HE D IRAC D ISTRIBUTION : THE LINK BETWEEN NEURAL NETWORKS AND PROBABILISTIC GRAPHICAL MODELS
[Figure]
y
In this section we explore the connection between neural networks and probabilistic graphical models following from the stochastic backpropagation rule of the Dirac delta distribution.
To this end, let us consider the probabilistic graphical model of ﬁgure 5.
The observed random variables in this model are denoted x and y representing the data and target variables.
We place the analysis in a supervised learning context, but the argument is valid for unsupervised models as well.
As usual the goal is to maximize the log likelihood for the data samples ( x, y ) , which is intractable, given that we need to integrate over the hidden variables.
However using variational inference, we can maximize an evidence lower bound of the form:
[Equation]
As suggested in the Dirac stochastic backpropagation rule, let us assume that the variational posteriors and priors are Dirac delta distribution of the form:
[Equation]
where, the a ( l ) , W ( l ) , and b ( l ) represent respectively the activation functions, the weights and biases for layer l , with the convention x :
= h (0) .
Under these assumptions, the Kullback-Leibler divergence term is equal to zero, and the evidence lower bound reduces to the the log-likelihood of a classical neural network:
[Equation]
with,
[Equation]
Thus, when using neural networks we are indirectly using a probabilistic graphical model and making the strong assumption that the hidden layers follow a parameterized Dirac distribution knowing the previous layer.
B E XPERIMENTS USING DISCRETE STOCHASTIC BACKPROPAGATION
We evaluate the Bernoulli and Categorical Stochastic Backpropagation estimators (BSB and CSB) of equations 17 and 18 on standard generative modeling benchmark tasks, using the MNIST and
Omniglot datasets (LeCun & Cortes, 2010; Lake et al., 2013).
We use the REBAR, RELAX, and Gumbel-softmax (or Concrete) estimators as baselines for our comparison (Jang et al., 2016;
Maddison et al., 2016; Tucker et al., 2017; Grathwohl et al., 2018).
The Bernoulli stochastic backpropagation is compared to the REBAR and RELAX estimators for three models: the sigmoid belief network of one and two stochastic hidden layers (Neal, 1992) and the variational autoencoder.
[Figure]
[Table]
In this case, we adopt the same architectures as (Grathwohl et al., 2018).
The categorical stochastic backpropagation estimator is compared to the Gumbel-softmax estimator (Maddison et al., 2016;
Jang et al., 2016) using two models: a variational autoencoder and a single layer belief network with categorical priors.
In this case, we set the dimension of the hidden layer to d = 20 and the number of modalities for each dimension to K = 10 .
All models are trained using the ADAM optimizer (Kingma & Ba, 2014) using a standard learning rate α = 10 − 4 and batch size of 100.
We train the models for 500 epochs on the MNIST dataset andepochs on the Omniglot dataset, longer learning epochs leads to overﬁtting and lower performance on the test sets for all estimators and models.
We perform 5 iterations of training in all experiments and we report the mean and standard deviation of each performance metric considered.
For all models and estimators, we report the mean marginal test likelihood in tables 1 and 2 for both datasets.
The test likelihood is estimated via importance sampling using 200 samples from the variational posterior.
In all cases the stochastic backpropagation estimator, a control variate free method outperforms the baselines.
In the case of the one layer sigmoid belief network the BSB estimator exhibits an increase of performance of about 4 nats in the case of the MNIST dataset and 10 nats in the case of the Omniglot dataset.
We also vary the number of samples used to estimate the expectation in the stochastic backpropagation rule S ∈ { 1 , 5 , 10 } .
We notice that using a single sample estimate does not hurt performance and leads to a faster training process.
We estimate the mean variance of the gradients w.r.t the parameters of the models using exponential moving averages of the ﬁrst and second moments computed by the ADAM optimizer.
The BSB estimator signiﬁcantly outperforms the REBAR, RELAX estimators in terms of variance reduction with a difference of about 2 nats in the case of sigmoid belief networks, and 1 nat in the case of the categorical variational autoencoder on the mnist dataset, leading to a more stable training process as shown in ﬁgures 6 and 7.
Finally, we evaluate the computational overhead of the categorical stochastic backpropagation estimator compared the Gumbel-softmax estimator.
We compare the two estimators in terms of execution time of one epoch of training.
The comparison is done using GPU implementations on
[Figure] [Figure]
a NVIDIA GeForce RTX 2080 Ti GPU, where the stochastic backpropagation rule of equation equation 18 is vectorized, thus leveraging the parallel batch treatment of the GPU.
As shown in table 3, the Gumbel-softmax method is faster than stochastic backpropagation ( S = 1 ) by a difference of about 3 seconds per training epoch.
This is due to the forward passes performed to compute each of the terms in equation equation 18.
The variance reduction and the increase in performance outweigh however the computational cost.