Recurrent neural networks: vanishing and exploding gradients are not the end of the story
Authors: Nicolas Zucchet
Antonio Orvieto
Venue: NeurIPS 2024 poster (2024)
Keywords: RNN
optimization
long-range dependencies
SSM
LSTM
Primary area: optimization_for_deep_networks

Abstract: Recurrent neural networks (RNNs) notoriously struggle to learn long-term memories, primarily due to vanishing and exploding gradients. The recent success of state-space models (SSMs), a subclass of RNNs, to overcome such difficulties challenges our theoretical understanding. In this paper, we delve into the optimization challenges of RNNs and discover that, as the memory of a network increases, changes in its parameters result in increasingly large output variations, making gradient-based learning highly sensitive, even without exploding gradients. Our analysis further reveals the importance of the element-wise recurrence design pattern combined with careful parametrizations in mitigating this effect. This feature is present in SSMs, as well as in other architectures, such as LSTMs. Overall, our insights provide a new explanation for some of the difficulties in gradient-based learning of RNNs and why some architectures perform better than others.

Recurrent neural networks: vanishing and exploding gradients are not the end of the story
Nicolas Zucchet
Antonio Orvieto
Department of Computer Science
ELLIS Institute Tübingen
ETH Zürich nzucchet@ethz.clt
MPI for Intelligent Systems
Tübingen AI Center antonio@tue ellil.ell....... eu
Abstract
Recurrent neural networks (RNNs) notoriously struggle to learn long-term mem- ories, primarily due vanishing and exploding gradients. The recent success of deep state-space models (SSMs), a subclass of RNNs, to overeome such difficulties challenges our theoretical understanding. tion challenges of RNNs and discover that, as the memory of a network increases, changes in parameters result in increasingly large output variations, making gradient- based learning highly sensitive, even without exploding gradients. Our analysis further reveals the importance of the element- wise recurrence design pat- tern combined with careful parametrizations in mitigating this effect. This feature is present in deep SSMs, as well as other architectures, such as LSTMs. Overall, our insights provide new explanation for some of the difficulties in gradient-based learning of RNNs and why some architectures perform better than others.
this paper, delve into the optimiza-
Recurrent neural networks [RNNs; 1, 2] have long been the canonical architecture for modeling temporal data [3, 4]. However, they are notoriously difficult to train on long sequences, as error signals flowing backward in time tend to either vanish explode [5-8]. Attention mechanisms [9], as featured in transformers [10], address these issues by enabling direct token-to-token communi- cation, considerably simplifying signal propagation across long time intervals. Yet, their superior performance comes with increased computational and memory costs, due to their quadratic scaling in the sequence length. This limitation has motivated significant research aimed at making transformers more efficient [11--.]
A promising line of research in this direction involves a new type of linear recurrent networks known deep state-space models [SSMs; 16-22]. These models trade expressivity for faster training speed, and they have been shown to be particularly effective at capturing long-range dependencies. In this paper, we wonder whether this effectiveness can be solely attributed to their ability to avoid vanishing and exploding gradients. The simplicity of such models presents an opportunity for in-depth theoretical analysis. We focus on signal propagation within these models.
After reviewing classical results on recurrent neural networks in Section 1, we demonstrate that they suffer from understudied problem: the recurrent network encodes longer memories, the network's activity becomes increasingly sensitive to changes in its parameters, even when its dynamics remains stable. In Section 3, we then show that SSMs, well other architectures such as LSTMs, are well equipped to mitigate this issue. We then analyze a simple teacher-student task (Section 4). This task already reveals the remarkable complexity underlying the learning of linear recurrent networks and enables us to verify empirically our theory. Finally, we discuss how our findings extend to more realistic scenarios (Section 5), both in terms of architectures and data. Overall, our paper provides theoretical insights into the training of recurrent neural networks, an area where such analysis rare. While vanishing and exploding gradients are well-known challenges, our results demonstrate that this is not the end of the story there exists an additional layer of complexity beyond them.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Vanishing and exploding gradients
Let us first introduce the notations we will be using throughout the rest of the paper. We consider recurrent neural network with hidden state bes update function fe parametrized by 0, and input sequence (Iz). The average performance of the network measured by a loss L. We have
T
h++1 = (o(h,, and L. = E
Σ Ta(ht))
The gradient of the instantaneous loss L. with respect to the parameters 0 is then equal to
dl+ aI, dhg
ah,
dhe 0So
db 00 (hc-1, TV)
(2)
d0 The d0
The
e'tt
In the equation above, we used @ to denote partial derivatives and for total derivatives. Using this notation enables us to distinguish between 0,, Lt, which corresponds to the error backpropagated from the current loss term to the hidden state through the readout function, and d,, I,, which accumulates the errors that are backpropagated through the future hidden state values. In particular, 0,, I. Ōh T and c,,, L = Ô, Le (ht) + Eess da, Ly (.r)). When stacking several recurrent layers on top each other, 0n, L corresponds to the current error being backpropagated to the hidden state he through the hierarchy of the network and d,, L to future error signals backpropagated through the recurrence.
Early work [5] highlighted the difficulty for gradient descent to make recurrent neural networks remember past inputs that will later be useful produce desired behavior. This is due the fact that error signals flowing backward in time tend to either explode or vanish. The key quantity is
dhg dhe
0h 0h,
[-1 II i=r"
0fo (h,, Z:41).
II
Öh
One remark that this quantity exponentially converges 0 when the spectral radius of the Jacobian Jo is upper bounded by a constant strictly smaller than 1, and can exponentially explode there exists some component bigger than 1. The error signal at time backpropagated time t' behaves similarly, as d Lt Ōn Lt ch. hr. Gradient-based learning of long-term memories is thus difficult: the contribution of past hidden states the current loss becomes either negligible or predominant as the time span considered increases.
Since then, the analysis has been refined [6 8] and the development of recurrent architectures has mostly been driven by the desire to solve this pathological issue. Most famously, the LSTM [3] unit, and later on the GRU [23], solve this problem by using memory neurons that facilitate direct information storage and retrieval, and thus facilitate error backpropagation. Other approaches to solving this problem, name few, involve gradient clipping [24, 8), activity normalization [25- 27], careful weight initialization [28, 29] or enforcing architectural constraints such as hierarchical processing [30, 31], orthogonal weight matrices [32-34] and oscillations [35-37].
The curse of memory
According common deep learning wisdom, it is often believed that solving the vanishing and exploding gradients problem enables recurrent neural networks to learn long-term dependencies. We challenge this view and question: solving those issues really enough to ensure well-behaved loss landscapes? We answer negatively by showing that gradients can explode as the memories of the network are kept for longer, even when the dynamics of the network remains stable.
2.1 Intuition
Recurrent neural networks have something special: the very same update function fo is applied over and over. Therefore, modifying the parameters will not only influence one update, as changing weights of given layer in a feedforward neural network would, but all of them. As the memory of the network gets longer, the hidden states keep a trace of the effects of more updates. Hidden states thus become increasingly sensitive to parameter changes. This is the curse of memory. We borrow

---

4
C
moment dyb
10ts
10a
104
10
0.0
10"
0.99
0.9
10-2
10
0.98
10
0.99
0.96
0.999
10"4
10
0.92
10a
1.0
10-6
10
0.68
0.36
0.99
0.9999
0.0
0.5
1.0
A
8-6
Figure 1: Optimization of recurrent neural networks gets harder as their memory increases. A. Evolution of the second moment of dx/, as function of the recurrent parameter À and of the input x auto-correlation decay rate P. when 14+1 = Xh, + 2y. As the memory of the network increases (A -> 1), hy becomes more sensitive to changes in X, particularly as the elements in the input sequence are more correlated (p + 1). The explosion of dyyh, is faster than the one of hy as highlighted with the grey line obtained for 1. See Section 2.2 for more detail. B, C. Illustration of the phenomenon on the toy one-dimensional teacher-student task of Section 4.1, in which the teacher is parametrized by a real number X and the student by complex number A. In B, À varies on the real axis, and varies on the circle of radius X parametrized by in C. The loss becomes sharper as information is kept longer in memory, making gradient- based optimization nearly impossible.
the term from [38, 391, although use it in different context, and note that Martens and Sutskever [40] hypothesized that such phenomenon could arise RNNs and hinder their optimization.
Let us formalize our intuition by considering the sensitivity of the hidden state h, on the parameters 0:
dhg Σ dhe 0se d0
ahe 00 (he (h
When information stays in the network's 's memory for longer, the number of non-neeee non-negligibbe Jacobian di,, h4 terms increases. As result, the magnitude of this sensitivity increases when the network encodes longer-term dependencies, and learning / becomes trickier. It crucial to observe that this phenomenon arises even when exploding gradients are removed from the picture by constraining the eigenvalues of the recurrent Jacobian to be smaller than one and ensuring that the network dynamics remains stable. The rest of this section will be dedicated to studying this behavior quantitatively.
2.2 Signal propagation in linear diagonal recurrent neural networks
We study how hidden state and gradient magnitudes evolve as the network encodes longer-term dependencies. Ideally, these quantities do not vanish or explode, as it improves the conditioning of the loss landscape [41] and cases optimization [42, 43]. We operate under the following assumptions:
a) Linear diagonal recurrent neural networks. We restrict ourselves to update functions of the form (h,,, 1++1) À h F++11 with vector of the size of h and @ the element-wise product. For ease of exposition, we present results for real-valued À here; see Appendix B.2 for the complex-valued setting. While this assumption is strong, it allows us identify some crucial mechanisms and models like S4 [17], S5 [19] and LRUs [20] satisfy it. We later show that our analysis can model some features of more sophisticated networks. Note that we do not consider the input and readout mappings usually featured in recurrent layers as they are feedforward layers and signal propagation within them is already well understood [e.g., 44, 45].
b) Infinite time horizon. We consider infinite sequences and initialize the network dynamics at to = -00 It simplifies our calculations while being a reasonable assumption when the sequences considered are longer than the characteristic timescales of the dependencies want to leam.
Wide-sense stationarity. We assume the different quantities that the network receives, which include the inputs 3t, to be wide sense stationary (WSS). random process X, said to be WSS if its auto-correlation function is independent of time, that is, for all t E Z and A E Z, Ex [X] =: Rx (A), where Ex denotes the expectation over the data. corresponds to assuming that the statistics of the data are invariant to time shifts. This is standard assump- tion when analyzing stochastic processes [46]. It keeps our calculations concise and does not qualitatively affect our conclusions (cf. Section 5). We discuss how to relax it in Appendix B.2.4.
3

---

We are now equipped to analyze signal propagation in one recurrent layer, both in the forward and backward passes. We show that both hidden states and backpropagated errors explode as |^| 1.
Forward pass. Here, we are interested in understanding how the hidden state second moment E|hğ] evolves as function of À and of the input auto-correlation function Rg. After calculation that we defer to Appendix B.2, we obtain
E h
(0) +2
Σ A R A21
(5)
1 X2
Importantly, this quantity goes to infinity as longer-term dependencies are encoded within the network, that is X - 1. Additionally, the divergence speed depends the input data distribution: it increases as consecutive time steps in the input distribution become more correlated (i.e., less of the R. (A) terms are negligible). This behavior already highlights potential difficulties of gradient-based learning of deep neural networks containing linear recurrent layers as the variance of neural activity can become arbitrarily large, hindering learning abilities of deeper layers.
Backward pass. Let us first derive the gradient of the loss with respect to A. Using the chain rule we have dxL= = E. Or, dyhg. We thus seek understand how dy behaves. We remark that dyhe=1 = Adx/t + he so that dy/r, is a low pass filtered version of the hidden state, which is itself a low pass filter version of the inputs. It therefore comes as no surprise that the second moment of dy diverges faster than the one of h2 when Al
+ 1. More precisely, we get
E
dhe Yp
2
1 1+
(0) +2
Σ 12) A21
Σ EAAAR((A)) (6) A>I
(1-3
(1 (1- X3-
We plot the exact behavior of this quantity when the auto-correlation of satisfies R. (A) p ام on Figure I and refer the interested reader to Appendix B.2 for derivation of Equation 6. More generally, the hidden state of the network, and thus its final output, becomes increasingly sensitive to changes in recurrent parameters as the network reaches the edge of dynamical stability (i) - 1).
The last quantity we need to consider is the error that is backpropagated to the inputs of the recurrent layer. It can be observed that the backward pass is dual to the forward pass in the sense that it is a recurrent process that receives backpropagated errors On, L and runs in reverse time:
dL dL Oh, dL the+1 OL Th dh
DL OL
+
dx, dhe dr. dhe+1 Oh,
ah.
in which we made use of 2z, h4 = 1. It follows that the analysis we did for the forward pass also holds here. Crucially, this implies that the explosion behavior will be most significant for the recurrent parameters rather than for potential input or readout weights.
2.3 Extending the analysis to the non diagonal case
We now generalize our results to fully connected linear recurrent neural networks of the form h++1 = Ah, + F₁. For the sake of the analysis, we assume that A is complex diagonalizable. that is there exists a complex-valued matrix /' and a complex-valued vector À such that A = Pd Note that this occurs with probability one under random initialization of A [20]. In this case,
he Ph
diag with ddiig diag(A)b) ding
p-120+1 I1+1
(8)
and
dh, Th, OP Th, dh, dieg dieg ax Th, dh, diak ap-1
(9)
dA aP AA
oh diag
dd 0A ding dP-1 0A
From the analysis above, we know that the dominating term in the limit |\] - among 0ph, dyhe and dp' h2 is dyhg, as P and p-1 act as readout and input weights. Given that all other terms do not directly depend on the magnitude of X, we have that d ~ Oggimchh he d d diag aĄA; cf. Appendix B.2.3 for formal statements. This has two consequences: First, the sensitivity of h on A will explode

---

as longer memories are encoded and this directly comes from the eigenvalues of A. Second. as each entry of A typically impacts all eigenvalues of the matrix, the explosion behavior will be distributed across all entries, whereas it was concentrated on the eigenvalues for the diagonal case. We will later observe that this has significant practical consequences and partly explains why fully connected linear RNNs are difficult to train. As side note, we remark that enforcing the matrix A to be orthogonal solves vanishing and exploding gradient issues but these weights may remain sensitive to learn because of the curse of memory.
3 Mitigating the curse of memory
We have discussed the sensitivity of recurrent networks to parameter updates. Given this problem, how can it be mitigated? We show that recurrent networks with diagonal connectivity are particularly well suited for this purpose. Besides enabling control over the Jacobian and avoiding exploding gradients, they facilitate the mitigation of the curse of memory. We additionally highlight that deep state-space models and gated RNNs inherently incorporate such mechanisms.
3.1 A solution: normalization and reparametrization
Both forward and backward passes explode as the network encodes longer memories. When h2+1 Xh, + Te+11 we argue that is relatively straightforward to mitigate this effect. We aim to keep E[hğ], (((d),,)))) and [[(dz,,,)² independent of X, similar to initialization schemes that maintain the magnitude of neural activity constant in deep networks [44, 45], regardless of the layer width [42, 47, 43].
Input normalization. A simple way to enforce E[? to stay constant is to introduce a scaling factor (()) applied to the inputs neuron receives, that satisfies (A))E[[]]] @(1). Given that the backward propagation of output errors to inputs is dual to the forward pass, the role of y in the backward pass will be similar. The value needs to take therefore both depends on the input distribu- tion to normalize the forward pass, as well as the output error distribution to normalize the backward pass. Perfect normalization is likely unrealistic, but some normalization can help, as shown in Figure 2.A.
A
10a*
without y
2nd moment. h
10s
102
0.0
0
0.99
0.9999
0.9
2nd moment 4,h @
without
0.99
10a
and exp
0.999
-.0
10a
10
Eigenvalue reparametrization. We are now left with keeping the gradient of the loss with respect to X under control. Input normalization partly reduces the memory- induced exploding effect, but not entirely as the variance of h is much larger than the one of h4 (cf. Fig. L.A). Reparametrization can close that gap. Indeed, if A is parametrized by w, we have that hh = A Choos- ing a parameterization that is more and more granular as goes to 1 thus helps in keeping the magnitude of h constant. Assuming is independent of À for simplicity, achieving (((d_,/, 2 ((1) requires solving the differ- ential equation (w) EE(x)))) = 1. While deriv- ing a universal optimal parametrization is again unrealistic due to dependency on the input distribution, reparametriza- tion definitely helps, as shown in Figure 2.B. Figure illustrates how it affects the loss landscape.
0
0.99
0.9999
A
Figure 2: Illustration of the effects of normalization and reparametrization. can effectively control the magnitude of A. [[h)]] and B. [((x)),)²) over all values when the input correlattno satisfies (A) اعهم with p but does not manage do to so for other type of distributions (p * 0). Here, we use 7()) VI X2, decouple from À when differentiating, and take exp(- exp(x)). as in [20]. The grey line indicates the value the two quantities take without any normaliza- tion and reparametrization, when 1.
The case of complex numbers. We have not yet discussed the case À E C, relevant for SSMs such as S4 [17]. We extend our analysis to complex-valie vvalued in Appendix B.3.2. Briefly, it reveals that changes in the magnitude of X have a similar impact as in the real case, but this similarity does not extend to the angle. To keep the sensitivity on the angle constant, its parametrization must depend on the magnitude of X. However, doing so hurts learning, particularly far from optimality, as we exemplify Appendix B.3.2. A key implication of this analysis is that the sensitivity of the hidden state on the angle of À explodes as its magnitude approaches 1.
5

---

3.2 Several RNN architectures implicitly alleviate the curse of memory
Deep state-space models, as well as gated RNNs, feature some form of normalization and reparametrization which help keeping signal propagation under control. We discuss how below.
Deep state-space models. Deep SSMs were originally motivated discretizations of the dif- ferential equation h Ah + Br [16]. Naïve discretization of the differential equation yields 14+1 = (Id + d + 1 which already provides some input normalization. More elaborate discretization schemes, such the zero-order hold, effectively reparametrize the A matrix, e-g. with exp(dlAA). Here, diagonalization arises from computational efficiency and simplicity reasons [18]. While such models can approximate any smooth mappings [48, 49], their expressivity remains limited [50]. The next generation of these models, including Mamba [21], incorporates input-dependent gates which modulate dt depending on the input It. The theory we developed above does not strictly apply to this setting as dt is no longer constant. However, since the recurrence Jacobian remains diagonal, we expect the qualitative behaviors we analyzed to remain.
Gated RNNs. While the original motivation behind gated RNNs such as LSTMs [3] GRUs [23] largely differs from the one of SSMs, they share similar mechanisms. In these networks, the memories stored in hidden neurons can be erased through forget gate, and incoming inputs can selectively be written in memory through an input gate. Mathematically, this corresponds to hidden state updates of the form h++1 fi+1 @ h4 + f++1 1 21,1, with the forget fa+1 and input 4++1 gates being independent non-linear functions of 2++1 and he The forget gate is akin to À and usually involves a sigmoid non-linearity, which has similar effect as reparametrizing À in the backward pass. The input gate can act as an input normalization depending on the initialization of the network or if is coupled to the forget gate in the GRU (/2 = in) [29]. Importantly, the gates here depend on the hidden states and thus make the Jacobian 0n h non diagonal. Yet, we argue that these architectures still have a bias towards diagonality. Indeed, the contributions of the hidden state through the forget and input gates are indirect, and they can be ignored when the weights connecting the hidden states to the gates are small. Empirically, we find that GRUs lie in this regime at initialization, ef. Section D.2, so that our theory accurately captures signal propagation in GRUs We additionally confirm that signal propagation is well behaved in gated RNNs in Section 5. In regimes which this approximation does not hold, studying signal propagation requires much more sophisticated analysis than the one we have done here [51].
4 A linear teacher-student analysis
We start our empirical analysis with teacher-student task using linear recurrent networks [52]. This is arguably the simplest setting in which one can train recurrent networks. Yer, as we shall see, it is already remarkably complex and captures some of the differences between different architectures observed in more realistic settings [20]. It additionally makes it possible to control the characteristic time constants of the data, which is only possible with synthetic data.
Our investigation starts with the one-dimensional setting to provide an intuitive illustration of the consequences of the curse of memory on the loss landscape. We then address the general setting and observe that linear networks indeed suffer from the curse of memory, and that the remedies we studied in the last section are effective. We additionally find that diagonality greatly modifies the structure of the loss landscape and helps optimizers with adaptive learning rates to compensate for eventual increased sensitivities.
4.1 The one-dimensional case
We first consider student and a teacher following the mieddmmesiional dynamics h441 Xh+ + It-1, with complex-valued parameter À for the student and X for the teacher. For simplicity, we independently draw each 2t+1 from a unit normal distribution (its autocorrelation function is R,, () âa=o) and note that other input distributions do not qualitatively change the results.
The performance of the student measured by loss L. that averages the per time-step losses
L, := A hh h over the entire sequence.

---

B
Figure 3: LRUs are better at replicating a teacher's behavior than linear RNNs. A. As the teacher encodes longer dependen- cies (v) + 1), the linear RNN struggles to reproduce it, but not the LRU. B. An ab- lation study (uo = 0.99) reveals that this gap mainly comes from having close to diagonal recurrent connectivity matrix. See Section 4.2 for more detail.
10
RMN
10-1
LRU
10-3
10
10-1
10
10-7
10
10-3
0.32
0.84
0.96
0.99
RNN
block agonaa
more
neurons
complex
numbers
norm.
LRU
This simple model already captures two key difficulties of gradicnt- based learning of recurrent neural networks. In Figure I, we plot the resulting loss landscape for different X values, when À evolves on the positive part of the real axis (Fig. 1.B) and when it evolves on the circle of radius [Λ*] in the complex plane (Fig. 1.C). We restrict As to have absolute values smaller than one: exploding gradients are out of the picture. Still, two difficulties for gradient-based learning appear here. On one side, vanishing gradients lead to flat loss regions that are hard to escape. On the other side, the loss sharpens as the student encodes longer memories because of the curse of memory. As a consequence, gradient-based optimization is extremely tedious, already in this simple example.
4.2 Diagonal connectivity simplifies optimization
We now move to the general case in which the teacher evolves according to
h++1 Ab, + Bra+1 and yt Clh, + DEt-
(10)
with he E RR, 3z E R drawn i.i.d. from N( 1), A E Raxx, Here both inputs and outputs are scalars.
E and D E R
Given the intuition we have developed far, we expect fully connected linear recurrent neural networks to struggle solving the task when the teacher encodes longer memories, not only because of exploding gradients but also due the curse of memory. Conversely, diagonality facilitates the eigenvalue reparametrizzaioon needed to avoid exploding gradients and make them better behaved. We run the following experiment to verify this intuition. We draw random teachers with hidden dimension = 10 and transform the complex eigenvalues of the recurrent matrix A to have magnitudes close to value 7/ that we control'. The larger Vo is, the longer the memories encoded by the teacher are. We train a linear RNN, as well as an LRU [20], with hidden dimension 64 on this task. The students are therefore largely overparametrized. We chose the LRU architecture to represent deep SSMs due to its simplicity. This architecture uses input normalization and an exponential reparametrization of the eigenvalues, similar to what we analyze in Section 3. Both networks are trained using the Adam optimizer [53] and cosine annealing schedule for 10k steps, on batches of size 128. To ensure that we are the infinite sequence length regime, we take the sequences to of length 300, that is three times longer than the characteristic time scale of the teacher. Learning rates are tuned separately for each method and training distribution. The results, which we plot in Figure 3.A, confirm our intuition: LRUs significantly outperform linear RNNs when long memories have to be learned, despite having 10 times fewer parameters.
Next, we wonder which design choices behind the LRU architecture are crucial to this performance improvement. To this end, we interpolate between a linear RNN and an LRU in the following way: First, we restrict the weight matrix of the linear RNN to a block diagonal with blocks of size 2. Each block can represent complex number, so the network can represent 32 complex numbers in total. We additionally double the number of hidden neurons. Second, we change those 2 x 2 blocks (and their input and output weights) to be complex numbers. Finally, we add the y input normalization and the exponential parametrization to obtain the final LRU architecture. We report the results of this experiment in Figure 3.B. Surprisingly, we find the gap comes from making the weight matrix block diagonal (4 x 4 blocks are here enough, cf. Figure 11 in the Appendix). Interestingly, this change reduces the number of parameters the model has and slightly reduces the model expressivity. An explanation of this behavior is therefore likely to be related to the optimization properties of those models. We confirm this hypothesis in the next section.
We draw each entry of A from (0 1/vn), complex diagonalize it, and apply the transformation + Mo + (1 vo) tanh(z) to the absolute values of the eigenvalues.
7

---

A
B
C
D
B
A
1
8
8M
C
C
10*
A
10
10a
10
Hessian
10
8m
10
-10
10-16
-10-
-10a
-10+
-10
Top 10
10-2
Erfectivve
10
10"
107
10a
10a
10
10-7
10
20
eigenvalue
eigenvalue
Figure 4: Differences in learning abilities between fully connected and complex diagonal linear RNNs are due to a better structure the loss landscape. A, B. Hessian of the loss at optimality, its 10 eigenvectors with greatest eigenvalues and its eigenspectra for a fully connected RNN (A) and complex diagonal one (B). The spectra are almost the same. However, the top eigenvectors are concentrated on few coordinates for the complex diagonal one but not for the fully connected one. C, D. This structure makes it possible for Adam to efficiently deal with the extra sensitivity, as shown with the effective learning rates that it uses at the end of learning. For the fully connected one (C), Adam uses small learning rates to compensate for the sensitivity, whereas can use larger ones for the complex diagonal one without hindering training stability. The horizontal grey line shows the learning rate used, which here 10 33
4.3 On the importance of adaptive learning rates
So far, our results highlight the importance of having a close to diagonal recurrent connectivity matrix. In this section, show that this parametrization alone does not mitigate any exploding behavior but modifies the structure of the loss landscape, making it possible for optimizers with adaptive learning rates compensate for these behaviors.
To demonstrate this, we consider the Hessian of the loss:
d² d@2
Σ E dh, 8 T,, dh,
aL, d2 de dh? de ah, d02
(11)
the network can perfectly fit the target data, which the case in experiments above, the second term vanishes at optimality. We plot the Hessian at optimality in Figure 4.A and B for standard linear recurrent network and one with complex diagonal parametrization, both with hidden neurons (Vo = 0.99). We observe that the eigenvalue spectra are similar for the two architectures, both exhibiting large terms that are characteristic of the curse memory, which makes learning with stochastic gradient descent almost impossible". However, their structures differ. For the fully connected linear RNN, the top eigenvectors are distributed over many coordinates, whereas they concentrated few coordinates for the complex diagonal optimization [e.g.. 57]: adapting to large curvature is much easier for Adam when the pathological directions are aligned to the canonical basis. This is what observe in practice.
This feature aids adaptive
In Figure 4.C and D, we compare the effective learning rate used by Adam, which we compute by providing vector of ones to the optimizer. For the dense linear RNN, the adaptive learning rates cannot compensate for the intricate coupling between components, resulting in small learning rates overall. Conversely, the sensitive directions of complex diagonal RNNs are concentrated on few
'The gradient Lipschitz constant L of the loss equals the maximum Hessian eigenvalue [54]. This quantity sets a bound 2/L for the maximum globally stable learning rate. While convergence might happen in subspace, it is generally aligned with the top Hessian eigenspace near the solution [55, 56].

---

A
105
cRNN
B
10a
cRNN
C
2nd mament
10
lay. 1
lay. 4
2nd moment 4
10a
LRU
A
rest
2nd moment det
10a
10a
na norm
A
rest
10*
LRU
lay. 1
lay. 4
10-1
10-1
magn. rest
angle
10-1
10-3
narm
layer
GRU
GRU
rest
0
0.9
0.99
lay.
0
0.9
0.99
GRU
0.9
0.99
lay. 4
Figure 5: Signal propagation in deep recurrent networks at initialization is consistent with our theory. A. E|, after the first and the fourth layer, function of the exponential decay parameter Vn, for complex-valued diagonal RNN (cRNN), LRU, and GRU recurrent layers. The input normalization present in the LRU and in the GRU effectively keeps neural activity constant across Va values. B. Comparison of the evolution of the loss gradient E ) for the different recurrent layers and specific groups of parameters. For the complex diagonal RNN, the gradients of all parameters explode and in particular the ones of the recurrent parameters, whereas only the ones of the angle of À explode for the LRU, consistently with the theory. Error signal propagation in GRUs is under control: the magnitude of the gradients is independent of vo. The GRU-specific parameters exhibit smaller gradients than the feedforward (ff) ones. C. Layer normalization keeps the overall gradient magnitude under control in cRNNs. Batch normalization yields similar results.
parameters, which adaptive learning rates can compensate for. This leads to more targeted and overall larger learning rates, significantly speeding up learning. As side note, the complex eigenvalues of the teacher come in conjugate pairs. However, during training, the complex values of the complex RNN are not conjugates of each other, thereby increasing Hessian diagonality. Finally, performing this analysis for the LRU, we find that the Hessian spectrum is similar to the diagonal setting and that the exploding dimensions of the Hessian are almost exclusively due to the angle parameter, consistently with our theoretical analysis; see Figure 9.A and C. The loss landscape for S4 model [17] can be qualitatively similar to the complex diagonal RNN or to the LRU, depending on which regime it is in; see Figure 9.B and D.
Before concluding this section, we investigate whether certain eigenvalue distributions can break the diagonal structure of the Hessian, thereby complicating optimization and increasing the need for eigenvalue reparametrization.
Appendix C.2, we provide theoretical quantification of the intuitive result that more concentrated eigenvalues lead to less diagonal Hessian. Consequently, the performance gap between complex-valee diagonal networks and LRUs widens, although the former still greatly outperform their fully- colleced counterpart (see Figure 10). An important corollary is that increasing the number of hidden neurons breaks the diagonal structure of the loss landscape, thus reducing the effectiveness of optimizers with adaptive learning rates in mitigating the curse memory. This observation may explain why Orvieto et al. [20] reported a more substantial performance improvement from eigenvalue reparametrization than what we observe in our study (cf. block diagonal vs. LRU in Figure 3.B).
5 Signal propagation in deep recurrent networks at initialization
The ultimate goal of our theoretical quest is to gain insights into the training of practical recurrent network architectures. Specifically, we aim to verify whether the trends established theoretically and in controlled experiments hold in practice, by studying signal propagation at initialization realistic next-token prediction natural language processing task.
To that matter, provide sentences as input deep recurrent networks that contain four blocks and use a next-token prediction loss to measure their performance. Each block consists of a recurrent layer followed by feedforward gated linear unit [58]. By default, there are normalization layers in this architecture. More details can be found in Appendix D. We empirically study how E[h?] and [) evolve when the characteristic time scale of the recurrent layers, controlled through Vo- increases. We compare three different recurrent layers: a complex-valle
diagonal RNN (cRNN),
LRU and a GRU initialized with the Chrono initialization [29].
The results are consistent with our theory. Complex-valued RNNs suffer from the curse of memory and recurrent parameters grow faster to infinity than the rest as v/ goes to 1. Perhaps more surprisingly.

---

finer grain analysis reveals that the gradient magnitude is independent of the layer; see Figure 13. LRUs almost perfectly mitigate exploding behaviors in the forward pass (Figure 5.A) well as in the backward pass (Figure 5.B), except for the angle parameter, consistently with our previous analysis. We also wonder whether layer normalization can replace the input normalization and reparametrization of the LRU. We find that it mitigates the memory- memory-induced gradient explosion at the macroscopic level (Figure 5.C., but it likely kills any learning signal for the smallest eigenvalues [20]. Finally, the GRU manages to keep the gradient magnitude constant over different characteristic time constants, consistently with the intuition we developed in Section 3.2. Preliminary experiments revealed that same trends also hold for LSTMs.
The results presented above for the GRU align qualitatively with the intuition developed throughout the paper. We now consider how well our theory can quantitatively explain this behavior. The primary difference between our simple model and GRUs that the À values, referred to as forget gates in the GRU terminology, depend on both inputs and hidden states, and are therefore not constant. Interestingly, we find that GRUs almost behave like the diagonal linear RNNs we have focused on in this paper, particularly for slowly decaying recurrent neurons with high vo values (see Appendix D.2). Consequently, applying our theory as if this context-dependencc does not exist only introduces minor approximation errors, which we confirm empirically in Appendix D.3. Given the similarity of the Chrono initialization to those used modern architectures like Mamba [21] and Hawk [22], we expect our theory to also serve good proxy for studying signal propagation in these models at initialization.
Conclusion
Vanishing and exploding gradients complicate the learning of recurrent networks, but solving these problems is not enough. We uncovered yet another difficulty of training such networks, which is rooted in their iterative nature and arises at the edge of dynamical stability. Reparamettizations and adaptive learning rates can effectively mitigate this behavior in practice, and diagonalizing the recurrence simplifies both. Our analysis additionally reveals the complexity of learning the angle of complex eigenvalues, which may explain why complex numbers were not found to be useful in most recent state-space model architectures [21, 22].
A side finding of our study is the symbiosis between independent modules, which are here neurons and be more generally small heads, with adaptive learning rate optimizers in linear recurrent networks. Such a design pattern has promising properties: it facilitates online learning [59] and compositional generalization [60], allows for high level of parallelization [22], and matches, high level, the modular organization of the cortex in cortical columns [61]. Understanding how to increase the expressivity of small linear modules while keeping their great optimization properties constitutes a promising avenue for future research.
Limitations
The theory we introduced, with focus on signal propagation, only addresses the training dynamics of recurrent neural networks. Consequently, it does not provide insights into other important questions such as generalization abilities or memory capacities of these networks. The main assumption underlying this analysis is that the recurrence is both diagonal and linear. While this approach offers valuable insights, it can only approximate signal propagation in more sophisticated architectures. Given the simplicity of the analytical tools employed, there is little hope that this framework can be extended to more general settings without sigmificant modifications.
Acknowledgments
The authors thank Robert Meier, João Sacramento, Guillaume Lajoie, Ezekiel Williams, Razvan Pascanu, Imanol Schlag and Bobby He for insightful discussions. Nicolas Zucchet was supported by an ETH Research Grant (ETH-23 21-1) and Antonio Orvieto acknowledges the financial support of the Hector Foundation.
10

---

References
[1] David E Rumelhart, Paul Smolensky, James L McClelland, and G Hinton. Sequential thought processes in PDP models. Parallel distributed processing: explorations in the microstructures of cognition. 2, 1986.
[2] Jeffrey Elman. Finding structure in time. Cognitive science, 14(2). 1990.
[3] Sepp Hochreiter and Jürgen Schmidhuber. Long Short-Term Memory. Neural Computation. 9(8), 1997.
[4] Ilya Sutskever, Oriol Vinyals, and Quoe Le. Sequence to sequence learning with neural networks. In Advances in Neural Information Processing Systems, 2014.
[5] Sepp Hachreiter. Untersuchungen zu dynamischen neuronalen Netzen. Diploma thesis, Technische Universität München, 1991.
[6] Yoshua Bengio. Patrice Simard, and Paolo Frasconi. Learning long-term dependencies with gradient descent is difficult. IEEE transactions on neural networks, 5(2), 1994.
Sepp Hochreiter, Yoshua Bengio, Paolo Frasconi, and Jurgen Schmidhuber. Gradient flow in recurrent nets: the difficulty of learning long-term dependencies. In A field guide dynamical recurrent networks. IEEE, 2001.
[8] Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio. On the difficulty of training recurrent neural networks. In International Conference on Machine Learning, 2013.
[9] Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. In International Conference Learning Representations, 2015.
[10] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances Neural Information Processing Systems, 2017.
[11] Quentin Fournier, Gaétan Marceau Caron, and Daniel Aloise. A practical survey on faster and lighter transformers. ACM Computing Surveys, 55( 14s), 2023.
[12] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Flcuret. Transformers are RNNs: fast autoregressive Transformers with linear attention. In International Conference Machine Learning, 2020.
[13] Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention; fast and memory- efficient exact attention with 10-awareness. arXiv preprint arXiv:2205. 14135, 2022.
[14] William Fedus, Barret Zoph, and Noam Shazeer. Switch Transformers: scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 2022.
[15] Shuming Ma, Hongyu Wang, Lingxiao Ma, Lei Wang, Wenhui Wang, Shaohan Huang, Li Dong, Ruiping Wang, Jilong Xue, and Furu Wei. The era of 1-bit LLMs: all Large Language Models are in 1.58 bits. arKiv preprint 2rXiv:2402.17766 2024.
[16] Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, and Christopher Ré. Combining recurrent, convolutional, and continuous-time models with linear state-space layers. In Advances in Neural Information Processing Systems, 2021.
[17] Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. In International Conference on Learning Representations, 2022.
[18] Ankit Gupta, Albert Gu, and Jonathan Berant. Diagonal state spaces are effective structured states spaces. In Advances in Neural Information Processing Systems, volume 35, 2022.
[19] Jimmy T.H. Smith, Andrew Warrington, and Scott W. Linderman. Simplified state space layers for sequence modeling. In International Conference on Learning Representations, 2023.
[20] Antonio Orvieto, Samuel Smith, Albert Gu. Anushan Fernando, Çağlar Gülççhree Razvan Pascanu, and Soham De. Resurrecting recurrent neural networks for long sequences. In International Conference on Machine Learning, 2023.
[21] Albert Gu and Tri Dao. Mamba: time sequence modeling with selective state spaces. arXiv preprint arXiv:2312. 00752 2023.
11

---

[22] Soham De. Samuel L. Smith, Anushan Fernando, Aleksandar Botev, George Cristian-Muaau, Albert Gu, Ruba Haroun, Leonard Berrada, Yutian Chen, Srivatsan Srinivasan, Guillaume Desjardins, Arnaud Doucet, David Budden, Yee Whye Teh, Razvan Pascanu, Nando Freitas, and Caglar Gulcehre. Griffin: mixing gated linear recurrences with local attention for efficient language models. arXiv preprint arXiv:2402211442,, 2024.
[23] Kyunghyun Cho, Bart van Merrienboer, Dzmitry Bahdanau, and Yoshua Bengio. On the properties of neural machine translation: encoder-decoder approaches. In Proceedings of SSS7-8, Eighth Workshop on Syntax, Semantics and Structure in Statistical Translation, 2014.
[24] Tomas Mikolov. Statistical language models based on neural networks. PhD thesis, Brn University of Technology. 2012.
[25] Sergey loffe and Christian Szegedy. Batch normalization: accelerating deep network training by reducing internal covariate shift. In International Conference on Machine Learning, 2015.
[26] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer normalization. In Neural Information Processing Systems- Deep Learning Symposium, 2016.
[27] Tim Cooijmans, Nicolas Ballas, César Laurent, Çağlar Gülgehre, and Aaron Courville. Recurrent batch normalization. In International Conference Learning Representations, 2017.
[28] Quoc V. Le. Navdeep Jaitly, and Geoffrey E. Hinton. A simple way to initialize recurrent networks of rectified linear units. arXiv preprint 1004v.150400044,, 2015. 2015.
[29] Corentin Tallec and Yann Ollivier. Can recurrent neural networks warp time? In International Conference Learning Representations, 2018
[30] Salah EI Hihi and Yoshua Bengio. Hierarchical recurrent neural networks for long-term dependencies. In Neural Information Processing Systems, 1995.
[31] Asier Mujika, Florian Meier, and Angelika Steger. Fast-slow recurrent neural networks. In Neural Information Processing Systems, 2017.
[32] Martin Arjovsky, Amar Shah, and Yoshua Bengio. Unitary evolution recurrent neural networks. In International Conference on Machine Learning, 2016.
[33] Eugene Vorontsov, Chiheb Trabelsi, Samuel Kadoury, and Chris Pal. On orthogonality and learning recurrent networks with long term dependencies. In International Conference on Machine Learning, 2017.
[34] Kyle Helfrich, Devin Willmott, and Qiang Ye. Orthogonal recurrent neural networks with scaled Cayley transform. In International Conference Machine Learning, 2018
[35] T. Konstantin Rusch and Siddhartha Mishra. Coupled oscillatory recurrent neural network (CoR an accurate and (gradient) stable architecture for learning long time dependencies. In International Conference on Learning Representations, 2021.
[36] T Anderson Keller and Max Welling. Neural wave machines: learning spatiotemporally structured representations with locally coupled oscillatory recurrent neural networks. In International Conference on Machine Learning, 2023.
[37] Memming Park, Ábel Ságodi, and Piotr Aleksander Sokół. Persistent learning signals and working memory without continuous attractors. arXiv preprint arXiv:2308. 12585, 2023.
[38] Zhong Li, Jiequn Han, Weinan E, and Qianxiao Li. Approximation and optimization theory for lincar continuous-ting time recurrent neural networks. Journal of Machine Learning Research, 23(42), 2022.
[39] Shida Wang. Zhong Li, and Qianxiao Li. Inverse approximation theory for nonlinear recurrent neural networks. In International Conference on Learning Representations, 2024.
[40] James Martens and Ilya Sutskever. Learning recurrent neural networks with hessian-free optimization. In International Conference on Machine Learning. 2011.
[41] Lorenzo Noci. Alexandru Meterez, Thomas Hofmann, and Antonio Orvieto. Why do learning rates transfer? Reconciling optimization and scaling limits for deep learning. arXiv preprint 2707 2024.
[42] Greg Yang and Edward J Hu. Feature learning in infinite-width neural networks. arXiv preprint arXiv:2011.1452.220 2020.
12

---

[43] Greg Yang. James B Simon, and Jeremy Bernstein. A spectral condition for feature learning. arXiv preprint 1rXiv:2310.17813, 2023.
[44] Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. Artificial Intelligence and Statistics, 2010.
[45] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification. In International Conference on Computer Vision, 2015.
[46] Grigorios A. Pavliotis. Stochastic processes and applications, volume 60 of Texts in Applied Mathematics. 2014.
[47] Greg Yang, Edward Hu, Igor Babuschkin, Szymon Sidor, Xiaodong I Liu, David Farhi, Nick Ryder, Jakuh Pachocki, Weizhu Chen, and Jianfeng Gao. Tensor Programs V: Tuning Large Neural Networks via Zero-Shol Hyperparameter Transfer. arXiv preprint arXiv:2203.03466 2022.
[48] S. Boyd and L. Chua. Fading memory and the problem of approximating nonlinear operators with Voltera series. IEEE Transactions on Circuits and Systems, 32(11), 1985.
[49] Antonio Orvieto, Soham De. Çaglar Gülçehre, Razvan Pascanu. and Samuel Smith. Universality of linear recurrences followed by non-linear projections: unite-width guarantees and benefits of complex eigenvalues. In International Conference on Machine Learning. 2024.
[50] William Merrill, Jackson Petty, and Ashish Sabharwal. The illusion of state in state-space models. In International Conference on Muchine Learning, 2024.
[51] Minmin Chen, Jeffrey Pennington, and Samuel Schoenholz. Dynamical isometry and mean field theory of RNNs: gating enables signal propagation in recurrent neural networks. 2018.
[52] Moritz Hardt, Tengyu Ma, and Benjamin Recht. Gradient descent learns linear dynamical systems. Journal of Machine Leurning Research, 19(29). 2018.
[53] Diederik P. Kingma and Jimmy Ba. Adam: method for stochastic optimization. In International Conference on Learning Representations, 2015.
[54] Yurii Nesterov and others. Lectures on convex optimization, volume 137, Springer, 2018.
[55] Jeremy Cohen, Simran Kaur, Yuanzhi Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. In International Conference on Learning Representations, 2020.
[56] Guy Gur-Ari, Daniel A Roberts, and Fthan Dyer. Gradient descent happens in tiny subspace. arXiv preprint arXiv: 2018.
[57] Yan Pan and Yuanzhi Li. Toward understanding why Adam converges faster than SGD for transformers. arXiv preprint arxivv2306.000244 2023.
[58] Yann N. Dauphin, Angela Fan, Michael Auli, and David Grangier. Language modeling with gated convolutional networks. In International Conference on Machine Learning, 2017.
[59] Nicolas Zucchet, Robert Meier, Simon Schug, Asier Mujika, and João Sacramento. Online learning of long-range dependencies. In Advances in Neural Information Processing Systems, 2023.
[60] Anirudh Goyal, Alex Lamb, Jordan Hoffmann, Shagun Sodhani, Sergey Levine, Yoshua Bengio, and Bernhard Schölkopf. Recurrent independent mechanisms. In International Conference on Learning Representations, 2021.
[61] V B Mountcastle. The columnar organization of the neocortex. Brain, 120(4), 1997.
[62] Christoph Boeddeker, Patrick Hanebrink, Lukas Drude, Jahn Heymann, and Reinhold Hacb-Umbar On the computation of complex-valued gradients with application to statistically optimum bcamforming. arXiv preprint Xiv: 00030009,, 2017.
[63] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Midderman--iinee and Qiao Zhang. JAX: compos- able transformations of Python+NumPy programs, 2018. URL http: //g http://gtttuuuuoooooeoeeee/ jon//oogle//ax.
13

---

[64] Jonathan Heek, Anselm Levskaya, Avital Oliver, Marvin Ritter. Bertrand Rondepierre, Andreas Steiner. and Marc Zee. Flax: neural network library and ecosystem for JAX, 2023. URL http://github. com/ ham//google/ffaan
[65] Nicolas Zucchet, Robert Meier, and Simon Schug. Minimal LRU, 2023. URL SiRLtttss//gAttub..com NicolaasZZccchheemmmmmmaaaaa-L
[66]
Wikimedia Foundation. Wikimedia Downloads. URL atttpa://dump.. wikimedia. org.
[67] Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. In Jill Burstein, Christy Doran, and Thamar Solorio, editors, Proceedings of the 2019 Conference of North American Chapter of the Association for Computational Linguistics: Human Language Techmologies, Volume (Long and Short Papers), 2019.
14

---

Appendix
Table of contents
Definition of recurrent networks use
A.1 Linear recurrent neural network
16
A.2 Complex-valued RNN and linear recurrent unit (LRU)
16
A.3 S4
16
A.4 GRU
17
B
Theory
18
Useful lemmas
................................
The curse memory: signal propagation analysis
................................
B.2.1
Forward pass
19
B.2.2
Backward pass
19
B.2.3
Extension to fully- connected networks
21
B.2.4
On the wide-sense stationarity assumption
22
B.3
Impact of input normalization and parametrization
22
B.3.1
Real case
22
B.3.2
On difficulty parametrizing complex numbers
23
C Linear teacher-student task
24
C.I
ID setting
24
C.I.1
Calculation of the loss
24
C.1.2
Optimal normalization and reparamctrization with uncorrelated inputs
24
C.1.3
Visualization the effect of input normalization and reparamctrization
26
C.1.4
Learning angle difficult practice: an example
C.2
Structure of the Hessian optimality
27
C.2.1 Hessian for complex-valued variables
28
C.2.2 Hessian with respect the recurrent eigenvalues
29
C.2.3 Hessian for different parametrizations
30
C.3
Experimental details
32
Additional analyses
32
C.4.1 Structure of the loss landscape for IRRUs and
32
C.4.2 Concentrating cigenvalue distributions
34
C.4.3 Impact of the number of heads in fully connected linear recurrent networks
35
D Signal propagation in randomly initialized deep recurrent neural networks
36
D.1 Experimental setup
36
D.2 Can gated RNNs he considered diagonal?
37
D.3 Does our theory apply to gated RNNs?
39
15

---

A Definition of the recurrent networks we use
In this section, we rigorously define all the architectures use in the main text and in the appendix, as well as precisely describe how we initialize them.
A.1 Linear recurrent neural network
Let us start by introducing the linear recurrent neural network we are using. It satisfies
ho h141 Aht + Brt+1
(12)
(13)
Y1 Chi + Dr
(14)
with 2; E [Rd", h4 E 1R" Y/ Rdout, A E RTx", B Raxdi,, C Rd and D Rd xdi.. We draw the entries of B, C and D from independent truncated normal distributions with fan_in in scaling. We draw each entry of A from AV 1/vn) independently and then apply the following posiprocessing to it: First we complex diagonalize A, which we can do almost surely. Note its eigenvalues. We then transform them according to
(aagle(A) 60
+ (vo + (1 vo) tanh( exp
(15)
π
with Hy and 9g two scalars that we control. This transformation has several benefits: we are guaranteed that the magnitude of X is within (vo 1] (and in (vo, Vo + (1 (1)))) tanh(1)) in the limit n - co the eigenvalues of A stay within the unit circle in that limit), and conjugate pairs of eigenvalues remain conjugate. This last point ensures that the resulting matrix remains real without having to change the eigenvectors. When we split the connectivity matrix A into several independent blocks, we initialize each block separately with the scheme described above.
A.2 Complex-valued RNN and linear recurrent unit (LRU)
Both architectures have recurrence of the form
ho 0 h++1 = x h4 + @ Br 111 Re(Ch, + Dx..]
(16)
(17)
(18)
with C the element-wise product, It. E Rdin, h1 E C", Yt Rdout C", B Coxdm 7 E R",
C' Cdoen and D E R Rddot Xdin.
For the complex-valle valued linear RNN, we take
A= re +ixim 7 1=1
(19)
(20)
so that it can be considered parametrizing the diagonalized version of the linear RNN,
For the LRU [20], take
À exp((exx(wx) exp(( exp(wo))
(21)
7 = -x
(22)
For both architectures, we use the LRU initialization that À uniformly distributed on the ring between the circles of radii vo and 1, with absolute angle restricted to be below 0p. For the LRU, we initialize to /1- [X
A.3 S4
We consider slightly modified version of S4 here:
A = softplus(ws)
(23)
h+11 exp ( (wii + ( o h4 + A o Br
(24)
Re(Ch, + Dz.]
(25)
16

---

with E RAti he E C", E RRou,, - E C", B E Ca WA E R", C'E C xn and D E Rdout Xdii The main difference with the standard architecture is that we do not consider any state expansion so that there is no parameter sharing. The makes this architecture closer the linear RNN architecture mostly focus on in this paper, while keeping the kind of parametrization it uses. In the same spirit, we do not couple the input and the readout matrices in any way.
The initialization we use also differs from existing ones as we initialize to and exp (A i)) in the same way in the LRU.
A.4
GRU
The GRU version we use the following:
T++1 o (W F++1 + Wrihe + b,)
(26)
fi+1 o(W + W Wyhh + by)
(27)
B++1 tanh W N111 + Unx r @ (Wanh, bnh))
(28)
h++1 = fi+1 @ he + (1 fet1) @ n4+1
(29)
(30)
with a the sigmoid function, h4 E RR, T2 E I and all the other matrices appropriately sized.
We initialize the parameters with the Chrono initialization [29]: all parameters initialized using standard practice (orthogonal initialization for the weights taking as input, Lecun initialization for the rest, biases initialized at 0) except for by, which initialized
by log (U (Tmin: Tmax))
(31)
Here, T Tii and T max denote the minimum and maximal characteristic time scale. In the experiments of Section 5 we take
Tiin 1
(32)
V/D 2
T Fmax
(33)
1/0
to enable the comparison with other architectures. Indeed, when ignoring the dependence of the forget gate f on the input æ and on the hidden state h. this corresponds having f E [Vo. 2
1+
17

---

B Theory
This section introduces all the theoretical results we directly or indirectly mention the main text, as well as provides proof for them.
B.1
Useful lemmas
Most, all the calculations, that we will be doing in this section involves infinite sums. We state and prove two useful lemmas to simplify later calculations.
Lemma 1. For , B E C satisfying |oo| < and [3] < 1, and (un)nez a bounded sequence satisfying
u- Uns have
Σ m,
a Un
uo
Σ A>1
(a
3
u
(34)
aB
Proof. The proof naturally comes from separating the indices n and m in three sets: one in which the are equals, one in which n larger and one in which m is larger. This gives
Σ
a a"3 9 un
Σ 3m
m
Σ
8mhh Un-m +
Σ a" 8m
Un-m
(35)
2,
n=m
m>
ncm
Σ
n 3 un
+ Σ
a
m gm
us + Σ
a 3
3 U-A
(36)
Σ
a g"
uo +
Σ
(a + )
(37)
n
A21
Ug + Σ
(a + 3
(38)
aBa
-
A21
Lemma 2. In the same conditions as Lemma 1, we have
d d
Σ nmma" 11 gmh Un da d, 1- ag
+
Σ
(a + BAJuS))
(39)
A21
Proof. This follows from remarking that
d da dB
d
Σ am am un
Σ
max 8 Un
(40)
do
n.m20
nma
gm
Un-m
(41)
am
and using Lemma I to get the final result.
B.2 The curse of memory: signal propagation analysis
We recall the assumptions that we stated in Section 2.2:
a) Linear diagonal recurrent neural networks. We restrict ourselves to networks satisfying h++1 X @ h, + Iz+1 with X, h, and z, complex numbers. Without loss of generality, we focus on the one dimensional setting. We additionally consider As with absolute values smaller than 1. Infinite time horizon. We consider infinite sequences and initialize the network dynamics at to -00.
18

---

Wide-sense stationarity. We assume the different quantities that the network receives, which includes the inputs 2,, to be wise-sense stationary (WSS). A random process X, said to be WSS if its auto-correlation function is independent of time, that is, for all i. E IR and A E R, [ == Rx(A).
B.2.1 Forward pass
Without loss of generality, we can take t = 0 given the wide-sense stationarity and infinite time horizon assumptions. Let us first remark that we have
ho = Σ x-xaa
(42)
n20
so that
hh =
Σ
X"XTE B [- p
(43)
n0
Σ n,m20
X" Re m)
(44)
T- R2(0) + Σ + R2) A21
(45)
We used Lemma to obtain the last equality. In Section 2.2, we focused on the real case X, so this formula becomes Equation 5. If we further assume that the cuto-correlation of z decreases exponentially with decay rate p. that is R2 (A) ااام, we can further simplify the last expression:
E hohol 1 =
1+ Σ
(1- +
(46)
- 1A/2
A21
1
Xp
Ap
1+
+
(47)
Xp
1 Ap
p 1112
(48)
1 pA (1 11
follows that if the inputs are i.i.d. (p=0) 0), we have E Ei 21 =
(AP) and if the inputs are
constant equal to 1 (p 1), have E[[ho [=]
B.2.2 Backward pass
Differentiating the update h4+1 = Al, + F4+1 with respect to gives
dhut1
= dhe
+ he
(49)
dx
dA
so that
dho
Σ X"han- n-1
(50)
yp
x x-n-m-1
(51)
n20
m20
=
Σ
r (x+m+1)
(52)
n,m20
nA"
nX
r 1-
(53)
n20
Note that some extra technicalities are needed to justify these equations as À and hy are complex valued: these formulas hold as they would in the real-valued case as by is an holomorphic function of A.
19

---

We can now compute the variance of the sensitivity of the hidden state with respect to the parameters.
dhp
2
E
mmA"-1imm- Rn,(n - -m).
(54)
Yd
n20 m20
Using Lemma gives
E
dhg
27
d
d
1
Rr( + (a (0AA
(55)
dd
da
di
- 1
ag
A21
a=
Differentiating this quantity product gives
E
dhp
1 + aB
R₂(0)+ Σ
(a
+ (A) +0
dx
a3)3
A21
Q
Σ
Aa Ra()
+
B
Σ
CBs--1 R2(4)
(56)
oB)²
(1 (1-aft ag)²
A21
A21
a=1,==
which then simplifies as
dhe
1+ 1212
E
(o) + Ea0+ +X*)RA(A)
dd
A
1
4(1- - + (57)
Σ
(1 AP-
A21
Note that Equation 6 in the main text the real-valued version of that formula.
Let us now further simplify this equation when Rs (A) ااامم If we use this in the differentiated quantity before differentiating it, we get
dhp dx
d d da d3
1 1
E
(58)
a- (1 (1-pox((1 p3) ==
Calculating this quantity manually is painful. Instead, is rather easy to compute, it is equal to (1 (1 pa) (1- p3)². We thus multiply it to the derivative of the function we want to compute in order to obtain polynomial with unknown factors, and use polynomial regression tools to derive the resulting coefficients. Massaging the obtained expression to make easier to compute the closed-form value of this quantity when p and p we get
the following trick. Its denominator
dh
(1 (1 + (AR) + p² 11/3)3 + p(1 - - + +1/)) - 2A 21)
E
dd
(1 - [A|3)3|1 pAd
(59)
This is the quantity we plot on Figure 1.A, when À is real-valued. When 0, this quantity becomes
dh
+ [A]2
E
(60)
dx
MP3s
and it is equal to
dhe
27
(61)
1 - 214
when p 1. Additionally, it will diverge whenever -=
A 1 when p 1, and when À + when
Regarding the backpropagation of emror to the inputs, the analysis we did in the main text also holds for complex number given that h2 is an holomorphic function of s; and it thus behaves as the forward pass once replacing the input distribution with the one of output errors O Lt.
20

---

B.2.3 Extension to fully-connected networks
We now turn to the non-diagonal case. For the sake of simplicity, we assume that recurrent matrix is complex diagonalizable and that its eigenvalues are all different. This will enable us differentiate the eigenvalues and the eigenvectors. We consider dynamics of the form
h++1 = Ah; + It+1
(62)
As A is complex diagonalizable, there exists cemplex--valuee matrix P and complex-valued vector such that
= Pdiag(A)P
(63)
P.P. = Vi.
(64)
The linear recurrent neural network considered above is equivalent its diagonal version
diag Ahς diaag + p-pppam P-12441 Ø141
(65)
h₂ Phi
(66)
We now differentiate h4 w.r.t. to using the diagonal parametrization and obtain
dh, The OP
Th, dndiis ax The dh, diag Op-1
(67)
dA aP 0A
Th diag dA
aA
Thg Ohdiag diag dP-1
0A
d d, and A can be considered constant. Intuitively, the eigenvalues and eigenvectors move smoothly as we restricted ourselves to the case in which eigenvalues are singular. If this is the case, math becomes trickier the eigenvectors are not uniquely defined. We can study the behavior of those quantities in more detail, following Boeddeker et al. [62]]
0x a
OA
diag
P
(68)
Aii
dP
0A DAi,"
P))
P
FO
p-1
P.)
(69)
dAi
The F' introduced in the last equation is equal to
if otherwisc.
Fgj
(70)
Importantly, those two quantities do not grow to infinity as the absolute value of the eigenvalues goes to 1, which means that we can consider those derivatives to be independent of A for the sake of our analysis. Note that the previous argument assumes that eigenvalues do not collapse.
d d ding the dominating term in da We wonder which of the three different terms that appear in dy (Equation 67) will be the dominating one as |X| (or A) goes 1. In the previous paragraph. we have shown that the derivative of p-1, P and À can be considered constant for the sake of our analysis. We thus focus on the other terms.
First, we have
Thu.1 OP
j=l
(71)
so the magnitude of this quantity is roughly the one of diag which corresponds to the low pass filtering of the inputs with different values.
Second, we know that θ Diiinn hy does not change in magnitude as X changes, as P remains bounded.
So, for the third term of the sum, we are left to study the behavior of d diag We can show that it evolves according to
dh dhding
+ J141
if =
(72)
dP-
dP.
diag
0 otherwise.
(73)
d
21

---

It follows that the third term in the sum also corresponds to a low pass filtering of the inputs.
Finally, we know that the second term, the one in dy dias will grow faster to infinity as corresponds two consecutive low pass filters with the same values (c.f. calculation above). It will thus be the dominating term in the infinite memory limit.
B.2.4 On the wide-sense stationarity assumption
In our analysis, we make the assumption that the different quantities given to the network are wide- temporal shift. In practice, this assumption
sense stationary, that is their statistics are invariant will likely never be satisfied, as sequences are finite and as parts of the sequence (c.g., the beginning of text) can have different statisties than other parts (c.g., the end of a text).
should be noted that the analysis we have done in the wide-sense stationary case can provide an upper bound on the different quantity study. Indeed, if there exists function U such that
[[ S U(n mm))
(74)
minor modifications to our analysis enable upper bound the different quantities we study, replacing RS by U.
In our experiments of Section 5, plug the empirical covariance defined
T-IAI
1
pempirical
Σ
J+J++AA
(75)
T
A
into analytical expressions. It comes with approximation it averages the autocorrelation for all positions, which are not equal when wide-sense stationarity is not met.
B.3 Impact of input normalization and parametrization
In this section, we consider a diagonal linear recurrent neural network of the form
h+11 = X(w)h, + +(A2a+11
(76)
with 7(A) the input normalization factor and À parametrized by a vector w. Next, we study the effect of input normalization and reparametrization, first in the real-valued setting and then in the complex-valued one.
B.3.1 Real case
Let us start with the forward pass: as
1y Σ X")((Aor
(77)
n20
Y rescales the value the hidden state takes. To avoid any explosion behavior, thus ideally want to be the inverse of value of E((h))² without normalization, which we have computed in Equation 45. The same behavior holds for the backpropagation of errors to the inputs
dL dæ,
x(2) (*
dL aL
(78)
dx
Th.
We now move to the impact of the parametrization. To simplify the calculation, will ignore the dependency of on À when differentiating it. This can easily be done in automatic differentiation software by removing this dependency from the computational graph with cistop_gradient(]) We then have
dh, dh dd
(79)
dw yx dw
and dy which is rescaled by 7 compared the calculation we did above. As a consequence, both the input normalization and the parametrization can help to mitigate the curse of memory.
22

---

B.3.2 On the difficulty of parametrizing complex numbers
We now extend the previous analysis to the complex case, and take polar parametrization of A: X(w) v(w) cxp(if(w)). The effect of the input normalization does not change when moving to complex numbers. The role of the reparametrization is however bit more subtle. As h is an holomorphic function of X, we have dx = 0 and
dhe dhy yb dw
dh, dv
exp(i0) + 2 dw
d0
exp(i0))
(80)
dA dw dd
2 dw
follows that
dhg
27
dhg
2
dv
+ iv do dw
E
-
E
exp(10)
exp(i0)
(81)
dw
4
yx
dw
dhe
27
dv
d0 2
4
E
yx
dw + iv
dw
(82)
dh,
2
du
v2 d0 2
E
(83)
dd
dw
dw
To simplify the analysis, we will further assume that dx | is only a function of v. This asumption holds in the case of 0 and p c.f. Section B.2.2, but not necessarily otherwise. To ensure that that this quantity does not depend on X, we thus want d E(v) @(1) and 0²d_ O2_0FF(()) E(v) = (()).
The second means that the angle parametrization must depend the value takes. Let take the = example to get an idea of what the ideal parametrization should be. First, we have ((A) = vI D2 so that
E(v) = rA)} 1 1 v2
1 v²
(84)
(1 (1- v²)3 (1 (1- -2)²
We are left with the differential equation D/ = 0(1 - -2²). which is for example solved with tanh(w,,). Now let look the parametrization of 0. we ignore the v² term for simplicity, the approximate differential equation it needs to solve is d O(1 v²), which can be solved by = stop_gre graddgradennn (1 12) The exact detail of this calculation do not really matter as this heavily input distribution dependent. However, the interesting part here is that the angle parameter must be rescaled by a function of v. This makes intuitive sense when looking looking at the sharpness of the loss around optimality in Figure 1.C, but this also makes the loss even flatter further away from optimality. We will come back to this point in Section C.1.4, showing that in practice, such parametrization complicates the learning of the 0. Learning complex numbers is thus difficulty, because of the angle.
23

---

C Linear teacher-student task
This section is dedicated to detail the theoretical results behind analysis of the teacher-student task, present all the details necessary to reproduce our empirical experiments, and provide additional analysis.
C.1 ID setting
C.1.1
Calculation of the loss
In this toy example, we are interested in learning a simple I-dimensional linear recurrent neural network which follows the dynamics
h4+1 Ah, + Iz41
(85)
reproduce the hidden state h of a teacher with recurrent parameter X Note that we here allow all variables to be complex-valued, We take the loss to be
T(A, 2T
Ry h4
(86)
We assume x to be drawn from a wide-sense stationary distribution so that we can focus on studying h21² to understand the behavior of the full loss L, in
the behavior of one (x) == f [ the limit of infinitely long sequences (? + oo). Moreover, to further simplify the calculations,
assume that is real-valued and that (A) = إ
Let us now proceed with the calculation:
1 B. [hehi hih, h44 k,hj].
L.(A,*)== =
(87)
2
We have shown Section B.2 that in the limit of L + DO.
1 + pd
pt
Ex [halie]
(88)
-
pd
(89)
Similar derivations hold for the other three terms in the loss. Grouping them gives the exact value of the loss. We omit the formula as it is not particularly insightful. In the case of constant inputs (p = 1), we have
1 X
1
L((A,))
(90)
1
X
In the case i..dd inputs (p = 0), we have
1
2
1)
(A. 2 212 1
Re
(91)
|x
1 XA*
This is the loss plot Figure 1.B and C.
C.1.2 Optimal normalization and reparametrization with uncorrelated inputs
Having a simple closed-form solution for the value the loss takes gives us the possibility to investigate in more detail what an optimal normalization and parametrization should be. We focus on the case
p =
For p the optimal normalization is 7()
Given that we add an input normalization to the student, we must also add it to the teacher for the student to be able to fit it. The loss becomes
7())
7(*))
27(A)(21)
L
+ 1 1
Re
(92)
1 1
1-
1
({*)
= -1-Ro
(93)
AA=

---

Next, we parametrize À as À = r(w,) exp (wa)) and seek to find parametrization such that, at optimality, Bi(d_,,,, ) = and [[[d.,, h = Given that the student perfectly fit the teacher- generated data at optimality and that the loss we use is the mean-squared error, this corresponds to having d² T = and d2 J 1. 1.
Deriving the optimal y/ parametrization. We now compute the Hessian of the loss w.r.t. wp. First, we can simplify our calculations by restricting ourselves to the case @ = l". The loss becomes
(())(())
L 1 1 10
(94)
Differentiating this function a first time, we obtain
db = 1 -v0
7(v))((V) 7(1"))²
(95)
dv
(1 vv-)2
Differentiating it second time gives
d'
((1))((1) 2)(v")ur>/(0
2)(v")u+
(96)
du²
1 1-v0" V0
(1 (1-vi) vv²)
(1 - -vv=)* )3
Leveraging the fact that
and )0(v) = -)(1}- 12
7(v)= = 7V))
-V
(97)
71133
we finally get, when v= v",
d²
(98)
dv2 (1 v272
Given that we are at optimality, we have dLL,
dh, d² L. dh,
chi d2 dh,
= v (((,))
1/ (ww,)
(99)
dw3
dww dh dwi.
dv dii dv
duz
To keep that quantity constant, thus have to solve the differential equation
v/ (www))= = - (1-x²),
(100)
which gives v(w.)- tanh(ww)).
Deriving the optimal 0 parametrization.
We now move to the parametrization of 0. We have
L = 1 Re (()h(e)))
=1 7())(00)(1 vv cos(0 0*))
(101)
11 11 1x
For notational convenience, we denote
a(O -
= -(0-0))==11- 1x (1 - vv" cos(0- - 0²)2 + v² sin(O - 0772.
(102)
We have
dl db = (()((())
(- (-u=siin00-0) 0²) (1 - vv* cos 0') (0- -
(0- -0))
(103)
a(l 10000%% 8r)
alB BO2
and
dg2 = ((})((²) ( (- Vv cos(O - 0²) vv sin - 0 )a' (0- - 0*)
1.
+2
all - B*)
all A+) )2
- (1-v" Vv" cos(0- - (0 04)
2 - vv cos(B 0 - 8") )a' (8-8")2)
(104)
all )2
a(O 0r
At optimality (0 = 0* and v = v*), we have (0)) = (1 - 222) a' (0) 0 and a²(0) 2v², so that
d'L,
v2(1+12)
(105)
d02
(1 22²²
25

---

The optimal parametrization thus has to satisfy
1 1 v²
uV + v2 8 (www))
(106)
that is
1 - v22
(wwo) we
(107)
DV1 + v2
There are two things we can remark:
First, the parametrization that derived for the general case in Section B.3.2, which additionally ignored the dependence of y on X, is relatively accurate. The only difference is the apparition of the extra term, which becomes insignificant in the long memory limit v/ - 1.
Second, the optimal 8 parametrization has to be function of v, and thus wy. so the differential equation y needs to satisfy changes. Yet, this considerably simplifies the calculation and there is no simple solution that problem. One could still argue that the initial choice we made, that is to use polar parametrization, is the issue. It could be, but most practical models end up using that choice so highlighting the limitations of this choice has important practical consequences.
In the rest of this section, ignore the dependency of A on and consider the optimal parametriza- tion this setting to be
D(w = tanh( (w.
(108)
((wgPt) = opt v2
(109)
vV + v2
C.1.3 Visualization of the effect of input normalization and reparametrization
We now visualize the effect of input normalization and reparametrization on the loss landscape. We focus on two such reparametrizations:
the one used in the LRU [20, 22] with (AA) = -T- = cxp(- -xxp(cxxp(((x)) and = exp(w,
the optimal one we derived in the previous Section (c.f. Equations 108 and 109), which is tailored
to this specific setting.
C.1.4 Learning the angle is difficult in practice: example
We use this one-dimensional teacher-student setting to test whether having parametrization that avoids exploding behaviors at optimality, such as the one we derived in Section C.1.2, facilitates learning. Figure already hints towards the fact the basin of attraction of the global minima either extremely narrow or that their number decreases longer memories are considered, making learning more tedious. Figure 7 confirms it. In this figure, we plot the learning dynamics obtained using the Adam optimizer with learning rate of 10 for 50k steps, starting from Å 0.99 exp(in/ i/) We consider three different parametrizations of the angle:
(lupolar) w/ polar
(110)
O(w log (og(up))) (xP)
(111)
0(wgrt) (1 v²)
vV1 + v² Wg.
(112)
The first one does not reparametrize the angle, the second one is the one used the LRU and the third one is the optimal one we derived above. We use v/ tanh(w,,) to parametrize the magnitude X" = 0.99 exp(in/110) The landscape when v correct therefore
the three cases. We set corresponds to the ones plotted in the last two columns of Figure 6. This example shows that efforts to reduce the sharpness of the loss at optimality, as done the last parametrization, inevitably make the loss flatter elsewhere and optimization impossible.
26

---

10'
10
10-1
10
Hessian
10
10
Locs
10-3
Hessian
10a
101
10-7
10
10-5
10-1
0.0
0.5
1.0
0.0
0.5
1.0
8
10-1
10
10
10a
10
Hessian
10-1
Locs
10-6
Hessian
10-1
10-7
10
10
10
10
100
10-3
Hessian
10
Loos
10
Hessian
109
10-6
10
10-7
10-2
10
5
-5
0
5
-5
0
AA 0.99
0.98
0.96 0.92
0.84 0.68
0.36
Figure 6: Visualization of the loss landscape with input normalization, in the teacher and the student, for different parametrizations. The teacher satisfies X exxeexi for different [X*] values. The first two columns correspond to students with correct angle 0 = 8 but wrong absolute value and the last two columns students with correct absolute value v |X but wrong angle. When we fix one variable, we ignore how it affects the loss for the Hessian caclulation. Each line corresponds
different parametrization: the first line uses a polar parametrization (A exp(if))) the second line uses the double exponential parametrization used in the LRU (exp) and the third one is the optimal parametrization for that task (tanh). Overall, both reparametrizations enable to control the explosion of the Hessian. However, the size of basins of attraction around optimality, or their number, shrinks as [X* goes to I for the angle, but not for the absolute value, highlighting how difficult learning the angle can be.
C.2 Structure of the Hessian at optimality
In Section 4, argue that the Hessian optimality is important object understand the learning dynamics in the linear teacher-student task we consider. We here provide some theoretical analysis of its structure the complex diagonal setting, that is we consider a recurrent network of the form
h441 X h4 bxr
(113)
8/2 = Re(c ht] dz,.
(114)
with X, 6 and e complex vectors of size n, with the number of hidden neurons, and d scalar. We additionally take the loss be the mean-square error, which also the one in our numerical experiments. Note that, as in our theoretical analysis of Section 2, we consider infinitely long sequences and wide-sense stationary inputs.
Recall that the Hessian of the loss is equal to
d d0a
dh, 0 L dh,
θh, he +
Ex
(115)
do Oh do
Th. d02
At optimality, only the first term remains, as 3n L is 0 for all data points. Given that we have shown earlier, e.g. Section B.2, that the most sensitive parameters to learn are the recurrent ones X, we focus on the Hessian with respect to these parameters in the following.
27

---

A
B
80
10a
1,00
60
10a
0.75
40
0.50
8g -2
10e
Loss
20
0.25
10-1
0.00
-4.
-5
-20
10
1.5
2.0
2.5
3,0
3
of"
10a
10
10
10-1
10-3
20k
40k
20k
40k
20k
40k
Step
Step
Step
Figure 7: Learning the angle is difficult, even in simple one-dimensional task. The target value is equal to 0.99 exp(ix/ (100 and is plotted in yellow. The black lines correspond the Adam learning dynamics. A. When the angle is not reparametrized (0 wi),, the loss landscape is extremely sharp the we direction, but Adam compensates for it. B. When the angle is parametrized exponentially (0 = exp(wa)) )), the loss landscape becomes smoother. However, this only hold when the considered angles are small enough, as the exponential parametrization does not bring extra granularity elsewhere. C. When reparametrizing the angle to reduce the gradient explosion as |^] 1, the loss becomes extremely tricky to navigate. The parameters are first attracted a nearby valley, which is flat the we direction and only indirectly connected to the global minimum. Such reparametrization thus hinders optimization far away from optimality. See Section C.1.4 for more detail.
C.2.1 Hessian for complex-valued variables
Before delving into more specific calculations, we make a few remarks on how to deal the Hessian when having complex-valued parameters. We will mostly leverage the fact that the loss L is real- valued.
Before that, recall few facts about Wirtinger derivatives:
- For f2) a complex-valiee valued function of 2, the Wirtinger derivatives are defined as:
df [Re[f]
ddm
(116)
dz [Re[2]]
dlm[z]
df [dRe[f]
dIm[f]
(117)
dz [Rc[2]] dlm[2
We have
df df
(118)
dz dz'
Leveraging the fact that L is real-valued that L L,
have
d²L
d
dp
(119)
dxz
dx
d
dr, T
(120)
yx
dス
12P
(121)
dx²
28

---

and. similarly. dxdxL = dzdzL. Second derivatives are symmetric, so we additionally have dxdرL= dxdzL which means that the complex Hessian is Hermitian matrix.
Taken all together, this shows that the full complex Hessian, which contains all cross derivatives, has a similar structure to the real case.
C.2.2 Hessian with respect to the recurrent eigenvalues
In this section, we compute the full complex Hessian with respect to the recurrent eigenvalue À and defer the analysis of reparametrization to the next section.
First, let us remark that
dL.
0L.
dhg
c
(122)
Yp
Zy:
yd
(123)
so that
de
d
dhi
T
(124)
dA
dd
Qy:
hh
θL,
dhi
L1 dhg Oy:
c
(125)
ax2
By:
dx
ax
We assumed that we are at optimality so that the network perfectly fits the target trajectories and Oy L, = O. Additionally, L, the mean-squared error loss that 0 L, Id. It follows that
L dx2
dhi
dhp
TC
(126)
dA
ij
ij
T dhe dd
dhe
(127)
dA
dhea
dhes
c dd;
Gj
(128)
dd,
In the last equation, made of the fact that the parameter X only affects the hidden state he and not the others, so d,, h,, = 0 0f j.
The previous calculation applied to one sequence, we now take the expectation over the data:
d L
dhy dhg T tooc dd dA
dx² = (cc O En.y
lim
(129)
Note that we introduced slight abuse of notation in the previous equation dy is in general matrix. However, given that the hidden neurons are independent here due to the diagonal connectivity. is effectively a vector, and we treat that way. Let us now compute the expectation, using similar calculation techniques to the one we used in Section B.2:
lim dhe, dhus = 1 dd; ddA
Eg.y
ber nmnmp-11 1030- Σ nA n.m20
(130)
= b,b,E
Σ
nd 1--11 x x_nmd -2-1
xm
(131)
n.m20
b,b,
amn
X3 1-
Ro m)
(132)
n,m20
We can now remark that this quantity is very similar the one we have encountered in Section B.2, up to the presence of b,by, and can be simplified using Lemma 2. For conciseness, note S(A, Xÿ)
29

---

the right-hand side of the last equation without the byby factor. Putting this result back in the Hessian, we get
T d (.b.Ge,5((...))
(133)
To gain further intuition of the behavior of this quantity, we take RA(A) = اا being real number. A similar calculation to the one we did in Section B.2 gives
S(A, A1) = (1 ()- + ( p²
(1 - -((--))AAAAA 111,) 22, 22,)
(134)
Given the complexity of this formula, we visualize the magnitude of S(A,, 23) on Figure 8. Interest- ingly, observe this quantity is large when X; and '₃ conjugate to each other and inputs uncorrelated. However, as elements in the input sequence get more correlated (p - 1), this effect disappears and |S| increases one of the two eigenvalue gets closer to I in the complex plane. In both cases, the effect gets amplified as the magnitude of the eigenvalue increases.
p 0.0
P p=
P 0.99
P 0.999
p=1.0
105
10a
0
10A
-1
-1
-1
-1
10-1
Re(A]
Re(A]
Re(A]
Re(A]
Rej A]
Figure 8: Visualization of A + Ao)) for Ào = 0.99 exp((ir44) This term measures how "similar" eigenvalues are in the Hessian. When p = 0, eigenvalues are mostly "similar" when they are conjugate to each other. As p increases, this effect decreases and eigenvalues become more "similar" one of them gets close to
We also need to compute d T. to get the full complex Hessian. Similarly to the previous calculation, we first have
dL,
aL,
aL.
OL. dyr.
dh, dA
(135)
di yb di
follows that
dLL
dhu dhu dii
EE
(136)
dA,dY
dA; b,,
(137)
Using the symmetry with the complex Hessian matrix, we now have all its components.
C.2.3 Hessian for different parametrizations
So far, we have computed the complex Hessian, which is not of direct use as we end up optimizing real numbers in practice. Here, study the impact of different parametrizations of X on the Hessian. Given that this parametrization only affects and not the other parameters in the network and that we only consider the Hessian at optimality here, computing the Hessian of those parameters reduces to left and right multiplying the Hessian by derivatives of À and Ā with respect to these parameters. For future reference, introduce
d8 dypdA dA,dA
d3
11:3
dh
afr. )
Agj Bij
Bij Agg.
(138)
d',
with Au = (A1, and By Sib,cie,S(A.....
30

---

Real-imaginary parametrization: X = Whe + Wim-
We aim at computing the matrix
d dure.
de dw deL dwian
(139)
dLL
dwlan
which is the building block to compute the full Hessian. First, let remark that dus Ao dwre, X 1/2, dwnm. d 1/2 and duim; 11/2. follows that
1/2,
L
(dueeo A3 dura, Hj) (dhi As dumn Ax)
(140)
dure, clure.j
(1 (1 1)T
(141)
(Re[Ay] + Re(B,,]).
(142)
Once again we emphasize that the first line only holds as we are optimality, Similar calculations give the rest of the elements of 11
(
Re(Ai + Buy Im| A.j Biyl
Im|- + B,,) Re[- - Aij + B.yl
(143)
Given the intuition we gained the structure of S previously, and the fact that As x S(A, Xÿ) and By & S(A,, Xÿ), we know that this block will have large components if the two corresponding eigenvalues are conjugate of each other or aligned to each other, or if one of them is close to 1.
One other quantity that we can calculate is the trace of the Hessian H which is equal to the sum of its eigenvalues. Note that this does not correspond the eigenvalues of the full Hessian matrix, it additionally contains entries for other parameters. Yet it already provides some idea of how large the Hessian will be, as the value of this trace appears in the value of the full trace. We have
Tr = Σ (Re(As + Bu] + Rel-A.t Ais + Bxl)
(144)
Σ Re [x]
(145)
(146)
where we used that S(,, A,) is real-valued in the last line. As a side note, this formula partly justifies why studying the expected squared magnitude of dy in Section 2 makes general sense,
dhu d0
E
(b,1,5(As, 7.).
(147)
Magnitude-angle parametrization: A = p(w))) (xp() The calculations for this parametrization similar the previous one, with the following differences:
dA
v' v'(w)) exp(if(wo))
(148)
dw,
dî
v v'(w) exp( (wp(-))(wa))
(149)
dw,
dA
in(w))) (wo) exp(ii(we))
(150)
dwig
dλ
iu(ww) (wo) exp( (oon(-()())))))))
(151)
dwi
31

---

After some calculations we obtain
dL
V (Wan)) Rele" Rele"(ounn (Otmn,)) As + el(0(un,)) e -(wo 5)3 Biyl
(152)
d durng
d
v' (wen) (wwn))((uung))) )0¹ (wois) Tm| (+owoo)))) Aij e (10(wo,,) Bayl
(153)
d dwo.j
I²I
v(w,, )0' ( (W),))/ (wem - - Ais + c'(0(un.)) -((w)))))) -(00(s0.) Biyl
(154)
dwo, duinss
2
d
u(wa,)) (w)))((w))))wwoo (wwoo)) ( As, e(10(uo.)) Bu]
d dwo.j
2
(155)
C.3 Experimental details
We use the linear RNN architecture defined in Appendix A. 1 as teacher and implement our experi- ments in JAX [63], using the default Flax [64] implementation of RNNs and the LRU implementation of Zucchet et al. [65]. Code is available here.
We initialize RNNs in the same way we initialized the teacher, and initialize the eigenvalues of the LRU and other complex-valued networks with magnitude in [uo 1] and angle within [-0g, Bu).
Given that we interested in the optimization properties of the different architectures, only report training losses and do not perform any cross validation.
Here are additional details related to the different figures:
Figure 3: see Tables and 2.
Figure 4 and 9: for panels and B, we use 1) 0.99 and draw in slightly different manner to the one described above (we directly draw the eigenvalues and eigenvectors so that we have two pairs of complex eigenvalues). We use automatic differentiation to compute the Hessian. For panels C and D, we use the same setup described in Table 2, but keep the learning rate constant over the course of learning. We report the effective learning rate at the end learning.
Figure 10: for panels A, B and C, we draw the magnitude and angle of 10 À values independently, uniformly in [Po. =t and [-0g, 0o]. Importantly, this means that there are no conjugate pairs, which leads to more diagonal Hessian matrices optimality than in Figure 4. For panel D, see Table
Figure 11: same setup as for Figure 3.
As rule of thumb, each LRU (or complex-valued diagonal network) experiment takes 3 minutes on a consumer-scale GPU (NVIDIA GeForce RTX 3070) and each RNN experiment takes 10 minutes CPU. The scans behind the results reported in the different figures require the order of few hundreds run each. Including our preliminary explorations, the results we report in this section required 30 days of compute, one third of it on GPUs and two thirds CPUs.
C.4 Additional analyses
C.4.1 Structure of the loss landscape for LRUs and S4
In the main text, we only provide an analysis of the loss landscape for the fully connected linear recurrent neural network and its complex-valus by performing the same analysis for the LRU and S4. Given that S4 involves some form of parameter sharing between the magnitude and the angle of the recurrence complex eigenvalues through A, which is required to avoid the explosion of the angle gradients, we are particularly interested in observing whether fully mitigates or not the gradient explosion effect. The results of Figure 9.B and D do not reveal such benefits: the loss landscape does not have high curvature on the w direction, but it is moved the wa direction. We haven't investigated whether qualitative changes arise when changing the input distribution.
diagonal counterpart. We here complete this result
32

---

RNN
LRU
Batch size
128
Sequence length
300
Hidden neurons (teacher)
10
Input / output dimension
Mo Oo
{0.32, 0.68, 0.84, 0.92, 0.96, 0.98, 0.99)
π
Hidden neurons (student)
64
log learning rate Optimizer (schedule) Initialization
[-5, 4.5, 4, -3.5, -3, -.5
|-2.5,
-2, 2, -15, -1, 0.5)
Adam (cosine)
yo teacher, un = 0]
Vo teacher
Number iterations Seeds
10k
10
Table I: Experimental details for Figure 3.A. We use |--.] to denote hyperparameters that were scanned over with grid search and }.... } denote the variables of interest for the figure. We chose the learning rates for the two architectures on preliminary scans and verified that non of the extreme learning rates were optimal in the final scan. For the RNN, we found that initializing with Po gave better results than initializing with the same distribution the teacher has, so we included this choice in the scan.
RNN BLOCK DIAG. RNN
COMPLEX DIAG. RNN LRU
Batch size
128
Sequence length
300
Hidden neurons (teacher)
10
Input/ / output dimension
0.99
0o
π
Hidden neurons (student)
64 / 64 and 128
64 [-2.5, -2,
log learning rate
[-5, -.5, (P teacher, 1// 0]
-3.,, -3 225
1.5, -1, 0.5)
Optimizer (schedule) Initialization
Adam (cosine)
1/ teacher
Number iterations Seeds
10k
10
Table 2: Experimental details for Figure 3.B. We use [...] to denote hyperparameters that were scanned over with grid search and f... } to denote the variables of interest for the figure. We chose the learning rates for the two architecture types on preliminary scans and verified that non of the extreme learning rates
optimal in the final scan. For the RNN, found that initializing with vn = gave better results than initializing with the same distribution the teacher has, so we included this choice in the scan. For the RNNs, used 64 neurons for the "RNN" entry, 64 for the "block diagonal" one, and 128 for the "more neurons" one.
33

---

RNN
COMPLEX DIAG. RNN LRU
Batch size
128
Sequence length
300
Hidden neurons (teacher)
10
Input/ / output dimension
Vo log(0u/1r))
0.99
{-2, - 1.5, -1, -0.5, 0}
Hidden neurons (student)
64
log learning rate Optimizer (schedule) Initialization
|-4.5, 4, -3.5, -3]
|-3.5, -3,
-0.5, 0]
Adam (cosine)
(P teacher, 0] Pg teacher
Up teacher 0o teacher
Number iterations Seeds
10k
10
Table 3: Experimental details for Figure 10. We use [...] denote hyperparameters that were scanned over with grid search and {...} to denote the variables of interest for the figure. We chose the learning rates for the two architectures learning rates were optimal in the final scan. For the RNN, we found that initializing with Po 0
preliminary
and verified that non of the extreme
gave better results than initializing with the same distribution the teacher has, we included this choice in the scan.
A
8 8
CM
C
B
4/2 w/x
8
8
C
C
04
10
105
10a
47
103
Hessian
84
10a
46
10a
Effectivve UR
B
0
B
0
-100
-100
C*
-103
C
--00
10-+
-10"
D
-10*
My
Top 10 vect.
D
0
Effective LR
10-1
10
107
107-
10-²
109
109
10-7
10-7
67
0m
8m
10
20
30
0
10
20
30
Index eigenvalue
Index eigenvalue
Figure 9: Equivalent of Figure for the LRU (A, C) and S4 (B, D). In the LRU, the exponential parametrization of the magnitude but not the one of the angle have accumulated leads to an Hessian at optimality that has large entries on all recurrent parameters w W and wa. similarly the behavior panel D, we initialized A at I given that this is the initialization that yielded best performance, We note that we didn't find qualitative changes in this plot when changing the initialization scheme of S4.
(xx(w,)) efficiently mitigates the Hessian explosion exp(wi), consistently with the theoretical and empirical evidence we far. In B, A is set 0.01 in the $4 architecture. Setting it to 1 (not plotted here)
observed for the complex diagonal RNN studied in the main text. For
C.4.2 Concentrating eigenvalue distributions
The goal of this experiment is to better understand how the concentration of eigenvalues À affect the learning dynamics. For fully connected RNNs, there no reason to expect major change in behavior. However, it is different for diagonal RNNs. The theoretical analysis we have done Section C.2 provides us with the following insights. When the elements the input sequence are uncorrelated, as it is the case here, the entries in the Hessian corresponding to two different eigenvalues increase if they are aligned or conjugate to each other, and if their magnitude large. We therefore expect that, as the interval on which the angle of the teacher's cigenvalues shrinks (0p - 0),
34

---

those eigenvalues will be more likely to be "similar" to each other. This results in large non non-diagonal terms, we confirm in Figure 10.A, and C. The LRU suffers less from this problem thanks its reparametrization, which reduces the overall magnitude of Hessian entries related to the magnitude, and partly the one of angle parameters (when it small positive number). As consequence, the performance between these two architectures increases as 0 + 0, as seen on Figure 10.D.
6g x/ /100
8g = x/10
8g = π
10a
10c
10a
104
A
10a
Hessian
10
10"
100
0
100
-109
-10
-10
4
-00
-10
-10
106
10
elg. vect.
Eigenvalue
10
10a
10
10a
10*
10a
10'
eigenvalue
eigenvalue
Index eigenvalue
10-1
10-3
RNN
Complex
LRU
10-5
10-2
10
Figure 10: Concentrating eigenvalues make the Hessian less diagonal (0o - 0) and consequently increases the gap between the LRU and the complex-valle diagonal RNN. A, B, C. Hessian of the loss with respect to the À parameters in the complex-values diagonal RNN. The Hessian is computed through the theoretical formula of Equation 143; computing it numerically marginally developed in Section C.2, concentrating the
affects the results. Consistently with the intuition eigenvalues affect the structure of the loss landscape. It makes the Hessian at optimality less diagonal and Adam cannot efficiently compensate it. The LRU does not suffer much from this problem, and the gap between the two architecture widens as 0g + 0.
C.4.3 Impact of the number of heads in fully connected linear recurrent networks
In Figure 3, we have shown that constraining the connectivity matrix to be block diagonal with blocks of size x 2 lead to critical boost in performance. Further analysis revealed that this arises the Hessian becomes more diagonal and Adam can thus better compensate for gradients explosion. Here, we study this behavior in more detail by interpolating between the fully connected case and the block diagonal one. This can be achieved by increasing the number of independent heads from (fully connected case) to 32 (2 particular, we are interested in understanding how big heads can be while keeping this performance boost. We plot the final performance, as well as the evolution of the effective learning rate for the A matrix over the course of learning Figure 11. We find that slightly bigger heads, until 4 x 4 (which corresponds to 16 heads), yield similar benefits. Additionally, the learning rate analysis reveals that the adaptive learning rates of Adam can more selectively compensate for potential gradient explosion cases as the number of heads increases, allowing for bigger learning rates overall and better performance.
Mlok-diagonal connectivity matrix, as we have 64 hidden neurons). In
35

---

10a
Effective LR (A)
Heads
10-+
10-2
10-s
10-1
Number of heads
Epoch
Figure Evolution of the performance (A) and effective learning rates for the A connectivity matrix (B) of linear recurrent neural network as vary the number of heads, keeping the overall number of hidden neurons fixed. should be noted that increasing the number of heads decrease the total number of parameters, the matrix A gets sparser.
Autocorrelation
10
10-2
y x-9.44
10
200
400
Figure 12: The empirical autocorrelation function (averaged over feature dimensions) of the BERT embeddings used in Section 5 can be approximated as sum of two exponen- tially decaying functions. The blue line represents the autocorrelation function R Rempineal (A) of the BERT embeddings of the Wikipedia dataset. As a first approximation, Re rempirical(A) z 0.3765x=0- For a more refined approximation, we perform linear regression of the log autocorrelation against A, shown by the black line. This yields the following approximation: (A) A) 0.3326 + 0.044 0.997
equal to
pempirical
D Signal propagation in randomly initialized deep recurrent neural networks
D.1 Experimental setup
We detail the experimental setup used Section 5. We select the first 512 tokens from 1024 random sequences in the Wikipedia dataset [66] and pass them through the BERT [67] tokenizer and embedding layer. This yields dataset of 1024 examples with length 512 and feature dimension 724. Figure 12 shows the autocorrelation function of these inputs, revealing that the i.i.d. assumption serves (p = 0) as good first order approximation. This validates the relevance of our toy experiments for studying signal propagation in more realistic settings. To refine this approximation, we can include a high correlation term (p close to 1).
We examine realistic networks comprising 4 blocks with the following structure: recurrent layer, non-linearity, a gated linear unit [58, GLU) and a skip connection. By default, we omit normalization layers, but when included, as in Figure 5.C, we place one normalization layer before the recurrent layer and another one before the GLU. All the layers involved contain 256 neurons. We also incorporate a linear encoder the beginning the network and linear decoder its end.
The loss that we use a next-token mean-squared error, defined as
L S (*11-1) x ||
(156)
where é; (x1:-11) represents the prediction the network. Figure 5 reports the average squared value of the hidden state the gradient. This average is computed over sequences, but also over all neurons parameters and over all time steps. We compute gradients using batches of size 8.
In Figure 5 we vary Pg. which controls the magnitude of the eigenvalues of the recurrent Jacobian. Specifically, we sample those magnitudes in the interval (Po, (1 + vn)/2). For the complex-valued diagonal RNN and the LRU, we apply the LRU initialization. For the LSTM,
the Chrono
36

---

LRU
CRNN
GRU
2nd moment del
2nd moment del
2nd moment dad.
106
10a
10-6
0.9
0.99
0.9
0.99
0.9
0.99
angle
2
GRU
1
2
3
magn.
rest
2 3 rest
2
Figure 13: Gradient magnitudes are independent of the layer. This figure presents similar plots to Figure 5.B, except that parameters from different layers are no longer aggregated together. Instead, each parameter group of each layer has its own line. The indices in the legend correspond to layer indices.
initialization proposed by Tallec and Ollivier [29]: it initializes the forget gate biases such that, when the input x and the hidden state h are equal to O, the time constant associated to f is uniformly sampled from 1-vo 1 1-v0 and the input gate is equal to 1-
While Figure 5.B presents aggregated gradients over layers, Figure 13 offers layer- version of this analysis. It reveals that the layer index does not significantly impact gradient magnitude. Surpisingly, given that the hidden states of the complex RNN gets larger with depth (c.f. Figure 5.A), this result might seem unexpected for cRNNs. We can attribute this to the backpropagated error signals also being amplified during the backward pass, as discussed in Section 2.2. In the first layers, hidden states are small and errors are large, while in the in last layers, errors are small and hidden states are large. Consequently, the gradient magnitude remains relatively constant accross layers. For the GRU, the gradient magnitude reported in Figure 5 for the non-GRU parameters included the linear encoder and decoder. As the encoder gradients are the dominating ones, this explains why the gradient magnitude for the non-GRU parameters is smaller in the layer-wist analysis.
D.2 Can gated RNNs be considered diagonal?
In Section 3.2, we argued that the diagonal linear setting we focused our theory on can be decent proxy for more general gated RNNs, whose À values can depend on both the inputs x and the hidden state h. Here, we assess whether this holds true at initialization. To that end, we study how the Jacobian d of a GRU evolves A grows. For the linear diagonal regime to be good proxy, two dh, necessary conditions must be met: First, all the non-diagonal terms should be negligible compared to the diagonal ones. Second, the diagonal terms should decay similarly as their À values were independent of and h.
In Figure 14, we report the evolution of all the diagonal terms of this Jacobian and random subset of the non-diagonal ones. Our findigns indicate that the non-diagonal terms are much smaller than the diagonal ones under the standard initialization, supporting the first necessary condition. Additionally, input and hidden state-dependent gates do not qualitatively change the decay of the diagonal terms, particularly when the network is initialized with À values close to 1 (which correspond large time constants). Furthermore, we find that increasing the strength of all the hidden state gate connections (Whn. Whr Whe) breaks the diagonal- like behavior and eventually leads to exploding Jacobians. However, this only occurs at values that are much larger (3 times more in this case) than the default initialization of these weights (orthogonal initialization).
In conclusion, those results confirm that the theoretical setting we have considered in this paper is good proxy for studying signal propagation in realistic recurrent networks. While we have focused our analysis on GRUs, we expect these results to hold for other architectures such as LSTMs, Mamba, or Hawk. For the last two, given that the different gates only depend on the input, we expect the matching with our theory to be even stronger (c.f. Figure 14.C σn 0, which captures this regime).
37

---

A
Constant gates
TE[1.1]
TE(1, 16]
Te(1,
1.0
Entries of
0.5
0.0
Input hidden state dependent gates [=[1.16]
TE(1.1]
Te [1,66))
1.0
Entries of se
0.5
A
4
C
Varying strength of hidden state to gates connections
1.0
Entries of tt
0.5
0.0
10
101
10²
10a
101
10?
10a
10'
107
Diagonal terms
Non-diagonal terms
Figure 14: GRUs behave like diagonal linear networks. This figure illustrates the evolution of the recurrent Jacobian d of a GRU, when provided with the BERT embeddings of a sentence extracted a/o from the Wikipedia dataset. A. In the first row, we take the forget gates to be independent and we sample them with the Chrono initialization [29] for different intervals. The resulting network is linear and diagonal, similar to what we have studied in the theory. These plots therefore serve as visual reference for comparison with the realistic case. B. The second row shows the same plots as the previous row, except that the gates are now dependent on the inputs æ and on the hidden states h. As mentioned in A.4, we initialize all the linear layers taking r as input with LeCun initialization and the layers taking h as input with orthogonal initialization. The recurrent Jacobian evolves similarly than in the constant case, particularly on slowly decaying dimensions (large T values). C. In the row, we aim to break the diagonality of the model by increasing the strength d of the hidden state to gate connections, for T [1, 16). The case 0 0 corresponds similar to architectures like Mamba or Hawk. The plot with σ = 1 is the same as the middle one in B. For σ₁ = and higher, diagonality progressively breaks and the recurrent Jacobian eventually explodes. We note that these plots were obtained from single example. Yet, we have found the behavior we report here to typical of the general behavior.
and
gates that only depend on æ,
38

---

D.3 Does our theory apply to gated RNNs?
Having established that gated RNNs behave similarly to the linear diagonal RNNs considered in our theoretical investigation, a question arises: How well can our theory describe signal propagation in gated RNNs on realistic data? To address this, we study simplified version of the GRU network (similar to the one studied by Chen et al. [51], which incorporates realistic gating mechanism:
fi+1 = O[ F1 + Wyhhe + by)
(157)
h+11 fe+1 he S141) T++1-
(158)
As in the rest of this section, provide BERT embeddings of sentences from the Wikipedia dataset as inputs.
To apply our theory to this architecture, we must address two main challenges:
The gate fi+1 depends on both h++1 and zt, making non-constant. Based on our empirical results from the previous section, we reasonably approximate fa+1 s o(by) =: X, ignoring this dependency.
2. Our theoretical derivations have overlooked cases where the recurrence strength À normalizes the input zt. When considered, we detached the normalization factor from the computational graph (as in Section B.3.2). However, can extend our calculations from Section B.2 to accommodate this setting:
((, = (1 (1 1 a.
(o) + Σ (a (2)) A>1
(159)
E|hi| = f(A,
(160)
Ohe
8 (ax 3)
E
ax
=
dadß
(161)
-=
Note that we obtain [E[( ( )²] by multiplying E[( the )²) by o' (by)².
With these adjustments in place, we can now empirically test our theoretical predictions. For simplicity, we approximate the auto-corre correlatooltnnn function Rg (blue line in Figure 12) as (x) x 0.3326 + 0.044 x 0.997 Figure 15 presents our results, which reveal:
An almost perfect match between theory and practice for constant gates, confirming that our sample size is large enough.
A very precise, though not perfect, match for context-dependent gates.
The variance of h4 and Oh. shows minimal dependence X, indicating that the magnitude bby of error signals received by by. Wys, fb, and fz are largely independent of the time constants encoded in the network.
39

---

Constant gates
Input-dependent gates
Input and hidden-dependent gates
10a
2nd moment h
10-2
10-4
10-6
10-4
10-10
10a
10a
10a
10a
10a
10a
10a
2nd moment deh
10-2
10"4
10-4
10-8
10
10A
10?
10t
10?
10a
102
Samples
Theoretical prediction
Empirical mean
Figure 15: The theory developed for linear diagonal recurrent networks captures signal propaga- tion within gated recurrent neural networks. The different samples were obtained as follows: 100 different randomly initialized networks are given different input sequence length 512. The biases of the forget gates by are initialized with Chrono initialization for T E [1, 256]. For each of these models / sequences, we measure 1512, and do db,., 1112, (i being the index of one of the 256 hidden neurons). We report this measurement (A T/1 + 7"). The empirical mean is obtained with kernel regression with the Gaussian kernel C(aa, b) := exp(-(( - 6)2 /100). The theoretical prediction comes from the approach described in Section D.3.
function of the time constant encoded by the neuron
40

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We provide detailed analysis of the claims we make in the abstract in the introduction in controlled setting and show that our conclusion partly extend to reasonable setting.
Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We clearly state the assumptions of our theory and our experiments are performed in controlled setting.
3. Theory Assumptions and Proofs
Question: For each thcoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: the assumptions of the theory clearly stated in the main text and all derivations are included in the appendix.
Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: The appendix provides thorough description of the experiments perform.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described supplemental material?
Answer: [Yes]
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary understand results?
Answer: [Yes]
Justification: The appendix provides thorough description of the experiments perform.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer:
[Yes]
Justification: We include error bars for all our learning experiments.
8. Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Answer: [Yes]
Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurIPS httpg:/pppp..eeeo.......opo//oat Code of Ethics https: /tte:://neurpspe http///bicchhhiicciiii..........
41

---

Answer: [Yes]
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: The nature of this work is highly theoretical, we don't foresee any negative nor positive societal impact.
11. Safeguards
Question:
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
12. Licenses for existing assets
Question: Are the creators or original owners of assets (c.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [NA].
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, well as details about compensation (if any)?
Answer: [NA]
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Question:
Answer: [NA]
42


Reviews:
[{"review_id": "7P0iAiSj4a", "reviewer": "Reviewer_zt6Z", "summary": "The paper provides a theoretical and practical analysis on the difficulties of training recurrent neural networks. In particular, given the current rise of interest in leveraging recurrent mechanisms for long sequence processing -- due to novel architectural components and solutions (deep state-space models, linear/diagonal recurrences, exponential parametrizations), identifying the crucial components allowing to preserve long range information is a very important research avenue.", "questions": "1) (Minor) While I understand that recent literature has completely overriden this notion, the term \"state-space model\" refers [1,2,3] to a very general broad concept which is common in control theory and dynamical systems -- which simply refers to a model equipped with state variables controlled by first-order differential equations or difference equations. Such definition  covers the typical feedback connection of Recurrent Models. Thus, stating in the abstract  and in the main text \"state-space models, a subclass of RNNs\" seem a bit too strong to me -- even if several recent works highlighted the connections. I would be completely fine with specifying \"Deep\" SSMs (just to highlight that the authors are referring to modern models) or \"Structured\" SSMs and to mention the suggested references, to better frame the paper to the readers. \n\n2) In Eq. 1, the authors consider a model not equipped with an output mapping (that was instead considered in Eq. 10, i.e. y_t = Ch_t + Dx_t) -- the hidden state h_t is directly provided to the loss function.  Why is that the case, and do the considered analysis still hold when considering such mapping (that is common in RNNs and SSMs)?\n\n3) The authors refers to an Infinite time horizon/infinite sequences (lines 98/100) as the basis for their analysis. Practically, how long the learning process necessitates to be to the considered issue to emerge/is this something that also the experimental analysis considers? In such case, does this have relation with recent literature highlighting the difficulties of sequential models when dealing with Infinite length sequences [3, 4]? I would like a comment by the authors on this.\n \n4) Line (40): the notation (x_t)_t is not very clear to me.\n\n5) Could the authors give an intuition on lines (113/116)?  Why an higher correlation in the input data should imply an higher variance in the model state (in my understanding, similar inputs imply similar loss function, thus smaller gradients)\n\n6) In line 135, the authors use yet another form for the state update that do not consider the B linear mapping. Why is that the case?\n\n7) The teacher-student task should be better described and justified. Why did the authors chose to tackle such task instead of other synthetic tasks (such as selective copy, induction heads etc. see recent surveys [3] for references on synthetic datasets ). And do the authors believe that the teacher-student task is general enough (i.e., points sampled from a normal distribution) to represent a setting where the state has to learn meaningful long range dependencies? Do the conclusions still hold when dealing with tasks  requiring to store more informative data? I believe that  experiments could benefit from experimentations on more \"difficult\" settings (selective copy, Long Range Arena) to confirm the theoretical considerations with real learning scenarios. \n \n[1]  Genshiro Kitagawa. A self-organizing state-space model. Journal of the American Statistical Association, pages 1203–1215, 1998. 4\n\n[2]  Anton Schafer and Hans Georg Zimmermann. Recurrent neural networks are universal approximators. In Proceedings of ICANN,  pages 632–640. Springer, 2006.\n\n[3] Tiezzi, Matteo, et al. \"State-Space Modeling in Long Sequence Processing: A Survey on Recurrence in the Transformer Era.\" arXiv preprint arXiv:2406.09062 (2024).\n\n[4] Tiezzi, Matteo, et al. \"On the resurgence of recurrent models for long sequences: Survey and research opportunities in the transformer era.\" arXiv preprint arXiv:2402.08132 (2024).", "limitations": "The paper does not discuss limitations.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper provides an interesting theoretical analysis that highlights which are the key components allowing modern solutions (Deep SSMs) and gated models (LSTMs/GRU) to achieve good performances. Moreover, it provides some proof-of-concept practical results that help to confirm the analysis. The paper is well written (apart from few typos, I suggest a proofreading).", "weaknesses": "I have some concerns with some of the paper initial framing and setting (see the Questions section) and the experimental analysis, which has been carried on very synthetic tasks which I am not sure are capable to represent real learning scenarios, thus possibly hindering the analysis validity.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their encouraging review. We answer their questions below:\n\n1. **SSM terminology.** We agree with the reviewer that the SSM term has been overloaded. This is the reason why we try to use the RNN terminology as much as possible. That said, we appreciate the suggestion of the reviewer and will change the terminology to “deep SSMs”.\n2. **Missing B, C and D?** We focus on the recurrence because the $B$, $C$ and $D$ mappings (from the usual SSM notations) are feedforward layers and signal propagation has been extensively studied in those layers (e.g. Glorot and Bengio 2010).\n3. **Infinite sequences.** In our preliminary investigations, we found that having time sequences longer than 3 times the characteristic time constant of the network ($1 / 1 - \\lambda$) is enough to consider the sequences infinitely long. The issue we mention appears from the very first parameter update and therefore is not directly related, to the best of our understanding, to the issues reported in the review (online learning and catastrophic forgetting).\n4. Would changing it to $(x_t)$ make it clearer? We want to emphasize that this is a sequence of numbers here.\n5. > Could the authors give an intuition on lines (113/116)? \n\nThe point we are trying to make in these lines is the following: if we want to train a deep network, we need neural activity to stay of constant magnitude over the network depth (following the classical results on signal propagation in feedforward neural networks). It follows that a recurrent layer that blows up the magnitude of the input will hinder the learning of the rest of the network. \n\n> Why an higher correlation in the input data should imply an higher variance in the model state\n\nWe do not have a proper intuition for that result. We would like to emphasize that correlation is here between time steps and not samples, which might break some intuition. Additionally, the intuition the reviewer has does not always apply: take a L2 loss, if the predictions are all the same (thus all similar) but extremely large, the gradients will also be large.\n\n6. We do not use any $B$ mapping for the same reasons as described in 2.\n\n7. **Justification for the teacher-student experiment.** That is indeed an important point that we didn’t comment on enough in the current version of the paper. We focused on this teacher-student task for 2 main reasons: First, we wanted to have an experiment in which we can easily control important details for signal propagation (e.g. magnitude of eigenvalues or the concentration of eigenvalues), which is almost impossible to do in other synthetic tasks. Second, we wanted to disambiguate signal propagation from other capabilities of architectures like memory. We believe that this setup is general enough for studying signal propagation as we can reproduce the same qualitative order between the architectures as in the LRU paper (on the LRA benchmark) and as our new analysis reveals that, up to a first approximation, inputs can be considered i.i.d. in our Wikipedia experiment (c.f. new results in the global answer). \n\nWe believe that our rebuttal answers the questions of the reviewer and hope that the reviewer may update their score accordingly. We remain available during the discussion period to clarify additional points.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "lY21hBWTLj", "reviewer": "Reviewer_uPdc", "summary": "The paper discusses \"the curse of memory\" which is the hypersensitivity of hidden states to parameters as the memory increases. This can lead to optimization issues, even if the exploding/vanishing gradient issue is addressed. The authors discuss solutions to this problem: complex diagonalization, normalization, and reparametrization. To this end, the authors derive closed-form solutions for the variance of the hidden state and its sensitivity to the parameters and discuss how the solutions can prevent the variances from diverging. This can give insights into why certain families of RNNs are performant. In particular, they show that both the variance of the hidden state and its derivative wrt to the parameters diverges as the memory goes to infinity (although at different rates). Therefore, normalizing the hidden state and reparametrizing the variable representing the memory's temporal scale can avoid the divergence issue.", "questions": "- What are the main contributions of the paper? Which results and which sections?\n- What are other types of normalization and reparametrization, other than the ones employed by LRUs, that can address the curse of memory?", "limitations": "NA", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- The paper puts forward closed-form solutions for the variance of hidden state and the sensitivity of hidden state to parameters. This is interesting and can provide insights into RNN dynamics and possibly network design.\n- The results and the discussion on adaptive learning rate and how it relates to the Hessian of fully connected and complex diagonal RNNs are interesting and, to my knowledge, novel.", "weaknesses": "- The main weakness of the paper is the presentation. It is unclear what the contribution of papers is. At many points in the manuscript, it is unclear if the authors are discussing previous results on LRUs or putting forward new results. It would be good if the authors explicitly spell out their contributions.\n- Although the paper is presented as discussing the curse of memory and the solutions to it, it reads more like \"why LRUs are performant\" The solution discussed: exponential reparametrization, normalization, and diagonalization, are all introduced and employed in LRU architecture [20]. However, the authors only discuss SSMs and gated RNNs as networks that implicitly address \"the curse of memory\". Meanwhile, LRUs employ the same solution discussed in the paper explicitly. Although LRUs can be thought of as a subclass of SSMs, vanilla SSMs [17] do not use the normalization, diagonalization, and exponential reparametrization discussed in the paper. Therefore, discussing LRUs only in experiments \"to represent SSMs due to its simplicity\" is not justified. \n- If the solution to the curse of memory is not the contribution of the paper, then the novelty and contribution become limited. As I mentioned earlier, LRUs employ the same solution and sometimes with the same justification, for example, exponential parametrization is used for higher granularity around 1 [20]. \n- To motivate their results, the authors mention that \"We answer negatively by showing that gradients can explode as the memory of the network increases, even when the dynamics of the network remains stable.\" but it is unclear where this is shown. \n\n- minor comments:\nL117: the paragraph on backward pass is confusing and hard to understand.\nL139: check wording. Could not understand it.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their review and address their concerns below.\n\n**On the contributions of the paper**\n\n> What are the main contributions of the paper?\n\nWe provide a detailed list of our contributions in the global answer.\n\n> If the solution to the curse of memory is not the contribution of the paper, then the novelty and contribution become limited.\n\nIn this paper, our objective is not to design a new architecture capable of avoiding the curse of memory — but instead to describe this issue thoroughly (it has never been reported before), and argue why recent architectures (e.g. SSMs) and classical ones (LSTMs, GRUs) have inner mechanisms alleviating the curse of memory. As we describe in the paper, the issue we study is deeply linked with the optimization of any recurrent model capturing long-range interactions, and translates into challenging loss landscape properties that we show can be alleviated by diagonality and reparametrization. Indeed, as the LRU paper (and others) shows, linear RNNs that are functionally equivalent (i.e. real dense linear RNNs, complex diagonal linear RNNs, reparametrized linear RNNs) have drastically different performances on tasks such as the ones within the long-range arena benchmark. Our paper dives deep into this intriguing finding and finds a direct effect of some of the commonly used tricks on the loss landscape.\n\n**On the different architectures**\n\n> What are other types of normalization and reparametrization, other than the ones employed by LRUs, that can address the curse of memory?\n\nWe decided to focus on the LRU due to its simplicity, but it is far from being the only architecture capable of addressing the curse of memory. As we argue in Section 3, the discretization of a continuous-time ODE used in SSMs (more detail on that in the next paragraph) and the input and forget gates of LSTMs / GRUs partly serve the same purpose.\n\n> Meanwhile, LRUs employ the same solution discussed in the paper explicitly.\n\n> vanilla SSMs [17] do not use the normalization, diagonalization, and exponential reparametrization discussed in the paper.\n\n> Therefore, discussing LRUs only in experiments \"to represent SSMs due to its simplicity\" is not justified.\n\nWhile the LRU paper discusses many of the techniques we study in our paper, their investigation is mainly empirical. The objective of our paper is instead to deepen our theoretical understanding of RNN architectures like LRUs, SSMs, as well as the classical LSTMs/GRUs. Our path is of course guided by the design of such architectures, but our findings do not simply revisit the relative papers: we go one step further and show how these models can partially solve the curse of memory issue through reparametrization, normalizations, and diagonal structure. The LRU is arguably the simplest model capable of alleviating the curse of memory. While we agree that the structure of SSMs such as S4 is quite different, it is already pointed out by the LRU authors themselves (see the last section in their paper), that many of the tricks employed in their design can be traced back to mechanisms used in S4: it uses a diagonal parametrization of the recurrence, delta serves as normalization, and ZOH discretization induces an exponential stable parametrization.\n\n**On the writing**\n\n> […] it is unclear where this is shown.\n\nThis sentence refers both to our analytical and empirical results. In our analytical results, we show that $\\mathrm{d}_\\theta h_t$, and therefore the loss gradient, grows to infinity as $\\lambda$ goes to $1$ (”the memory increases”). Figure 5.B shows that the gradient of the loss explodes in practice (there is a typo in the legend of this figure: the quantity reported here is the gradient of the loss). In all cases, the dynamics of the network are stable as $\\lambda \\in [0, 1)$.\n\nWe are confident that our remarks address the reviewer's concerns and provide necessary clarification, and we remain available during the discussion period for any additional questions. We encourage the reviewer to examine the new results reported in our global answer, which substantially extend the scope of our work, and we hope that our rebuttal will prompt a favorable reassessment of our paper's contributions and impact.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "4xUx13brNh", "reviewer": "Reviewer_Bgq8", "summary": "The paper explores challenges RNNs face in learning long-term dependencies. While generally attributed to the exploding and vanishing gradients problem (EVGP), the authors reveal that as a network's memory increases, parameter changes cause an explosion of the second moment of hidden states and their gradients, making gradient-based learning highly ill-posed. They study different design choices in recurrent architectures such as element-wise recurrence, gating mechanisms and careful parameterization with respect to the curse of memory. They apply their theoretical findings to successful architectures such as state-space models (SSMs) and LSTMs, offering new insights into why some RNN architectures perform better in gradient-based learning.", "questions": "- in l. 180 you assume $\\gamma$ is independent of $\\lambda$ yet in the Diff. eq. in l. 182 you introduce the dependency again ($\\gamma(\\lambda)$)? Is this connected to what you write in the caption of Fig. 2 where you decouple $\\gamma$ from $\\lambda$ when differentiating? Can you explain the reason for this in more detail?\n- in l. 289 you mention that you probe ADAM’s effective learning rate by providing a vector of ones to the optimizer - does that mean you simply read out ADAM’s (corrected) second-moment estimates and report $\\frac{\\eta}{\\sqrt{v_n} + \\epsilon}$, where $\\eta$ is the global learning rate and $v_n$ the second moment estimates at iteration $n$? If so, at which iteration/epoch do you query the effective learning rate?", "limitations": "While an explicit section on limitations is missing, the authors mention their limitations and specific assumptions throughout the manuscript.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "- The paper addresses a fundamental problem, the curse of memory, and connects it through theoretical analysis and experimental validation to the success of recent architectures, which makes the insights incredibly valuable to the machine learning community\n- I like how the authors connect the dots between multiple facets of model training, e.g. investigating the implications of RNN recurrence parametrization on adaptive LR optimizers, which are actually used in practice\n- The paper is also well structured and generally well written.", "weaknesses": "- To add to EVGP mitigation techniques (l. 61 - 64): There also exist regularization techniques based on dynamical systems theory, see especially [1]\n- There are quite some typos across the document, see the list below. I recommend running the paper through some sort of language checker to improve the manuscript from a readability and presentation side.\n- The authors repeatedly call $\\mathbb{E}[X^2]$ a variance (where $X$ here is a placeholder for any RV, i.e.  $h_t$ and $d_\\theta h_t$ in the paper), which is technically not correct. It is the second moment. I’d consider fixing this or explaining why the missing term for the variance, $\\mathbb{E}[X]^2$ (first moment squared), can be neglected.\n\nTypos / text-bugs:\n- l. 119: “[...] a low a pass filtered [...]\n- l. 191: “[...] does not hurts [...]”\n- l. 192 “Several RNN architectures implicitly alleviates [...]”\n- l. 200 “While such models can can approximate any smooth mappings [...]” \n- l. 219 “anaylsis”\n- l. 257 “Both network are trained [...]”\n- word missing in l. 283 “[...] that characteristic of the curse of memory, [...]”?\n- l. 298 consistently -> consistent\n\n[1] Schmidt et al., Identifying nonlinear dynamical systems with multiple time scales and long-range dependencies (ICML 2021)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer’s very positive feedback and their valuable input regarding these additional references and important details. We address their questions below.\n\n> in l. 180 you assume $\\gamma$ is independent of $\\lambda$ yet in the Diff. eq. in l. 182 you introduce the dependency again ($\\gamma(\\lambda)$)? Is this connected to what you write in the caption of Fig. 2 where you decouple from $\\lambda when differentiating? Can you explain the reason for this in more detail?\n\nIndeed, $\\gamma$ must be a function of $\\lambda$ for the normalization to be effective. However, including $\\lambda$ inside the normalization significantly complicates our analytical formula, making it harder to reason about them in the backward pass. This is why, for that section, we reason about $\\gamma(\\mathrm{stopgradient}(\\lambda))$.\n\n> does that mean you simply read out ADAM’s (corrected) second-moment estimates?\n\nThe quantity we report incorporates both the first-moment and second-moment estimates from Adam, not just the second-moment one.\n\n> at which iteration/epoch do you query the effective learning rate?\nIn the plots presented, we query the effective learning rate at the very end of training, using a constant learning rate (in contrast to the cosine scheduler used in our other experiments). However, our experimental pipeline monitors the effective learning rate at every epoch, and we have observed that the trends remain consistent throughout training.\n\nWe recommend the reviewer to take a look at the new experiments we report in the global answer. They show that our theory extends to the gated case and to a larger class of practically relevant architectures. We hope the reviewer may take these additional results into account and consider adjusting their assessment accordingly.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "71UBk4zSp3", "reviewer": "Reviewer_Bq7Y", "summary": "In this work, the authors present an analysis of Recurrent Neural Networks (RNNs), focusing on problems which impede optimization as the length of input sequences and model memory increase. The authors first focus on the well-studied problem of vanishing/exploding gradients, arguing that this problem by itself does not explain the difficulties in optimization. The authors turn to the less studied \"curse of memory\", and argue that this issue persists for RNNs even when gradients are controlled to not explode or vanish. Finally, the authors study how some of the alternative recurrent architectures such as LSTMs and SSMs mitigate the curse of memory, and provide some suggestions for how optimization can be improved.", "questions": "See above - why is the assumption of Wide sense stationarity needed, and how is your analysis affected by not adopting the assumption?", "limitations": "No limitations are presented or discussed in the paper as far as I could find. \n\nI see no ethical issues with this work.", "rating": 5, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "I am very intrigued by the approach taken in this work. \n\nThe analysis seems novel and the arguments provided are given in a straightforward and convincing manner. The theoretical and empirical results in the main body of the work all seem sound, and are augmented by substantial additional experiments in the appendices.\n\nThe authors' analysis results in clear explanations both for how current architectures mitigate the issues studied in this work, and the analysis suggests some directions that may lead to new architectural innovations in the future.  \n\nDespite many technical issues in the writing, the argument is presented in a clear way and the flow of the paper makes sense. See below.", "weaknesses": "I have three major issues with the paper in its current form. If issue 1 is addressed, I am willing to increase the score of the presentation from fair to excellent, which would bring the paper to a borderline reject. Issues 2 and 3 are more relevant to the soundness of the work, and I've lowered the score to a 2. If the authors address these, I will consider raising the score; however, I can't go above a 3 without a more substantial discussion of limitations, which I think might not be feasible within the rebuttal period; however, I am still open to adjusting this score depending on the discussion.\n\n1. One of the more glaring issues with this work is with the writing. The authors need to take some time to carefully check for typos, incomplete sentences, and correct grammar. I have listed a few particular issues here:\n  * line 30: remains -> remain\n  * line 69: You say here that you are showing the gradients can explode, but in the next section say you are controlling the gradients so that they do not explode (lines 83-86). Am I missing something conceptually here that makes this not an inconsistency?\n   * line 70: remains -> remain\n   * Figure 1: the last sentence of the caption is unclear and seems to be missing some words.\n   * line 86: should be \"studying this behavior quantitatively\", not the other way around\n   * line 192: alleviates -> alleviate\n   * line 283: characteristic -> are characteristic\n\nif these issues are all fixed and the authors will take some time to carefully edit, I will update my score to reflect this effort. Overall the typos only affect clarity in a few spots.\n\n2. In section 2.2 - you justify assumptions A) and B), but I couldn't find any justification of assumption C). Unless I am missing something, this needs to be provided, as assuming Wide-sense stationarity is not necessarily standard. At the very least, this assumption could be acknowledged as a limitation and room for future work. \n\n3. There is no real discussion section at all, and there is particularly no discussion of limitations.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer’s positive feedback on our paper and address their concerns below.\n\n1. **Typos**. We thank the reviewer for their detailed report of our typos. It will help us make the manuscript as clear as possible and without typos. Even though the system does not allow for the upload of a revised version of the paper at this stage, we already implemented all the corrections in our latex file and made several additional passes to make sure everything is in the best shape. That said, we are happy to provide clarifications for the typos that affect the understanding of our argument.\nRegarding line 69: we mean that the gradients’ magnitude can explode when the memory increases (i.e., $\\lambda \\rightarrow 1$), even when we are outside the traditional exploding gradient regime (i.e., we enforce $\\lambda \\leq 1$). It should be noted that in the first case the gradients explode as the memory increases, whereas in the second case, they explode as the interval considered increases.\nRegarding Figure 1: “as” is missing (”The loss becomes sharper as information is kept…”).\n2. **Wide-sense stationarity assumption.** A discussion of this assumption was indeed missing in the paper. We address it in the global response. In short, this assumption is standard in the analysis of linear time filters (see e.g. the Proakis and Manolakis book) and is mostly for convenience in making calculations easier (see e.g. the classical Wiener–Khinchin theorem). When it is violated, our theory does not perfectly predict signal propagation, but our results hold qualitatively and almost quantitatively (c.f. Figures 1 and 3 in the rebuttal PDF, which test our theoretical predictions in a practically relevant setting). Overall, this assumption is not a significant limitation of our theory, and using the empirical autocorrelation function leads to very similar results (see Figure 3 of the rebuttal PDF, left column).\n3. **Discussion of the limitations.** As we discussed in the global response, the limitations of our work are mostly due to the simplicity of our model and the nature of the approach we are following.\nRegarding the model we study: While our results only apply to the non-gated case, we show that they still capture the gated case relatively accurately, despite the dynamics being significantly nonlinear (c.f. Figure 3 in the rebuttal PDF). However, there is little hope that our analysis extends to nonlinear vanilla RNNs.\nRegarding what we can achieve with such an approach: Studying signal propagation has its limits regarding the insights it can bring. We can mainly study the loss landscape through such an approach, and cannot say anything about, e.g., the memory capabilities of a given architecture.\n\nThe reviewer's thoughtful comments have helped us clarify key aspects of our work. We are confident that these explanations address the concerns raised and provide a more comprehensive view of the contributions and limitations of our work. We hope that these clarifications reassure the reviewer to revise their score to better reflect their positive opinion of our paper.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "SW7hz8AuFt", "reviewer": "Reviewer_xzKN", "summary": "This paper studies the vanishing and exploding gradient phenomenon in RNNs. In particular, it answer the question: \"is solving those issues really enough to ensure well-behaved loss landscapes?\" This paper reveals the importance of the element-wise recurrence design pattern combined with careful parametrizations in mitigating the scale issue of hidden states.", "questions": "1. On Page 4, Equation 5 focuses on the second-order expectation of $h_t$​. It would be helpful to understand the specific reason for choosing a second-order expectation over a first-order one. Generally, implementations do not involve the second moment of hidden states. Why is the first-order expectation not considered in this case? Does first-order suffer from the same problem? \n\t1. A similar forward-pass blow-up has been proved in the LRU(Orvieto, 2023) paper, proposition 3.3. What is the difference in terms of assumption and technique between your result and the LRU result? \n2. On Page 5, in the paragraph titled \"What About Complex Numbers?\", you discuss the concept of polar parametrization. This raises an important question: do all parametrizations encounter the same issues as the polar parametrization? If they do, demonstrating this could effectively highlight the limitations of using complex numbers in a polar setting. If not, it may be worthwhile to find the good parameterization.  Currently, the exploration of polar parametrization seems quite restricted.\n3. Your title is \"vanishing and exploding gradients are not the end of the story\". A question based on the saying in Page 9 Figure 5 Panel C is: \n\t1. Does layer normalization and batch normalization resolves the all problems?\n\t2. If not, I think it's hard to claim \"Layer normalization keeps the overall gradient magnitude under control\".", "limitations": "1. This paper investigates how the hidden states of recurrent models contribute to learning long-term memory. There is very little research available to explain why the mamba model, which uses both gating and state-space models, is successful.\n\t1. The discussion on Gated RNNs in lines 205-220 lacks rigorous substantiation, relying primarily on speculative assertions to illustrate the impact of gating mechanisms.\n\t2. The paper does not clearly establish whether the proposed methodology is effective within a gated framework.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. It is good to see people discuss the complex eigenvalue setting as previous work seems to be limited to real eigenvalues (such as HiPPO)\n2. The teacher-student analysis is well-presented and clarifies the understanding of complex systems within this context.", "weaknesses": "1. The assumption of inputs being wide-sense stationarity seems strong. Does it include the setting such as the Cifar10 dataset? \n2. The argument of input normalization seems hand-waving. \n\t1. When $\\lambda \\to 1$, the forward pass and backward pass are exploding with a different speed: Forward is asymptotically $1/(1-\\lambda)$ while backward is (at least) asymptotically $1/(1-\\lambda)^2$. \n\t2. It should be proper to say normalization can relax one scale issue but it cannot achieve the scaling both at the same time. \n3. The paragraph \"On the Importance of Adaptive Learning Rates\" addresses an essential topic but lacks direct training curves.\n\t1. While the local Hessian analysis provided is valuable, it doesn’t assure improved global performance.\n\t2. It would be more effective and indeed necessary to include training curves comparing results using SGD and adaptive learning rates (such as AdamW). This would clearly illustrate the advantages and need for adaptive optimizers.\n4. Related works on parameterizations and orthogonal matrix is limited:\n\t1. https://proceedings.mlr.press/v80/zhang18g.html\n\t2. https://arxiv.org/abs/1909.09501: This work proposes a parameterised approach to train orthogonal matrix. \n\t3. https://arxiv.org/abs/1901.08428\n\t4. https://openreview.net/pdf?id=ryxepo0cFX\n\t5. https://arxiv.org/abs/2311.14495: This work proposes the relationship between parameterization and gradient in RNNs/SSMs. \n\t6. https://arxiv.org/pdf/2006.12070\n5. I am willing to increase the score if the above issues are addressed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their thorough feedback. We address their concerns below.\n\n**Weaknesses**\n\n1. **Wide-sense stationarity.** A discussion of this assumption was indeed missing. We address it in the global response. In short, this assumption is standard in the analysis of linear time filters (see e.g. the Proakis and Manolakis book) and is mostly for convenience in making calculations easier (see e.g. the classical Wiener–Khinchin theorem). When it is violated, our theory does not perfectly predict signal propagation, but our results hold qualitatively and almost quantitatively (c.f. Fig 1 and 3 of the PDF, which test our predictions in a practically relevant setting). This assumption is not a significant limitation of our theory, as using the empirical autocorrelation function leads to very similar results (see Fig 3 of the PDF, left col).\n2. **Input normalization.** The reviewer accurately points out that the forward and backward passes (for the $\\lambda$ parameters, but not for the inputs $x$) explode at different speeds and that normalization alone cannot solve it. Normalization is a way to ensure that signals, both hidden states and errors, stay constant over the network's depth. Therefore, its role is to keep the magnitude of the hidden states bounded. However, it does not entirely avoid the explosion of the recurrent parameter gradients, as hinted by the reviewer. This is where reparametrization is required. We hope that this clarifies our argument. If not, can the reviewer point out the source of confusion in this argument?\n3. **Adaptive learning rates.** The Hessian analysis indeed only captures the geometry of the loss landscape around optimality. Outside these regions, this only corresponds to part of the Hessian and includes some approximation. However, we would like to point out that this approximation is very standard, e.g., [Sagun et al. 2018](https://arxiv.org/pdf/1706.04454), and we find that the insights we gained by studying the Hessian at optimality hold for the rest of training.\nRegarding the training curves of SGD/Adam: This section was indeed missing some discussion of SGD. We found that it is almost impossible to train such recurrent networks with SGD. The Hessian at optimality gives us a hint as to why this happens: the Hessian on the recurrent parameters (e.g., $\\lambda$) is much larger than the one on feedforward parameters (e.g., $D$). As a result, to keep SGD stable, we need tiny learning rates, and the feedforward parameters barely change. The training curves for SGD are therefore almost flat. On the contrary, Adam uses adaptive LRs for different parameters, thereby avoiding this problem. We will mention this point and include some training curves in the next version of the paper.\n4. **References**. We thank the reviewer for these references.\n\n**Questions**\n\n1. **Why second-moment?** We focus on the second moment for multiple reasons. First, the first moment would be 0 as long as the input has a mean of 0, which is the traditional regime in deep learning. Second, this enables us to partially capture the norm of the different quantities at hand. Third, this quantity appears, up to some kind of data-dependent change of norm, in the Hessian, see e.g. Eq. 11. Finally, we would like to note that studying second-order quantities is standard in deep learning, a classical example is the seminal Glorot & Bengio initialization paper (2010).\nRegarding the comparison with Prop 3.3 of the LRU: our Eq. 4 is indeed very close to their result. Here, we generalize it to more general input distributions.\n2. **On complex numbers.** In our paper, we discuss both polar and real-imaginary parametrizations, which are arguably the most usual ones, and show that they suffer from explosion behaviors. We tried our best to find an \"optimal\" parametrization for the polar case, given that we know how to control the magnitude of the gradient w.r.t. $|\\lambda|$. However, we found that the most natural choices all have some issues (cf. A.3.2). It does not mean that it is impossible to find a better parametrization, just that we could not. We are happy to consider any suggestions that the reviewer may have.\n3. **On normalization.** In Fig 5, we show that overall gradient magnitudes remain bounded with layer normalization, hence “the overall gradient magnitude [is] under control”. However, this is not a proper solution as different hidden neurons share the same normalizers. It follows that a neuron with a high $\\lambda$ will heavily influence the normalization factor and thus neurons with smaller $\\lambda$ values will be ignored. We would also like to highlight that the LRU paper already includes a similar argument.\n\n**Limitations**\n\nFirst, we would like to point out how rich the learning dynamics of the simple time-invariant recurrent network are and that analyzing it in depth is not a trivial task. We are convinced that this analysis alone is already a strong contribution to the field. That said, the reviewer’s remark prompted us to analyze more carefully the gated case on our task of Sec 5. We study a GRU as it is the simplest gated architecture. We report those results in the global answer. Overall, we find that diagonality almost holds under standard initialization, for the reasons we mention in lines 205-220, and that our theory captures signal propagation in this regime. However, it is likely that, during training, the architecture becomes less linear and that our theory cannot capture signal propagation anymore. Such architectures have been studied using the mean field framework by Chen et al. 2018 [50]. Their result captures non-linearity but is arguably harder to interpret. We believe that our work provides a complementary and more intuitive understanding of signal propagation while being less general.\n\nWe hope that our rebuttal convinces the reviewer of the relevance and soundness of our work and that they may consider reevaluating our paper.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We are grateful to the reviewers for their high-quality feedback and their positive comments about our paper. We answer the main points below and respond to each reviewer more specifically in corresponding threads.

**Contributions-limitations.** Many reviewers highlighted that the contributions and the limitations of the paper should be made clearer. We list them below:

- **Contributions**
    - We identify, **for the first time**, a crucial issue arising in the training of recurrent neural networks. We study it analytically in great detail in a simplified, yet realistic, setting. We then highlight that **successful RNN architectures** all share a common feature: they **mitigate this issue**.
    - We confirm that our theory accurately **captures learning dynamics in a controlled teacher-student regime**. We additionally show how our analysis enables understanding parts of the Hessian structure, thereby improving our understanding of the critical differences between parametrizations/architectures and how they **interplay with different optimizers**.
    - Overall, our paper brings theoretical insights into the training of recurrent neural networks that are of **great practical relevance**. Those results are particularly important as the interest in RNNs has been rising in recent years and there is **still very little literature** on the topic.
    - Reviewers uPdc and xzKN rightly pointed out that we should better position our paper in comparison to the LRU paper. The **key distinction lies in the question we study**: the LRU paper aims to extract **which** components are critical to the great performance of SSMs, whereas we aim to understand **why** (some of) these components are critical. Their paper is therefore mostly empirical while ours is theoretical. That said, our work **naturally builds on the LRU**: First, we found their architecture to be more amenable to theoretical analysis than SSMs. Second, we designed our teacher setup task to reproduce the qualitative findings of the LRU paper on a more sophisticated task (a discussion of this point in Sec. 4 is missing). Third, our Eq. 5 generalizes their Eq. 6 to more general input distributions.
- **Limitations**
    - Our theoretical analysis makes three main assumptions that are reasonable, yet have some limitations.
        1. The recurrence is **diagonal and time-invariant**. Our analysis **does not exactly capture signal propagation in the gated case**. However, we provide new results (see below + PDF) demonstrating that our theory is a good proxy for gated RNNs at initialization.
        2. The sequences are **infinitely long.** In practical terms, our results only **apply whenever the sequence length is larger than the largest characteristic timescale of the network** (in our preliminary investigations we found x3 to be large enough).
        3. Input sequences are **wide-sense stationary.** This is a very **standard assumption** in the analysis of linear time-invariant filters. Yet, many real-world processes are not stationary, e.g. the Wikipedia data we are considering in Sec. 5. We provide new results (c.f. next paragraph) that show that our **analysis is still accurate enough** when this assumption is not met.
    - The technical tools we use for our theoretical analysis are rather elementary and there is little to no hope that they generalize to highly nonlinear dynamics.
    - Signal propagation that is independent of the time scale considered is necessary for efficient gradient-based learning but not sufficient. There are many questions that we thus cannot answer with such an analysis. For example, our theory cannot give any insight into how much memory a model can store, or tell us how learned networks will generalize to longer sequences than the training ones.

**New results.** Following the reviewers’ questions regarding 1) whether our analysis holds when the inputs are not wide-sense stationary and 2) how much it holds for gated RNNs, we provide a refined empirical analysis of signal propagation in the setting of Sec. 5.

1. The Wikipedia dataset is not wide-sense stationary and of finite length so we cannot directly apply the results of our analysis. Instead, we compute the empirical auto-correlation function (PDF Fig. 1). We find that we can approximate it well with one i.i.d. component ($\rho = 0$) and a slowly decaying one ($\rho \approx 1$). This is the autocorrelation function we plug into our analytical formulas when needed.
2. In Sec. 3, we argued that gated RNNs can behave like the diagonal networks we study theoretically. We verify this empirically in Fig. 2 of the PDF by looking at the diagonal and non-diagonal components of the recurrent Jacobian $\mathrm{d}h_t/ \mathrm{d}h_0$ of a GRU. We find that the non-diagonal components are indeed negligible in standard initialization schemes (this breaks when we multiply the hidden to hidden connections by a factor ~3).
3. We study signal propagation in a simplified GRU network (same as in Chen et al. 2018) and compare it with our theory (slightly modified to take into account normalization). We find that our theory almost perfectly captures the time-invariant case and adequately captures the gated case.

We are confident that our rebuttal addresses the concerns raised by the reviewers and hope that they will consequently update their scores. Given that we cannot update the paper, we have integrated those remarks, as well as fixed typos, in what will be the next version of the paper.

Meta-review: The paper analyses what they call the "curse of memory" in recurrent neural networks (a term borrowed from prior research), i.e., the fact that the variance of the activations can exponentially diverge with respect to the size of the memory for very long sequences. In a sense, the analysis is dual to the one over gradients, as both exploit long chain of Jacobians over the sequence time steps.

The paper had a very live rebuttal period, at the end of which it moved from a potential reject (3/4/4/5/6), to a scenario where one reviewer is recommending acceptance and all the other reviewers remain non-committal with borderline scores. Some parts of the discussion have been correctly summarized by the authors here on OpenReview. Before proceeding to the decision, I would like to weigh on some remaining concerns:

* Reviewer `zt6Z` is concerned that the analysis only involves the transition matrix of SSMs ($\mathbf{A}$ in common notation), and not a full-fledged model. In addition, they were concerned about the use of an artificial scenario to validate the results. From what I understand, both points were addressed but the reviewer is not convinced they can be integrated easily in a camera-ready version. After going through the paper carefully, I agree on this with the authors that both points are relatively minor modifications to the paper, and I respectfully disagree with the reviewer (who, in any case, voted for a borderline accept).

* Reviewer `uPdc` is concerned about the relation of the paper with the LRU paper by Orvieto et al. In their opinion, the paper should be made more "LRU-centric", and they are also concerned that "*LRUs are not mentioned at all until the experriment*". The authors have argued that this paper complements the LRU paper, since most of the techniques of LRU (e.g., input normalization) are introduced empirically, and they are motivated from a signal propagation point of view here. I note that the authors have been very careful with prior literature (e.g., page 3, "*Martens and Sutskever [40] hypothesized that such a phenomenon could arise in RNNs*"), and LRUs are indeed cited in Section 2.2a as an example of application of the paper. However, the analysis does hold for other types of architectures. Hence, after careful analysis, I side with the authors also on this point.

* Reviewer `Bq7Y` has a few concerns (e.g., typos, discussion of the limitations) which were solved in the rebuttal period. They are now in a position similar to `zt6Z` concerning a potential inability to fit the material for the camera ready. On this, I have the same opinion as for `zt6Z`, since most modifications appear to be small.

* Reviewer `xzKN` is concerned about the assumption of wide sense stationarity made by the authors. This does not invalidate their claim, but they propose a different condition that should be less restrictive. This does not appear to be especially critical with respect to the results, as the authors also argue.

Importantly, all reviewers agree on the novelty of the paper and the potential for further investigation for RNNs, both in the original reviews and during the rebuttal period. I believe this is a strong point in favor of the paper. In addition, my own evaluation is that the paper is very well written, easy to follow (despite being fundamentally a theoretical paper), well organized, and with clear and easy to read figures and experiments that complement the analysis. None of the modifications discussed above would modify this substantially, and I do not see a strong reason for waiting another review cycle for incorporating them. Hence, I suggest acceptance for NeurIPS due to the importance of the topic and the results.
Decision: Accept (poster)