Initializing ReLU networks in an expressive subspace of weights

1 INTRODUCTION . Rectified Linear Unit ( ReLU ) Fukushima ( 1969 ) ; Fukushima & Miyake ( 1982 ) is the most widely used non-linear activation function in Deep Neural Networks ( DNNs ) LeCun et al . ( 2015 ) ; Ramachandran et al . ( 2018 ) ; Nair & Hinton ( 2010 ) , applied to various tasks like computer vision Glorot et al . ( 2011b ) ; Krizhevsky et al . ( 2012 ) ; He et al . ( 2015 ) , speech recognition Maas et al . ( 2013 ) ; Tóth ( 2013 ) ; Hinton et al . ( 2012 ) , intelligent gaming Silver et al . ( 2016 ) , and solving scientific problems Seif et al . ( 2019 ) . ReLU , φ ( x ) = max ( 0 , x ) , outperforms most of the other activation functions proposed Glorot et al . ( 2011a ) . It has several advantages over other activations . ReLU activation function is computationally simple as it essentially involves only a comparison operation . ReLU suffers less from the vanishing gradients , a major problem in training networks with sigmoid-type activations that saturate at both ends Glorot et al . ( 2011b ) . They generalize well even in the overly parameterized regime Maennel et al . ( 2018 ) . Despite its success , ReLU also has a few drawbacks , one of which is the dying ReLU problem He et al . ( 2015 ) ; Trottier et al . ( 2017 ) . The dying ReLU is a type of vanishing gradient problem in which the network outputs zero for all inputs and is dead . There is no gradient flow in this state . ReLU also suffers from exploding gradient problem , which occurs when backpropagating gradients become large Hanin ( 2018 ) . Several methods are proposed to overcome the vanishing/exploding gradient problem ; these can be classified into three categories Lu et al . ( 2020 ) . The first approach modifies the architecture , which includes using modified activation functions Ramachandran et al . ( 2018 ) ; He et al . ( 2015 ) ; Trottier et al . ( 2017 ) ; Clevert et al . ( 2016 ) ; Klambauer et al . ( 2017 ) ; Hendrycks & Gimpel ( 2016 ) , adding connections between non-consecutive layers ( residual connections ) He et al . ( 2016 ) , and optimizing network depth and width . The proposed activations are often computationally less efficient and require a fine-tuned parameter Lu et al . ( 2020 ) . The second approach relies on normalization techniques Ba et al . ( 2016 ) ; Ioffe & Szegedy ( 2015 ) ; Salimans & Kingma ( 2016 ) ; Ulyanov et al . ( 2016 ) ; Wu & He ( 2018 ) , the most popular one being batch normalization Ioffe & Szegedy ( 2015 ) . Batch normalization prevents the vanishing and exploding gradients by normalizing the output at each layer but with an additional computational cost of up to 30 % Mishkin & Matas ( 2016 ) . A related strategy involves using the self-normalizing activation ( SeLU ) , which by construction ensures output with zero mean and unit variance Klambauer et al . ( 2017 ) . The third approach focuses on the initialization of the weights and biases . As local ( gradient-based ) algorithms are used for optimization Kingma & Ba ( 2015 ) ; Zeiler ( 2012 ) ; Duchi et al . ( 2011 ) , it is challenging to train deep networks with millions of parameters Du et al . ( 2019 ) ; Srivastava et al . ( 2015 ) , and optimal initialization is essential for efficient training Nesterov ( 2014 ) . He-initialization He et al . ( 2015 ) is a commonly used strategy that uses uncorrelated Gaussian weights with variance 2N , where N is the width of the network . Recently , Lu et al . ( 2020 ) proposed Random asymmetric initialization ( RAI ) , which reduces the probability of dead ReLU at the initialization . In this paper , we aim to further improve the initialization scheme for ReLU networks . A growing body of work has analyzed signal propagation in infinitely wide networks to understand the phase diagram of forward-propagation in DNNs Saxe et al . ( 2014 ) ; Poole et al . ( 2016 ) ; Raghu et al . ( 2017 ) ; Schoenholz et al . ( 2017 ) ; Lee et al . ( 2018 ) ; Hayou et al . ( 2019a ) ; Li & Saad ( 2018 ; 2020 ) ; Bahri et al . ( 2020 ) . We mention a few results for ReLU networks . Hayou et al . ( 2019a ) showed that correlations in input signals propagating through a ReLU network always converge to one . Many other works found that ReLU networks are in general biased towards computing simpler functions De Palma et al . ( 2019 ) ; Rahaman et al . ( 2019 ) ; He et al . ( 2020 ) ; Valle-Perez et al . ( 2019 ) ; Hanin & Rolnick ( 2019 ) , which may account for their better generalization properties even in the overly parameterized regime . However , from their successful application in different domains , one may guess that they should be capable of computing more complex functions . There might be a subspace of the parameters where the network can represent complex functions . Li & Saad ( 2018 ; 2020 ) applied weight and input perturbations to analyze the function space of ReLU networks . They found that ReLU networks with anti-correlated weights compute richer functions than uncorrelated/positively correlated weights . Consistent with this , Shang et al . ( 2016 ) found that ReLU CNN ’ s produce anti-correlated feature matrices after training . These studies motivated us to analyze the phase diagram of signal propagation in ReLU networks with anti-correlated weights . Following the mean-field theory of signal propagation proposed by Poole et al . ( 2016 ) , we found that ReLU networks with anti-correlated weights have a chaotic phase , which implies higher expressivity . In contrast , ReLU networks with uncorrelated weights do not have a chaotic phase . Furthermore , we find that initializing ReLU networks with anti-correlated weights results in faster training . We call it Anti-correlated initialization ( ACI ) . Additional improvement in performance is achieved by incorporating RAI , which reduces the dead node probability . This combined scheme , which we call Random asymmetric anti-correlated initialization ( RAAI ) , is the main result of this work and is defined as follows . We pick weights and bias incoming to each node from anti-correlated Gaussian distribution and replace one randomly picked weight/bias with a random variable drawn from a beta distribution . The code to generate weights drawn from the RAAI distribution is given in Appendix G. We analyze the correlation properties of RAAI and show that it performs better than the best-known initialization schemes on tasks of varying complexity . It may be of concern that initialization in an expressive space may lead to overfitting , and we do observe the same for ACI for deeper networks and complex tasks . In contrast , RAAI shows no signs of overfitting and performs consistently better than all other initialization schemes . We organize the article as follows . First , we contrast the mean-field analysis of ReLU networks with correlated weights with uncorrelated in Section 2 . Next , Section 3 analyzes the critical properties of correlations in input signals for RAI and RAAI . Then , in Section 4 , we describe the various tasks used to validate the performance of different initialization schemes in Section 5 . Lastly , Section 6 concludes the article . 2 MEAN-FIELD ANALYSIS OF SIGNAL PROPAGATION WITH CORRELATED WEIGHTS . This section presents the mean-field theory of signal propagation ( proposed by Ref . Poole et al . ( 2016 ) ) in ReLU networks with correlated weights and compares it with uncorrelated weights . Unlike Ref . Li & Saad ( 2018 ; 2020 ) , which study perturbation to a ReLU network , we aim to understand the phase diagram of the signal propagation . Furthermore , we provide numerical results to corroborate the mean-field results . Consider a fully connected neural network with L layers ( in addition to the input layer ) and Nl nodes in layer l. The layer index ranges between 0 and L. For an input signal s0 = x , we denote the pre-activation at node i in layer l by hli ( x ) and activation by s l i ( x ) . A signal ( s l−1 1 , . . , s l−1 Nl ) at layer l − 1 propagates to layer l by the rule hli ( x ) = Nl−1∑ j=1 wlijs l−1 j ( x ) + b l i where l ∈ { 1 , L } sli ( x ) = φ ( h l i ( x ) ) , where φ is the non-linear activation function and wlij , b l i are the weights and biases . We consider correlations within the set of weights ( wli ) incoming to each node i . The correlated Gaussian distribution is P ( wl1 , w l 2 , w l 3 . . . ) = Nl∏ i e ( − 1 2 ( w l i ) TA−1wli ) √ ( 2π ) Nl−1 |A| with A = σ2w Nl−1 ( I− k 1 + k J Nl−1 ) , ( 1 ) where I is the identity matrix , J is an all-ones matrix , and k parameterizes the correlation strength . Positively correlated and anti-correlated regimes correspond to the regions −1 < k < 0 and k > 0 , respectively , whereas k = 0 generates uncorrelated weights . The overall scaling by 1/Nl−1 in the covariance matrix ensures that the input contribution from the last layer to each node is O ( 1 ) . The bias is drawn from a Gaussian distribution bli ∼ N ( 0 , σ2b ) . Note that weights reaching two different nodes are uncorrelated , and also the bias is uncorrelated with the weights . To track the layer-wise information flow , consider the squared length and overlap of the preactivations for two input signals , s0 = x1 and s0 = x2 , after propagating to layer l qlh ( xa ) = 1 Nl Nl∑ i=1 ( hli ( xa ) ) 2 where a ∈ { 1 , 2 } and 1 ≤ l ≤ L qlh ( x1 , x2 ) = 1 Nl Nl∑ i=1 hli ( x1 ) h l i ( x2 ) . Assuming self averaging , consider an average over the weights and the bias incoming to layer l. For simplicity of notations later , we use the same symbol for averaged qlh . qlh ( xa ) = σ2w Nl−1 Nl−1∑ j , m=1 ( δj , m − k 1 + k 1 Nl−1 ) φ ( hl−1j ( xa ) ) φ ( h l−1 m ( xa ) ) + σ 2 b qlh ( x1 , x2 ) = σ2w Nl−1 Nl−1∑ j , m=1 ( δj , m − k 1 + k 1 Nl−1 ) φ ( hl−1j ( x1 ) ) φ ( h l−1 m ( x2 ) ) + σ 2 b , For large width , each hl−1i is a weighted sum of a large number of zero-mean random variables . Thus , we expect the joint distribution of hl−1i ( x1 ) and h l−1 i ( x2 ) to converge to a zero-mean Gaussian with a covariance matrix with diagonal entries ql−1h ( x1 ) , q l−1 h ( x2 ) and off-diagonal entries ql−1h ( x1 , x2 ) . On replacing the average over h l−1 i ( this is equivalent to considering an average over all previous layers ) in the last layer with an average over this Gaussian distribution , we obtain iterative maps for the length and overlap . Specializing to ReLU activation yeilds the equations qlh ( x ) = σ2w 2 ( 1− k 1 + k 1 π ) ql−1h ( x ) + σ 2 b qlh ( x1 , x2 ) = σ2w 2 ( f ( cl−1h ) − k 1 + k 1 π ) √ ql−1h ( x1 ) q l−1 h ( x2 ) + σ 2 b f ( cl−1h ) = cl−1h 2 + cl−1h π sin−1 ( cl−1h ) + 1 π √ 1− ( cl−1h ) 2 , ( 2 ) where clh = qlh ( x1 , x2 ) √ qlh ( x1 ) q l h ( x2 ) is the correlation coefficient between the two signals reaching layer l ( for details of the derivation , see Appendix A ) . Poole et al . ( 2016 ) found that the signal ’ s length reaches its fixed point within a few layers , and the fixed point of the correlation coefficient , c∗h , can be estimated with the assumption that q l h ( x ) has reached its fixed point q∗h . We can check that c ∗ h = 1 is always a fixed point of the recursive map ( Eqn . 2 ) under this assumption . The stability of the fixed point c∗h = 1 is determined by χ1 ≡ ∂clh ∂cl−1h ∣∣∣ cl−1h =1 , which evaluates to σ 2 w 2 . χ1 separates the parameter space into two phases —first , an ordered phase with χ1 < 1 , where the c∗h = 1 fixed point is stable ; and second , a chaotic phase with χ1 > 1 , where the c∗h = 1 fixed point is unstable . χ1 = 1 defines the phase boundary line . In the ordered phase , two distinct signals will become perfectly correlated asymptotically . In the chaotic phase , the correlations converge to a stable fixed point below unity . Two closely related signals will eventually lose correlations in this phase . This suggests that initializing the network with parameters ( σ2w , σ 2 b ) at the phase transition boundary ( corresponding to an infinite depth of correlations ) allows for an optimal information flow through the network Poole et al . ( 2016 ) ; Schoenholz et al . ( 2017 ) ReLU networks with uncorrelated weights ( k = 0 ) The above analysis is applied assuming q∗h is finite shows that ReLU networks with uncorrelated weights do not have a chaotic phase , and any two signals propagating through a ReLU network become asymptotically correlated for all values of ( σ2w , σ 2 b ) . In other words , c ∗ h = 1 is always a stable fixed point . However , the parameter space can be classified into two phases based on the boundedness of the fixed point q∗h of the length map ( Eqn . 2 ) - first , a bounded phase where q∗h is finite and non-zero ; second , an unbounded phase , where q ∗ h is either zero or infinite Lee et al . ( 2018 ) ; Hayou et al . ( 2019b ) . The two phases are separated by the boundary σ2w = 2 . Figure 1a depicts the phase diagram for ReLU networks with uncorrelated weights . Note that the analysis of the stability of c∗h = 1 fixed point in ReLU networks is valid only in the bounded phase . However , numerical results presented in Fig . 2 indicate that the fixed point remains stable even in the unbounded phase . ReLU networks with correlated weights The phase diagram for ReLU networks with correlated weights can be analyzed similarly . The length is bounded if σ2w < gk = 2 ( 1− k1+k 1 π ) . Thus , for anti- correlated weights ( k > 0 ) , the boundary gk moves upwards relative to the k = 0 case ( see Fig . 1a ) . The c∗h = 1 fixed point of the correlations is unstable in this region of the bounded phase . In summary , anti-correlations induce a bounded chaotic phase in 2 < σ2w < gk ( see Fig . 1b ) . We demonstrate these results numerically in Fig . 3 for a correlation strength of k = 100 . As predicted by the above equations , the stability of the fixed point c∗h = 1 changes at σ 2 w = 2 , and the length diverges at gk=100 = 2.92 . In contrast , for positively correlated weights , the length ’ s fixed point boundary shifts downward resulting in a similar phase diagram as uncorrelated weights . As a result , a ReLU network with anti-correlated weights can be more expressive by taking advantage of a chaotic phase , and it may be beneficial for a ReLU network to remain in this subspace . Thus , we propose initializing ReLU networks with anti-correlated weights at the order to chaos boundary ( σ2w , σ 2 b ) = ( 2 , 0 ) . We call it Anti-Correlated Initialization ( ACI ) . Appendix B demonstrates that ReLU networks initialized with anti-correlated weights give an advantage over He initializations for a range of tasks . Many alternatives are proposed to improve ReLU networks Trottier et al . ( 2017 ) ; Lu et al . ( 2020 ) ; Clevert et al . ( 2016 ) . Of particular interest is Random asymmetric initialization ( RAI ) , which aims to increase expressivity through an independent strategy of reducing the dead node probability . In the next section , we analyze the correlation properties of RAI and then combine it with ACI to propose a new initialization scheme RAAI , which has both a chaotic phase and low dead node probability .

Summary: The paper studies the signal propagation properties of deep ReLU networks at initialization time. For a standard ReLU network, the correlation between activations of two training examples converges to one as the depth goes to infinity. To resolve that, the authors proposed to use correlated weights at initialization time to avoid this pathology. With correlated weights, the correlation could instead converge to a value smaller than one (so-called chaotic phase). Empirically, the authors found the networks initialized with correlated weights train faster.