SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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 \to \infty$  and usually do not capture finite-size effects [58] (see appendix G).
6 Initial Gradient Flossing Improves Trainability
Figure 3: Gradient flossing improves trainability on tasks that involve long time horizons A) Test error for Vanilla RNNs trained on delayed copy task  $y_{t} = x_{t - d}$  for  $d = 40$  with and without gradient flossing flossing. Solid lines are medians across 5 network realizations. B) Same as A for delayed XOR task with  $y_{t} = |x_{t - d / 2} - x_{t - d}|$ . C) Mean final test loss as a function of task difficulty (delay  $d$ ) for delayed copy task. D) Mean final test loss as a function of task difficulty (delay  $d$ ) for delayed 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. Shaded regions in C and D indicate the  $20\%$  and  $80\%$  percentiles and solid line shows mean. Dots are individual runs. Task loss:  $\mathrm{MSE}(y, \hat{y})$ .
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}_{\mathrm{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 - 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
Figure 4: Gradient flossing during training further improves trainability
A) Test accuracy for Vanilla RNNs trained on delayed temporal binary XOR task  $y_{t} = x_{t - d / 2} \oplus 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_{t - d}^{1}\oplus x_{t - d}^{2}\oplus x_{t - d}^{3}$  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  $\hat{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}\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_{t - d}^{1}\oplus x_{t - d}^{2}\oplus x_{t - d}^{3}$  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.
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.
Figure 5: Gradient flossing for different numbers of flossed Lyapunov exponents  
A) Test accuracy for delayed temporal XOR task as a function of delay  $d$  with different numbers flossed Lyapunov exponents  $k$ . B) Same data as A but here test accuracy as a function of number of flossed Lyapunov exponents  $k$ . Parameters:  $g = 1$ , batch size  $b = 16$ ,  $N = 80$ , epochs  $= 10^4$  for delayed temporal XOR, epochs  $= 5000$  for delayed spatial XOR,  $T = 300$ , gradient flossing for  $E_f = 500$  epochs before training and during training for A, B. Shaded areas are  $25\%$  and  $75\%$  percentile, solid lines are means, transparent dots are individual simulations, task loss: cross-entropy between  $y$  and  $\hat{y}$ .
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 [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.
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.
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 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).
[1] P. WERBOS. Beyond Regression .. Ph. D. dissertation, Harvard University, 1974.
[2] DB Parker. Learning-logic (TR-47). Center for Computational Research in Economics and Management Science. MIT-Press, Cambridge, Mass, 8, 1985.
[3] Y. LECUN. Une procedure d'apprentissage pour reseau a seuil asymetrique. Proceedings of Cognitive 85, pages 599-604, 1985.
[4] David E. Rumelhart, Geoffrey E. Hinton, and Ronald J. Williams. Learning representations by backpropagating errors. Nature, 323(6088):533, October 1986.
[5] Sepp Hochreiter. Untersuchungen zu dynamischen neuronalen Netzen. PhD thesis, 1991.
[6] Sepp Hochreiter and Jürgen Schmidhuber. Long Short-Term Memory. Neural Computation, 9(8):1735-1780, 1997.
[7] 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.
[8] Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio. On the difficulty of training Recurrent Neural Networks. arXiv:1211.5063 [cs], November 2012. arXiv: 1211.5063.
[9] 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.
[10] 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.
[11] 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.
[12] Rainer Engelken, Fred Wolf, and L. F. Abbott. Lyapunov spectra of chaotic recurrent neural networks. Physical Review Research, 5(4):043044, October 2023.
[13] J. P. Eckmann and D. Ruelle. Ergodic theory of chaos and strange attractors. Reviews of Modern Physics, 57(3):617-656, July 1985.
[14] Arkady Pikovsky and Antonio Politi. Lyapunov Exponents: A Tool to Explore Complex Dynamics. Cambridge University Press, Cambridge, February 2016.
[15] 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.
[16] 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.
[17] Tomáš Mikolov. Statistical language models based on neural networks. Ph.D. thesis, Brno University of Technology, Faculty of Information Technology, Brno, CZ, 2012.
[18] Ian Goodfellow, Yoshua Bengio, and Aaron Courville. Deep Learning. MIT Press, November 2016. Google-Books-ID: omivDQAAQBAJ.
[19] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer Normalization, July 2016. arXiv:1607.06450 [cs, stat].
A Backpropagation Through QR Decomposition
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 1 t u} (\mathbf {M}) \right] \mathbf {R} ^ {- T} \tag {11}
$$
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  $[\mathrm{copyltu}(M)]_{ij} = M_{\max (i,j),\min (i,j)}$  [51, 52, 53, 54]. Adjoint variable are written here as  $\overline{T} = \partial \mathcal{L} / \partial 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 BackwardsLinalg.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
Figure 6: Gradient flossing for recurrent LSTM networks and recurrent ReLU networks  
A) First Lyapunov exponent of LSTM network as a function of training epochs. Minimizing the mean squared error between estimated first Lyapunov exponent and target Lyapunov exponent  $\lambda_1 = 0$  by gradient descent. First Lyapunov exponent of LSTM network (solid lines) converges to target value (thick dashed lines) within less than 100 epochs. 10 random LSTM RNNs were initialized with Gaussian recurrent weights, where standard deviations of weight scaling were drawn  $g \sim \mathrm{Unif}(0,1)$ .  
B) Gradient flossing minimizes the square of the first Lyapunov exponent of random recurrent LSTM networks over epochs. C) Same as A for recurrent ReLU network. Here networks were initialized with Gaussian recurrent weights  $W_{ij} \sim \mathcal{N}(-0.1, g^2/N)$  where values of  $g$  were drawn  $g \sim \mathrm{Unif}(0,1)$ .  
D) B) for recurrent ReLU network. Parameters: network size  $N = 32$  with 10 network realizations. Shaded regions in B, D are  $25\%$  and  $75\%$  percentiles, solid line shows median.
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]:
$$
f _ {t} = \sigma_ {g} \left(U _ {f} h _ {t - 1} + W _ {f} x _ {t} + b _ {f}\right) \tag {12}
$$
$$
o _ {t} = \sigma_ {g} \left(U _ {o} h _ {t - 1} + W _ {o} x _ {t} + b _ {o}\right) \tag {13}
$$
$$
i _ {t} = \sigma_ {g} \left(U _ {i} h _ {t - 1} + W _ {i} x _ {t} + b _ {i}\right) \tag {14}
$$
$$
\tilde {c} _ {t} = \sigma_ {h} \left(U _ {c} h _ {t - 1} + W _ {c} x _ {t} + b _ {c}\right) \tag {15}
$$
$$
c _ {t} = f _ {t} \odot c _ {t - 1} + i _ {t} \odot \tilde {c} _ {t} \tag {16}
$$
$$
h _ {t} = o _ {t} \odot \phi \left(c _ {t}\right) \tag {17}
$$
where  $\odot$  denotes the Hadamard product,  $\sigma_g(x) = \frac{1}{1 + \exp(-x)}$  is the sigmoid function,  $\sigma_h(x) = \tanh (x)$  and entries of the matrices  $U_{x}$  are drawn from  $U_{x}\sim \mathcal{N}(0,g_{x}^{2} / 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.