SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Rainer Engelken
Zuckerman Mind Brain Behavior Institute
re2365@columbia.edu
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.
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.
- 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  $\theta$  are iteratively updated by stochastic gradient descent such that a loss  $L_{t}$  is locally reduced [1, 2, 3, 4]. 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
$^{1}$ Code in Julia using Flux [77, 78] is available at https://github.com/RainerEngelken/GradientFlossing
unrolling the network dynamics in time. The resulting expression for the gradient is given by:
$$
\frac {\partial \mathcal {L} _ {t}}{\partial \theta} = \frac {\partial \mathcal {L} _ {t}}{\partial \mathbf {h} _ {t}} \sum_ {\tau = t - l} ^ {\tau = t - 1} \left(\prod_ {\tau^ {\prime} = \tau} ^ {t - 1} \frac {\partial \mathbf {h} _ {\tau^ {\prime} + 1}}{\partial \mathbf {h} _ {\tau^ {\prime}}}\right) \frac {\partial \mathbf {h} _ {\tau}}{\partial \theta} = \frac {\partial \mathcal {L} _ {t}}{\partial \mathbf {h} _ {t}} \sum_ {\tau} \mathbf {T} _ {t} (\mathbf {h} _ {\tau}) \frac {\partial \mathbf {h} _ {\tau}}{\partial \theta} \tag {1}
$$
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}$ :
$$
\mathbf {T} _ {t} \left(\mathbf {h} _ {\tau}\right) = \prod_ {\tau^ {\prime} = \tau} ^ {t - 1} \frac {\partial \mathbf {h} _ {\tau^ {\prime} + 1}}{\partial \mathbf {h} _ {\tau^ {\prime}}} = \prod_ {\tau^ {\prime} = \tau} ^ {t - 1} \mathbf {D} _ {\tau^ {\prime}} \tag {2}
$$
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$  depends on inputs  $\mathbf{x}$  or states  $\mathbf{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  $\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 [9, 12]: 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 [13, 48, 49]
$$
\lambda_ {i} = \lim  _ {t \rightarrow \infty} \frac {1}{t - \tau} \log \left(\sigma_ {i, t}\right) \tag {3}
$$
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.
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  $\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 \ldots, \lambda_k$  as a loss function:
$$
\mathcal {L} _ {\text {f l o s s i n g}} = \sum_ {i = 1} ^ {k} \lambda_ {i} ^ {2} \tag {4}
$$
and evaluate the gradient obtained from backpropagation through time:
$$
\frac {\partial \mathcal {L} _ {\text {f l o s s i n g}}}{\partial \theta} = \sum_ {i = 1} ^ {k} \frac {\partial \lambda_ {i} ^ {2}}{\partial \theta} \tag {5}
$$
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 [12, 50].
First, we evolve an initially orthonormal system  $\mathbf{Q}_s = [\mathbf{q}_s^1, \mathbf{q}_s^2, \dots, \mathbf{q}_s^{k_1}]$  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
$$
\widetilde {\mathbf {Q}} _ {s + 1} = \mathbf {D} _ {s} \mathbf {Q} _ {s} \tag {6}
$$
at every time-step. Second, we extract the exponential growth rates using the QR decomposition,
$$
\widetilde {\mathbf {Q}} _ {s + 1} = \mathbf {Q} _ {s + 1} \mathbf {R} ^ {s + 1},
$$
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}_{s+1}^{\top} \mathbf{Q}_{s+1} = \mathbb{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_{\mathrm{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$  [49, 50]:
$$
\lambda_ {i} = \lim  _ {t \rightarrow \infty} \frac {1}{t} \log \prod_ {s = 1} ^ {t} \mathbf {R} _ {i i} ^ {s} = \lim  _ {t \rightarrow \infty} \frac {1}{t} \sum_ {s = 1} ^ {t} \log \mathbf {R} _ {i i} ^ {s}. \tag {7}
$$
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 [51]: The backward pass of the QR decomposition is given by [51, 52, 53, 54]
$$
\overline {{\mathbf {Q}}} = \left[ \overline {{\mathbf {Q}}} + \mathbf {Q} \operatorname {c o p y l t u} (\mathbf {M}) \right] \mathbf {R} ^ {- T}, \tag {8}
$$
where  $\mathbf{M} = \mathbf{R}\overline{\mathbf{R}}^T -\overline{\mathbf{Q}}^T\mathbf{Q}$  and the copy1tu function generates a symmetric matrix by copying the lower triangle of the input matrix to its upper triangle, with the element  $[\mathrm{copy1tu}(M)]_{ij} = M_{\max (i,j),\min (i,j)}$  [51, 52, 53, 54]. We denote here adjoint variable as  $\overline{T} = \partial \mathcal{L} / \partial T$ . A simple implementation of this algorithm in pseudocode is:
Algorithm 1 Algorithm for gradient flossing of  $k$  tangent space directions
initialize h, Q  
for  $e = 1 \rightarrow E$  do  
for  $t = 1 \rightarrow T$  do  
 $\mathbf{h} \leftarrow \mathbf{f}_{\theta}(\mathbf{h}, \mathbf{x})$ $\mathbf{D} \leftarrow \frac{\mathrm{d}\mathbf{h}_t}{\mathrm{d}\mathbf{h}_{t-1}}$ $\mathbf{Q} \leftarrow \mathbf{D} \cdot \mathbf{Q}$   
if  $t \equiv 0 \pmod{t_{\mathrm{ONS}}}$  then  
 $\mathbf{Q}, \mathbf{R} \leftarrow \operatorname{qr}(\mathbf{Q})$ $\gamma_i + = \log(R_{ii})$   
end if  
end for  
 $\lambda_i = \gamma_i / T$ $\theta_{e+1} \leftarrow \theta_e - \eta \frac{\partial \mathcal{L}_{\text{flossing}}}{\partial \theta}$   
end for
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 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  $\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$  [9, 12, 46]. The Jacobian matrix  $\mathbf{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
$$
\mathbf {h} _ {s + 1} = \mathbf {f} \left(\mathbf {h} _ {s}, \mathbf {x} _ {s + 1}\right) = \mathbf {W} \phi \left(\mathbf {h} _ {s}\right) + \mathbf {V} \mathbf {x} _ {s + 1}. \tag {9}
$$
Figure 1: Gradient flossing controls Lyapunov exponents and gradient signal propagation
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_1^{\mathrm{target}} = \{-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^{\mathrm{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(Nk^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.
5 Gradient Flossing: Condition Number of the Long-Term Jacobian
Figure 2: Gradient flossing reduces condition number of the long-term Jacobian A) Condition number  $\kappa_{2}$  of long-term Jacobian  $\mathbf{T}_t(\mathbf{h}_{\tau})$  as a function of time horizon  $t - \tau$  at initialization (blue) and after gradient flossing (orange). Direct numerical simulations are done with arbitrary precision floating point arithmetic (transparent lines) with 256 bits per float, asymptotic theory based on Lyapunov exponents (dashed lines) (Eq 10). B) Condition number for different number of tangent space dimensions  $m$ . Simulations (dots) and Lyapunov exponent based theory (dashed lines) at initialization (blue) and after gradient flossing (orange). Gradient flossing increases the number of tangent space dimensions available for backpropagation for a given condition number (Grey dotted line as a guide for eye for  $\kappa_{2} = 10^{5}$ .) First 15 Lyapunov exponents were flossed. C) Comparison of condition number obtained via direct numerical simulations vs. Lyapunov exponent-based. Colors denote the number of flossed Lyapunov exponents  $k$ . Parameters:  $g = 1$ , batch size  $b = 1$ ,  $N = 80$ , epochs  $= 500$ ,  $T = 500$ , gradient flossing for  $E_{f} = 500$  epochs. Input  $\mathbf{x}_{s}$  identical to delayed XOR task in Fig 3D.
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  $\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_{\mathrm{max}}$  and the smallest singular values  $\sigma_{\mathrm{min}}$ , so  $\kappa_{2}(A) = \frac{\sigma_{\mathrm{max}}(A)}{\sigma_{\mathrm{min}}(A)}$ . According to the rule of thumb given in [57], 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:
$$
\kappa_ {2} \left(\widetilde {\mathbf {Q}} _ {t + \tau}\right) = \kappa_ {2} \left(\mathbf {T} _ {t} \left(\mathbf {h} _ {\tau}\right) \mathbf {Q} _ {t}\right) = \frac {\sigma_ {1} \left(\mathbf {T} _ {t} \left(\mathbf {h} _ {\tau}\right)\right)}{\sigma_ {m} \left(\mathbf {T} _ {t} \left(\mathbf {h} _ {\tau}\right)\right)} \approx \exp \left(\left(\lambda_ {1} - \lambda_ {m}\right) (t - \tau)\right). \tag {10}
$$
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).