The Dynamics of Signal Propagation in Gated Recurrent Neural Networks

1 INTRODUCTION . A common paradigm for research and development in deep learning involves the introduction of novel network architectures followed by experimental validation on a selection of tasks . While this methodology has undoubtedly generated significant advances in the field , it is hampered by the fact that the full capabilities of a candidate model may be obscured by difficulties in the training procedure . It is often possible to overcome such difficulties by carefully selecting the optimizer , batch size , learning rate schedule , initialization scheme , or other hyperparameters . However , the standard strategies for searching for good values of these hyperparameters are not guaranteed to succeed , especially if the trainable configurations are constrained to a low-dimensional subspace of hyperparameter space , which can render random search , grid search , and even Bayesian hyperparameter selection methods unsuccessful . In this work , we argue that for long sequence tasks , the trainable configurations of initialization hyperparameters for LSTMs and GRUs lie in just such a subspace , which we characterize theoretically . In particular , we identify precise conditions on the hyperparameters governing the initial weight and bias distributions that are necessary to ensure trainability . These conditions derive from the observation that in order for a network to be trainable , ( a ) signals from the relevant parts of the input sequence must be able to propagate all the way to the loss function and ( b ) the gradients must be stable ( i.e. , they must not explode or vanish exponentially ) . As shown in Figure 1 , training of recurrent networks with standard initialization on certain tasks begins to fail as the sequence length increases and signal propagation becomes harder to achieve . However , as we shall show , a suitably-chosen initialization scheme can dramatically improve trainability on such tasks . We study the effect of the initialization hyperparameters on signal propagation for a very broad class of recurrent architectures , which includes as special cases many state-of-the-art RNN cells , including the GRU ( Cho et al . ( 2014 ) ) , the LSTM ( Hochreiter & Schmidhuber ( 1997 ) ) , and the peephole LSTM ( Gers et al . ( 2002 ) ) . The analysis is based on the mean field theory of signal propagation developed in a line of prior work ( Schoenholz et al . ( 2016 ) ; Xiao et al . ( 2018 ) ; Chen et al . ( 2018 ) ; Yang et al . ( 2019 ) ) , as well as the concept of dynamical isometry ( Saxe et al . ( 2013 ) ; Pennington 200 400 600 800 Sequence length 0.0 0.2 0.4 0.6 0.8 1.0 Ac cu ra cy Critical initialization Standard initialization Figure 1 : Critical initialization improves trainability of recurrent networks . Test accuracy for peephole LSTM trained to classify sequences of MNIST digits after 8000 iterations . As the sequence length increases , the network is no longer trainable with standard initialization , but still trainable using critical initialization . et al . ( 2017 ; 2018 ) ) that is necessary for stable gradient backpropagation and which was shown to be crucial for training simpler RNN architectures ( Chen et al . ( 2018 ) ) . We perform a number of experiments to corroborate the results of the calculations and use them to motivate initialization schemes that outperform standard initialization approaches on a number of long sequence tasks . 2 BACKGROUND AND RELATED WORK . 2.1 MEAN FIELD ANALYSIS OF NEURAL NETWORKS . Signal propagation at initialization can be controlled by varying the hyperparameters of fullyconnected ( Schoenholz et al . ( 2016 ) ; Yang & Schoenholz ( 2017 ) ) and convolutional ( Xiao et al . ( 2018 ) ) feed-forward networks , as well as for simple gated recurrent architectures ( Chen et al . ( 2018 ) ) . In all these cases , such control was used to obtain initialization schemes that outperformed standard initializations on benchmark tasks . In the feed-forward case , this enabled the training of very deep architectures without the use of batch normalization or skip connections . By forward signal propagation , we specifically refer to the propagation of correlations between inputs at the start of an input sequence through the hidden states of the recurrent network at later times , as will be made precise in Section 4.2 . This appears to be a necessary but not sufficient condition for trainability on certain tasks . Another is the stability of the gradients , which depend on the state-to-state Jacobian matrix , as discussed in Bengio et al . ( 1994 ) . In our context , the goal of the backward propagation analysis is to improve the conditioning of the Jacobian by controlling the first two moments of its squared singular values . Forward signal propagation and the spectral properties of the Jacobian at initialization can be studied using mean field theory and random matrix theory ( Poole et al . ( 2016 ) ; Schoenholz et al . ( 2016 ) ; Yang & Schoenholz ( 2017 ; 2018 ) ; Xiao et al . ( 2018 ) ; Pennington et al . ( 2017 ; 2018 ) ; Chen et al . ( 2018 ) ; Yang et al . ( 2019 ) ) . As neural network training is a nonconvex problem , using a modified initialization scheme could lead to convergence to different points in parameter space in a way that adversely affects the generalization error . We provide some empirical evidence that this does not occur , and in fact , the use of initialization schemes satisfying these conditions has a beneficial effect on the generalization error . 2.2 THE EXPLODING/VANISHING GRADIENT PROBLEM AND SIGNAL PROPAGATION IN RECURRENT NETWORKS . The exploding/vanishing gradient problem is a well-known phenomenon that hampers training on long time sequence tasks ( Bengio et al . ( 1994 ) ; Pascanu et al . ( 2013 ) ) . Apart from the gating mechanism , there have been numerous proposals to alleviate the vanishing gradient problem by constraining the weight matrices to be exactly or approximately orthogonal ( Pascanu et al . ( 2013 ) ; Wisdom et al . ( 2016 ) ; Vorontsov et al . ( 2017 ) ; Jose et al . ( 2017 ) ) , or more recently by modifying some terms in the gradient ( Arpit et al . ( 2018 ) ) , while exploding gradients can be handled by clipping ( Pascanu et al . ( 2013 ) ) . Another recently proposed approach to ensuring signal propagation in long sequence tasks introduces auxiliary loss functions ( Trinh et al . ( 2018 ) ) . This modification of the loss can be seen as a form of regularization . Chang et al . ( 2019 ) study the connections between recurrent networks and certain ordinary differential equations and propose the AntisymmetricRNN that can capture long term dependencies in the inputs . While many of these approaches have been quite successful , they typically require modifying the training algorithm , the loss function , or the architecture , and as such exist as complementary methods to the one we investigate here . We postpone the investigation of a combination of techniques to future work . 2.3 DYNAMICS BEYOND INITIALIZATION . While the analysis below applies to the network at initialization , it was recently shown that as the network width becomes large , parameters move very little from their initial values ( Jacot et al . ( 2018 ) ; Lee et al . ( 2019 ) ) . A non-trivial and surprising consequence is that covariance matrices like those studied here remain constant during training . While Jacot et al . ( 2018 ) ; Lee et al . ( 2019 ) do not specifically study RNNs , the analysis of Yang ( 2019 ) suggests the conclusions are likely to carry over . Moreover , even for finite-width networks , simply satisfying the conditions derived in this work at initialization can lead to a dramatic improvement in trainability as we show below . 3 NOTATION . We denote matrices by bold upper case Latin characters and vectors by bold lower case Latin characters . Dx denotes a standard Gaussian measure . The normalized trace of a random N × N matrix A , 1NEtr ( A ) , is denoted by τ ( A ) . ◦ is the Hadamard product , σ ( · ) is a sigmoid function and both σ ( · ) , tanh ( · ) act element-wise . We denote by Da a diagonal matrix with a on the diagonal . 4 MEAN FIELD ANALYSIS OF SIGNAL PROPAGATION AND DYNAMICAL ISOMETRY . 4.1 MODEL DESCRIPTION AND IMPORTANT ASSUMPTIONS . We begin with a general description of recurrent architectures that can be specialized to the GRU , LSTM and peephole LSTM among others . We denote the state of a recurrent network at time t by st ∈ RN with s0i ∼ D0 , and a sequence of inputs to the network by { z1 , ... , zT } , zt ∈ RN . The state evolution of the network is given by st = f ( st−1 , { u1k } , ... , { utk } , zt ) ( 1a ) where f is an element-wise , affine function of st−1 and { utk } is a set of pre-activations utk ∈ RN , k ∈ K defined for a set of subscripts K. They are given by utk = Wks t−1 + Ukz t + bk ( 1b ) where Wk , Uk ∈ RN×N , bk ∈ RN . We define additional pre-activations utk2 = Wk2Dgk ( utk ) s t−1 + Uk2z t + bk2 ( 1c ) where gk : RN → RN is an element-wise function and utk is defined as in eqn . ( 1b ) 1 . In cases where there is no need to distinguish between variables of the form ( 1b ) and ( 1c ) we will refer to both as 1Variables of the form ( 1c ) will be present only in the GRU ( see Table 1 ) . utk . This state update equation describes all the architectures studied in this paper , as well as many others , as detailed in Tables 1,3 . The pre-activations utk , which are typically used to construct gates in recurrent network with the application of appropriate nonlinearities , will be of interest because they will tend in distribution to Gaussians as the network width is taken to infinity . We assume Wk , ij ∼ N ( 0 , σ2k/N ) , Uk , ij ∼ N ( 0 , ν2k/N ) , bk , i ∼ N ( µk , ρ2k ) i.i.d . and denote Θ =⋃ k { σ2k , ν2k , ρ2k , µk } . As in Chen et al . ( 2018 ) , we make the untied weights assumption that Wk is independent of st . Although the assumption of untied weights may seem counter-intuitive in the context of recurrent networks , recall that we are only interested in characterizing various statistics of the network , not in computing the entire input-output map , for which using tied weights would certainly be crucial . For merely characterizing certain statistics , the untied weights assumption actually has long history of yielding correct predictions , and therefore it is not at all unnatural to adopt : 1 ) when computing the covariance of gradients in an MLP , the weights during backpropagation can be treated as independently sampled from the weights used during the forward propagation ( Schoenholz et al. , 2016 ; Yang & Schoenholz , 2017 ) ; 2 ) when computing the Jacobian singular value distribution of an MLP , the pre-activations h = Wx can be treated as the result of multiplying an iid copy W′ of W ( i.e . h = W′x ) ( Pennington et al. , 2017 ; 2018 ; Hayase , 2019 ) ; 3 ) a wide , randomly initialized MLP corresponds to the same Gaussian process whether or not its weights across layers are tied ( Poole et al. , 2016 ) , and similarly , a simple RNN evolves the same whether or not its weights are tied across time , as long as its width is large ( Chen et al. , 2018 ) . Indeed , we provide empirical evidence that calculations performed under this assumption still have considerable predictive power , even in cases where it is violated .

Summary: The aim of this paper is to suggest randomized initializations for the various weights of a recurrent neural network (GRUs and various LSTMs are covered), such that training these networks gets to a successful start, when the model is trained on long sequences. Instead of being heuristic, their approach follows first principles of analyzing signal propagation through time, using ideas from statistical thermodynamics (mean field approximations). Some experiments, on toy datasets, validate their approach.