L OW B IAS G RADIENT E STIMATES FOR V ERY D EEP B OOLEAN S TOCHASTIC N ETWORKS
Anonymous authors
Paper under double-blind review
A BSTRACT
Stochastic neural networks with discrete random variables are an important class of models for their expressiveness and interpretability.
Since direct differentiation and backpropagation is not possible, Monte Carlo gradient estimation techniques have been widely employed for training such models.
Efﬁcient stochastic gradient estimators, such Straight-Through and Gumbel-Softmax, work well for shallow models with one or two stochastic layers.
Their performance, however, suffers with increasing model complexity.
In this work we focus on stochastic networks with multiple layers of Boolean latent variables.
To analyze such networks, we employ the framework of harmonic analysis for Boolean functions.
We use it to derive an analytic formulation for the source of bias in the biased Straight-Through estimator.
Based on the analysis we propose FouST , a simple gradient estimation algorithm that relies on three simple bias reduction steps.
Extensive experiments show that FouST performs favorably compared to state-of-the-art biased estimators, while being much faster than unbiased ones.
To the best of our knowledge FouST is the ﬁrst gradient estimator to train very deep stochastic neural networks, with up to 80 deterministic and 11 stochastic layers in our experiments.
I NTRODUCTION
Stochastic neural networks with discrete latent variables have been an alluring class of models for their expressivity and interpretability, dating back to foundational work on Helmholtz machines (Dayan et al., 1995) and sigmoid belief nets (Neal, 1992).
Since they are not directly differentiable, discrete random variables do not mesh well with the workhorse of modern Deep Learning, that is the backpropagation algorithm.
Monte Carlo gradient estimation is an effective solution where, instead of computing the true gradients, one can sample gradients from some distribution.
The sample estimates can be either biased or unbiased.
Unbiased gradient estimates like score function estimators (Williams, 1992) come typically at the cost of high variance leading to slow learning.
In contrast, biased gradient estimates such Straight-Through (Bengio et al., 2013), while efﬁcient, run the risk of convergence to poor minima and unstable training.
To this end several solutions have recently been proposed that either reduce variance in unbiased estimators (Mnih & Gregor, 2014; Gu et al., 2015; Tucker et al., 2017; Rezende et al., 2014; Grathwohl et al., 2017) or control bias in biased estimators (Jang et al., 2016; Maddison et al., 2016).
These methods, however, have difﬁculty scaling up to complex neural networks with multiple stochastic layers: low-variance unbiased estimators are too expensive 1 , while the compounded bias from the continuous relaxations on multiple stochastic layers leads to poor minima.
In this work we focus on biased estimators.
Our goal in this paper is a gradient estimator for Boolean random variables that works for any complex –deep or wide– neural network architecture.
We resort to the term Boolean instead of binary to emphasize that we work directly on the Boolean space {− 1 , +1 } , without any continuous relaxations or quantizations.
With this in mind we re-purpose the framework of harmonic analysis of Boolean functions, widely used in computational learning and computational complexity theory (O’Donnell, 2014; Linial et al., 1993; Mossel et al., 2003; Mansour, 1994).
We cast stochastic neural networks as Boolean functions f ( z ) over Boolean latent variables z sampled from probability
Require: Parameters θ ∈ R K , Bernoulli Representation {− a, a } , Interval Parameter b ∈ [0 , a ] , Constant scaling parameter γ distributions p ( z ) .
We then use harmonic analysis to determine that the bias in the Straight-Through gradient estimates corresponds to the weighted sum of higher-order Taylor coefﬁcients of f ( z ) .
The direct consequence is that one can control the bias in the Straight-Through estimator by manipulating the higher-order Taylor coefﬁcients of f ( z ) .
Building upon the harmonic analysis of existing gradient estimators, we present an algorithm, FouST , that admits low bias gradient estimates for Boolean latent variable models.
In experiments, we were able to scale up the stochastic depth and width of neural networks, training deep stochastic residual networks with up to 80 deterministic and 11 stochastic layers with little difﬁculty.
We summarize FouST in Algorithm 1.
With this work we make the following three contributions
We introduce the framework of harmonic analysis of Boolean functions to analyze discrete stochastic neural networks and their REINFORCE and Straight-Through gradients.
We show that stochastic gradients compute Fourier coefﬁcients.
Based on the above harmonic analysis we present FouST – a low-bias gradient estimator for Boolean latent variables based on three bias reduction steps.
As a side contribution, we show that the gradient estimator employed with DARN (Gregor et al., 2013), originally proposed for autoregressive models, is a strong baseline for gradient estimation in large and complex models with many stochastic layers.
We show that FouST is amenable to complex stochastic neural networks with Boolean random variables.
To the best of our knowledge, FouST is the ﬁrst gradient estimate algorithm that can train very deep stochastic neural networks with Boolean latent variables . The practical
outcome is a simple gradient estimate algorithm that can be plugged
in complex
stochastic neural networks
with
multiple layers
of
Boolean
random variables
H ARMONIC A NALYSIS OF B OOLEAN F UNCTIONS
We consider Boolean functions on the n -dimensional Boolean cube, f : {− 1 , +1 } n → R .
The setting of Harmonic Analysis for Boolean functions is the space of Boolean functions f with a product probability distribution on the Boolean input, that is p ( z )
= (cid:81) ni =1 p i ( z ) .
We denote p i as the probability of the i -th dimension being one, i.e. , p i :=
p i ( z i = +1) .
We denote the mean and variance of z i by µ i and σ i , respectively.
An example of a Boolean function in this setting is a generative neural network f : z (cid:55)→ y with a factorized latent distribution, as commonly done in representation learning (Kingma & Welling, 2013;
Higgins et al., 2017).
In this example, z is the stochastic input – also known as the latent code in stochastic neural networks – taking only two possible values and y is the output, like cross entropy loss.
Often, the goal of a generative neural network is to learn or approximate the latent distribution given input data x , i.e. , p ( z | x ) , as we will also explore in the experiments.
We ﬁrst introduce a few basic operations in the context of Harmonic Analysis of Boolean functions, which we shall use further on.
Further necessary details are in the Appendix A. For a more comprehensive introduction, however, we refer the reader to O’Donnell (2014).
Inner product.
The inner product of two Boolean functions f
and g is: (cid:104) f, g (cid:105) =
E p ( z ) [ f ( z ) g ( z )] .
Fourier expansion.
Let S be any subset of dimensions of the n -dimensional Boolean cube, S ⊆
[ n ] = { 1 , ..., n } .
Per S we deﬁne a basis function, φ S ( z ) :=
(cid:81)
i ∈ S φ
i ( z i ) , where for the empty set φ ∅ ( z ) = 1 and φ i is the z -score normalized dimension, i.e. , φ i :=
z i − µ i σ i .
For example, under the uniform Bernoulli distribution for the i -th dimension, p i ( z i = +1)
=
p i ( z i = − 1)
= 1 / 2 , we have φ i ( z i )
=
z i .
The 2 n functions φ S form an orthonormal basis for the space of Boolean functions f ,
[Equation]
for i (cid:54) =
j , where the expectations compute the inner product between two Boolean functions.
The last identity derives from the independence of any dimensions i (cid:54)
=
j .
We can then expand the Boolean function f on the set of 2 n orthonormal basis functions,
[Equation]
also known as the p -biased Fourier expansion of f .
The ˆ f ( p ) ( S ) are the Fourier coefﬁcients computed by the inverse Fourier expansion,
[Equation]
That is, the inverse expansion is computed by the inner product for Boolean functions deﬁned above.
The cardinality of S is the degree of the coefﬁcient ˆ f ( p ) ( S ) .
For instance, we have only one degree- 0 coefﬁcient, ˆ f ( p ) ( ∅ ) , which equals to the expected value of f under the distribution p ( z ) , ˆ f ( p ) ( ∅ ) =
E [ f ( z )] since φ ∅ = 1 .
Further, we have n degree-1 coefﬁcients ˆ f ( p ) ( i ) =
(cid:104) f, φ i (cid:105) and so on.
H ARMONIC A NALYSIS ON B IASED S TRAIGHT -T HROUGH G RADIENTS
We examine Straight-Through gradient estimates using the framework of Harmonic Analysis of
Boolean functions.
For training the model parameters with a loss function L := E p ( z )
[ f ( z )] we want to compute the gradient ∂ p i L =
∂ p i E p ( z )
[ f ( z )] for the i -th dimension.
As the random sampling operation in the expectation is not differentiable, Bengio et al.
(2013) propose the Straight-Through estimator that approximates the gradient with ∂ p i L ≈ E p ( z ) [ ∂ z i f ( z )] .
Clearly, the Straight-Through computes a biased approximation to the true gradient.
Next, we quantify the bias in the Straight-Through estimator using the harmonic analysis of f ( x ) .
For the quantiﬁcation of bias we ﬁrst need the following lemma that connects the REINFORCE gradients with the degree-1 Fourier coefﬁcients.
The lemma is an extension of Margulis-Russo (Margulis, 1974; Russo, 1982; O’Donnell, 2014) formula.
Lemma 1.
Let f be a Boolean function.
Then, the REINFORCE gradient estimates the degree Fourier coefﬁcients ˆ f ( p ) ( i ) of f under p ( x ) ,
[Equation]
Proof.
For compactness and clarity, we provide the proof in the Appendix B.1.
We introduce new notation.
In the following lemma bias ( p ) ( g i ST ) denotes the bias of the i -th gradient estimate under the distribution p .
Also, given the distribution p ( z ) , p i → 1 / 2 ( z ) is the distribution for which we set p ( z i = +1)
=
p ( z i = − 1) = 1 / 2 for a given dimension i .
Lemma 2.
Let f be a Boolean function, g i =
∂ p i E p ( z )
[ f ( z )] its true gradients and g i ST
=
E p ( z ) [ ∂
z
i f ( z )] its Straight-Through gradient estimator.
The bias in g i ST is due to a mismatch between the higher-order odd terms in the Taylor expansion of g i and g i ST and is equal to
[Equation]
where c k are the Taylor coefﬁcients for the i -th dimension on f ( z ) , that is z
i , around 0 and bias
[Equation]
Proof.
For compactness and clarity, we provide only a proof sketch here showing the basic steps.
These steps are also needed later in the description of the proposed gradient estimate algorithm.
For the detailed proof please refer to the Appendix B.2.
The proof sketch goes as follows.
First, we derive a relation between the Fourier coefﬁcients under the unknown distribution p ( z ) and under the uniform Bernoulli distribution p i
→ 1 / 2 ( z ) .
Then, using this relation we derive the Taylor expansions for the true gradient as well as the Straight-Through gradient estimator.
Last, to prove the lemma we compare the two Taylor expansions.
Relation between Fourier coefﬁcients under p ( z ) and p i → 1 / 2 ( z ) .
If we expand the function f in terms of its φ S basis as in equation 2 and focus on the i -th dimension, by Lemma 1 we can show that the REINFORCE gradient is given by
[Equation]
Taylor expansions of the true and the Straight-Through gradients.
The Taylor expansion of f ( z ) for z i around 0 is given by f ( z ) =
c
+ c 1 z
i + c 2 z 2 i + c 3 z 3 i + c 4 z 4 i + c 5 z 5 i + · · · , where c k k ! =
∂
kz
i f
( z ) |
z
i =0 are the Taylor coefﬁcients.
All c k are a function of z j , j (cid:54) =
i .
Let’s ﬁrst focus on the true gradient.
Since we work with Boolean ± 1 values, we have that z ki
=
for even k and z ki
= z i for odd k .
This will inﬂuence the even and the odd terms of the Taylor expansions.
Speciﬁcally, for the Taylor expansion of the true gradient we can show that
[Equation]
The expression in equation 7 implies that the true gradient with respect to the p i is the expected sum of the odd Taylor coefﬁcients.
Here we note that although the ﬁnal expression in equation 7 can also be derived by a ﬁnite difference method, it does not make explicit, as in equation 31, the dependence on z i and µ i of the term inside the expectation.
Now, let’s focus on the Straight-Through gradient.
Taking the derivative of the Taylor expansion w.r.t.
to z i , we have
[Equation]
The Straight-Through gradient is the expectation of equation 8 in the i -th dimension, that is
[Equation]
{+where z 2 j i = 1 , z 2 j +
+}
i = z i , and E p ( z i ) [ z i ] = µ i .
Comparing the Taylor expansions.
By comparing the expansion of the Straight-Through gradient in equation 10 and the expansion of the true gradient in equation 7 and given that bias ( p ) ( g i ST ) = E p [Equation] z ) [
[Equation]
i
[Equation]
z
[Equation]
] −
[Equation]
p i E p
[Equation]
[ f
[Equation]
] , bias (
[Equation]
[Equation]
(cid:88)
[Equation]
(cid:88) k
[Equation]
kc
[Equation]   µ [Equation] Taking the expectation in equation 9 under p i → 1 / 2 causes the ﬁnal term in equation 11 to vanish leaving bias ( p i → 1 / 2 ) ( g i ST ) =
E p ( z \ i ) (cid:104)(cid:80) k
=
j +1 ,j> 0 ( k − 2)
c k (cid:105) .
Combining this expression with equation 9 gives the ﬁnal expression (equation 4) from the lemma.
L OW - BIAS G RADIENT E STIMATES 4. F OU ST G RADIENT E STIMATION A LGORITHM
Inspired by the Harmonic Analysis of the Straight-Through gradient estimates, we present a gradient estimate algorithm for deep Boolean latent models, FouST , for Fourier Straight-Through estimator.
The algorithm relies on three bias reduction steps on the Straight-Through, lines 2, 3, 5 in Algorithm 1.
4. L OWERING B IAS BY I MPORTANCE S AMPLING C ORRECTION
As detailed earlier, the bias in the Straight-Through estimator is the sum of the bias under the uniform
Bernoulli distribution plus extra bias due to non-zero expectation terms in higher-order harmonics when sampling from p ( z ) .
Sampling from p i
→ 1 / 2 instead of p ( z ) would decrease the total bias from the form in equation 11 by setting the ﬁnal term to 0.
As a ﬁrst bias reduction step, therefore, we do importance sampling.
Speciﬁcally, after getting samples from p ( z ) and computing the gradients ∂
z i f ( z ) with the Straight-Through, we estimate the expectation under p i
→ 1 / 2 as
[Equation]
Interestingly, Gregor et al.
(2013) arrive at equation 12 in the context of unbiased control variates for quadratic functions.
4. L OWERING B IAS BY M ATCHING U NIFORM D ISTRIBUTION M OMENTS
Lemma 2 shows that part of the bias in the Straight-Through estimator is due to the presence of extra factors in the Taylor coefﬁcients.
We can reduce the effect of these factors by taking advantage of the moments of the uniform distribution.
Recalling that (cid:82) b 0 1 b u k du = b k k +1 , we can attempt to correct the coefﬁcients in equation 9, which for z k have the form ( k + 1) c k , with the same extra factor of k + 1 that appears in the denominator of the k th moment.
This suggests that we can sample from an auxiliary variable u and then use the auxiliary variable u with f instead of z and exploit the higher moments of the uniform distribution to reduce bias.
For brevity, we illustrate the method with a case study of a two-dimensional z , and a bivariate f ( z 1 , z 2 ) .
As in Lemma 2, the partial true gradient of f ( z 1 , z 2 ) w.r.t.
the ﬁrst distribution parameter p 1 equals to
[Equation]
“Bernoulli splitting uniform” trick.
Assume an auxiliary variable u
= ( u 1 , u 2 ) , which we choose as follows.
First, we sample z
=
( z 1 , z 2 ) from the uniform Bernoulli distribution p 1 → 1 / 2 with ( i set to 1).
Then we take a uniform sample ( u 1 , u 2 ) with u i sampled from either [0 , 1] for z i
= +1 or from [ − 1 , 0]
if z i = − 1 .
The expectation of the gradient under such random sampling is
[Equation]
[Equation]
Further detail is in Appendix C.1.
We compare equation 14 with equation 13.
In equation 14 we observe that the pure terms in z 1 , namely terms with j = 0 , always match those of the true gradient in equation 13.
For j > 0 we obtain mixed terms with coefﬁcients that do not match those of the true gradient in equation 13.
Due to the 1 j +1 factor, for small j the mixed-degree terms are closer to the original ones in equation 13.
For functions with small mixed degree terms, this can lead to bias reduction, at the cost of an increased variance because of sampling an auxiliary variable.
In practice, to manage this bias-variance trade-off and to deal with functions that have greater dependence on mixed degree terms, we use smaller intervals for the random samples as in Algorithm 1.
To summarize, for appropriate functions, the “Bernoulli splitting uniform” relies on the continuous variable u conditioned on the binary sample to reduce the bias.
However, it is important to emphasize that u is only an auxiliary variable; the actual latent variable z is always binary .
Thus, the “Bernoulli splitting uniform” trick does not lead to a relaxation of the sort used by Gumbel-Softmax (Jang et al., 2016), where there are no hard samples.
Lastly we note that for a univariate f the “Bernoulli splitting uniform” trick leads to an unbiased estimator with an increased variance.
4. L OWERING B IAS BY R EPRESENTATION S CALING
The Fourier basis does not depend on the particular input representation and any two-valued set, say {− t, t } can be used as the Boolean representation.
The choice of a representation, however, does affect the bias as we show next.
As a concrete example, we let our input representation be z i ∈ {− 1 / 2 , 1 / 2 } n , where p i = p ( z i
= +1 / 2) .
While we can change the input representation like that, in general the Fourier coefﬁcients in equation 3 will be different than for the {− 1 , +1 } representation.
We give the ﬁnal forms of the gradients here.
Details are given in Appendix C.2.
Under the p i → 1 / 2 distribution the degree-1 Fourier coefﬁcients are:
[Equation]
Note that compared to equation 7, in equation 15 we still get the odd terms c 1 , c 3 albeit decayed by inverse powers of 2.
Following the same process for the Straight-Through gradient as in equation 10, we have that
[Equation]
While this is still biased, compared to equation 7 the bias is reduced by damping higher order terms by inverse powers of 2.
4. A LGORITHMIC C OMPLEXITY
The algorithm, described in algorithm 1, is a Straight-Through gradient estimator with the bias reduction steps described above, where a single sample is used to estimate the gradient.
We emphasize that the algorithm uses a single sample and a single evaluation of the decoder per example and latent vector sample.
Thus, the algorithm has the same complexity as that of the original Straight-Through estimator.
R ELATED W ORK
Monte Carlo gradient estimators for training models with stochastic variables can be biased or unbiased.
Perhaps the best known example of an unbiased gradient estimator is the REINFORCE algorithm (Williams, 1992).
Unfortunately, REINFORCE gives gradients of high variance.
For continuous stochastic variables Kingma & Welling (2013) propose the reparameterization trick, which transforms the random variable into a function of deterministic ones perturbed by a ﬁxed noise source, yielding much lower variance gradient estimates.
For discrete stochastic variables,
REINFORCE is augmented with control variates for variance reduction.
A number of control variate schemes have been proposed: NVIL (Mnih & Gregor, 2014) subtracts two baselines (one constant and one input-dependent) from the objective to reduce variance.
MuProp (Gu et al., 2015) uses the ﬁrst-order Taylor approximation of the function as a baseline.
REBAR (Tucker et al., 2017) uses the
Gumbel-Softmax trick to form a control variate for unbiased gradient estimates.
RELAX (Grathwohl et al., 2017) generalizes REBAR to include an auxiliary network in the gradient expression and uses continuous relaxations and the reparameterization trick to give unbiased gradients.
Regarding biased estimators, a simple choice is the Straight-Through estimator (Bengio et al., 2013) which uses the gradient relative to the sample as that relative to the probability parameter.
Another recent approach is to use continuous relaxations of discrete random variables so that the reparameterization trick becomes applicable.
The most common example of this being the GumbelSoftmax estimator (Maddison et al., 2016; Jang et al., 2016).
Although this is a continuous relaxation, it has been used to deﬁne the Gumbel Straight-Through estimator with hard samples.
This uses arg max in the forward pass and the Gumbel-Softmax gradient is used as an approximation during in the backward pass.
DARN (Gregor et al., 2013), like MuProp, also uses the ﬁrst-order Taylor expansion as a baseline but does not add the analytical expectation, making the estimator biased for non-quadratic functions.
In this work we focus on biased Straight-Through gradient estimators.
Speciﬁcally, we analyse how to reduce bias via Fourier expansions of Boolean functions.
The Fourier expansion itself is widely used in computational learning theory with applications to learning low-degree functions (Kushilevitz & Mansour, 1993), decision trees (Mansour, 1994), constant-depth circuits (Linial et al., 1993) and juntas (Mossel et al., 2003).
To the best of our knowledge we are the ﬁrst to explore Fourier expansions for bias reduction of biased stochastic gradient estimators.
E XPERIMENTS
Experimental Setup.
We ﬁrst validate FouST on a toy setup, where we already know the analytic expression of f ( z ) .
Next we validate FouST by training generative models using the variational autoencoder framework of Kingma & Welling (2013).
We optimize the single sample variational lower bound (ELBO) of the log-likelihood.
We train variational autoencoders exclusively with
Boolean latent variables on OMNIGLOT, CIFAR10, mini-ImageNet (Vinyals et al., 2016) and
MNIST (Appendix D.1).
We train all models using a regular GPU with stochastic gradient descent with a momentum of 0.9 and a batch size of 128.
We compare against Straight-Through, GumbelSoftmax, and DARN, although on more complex models some estimators diverge.
The results were consistent over multiple runs.
Details regarding the architectures and hyperparameters used are in
Appendix E. Upon acceptance we will open source all code, models, data and experiments.
6. B IASED E STIMATORS ON A T OY P ROBLEM
[Figure]
To validate the excessive bias in the Straight-Through estimator, as well as to explore the beneﬁts of FouST, we ﬁrst explore a toy problem.
Similar to Tucker et al.
(2017), we minimize E p ( z ) [( z − t ) 2 ] , where t ∈ (0 , 1) is a continuous target value and z is a sample from the Bernoulli distribution p ( z ) .
The optimum is obtained for p ( z = +1) ∈ { 0 , 1 } .
Figure 1, shows a case with t = 0 .
45 , where the minimizing solution is p ( z = +1)
= 0 .
We observe that unlike the
Straight-Through estimator, FouST converges to the minimizing deterministic solution (lower is better).
6. F OU ST VS . S TATE - OF - THE -A RT G RADIENT E STIMATORS
Training Stochastic MLPs.
We train MLPs with one and two stochastic layers on OMNIGLOT, following the non-linear architecture of Tucker et al.
Each stochastic Boolean layer is preceded by two deterministic layers of 200 tanh units.
All hyperparameters remain ﬁxed throughout the training.
All estimators use one sample per example and a single decoder evaluation .
We present results in Fig.
FouST outperforms other biased gradient estimators in both datasets and architectures.
FouST is clearly better than the StraightThrough estimator.
Despite the complicated nature of the optimized neural network function f ( z ) the bias reduction appears fruitful.
With one or two stochastic layers we can also use the unbiased
REBAR.
REBAR is not directly comparable to the estimators we study, since it uses multiple decoder evaluations and for models with multiple stochastic layers, multiple passes through later layers.
Nevertheless, as shown in appendix D.1 for MNIST, with two stochastic layers REBAR reaches a worse test ELBO of -94.43 v. -91.94 for FouST.
A possibility of worse test than training ELBOs for REBAR was also suggested in the original work (Tucker et al., 2017).
Training Stochastic ResNets.
We further validate FouST in a setting where the encoder and decoder are stochastic ResNets, S-ResNets, which are standard ResNets with stochastic layers inserted between ResNet blocks.
Similar to MLPs, FouST outperforms other biased estimators in this setting on CIFAR-10 (left in Figure 3).
Note that despite the hyperparameter sweep, we were unable to train S-ResNet’s with Gumbel-Softmax.
So we compare against DARN and Straight-Through only.
With an S-ResNet with 12 ResNet blocks and 4 stochatic layers FouST yields a score of 5.08 bits per dimension (bpd).
This is comparable to the 5.14 bpd with the categorical VIMCO-trained model (Mnih & Rezende, 2016) obtained with 50 samples per example, as reported by van den Oord et al. [Table]
In the plots, we observe sharp spikes or slower curving.
We hypothesize these are due, respectively, to stochasticity and bias, and are corrected to some degree along the way.
Efﬁciency.
We compare the efﬁciency of different estimators in Tab. 1.
Like other biased estimators, FouST requires a single sample for estimating the gradients and has similar wallclock times.
On MNIST, the unbiased REBAR is 15x and 40x slower than the biased estimators for two and ﬁve stochastic layer MLP’s respectively.
From the above experiments we conclude that FouST allows for efﬁcient and effective training of fully connected and convolutional neural networks with Boolean stochastic variables.
6. A DDING S TOCHASTIC D EPTH AND W IDTH
Last, we evaluate FouST on more complex neural networks with deeper and wider stochastic layers.
We perform experiments with convolutional architectures on the larger scale and more realistic mini-ImageNet (Vinyals et al., 2016).
As the scope of this work is not architecture search, we present two architectures inspired from residual networks (He et al., 2016) of varying stochastic depth and width.
The ﬁrst one is a wide S-ResNet, S-ResNet-40-2-800, and has 40 deterministic (with encoder and decoder combined), 2 stochastic layers, and 800 channels for the last stochastic layer.
The second,
S-ResNet-80-11-256, is very deep with 80 deterministic and 11 stochastic layers, and a last stochastic layer with 256 channels.
Architecture details are given in Appendix E.2.
In this setup, training with existing unbiased estimators is intractable.
We present results in Fig.
We compare against DARN, since we were unable to train the models with Gumbel-Softmax.
Incomplete lines indicate failure.
We observe that FouST is able to achieve better training ELBO’s in both cases.
We conclude that
FouST allows for scaling up the complexity of stochastic neural networks in terms of stochastic depth and width. [Figure] [Figure]
C ONCLUSION
In this work we introduce the framework of harmonic analysis for Boolean functions.
We then use the harmonic analysis to derive the source of bias in the Straight-Through estimator for Boolean random variables.
Based on the analysis we propose FouST, which is a novel gradient estimate algorithm.
We use FouST to train neural networks with Boolean random variables in stochastic layers.
FouST outperforms state-of-the-art biased gradient estimators, while maintaining efﬁciency that is orders of magnitude higher than unbiased estimators.
Importantly, to the best of our knowledge FouST is the ﬁrst gradient estimate algorithm, biased or unbiased, to be able to train very deep and wide neural networks with Boolean random variables.
In our experiments FouST successfully trained networks with up to 80 deterministic and 11 stochastic layers.
We conclude that the Harmonic
Analysis framework for Boolean function is a useful methodological tool for analyzing Boolean neural networks.
Also, we conclude that FouST is a practical gradient estimate algorithm that can train complex stochastic neural networks with multiple layers of Boolean random variables.
A PPENDICES A F URTHER D ETAILS FROM THE H ARMONIC A NALYSIS F RAMEWORK FOR B OOLEAN F UNCTIONS A. D ISCRETE DERIVATIVE
For a Boolean function f the discrete derivative on the i -th latent dimension with a basis function φ i is deﬁned as
[Equation]
The Fourier expansion of the discrete derivative equals D φ i f ( z )
= (cid:80) S (cid:51) i
ˆ f ( p ) ( S ) φ S \
i ( z ) , The Fourier expansion of the discrete derivative is derived by equation 2: (i) all bases that do not contain the i -th dimension are constant to φ i and thus set to zero, while (ii) for the rest of the terms ∂φ S
dφ i =
φ S \ i from the deﬁnition of basis functions φ S ( z ) .
In the following we differentiate partial derivatives on continuous functions noted with ∂ · from discrete derivatives on Boolean functions noted with D · .
The i -th discrete derivative is independent of z i both in the above deﬁnitions.
B H ARMONIC A NALYSIS OF E XISTING G RADIENT E STIMATES B. P ROOF OF L EMMA
Proof.
We follow O’Donnell (2014, §8.4).
In this proof we work with two representations of the Boolean function f .
The ﬁrst is the Fourier expansion of f under the uniform Bernoulli distribution.
This is also the representation obtained by expressing f as a polynomial in z i .
Since the domain of the function f is the Boolean cube, the polynomial representation is multilinear.
That is f
( z ) = (cid:80) S ⊆
[ n ] ˆ f ( S ) (cid:81) j ∈ S z j .
In what follows, f ( z ) denotes this representation.
Note that since this representation is a polynomial it is deﬁned over any input over R n .
In particular,
[Equation] j [Equation] The second representation we use is the Fourier expansion of the Boolean function f under p ( x ) .
We denote this by f ( p ) .
The following relation follows from the fact that when working with the Fourier representation, f ( z ) is multilinear, E p ( z )
[ z i ] = µ i and the linearity of expectation.
[Equation]
As the partial derivative of f w.r.t.
µ i is equivalent to discrete derivative of f w.r.t.
z i , ∂ µ i f ( µ ) =
D z i f ( µ ) , and keeping in mind that φ i =
( z i − µ i ) /σ
i , we have that
[Equation]
We then note that the discrete derivative of f w.r.t.
z i , D
z
i f ( µ ) , from the left hand side of equation 19, is equivalent to the partial derivative of f w.r.t.
µ i , ∂ µ i f ( µ ) .
[Equation]
We complete the proof by combining equation ?? , equation ?? , equation 22, equation 24.
B. A NALYSIS OF S TRAIGHT -T HROUGH
The detailed proof of Lemma 2 is as follows.
Proof.
We ﬁrst derive the Taylor expansions for the true gradient as well as the Straight-Through gradient estimator.
Then, we prove the lemma by comparing the two Taylor expansions.
By expanding the function f
in terms of its φ S basis as in equation 2 and focusing on the i -th dimension, we have that
[Equation]
The ﬁrst term, ˆ f ( p ) ( i ) z i σ i , is the term corresponding to { i } in the Fourier expansion of f under p i
→ 1 / 2 ( z ) .
That is
[Equation]
This follows from the fact that when moving from p ( z ) to p i → 1 / 2 ( z ) (i we have that φ i =
z i , and (ii) no other term under the p ( z ) expansion contributes to the z term the p
→ / 2 ( z ) expansion.
As a consequence of Lemma 1 the REINFORCE gradient for the i -th dimension is given by
[Equation]
Next, we will derive the Taylor expansions of the true and the Straight-Through gradients.
The Taylor expansion of f ( z ) for z i around 0 is
[Equation]
where c k k !
=
∂
kz
i f
( z ) |
z
i =0 are the Taylor coefﬁcients.
All c k are a function of z j , j (cid:54) =
i .
Let’s ﬁrst focus on the true gradient.
Since we work with binary ± 1 values, we have that z
ki
= 1 for even k and z ki =
z i for odd k .
This will inﬂuence the even and the odd terms of the Taylor expansions.
Speciﬁcally, for the Taylor expansion of the true gradient we have from equation 27 and equation 3 that
[Equation]
The expression in equation 32 implies that the true gradient with respect to the p i is the expected sum of the odd Taylor coefﬁcients.
Here we note that the although ﬁnal expression in equation 32 can also be derived by a ﬁnite difference method, it does not make explicit, as in equation 31, the dependence on z i and µ i of the term inside the expectation.
Now, let’s focus on the Straight-Through gradient.
Taking the derivative of the Taylor expansion in equation 28 w.r.t.
to z i , we have
[Equation]
The Straight-Through gradient is the expectation of equation 33 in the i -th dimension, that is
[Equation]
By comparing the expansion of the Straight-Through gradient in equation 35 and the expansion of the true gradient in equation 32,
[Equation]
Taking the expectation in equation 34 under p i → 1 / 2 causes the ﬁnal term in equation 37 to vanish leaving bias
[Equation]
C L OW - BIAS G RADIENT E STIMATES C. L OWERING B IAS BY M ATCHING U NIFORM D ISTRIBUTION M OMENTS
We describe the case of a bivariate function in detail.
For brevity, we focus on a case study of a two-dimensional z , and a bivariate f ( z 1 , z 2 ) with the bivariate Taylor expansion f
( z 1 , z 2 ) =
(cid:80) i,j c i,j z
i 1 z j 2 .
As in Lemma 2, the partial true gradient of f ( z 1 , z 2 ) w.r.t.
the ﬁrst distribution parameter p 1 equals to
[Equation]
[Equation]
Further, the Taylor expansion of ∂ p 1 f ( z 1 , z 2 ) is ∂ z 1 f
[Equation]
z 1 , [Equation] 21 + · · ·
“Bernoulli splitting uniform” trick.
Assume an auxiliary variable u
= ( u 1 , u 2 ) , which we choose as follows.
First, we sample z
=
( z 1 , z 2 ) from the uniform Bernoulli distribution p 1 → 1 / 2 .
Then we take a uniform sample ( u 1 , u 2 ) with u i sampled from either [0 , 1] for z i
= +1 or from [ − 1 , 0]
if z i = − 1 .
At this point it is important to note that the moments of the uniform distribution in [ a, b ] are (cid:82) b a 1 b − a u k = 1 b − a b k +1 − a k +1 k
+ , which simpliﬁes to b/ 2 , b 2 / 3 , b 3 / 4 , . . .
for a = 0 , and where we think of b as a binary sample i.e., b ∈ {− 1 , 1 } .
The expectation of the gradient under such random sampling is
[Equation]
We then compare equation 40 with equation 39.
In equation 40 we observe that the pure terms in z 1 , namely terms with j = 0 , always match those of the true gradient in equation 39.
For j > 0 we obtain mixed terms with coefﬁcients that do not match those of the true gradient in equation 39.
However, the partial gradient obtained with the auxiliary variables in equation 40 has coefﬁcients following a decaying trend due to the 1 j +1 .
For small j , that is, the mixed-degree terms are closer to the original ones in equation 39.
For functions with smaller mixed degree terms this leads to bias reduction, at the cost of an increased variance due to additional sampling. In practice many functions would have greater dependence on mixed degree terms.
For such functions and to manage the bias-variance trade-off we choose smaller intervals for the uniform samples, that is a → b .
C. L OWERING B IAS BY R EPRESENTATION S CALING
The Fourier basis does not depend on the particular input representation and any two-valued set, say {− t, t } can be used as the Boolean representation.
The choice of a representation, however, does affect the bias as we show next.
As a concrete example, we let our input representation be z i ∈ {− 1 / 2 , 1 / 2 } n , where p i = p ( z i
= +1 / 2) .
While we can change the input representation like that, in general the Fourier coefﬁcients in equation 3 will be different than for the {− 1 , +1 } representation.
Letting h ( z i ) =
2 z i ∈ {− 1 , 1 } , the functions φ i are now given as φ i
= h ( z i ) − µ i σ i .
Next, we write the Taylor series of f in terms of h ( z i ) , [Equation]
[Equation]
Under the p i → 1 / 2 distribution, we still have that E p → 1 / 2 [ h ( z i )]
= 0 and the degree-1 Fourier coefﬁcients are:
[Equation]
Note that compared to equation 7, in equation 43 we still get the odd terms c 1 , c 3 albeit decayed by inverse powers of 2.
Following the same process like for equation [Figure] we have that [Figure]
[Equation]
[Equation]
D E XTRA E XPERIMENTS D. E XPERIMENTS ON MNIST
[Figure]
Method
Stochastic Layers
Hidden Units per Layer
Test ELBO
Rebar
[Table]
D. A BLATION E XPERIMENTS
To further judge the effect of our proposed modiﬁcations to Straight-Through, we performed ablation experiments where we separately applied scaling and noise to the importance-corrected Straight-Through. These experiments were performed on the single stochastic layer MNIST and OMNIGLOT models. The results of the ablation experiments are shown in ﬁgure 5.
From the ﬁgure it can be seen that scaling alone improves optimization in both cases and noise alone helps in the case of MNIST.
Noise alone results in a worse
ELBO in the case of OMNIGLOT, but gives an improvement when combined with scaling.
From these results we conclude that the proposed modiﬁcations are effective.
[Figure]
E A RCHITECTURES U SED IN THE EXPERIMENTS E. A RCHITECTURES FOR MNIST AND O MNIGLOT
The encoder and decoder networks in this case are MLP’s with one or more stochastic layers.
Each stochastic layer is preceded by 2 deterministic layers with a tanh activation function.
We chose learning rates from { 1 × 10 − 4 , 2 × 10 − 4 , 4 × 10 − 4 , 6 × 10 − 4 } , Gumbel-Softmax temperatures from { 0 . 1 , 0 . 5 } , and noise interval length for FouST from { 0 . 1 , 0 . 2 } .
E. A RCHITECTURES FOR CIFAR-10 AND MINI -I MAGE N ET
For these dataset we use a stochastic variant or ResNets (He et al., 2016).
Each network is composed of stacks of layers .
Each layer has (i) one regular residual block as in He et al.
(2016), (ii) followed by at most one stochastic layer, except for the CIFAR architecture B in ﬁgure 3 where we used two stochastic layers in the last layer.
The stacks are followed by a ﬁnal stochastic layer in the encoder.
We do downsampling at most once per stack.
We used two layers per stack.
For CIFAR we downsample twice so that the last stochastic layer has feature maps of size 8x8.
We chose learning rate from { 9 × 10 − 7 , 1 × 10 − 6 , 2 × 10 − 6 , 4 × 10 − 6 } , the FouST scaling parameter from { 0 . 5 , 0 . 8 , 0 . 9 } , and the uniform interval was scaled by a factor from { 0 . 01 , 0 . 05 , 0 . 1 }
For mini-ImageNet we downsample thrice. We chose the learning rate from { 2 × 10 − 7 , 3 × 10 − 7 , 4 × 10 − 7 , 5 × 10 − 7 } .
[Table]