Critical initialisation in continuous approximations of binary neural networks

1 INTRODUCTION . The problem of learning with low-precision neural networks has seen renewed interest in recent years , in part due to the deployment of neural networks on low-power devices . Currently , deep neural networks are trained and deployed on GPUs , without the memory or power constraints of such devices . Binary neural networks are a promising solution to these problems . If one is interested in addressing memory usage , the precision of the weights of the network should be reduced , with the binary case being the most extreme . In order to address power consumption , networks with both binary weights and neurons can deliver significant gains in processing speed , even making it feasible to run the neural networks on CPUs Rastegari et al . ( 2016 ) . Of course , introducing discrete variables creates challenges for optimisation , since the networks are not continuous and differentiable . Recent work has opted to train binary neural networks directly via backpropagation on a differentiable surrogate network , thus leveraging automatic differentiation libraries and GPUs . A key to this approach is in defining an appropriate differentiable surrogate network as an approximation to the discrete model . A principled approach is to consider binary stochastic variables and use this stochasticity to “ smooth out ” the non-differentiable network . This includes the cases when ( i ) only weights , and ( ii ) both weights and neurons are stochastic and binary . In this work we study two classes of surrogates , both of which make use of the Gaussian central limit theorem ( CLT ) at the receptive fields of each neuron . In either case , the surrogates are written as differentiable functions of the continuous means of stochastic binary weights , but with more complicated expressions than for standard continuous networks . One approximation , based on analytic integration , yields a class of deterministic surrogates Soudry et al . ( 2014 ) . The other approximation is based on the local reparameterisation trick ( LRT ) Kingma & Welling ( 2013 ) , which yields a class of stochastic surrogates Shayer et al . ( 2017 ) . Previous works have relied on heuristics to deal with binary neurons Peters & Welling ( 2018 ) , or not backpropagated gradients correctly . Moreover , none of these works considered the question of initialisation , potentially limiting performance . The seminal papers of Saxe et al . ( 2013 ) , Poole et al . ( 2016 ) , Schoenholz et al . ( 2016 ) used a mean field formalism to explain the empirically well known impact of initialization on the dynamics of learning in standard networks . From one perspective the formalism studies how signals propagate forward and backward in wide , random neural networks , by measuring how the variance and correlation of input signals evolve from layer to layer , knowing the distributions of the weights and biases of the network . By studying these moments the authors in Schoenholz et al . ( 2016 ) were able to explain how heuristic initialization schemes avoid the “ vanishing and exploding gradients problem ” Glorot & Bengio ( 2010 ) , establishing that for neural networks of arbirary depth to be trainable they must be initialised at “ criticality ” , which corresponds to initial correlation being preserved to any depth . The paper makes three contributions . The first contribution is the presentation of new algorithms , with a new derivation able to encompass both surrogates , and all choices of stochastic binary weights , or neurons . The derivation is based on representing the stochastic neural network as a Markov chain , a simplifying and useful development . As an example , using this representation we are easily able to extend the LRT to the case of stochastic binary neurons , which is new . This was not possible in Shayer et al . ( 2017 ) , who only considered stochastic binary weights . As a second example , the deterministic surrogate of Soudry et al . ( 2014 ) is easily derived , without the need for Bayesian message passing arguments . Moreover , unlike Soudry et al . ( 2014 ) we correctly backpropagate through variance terms , as we discuss . The second contribution is the theoretical analysis of both classes of surrogate at initialisation , through the prism of signal propagation theory Poole et al . ( 2016 ) , Schoenholz et al . ( 2016 ) . This analysis is achieved through novel derivations of the dynamic mean field equations , which hinges on the use of self-averaging arguments Mezard et al . ( 1987 ) . The results of the theoretical study , which are supported by numerical simulations and experiment , establish that for a surrogate of arbitrary depth to be trainable , it must be randomly initialised at “ criticality ” . In practical terms , criticality corresponds to using initialisations that avoid the “ vanishing and exploding gradients problem ” Glorot & Bengio ( 2010 ) . We establish the following key results : • For networks with stochastic binary weights and neurons , the deterministic surrogate can achieve criticality , while the LRT can not . • For networks with stochastic binary weights and continuous neurons , the LRT surrogate can achieve criticality ( no deterministic surrogate exists for this case ) In both cases , the critical initialisation corresponds to randomly initialising the means of the binary weights close to ±1 , a counter intuitive result . A third contribution is the consideration of the signal propagation properties of random binary networks , in the context of training a differentiable surrogate network . We derive these results , which are partially known , and in order to inform our discussion of the experiments . This paper provides insights into the dynamics and training of the class of binary neural network models . To date , the initialisation of any binary neural network algorithm has not been studied , although the effect of quantization levels has been explored through this perspective Blumenfeld et al . ( 2019 ) . Currently , the most popular surrogates are based on the so-called “ Straight-Through ” estimator Bengio et al . ( 2013 ) , which relies on heuristic definitions of derivatives in order to define a gradient . However , this surrogate typically requires the use of batch normalization , and other heuristics . The contributions in this paper may help shed light on what is holding back the more principled algorithms , by suggesting practical advice on how to initialise , and what to expect during training . Paper outline : In section 2 we present the binary neural network algorithms considered . In subsection 2.1 we define binary neural networks and subsection 2.2 their stochastic counterparts . In subsection 2.3 we use these definitions to present new and existing surrogates in a coherent framework , using the Markov chain representation of a neural network to derive variants of both the deterministic surrogate , and the LRT-based surrogates . We derive the LRT for the case of stochastic binary weights , and both LRT and deterministic surrogates for the case of stochastic binary weights and neurons . In section 3 we derive the signal propagation equations for both the deterministic and stochastic LRT surrogates . This includes deriving the explicit depth scales for trainability , and solving the equations to find the critical initialisations for each surrogate , if they exist . In section 4 we present the numerical simulations of wide random networks , to validate the mean field description , and experimental results to test the trainability claims . In section 5 we summarize the key results , and provide a discussion of the insights they provide . 2 BINARY NEURAL NETWORK ALGORITHMS . 2.1 CONTINUOUS NEURAL NETWORKS AND BINARY NEURAL NETWORKS . A neural network model is typically defined as a deterministic non-linear function . We consider a fully connected feedforward model , which is composed ofN ` ×N ` −1 weight matricesW ` and bias vectors b ` in each layer ` ∈ { 1 , . . . , L } , with elements W ` ij ∈ R and b ` i ∈ R. Given an input vector x0 ∈ RN0 , the network is defined in terms of the following recursion , x ` = φ ` ( h ` ) , h ` = 1√ N ` −1 W ` x ` −1 + b ` ( 1 ) where the pointwise non-linearity is , for example , φ ` ( · ) = max ( 0 , · ) . We refer to the input to a neuron , such as h ` , as the pre-activation field . A deterministic binary neural network simply has weights W ` ij ∈ { ±1 } and φ ` ( · ) = sign ( · ) , and otherwise the same propagation equations . Of course , this is not differentiable , thus we instead consider stochastic binary variables in order to smooth out the non-differentiable network . Ideally , the product of training a surrogate of a stochastic binary network is a deterministic ( or stochastic ) binary network that is able to generalise from its training set . 2.2 STOCHASTIC BINARY NEURAL NETWORKS . In stochastic binary neural networks we denote the matrices as S ` with all weights1 S ` ij ∈ { ±1 } being independently sampled binary variables with probability is controlled by the mean M ` ij = ES ` ij . Neuron activation in this model are also binary random variables , due to pre-activation stochasticity and to inherent noise . We consider parameterised neurons such that the mean activation conditioned on the pre-activation is given by some function taking values in [ −1 , 1 ] , i.e . E [ x ` i |h ` i ] = φ ( h ` i ) , for example φ ( · ) = tanh ( · ) . We write the propagation rules for the stochastic network as follows : S ` ∼ p ( • ; M ` ) ; h ` = 1√ N ` −1 S ` x ` −1 + b ` ; x ` ∼ p ( • ; φ ( h ` ) ) ( 2 ) Notice that the distribution of x ` factorizes when conditioning on x ` −1 . The form of the neuron ’ s mean function φ ( · ) depends on the underlying noise model . We can express a binary random variable x ∈ { ±1 } with x ∼ p ( x ; θ ) via its latent variable formulation x = sign ( θ + αL ) . In this form θ is referred to as a “ natural ” parameter , and the term L is a latent random noise , whose cumulative distribution function σ ( · ) determines the form of the non-linearity since φ ( · ) = 2σ ( · ) −1 . In general the form of φ ( · ) will impact on the surrogates ’ performance , including within and beyond the mean field description presented here . However , a result from the analysis in Section 3 is that choosing a deterministic binary neuron , ie . the sign ( · ) function , or a stochastic binary neuron , produces the same signal propagation equations , up to a scaling constant . 1We denote random variables with bold font . Also , following physics ’ jargon , we refer to binary ±1 variables as Ising spins or just spins . 2.3 DERIVATIONS OF NEW AND EXISTING SURROGATE NETWORKS . The idea behind several recent papers Soudry et al . ( 2014 ) , Baldassi et al . ( 2018 ) , Shayer et al . ( 2017 ) , Peters & Welling ( 2018 ) is to adapt the mean of the binary stochastic weights , with the stochastic model essentially used to “ smooth out ” the discrete variables and arrive at a differentiable function , open to the application of continuous optimisation techniques . We now derive both the deterministic surrogate and LRT-based surrogates , in a common framework . We consider a supervised classification task , with training set D = { xµ , yµ } Pµ=1 , with yµ the label . we define a loss function for our surrogate model via L ( M , b ) = − 1 P P∑ µ=1 logES , x p ( yµ |xµ , S , x , b ) , ( 3 ) For a given input xµ and a realization of weights , neuron activations and biases in all layers , denoted by ( S , x , b ) , the stochastic neural network produces a probability distribution over the classes . Expectations over weights and activations are given by the mean values , ES ` = M ` and E [ x ` |h ` ] = φ ( h ` ) . This objective can be recognised as a ( minus ) marginal likelihood , thus this method could be described as Type II maximum likelihood , or empirical Bayes . The starting point for our derivations comes from rewriting the expectation equation 3 as the marginalization of a Markov chain , with layers ` indexes corresponding to time indices . Markov chain representation of stochastic neural network : ES , x p ( yµ |xµ , S , b , x ) = ∑ S , x : x0=xµ p ( yµ |xL ) L∏ ` =1 p ( x ` |x ` −1 , S ` ) p ( S ` ; M ` ) = ∑ SL , xL−1 p ( yµ ∣∣SL , xL−1 ) p ( SL ) ∑ SL−1 , xL−2 p ( xL−1|xL−2 , SL−1 ) p ( SL−1 ) · · · ∑ S1 p ( x1|xµ , S1 ) p ( S1 ) ( 4 ) where in the second line we dropped from the notation p ( S ` ; M ` ) the dependence onM ` for brevity . Therefore , for a stochastic network the forward pass consists in the propagation of the joint distribution of layer activations , p ( x ` |xµ ) , according to the Markov chain . We drop the explicit dependence on the initial input xµ from now on . In what follows we will denote with φ ( h ` ) the average value of x ` according to p ( x ` ) . The first step to obtaining a differentiable surrogate is to introduce continuous random variables . We take the limit of large layer width and appeal to the central limit theorem to model the field h ` as Gaussian , with mean h̄ ` and covariance matrix Σ ` . Assumption 1 : ( CLT for stochastic binary networks ) In the large N limit , under the Lyapunov central limit theorem , the field h ` = 1√ N ` −1 S ` x ` −1 + b ` converges to a Gaussian random variable with mean h̄ ` i = 1√ N ` −1 ∑ jM ` ijφ ( h ` −1 j ) + b ` i and covariance matrix Σ ` with diagonal Σ ` ii = 1 N ` −1 ∑ j 1− ( M ` ijφ ( h ` −1 j ) ) 2 . While this assumption holds true for large enough networks , due to S ` and x ` −1 independency , the Assumption 2 below , is stronger and tipically holds only at initialization . Assumption 2 : ( correlations are zero ) We assume the independence of the pre-activation field h ` between any two dimensions . Specifically , we assume the covariance Σ = Cov ( h ` , h ` ) to be well approximated by Σ ` MF ( φ ( h ` −1 ) ) , with MF denoting the mean field ( factorized ) assumption , where ( Σ ` MF ( x ) ) ii′ = δii′ 1 N ` −1 ∑ j 1− ( M ` ijφ ( h ` −1j ) ) 2 ( 5 ) This assumption approximately holds assuming the neurons in each layer are not strongly correlated . In the first layer this is certainly true , since the input neurons are not random variables2 . In subsequent layers , since the fields h ` i and h ` j share stochastic neurons from the previous layer , this can not be assumed to be true . We expect this correlation to not play a significant role , since 2In this case the variance is actually 1 N ` −1 ∑ j ( 1− ( M1ij ) 2 ) ( xµ , j ) 2. the weights act to decorrelate the fields , and the neurons are independently sampled . However , the choice of surrogate influences the level of dependence . The sampling procedure used within the local reparametrization trick reduces correlations since variables are sampled , while the deterministic surrogate entirely discards them . We obtain either surrogate model by successively approximating the marginal distributions , p ( x ` ) =∫ dh ` p ( x ` |h ` ) ≈ p̂ ( x ` ) , starting from the first layer . We can do this by either ( i ) marginalising over the Gaussian field using analytic integration , or ( ii ) sampling from the Gaussian . After this , we use the approximation p̂ ( x ` i ) to form the Gaussian approximation for the next layer , and so on . Deterministic surrogate : We perform the analytic integration based on the analytic form of p ( x ` +1i |h ` ) = σ ( x ` ih ` i ) , with σ ( · ) a sigmoidal function . In the case that σ ( · ) is the Gaussian CDF , we obtain p̂ ( x ` i ) exactly 3 by the Gaussian integral of the Gaussian cumulative distribution function , p̂ ( x ` i ) = ∫ dh σ ( x ` ih ) N ( h ; h̄ ` i , Σ ` MF , ii ) = Φ ( h̄ ` i ( 1 + Σ ` MF ) 1/2 ii x ` i ) ( 6 ) Since we start from the first layer , all random variables are marginalised out , and thus h̄ ` i has no dependence on random h ` −1j via the neuron means φ ( h ` ) as in Assumption 1 . Instead , we have dependence on means x̄ ` = Eh ` E [ x ` |h ` ] = Eh ` φ ( h ` ) . Thus it is convenient to define the mean under p̂ ( x ` i ) as ϕ ` ( h̄ , σ2 ) = ∫ dh φ ` ( h ) N ( h ; h̄ , σ2 ) . In the case that σ ( · ) is the Gaussian CDF , then ϕ ` ( · ) is the error function . Finally , the forward pass can be expressed as x̄ ` = ϕ ` ( h ` ) h ` = ( 1 + Σ ` MF ) − 12 h̄ ` h̄ ` = 1√ N ` −1 M ` x̄ ` −1 + b ` , ( 7 ) This is a more general formulation than that in Soudry et al . ( 2014 ) , which considered sign activations , which we obtain in the appendices as a special case . Furthermore , in all implementations we backpropagate through the variance terms Σ− 1 2 MF , which were ignored in the previous work of Soudry et al . ( 2014 ) . Note that the derivation here is simpler as well , not requiring complicated Bayesian message passing arguments , and approximations therein . LRT surrogate : The basic idea here is to rewrite the incoming Gaussian field h ∼ N ( µ , Σ ) as h = µ+ √ Σ where ∼ N ( 0 , I ) . Thus expectations over h can be written as expectations over and approximated by sampling . The resulting network is thus differentiable , albeit not deterministic . The forward propagation equations for this surrogate are h ` = 1√ N ` −1 M ` x̄ ` −1 + b ` + √ Σ ` MF ( x̄ ` −1 ) ` , x̄ ` = φ ` ( h ` ) . ( 8 ) The local reparameterisation trick ( LRT ) Kingma & Welling ( 2013 ) has been previously used to obtain differentiable surrogates for binary networks . The authors of Shayer et al . ( 2017 ) considered only the case of stochastic binary weights , since they did not write the network as a Markov chain . Peters & Welling ( 2018 ) considered stochastic binary weights and neurons , but relied on other approximations to deal with the neurons , having not used the Markov chain representation . The result of each approximation , applied successively from layer to layer by either propagating means and variances or by , produces a differentiable function of the parameters M ` ij . It is then possible to perform gradient descent with respect to the M and b . Ideally , at the end of training we obtain a binary network that attains good performance . This network could be a stochastic network , where we sample all weights and neurons , or a deterministic binary network . A deterministic network might be chosen taking the most likely weights , therefore setting W ` ij = sign ( M ` ij ) , and replacing the stochastic neurons with sign ( · ) activations .

Summary: The paper provides an in-depth exploration of stochastic binary networks, continuous surrogates, and their training dynamics with some potentially actionable insights on how to initialize weights for best performance. This topic is relevant and the paper would have more impact if its structure, presentation  and formalism could be improved. Overall it lacks clarity in the presentation of the results, the assumptions made are not always clearly stated and the split between previous work and original derivations should be improved.