Title: Gradient Flossing: Improving Gradient Descent through Dynamic Control of Jacobians
Venue: NeurIPS 2023
Accepted: yes
Keywords: exploding/vanishing gradients

Lyapunov exponents

Lyapunov spectrum

chaos

RNN

condition number

Jacobian
Abstract: Training recurrent neural networks (RNNs) remains a challenge due to the instability of gradients across long time horizons, which can lead to exploding and vanishing gradients. Recent research has linked these problems to the values of Lyapunov exponents for the forward-dynamics, which describe the growth or shrinkage of infinitesimal perturbations. Here, we propose gradient flossing, a novel approach to tackling gradient instability by pushing Lyapunov exponents of the forward dynamics toward zero during learning. We achieve this by regularizing Lyapunov exponents through backpropagation using differentiable linear algebra. This enables us to "floss" the gradients, stabilizing them and thus improving network training. We show that gradient flossing controls not only the gradient norm but also the condition number of the long-term Jacobian, facilitating multidimensional error feedback propagation. We find that applying gradient flossing before training enhances both the success rate and convergence speed for tasks involving long time horizons.
For challenging tasks, we show that gradient flossing during training can further increase the time horizon that can be bridged by backpropagation through time. Moreover, we demonstrate the effectiveness of our approach on various RNN architectures and tasks of variable temporal complexity. Additionally, we provide a simple implementation of our gradient flossing algorithm that can be used in practice. Our results indicate that gradient flossing via regularizing Lyapunov exponents can significantly enhance the effectiveness of RNN training and mitigate the exploding and vanishing gradients problem.

Gradient Flossing: Improving Gradient Descent through Dynamic Control of Jacobians

Abstract

Training recurrent neural networks (RNNs) remains a challenge due to the instability of gradients across long time horizons, which can lead to exploding and vanishing gradients. Recent research has linked these problems to the values of Lyapunov exponents for the forward-dynamics, which describe the growth or shrinkage of infinitesimal perturbations. Here, we propose gradient flossing , a novel approach to tackling gradient instability by pushing Lyapunov exponents of the forward dynamics toward zero during learning. We achieve this by regularizing Lyapunov exponents through backpropagation using differentiable linear algebra. This enables us to "floss" the gradients, stabilizing them and thus improving network training. We demonstrate that gradient flossing controls not only the gradient norm but also the condition number of the long-term Jacobian, facilitating multidimensional error feedback propagation. We find that applying gradient flossing prior to training enhances both the success rate and convergence speed for tasks involving long time horizons. For challenging tasks, we show that gradient flossing during training can further increase the time horizon that can be bridged by backpropagation through time. Moreover, we demonstrate the effectiveness of our approach on various RNN architectures and tasks of variable temporal complexity. Additionally, we provide a simple implementation of our gradient flossing algorithm that can be used in practice. Our results indicate that gradient flossing via regularizing Lyapunov exponents can significantly enhance the effectiveness of RNN training and mitigate the exploding and vanishing gradients problem.

1 Introduction

Recurrent neural networks are commonly used both in machine learning and computational neuroscience for tasks that involve input-to-output mappings over sequences and dynamic trajectories. Training is often achieved through gradient descent by the backpropagation of error information across time steps [1, 2, 3, 4]. This amounts to unrolling the network dynamics in time and recursively applying the chain rule to calculate the gradient of the loss with respect to the network parameters. Mathematically, evaluating the product of Jacobians of the recurrent state update describes how error signals travel across time steps. When trained on tasks that have long-range temporal dependencies, recurrent neural networks are prone to exploding and vanishing gradients [5, 6, 7, 8]. These arise from the exponential amplification or attenuation of recursive derivatives of recurrent network states over many time steps. Intuitively, to evaluate how an output error depends on a small parameter change at a much earlier point in time, the error information has to be propagated through the recurrent network states iteratively. Mathematically, this corresponds to a product of Jacobians that describe how changes in one recurrent network state depend on changes in the previous network state. Together, this product forms the long-term Jacobian. The singular value spectrum of the long-term Jacobian reg-

ulates how well error signals can propagate backwards along multiple time steps, allowing temporal credit assignment. A close mathematical correspondence of these singular values and the Lyapunov exponents of the forward dynamics was established recently [9, 10, 11, 12]. Lyapunov exponents characterize the asymptotic average rate of exponential divergence or convergence of nearby initial conditions and are a cornerstone of dynamical systems theory [13, 14]. We will use this link to improve the trainability of RNNs.

Previous approaches that tackled the problem of exploding or vanishing gradients have suggested solutions at different levels. First, specialized units such as LSTM and GRU were introduced, which have additional latent variables that can be decoupled from the recurrent network states via multiplicative (gating) interactions. The gating interactions shield the latent memory state, which can therefore transport information across multiple time steps [5, 6, 15]. Second, exploding gradients can be avoided by gradient clipping, which re-scales the gradient norm [16] or their individual elements [17] if they become too large [18]. Third, normalization schemes like batch normalization, layer norm and group norm prevent saturated nonlinearities that contribute to vanishing gradients [19, 20, 21]. Fourth, it was suggested that the problem of exploding/vanishing gradients can be ameliorated by specialized network architectures, for example, antisymmetric networks [22], orthogonal/unitary initializations [23, 24, 25], coupled oscillatory RNNs [26], Lipschitz RNNs [27], structured state space models [28, 29, 30, 31], echo state networks [32, 33], (recurrent) highway networks [34, 35], and stable limit cycle neural networks [11, 36, 37]. Fifth, for large networks, a suitable choice of weights can guarantee a well-conditioned Jacobian at initialization [23, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47]. These initializations are based on mean-field methods, which become exact only in the large-network limit. Such initialization schemes have also been suggested for gated networks [45]. However, even when initializing the network with well-behaved gradients, gradients will typically not retain their stability during training once the network parameters have changed.

Here, we propose a novel approach to tackling this challenge by introducing gradient flossing , a technique that keeps gradients well-behaved throughout training. Gradient flossing is based on a recently described link between the gradients of backpropagation through time and Lyapunov exponents, which are the time-averaged logarithms of the singular values of the long-term Jacobian [9, 11, 12, 37]. Gradient flossing regularizes one or several Lyapunov exponents to keep them close to zero during training. This improves not only the error gradient norm but also the condition number of the long-term Jacobian. As a result, error signals can be propagated back over longer time horizons. We first demonstrate that the Lyapunov exponents can be controlled during training by including an additional loss term. We then demonstrate that gradient flossing improves the gradient norm and effective dimension of the gradient signal. We find empirically that gradient flossing improves test accuracy and convergence speed on synthetic tasks over a range of temporal complexities. Finally, we find that gradient flossing during training further helps to bridge long-time horizons and show that it combines well with other approaches to ameliorate exploding and vanishing gradients, such as dynamic mean-field theory for initialization, orthogonal initialization and gated units. Our contributions include:

Gradient flossing , a novel approach to the problem of exploding and vanishing gradients in recurrent neural networks based on regularization of Lyapunov exponents. 1

Analytical estimates of the condition number of the long-term Jacobian based on Lyapunov exponents.

Empirical evidence that gradient flossing improves training on tasks that involve bridging long time horizons.

2 RNN Gradients and Lyapunov Exponents

We begin by revisiting the established mathematical relationship between the gradients of the loss function, computed via backpropagation through time, and Lyapunov exponents [9, 12], and how it relates to the problem of vanishing and exploding gradients. In backpropagation through time, network parameters θ are iteratively updated by stochastic gradient descent such that a loss L t is locally reduced [1, 2, 3, 4]. For RNN dynamics h s +1 = f θ ( h s , x s +1 ) , with recurrent network state h , external input x , and parameters θ , the gradient of the loss L t with respect to θ is evaluated by

unrolling the network dynamics in time. The resulting expression for the gradient is given by:

where T t ( h τ ) is composed of a product of one-step Jacobians D s = ∂ h s +1 ∂ h s :

Due to the chain of matrix multiplications in T t , the gradients tend to vanish or explode exponentially with time. This complicates training particularly when the task loss at time t dependents on inputs x or states h from many time steps prior which creates long temporal dependencies [5, 6, 7, 8]. How well error signals can propagate back in time is constrained by the tangent space dynamics along trajectory h t , which dictate how local perturbations around each point on the trajectory stretch, rotate, shear, or compress as the system evolves.

The singular values of the Jacobian's product T t , which determine how quickly gradients vanish or explode during backpropagation through time, are directly related to the Lyapunov exponents of the forward dynamics [9, 12]: Lyapunov exponents λ 1 ≥ λ 2 · · · ≥ λ N are defined as the asymptotic time-averaged logarithms of the singular values of the long-term Jacobian [13, 48, 49]

where σ i,t denotes the i th singular value of T t ( h τ ) with σ 1 ,t ≥ σ 2 ,t . . . σ N,t (See Appendix I for details). This means that positive Lyapunov exponents in the forward dynamics correspond to exponentially exploding gradient modes, while negative Lyapunov exponents in the forward dynamics correspond to exponentially vanishing gradient modes.

In summary, the Lyapunov exponents give the average asymptotic exponential growth rates of infinitesimal perturbations in the tangent space of the forward dynamics, which also constrain the signal propagation in backpropagation for long time horizons. Lyapunov exponents close to zero in the forward dynamics correspond to tangent space directions along which error signals are neither drastically attenuated nor amplified in backpropagation through time. Such close-to-neutral modes in the tangent dynamics can propagate information reliably across many time steps.

3 Gradient Flossing: Idea and Algorithm

We now leverage the mathematical connection established between Lyapunov exponents and the prevalent issue of exploding and vanishing gradients for regularizing the singular values of the long-term Jacobian. We term this procedure gradient flossing . To prevent exploding and vanishing gradients, we constrain Lyapunov exponents to be close to zero. This ensures that the corresponding directions in tangent space grow and shrink on average only slowly. This leads to a better-conditioned long-term Jacobian T t ( h τ ) . We achieve this by using the sum of the squares of the first k largest Lyapunov exponent λ 1 , λ 2 . . . λ k as a loss function:

and evaluate the gradient obtained from backpropagation through time:

This might seem like an ill-fated enterprise, as the gradient expression in Eq 5 suffers from its own problem of exploding and vanishing gradients. However, instead of calculating the Lyapunov exponents by directly evaluating the long-term Jacobian T t (Eq 2), we use an established iterative reorthonormalization method involving QR decomposition that avoids directly evaluating the illconditioned long-term Jacobian [12, 50].

First, we evolve an initially orthonormal system Q s = [ q 1 s , q 2 s , . . . q k s ] in the tangent space along the trajectory using the Jacobian D s = ∂ h s +1 ∂ h s . This means to calculate

at every time-step. Second, we extract the exponential growth rates using the QR decomposition,

which decomposes ˜ Q s +1 uniquely into the product of an orthonormal matrix Q s +1 of size N × k so Q glyph[latticetop] s +1 Q s +1 = 1 k × k and an upper triangular matrix R s +1 of size k × k with positive diagonal elements. Note that the QR decomposition does not have to be applied at every step, just sufficiently often, i.e., once every t ONS such that ˜ Q does not become ill-conditioned.

The Lyapunov exponents are given by time-averaged logarithms of the diagonal entries of R s [49, 50]:

This way, the Lyapunov exponent can be expressed in terms of a temporal average over the diagonal elements of the R s -matrix of a QR decomposition of the iterated Jacobian. To propagate the gradient of the square of the Lyapunov exponents backward through time in gradient flossing , we used an analytical expression for the pullback of the QR decomposition [51]: The backward pass of the QR decomposition is given by [51, 52, 53, 54]

where M = RR T -Q T Q and the copyltu function generates a symmetric matrix by copying the lower triangle of the input matrix to its upper triangle, with the element [ copyltu ( M )] ij = M max( i,j ) , min( i,j ) [51, 52, 53, 54]. We denote here adjoint variable as T = ∂ L /∂T . A simple implementation of this algorithm in pseudocode is:

Algorithm 1 Algorithm for gradient flossing of k tangent space directions

For clarity, we described gradient flossing in terms of stochastic gradient descent, but we actually implemented it with the ADAM optimizer using standard hyperparameters η , β 1 and β 2 . An example implementation in Julia [77] using Flux [78] is available here . Note that this algorithm also works for different recurrent network architectures. In this case, the Jacobians D has size n × n , where n is the number of dynamic variables of the recurrent network model. For example, in case of a single recurrent network of N LSTM units, the Jacobian has size 2 N × 2 N [9, 12, 46]. The Jacobian matrix D can either be calculated analytically or it can be obtained via automatic differentiation.

4 Gradient Flossing: Control of Lyapunov Exponents

In Fig 1, we demonstrate that gradient flossing can set one or several Lyapunov exponents to a target value via gradient descent with the ADAM optimizer in random Vanilla RNNs initialized with different weight variances. The N units of the recurrent neural network follow the dynamics

0

20

40

60

80

100

Epochs

10

5

10

3

10

1

1

k

k

i

=1

2

i

C

number of flossed

i

k = 32

k = 16

k = 1

T

t

(

h

τ

) =

t

-

1

∏

τ

′

=

τ

∂

h

τ

′

+1

∂

h

τ

′

0

10

20

30

40

50

60

70

80

Epochs

2.0

1.5

1.0

0.5

0.0

1(1/

)

B

target

1

actual

1

0.0

0.2

0.4

0.6

0.8

1.0

i/N

6

4

2

0

i

(1/

)

D

number of flossed

i

k = 32

k = 16

k = 1

The initial entries of W are drawn independently from a Gaussian distribution with zero mean and variance g 2 /N , where g is a gain parameter that controls the heterogeneity of weights. We here use the transfer function φ ( x ) = tanh( x ) . (See appendix B for gradient flossing with ReLU and LSTM units). x s is a sequence of inputs and V is the input weight. x s is a stream of i.i.d. Gaussian input x s ∼ N (0 , 1) and the input weights V are N (0 , 1) . Both W and V are trained during gradient flossing .

In Fig 1B, we show that for randomly initialized RNNs, the Lyapunov exponent can be modified by gradient flossing to match a desired target value. The networks were initialized with 10 different values of initial weight strength g chosen uniformly between 0 and 1. During gradient flossing , they quickly approached three different target values of the first Lyapunov exponents λ target 1 = {-1 , -0 . 5 , 0 } within less than 100 training epochs with batch size B = 1 . We note that gradient flossing with positive target λ target 1 seems not to arrive at a positive Lyapunov exponent λ 1 .

Fig 1C shows gradient flossing for different numbers of Lyapunov exponents k . Here, during gradientdescent, the sum of the squares of 1, 16, or 32 Lyapunov exponents is used as loss in gradient flossing (see Fig 1A). Fig 1D shows the Lyapunov spectrum after flossing, which now has 1, 16, or 32 Lyapunov exponents close to zero. We conclude that gradient flossing can selectively manipulate one, several, or all Lyapunov exponents before or during network training. Gradient flossing also works for RNNs of ReLU and LSTM units (See appendix B. Further, we find that the computational bottleneck of gradient flossing is the QR decomposition, which has a computational complexity of O ( N k 2 ) , both in the forward pass and in the backward pass. Thus, gradient flossing of the entire Lyapunov spectrum is computationally expensive. However, as we will show, not all Lyapunov exponents need to be flossed and only short episodes of gradient flossing are sufficient for significantly improving the training performance.

5 Gradient Flossing: Condition Number of the Long-Term Jacobian

0

100

200

300

time horizon

t

10 7

10 16

10 25

10 34

condition number

2(

Tt

(

))

A

initial

after flossing

theory

simulations

5

10

15

tangent space dimension

m

10 5

10 12

10 19

10 26

2(

Tt

(

))

B

initial

after flossing

theory

simulations

10 0

10 10

10 20

10 30

2 (numerical)

10 0

10 10

10 20

10 30

2 (theory)

C

k

5

10

15

A well-conditioned Jacobian is essential for efficient and fast learning [23, 55, 56]. Gradient flossing improves the condition number of the long-term Jacobian which constrains the error signal propagation across long time horizons in backpropagation (Fig 2). The condition number κ 2 of a linear map A measures how close the map is to being singular and is given by the ratio of the largest singular value σ max and the smallest singular values σ min, so κ 2 ( A ) = σ max ( A ) σ min ( A ) . According to the rule of thumb given in [57], if κ 2 ( A ) = 10 p , one can anticipate losing at least p digits of precision when solving the equation Ax = b . Note that the long-term Jacobian T t is composed of a product of Jacobians, which generically makes it ill-conditioned. To nevertheless quantify the condition number numerically, we use arbitrary-precision arithmetic with 256 bits per float. We find numerically that the condition number of T t exponentially diverges with the number of time steps (Fig 2A). We compare the numerically measured condition number κ 2 with an asymptotic approximation of the condition number based on Lyapunov exponents that are calculated in the forward pass and find a good match (Fig 2A).

Our theoretical estimate of the condition number κ 2 of an orthonormal system Q of size N × m that is temporally evolved by the long-term Jacobian T t is:

where σ 1 ( T t ( h τ )) and σ m ( T t ( h τ )) are the first and m th singular value of the long-term Jacobian. We note that this theoretical estimate of the condition number follows from the asymptotic definition of Lyapunov exponents and should be exact in the limit of long times. We find that gradient flossing reduces the condition number by a factor whose magnitude increases exponentially with time (orange in Fig 2A). Thus, we can expect that gradient flossing has a stronger effect on problems with a long time horizon to bridge. We will later confirm this numerically.

Moreover, Lyapunov exponents enable the estimation of the number of gradient dimensions available for the backpropagation of error signals. Generally, the long-term Jacobian is ill-conditioned, however, the Lyapunov spectrum provides for a given number of tangent space dimensions an estimate of the condition number. This indicates how close to singular the gradient signal for a given number of tangent space dimensions is. Given a fixed acceptable condition number-determined, for example, by noise level or floating-point precision-we observe that gradient flossing increases the number of usable tangent space dimensions for backpropagation (Fig 2B).

Finally, we show that the asymptotic estimate of the condition number based on Lyapunov exponents can even predict differences in condition number that originate from finite network size N (Fig 2C). We emphasize that this goes beyond mean-field methods, which become exact only in the largenetwork limit N →∞ and usually do not capture finite-size effects [58] (see appendix G).

6 Initial Gradient Flossing Improves Trainability

0

2000

4000

6000

8000

10000

Epochs

10

4

10

3

10

2

test loss

A

delayed copy task

no gradient flossing

gradient flossing

20

40

60

80

task difficulty (delay d)

10

4

10

3

10

2

test loss

C

delayed copy task

0

2000

4000

6000

8000

10000

Epochs

10

3

10

2

test loss

B

delayed XOR task

no gradient flossing

gradient flossing

10

20

30

40

50

60

70

80

task difficulty (delay d)

10

4

10

3

10

2

test loss

D

delayed XOR task

We next present numerical results on two tasks with variable spatial and temporal complexity, demonstrating that gradient flossing before training improves the trainability of Vanilla RNNs. We call gradient flossing before training in the following preflossing. For preflossing, we first initialize the network randomly, then minimize L flossing = ∑ k i =1 λ 2 i using the ADAM optimizer and subsequently train on the tasks. We deliberately do not use sequential MNIST or similar toy tasks commonly used to probe exploding/vanishing gradients, because we want a task where the structure of long-range dependencies in the data is transparent and can be varied as desired.

First, we consider the delayed copy task, where a scalar stream of random input numbers x must be reproduced by the output y delayed by d time steps, i.e. y t = x t -d . Although the task itself is trivial and can be solved even by a linear network through a delay line (see appendix E), RNNs encounter vanishing gradients for large delays d during training even with 'critical' initialization with g = 1 . Our experiments show that gradient flossing can substantially improve the performance of RNNs on this task (Fig 3A, C). While Vanilla RNNs without gradient flossing fail to train reliably beyond d = 20 , Vanilla RNNs with gradient flossing can be reliably trained for d = 40 (Fig 3C). Note that we flossed here k = 40 Lyapunov exponents before training. We will later investigate the role of the number of flossed Lyapunov exponents.

Second, we consider the temporal XOR task, which requires the RNN to perform a nonlinear inputoutput computation on a sequential stream of scalar inputs, i.e., y t = | x t -d / 2 -x t -d | , where d denotes a time delay of d time steps (For details see appendix H). Fig 3D demonstrates that gradient flossing helps to train networks on a substantially longer delay d . We found similar improvements through gradient flossing for RNNs initialized with orthogonal weights (see appendix G).

7 Gradient Flossing During Training

0

2000

4000

6000

8000

10000

Epochs

50

60

70

80

90

100

test accuracy (%)

A

delayed temporal XOR

flossing during training

preflossing

no flossing

10

20

30

40

50

60

70

complexity (delay d)

50

60

70

80

90

100

test accuracy (%)

C

0

2000

4000

6000

8000

10000

Epochs

50

60

70

80

90

100

test accuracy (%)

B

delayed spatial XOR

10

20

30

40

50

60

70

complexity (delay d)

50

60

70

80

90

100

test accuracy (%)

D

flossing during training

preflossing

no flossing

A) Test accuracy for Vanilla RNNs trained on delayed temporal binary XOR task y t = x t -d/ 2 ⊕ x t -d with gradient flossing during training (green), preflossing ( gradient flossing before training) (orange), and with no gradient flossing (blue) for d = 70 . Solid lines are mean across 20 network realizations, individual network realizations shown in transparent fine lines. B) Same as A for delayed spatial XOR task with y t = x 1 t -d ⊕ x 2 t -d ⊕ x 3 t -d . Parameters ( g = 1 , batch size b = 16 ). C) Test accuracy as a function of task difficulty (delay d ) for delayed temporal XOR task. D) Test accuracy as a function of task difficulty (delay d ) for delayed spatial XOR task. Parameters: g = 1 , batch size b = 16 , N = 80 , epochs = 10 4 , T = 300 , gradient flossing for E f = 500 epochs on k = 75 before training and during training for green lines, and only before training for orange lines. Same plotting conventions as previous figure. Task loss: cross-entropy between y and ˆ y .

We next investigate the effects of gradient flossing during the training and find that gradient flossing during training can further improve trainability. We trained RNNs on two more challenging tasks with variable temporal complexity and performed gradient flossing either both during and before training, only before training, or not at all.

Fig 4A shows the test accuracy for Vanilla RNNs training on the delayed temporal XOR task y t = x t -d/ 2 ⊕ x t -d with random Bernoulli process x ∈ { 0 , 1 } . The accuracy of Vanilla RNNs falls to chance level for d ≥ 40 (Fig 4C). With gradient flossing before training, the trainability can be improved, but still goes to chance level for d = 70 . In contrast, for networks with gradient flossing during training, the accuracy is improved to > 80% at d = 70 . In this case, we preflossed for 500 epochs before task training and again after 500 epochs of training on the task. In Fig 4B, D the networks have to perform the nonlinear XOR operation y t = x 1 t -d ⊕ x 2 t -d ⊕ x 3 t -d on a three-dimensional binary input signal x 1 , x 2 , and x 3 and generate the correct output with a delay of d steps. While the solution of the task itself is not difficult and could even be implemented by hand (see appendix), the task is challenging for backpropagation through time because nonlinear temporal associations bridging long time horizons have to be formed. Again, we observe that gradient flossing before training improves the performance compared to baseline, but starts failing for long delays d > 60 . In contrast, networks that are also flossed during training can solve even more difficult tasks (Fig 4D). We find that after gradient flossing , the norm of the error gradient with respect to initial conditions h 0 is amplified (appendix C). Interestingly, gradient flossing can also be detrimental to task performance if it is continued throughout all training epochs (appendix C)

Wenote that merely regularizing the spectral radius of the recurrent weight matrix W or the individual one-step Jacobians D s numerically or based on mean-field theory does not yield such a training improvement. This suggests that taking the temporal correlations between Jacobians D s into account is important for improving trainability.

7.1 Gradient Flossing for Different Numbers of Flossed Lyapunov Exponents

We investigated how many Lyapunov exponents k have to be flossed to achieve an improvement in training success (Fig 5). We studied this in the binary temporal delayed XOR task with gradient flossing during training (same as Fig 3) and varied the task difficulty by changing the delay d .

We found that as the task becomes more difficult, networks where not enough Lyapunov exponents k are flossed begin to fall below 100% test accuracy (Fig 5A). Correspondingly, when measuring final test accuracy as a function of the number of flossed Lyapunov exponents, we observed that more Lyapunov exponent k have to be flossed to achieve 100% accuracy as the tasks become more difficult (Fig 5B). We also show the entire parameter plane of median test accuracy as a function of both number of flossed Lyapunov exponents k and task difficulty (delay d ), and found the same trend (Fig 5B). Overall, we found that tasks with larger delay d require more Lyapunov exponents close to zero. We note that this might also partially be caused by the 'streaming' nature of the task: in our tasks, longer delays automatically imply that more values have to be stored as at any moment all the values in the 'delay line' have to be remembered to successfully solve the tasks. This is different from tasks where a single variable has to be stored and recalled after a long delay. It would be interesting to study tasks where the number of delay steps and the number of items in memory can be varied independently.

Finally, we did the same analysis on networks with only preflossing ( gradient flossing before training) and found the same trend (supplement Fig 7D), however, in that case even if all N Lyapunov exponents were flossed, thus k = N , they were not able to solve the most difficult tasks. This seems to indicate that gradient flossing during training cannot be replaced by just gradient flossing more Lyapunov exponents before training.

8 Limitations

The mathematical connection between Lyapunov exponents and backpropagation through time exploited in gradient flossing is rigorously established only in the infinite-time limit. It would be interesting to extend our analysis to finite-time Lyapunov exponents.

Furthermore, the backpropagation through time gradient involves a sum over products of Jacobians of different time periods t -τ , but the Lyapunov exponent only considers the asymptotic longest product. Additionally, Lyapunov exponents characterize the asymptotic dynamics on the attractor of the dynamics, whereas RNNs often exploit transient dynamics from some initial conditions outside or towards the attractor.

Although our proposed method focuses on exploiting Lyapunov exponents, it neglects the geometry of covariant Lyapunov vectors [59], which could be used to improve training performance, speed, and reliability. Additionally, it is important to investigate how sensitive the method is to the choice of orthonormal basis employed because it is only guaranteed to become unique asymptotically [60].

Finally, the computational cost of our method scales with O ( Nk 2 ) , where N is the network size and k is the number of Lyapunov exponents calculated. To reduce the computational cost, we suggest doing QR decomposition only sufficiently often to ensure that the orthonormal system is not ill-conditioned and using gradient flossing only intermittently or as pretraining. One could also calculate the Lyapunov spectrum for a shorter time interval or use a cheaper proxy for the Lyapunov spectrum and investigate more efficient gradient flossing schedules.

9 Discussion

We tackle the problem of gradient signal propagation in recurrent neural networks through a dynamical systems lens. We introduce a novel method called gradient flossing that addresses the problem of gradient instability during training. Our approach enhances gradient signal stability both before and during training by regularizing Lyapunov exponents. By keeping the long-term Jacobian well-conditioned, gradient flossing optimizes both training accuracy and speed. To achieve this, we combine established dynamical systems methods for calculating Lyapunov exponents with an analytical pullback of the QR factorization. This allows us to establish and maintain gradient stability in a in a manner that is memory-efficient, numerically stable, and exact across long time horizons. Our method is applicable to arbitrary RNN architectures, nonlinearities, and also neural ODEs [61]. Empirically, pre-training with gradient flossing enhances both training speed and accuracy. For difficult temporal credit assignment problems, gradient flossing throughout training further enhances signal propagation. We also demonstrate the versatility of our method on a set of synthetic tasks with controllable time-complexity and show that it can be combined with other approaches to tackle exploding and vanishing gradients, such as dynamic mean-field theory for initialization, orthogonal initialization and specialized single units, such as LSTMs.

Prior research on exploding and vanishing gradients mainly focused on selecting network architectures that are less prone to exploding/vanishing gradients or finding parameter initializations that provide well-conditioned gradients at least at the beginning of training. Our introduced gradient flossing can be seen as a complementary approach that can further enhance gradient stability throughout training. Compared to the work on picking good parameter initializations based on random matrix theory [46] and mean-field heuristics [45], gradient flossing provides several improvements: First, mean-field theory only considers the gradient flow at initialization, while gradient flossing can maintain gradient flow and well-conditioned Jacobians throughout the training process. Second, random matrix theory and mean-field heuristics are usually confined to the limit of large networks [58], while gradient flossing can be used for networks of any size. The link between Lyapunov exponents and the gradients of backpropagation through time has been described previously [9, 12] and has been spelled out analytically and studied numerically [10, 11, 62, 63, 64]. In contrast, we use Lyapunov exponents here not only as a diagnostic tool for gradient stability but also to show that they can directly be part of the cure for exploding and vanishing gradients.

Future investigations could delve further into the roles of the second to N th Lyapunov exponents in trainability, and how it is related to the task at hand, the rank of the parameter update, the dimensionality of the solution space, as well as the network dynamics (see also [37, 65, 66]). Our results suggest a trade-off between trainability across long time horizons and the nonlinear task demands that is worth exploring in more detail (appendix C). Applying gradient flossing to real-time recurrent learning and its biologically plausible variants is another avenue [67]. Extending gradient flossing to feedforward networks, state-space models and transformers is a promising avenue for future research (see also [31, 47, 68, 69, 70, 71, 72]). While Lyapunov exponents are only strictly defined for dynamical systems, such as maps or flows that are endomorphisms, the long-term Jacobian of deep feedforward networks could be treated similarly. This could also provide a link between the stability of the network against adversarial examples and its dynamic stability, as measured by Lyapunov exponents. Given that time-varying input can suppress chaos in recurrent networks [9, 12, 73, 74, 75, 76], we anticipate they may exacerbate vanishing gradients. Gradient flossing could also be applied in neural architecture search, to identify and optimize trainable networks. Finally, gradient flossing is applicable to other model parameters, as well. For instance, gradients of Lyapunov exponents with respect to single-unit parameters could optimize the activation function and single-neuron biophysics in biologically plausible neuron models.

Acknowledgments and Disclosure of Funding

I thank E. Izhikevich, F. Wolf, S. Goedeke, J. Lindner, L.F. Abbott, L. Logiaco, M. Schottdorf, G. Wayne and P. Sokol for fruitful discussions and J. Stone, L.F. Abbott, M. Ding, O. Marshall, S. Goedeke, S. Lippl, M.P. Puelma-Touzel, J. Lindner and the reviewers for feedback on the manuscript. I thank Jinguo Liu for the Julia package B ackwardsLinalg.jl and also the developers of Julia and Flux [77, 78]. Research supported by NSF NeuroNex Award (DBI-1707398), the Gatsby Charitable Foundation (GAT3708), the Simons Collaboration for the Global Brain (542939SPI), and the Swartz Foundation (2021-6).

References

P. WERBOS. Beyond Regression :. Ph. D. dissertation, Harvard University , 1974.

DB Parker. Learning-logic (TR-47). Center for Computational Research in Economics and Management Science. MIT-Press, Cambridge, Mass , 8, 1985.

Y. LECUN. Une procedure d'apprentissage ponr reseau a seuil asymetrique. Proceedings of Cognitiva 85 , pages 599-604, 1985.

David E. Rumelhart, Geoffrey E. Hinton, and Ronald J. Williams. Learning representations by backpropagating errors. Nature , 323(6088):533, October 1986.

Sepp Hochreiter. Untersuchungen zu dynamischen neuronalen Netzen . PhD thesis, 1991.

Sepp Hochreiter and Jürgen Schmidhuber. Long Short-Term Memory. Neural Computation , 9(8):17351780, 1997.

Y. Bengio, P. Simard, and P. Frasconi. Learning long-term dependencies with gradient descent is difficult. IEEE Transactions on Neural Networks , 5(2):157-166, March 1994.

Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio. On the difficulty of training Recurrent Neural Networks. arXiv:1211.5063 [cs] , November 2012. arXiv: 1211.5063.

Rainer Engelken, Fred Wolf, and L. F. Abbott. Lyapunov spectra of chaotic recurrent neural networks. arXiv:2006.02427 [nlin, q-bio] , June 2020. arXiv: 2006.02427.

Jonas Mikhaeil, Zahra Monfared, and Daniel Durstewitz. On the difficulty of learning chaotic dynamics with RNNs. Advances in Neural Information Processing Systems , 35:11297-11312, December 2022.

Il Memming Park, Ábel Ságodi, and Piotr Aleksander Sokół. Persistent learning signals and working memory without continuous attractors. ArXiv , page arXiv:2308.12585v1, August 2023.

Rainer Engelken, Fred Wolf, and L. F. Abbott. Lyapunov spectra of chaotic recurrent neural networks. Physical Review Research , 5(4):043044, October 2023.

J. P. Eckmann and D. Ruelle. Ergodic theory of chaos and strange attractors. Reviews of Modern Physics , 57(3):617-656, July 1985.

Arkady Pikovsky and Antonio Politi. Lyapunov Exponents: A Tool to Explore Complex Dynamics . Cambridge University Press, Cambridge, February 2016.

Kyunghyun Cho, Bart van Merrienboer, Dzmitry Bahdanau, and Yoshua Bengio. On the Properties of Neural Machine Translation: Encoder-Decoder Approaches. Technical report, September 2014. ADS Bibcode: 2014arXiv1409.1259C Type: article.

Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio. On the difficulty of training recurrent neural networks. In Proceedings of the 30th International Conference on International Conference on Machine Learning - Volume 28 , ICML'13, pages III-1310-III-1318, Atlanta, GA, USA, June 2013. JMLR.org.

Tomáš Mikolov. Statistical language models based on neural networks . Ph.D. thesis, Brno University of Technology, Faculty of Information Technology, Brno, CZ, 2012.

Ian Goodfellow, Yoshua Bengio, and Aaron Courville. Deep Learning . MIT Press, November 2016. Google-Books-ID: omivDQAAQBAJ.

Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer Normalization, July 2016. arXiv:1607.06450 [cs, stat].

Sergey Ioffe and Christian Szegedy. Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift. In Proceedings of the 32nd International Conference on Machine Learning , pages 448-456. PMLR, June 2015.

Yuxin Wu and Kaiming He. Group Normalization. pages 3-19, 2018.

Bo Chang, Minmin Chen, Eldad Haber, and Ed H. Chi. AntisymmetricRNN: A Dynamical System View on Recurrent Neural Networks. International Conference on Learning Representations, December 2018.

Andrew M. Saxe, James L. McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. arXiv:1312.6120 [cond-mat, q-bio, stat] , December 2013. arXiv: 1312.6120.

Martin Arjovsky, Amar Shah, and Yoshua Bengio. Unitary Evolution Recurrent Neural Networks. In Proceedings of The 33rd International Conference on Machine Learning , pages 1120-1128. PMLR, June 2016.

Kyle Helfrich, Devin Willmott, and Qiang Ye. Orthogonal Recurrent Neural Networks with Scaled Cayley Transform. In Proceedings of the 35th International Conference on Machine Learning , pages 1969-1978. PMLR, July 2018.

T. Konstantin Rusch and Siddhartha Mishra. Coupled Oscillatory Recurrent Neural Network (coRNN): An accurate and (gradient) stable architecture for learning long time dependencies. arXiv e-prints , page arXiv:2010.00951, October 2020.

N. Benjamin Erichson, Omri Azencot, Alejandro Queiruga, Liam Hodgkinson, and Michael W. Mahoney. Lipschitz Recurrent Neural Networks. International Conference on Learning Representations, January 2021.

Albert Gu, Karan Goel, and Christopher Ré. Efficiently Modeling Long Sequences with Structured State Spaces, August 2022. arXiv:2111.00396 [cs].

Jimmy T. H. Smith, Andrew Warrington, and Scott W. Linderman. Simplified State Space Layers for Sequence Modeling, March 2023. arXiv:2208.04933 [cs].

Antonio Orvieto, Samuel L Smith, Albert Gu, Anushan Fernando, Caglar Gulcehre, Razvan Pascanu, and Soham De. Resurrecting Recurrent Neural Networks for Long Sequences. Technical report, March 2023. ADS Bibcode: 2023arXiv230306349O Type: article.

Albert Gu and Tri Dao. Mamba: Linear-Time Sequence Modeling with Selective State Spaces, December 2023. arXiv:2312.00752 [cs].

Herbert Jaeger. The 'echo state' approach to analysing and training recurrent neural networks-with an erratum note. Bonn, Germany: German National Research Center for Information Technology GMD Technical Report , 148:34, 2001.

Mustafa C. Ozturk, Dongming Xu, and José C. Príncipe. Analysis and Design of Echo State Networks. Neural Computation , 19(1):111-138, January 2007.

Rupesh K Srivastava, Klaus Greff, and Jürgen Schmidhuber. Training Very Deep Networks. In Advances in Neural Information Processing Systems , volume 28. Curran Associates, Inc., 2015.

Julian Georg Zilly, Rupesh Kumar Srivastava, Jan Koutnik, and Jürgen Schmidhuber. Recurrent Highway Networks. In Proceedings of the 34th International Conference on Machine Learning , pages 4189-4198. PMLR, July 2017.

Piotr A. Sokół, Ian Jordan, Eben Kadile, and Il Memming Park. Adjoint Dynamics of Stable Limit Cycle Neural Networks. In 2019 53rd Asilomar Conference on Signals, Systems, and Computers , pages 884-887, November 2019. ISSN: 2576-2303.

Piotr A. Sokół. Geometry of Learning and Representations in Neural Networks . PhD thesis, Stony Brook University, May 2023.

Xavier Glorot, Antoine Bordes, and Yoshua Bengio. Deep Sparse Rectifier Neural Networks. volume 15, pages 315-323, April 2011.

Ben Poole, Subhaneil Lahiri, Maithra Raghu, Jascha Sohl-Dickstein, and Surya Ganguli. Exponential expressivity in deep neural networks through transient chaos. arXiv:1606.05340 [cond-mat, stat] , June 2016. arXiv: 1606.05340.

Minmin Chen, Jeffrey Pennington, and Samuel S. Schoenholz. Dynamical Isometry and a Mean Field Theory of RNNs: Gating Enables Signal Propagation in Recurrent Neural Networks. arXiv:1806.05394 [cs, stat] , August 2018. arXiv: 1806.05394.

Boris Hanin and Mihai Nica. Products of Many Large Random Matrices and Gradients in Deep Neural Networks. December 2018.

Samuel S. Schoenholz, Justin Gilmer, Surya Ganguli, and Jascha Sohl-Dickstein. Deep Information Propagation. November 2016.

Boris Hanin and David Rolnick. How to Start Training: The Effect of Initialization and Architecture. In Advances in Neural Information Processing Systems , volume 31. Curran Associates, Inc., 2018.

Piotr A. Sokol and Il Memming Park. Information Geometry of Orthogonal Initializations and Training. Technical report, October 2018. ADS Bibcode: 2018arXiv181003785S Type: article.

Dar Gilboa, Bo Chang, Minmin Chen, Greg Yang, Samuel S. Schoenholz, Ed H. Chi, and Jeffrey Pennington. Dynamical Isometry and a Mean Field Theory of LSTMs and GRUs. January 2019.

Tankut Can, Kamesh Krishnamurthy, and David J. Schwab. Gating creates slow modes and controls phase-space complexity in GRUs and LSTMs. arXiv:2002.00025 [cond-mat, stat] , January 2020. arXiv: 2002.00025.

Darshil Doshi, Tianyu He, and Andrey Gromov. Critical Initialization of Wide and Deep Neural Networks using Partial Jacobians: General Theory and Applications. November 2023.

Valery Iustinovich Oseledets. A multiplicative ergodic theorem. Characteristic Ljapunov, exponents of dynamical systems. Trudy Moskovskogo Matematicheskogo Obshchestva , 19:179-210, 1968.

Karlheinz Geist, Ulrich Parlitz, and Werner Lauterborn. Comparison of Different Methods for Computing Lyapunov Exponents. Progress of Theoretical Physics , 83(5):875-893, May 1990.

Giancarlo Benettin, Luigi Galgani, Antonio Giorgilli, and Jean-Marie Strelcyn. Lyapunov Characteristic Exponents for smooth dynamical systems and for hamiltonian systems; A method for computing all of them. Part 2: Numerical application. Meccanica , 15(1):21-30, March 1980.

Hai-Jun Liao, Jin-Guo Liu, Lei Wang, and Tao Xiang. Differentiable Programming Tensor Networks. Physical Review X , 9(3):031041, September 2019. arXiv:1903.09650 [cond-mat, physics:quant-ph].

S. F. Walter and L. Lehmann. Algorithmic Differentiation of Linear Algebra Functions with Application in Optimum Experimental Design (Extended Version). Technical report, January 2010. ADS Bibcode: 2010arXiv1001.1654W Type: article.

Robert Schreiber and Charles Van Loan. A Storage-Efficient $WY$ Representation for Products of Householder Transformations. SIAM Journal on Scientific and Statistical Computing , 10(1):53-57, January 1989.

Matthias Seeger, Asmus Hetzel, Zhenwen Dai, Eric Meissner, and Neil D. Lawrence. Auto-Differentiating Linear Algebra. Technical report, October 2017. ADS Bibcode: 2017arXiv171008717S Type: article.

Jeffrey Pennington, Samuel Schoenholz, and Surya Ganguli. Resurrecting the sigmoid in deep learning through dynamical isometry: theory and practice. In Advances in Neural Information Processing Systems , volume 30. Curran Associates, Inc., 2017.

Jeffrey Pennington, Samuel S. Schoenholz, and Surya Ganguli. The Emergence of Spectral Universality in Deep Networks. arXiv:1802.09979 [cs, stat] , February 2018. arXiv: 1802.09979.

E. Cheney and David Kincaid. Numerical Mathematics and Computing . Cengage Learning, August 2007. Google-Books-ID: ZUfVZELlrMEC.

Yasaman Bahri, Jonathan Kadmon, Jeffrey Pennington, Sam S. Schoenholz, Jascha Sohl-Dickstein, and Surya Ganguli. Statistical Mechanics of Deep Learning. Annual Review of Condensed Matter Physics , 11(1):501-528, 2020.

F. Ginelli, P. Poggi, A. Turchi, H. Chaté, R. Livi, and A. Politi. Characterizing Dynamics with Covariant Lyapunov Vectors. Physical Review Letters , 99(13):130601, September 2007.

Sergey V. Ershov and Alexei B. Potapov. On the concept of stationary Lyapunov basis. Physica D: Nonlinear Phenomena , 118(3):167-198, July 1998.

Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural Ordinary Differential Equations. In Advances in Neural Information Processing Systems , volume 31. Curran Associates, Inc., 2018.

Javed Lindner. Investigating the exploding and vanishing gradients problem with Lyapunov exponents. Master's thesis, RWTH Aaachen, Aaachen/Juelich, 2021.

Ryan Vogt, Maximilian Puelma Touzel, Eli Shlizerman, and Guillaume Lajoie. On Lyapunov Exponents for RNNs: Understanding Information Propagation Using Dynamical Systems Tools. Frontiers in Applied Mathematics and Statistics , 8, 2022.

Kamesh Krishnamurthy, Tankut Can, and David J. Schwab. Theory of Gating in Recurrent Neural Networks. Physical Review X , 12(1):011011, January 2022.

L. S. Pontryagin. Mathematical Theory of Optimal Processes: The Mathematical Theory of Optimal Processes . Routledge, New York, 1st edition edition, March 1987.

Daniel Liberzon. Calculus of Variations and Optimal Control Theory: A Concise Introduction. In Calculus of Variations and Optimal Control Theory . Princeton University Press, December 2011.

Owen Marschall, Kyunghyun Cho, and Cristina Savin. A unified framework of online learning algorithms for training recurrent neural networks. The Journal of Machine Learning Research , 21(1):135:5320135:5353, January 2020.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is All you Need. In Advances in Neural Information Processing Systems , volume 30. Curran Associates, Inc., 2017.

Judy Hoffman, Daniel A. Roberts, and Sho Yaida. Robust Learning with Jacobian Regularization. Technical report, August 2019. ADS Bibcode: 2019arXiv190802729H Type: article.

Hongyi Zhang, Yann N. Dauphin, and Tengyu Ma. Fixup Initialization: Residual Learning Without Normalization. Technical report, January 2019. ADS Bibcode: 2019arXiv190109321Z Type: article.

Thomas Bachlechner, Bodhisattwa Prasad Majumder, Henry Mao, Gary Cottrell, and Julian McAuley. ReZero is all you need: fast convergence at large depth. In Proceedings of the Thirty-Seventh Conference on Uncertainty in Artificial Intelligence , pages 1352-1361. PMLR, December 2021. ISSN: 2640-3498.

Tianyu He, Darshil Doshi, and Andrey Gromov. AutoInit: Automatic Initialization via Jacobian Tuning, June 2022. arXiv:2206.13568 [cond-mat, stat].

L. Molgedey, J. Schuchhardt, and H. G. Schuster. Suppressing chaos in neural networks by noise. Physical Review Letters , 69(26):3717-3719, December 1992.

Kanaka Rajan, L. F. Abbott, and Haim Sompolinsky. Stimulus-dependent suppression of chaos in recurrent neural networks. Physical Review E , 82(1):011903, July 2010.

Jannis Schuecker, Sven Goedeke, David Dahmen, and Moritz Helias. Functional methods for disordered neural networks. arXiv:1605.06758 [cond-mat, q-bio] , May 2016. arXiv: 1605.06758.

Rainer Engelken, Alessandro Ingrosso, Ramin Khajeh, Sven Goedeke, and L. F. Abbott. Input correlations impede suppression of chaos and learning in balanced firing-rate networks. PLOS Computational Biology , 18(12):e1010590, December 2022.

J. Bezanson, A. Edelman, S. Karpinski, and V. Shah. Julia: A Fresh Approach to Numerical Computing. SIAM Review , 59(1):65-98, January 2017.

Innes. Flux: Elegant machine learning with Julia. Journal of Open Source Software , 3(25):602, May 2018.

Edward Ott. Chaos in Dynamical Systems . Cambridge University Press, August 2002. Google-Books-ID: PfXoAwAAQBAJ.

Angelo Vulpiani, Fabio Cecconi, and Massimo Cencini. Chaos: From Simple Models to Complex Systems . World Scientific Pub Co Inc, Hackensack, NJ, September 2009.

A Backpropagation Through QR Decomposition

The backward pass of the QR decomposition is given by [51, 52, 53, 54]

where M = RR T -Q T Q and the copyltu function generates a symmetric matrix by copying the lower triangle of the input matrix to its upper triangle, with the element [ copyltu ( M )] ij = M max( i,j ) , min( i,j ) [51, 52, 53, 54]. Adjoint variable are written here as T = ∂ L /∂T .

Using an analytical pullback is more memory-efficient and less computationally costly than directly doing automatic differentiation through the QR-decomposition. Moreover, from a practical perspective, for QR decomposition, often BLAS/LAPACK routines are utilized which are not amenable to common differentiable programming frameworks like TensorFlow, PyTorch, JAX and Zygote. In our implementation of gradient flossing , we used the Julia package B ackwardsLinalg.jl by Jinguo Liu available at here .

B Further Details and Analysis of Gradient Flossing

An example implementation of gradient flossing in Flux [78], a machine learning library in Julia [77] is available at https://github.com/RainerEngelken/GradientFlossing. We are actively developing implementations for other widely used differentiable programming frameworks.

B.1 Gradient Flossing for recurrent LSTM and ReLU networks

0

20

40

60

80

100

Epochs

0.4

0.2

0.0

0.2

1

A

LSTM

target

1

actual

1

0

20

40

60

80

100

Epochs

10

6

10

5

10

4

10

3

10

2

10

1

2

1

B

0

20

40

60

80

100

Epochs

2

0

2

4

6

1

C

ReLU

target

1

actual

1

0

20

40

60

80

100

Epochs

10

5

10

3

10

1

10 1

2

1

D

We demonstrate that gradient flossing can also be applied to recurrent LSTM and ReLU networks in Fig 6. To this end, we generated random LSTM networks where the weights of all the different gates and biases were independently and identically distributed (i.i.d.) and sampled from Gaussian distributions of different variance. Our results show that gradient flossing can also constrain the Lyapunov exponent to be close to zero. The dynamics of each of the N LSTM units follows the map [6]:

where glyph[circledot] denotes the Hadamard product, σ g ( x ) = 1 1+exp( -x ) is the sigmoid function, σ h ( x ) = tanh( x ) and entries of the matrices U x are drawn from U x ∼ N (0 , g 2 x /N ) . For simplicity, the bias terms b x are scalars. Subscripts f , o and i denote respectively the forget gate, the output gate, the input gate, and c is the cell state. In each LSTM unit, there are two dynamic variables c and h , and three gates f , o , and i that control the flow if signals into and out of the cell c . We set the values g ih , g ix , g fx , b f , g ch , g cx , g cx , g ox to be uniformly distributed between 0 and 1 and initialize b i , g fh , b c , b 0 as zero.

During gradient flossing , the actual Lyapunov exponents of different random network realizations converge close to the target Lyapunov exponent λ target 1 = 0 in fewer than 100 epochs as shown in Fig 6A. Fig 6B shows that the squared Lyapunov exponents converge towards zero. We note that for LSTM networks, a target Lyapunov exponent of λ target 1 = -1 is achieved after 100 gradient flossing steps only for a subset of random network realizations (not shown). We speculate that behavior is influenced by the gating structure of LSTM units, which seems to naturally place the first Lyapunov exponent close to zero for certain initializations (See also [9, 12, 46, 64]).

For the recurrent ReLU networks, we considered the same Vanilla RNN dynamics as in the main manuscript in Eq 9

The initial entries of W are drawn independently from a Gaussian distribution with a negative mean of -0 . 1 and variance g 2 /N , where g is a gain parameter that controls the heterogeneity of weights. We use the transfer function φ ( x ) = max( x, 0) . x s is a sequence of inputs and V is the input weight. x s is a stream of i.i.d. Gaussian input x s ∼ N (0 , 1) and the input weights V are N (0 , 1) . Both W and V are trained during gradient flossing . We found that some ReLU network had initially unstable dynamics with positive Lyapunov exponents Fig 6C. However, during gradient flossing , these unstable networks were quickly stabilized. Fig 6D shows that the squared Lyapunov exponents of ReLU networks converge towards zero.

B.2 Additional Results for Different Numbers of Flossed Lyapunov Exponents

Additionally to the main Fig 5, we did the same analysis on networks with only preflossing ( gradient flossing before training) and found that more Lyapunov exponent k have to be flossed to achieve 100% accuracy as the tasks become more difficult (Fig 7D), however, in that case even if all N Lyapunov exponents were flossed, thus k = N , they were not able to solve the most difficult tasks. This seems to indicate that gradient flossing during training cannot be replaced by just gradient flossing more Lyapunov exponents before training.

C Additional Results on Gradient Flossing Throughout Training

We now discuss some additional results on gradient flossing throughout training. First, we analyze how gradient flossing affects the gradients and find that during gradient flossing , the norm of gradients that bridge many time steps are boosted. Moreover, subordinate singular values of the error norm of the recurrent weights are also boosted, indicating that gradient flossing can increase the effective rank of the parameter update. Additionally, we show that if gradient flossing is continued throughout training it can be detrimental to the accuracy. Finally, we show that Lyapunov exponents of successfully trained networks after training for the spatial delayed XOR task have a simple relationship to the delay d .

D Gradient Flossing boosts the Gradient Norm for Long Time Horizons

In this section, we investigate the impact of gradient flossing on the norm and structure of the gradient. It is important to note that the complete error gradient of backpropagation through time is composed of a summation of products of one-step Jacobians, reflecting the number of "loops" the error signal traverses through the recurrent

20

40

60

complexity (delay d)

60

80

100

test accuracy (%)

A

k

1

20

40

60

80

1 10 20 30 40 50 60 70 80

number of flossed

i

k

10

20

30

40

50

60

70

complexity (delay d)

C

50

100

test accuracy (%)

1 10 20 30 40 50 60 70 80

number of flossed

i

k

10

20

30

40

50

60

70

complexity (delay d)

D

50

100

test accuracy (%)

0

20

40

60

80

number of flossed

i

k

60

80

100

test accuracy (%)

B

delay

30

50

70

dynamics before reaching its target. Consequently, when the singular values of the long-term Jacobian are smaller than 1, the influence of the shorter loops typically dominates the long-term Jacobian.

In our tasks, we have full control over the correlation structure of the task and thus know exactly which loop length of backpropagation through time is necessary for finding the correct association. We were moreover careful in our task design not to have any additional signals in our task that might help to bridge the long time scale. In the case of vanishing gradients, the gradient norm is predominantly influenced by the shorter loops, even though the actual signal in the gradient originates solely from the loop of length d in our task. To mitigate the contamination of spurious signals from shorter loops and effectively extract the gradient that spans long time horizons, we focus on the gradient with respect to the initial conditions h 0 .

We note that the sum conveniently drops as only the longest 'loop', in other words, the only summand that contributes is the product of Jacobians going from 0 to t . By considering this gradient, we can therefore ensure that no undesired signals stemming from shorter loops interfere with the analysis. Moreover, we note that we use the binary cross entropy loss which makes the derivative ∂ L t ∂ h t trivial.

In Fig 8 we show that gradient flossing boosts the gradient with respect to the initial conditions. Specifically, we compare two identical networks trained on the binary delayed temporal XOR task with a loop length of d = 70 .

One network is trained with gradient flossing at epochs e ∈ { 0 , 100 , 200 , 300 , 400 } ) , while the other is trained without gradient flossing .

For the network without gradient flossing , the gradient norm of | d L dh 0 | diminishes to extremely small values ( < 10 -6 ) and remains small throughout training. In contrast, for the network trained with gradient flossing , each episode of gradient flossing causes the norm | d L dh 0 | to spike, surpassing values larger than 10 -2 . These findings are direct evidence that gradient flossing boosts the gradient norm, facilitating to bridge long time horizons in challenging temporal credit assignment tasks. We observe that after several episodes of gradient flossing , the gradient | d L dh 0 | of the networks stays around 10 -4 and eventually rise up to values around 10 -2 . Subsequent in training, the test accuracy surpasses chance level (Fig 8B). We observed this temporal relationship between gradient norm | d L dh 0 | and training success consistently across numerous network realizations (Fig 8C and D). These findings suggest that the gradient norm | d L dh 0 | can be a good predictor of learning success, sometimes hundreds of epochs before the accuracy exceeds the chance level of 50%. Indeed, when depicting the gradient norm aligned to the last epoch where accuracy was ≤ 50% , we see for many network realizations a gradual growth of gradient norm oven epochs before accuracy surpasses chance level (Fig 9A). Analogously, when plotting the accuracy as a function of epoch aligned with the last epoch with | d L dh 0 | < 0 . 001 , we observe for this task that the increase of gradient norm | d L dh 0 | reliably precedes the epoch at which the accuracy surpasses the chance level (See Fig 9B). We note that when measuring the overlap of the orientation of the gradient vector d L dh 0 with the first covariant Lyapunov vector of the forward dynamics, we found a significant increase in overlap around the training epoch where the accuracy surpasses the chance level both in networks with and without gradient flossing. This does not come as a surprise as the covariant Lyapunov vector measure the most unstable (or least stable) direction in the tangent space of a trajectory and perturbations of h 0 that have to travel over many epochs align

D.1 Gradient Flossing Boosts Effective Dimension of Error Gradient

To further investigate the effect of gradient flossing on training, we investigated the structure of the error gradient and how it is changed by gradient flossing . To this end, we decompose the recurrent weight gradient σ i ( d L dW ) into in weighted sum of outer products using singular value decomposition (Fig 10).

As the Lyapunov exponents are the time-averaged logarithms of the singular values of the asymptotic long-term Jacobian T t ( h τ ) , this allows us to directly link the effect of pushing Lyapunov exponents toward zero during gradient flossing to the structure of the error gradient of the recurrent weights, as they are intimately linked:

We again note that different 'loops' contribute to the total gradient expression and the Lyapunov exponents only characterize the longest loop. Further, we note that in our controlled tasks, depending on delay d , only few of the summands are relevant for solving the task. We thus expect the relevant gradient summand that carries important signals about the task to be contaminated by summands of both shorter and longer chains, which contribute irrelevant fluctuations.

The singular values of the recurrent weight gradient σ i ( d L dW ) as a function of training epoch reveal that the subordinate singular values subordinate singular value σ 20 and σ 40 exhibit peaks at the times of gradient flossing, while the first singular value σ 1 only shows a slight peak (Fig 10A). This indicates that gradient flossing increases the effective rank of the recurrent weight gradient d L dW . In other words, gradient flossing facilitates high-dimensional parameter updates. Our interpretation is, as gradient flossing pushes Lyapunov exponents to zero, the different summands in the total gradient contribute more equitable as long loops have neither a dominant contribution (which would happen for exploding gradients) nor a vanishing contribution (which would happen for vanishing gradients). This way, the sum of gradient terms has a higher effective rank.

In contrast, without gradient flossing , the subordinate singular values (in Fig 10A σ 20 and σ 40 ) rapidly diminish to extremely small values over training epochs and remain very small throughout training. Note however that the leading singular values σ 1 are of comparable size irrespective whether gradient flossing was performed or not.

We note that similar to the gradient norm of the loss with respect to the initial condition | d L dh 0 | , the subordinate singular values seem to predict when the test accuracy of networks with gradient flossing grows beyond chance level (Fig 10B). We confirmed this in multiple other network realizations and give here another example we the accuracy grows beyond chance only later during training (Fig 11).

D.2 Gradient Flossing Throughout Training Can Be Detrimental

We find that gradient flossing continued throughout all training epochs can be detrimental for performance (Fig 12). We demonstrate this again in the binary delayed temporal XOR task. We compare three different

0

500

1000

1500

2000

2500

3000

Epochs

0.00

0.01

0.02

0.03

0.04

0.05

|

d

dh

0

|

C

0

500

1000

1500

2000

2500

3000

Epochs

50

100

accuracy (%)

D

0

500

1000

1500

2000

2500

3000

Epochs

10

12

10

10

10

8

10

6

10

4

10

2

|

d

dh

0

|

A

with flossing during training

without flossing

0

500

1000

1500

2000

2500

3000

Epochs

50

100

accuracy (%)

B

with flossing during training

without flossing

conditions: Either, we floss throughout the training every 100 training epochs for 500 flossing epochs (red), or we floss only early during training at training epochs e ∈ { 0 , 100 , 200 , 300 , 400 } ) (green) or we do not floss at all (blue).

We observe that after every episode of gradient flossing , the accuracy drops down close to chance level of 50% (Fig 12A red line). Between flossing , the accuracy quickly recovers but never reaches 100%. Simultaneously, when the accuracy drops the test error jumps up (Fig 12B). We also observed that the gradient norm | d L dh 0 | is initially boosted by gradient flossing , but stays close to indistinguishable once the gradient norm | d L dh 0 | becomes macroscopically large (Fig 12C). This suggests that once gradient flossing facilitates signal propagation across long time horizons and the network picks up the relevant gradient signal, further gradient flossing can be harmful to the actual task execution. We hypothesize that there might be (at least for the Vanilla networks considered here), a trade-off between the ability to bridge long time scales which seems to require one or several Lyapunov exponents of the forward dynamics close to zero and nonlinear tasks requirements, which require at least a fraction of the units to be in the nonlinear regime of the nonlinearity φ , where φ ′ ( x ) < 1 . It would be an interesting future research avenue to further investigate this potential trade-off also in other network architectures.

D.3 Lyapunov Exponents after Training With and Without Gradient Flossing

In Fig 13, we show the first (Fig 13A, B) and the tenth (Fig 13C, D) Lyapunov exponent after training on the spatial delayed XOR task both with and without gradient flossing . We find for successful networks with

3000

2000

1000

0

1000

Epochs

10

5

10

4

10

3

10

2

|

d

dh

0

|

A

0

100

200

300

400

Epochs

45

50

55

60

65

70

75

80

accuracy (%)

B

gradient flossing a systematic relationship between the first Lyapunov exponent and the delay, that can be fitted by approximately by λ 1 ( d ) = -0 . 2 exp. ( -0 . 03 delay ) . Unsuccessful networks with accuracy at chance level have a much smaller largest Lyapunov exponent. The same seems to hold true for the tenth Lyapunov exponent. In a previous study [63], a similar trend was observed, albeit in the context of a task that did not possess an analytically tractable temporal correlation structure, which might partially explain the less conclusive results. It is important to note that the numerical evaluation of Lyapunov exponents in recurrent LSTM networks in [63] was based solely on the N × N Jacobian of the memory state. From a dynamical systems standpoint, a 2 N × 2 N Jacobian matrix encompassing interactions between both memory and cell states into account is required [9, 12, 64].

E Gradient Flossing for Linear Network

We provide code for gradient flossing in linear networks here. We find that gradient flossing also helps to train linear networks on tasks with many time steps that can be solved by linear networks, for example the copy task, but not for tasks the require a nonlinear input-output operation like the temporally delayed XOR task. Full analytical description of gradient flossing for linear networks would be a promising avenue for future research as networks with linear dynamics can still have nonlinear learning dynamics [23]. However this is beyond the cope of the presented work.

F Computational Complexity of Gradient Flossing

We present here a more in-depth scaling analysis of the computational cost of gradient flossing . There are three main contributors to the computational cost (table 1): First the RNN step, which has a computational complexity of O ( N 2 b ) per time step, where N is the dimension of the recurrent network state (which in case of Vanilla networks equals the number of units) and b is the batch-size both in the forward and backward pass. Second, the Jacobian step which scales with O ( N 2 k ) per time step, where k is the number of flossed Lyapunov exponents. Third, the QR decomposition, which scales with O ( N k 2 ) , where k is the number of Lyapunov exponents considered.

0

500

1000

1500

2000

2500

3000

Epochs

50

100

accuracy (%)

B

early training flossing

without flossing

0

500

1000

1500

2000

2500

3000

10

4

10

2

i

(

d

dW

)

A

Together, this results in a total amortized cost of O ( N 2 b T ) per training epoch, where T is the number of training time steps and a total amortized costs per flossing epoch of O ( N 2 T f (1 + k/t ONS + k ) ) where T f is the number of flossing time steps.

In case of preflossing , thus, the total computation cost scale with O ( N 2 [ EbT + E p T f (1 + k/t ONS + k )] ) , where E is the number of training epochs and E p is the number of preflossing epochs.

For gradient flossing during training (assuming that there is also preflossing done), the amortized cost scale with O ( N 2 [ EbT + E p T p + E f T f (1 + k/t ONS + k )] ) , where E f is the total number of flossing epochs during training.

Empirically, we find that both the number of preflossing epochs E p and flossing episodes E f necessary for training success is much smaller than the total number of training epochs E . For example, the preflossing for 500 epochs in the numerical experiment of Fig 3 took ∼ 37 seconds, while the overall training on 10000 training epochs with batch size b = 16 took ∼ 1680 seconds. Thus only approximately 2.2% of the total training time was spent on gradient flossing . Moreover, T p can be smaller than T , it just has to be long enough such that the temporal correlations in the task can be bridged. In case of the tasks discussed in the manuscript, this would be the delay d . It remains an important challenge to infer the suitable number of flossing time steps T f for tasks with unknown temporal correlation structure.

It would also be interesting to investigate how the CPU hours/wall-clock time/flops/Joule/CO2-emission spent on gradient flossing vs on training networks with larger N are trading off against each other. For this, we would suggest to first find the smallest network that on median successfully trains on a binary temporal XOR task for a fixed given delay d and measure the computational resources involved in training it, e.g. in terms of CPU hours. Then compare it to a network with gradient flossing . This would be a promising analysis but is beyond our current computational budget. We will start such experiments an might be able to provide results during the reviewer period.

0

500

1000

1500

2000

2500

3000

Epochs

50

100

accuracy (%)

B

early training flossing

without flossing

0

500

1000

1500

2000

2500

3000

10

4

10

2

i

(

d

dW

)

A

G Additional controls

We also investigate the effects of gradient flossing during the training with orthogonal weight initializations and confirm our finding that gradient flossing improves trainability on tasks that have long time horizons to bridge. Moreover, we find that gradient flossing during training can further improve trainability. We replicated the two more challenging tasks from the main paper (Fig 4) for orthogonal initialization with variable temporal complexity and performed gradient flossing either both during and before training, only before training, or not at all.

Fig 14A shows the test accuracy for Vanilla RNNs with orthogonal initialization trained on the delayed temporal XOR task y t = x t -d/ 2 ⊕ x t -d with random Bernoulli process x ∈ { 0 , 1 } . The accuracy of orthogonal Vanilla RNNs falls to chance level for d ≥ 40 (Fig 14C). With gradient flossing before training, the trainability can be improved, but still falls close to chance level for d = 70 . In contrast, for initially orthogonal networks with gradient flossing during training, the accuracy is improved to > 80% at d = 70 . In this case, we preflossed for 500 epochs before task training and again after 500 epochs of training on the task. In Fig 14B, D the networks have to perform the nonlinear XOR operation y t = x 1 t -d ⊕ x 2 t -d ⊕ x 3 t -d on a three-dimensional binary input signal x 1 , x 2 , and x 3 and generate the correct output with a delay of d steps identical to Fig 4 in the main text. Again, we observe similar to networks with Gaussian initialization that flossing before training improves the performance compared to baseline, but starts failing for long delays d > 60 . In contrast, orthogonal networks that are also flossed during training can solve even more difficult tasks (Fig 14D). We note that for Fig 14B and D, we trained the network only on 5000 epochs, compared to 10000 epochs in networks with random Gaussian initialization because for 10000 epochs, both networks with gradient flossing only before training and with gradient flossing before and during training were able to bridge d = 70 . These results suggest that orthogonal initialization does seem to slightly improve performance for tasks with long time horizons to bridge and gradient flossing and additionally boost the performance. Thus orthogonal initialization and gradient flossing seems to go well together. It would be interesting to study if orthogonal initialization also reduces the number of gradient flossing steps necessary to improve performance.

H Additional Details on Training Tasks

In this section, we provide a more rigorous definition of the tasks used for training, as discussed in Section 3:

A

B

accuracy (%)

test error

100

50

10 0

10

1

0

0

500

500

1000

1000

1500

1500

early training flossing with flossing throughout training

without flossing

2000

2500

2000

2500

3000

3000

0

500

1000

1500

2000

2500

3000

Epochs

10

7

10

3

|

d

dh

0

|

C

with flossing throughout training

early training flossing

without flossing

H.1 Copy task

For the copy task, the target network readout at time t is y t = x t -d , where d denotes the delay. We chose the input to be sampled i.i.d. from a uniform distribution between 0 and 1.

H.2 Temporal XOR task

The temporal XOR task requires the target network readout y t at time t to be computed as follows:

where again d denotes a time delay of d time steps. In the case of x ∈ { 0 , 1 } 2 and y ∈ { 0 , 1 } , the output y t follows the truth table of the XOR digital logic gate (Table 2). Thus, the function f ( x a , x b ) = | x a -x b | can be seen as an analytical representation of the XOR gate. It is important to note that f ( x, 0) = x only for x ≥ 0 , and that this task requires a nonlinearity. The implementation can easily be constructed analytically, for example, using two rectified linear units φ ( x ) = max( x, 0) the outbut can be constructed by

Together with a delay line to transmit the signal x t -d over time, this can solve the task.

20

40

60

complexity (delay d)

0.2

0.1

1

A

no gradient flossing

20

40

60

complexity (delay d)

0.2

0.1

1

B

gradient flossing during training

20

40

60

complexity (delay d)

0.3

0.2

0.1

10

C

20

40

60

complexity (delay d)

0.2

0.1

10

D

I Additional Background on Lyapunov Exponents of RNNs

An autonomous dynamical system is usually defined by a set of ordinary differential equations d h / d t = F ( h ) , h ∈ R N in the case of continuous-time dynamics, or as a map h s +1 = f ( h s ) in the case of discrete-time dynamics. In the following, the theory is presented for discrete-time dynamical systems for ease of notation, but everything directly extends to continuous-time systems [49]. Together with an initial condition h 0 , the map forms a trajectory. As a natural extension of linear stability analysis, one can ask how an infinitesimal perturbation h ′ 0 = h 0 + glyph[epsilon1] u 0 evolves in time. Chaotic systems are sensitive to initial conditions; almost all infinitesimal perturbations glyph[epsilon1] u 0 of the initial condition grow exponentially with time | glyph[epsilon1] u t | ≈ exp( λ 1 t ) | glyph[epsilon1] u 0 | . Finite-size perturbations, therefore, may lead to a drastically different subsequent behavior. The largest Lyapunov exponent λ 1 measures the average rate of exponential divergence or convergence of nearby initial conditions:

In dynamical systems that are ergodic on the attractor, the Lyapunov exponents do not depend on the initial conditions as long as the initial conditions are in the basins of attraction of the attractor. Note that it is crucial to first take the limit glyph[epsilon1] → 0 and then t →∞ , as λ 1 ( h 0 ) would be trivially zero for a bounded attractor if the limits are exchanged, as lim t →∞ log || glyph[epsilon1] u t || || glyph[epsilon1] u 0 || is bounded for finite perturbations even if the system is chaotic. To measure k Lyapunov exponents, one has to study the evolution of k independent infinitesimal perturbations u s spanning the tangent space:

Table 1: Computational cost for gradient flossing and training of RNNs

N denotes number of neurons, b is the batch size, T is the number of time steps in forward pass of training, T f is the number of time steps in forward pass of flossing , t ONS is the reorthonormalization interval, k is the number of flossed Lyapunov exponents, E is the number of training epochs, E p is the number of preflossing epochs, E f is the number of flossing epochs during training. Empirically, we find that the necessary number of preflossing epochs E p and flossing episodes E f is much smaller than both the total number of training epochs E . Moreover, T p can be smaller than T .

where the N × N Jacobian D s ( h s ) = d f ( h s ) / d h characterizes the evolution of generic infinitesimal perturbations during one step. Note that this Jacobian along the trajectory is equivalent to a stability matrix only at a fixed point, i.e., when h s +1 = f ( h s ) = h s .

We are interested in the asymptotic behavior, and therefore we study the long-term Jacobian

Note that T t ( h 0 ) is a product of generally noncommuting matrices. The Lyapunov exponents λ 1 ≥ λ 2 · · · ≥ λ N are defined as the logarithms of the eigenvalues of the Oseledets matrix

where glyph[latticetop] denotes the transpose operation. The expression inside the brackets is the Gram matrix of the long-term Jacobian T t ( h 0 ) . Geometrically, the determinant of the Gram matrix is the squared volume of the parallelotope spanned by the columns of T t ( h 0 ) . Thus, the exponential volume growth rate is given by the sum of the logarithms of its first k (sorted) eigenvalues. Oseledets' multiplicative ergodic theorem guarantees the existence of the Oseledets matrix Λ ( h 0 ) for almost all initial conditions h 0 [48]. In ergodic systems, the Lyapunov exponents λ i do not depend on the initial condition h 0 . However, for a numerical calculation of the Lyapunov spectrum, Eq 25 cannot be used directly because the long-term Jacobian T t ( h 0 ) quickly becomes ill-conditioned, i.e., the ratio between its largest and smallest singular value diverges exponentially with time.

J Algorithm for Calculating Lyapunov Spectrum of Rate Networks

For calculating the first k Lyapunov exponents, we exploit the fact that the growth rate of a k -dimensional infinitesimal volume element is given by λ ( m ) = ∑ m i =1 λ i . Therefore, λ 1 = λ (1) , λ 2 = λ (2) -λ 1 , λ 3 = λ (3) -λ 1 -λ 2 , . . . [50]. The volume growth rates can be obtained via QR-decomposition.

orthogonal weight initialization

0

2000

4000

6000

8000

10000

Epochs

60

80

100

test accuracy (%)

A

delayed temporal XOR

flossing during training

preflossing

no flossing

10

20

30

40

50

60

70

complexity (delay d)

60

80

100

test accuracy (%)

C

0

1000

2000

3000

4000

5000

Epochs

60

80

100

test accuracy (%)

B

delayed spatial XOR

10

20

30

40

50

60

70

complexity (delay d)

60

80

100

test accuracy (%)

D

flossing during training

preflossing

no flossing

A) Test accuracy for orthogonally initialized vanilla RNNs trained on delayed temporal binary XOR task y t = x t -d/ 2 ⊕ x t -d with gradient flossing during training (green), preflossing (orange), and with no gradient flossing (blue) for d = 70 . Solid lines are mean, transparent thin lines are individual network realizations B) Same as A for delayed spatial XOR task with y t = x 1 t -d ⊕ x 2 t -d ⊕ x 3 t -d . C) Test accuracy as a function of task difficulty (delay d ) for delayed temporal XOR task. D) Test accuracy as a function of task difficulty (delay d ) for delayed spatial XOR task. Parameters: g = 1 , batch size b = 16 , N = 80 , epochs = 10 4 for delayed temporal XOR, epochs = 5000 for delayed spatial XOR, T = 300 , flossing for E f = 500 epochs on k = 75 Lyapunov exponents before training and during training for green lines, and only before training for orange lines. Shaded areas are 25% and 75% percentiles, solid lines are means, transparent dots are individual simulations, task loss is cross-entropy between y, ˆ y .

First, we evolve an initially orthonormal system Q s = [ q 1 s , q 2 s , . . . q s m ] in the tangent space along the trajectory using the Jacobian D s :

A continuous system can be transformed into a discrete system by considering a stroboscopic representation, where the trajectory is only considered at certain discrete time points. We use here the notation of discrete dynamical systems where this corresponds to performing the product of Jacobians along the trajectory ˜ Q s +1 = D s Q s . We study the discrete network dynamics in the limit of small time step ∆ t → 0 and for discrete time ∆ t = 1 . The notation can be readily extended to continuous systems [49].

Second, we extract the exponential growth rates using the QR-decomposition,

which uniquely decomposes ˜ Q s +1 into an orthonormal matrix Q s +1 of size N × k so Q glyph[latticetop] s +1 Q s +1 = 1 m × m and to an upper triangular matrix R s +1 of size k × k with positive diagonal elements. Geometrically, Q s +1 describes the rotation of Q s caused by D s and the diagonal entries of R s +1 describe the stretching and shrinking of the columns of Q s , while the off-diagonal elements represent the shearing. Fig 15 visualizes D s and the QR-decomposition for k = 2 .

The Lyapunov exponents are given by time-averaged logarithms of the diagonal elements of R s :

q

s

1

q

s

2

q

s+1

~1

q

s+1

~2

D s

QR

R

11

s+1

R

22

s+1

R

11

s+1

R

22

s+1

q

s+1

1

q

s+1

2

Note that the QR-decomposition does not need to be performed at every simulation step, just sufficiently often, i.e., once every s ONS steps such that ˜ Q s + s ONS = D s + s ONS -1 · D s + s ONS -2 . . . D s · Q s remains well-conditioned [50]. An appropriate reorthonormalization interval s ONS = t ONS / ∆ t thus depends on the condition number, the ratio of the smallest and largest singular value:

An initial transient should be disregarded in the calculation of the Lyapunov spectrum because h first has to converge towards the attractor and Q has to converge to the unique eigenvectors of the Oseledets matrix (Eq 25) [60]. A simple example of this algorithm in pseudocode is:

Algorithm 2 Jacobian-based algorithm for Lyapunov spectrum

initialize h , Q

It is guaranteed that under general conditions initially random orthonormal systems will exponentially converge towards a unique basis that is given by the eigenvectors of the Oseledets matrix Eq 25 [60]. A minimal example of this algorithm in pseudocode is shown in appendix 3. A feasible strategy to determine the reorthonormalization time interval t ONS is to get first a rough estimate of the Lyapunov spectrum using a short simulation time t sim and a small t ONS and repeat with a longer simulation time and a t ONS based on the Lyapunov spectrum of the rough estimate of the Lyapunov spectrum. Another strategy is, to first iteratively adapt t ONS on a short simulation run to get an acceptable condition number. It should be noted that there exists a diversity of other methods to estimate the Lyapunov spectrum [14, 49, 79, 80].

K Convergence of Lyapunov Exponents of RNNs

In Fig. 16, we demonstrate the convergence of the Lyapunov exponents. We show the estimate of the Lyapunov exponents λ i for i = 1 , 20 , 60 , 80 for different initial conditions but identical network realization.

10 0

10 1

10 2

steps

40

30

20

10

0

i

(1/

steps

)

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose a novel approach called \"gradient flossing\" to tackle the instability of gradients in the training of recurrent neural networks (RNNs). The process of pushing the Lyapunov exponents toward zero is referred to as \"flossing\" the gradients. This stabilizes the gradients and improves network training.\n\nSTRENGTHS:\n+ The authors show improvements in several toy sequence learning tasks: Delayed Copy Task, Temporal XOR Task, etc.\n\n+ Idea of using Lyapunov exponents towards zero to mitigate vanishing and exploding gradients is well justified and works as expected on toy tasks.\n\nWEAKNESSES:\nThe lack of any experiments on non-simulated data is problematic, and the explanation given by the authors is not convincing. The paper's authors explicitly state that they deliberately did not use sequential MNIST or similar non-toy tasks commonly used to probe exploding/vanishing gradients. \n\n> We deliberately do not use sequential MNIST or similar toy tasks commonly used to probe exploding/vanishing gradients, because we want a task where the structure of long-range dependencies in the data is transparent and can be varied ad libitum\n\nWhile toy tasks are more transparent, real-world data is important to understand the proposed methods' potential failure modes and expected characteristics in *uncontrolled* settings. Avoiding running experiments on the basis that the results are more difficult to interpret than simulated data is not an acceptable excuse. Transparency in the work itself is hindered when experiments are consciously avoided.\n\nQUESTIONS:\nCan the authors comment that such a technique of pushing the Lyapunov exponents to zero could limit the capacity to learn certain tasks? Maybe enabling the network to amplify or forget things is important for some task contexts, which has been a motivation for specialized architectures like the LSTM.\n\nLIMITATIONS:\nExperimental limitation mentioned above.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose a new technique for handling gradient instability (vanishing/exploding) in neural network models, especially sequence models such as RNNs. Specifically, by leveraging results from recent works that establish a connection between Lyapunov growth exponents (from dynamical systems) and the singular values of the model-Jacobian itself, the authors propose a (regularization) approach that they coin ‘gradient flossing’ to control the evolution of these singular values desirably. They demonstrate that, as a result, the gradient signal remains stable, thereby allowing for successful training of the model itself.\n\nAdditionally, the authors provide insights on the steps (schedule, pre/during training) to perform such a regularization leading to optimal gradient stabilization & computational benefits.\n\nSTRENGTHS:\nThe authors visit a longstanding problem in training NN models, esp. sequence models with long temporal dependencies. By leveraging recently established connections between Lyapunov exponents and the singular values of the (long-term) Jacobian, the authors propose a new technique in the form of:\n- Defining a loss function out of these Lyapunov exponents\n- Determine a clever strategy to evaluate this loss function through a QR decomposition\nto constrain these exponents directly to stable target values. This, in turn, allows control over the evolution of the Jacobian spectra itself.\n\nThe authors provide experimental evidence supporting this technique on two specific tasks that, by design, require long-time gradient signals to propagate efficiently.\n\nFurthermore, the authors provide practical insights into the (expensive) computational needs of the exact implementation of such a scheme and the place (pre/during training) of the same. They then propose a modified, intermittent regularization routine with preserved benefits yet having lower compute necessities.\n\nThe authors do an overall good job of satisfactorily articulating the problem statement, the proposed solution, observed benefits, and implementation limitations.\n\nWEAKNESSES:\nWhile I believe this is a good contribution, I feel the work can undergo a fair amount of extra polishing in terms of clarity. Please find below (Questions) my comments/questions on areas where this work can be further improved.\n\nQUESTIONS:\n- Eq. (2) does not necessarily need to refer to the ‘long-term Jacobian’ unless $\\Delta{t} \\equiv t - \\tau \\to \\infty$. That should be clarified.\n- Is the loss function from Eq. (4) used as an additional or the only loss to train?\n- Eq. (4) does not seem to contain a reference to the target $\\lambda$ values that the authors eventually aim to constrain them to\n- While the analytical $t\\to\\infty$ limit is understandable, can you provide an estimate of a reasonable sequence length for *practical data* that can serve as a finite size equivalent of this limit?\n- Is Eq. (5) subject to the assumption that parameter updates are gradient-based?\n- I believe the ’s’-index in Eq. (6) is supposed to run from $\\tau$ to be exact (?)\n- Clearly, in arriving at Eq. (6), the $Q$’s from *subsequent* time-steps $s$ and $s+1$ act on each other to reduce to unity ($Q_{s+1}^T Q_{s}$), leaving with the $R$’s only. Can the authors explain why that might be the case in general? Is there an assumption about a small learning rate here?\n- Since the authors claim the QR decomposition need not be applied at every time step, will the equivalent Eq. (6) thus arrived at still remain to be factually correct? The expression of $T_t$ contains an iterative product of $ D_s$, of which if one chooses to QR decompose a handful only, then that does not seem to yield $\\prod_s R^s$ for some period $t_{ons}$.\n- In Lines 140-141, you refer to using vanilla gradient descent for the experiments of Fig. 1 when you specify in Lines 132-133 that experiments were performed using ADAM.\n- Isn’t a target $\\lambda$ value of -1 indicative of vanishing gradients as per Eq. (3)? How is the model still reliably attaining that value regardless, then?\n- In Fig. 2 (A), the ‘theory’ and ‘simulation’ legends seem to be interchanged.\n- In Fig. 2 (A) & (C), the numerical and theoretical results should overlap in the long time-step regime, yet as time increases, the two appear to diverge in the former and maintain a uniform scaling (~10^2) in the latter. Can you explain why?\n- Can you explain what you mean by ‘flossing before training’? That does not seem to be well clarified in the text.\n\nLIMITATIONS:\nTo the best of my knowledge, the authors act commendably in attributing to the various limitations of their approach.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nAuthors propose gradient flossing for RNNs, which adds an additional regularization term that keeps the sum of Lyapunov exponents close to zero. This essentially encourages the singular values of the long-term Jacobian to be close to 1, hence addressing the vanishing/exploding gradient problem in RNNs. The authors point out that minimizing this regularization term again has the same gradient instability problem, and they purpose to apply QR decomposition once in a while to avoid the ill-conditioned long-term Jacobian pass. Finally, the authors show gradient flossing achieves faster convergence and a higher success rate on simple long-term tasks up to 500 steps.\n\nSTRENGTHS:\n* The idea of regularizing the Lyapunov exponent and using QR decomposition is novel and interesting.\n* The problem is well-motivated.\n\nWEAKNESSES:\n* Empirical evaluations are toy. Did not include results on standard long-time horizon benchmarks like Longe Range Arena [1].\n\n* The paper could have studied gradient flossing with architectures that are more suited for long horizon tasks such as state space models [2] and LRUs [3], instead of Vanilla RNNs and LSTMs.\n\n* Authors should compare to more recent methods as baselines such as HiPPO Matrices[2] and initialization methods used in LRU [3]. Those methods tend to be more competitive than a simple orthogonal initialization.\n\n* Clarity of the paper can be improved. For example, $\\tau$ in equation 1 is never defined, and equation 4 could be better explained. In addition, I did not quite understand the technical parts of the paper explaining how equation 4 is calculated in practice (see questions below). \n\n[1] Tay, Yi, et al. \"Long range arena: A benchmark for efficient transformers.\" arXiv preprint arXiv:2011.04006 (2020).\n\n[2] Gu, Albert, Karan Goel, and Christopher Ré. \"Efficiently modelling long sequences with structured state spaces.\" arXiv preprint arXiv:2111.00396 (2021).\n\n[3] Orvieto, Antonio, et al. \"Resurrecting recurrent neural networks for long sequences.\" arXiv preprint arXiv:2303.06349 (2023).\n\nQUESTIONS:\n* I’m still unclear how QR decomposition is able to avoid the long-term Jacobian being ill-conditioned. \n\n* Why can we not use QR decomposition directly on the original RNN objective in equation 1 to avoid the long-term gradient instability problem?\n\nLIMITATIONS:\nSee above.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe presented paper proposes a new method for tackling numerical instabilities during training for recurrent neural networks.\nThe proposed method exploits a theoretical link between Lyapunov exponents and the singular values of the long-term Jacobian.\nThe set of experiments is well-chosen to showcase the extent to which gradient flossing helps stabilize training for tasks involving long-range dependencies.\nThere is an in-depth coverage of the practical aspects of the methods in the appendix, which is very useful for both machine learning practitioners and researchers.\n\nBy steering the Lyapunov exponents to 0, either before and/or during training, the authors are able to show consistent improvements on a variety of synthetic learning tasks involving long-range dependencies in temporal sequences.\n\nSTRENGTHS:\n- There are experiments validating basic hypotheses, such as the fact that the proposed algorithm is able to control the Lyapunov exponents, and that estimating the Lyapunov exponents via the QR decomposition matches the theory.\n- Practical aspects such as when to perform gradient flossing, how often should it be used during training, and also how often should the QR decomposition should be computed when unrolling the neural network dynamics.\n\nWEAKNESSES:\n- There are no evaluation on real world complex benchmark to estimate the impact of the proposed contribution. However, the paper is already quite dense.\n\n- Minor remark: I don't understand the \"the finite network size fluctuations of the Lyapunov exponents\" in figure 1.\n\nQUESTIONS:\n- (line 197 - 201) How does fig 2.C shows that the estimated condition number based on Lyapunov exponents can predict differences in condition number originating from finite network size $N$ ?\nThe authors refer to appendix F for this matter, but it doesn't seem related to this claim.\n- Can the author explain in more details what are the \"the finite network size fluctuations of the Lyapunov exponents\" in figure 1 ?\n- Figure 11 in the appendix shows that too much gradient flossing during training can be detrimental, which means that the local minima of the gradient flossing procedure hardly coincides with the one corresponding to the actual learning task. It is also said in the introduction that appropriate initialization schemes can ensure well-behaved gradient at initialization, a property that can be lost during training. I suspect that the gradient flossing procedure prevents the weights from falling into an \"ill-conditioned\" region. Thus, how often should gradient flossing be used should depend on, given a fixed learning rate, how many iterations are needed to go from a \"well-behave\" region to an ill-condition one. Can the author elaborate on this specific intuition ?\nI am asking this question because I suspect that you already found regimes where the regularization objective conflicts with the actual learning task.\n\nLIMITATIONS:\nAs emphasized by the author, the Lyapunov exponent characterizes the singularity of the long-term Jacobian, and thus cannot characterize local-in-time behavior.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-07"}