Low Bias Gradient Estimates for Very Deep Boolean Stochastic Networks

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 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 1Training a nonlinear sigmoid belief network model with two stochastic layers on MNIST with REBAR took a day and a half on GPU . Algorithm 1 FouST Gradient Estimator Require : Parameters θ ∈ RK , Bernoulli Representation { −a , a } , Interval Parameter b ∈ [ 0 , a ] , Constant scaling parameter γ 1 : Sample xi ∼ pθi ( xi ) , i = 1 , ... , K . This and all steps below have constant complexity . 2 : Sample yi ∼ Unif ( b , a ) 3 : Set yi : = xi ∗ yi 4 : Compute ST gradient ∂θi : = ∂yif ( y ) 5 : Importance reweighing ∂θi : = 0.5/p ( xi ) · ∂θ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 pi ( z ) . We denote pi as the probability of the i-th dimension being one , i.e. , pi : = pi ( zi = +1 ) . We denote the mean and variance of zi by µi and σi , respectively . An example of a Boolean function in this setting is a generative neural network f : z 7→ 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〉 = Ep ( 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 ( zi ) , 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 , pi ( zi = +1 ) = pi ( zi = −1 ) = 1/2 , we have φi ( zi ) = zi . The 2n functions φS form an orthonormal basis for the space of Boolean functions f , Ep ( z ) [ φi ( zi ) ] = 0 , ( 1 ) Ep ( z ) [ φi ( zi ) 2 ] = 1 , Ep ( z ) [ φi ( zi ) φj ( zj ) ] = Ep ( z ) [ φi ( zi ) ] Ep ( z ) [ φj ( zj ) ] , for i 6= j , where the expectations compute the inner product between two Boolean functions . The last identity derives from the independence of any dimensions i 6= j . We can then expand the Boolean function f on the set of 2n orthonormal basis functions , f ( z ) = ∑ S⊂ [ n ] f̂ ( p ) ( S ) φS ( z ) , ( 2 ) also known as the p-biased Fourier expansion of f . The f̂ ( p ) ( S ) are the Fourier coefficients computed by the inverse Fourier expansion , f̂ ( p ) ( S ) = Ep ( z ) [ f ( z ) φS ( z ) ] . ( 3 ) 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 : = Ep ( z ) [ f ( z ) ] we want to compute the gradient ∂pi L = ∂piEp ( 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 ≈ Ep ( z ) [ ∂zif ( 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 ) , Ep ( z ) [ gi REINFORCE ] = ∂piEp ( 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 ) ( gi ST ) denotes the bias of the i-th gradient estimate under the distribution p. Also , given the distribution p ( z ) , pi→1/2 ( z ) is the distribution for which we set p ( zi = +1 ) = p ( zi = −1 ) = 1/2 for a given dimension i. Lemma 2 . Let f be a Boolean function , gi = ∂piEp ( z ) [ f ( z ) ] its true gradients and giST = Ep ( z ) [ ∂zif ( z ) ] its Straight-Through gradient estimator . The bias in giST is due to a mismatch between the higher-order odd terms in the Taylor expansion of gi and gi ST and is equal to bias ( p ) ( gi ST ) = bias ( p i→1/2 ) ( gi ST ) + E ∑ k=2j , j > 0 kck µi ( 4 ) where ck are the Taylor coefficients for the i-th dimension on f ( z ) , that is zi , around 0 and bias ( p i→1/2 ) = E [ ∑ k=2j+1 , j > 0 ( k − 2 ) ck ] . 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 pi→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 pi→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 ∂piEp ( z ) [ f ( z ) ] = 2 f̂ ( p ) ( i ) σi = 2f̂ ( p→1/2 ) ( i ) ( 5 ) Taylor expansions of the true and the Straight-Through gradients . The Taylor expansion of f ( z ) for zi around 0 is given by f ( z ) = c0 + c1zi + c2z2i + c3z 3 i + c4z 4 i + c5z 5 i + · · · , where ckk ! = ∂ k zif ( z ) |zi=0 are the Taylor coefficients . All ck are a function of zj , j 6= i. Let ’ s first focus on the true gradient . Since we work with Boolean ±1 values , we have that zki = 1 for even k and zki = 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 can show that ∂pi [ Ep ( z ) [ f ( z ) ] ] = 2f̂ ( p→1/2 ) ( i ) = 2Epi→1/2 ( z ) [ c0zi + c1z2i + ciz3i + · · · ] ( 6 ) = 2Ep ( z\i ) [ c1 + c3 + c5 + · · · ] ( 7 ) The expression in equation 7 implies that the true gradient with respect to the pi 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 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 w.r.t . to zi , we have ∂zif ( z ) = c1 + 2c2zi + 3c3z 2 i + 4c4z 3 i + 5c5z 4 i + ... ( 8 ) The Straight-Through gradient is the expectation of equation 8 in the i-th dimension , that is Ep ( z ) [ ∂zif ( z ) ] = Ep ( z ) [ c1 + 2c2zi + 3c3z2i + 4c4z3i + 5c5z4i + · · · ] ( 9 ) = Ep ( z\i ) [ c1 + 3c3 + 5c5 + · · · ] + Ep ( z\i ) [ 2c2 + 4c4 + · · · ] µi ( 10 ) where z2ji = 1 , z 2j+1 i = zi , and Ep ( zi ) [ zi ] = µ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 ) ( gi ST ) = Ep ( z ) [ ∂zif ( z ) ] − ∂piEp ( z ) [ f ( z ) ] , bias ( p ) ( gi ST ) = Ep ( z\i ) ∑ k=2j+1 , j > 0 ( k − 2 ) ck + Ep ( z\i ) ∑ k=2j , j > 0 kck µi . ( 11 ) Taking the expectation in equation 9 under pi→1/2 causes the final term in equation 11 to vanish leaving bias ( p i→1/2 ) ( gi ST ) = Ep ( z\i ) [ ∑ k=2j+1 , j > 0 ( k − 2 ) ck ] . Combining this expression with equation 9 gives the final expression ( equation 4 ) from the lemma .

Summary: The authors analyze the bias in the straight-through gradient estimator using the framework of harmonic analysis of boolean functions. Based on this analysis, they propose three methods to reduce the bias of the straight-through estimator, resulting in a less-biased estimator that is the same computational complexity as the original. They evaluate this estimator on a series of generative modeling tasks where they demonstrate improvements over existing methods, including the ability to train a very deep stochastic network.