Recurrent neural networks: vanishing and exploding gradients are not the end of the story

We thank the reviewer for their thorough feedback. We address their concerns below.

Weaknesses

1. 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).
2. 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?
3. 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, and we find that the insights we gained by studying the Hessian at optimality hold for the rest of training. Regarding 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.
4. References. We thank the reviewer for these references.

Questions

1. 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). Regarding 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.
2. 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.
3. 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.

Limitations

First, 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.

We hope that our rebuttal convinces the reviewer of the relevance and soundness of our work and that they may consider reevaluating our paper.

I have carefully reviewed both the reviews and the rebuttals.

However, I believe that the section on the wide-sense stationary assumption may not be essential for deriving the key results. Estimating hidden states could potentially be approached deterministically with bounded inputs, which might more effectively highlight the advantages of specific parameterizations, initializations, or normalizations.

While the inclusion of the stationary property facilitates a discussion of the second-order moment, this approach seems somewhat forced. The bound on the second-order moment within a random process context does not provide a significantly more valuable insight. Although generalizing to more diverse input distributions is a noteworthy contribution, it is crucial to clearly define the boundaries of this generalization. Without specifying the limits of input expansion, the contribution may not be fully strengthened. It remains unclear how this paper’s contribution stands apart from the LRU paper and what kinds of stronger understanding can be derived for complex model such as gated.

Therefore, I will increase my score from 4 to 5 and encourage the authors to address the issues related to the stationary assumption and to clarify the improvements over the LRU paper in the next version.

We appreciate the reviewer’s positive feedback on our paper and address their concerns below.

1. 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. Regarding 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. Regarding Figure 1: “as” is missing (”The loss becomes sharper as information is kept…”).
2. 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).
3. 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. Regarding 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. Regarding 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.

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

We appreciate the reviewer’s very positive feedback and their valuable input regarding these additional references and important details. We address their questions below.

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?

Indeed, $\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))$.

does that mean you simply read out ADAM’s (corrected) second-moment estimates?

The quantity we report incorporates both the first-moment and second-moment estimates from Adam, not just the second-moment one.

at which iteration/epoch do you query the effective learning rate? In 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.

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

We thank the reviewer for their review and address their concerns below.

On the contributions of the paper

What are the main contributions of the paper?

We provide a detailed list of our contributions in the global answer.

If the solution to the curse of memory is not the contribution of the paper, then the novelty and contribution become limited.

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

On the different architectures

What are other types of normalization and reparametrization, other than the ones employed by LRUs, that can address the curse of memory?

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

Meanwhile, LRUs employ the same solution discussed in the paper explicitly.

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.

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

On the writing

[…] it is unclear where this is shown.

This 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)$.

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

We thank the reviewer for their encouraging review. We answer their questions below:

1. 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”.
2. 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).
3. 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).
4. Would changing it to $(x_t)$ make it clearer? We want to emphasize that this is a sequence of numbers here.

5. Could the authors give an intuition on lines (113/116)?

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

Why an higher correlation in the input data should imply an higher variance in the model state

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

6. We do not use any $B$ mapping for the same reasons as described in 2.

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

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