Title: CRITICAL INITIALISATION IN CONTINUOUS APPROXI-MATIONS OF BINARY NEURAL NETWORKS

Abstract: The training of stochastic neural network models with binary (±1) weights and activations via continuous surrogate networks is investigated. We derive new surrogates using a novel derivation based on writing the stochastic neural network as a Markov chain. This derivation also encompasses existing variants of the surrogates presented in the literature. Following this, we theoretically study the surrogates at initialisation. We derive, using mean field theory, a set of scalar equations describing how input signals propagate through the randomly initialised networks. The equations reveal whether so-called critical initialisations exist for each surrogate network, where the network can be trained to arbitrary depth. Moreover, we predict theoretically and confirm numerically, that common weight initialisation schemes used in standard continuous networks, when applied to the mean values of the stochastic binary weights, yield poor training performance. This study shows that, contrary to common intuition, the means of the stochastic binary weights should be initialised close to ±1, for deeper networks to be trainable.

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

• 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 

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 of N × N -1 weight matrices W and bias vectors b in each layer ∈ {1, . . . , L}, with elements W ij ∈ R and b i ∈ R. Given an input vector x 0 ∈ R N0 , 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 weights 1 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:

EQUATION

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.

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

EQUATION

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:

E S,x p(y µ | x µ , S, b, x) = S,x : x 0 =xµ p(y µ | x L ) L =1 p(x | x -1 , S ) p(S ; M ) = S L ,x L-1 p(y µ S L , x L-1 )p(S L ) S L-1 ,x L-2 p(x L-1 |x L-2 , S L-1 )p(S L-1 ) • • • S 1 p(x 1 |x µ , S 1 )p(S 1 ) (4)

where in the second line we dropped from the notation p(S ; M ) the dependence on M 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 j M 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 Σ M F (φ(h -1 )), with MF denoting the mean field (factorized) assumption, where

Σ M F (x) ii = δ ii 1 N -1 j 1 -(M ij φ(h -1 j )) 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 variables 2 . In subsequent layers, since the fields h i and h j share stochastic neurons from the previous layer, this cannot be assumed to be true. We expect this correlation to not play a significant role, since 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 +1 i |h ) = σ(x i h i ), with σ(•) a sigmoidal function. In the case that σ(•) is the Gaussian CDF, we obtain p(x i ) exactly3 by the Gaussian integral of the Gaussian cumulative distribution function,

EQUATION

Since we start from the first layer, all random variables are marginalised out, and thus h i has no dependence on random h -1 j via the neuron means φ(h ) as in Assumption 1. Instead, we have dependence on means

x = E h E x | h = E h φ(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

EQUATION

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 M F , 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

EQUATION

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.

3. SIGNAL PROPAGATION THEORY FOR CONTINUOUS SURROGATES:

Since all the surrogates still retain the basic neural network structure of layerwise processing, crucially applying backpropagation for optimisation, it is reasonable to expect that surrogates are likely to inherit similar "training problems" as standard neural networks. In this section we apply this formalism to the surrogates considered, given random initialisation of the means M ij and biases b i . We are able to solve for the conditions of critical initialisation for each surrogate, which essentially allow signal to propagate forwards, and gradients to propagate backwards, without the effects such as neuron saturation. The critical initialisation for the surrogates, the key results of the paper, are provided in Claims 1 and 3.

3.1. FORWARD SIGNAL PROPAGATION FOR STANDARD CONTINUOUS NETWORKS:

We first recount the formalism developed in Poole et al. (2016) . Assume the weights of a standard continuous network are initialised with

W ij ∼ N (0, σ 2 w ), biases b ∼ N (0, σ 2 b )

, and input signal x 0 a has zero mean Ex 0 = 0 and variance E[x 0 a • x 0 a ] = q 0 aa , and with a denoting a particular input pattern. As before, the signal propagates via Equation 1 from layer to layer.

We are interested in computing, from layer to layer, the variance q aa = 1 N i (h i;a ) 2 from a particular input x 0 a , and also the covariance between the pre-activations q ab = 1 N i h i;a h i;b , arising from two different inputs x 0 a and x 0 b with given covariance q 0 ab . The mean field approximation used here replaces each element in the pre-activation field h i by a Gaussian random variable whose moments are matched. Assuming also independence within a layer; Eh i;a h j;a = q aa δ ij and Eh i;a h j;b = q ab δ ij , one can derive recurrence relations from layer to layer,

q aa = σ 2 w Dzφ 2 ( q -1 aa z) + σ 2 b = σ 2 w Eφ 2 (h -1 j,a ) + σ 2 b (9) with Dz = dz √ 2π e -z 2

2 the standard Gaussian measure. The recursion for the covariance is given by

EQUATION

)

where

u a = q -1 aa z 1 , u b = q -1 bb c -1 ab z 1 + 1 -(c -1 ab ) 2 z 2

, and we identify c ab as the correlation in layer . The other important quantity is the slope of the correlation recursion equation or mapping from layer to layer, denoted as χ, which is given by:

EQUATION

We denote χ at the fixed point c * = 1 as χ 1 . As discussed Poole et al. (2016) , when χ 1 = 1, correlations can propagate to arbitrary depth.

Definition 1: Critical initialisations are the points (σ 2 b , σ 2 w ) corresponding to χ 1 = 1. Furthermore, χ 1 is equivalent to the mean square singular value of the Jacobian matrix for a single layer J ij = ∂h i ∂h -1 j , as explained in Poole et al. (2016) . Therefore controlling χ 1 will prevent the gradients from either vanishing or growing exponentially with depth. We thus define critical initialisations as follows. This definition also holds for the surrogates which we now study.

3.2. SIGNAL PROPAGATION THEORY FOR DETERMINISTIC SURROGATES:

For the deterministic surrogate model we assume at initialization that the binary weight means M ij are drawn independently and identically from a distribution P (M ), with mean zero and variance of the means given by σ 2 m . For instance, a valid distribution could be a clipped Gaussian4 , or another stochastic binary variable, for example

P (M ) = 1 2 δ(M + σ m ) + 1 2 δ(M -σ m ), whose variance is σ 2 m .

The biases at initialization are distributed as b ∼ N (0, σ 2 b ). We show in Appendix B that the stochastic and deterministic binary neuron cases reduce to the same signal propagation equations, up to scaling constants. In light of this, we consider the deterministic sign(•) neuron case, since equation for the field is slightly simpler:

EQUATION

which we can be read from the Eq. 7. As in the continuous case we are interested in computing the variance q aa = 1 N i (h i;a ) 2 and covariance Eh i;a h j;b = q ab δ ij , via recursive formulae. The key to the derivation is recognising that the denominator Σ M F,ii is a self-averaging quantity Mezard et al. (1987) . This means it concentrates in probability to its expected value for large N . Therefore we can safely replace it with its expectation. Following this self-averaging argument, we can take expectations more readily as shown in the appendices. We find the variance recursion to be

EQUATION

Based on this expression, and assuming q aa = q bb , the correlation recursion can be written as

EQUATION

)

The slope of the correlation mapping from layer to layer, when the normalized length of each input is at its fixed point q aa = q bb = q * (σ m , σ b ), denoted as χ, is given by:

EQUATION

where u a and u b are defined exactly as in the continuous case. Refer to the appendices for full details of the derivation.

3.2.1. CRITICAL INITIALISATION: DETERMINISTIC SURROGATE:

The condition for critical initialisation is χ 1 = 1, since this determines the stability of the correlation map fixed point c * = 1. Note that for the deterministic surrogate this is always a fixed point. We can solve for the hyper-parameters (σ 2 b , σ 2 m ) that satisfy this condition, using the dynamical equations of the network.

Claim 1: The points (σ 2 b , σ 2 m ) corresponding to critical initialisation are given by σ

2 m = 1/E[ ϕ ( √ q * z) 2 ] + E[ϕ 2 ( √ q * z)

] and finding σ 2 b that satisfies

q aa = σ 2 b + (σ 2 b + 1) Eϕ 2 (h l-1 j,a ) E[ ϕ ( √ q * z) 2 ]

This can be established by rearranging Equations 13 and 15. We solve for σ 2 b numerically, as shown in Figure 3 , for different neuron noise models and hence non-linearities ϕ(•). We find that the critical initialisation for any of these design choices is close to the point (σ 2 m , σ 2 b ) = (1, 0). However, it is not just the singleton point, as for example in Hayou et al. (2019) for the ReLu case for standard networks. We plot the solutions in the Appendix.

3.2.2. ASYMPTOTIC EXPANSIONS AND DEPTH SCALES:

The depth scales, as derived in Schoenholz et al. ( 2016) provide a quantitative indicator to the number of layers correlations will survive for, and thus how trainable a network is. Similar depth scales can be derived for these deterministic surrogates. Asymptotically in network depth , we expect that |q aa -q * | ∼ exp(ξq ) and |c ab -c * | ∼ exp(ξc ), where the terms ξ q and ξ c define the depth scales over which the variance and correlations of signals may propagate. We are most interested in the correlation depth scale, since it relates to χ. The derivation is identical to that of Schoenholz et al. (2016) . One can expand the correlation c ab = c * + , and assuming q aa = q * , it is possible to write

+1 = 1 + q * 1 + σ 2 b σ 2 m Dzϕ (u 1 )ϕ (u 2 ) + O(( ) 2 ) (16)

Published as a conference paper at ICLR 2020

The depth scale ξ -1 c are given by the log ratio log

+1 . ξ -1 c = -log 1 + q * 1 + σ 2 b σ 2 m Dzϕ (u 1 )ϕ (u 2 ) = -log χ (17)

We plot this depth scale in Figure 2 . We derive the variance depth scale in the appendices, since it is different to the standard continuous case, but not of prime practical importance.

3.3. SIGNAL PROPAGATION THEORY FOR LOCAL REPARAMETERIZATION TRICK SURROGATES:

From Equation 8, the pre-activation field for the perturbed surrogate with both stochastic binary weights and neurons is given by,

EQUATION

where we recall that ∼ N (0, 1). The non-linearity φ(•) can of course be derived from any valid binary stochastic neuron model. Appealing to the same self-averaging arguments used in the previous section, we find the variance map to be

q aa = E (h l i,a ) 2 = σ 2 m Eφ 2 (h l-1 j,a ) + σ 2 b + (1 -σ 2 m Eφ 2 (h l-1 j,a )) = 1 + σ 2 b (19)

Interestingly, we see that the variance map does not depend on the variance of the means of the binary weights. This is not immediately obvious from the pre-activation field definition. In the covariance map we do not have such a simplification since the perturbation i,a is uncorrelated between inputs a and b. Thus the correlation map is given by

c l ab = σ 2 m Eφ(h l-1 j,a )φ(h l-1 j,a ) + σ 2 b 1 + σ 2 b (20)

3.4. CRITICAL INITIALISATION: LRT SURROGATES:

Claim 2: There is no critical initialisation for the local reparameterisation trick based surrogate, for a network with binary weights and neurons.

Proof: The conditions for a critical initialisation are that c * = 1 to be a fixed point and χ 1 = 1.

No such fixed point exists. We have a fixed point c * = 1 if and only if

σ 2 m = 1/E[φ 2 (h l-1 j,a )]. Note that σ 2 m ≤ 1.

For any φ(z) which is the mean of the stochastic binary neuron, the expectation E[φ 2 (z)] ≤ 1. For example, consider φ(z) = tanh(κz) for any finite kappa.

We also considered the LRT surrogate with continuous (tanh(•)) neurons and stochastic binary weights. The derivations are very similar to the previous case, as we show in the appendix. The variance and correlation maps are given by

EQUATION

)

This leads to the following result, Claim 3: The critical initialisation for the LRT surrogate, for the case of continuous tanh(•) neurons and stochastic binary weights is the singleton (σ 2 b , σ 2 m ) = (0, 1). Proof: From the correlation map we have a fixed point c * = 1 if and only if σ 2 m = 1, by inspection. In turn, the critical initialisation condition

χ 1 = 1 holds if E[(φ (h l-1 j,a )) 2 ] = 1 σ 2 m = 1.

Thus, to find the critical initialisation, we need to find a value of q aa = Eφ 2 (h l-1 j,a ) + σ 2 b that satisfies this final condition. In the case that φ(•) = tanh(•), then the function (φ (h l-1 j,a )) 2 ≤ 1, taking the value 1 at the origin only, this requires q aa → 0. Thus we have the singleton (σ 2 b , σ 2 m ) = (0, 1) as the solution. 

4. NUMERICAL AND EXPERIMENTAL RESULTS:



4.1. SIMULATIONS:

We first verify that the theory accurately predicts the average behaviour of randomly initialised networks. We present simulations for the deterministic surrogate in Figure 1 . We see that the average behaviour of random networks are well predicted by the mean field theory. Estimates of the variance and correlation are plotted, with dotted lines corresponding to empirical means and the shaded area corresponding to one standard deviation. Theoretical predictions are given by solid lines, with strong agreement for even finite networks. Similar plots can be produced for the LRT surrogate. In Appendix D we plot the depth scales as functions of σ m and σ b .

4.2. TRAINING PERFORMANCE FOR DIFFERENT MEAN INITIALISATION σ 2 m:

Here we experimentally test the predictions of the mean field theory by training networks to overfit a dataset in the supervised learning setting, having arbitrary depth and different initialisations. We consider first the performance of the deterministic and LRT surrogates, not their corresponding binary networks.

We use the MNIST dataset with reduced training set size (50%) and record the training performance (percentage of the training set correctly labeled) after 10 epochs of gradient descent over the training set, for various network depths L < 70 and different mean variances σ 2 m ∈ [0, 1). The optimizer used was SGD with Adam Kingma & Ba (2014) with a learning rate of 2 × 10 -4 chosen after simple grid search, and a batch size of 64. We see that the experimental results match the correlation depth scale derived, which are overlaid as dotted curves. A proportion of 3ξ c was found to indicate the maximum attenuation in signal strength before trainability becomes difficult, similarly to previous works Schoenholz et al. (2016) .

A reason we see the trainability not diverging in Figure 2 is that training time increases with depth, on top of requiring smaller learning rates for deeper networks, as described in Saxe et al. (2013) . The experiment here used the same number of epochs regardless of depth, meaning shallower networks actually had an advantage over deeper networks. Note that the theory does not specify for how many steps of training the effects of critical initialisation will persist. Therefore, the number of steps we trained the network for is an arbitrary choice, and thus the experiments validate the theory in a more qualitative way. Results were similar for other optimizers, including SGD, SGD with momentum, and RMSprop. Note that these networks were trained without dropout, batchnorm or any other heuristics.

In Figure 2 we present the training performance for the deterministic surrogate and its stochastic binary counterpart. The results for a deterministic binary network were similar to a single Monte Carlo sample. Once again, we test our algorithms on the MNIST dataset and plot results after 5 epochs. We see that the performance of the stochastic network matches more closely the performance of the continuous surrogate as the number of samples increases, from N = 5 to N = 100 samples. We can report that the number of samples necessary to achieve better classification, at least for more shallow networks, appears to depends on the number of training epochs. This is a sensible relationship, since during the course of training we expect the means of the weights to polarise, moving closer to the bounds ±1. Likewise, we expect that neurons, which initially have zero mean pre-activations, will also "saturate" during training, becoming either always "on" (+1) or "off" (-1). A stochastic network being "closer" to deterministic would require fewer samples overall. 

5. DISCUSSION:

This study of two classes of surrogate networks, and the derivation of their initialisation theories has yielded results of practical significance. Based on the results of Section 3, in particular Claims 1-3, we can offer the following advice. If a practitioner is interested in training networks with binary weights and neurons, one should use the deterministic surrogate, not the LRT surrogate, since the latter has no critical initialisation. If a practitioner is interested in binary weights only,the LRT in this case does have a critical initialisation (and is the only choice from amongst these two classes of surrogate). Furthermore, both networks are critically initialised when σ 2 b → 0 and by setting the means of the weights to ±1.

It was seen that during training, when evaluating the stochastic binary counterparts concurrently with the surrogate, the performance of binary networks was worse than the continuous model, especially as depth increases. We reported that the stochastic binary network, with more samples, outperformed the deterministic binary network, a reasonable result since the objective optimised is the expectation over an ensemble of stochastic binary networks.

A study of random deterministic binary networks, included in the Appendices, and published recently Blumenfeld et al. (2019) for a different problem, reveals unsurprisingly that binary networks are always in a chaotic phase. However a binary network which is trained via some algorithm will of course have different signal propagation behaviour. It makes sense that the closer one is to the early stages of the training process, the closer the signal propagation behaviour is to the randomly initialised case. We might expect that as training progresses the behaviour of the binary counterparts approaches that of the trained surrogate. Any such difference would not be observed for a heuristic surrogate as used in Courbariaux & Bengio (2016) or Rastegari et al. (2016) , which has no continuous forward propagation equations.

A DERIVATION OF DETERMINISTIC SURROGATE NETWORKS:



A.1 INTEGRATING OVER STOCHASTIC OR DETERMINISTIC BINARY NEURONS:

The form of each neuron's probability distribution depends on the underlying noise model. We can express a stochastic binary random variable S ∈ {±1} with S ∼ p(S; θ) via its latent variable formulation,

S = sign(θ + αL) (22)

In this form θ is referred to as a "natural" parameter, from the statistics literature on exponential families. The term L is a latent random noise, which determines the form of the probability distribution. We also introduce a scaling α to control the variance of the noise, so that as α → 0 the neuron becomes a deterministic sign function. Letting α = 1 for simplicity, we see that the probability of the binary variable taking a positive value is

EQUATION

)

where p(L) is the known probability density function for the noise L. The two common choices of noise models are Gaussian or logistic noise. The Gaussian of course has shifted and scaled erf(•) function as its cumulative distribution. The logistic random variable has the classic "sigmoid" or logistic function as its CDF, σ(z) = 1 1+e -z . Thus, the probability of a the variable being positive is a function of the CDF. In the Gaussian case, this is Φ(θ). By symmetry, the probability of p(S = -1) = Φ(-θ). Thus, we see the probability distribution for the binary random variable in general is the CDF of the noise L, and we write p(S) = Φ(Sθ). In the logistic noise case, we have p(S) = σ(Sθ)

For the stochastic neurons, the natural parameter is the incoming field h i = j S i,j x -1 j + b i . Assuming this is approximately Gaussian in the large layer width limit, we can successively marginalise over the stochastic inputs to each neuron, calculating an approximation of each neuron's probability distribution, p(x i ). This approximation is then used in the central limit theorem for the next layer, and so on.

For the case of neurons with latent Gaussian noise as part of the binary random variable model, the integration over the pre-activation field (assumed to be Gaussian) is exact. Explicitly,

EQUATION

)

where Φ(•) is the CDF of the Gaussian distribution. We have again Σ M F denoting the mean field approximation to the covariance between the stochastic binary pre-activations. The Gaussian expectation of the Gaussian CDF is a known identity, which we state in more generality in the next section, where we also consider neurons with logistic noise.

This new approximate probability distribution p(x i ) can then used as part of the Gaussian CLT applied at the next layer, since it determines the means of the neurons in the next layer,

EQUATION

If we follow these setps from layer to layer, we see that we are actually propagating approximate means for the neurons, combined non-linearly with the means of the weights. Given the approximately analytically integrated loss function, it is possible to perform gradient descent with respect to the means and biases, M ij and b i .

In the case of deterministic sign() neurons we obtain particularly simple expressions. In this case the "probability" of a neuron taking, for instance, positive is just Heaviside step function of the incoming field. Denoting the Heaviside with Θ(•), we have

EQUATION

We can write out the network forward equations for the case of deterministic binary neurons, since it is a particularly elegant result. In general we have

EQUATION

)

where φ(•) = erf(•) is the mean of the next layer of neurons, being a scaled and shifted version of the neuron's noise model CDF. The constant is η = 1 √ 2 , standard for the Gaussian CDF to error functin conversion.

A.2 EXACT AND APPROXIMATE GAUSSIAN INTEGRATION OF SIGMOIDAL FUNCTIONS:

We now present the integration of stochastic neurons with logistic as well as Gaussian noise as part of their latent variable models. The logistic case is an approximation built on the Gaussian case, motivated by approximating the logistic CDF with the Gaussian CDF. The reason we may be interested in using logistic CDFs, rather than just considering latent Gaussian noise models which integrate exactly, is not justified in any rigorous or experimental way. Any such analysis would likely consider the effect of the tails of the logistic versus the Gaussian distributions, where the logistic tails are much heavier than those of the Gaussian. One historic reason for considering the logistic function, we note, is the prevalence of logistic-type functions (such as tanh(•)) in the neural network literature. The computational cost of evaluating either logistic or error functions is similar, so there is no motivation from the efficiency side. Instead it seems a historic preference to have logistic type functions used with neural networks.

As we saw in the previous subsection, the integration over the analytic probability distribution for each neuron gave a function which allows us to calculate the means of the neurons in the next layer. Therefore, we directly calculate the expression for the means.

The Gaussian integral of the Gaussian CDF was used in the previous section to derive the exact probability distribution for the stochastic binary neuron in the next layer. The result is well known, and can be stated in generality as follows,

EQUATION

We can integrate a logistic noise binary neuron using this result as well. The idea is to approximate the logistic noise with a suitably scaled Gaussian noise. However, since the overall network approximation results in propagating means from layer to layer, we can equivalently need to approximate the tanh(•) with the with the erf. Specifically, if we have f (x; α) = tanh( x α ), an approximation is g(x; α) = erf( √ π 2α x), by requiring equality of derivatives at the origin. In order to establish this, consider

EQUATION

and

EQUATION

Published as a conference paper at ICLR 2020 Equating these, gives σ 2 = 4α 2 π , thus σ = 2α √ π . The approximate integral over the stochastic binary neuron mean is then

EQUATION

EQUATION

If we so desire, we can approximate this again with a tanh(•) using the tanh(•) to erf(•) approximation in reverse. The scale parameter of this tanh(•) will be

EQUATION

B EQUIVALENCE OF DETERMINISTIC AND STOCHASTIC NEURONS FOR DETERMINISTIC SURROGATE:

Assume a stochastic neuron with some latent noise, as per the previous appendix, with mean

x i = E p(xi) x i = φ(h -1 i ).

The field is given by

EQUATION

We see that the expression for the variance of the field simplifies as follows,

EQUATION

By similar steps, we find that in the deterministic binary neuron case, we would obtain the same expression, albeit with a different scaling constant. This is easily seen by inspection of the field term in the deterministic neuron case,

EQUATION

which again was derived in the previous appendix.

C DERIVATION OF SIGNAL PROPAGATION EQUATIONS IN DETERMINISTIC SURROGATE NETWORKS:

Here we present the derivations for the signal propagation in the continuous network models studied in the paper.

C.1 VARIANCE PROPAGATION:

We first calculate the variance given a signal:

EQUATION

Where for us:

EQUATION

and

EQUATION

Where, Eφ 2 h l-1 j,a can be written explicitly, taking into account that h l-1 j,a ∼ N (0, q aa ):

EQUATION

We can now perform the following change of variable:

EQUATION

Then:

EQUATION

In the first layer, input neurons are not stochastic: they are samples drawn from the Gaussian distribution x 0 ∼ N 0, q 0 : C.1.1 CORRELATION PROPAGATION To determine the correlation recursion we start from its definition:

EQUATION

where q l ab represents the covariance of the pre-activations h l i,a and h l i,b , related to two distinct input signals and therefore defined as:

EQUATION

Replacing the pre-activations with their expressions provided in eq. ( 41) and taking advantage of the self-averaging argument, we can then write:

EQUATION

At this point, given that q l aa and q l bb quite quickly approach the fixed point, we can conveniently assume q l aa = q l bb . Moreover, exploiting eq.( 47), we can finally write the expression for the correlation recursion:

EQUATION

C.2 DERIVATION OF THE SLOPE OF THE CORRELATIONS AT THE FIXED POINT:

To check the stability at the fixed point, we need to compute the slope of the correlations mapping from layer to layer at the fixed point:

EQUATION

)

where we get rid of σ b because independent from c l-1 ab . Replacing the definition of u a and u b provided in the continuous model, we can explicitly compute the derivative with respect to c l-1 ab :

EQUATION

where we have defined A and B as:

EQUATION

We can focus on B first. Integrating by parts over z b we get:

EQUATION

Then, integrating by parts over z a , we the get:

EQUATION

)

Replacing A and B in eq. ( 53), we then obtain the closest expression for the stability at the variance fixed point, namely:

EQUATION

C.3 VARIANCE DEPTH SCALE:

As pointed out in the main text, it should hold asymptotically that:

EQUATION

with ξ q defining the variance depth scale. To compute it we can expand over small perturbations around the fixed point, namely:

EQUATION

)

Expanding the square root for small l , we can then write:

EQUATION

We can now expand the activation function φ around small perturbations and then computing the square getting rid of higher order terms in l , thus finally obtaining:

q l+1 aa q * + 1 + q * √ q * σ 2 m Dzφ √ q * z φ √ q * z z 1 -σ 2 m Dzφ 2 √ q * z l (61)

Comparing this expression with the one in eq. ( 59), we can then write:

EQUATION

)

Integrating by parts over z, we then obtain:

EQUATION

)

Given that it holds eq. ( 47), and noticing that χ evaluated at the correlation fixed point c * = 1 is given by:

EQUATION

we can finally get:

EQUATION

)

Given that we expect (58) to hold asymptotically, that is:

EQUATION

we can finally obtain the variance depth scale: We see in Figure 3 that the set of critical initialisations exist in the plane, for but σ 2 b > 10 -20 all the corresponding mean variances σ + m 2 > 1 which is not possible.

EQUATION

D.2 DEPTH SCALES:

We see in Figure 4 the depth scales for the deterministic surrogate. Note the divergence as one expects following the simulations in Figure 3 . m → 1, whereas for standard continuous networks, there are an infinite number of such points, corresponding to various combinations of the weight and bias variances.

D.3 JACOBIAN MEAN SQUARED SINGULAR VALUE AND MEAN FIELD GRADIENT BACKPROPAGATION:

An alternative perspective on critical initialisation, to be contrasted with the forward signal propagation theory, is that we are simply attempting to control the mean squared singular value of the input-output Jacobain matrix of the entire network, which we can decompose into the product of single layer Jacobian matrices. In standard networks, the single layer Jacobian mean squared singular value is equal to the derivative of the correlation mapping χ as established in Poole et al. (2016) .

For the Gaussian model studied here this is not true, and corrections must be made to calculate the true mean squared singular value. This can be seen by observing the terms arising from denominator of the pre-activation field,

EQUATION

Since Σ ii is a quantity that scales with the layer width N , it is clear that when we consider squared quantities, such as the mean squared singular value, the second term, from the derivative of the denominator, will vanish in the large layer width limit. Thus the mean squared singular value of the single layer Jacobian approaches χ. We will proceed as if χ is the exact quantity we are interested in controlling.The analysis involved in determining whether the mean squared singular value is well approximated by χ essentially takes us through the mean field gradient backpropagation theory as described in Schoenholz et al. (2016) . This idea provides complementary depth scales for gradient signals travelling backwards.

E REPARAMETERISATION TRICK SURROGATE E.1 SIGNAL PROPAGATION EQUATIONS:

We present, in slightly more detail, the signal propagation equations for the case of continuous neurons and stochastic binary weights yields the variance map,

q aa = Eφ 2 (h l-1 j,a ) + σ 2 b (69)

Thus, once again, the variance map does not depend on the variance of the means of the binary weights. The covariance map however does retain a dependence on σ 2 m ,

EQUATION

)

with the same expression as before. The correlation map is given by

c l ab = σ 2 m Eφ(h l-1 j,a )φ(h l-1 j,a ) + σ 2 b Eφ 2 (h l-1 j,a ) + σ 2 b (71)

and we have the derivative of the correlation map given by

EQUATION

E.2 DETERMINING THE CRITICAL INITIALISATION CONDITIONS:

We recount the argument from the paper here. Since the mean variance σ 2 m does not appear in the variance map, we must once again consider different conditions for critical initialisation. Specifically, from the correlation map we have a fixed point c * = 1 if and only if

σ 2 m = 1 (73)

In turn, the condition

EQUATION

Thus, to find the critical initialisation, we need to find a value of q aa = Eφ 2 (h l-1 j,a ) + σ 2 b that satisfies this final condition. In the case that φ(•) = tanh(•), then the function (φ (h l-1 j,a )) 2 ≤ 1, taking the value 1 at the origin only, this requires q aa → 0. Thus the critical initialisation is the singleton point (σ 2 b , σ 2 m ) = (0, 1). This is confirmed by experiment, as we reported in the paper. It is of course possible to investigate this perturbed surrogate for different noise models. For example, given different noise scaling κ, as in the previous chapter, there will be a corresponding σ 2 b that satisfy the critical initialisation condition. We leave such an investigation to future work, given the case of binary weights and continuous neurons does not appear to be of a particular interest over the binary neuron case.

F SIGNAL PROPAGATION OF BINARY NETWORKS F.1 FORWARD SIGNAL PROPAGATION:

In this neural network, it should be understood that all neurons are simply sign(•) functions of their input, and all weights W ij ∈ {±1} are randomly distributed according to

EQUATION

(76) thus maintaining a zero mean.

The pre-activation field is given by

EQUATION

So, the length map is:

q aa = Dz(sign( q -1 aa z) 2 ) + σ 2 b (78)

= 1 + σ 2 b (79)

Interestingly, this is the same value as for the perturbed Gaussian with stochastic binary weights and neurons.

The covariance evolves as where as in the paper, u a = q -1 aa z 1 , u b = q -1 bb c -1 ab z 1 + 1 -(c -1 ab ) 2 z 2 . We can find this correlation in closed form. First we rewrite our integral with h, for a joint density p(h a , h b ), and then rescale the h a such that the variance is 1, so that dh a = √ q aa dv a 

q ab =

EQUATION

where p(v a , v b ) is a joint with the same correlation c ab (which is now equal to its covariance), and the capital P (v 1 , v 2 ) corresponds to the (cumulative) distribution function. A standard result for standard bivariate normal distributions with correlation ρ,

EQUATION

So we then have that dh a dh b φ(h a )φ(h b )p(h a , h b ) = √ q aa q bb 1 2 + sin -1 (c -1 ab ) π -cos -1 (c -1 ab ) π (85)

Thus the correlation map is:

c ab = 1 2 + sin -1 (c -1 ab ) π - cos -1 (c -1 ab ) π + σ 2 b q -1 aa q -1 bb (86) = 2 π sin -1 (c -1 ab ) + σ 2 b q -1 aa q -1 bb (87)

Since, from before we have q aa = 1 + σ 2 b , we then obtain

EQUATION

)

Recall that sin -1 (1) = π 2 , so we have that c * = 1 is a fixed point always.

We will now derive its slope, denoted as χ = ∂c ab ∂c -1 ab , but by first integrating over the φ() = sign() non-linearities, and then taking the derivative. Now we are in a place to take the derivative :

EQUATION

We can see that the derivative χ diverges at c ab = 1, meaning that there is no critical initialisation for this system. This of course means that correlations will not propagate to arbitrary depth in deterministic binary networks, as one might have expected.

F.2 STOCHASTIC WEIGHTS AND NEURONS:

We begin again with the variance map, q l aa = E (h l i,a ) 2 (90) where in this the field is given by

h l i,a = 1 √ N j W l ij x h l-1 j,a + b l i (91)

where x h l-1 j,a denotes a stochastic binary neuron whose natural parameter is the pre-activation from the previous layer.

The expectation for the length map is defined in terms of nested conditional expectations, since we wish to average over all random elements in the forward pass,

q aa = E h E x|h x h l-1 j,a + σ 2 b (92) = 1 + σ 2 b

(93) Once again, this is the same value as for the perturbed Gaussian with stochastic binary weights and neurons.

Similarly, the covariance map gives us, with phi(•) being the mean function, or a shifted and scaled version of the cumulative distribution function for the stochastic binary neurons, just as in previous Chapters. This expression is equivalent to the perturbed surrogate for stochastic binary weights and neurons, with a mean variance of σ 2 m = 1. Following the arguments for that surrogate, no critical initialisation exists.

EQUATION

:

We denote random variables with bold font. Also, following physics' jargon, we refer to binary ±1 variables as Ising spins or just spins.

In this case the variance is actually 1N -1 j 1 -(M 1 ij ) 2 (xµ,j) 2 .

In the Appendices we show that other sigmoidal σ(•) can be approximated by a Gaussian CDF.

That is, sample from a Gaussian then pass the sample through a function bounded on the interval [-1, 1].

Figures and Tables:

Figure 1: Dynamics of the variance and correlation maps, with simulations of a network of width N = 1000, 50 realisations, for various hyperparameter settings: σ 2 m ∈ {0.2, 0.5, 0.99} (blue, green and red respectively). (a) variance evolution, (b) correlation evolution. (c) correlation mapping (c in to c out ), with σ 2 b = 0.001

Figure 2: Top: Training performance of the deterministic surrogate (left) and the LRT surrogate for stochastic binary weights and continuous neurons (right). The vertical axis represents network depth against the variance of the means σ 2 m . Both surrogates were trained with σ 2 b = 0. Thus, as σ 2 m → 1 we approach criticality in both cases. Overlaid are curves proportional to the correlation depth scale ξ c . Bottom: Training performance of the deterministic surrogate and its binary counterparts after training on the MNIST dataset for 5 epochs. Left: performance of the continuous surrogate. Centre: the performance of the stochastic binary network, averaged over 5 Monte Carlo samples. Right: 100 Monte Carlo samples. The deterministic binary evaluation is similar to a single Monte Carlo sample, resembling the central figure.

Figure 3: Plots of the valid critical initialisations for the deterministic surrogate model, for stochastic binary weights and stochastic or deterministic binary neurons. Presented are the critical initialisations in the (σ 2 m , σ 2 b ), for both the a) stochastic neuron case with φ(z) = erf( 1 4 z), b) the deterministic sign neuron case with φ(z) = erf( 1 2 •), and (c) the logistic based stochastic neuron, with tanh() approximation. We see all lines are above σ 2 = 1 for all but small σ 2 b << 1.

Figure 4: Depth scales as σ 2 m is varied. (a) The depth scale controlling the variance propagation of a signal (b) The depth scale controlling correlation propagation of two signals. Notice that the correlation depth scale ξ c only diverges as σ 2m → 1, whereas for standard continuous networks, there are an infinite number of such points, corresponding to various combinations of the weight and bias variances.

Dz 1 Dz 2 sign(u a ) sign(u b ) + σ 2 b (80) we again have a correlation map: c ab = Dz 1 Dz 2 sign(u a ) sign(u b ) + σ 2 b

dh a dh b sign(h a ) sign(h b )p(h a , h b ) = dv a dv b sign(v a ) sign(v b )p(v a , v b )(82)

E ha,h b E x b |ha E x b |h b x h l-1 j,a x h l-1 j,b + σ 2 b = Eφ(h l-1 j,a )φ(h l-1 j,a ) + σ 2 b (95)

Peer review (iclr2020):
The paper addresses a very important and relevant topic of initialisation of weights of neural networks. It builds up on highly celebrated results of Poole, Schoenholz and others, using the language of mean field theory and an approach rooted in Dynamical Systems. What the authors propose is an extension of the approach to other settings. The paper is very scientific and math-heavy. A good practice in such cases is to adhere to a format of a scientific mathematical paper and organise the material using Theorema, Lemmata, Propositions and Corollaries , Definitions and Proofs. Such language and framework exists for a reason - to structure the material and make the paper readable. The paper as is a stream of equations and discussion making it very unclear what the point is. In order for this paper to be suitable for publication the reviewer would like to strongly suggest: - Organise the material in a way that would make it clear what is claimed, what is proven etc. - Make more specific what the added value of the paper is. Also - for the contribution of the paper to be less incremental it would be valuable to add more formality to the original results, for example - Gaussian approximation is claimed without any sort of verification of assumptions of any version of CLT or Law of Large Numbers.