Gradient Flossing: Improving Gradient Descent through Dynamic Control of Jacobians

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 werbos_beyond_1974 ; parker_learning-logic_1985 ; lecun_procedure_1985 ; rumelhart_learning_1986 . 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 hochreiter_untersuchungen_1991 ; hochreiter_long_1997 ; bengio_learning_1994 ; pascanu_difficulty_2012 . 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 regulates 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 engelken_lyapunov_2020 ; mikhaeil_difficulty_2022 ; park_persistent_2023 ; engelken_lyapunov_2023 . Lyapunov exponents characterize the asymptotic average rate of exponential divergence or convergence of nearby initial conditions and are a cornerstone of dynamical systems theory eckmann_ergodic_1985 ; pikovsky_lyapunov_2016 . 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 hochreiter_untersuchungen_1991 ; hochreiter_long_1997 ; cho_properties_2014 . Second, exploding gradients can be avoided by gradient clipping, which re-scales the gradient norm pascanu_difficulty_2013 or their individual elements mikolov_statistical_2012 if they become too large goodfellow_deep_2016 . Third, normalization schemes like batch normalization prevent saturated nonlinearities that contribute to vanishing gradients ioffe_batch_2015 . Third, it was suggested that the problem of exploding/vanishing gradients can be ameliorated by specialized network architectures, for example, antisymmetric networks chang_antisymmetricrnn_2018 , orthogonal/unitary initializations saxe_exact_2013 ; arjovsky_unitary_2016 ; helfrich_orthogonal_2018 , coupled oscillatory RNNs konstantin_rusch_coupled_2020 , Lipschitz RNNs erichson_lipschitz_2021 , linear recurrent units orvieto_resurrecting_2023 , echo state networks jaeger_echo_2001 ; ozturk_analysis_2007 , (recurrent) highway networks srivastava_training_2015 ; zilly_recurrent_2017 , and stable limit cycle neural networks sokol_adjoint_2019-1 ; sokol_geometry_2023 ; park_persistent_2023 . Fourth, for large networks, a suitable choice of weights can guarantee a well-conditioned Jacobian at initialization glorot_deep_2011 ; saxe_exact_2013 ; poole_exponential_2016 ; chen_dynamical_2018 ; hanin_products_2018 ; schoenholz_deep_2016 ; hanin_how_2018 ; sokol_information_2018 ; gilboa_dynamical_2019 ; can_gating_2020 . 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 gilboa_dynamical_2019 . 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 engelken_lyapunov_2020 ; sokol_geometry_2023 ; park_persistent_2023 ; engelken_lyapunov_2023 . 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.

• •

  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.

We begin by revisiting the established mathematical relationship between the gradients of the loss function, computed via backpropagation through time, and Lyapunov exponents engelken_lyapunov_2020 ; engelken_lyapunov_2023 , and how it relates to the problem of vanishing and exploding gradients. In backpropagation through time, network parameters $\theta$ are iteratively updated by stochastic gradient descent such that a loss $L_{t}$ is locally reduced werbos_beyond_1974 ; parker_learning-logic_1985 ; lecun_procedure_1985 ; rumelhart_learning_1986 . For RNN dynamics $\mathbf{h}_{s+1}=\mathbf{f}_{\theta}(\mathbf{h}_{s},\mathbf{x}_{s+1})$, with recurrent network state $\mathbf{h}$, external input $\mathbf{x}$, and parameters $\theta$, the gradient of the loss $\mathcal{L}_{t}$ with respect to $\theta$ is evaluated by unrolling the network dynamics in time. The resulting expression for the gradient is given by:

where $\mathbf{T}_{t}(\mathbf{h}_{\tau})$ is composed of a product of one-step Jacobians $\mathbf{D}_{s}=\frac{\partial\mathbf{h}_{s+1}}{\partial\mathbf{h}_{s}}$:

Due to the chain of matrix multiplications in $\mathbf{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 $\mathbf{x}$ or states $\mathbf{h}$ from many time steps prior which creates long temporal dependencies hochreiter_untersuchungen_1991 ; hochreiter_long_1997 ; bengio_learning_1994 ; pascanu_difficulty_2012 . How well error signals can propagate back in time is constrained by the tangent space dynamics along trajectory $\mathbf{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 $\mathbf{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 engelken_lyapunov_2020 ; engelken_lyapunov_2023 : Lyapunov exponents $\lambda_{1}\geq\lambda_{2}\dots\geq\lambda_{N}$ are defined as the asymptotic time-averaged logarithms of the singular values of the long-term Jacobian oseledets_multiplicative_1968 ; eckmann_ergodic_1985 ; geist_comparison_1990

where $\sigma_{i,t}$ denotes the $i$th singular value of $\mathbf{T}_{t}(\mathbf{h}_{\tau})$ with $\sigma_{1,t}\geq\sigma_{2,t}\dots\sigma_{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.

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 $\mathbf{T}_{t}(\mathbf{h}_{\tau})$. We achieve this by using the sum of the squares of the first $k$ largest Lyapunov exponent $\lambda_{1},\lambda_{2}\dots\lambda_{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 $\mathbf{T}_{t}$ (Eq 2), we use an established iterative reorthonormalization method involving QR decomposition that avoids directly evaluating the ill-conditioned long-term Jacobian engelken_lyapunov_2023 ; benettin_lyapunov_1980 .

First, we evolve an initially orthonormal system $\mathbf{Q}_{s}=[\mathbf{q}_{s}^{1},\,\mathbf{q}_{s}^{2},\dots\mathbf{q}_{s}^{k}]$ in the tangent space along the trajectory using the Jacobian $\mathbf{D}_{s}=\frac{\partial\mathbf{h}_{s+1}}{\partial\mathbf{h}_{s}}$. This means to calculate

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

which decomposes $\widetilde{\mathbf{Q}}_{s+1}$ uniquely into the product of an orthonormal matrix $\mathbf{Q}_{s+1}$ of size $N\times k$ so $\mathbf{Q}^{\top}_{s+1}\mathbf{Q}_{s+1}=\mathds{1}_{k\times k}$ and an upper triangular matrix $\mathbf{R}^{s+1}$ of size $k\times 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_{\textnormal{ONS}}$ such that $\widetilde{\mathbf{Q}}$ does not become ill-conditioned.

The Lyapunov exponents are given by time-averaged logarithms of the diagonal entries of $\mathbf{R}^{s}$ benettin_lyapunov_1980 ; geist_comparison_1990 :

This way, the Lyapunov exponent can be expressed in terms of a temporal average over the diagonal elements of the $\mathbf{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 liao_differentiable_2019 : The backward pass of the QR decomposition is given by walter_algorithmic_2010 ; liao_differentiable_2019 ; schreiber_storage-efficient_1989 ; seeger_auto-differentiating_2017

where $\mathbf{M}=\mathbf{R}\overline{\mathbf{R}}^{T}-\overline{\mathbf{Q}}^{T}\mathbf{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 $[\texttt{copyltu}(M)]_{ij}=M_{\max(i,j),\min(i,j)}$ walter_algorithmic_2010 ; liao_differentiable_2019 ; schreiber_storage-efficient_1989 ; seeger_auto-differentiating_2017 . We denote here adjoint variable as $\overline{T}={\partial\mathcal{L}}/{\partial T}$. A simple implementation of this algorithm in pseudocode is:

For clarity, we described gradient flossing in terms of stochastic gradient descent, but we actually implemented it with the ADAM optimizer using standard hyperparameters $\eta$, $\beta_{1}$ and $\beta_{2}$. An example implementation is available here. Note that this algorithm also works for different recurrent network architectures. In this case, the Jacobians $\mathbf{D}$ has size $n\times 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 $2N\times 2N$ can_gating_2020 ; engelken_lyapunov_2020 ; engelken_lyapunov_2023 . The Jacobian matrix $\mathbf{D}$ can either be calculated analytically or it can be obtained via automatic differentiation.

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

The initial entries of $\mathbf{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 $\phi(x)=\tanh(x)$. (See appendix B for gradient flossing with ReLU and LSTM units). $\mathbf{x}_{s}$ is a sequence of inputs and $\mathbf{V}$ is the input weight. $\mathbf{x}_{s}$ is a stream of i.i.d. Gaussian input $x_{s}\sim\mathcal{N}(0,\,1)$ and the input weights $\mathbf{V}$ are $\mathcal{N}(0,\,1)$. Both $\mathbf{W}$ and $\mathbf{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 $\lambda^{\textnormal{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 $\lambda_{1}^{\textnormal{target}}$ seems not to arrive at a positive Lyapunov exponent $\lambda_{1}$.

Fig 1C shows gradient flossing for different numbers of Lyapunov exponents $k$. Here, during gradient-descent, 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 $\mathcal{O}\left(N\,k^{2}\right)$, 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.

A well-conditioned Jacobian is essential for efficient and fast learning saxe_exact_2013 ; pennington_resurrecting_2017 ; pennington_emergence_2018 . 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 $\kappa_{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 $\sigma_{\textnormal{max}}$ and the smallest singular values $\sigma_{\textnormal{min}}$, so $\kappa_{2}(A)=\frac{\sigma_{\textnormal{max}}(A)}{\sigma_{\textnormal{min}}(A)}$. According to the rule of thumb given in cheney_numerical_2007 , if $\kappa_{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 $\mathbf{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 $\mathbf{T}_{t}$ exponentially diverges with the number of time steps (Fig 2A). We compare the numerically measured condition number $\kappa_{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 $\kappa_{2}$ of an orthonormal system $\mathbf{Q}$ of size $N\times m$ that is temporally evolved by the long-term Jacobian $\mathbf{T}_{t}$ is:

where $\sigma_{1}(\mathbf{T}_{t}(\mathbf{h}_{\tau}))$ and $\sigma_{m}(\mathbf{T}_{t}(\mathbf{h}_{\tau}))$ 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 large-network limit $N\rightarrow\infty$ and usually do not capture finite-size effects bahri_statistical_2020 (see appendix G).

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 $\mathcal{L}_{\textnormal{flossing}}=\sum_{i=1}^{k}\lambda_{i}^{2}$ 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 input-output computation on a sequential stream of scalar inputs, i.e., $y_{t}=|x_{t-\nicefrac{{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).

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}\oplus x_{t-d}$ with random Bernoulli process $x\in\{0,1\}$. The accuracy of Vanilla RNNs falls to chance level for $d\geq 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}\oplus x^{2}_{t-d}\oplus 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 $\mathbf{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)
We note that merely regularizing the spectral radius of the recurrent weight matrix $\mathbf{W}$ or the individual one-step Jacobians $\mathbf{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 $\mathbf{D}_{s}$ into account is important for improving trainability.

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-\tau$, 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 ginelli_characterizing_2007 , 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 ershov_concept_1998 . 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.

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 chen_neural_2018 . 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 can_gating_2020 and mean-field heuristics gilboa_dynamical_2019 , 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 bahri_statistical_2020 , 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 engelken_lyapunov_2020 ; engelken_lyapunov_2023 and has been spelled out analytically and studied numerically lindner_investigating_2021 ; vogt_lyapunov_2022 ; mikhaeil_difficulty_2022 ; park_persistent_2023 ; krishnamurthy_theory_2022 . 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 pontryagin_mathematical_1987 ; liberzon_calculus_2011 ; sokol_geometry_2023 ). 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 marschall_unified_2020 . Extending gradient flossing to feedforward networks, state-space models and transformers is a promising avenue for future research (see also hoffman_robust_2019 ; zhang_fixup_2019 ). 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 molgedey_suppressing_1992 ; rajan_stimulus-dependent_2010 ; schuecker_functional_2016 ; engelken_lyapunov_2020 ; engelken_input_2022 ; engelken_lyapunov_2023 , 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.

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