Title: LOW BIAS GRADIENT ESTIMATES FOR VERY DEEP BOOLEAN STOCHASTIC NETWORKS

Abstract: 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. Efficient 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 first gradient estimator to train very deep stochastic neural networks, with up to 80 deterministic and 11 stochastic layers in our experiments.

1. INTRODUCTION:

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 efficient, 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 difficulty scaling up to complex neural networks with multiple stochastic layers: low-variance unbiased estimators are too expensive1 , 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 Algorithm 1 FouST Gradient Estimator Require: Parameters θ ∈ R K , Bernoulli Representation {-a, a}, Interval Parameter b ∈ [0, a],

Constant scaling parameter γ 1: Sample x i ∼ p θi (x i ), i = 1, ..., K This and all steps below have constant complexity. 2: Sample y i ∼ Unif(b, a) 3: Set y i := x i * y i 4: Compute ST gradient ∂ θi := ∂ yi f (y) 5: Importance reweighing ∂ θi := 0.5/p(x i ) • ∂ θi 6: return partial gradient γ∂ θi , i = 1, ..., K 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 coefficients of f (z). The direct consequence is that one can control the bias in the Straight-Through estimator by manipulating the higher-order Taylor coefficients 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 difficulty. We summarize FouST in Algorithm 1.

With this work we make the following three contributions.

1. 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 coefficients.

2. 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.

3. We show that FouST is amenable to complex stochastic neural networks with Boolean random variables. To the best of our knowledge, FouST is the first 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.

2. HARMONIC ANALYSIS OF BOOLEAN FUNCTIONS:

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) = n i=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 → 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 first 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:

f, g = 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 define a basis function, φ S (z) := i∈S φ i (z i ), where for the empty set φ ∅ (z) = 1 and φ i is the z-score normalized dimension, i.e., φ i := zi-µ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

E p(z) [φ i (z i ) 2 ] = 1, E p(z) [φ i (z i )φ j (z j )] = E p(z) [φ i (z i )] E p(z) [φ j (z j )]

, for i = j, where the expectations compute the inner product between two Boolean functions. The last identity derives from the independence of any dimensions i = 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 coefficients computed by the inverse Fourier expansion,

EQUATION

That is, the inverse expansion is computed by the inner product for Boolean functions defined above. The cardinality of S is the degree of the coefficient f (p) (S). For instance, we have only one degree-0 coefficient, 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 coefficients f (p) (i) = f, φ i and so on.

3. HARMONIC ANALYSIS ON BIASED STRAIGHT-THROUGH GRADIENTS:

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

∂ pi L = ∂ pi 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

∂ pi L ≈ E p(z) [∂ zi 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 quantification of bias we first need the following lemma that connects the REINFORCE gradients with the degree-1 Fourier coefficients. 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 1 Fourier coefficients

f (p) (i) of f under p(z), E p(z) g i REINFORCE = ∂ pi E p(z) [f (z)] = 2 σ i f (p) (i).

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 = ∂ pi E p(z) [f (z)] its true gradients and g i ST = E p(z) [∂ zi 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

bias (p) g i ST = bias (p i→1/2 ) g i ST + E   k=2j,j>0 kc k   µ i (4)

where c k are the Taylor coefficients for the i-th dimension on f (z), that is z i , around 0 and bias (p i→1/2 ) = E k=2j+1,j>0 (k -2)c k .

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 coefficients 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 coefficients 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

∂ pi E p(z) [f (z)] = 2 f (p) (i) σ i = 2 f (p→1/2) (i) (5)

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 0 + 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! = ∂ k zi f (z)| zi=0

are the Taylor coefficients. All c k are a function of z j , j = i. Let's first focus on the true gradient. Since we work with Boolean ±1 values, we have that z k i = 1 for even k and z k i = z i for odd k. This will influence the even and the odd terms of the Taylor expansions. Specifically, for the Taylor expansion of the true gradient we can show that

∂ pi [E p(z) [f (z)]] = 2 f (p→1/2) (i) = 2E p i→1/2 (z) [c 0 z i + c 1 z 2 i + c i z 3 i + • • • ] (6) = 2E p(z \i ) [c 1 + c 3 + c 5 + • • • ] (7)

The expression in equation 7 implies that the true gradient with respect to the p i is the expected sum of the odd Taylor coefficients. Here we note that although the final expression in equation 7 can also be derived by a finite 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 2j i = 1, z 2j+1 i = z i , and E p(zi) [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

EQUATION

Taking the expectation in equation 9 under p i→1/2 causes the final term in equation 11 to vanish leaving bias

(p i→1/2 ) (g i ST ) = E p(z \i ) k=2j+1,j>0 (k -2)c k .

Combining this expression with equation 9 gives the final expression (equation 4) from the lemma.

4. LOW-BIAS GRADIENT ESTIMATES:



4.1. FOUST GRADIENT ESTIMATION ALGORITHM:

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.2. LOWERING BIAS BY IMPORTANCE SAMPLING CORRECTION:

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 final term to 0.

As a first bias reduction step, therefore, we do importance sampling. Specifically, after getting samples from p(z) and computing the gradients ∂ zi f (z) with the Straight-Through, we estimate the expectation under p i→1/2 as

E p i→1/2 (z) [∂ zi f (z)] = E p(z) 1 2p(z i ) ∂ zi f (z) (p i→1/2 (z i ) = 1/2) (12)

Interestingly, Gregor et al. (2013) arrive at equation 12 in the context of unbiased control variates for quadratic functions. , we can attempt to correct the coefficients 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 kth 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 first 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

E z,u [∂ z1 f (u 1 , u 2 )] = j c 1,j j + 1 E z [z j 2 ] + j c 3,j j + 1 E z [z j 2 ] + • • • (14)

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 coefficients 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.4. LOWERING BIAS BY REPRESENTATION SCALING:

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 coefficients in equation 3 will be different than for the {-1, +1} representation. We give the final forms of the gradients here. Details are given in Appendix C.2.

Under the p i→1/2 distribution the degree-1 Fourier coefficients are:

f (p i→1/2 ) (i) = E p i→1/2 (z) c 1 2 + c 3 2 3 + c 5 2 5 + • • • (15)

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.5. ALGORITHMIC COMPLEXITY:

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. (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 first-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.

5. RELATED WORK:

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 Gumbel-Softmax estimator (Maddison et al., 2016; Jang et al., 2016) . Although this is a continuous relaxation, it has been used to define 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 first-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. Specifically, 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 first to explore Fourier expansions for bias reduction of biased stochastic gradient estimators.

6. EXPERIMENTS:

Experimental Setup. We first 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, Gumbel-Softmax, 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. (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.1. BIASED ESTIMATORS ON A TOY PROBLEM:



6.2. FOUST VS. STATE-OF-THE-ART GRADIENT ESTIMATORS:

Training Stochastic MLPs. We train MLPs with one and two stochastic layers on OMNIGLOT, following the non-linear architecture of Tucker et al. (2017) . Each stochastic Boolean layer is preceded by two deterministic layers of 200 tanh units.

All hyperparameters remain fixed throughout the training. All estimators use one sample per example and a single decoder evaluation. We present results in Fig. 2 . FouST outperforms other biased gradient estimators in both datasets and architectures. FouST is clearly better than the Straight-Through 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 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. (2017) . 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.

Efficiency. We compare the efficiency 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 five stochastic layer MLP's respectively.

From the above experiments we conclude that FouST allows for efficient and effective training of fully connected and convolutional neural networks with Boolean stochastic variables.

6.3. ADDING STOCHASTIC DEPTH AND WIDTH:

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 first 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. 3 . 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. 

7. CONCLUSION:

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 efficiency that is orders of magnitude higher than unbiased estimators. Importantly, to the best of our knowledge FouST is the first 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.

APPENDICES A FURTHER DETAILS FROM THE HARMONIC ANALYSIS FRAMEWORK FOR BOOLEAN FUNCTIONS:

A.1 DISCRETE DERIVATIVE For a Boolean function f the discrete derivative on the i-th latent dimension with a basis function φi is defined as

EQUATION

)

The Fourier expansion of the discrete derivative equals

D φ i f (z) = S 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 definition 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 zi both in the above definitions.

B HARMONIC ANALYSIS OF EXISTING GRADIENT ESTIMATES B.1 PROOF OF LEMMA 1:

Proof. We follow O'Donnell (2014, §8.4) . In this proof we work with two representations of the Boolean function f . The first is the Fourier expansion of f under the uniform Bernoulli distribution. This is also the representation obtained by expressing f as a polynomial in zi. Since the domain of the function f is the Boolean cube, the polynomial representation is multilinear. That is f (z) = S⊆[n] f (S) j∈S zj. To avoid confusion and to differentiate the representation from the Boolean function, we use f (u) (z) to denote this representation in the following. Note that since this representation is a polynomial it is defined over any input in R n . In particular,

E[f (z)] = E[ S⊆[n] f (S) j∈S zj] = S⊆[n] f (S) j∈S E[zj] = f (u) (µ1, . . . , µn)

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) [zi] = µi and the linearity of expectation.

EQUATION

)

As the partial derivative of f (u) w.r.t. µi is equivalent to discrete derivative of f (u) w.r.t. zi, ∂µ i f (u) (µ) = Dz i f (u) (µ), and keeping in mind that φi = (zi -µi)/σi, we have that

EQUATION

We then note that the discrete derivative of f w.r.t. zi, Dz i f (u) (µ), from the left hand side of equation 19, is equivalent to the partial derivative of f w.r.t. µi, ∂µ i f (u) (µ).

EQUATION

We complete the proof by noting that the right hand side in equation 24 is 1 2 times the REINFORCE gradient.

B.2 ANALYSIS OF STRAIGHT-THROUGH:

The detailed proof of Lemma 2 is as follows.

Proof. We first 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 first 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 = zi, and (ii) no other term under the p(z) expansion contributes to the zi term under the p i→1/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 zi around 0 is

EQUATION

where

c k k! = ∂ k z i f (z)|z i =0

are the Taylor coefficients. All c k are a function of zj, j = i.

Let's first focus on the true gradient. Since we work with binary ±1 values, we have that z k i = 1 for even k and z k i = zi for odd k. This will influence the even and the odd terms of the Taylor expansions. Specifically, 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 pi is the expected sum of the odd Taylor coefficients. Here we note that the although final expression in equation 32 can also be derived by a finite difference method, it does not make explicit, as in equation 31, the dependence on zi 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 zi, 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 final term in equation 37 to vanish leaving

EQUATION

C LOW-BIAS GRADIENT ESTIMATES

C.1 LOWERING BIAS BY MATCHING UNIFORM DISTRIBUTION MOMENTS:

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 (z1, z2) with the bivariate Taylor expansion f (z1, z2) = i,j ci,jz i 1 z j 2 .

As in Lemma 2, the partial true gradient of f (z1, z2) w.r.t. the first distribution parameter p1 equals to

EQUATION

Further, the Taylor expansion of ∂p

1 f (z1, z2) is ∂z 1 f (z1, z2) = j c1,jz j 2 +2 j c2,jz j 2 x1 +3 j c3,jz j 2 z 2 1 + • • •

"Bernoulli splitting uniform" trick. Assume an auxiliary variable u = (u1, u2), which we choose as follows. First, we sample z = (z1, z2) from the uniform Bernoulli distribution p 1→1/2 . Then we take a uniform sample (u1, u2) with ui sampled from either [0, 1] for zi = +1 or from [-1, 0] if zi = -1. At this point it is important to note that the moments of the uniform distribution in We then compare equation 40 with equation 39. In equation 40 we observe that the pure terms in z1, namely terms with j = 0, always match those of the true gradient in equation 39. For j > 0 we obtain mixed terms with coefficients 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 coefficients 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.2 LOWERING BIAS BY REPRESENTATION SCALING:

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 zi ∈ {-1/2, 1/2} n , where pi = p(zi = +1/2). While we can change the input representation like that, in general the Fourier coefficients in equation 3 will be different than for the {-1, +1} representation. Letting h(zi) = 2zi ∈ {-1, 1}, the functions φi are now given as φi = h(z i )-µ i 

:

Training a nonlinear sigmoid belief network model with two stochastic layers on MNIST with REBAR took a day and a half on GPU.

Figures and Tables:

4.3 LOWERING BIAS BY MATCHING UNIFORM DISTRIBUTION MOMENTS Lemma 2 shows that part of the bias in the Straight-Through estimator is due to the presence of extra factors in the Taylor coefficients. We can reduce the effect of these factors by taking advantage of the moments of the uniform distribution. Recalling that b 0 1 b u k du = b k k+1

Figure 1: Biased estimators on a toy problem

Figure 2: Training ELBO for the one (left) and two (right) stochastic layer nonlinear models on OMNIGLOT

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 Ez,u[∂z 1 f (u1, u2)]

Figure 4: Training ELBO for the one (left) and two (right) stochastic layer nonlinear models on MNIST

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 variablesKingma & Welling (2013) propose the reparameterization trick, which transforms the random variable into a function of deterministic ones perturbed by a fixed 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

Table content: <table/>

Wall clock times for various gradient estimators on MNIST.

Table content: <table><tr><td colspan="4">Method #Eval. Walltime in sec./Epoch</td></tr><tr><td/><td/><td>2 layers</td><td>5 layers</td></tr><tr><td>REBAR</td><td>3</td><td>45.3</td><td>205.15</td></tr><tr><td>ST</td><td>1</td><td>2.86</td><td>4.96</td></tr><tr><td>Gumbel</td><td>1</td><td>3.27</td><td>6.14</td></tr><tr><td>DARN</td><td>1</td><td>2.96</td><td>5.36</td></tr><tr><td>FouST</td><td>1</td><td>3.1</td><td>5.67</td></tr></table>

Peer review (iclr2020):
***Score updated to weak accept after the rebuttal.*** Straight-Through is a popular, yet not theoretically well-understood, biased gradient estimator for Bernoulli random variables. The low variance of this estimator makes it a highly useful tool for training large-scale models with binary latents. However, the bias of this estimator may cause divergence in training, which is a significant practical issue. The paper develops a Fourier analysis of the Straight-Through estimator and provides an expression for the bias of the estimator in terms of the Fourier coefficients of the considered function. Motivated by this expression, the paper proposes two modifications of Straight-Through which may reduce the bias of the estimator, at the cost of the variance. The experimental results show advantage of this improved estimator over Gumbel-Softmax and DARN estimator. While I really like the premise of the paper, I feel that it needs a significant amount of additional work. The text is currently fairly hard to read. The theoretical part of the paper does not quantify the variance of the estimator. The experiments are a bit unfinished and do not include ablations of the proposed modifications of Straight-Through. Most importantly, I think that in the current form the theoretical and the empirical parts of the papers are not well-connected. Because of this, I believe that the paper should currently be rejected, but I encourage the authors to continue this line of work. Pros: 1. Theoretical analysis and empirical improvement of the Straight-Through estimator is an important avenue of work. 2. The paper makes a solid contribution of deriving the Fourier expansion of the Straight-Through estimator bias. 3. Based on this expansion, the paper proposes an algorithm with reduced bias. The algorithm is simple to implement, practical and appears to work slightly better than DARN. Cons: 1. The key weakness of the theoretical part of the paper is that it focuses on the bias of the estimator, but does not quantify the variance, especially after the modifications. If reducing the bias was the only goal, one could use unbiased (but high-variance) estimators such as REINFORCE or VIMCO. 2. The final algorithm appears to be the DARN estimator combined with relaxation by uniform noise (“Bernoulli splitting uniform”) and scaling. The paper does not have an ablation showing how the uniform noise and scaling perform on their own. 3. There are a few incorrect statements that I’ve noticed. * “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.” - MuProp paper compared to this estimator under the name 1/2-estimator * In Lemma 1 the “REINFORCE gradient” is just the exact gradient of the expectation, not a stochastic REINFORCE gradient. * “ To the best of our knowledge, FouST is the first gradient estimate algorithm that can train very deep stochastic neural networks with Boolean latent variables.” This paper uses up to 11 latent variable layers, while [1] has trained models with >20 latent variable layers (although their “layers” have just one unit). 4. The derivation of “Bernoulli splitting uniform” trick is confusing and contains a lot of typos. For instance, the text before eqn. (14) implies that the distribution of u_i is U[-1, 1], which cannot be right and does not correspond to Algorithm 1. The statement that this trick does not lead to a relaxation is odd, since the function is being evaluated at non-discrete points. 5. There are generally many typos and some poor formatting in the math. For example, in eqn. (6) the coefficients are off by one: it should be c0 + c1 z1 + c2 z2^2 + … . The equations (10) and (11) are poorly formatted. The notation partial_z1 f(u_1, u_2) in eqn. (14) is strange. In many places p^{i->½} is denoted as p^{1->½}. 5. I don’t think I understood the idea of representation scaling (Section 4.4). The eqn. (16) would suggest that the scaling should optimally be set to zero, which is just saying that the gradient is unbiased when the model does not use the latents. There is no other practical guidance on choosing this coefficient. Furthermore, one can always absorb the global scaling factor into the succeeding weights layer of the model, so this trick can probably be replaced by a modification of the weights initialization. 6. The experiments are missing a comparison to the Straight-Through Gumbel-Softmax estimator, introduced in the original Gumbel-Softmax paper. This is a popular biased estimator for Bernoulli latents, e.g. used in [1] [2]. Another interesting comparison would be [3] which proposes a lower-bias version of Gumbel-Softmax. 7. Figure 2 is missing the line for REBAR, even though this line is referred to on Page 8. Figure 2 and Figure 4 are both labeled as training ELBOs, despite the plots being different. [1] Andreas Veit, Serge Belongie “Convolutional Networks with Adaptive Inference Graphs” ECCV 2018 [2] Patrick Chen, Si Si, Sanjiv Kumar, Yang Li, Cho-Jui Hsieh “Learning to Screen for Fast Softmax Inference on Large Vocabulary Neural Networks” ICLR 2019 https://openreview.net/forum?id=ByeMB3Act7 [3] Evgeny Andriyash, Arash Vahdat, Bill Macready “Improved Gradient-Based Optimization Over Discrete Distributions” https://arxiv.org/abs/1810.00116