SlowGuess/NeurIPS_2023_OCR huggingface dataset:
During gradient flossing, the actual Lyapunov exponents of different random network realizations converge close to the target Lyapunov exponent  $\lambda_1^{\mathrm{target}} = 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  $\lambda_1^{\mathrm{target}} = -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
$$
\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},
$$
The initial entries of  $\mathbf{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  $\phi (x) = \max (x,0)$ .  $\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. 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
Figure 7: 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$ . C) Median test accuracy for delayed temporal XOR task as a function of delay  $d$  and  $k$  for networks with gradient flossing during training (500 steps of gradient flossing at epochs  $e \in \{0, 100, 200, 300, 400\}$ ). D) Same as B for preflossing only. 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, C, and only before training for C. Shaded areas are  $25\%$  and  $75\%$  percentiles, solid lines are means, transparent dots are individual simulations, task loss: cross-entropy btw.  $y$ ,  $\hat{y}$ .
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  $\mathbf{h}_0$ .
$$
\frac {\partial \mathcal {L} _ {t}}{\partial \mathbf {h} _ {0}} = \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 \mathbf {h} _ {0}} = \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} + l}}{\partial \mathbf {h} _ {\tau^ {\prime}}}\right) \delta_ {\tau 0} = \frac {\partial L _ {t}}{\partial \mathbf {h} _ {t}} \mathbf {T} _ {t} (\mathbf {h} _ {0}) \tag {18}
$$
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  $\frac{\partial\mathcal{L}_t}{\partial\mathbf{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 \in \{0, 100, 200, 300, 400\}$ , while the other is trained without gradient flossing.
For the network without gradient flossing, the gradient norm of  $\left|\frac{d\mathcal{L}}{dh_0}\right|$  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  $\left|\frac{d\mathcal{L}}{dh_0}\right|$  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  $\left|\frac{d\mathcal{L}}{dh_0}\right|$  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  $\left|\frac{d\mathcal{L}}{dh_0}\right|$  and training success consistently across numerous network realizations (Fig 8C and D). These findings suggest that the gradient norm  $\left|\frac{d\mathcal{L}}{dh_0}\right|$  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  $\leq 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  $\left|\frac{d\mathcal{L}}{dh_0}\right| < 0.001$ , we observe for this task that the increase of gradient norm  $\left|\frac{d\mathcal{L}}{dh_0}\right|$  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  $\frac{d\mathcal{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  $\sigma_{i}\left(\frac{d\mathcal{L}}{dW}\right)$  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  $\mathbf{T}_t(\mathbf{h}_{\tau})$ , 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:
$$
\frac {\partial \mathcal {L} _ {t}}{\partial \mathbf {W}} = \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 \mathbf {W}} = \frac {\partial L _ {t}}{\partial \mathbf {h} _ {t}} \sum_ {\tau} \mathbf {T} _ {t} (\mathbf {h} _ {\tau}) \frac {\partial \mathbf {h} _ {\tau}}{\partial \mathbf {W}} \tag {19}
$$
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  $\sigma_{i}\left(\frac{d\mathcal{L}}{dW}\right)$  as a function of training epoch reveal that the subordinate singular values subordinate singular value  $\sigma_{20}$  and  $\sigma_{40}$  exhibit peaks at the times of gradient flossing, while the first singular value  $\sigma_{1}$  only shows a slight peak (Fig 10A). This indicates that gradient flossing increases the effective rank of the recurrent weight gradient  $\frac{d\mathcal{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  $\sigma_{20}$  and  $\sigma_{40}$ ) rapidly diminish to extremely small values over training epochs and remain very small throughout training. Note however that the leading singular values  $\sigma_{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  $\left|\frac{d\mathcal{L}}{dh_0}\right|$ , 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
Figure 8: Gradient flossing boosts norm of long-term Jacobian A) Gradient norm of  $\left| \frac{d\mathcal{L}}{dh_0} \right|$  as a function of training epochs for networks without flossing (blue) and networks with flossing during training (orange). Error gradient norm is boosted after gradient flossing at epochs  $e \in \{0,100,200,300,400,500\}$ . In networks without gradient flossing, the gradient norms  $\left| \frac{d\mathcal{L}}{dh_0} \right|$  are much smaller overall. One out of ten random network realizations with solid line, the other 9 with transparent line. B) Accuracy as a function of epoch, same depiction and network realizations as in A. Note that accuracy of networks with gradient flossing grows beyond chance level approximately when the gradient norm  $\left| \frac{d\mathcal{L}}{dh_0} \right|$  becomes macroscopically large. C) Same as A in linear scale. Mean final test loss as a function of task difficulty (delay  $d$ ) for delayed copy task. Different colors are different network realizations with gradient flossing during training. Black lines are without any gradient flossing. D) Accuracy as a function of epochs, same colors as in C. Note that for all network realizations the moments where gradient norm  $\left| \frac{d\mathcal{L}}{dh_0} \right|$  becomes macroscopically large coincides with the moment the accuracy is beyond chance level. Parameters:  $g = 1$ , batch size  $b = 16$ ,  $N = 80$ , epochs  $= 10^4$ ,  $T = 300$ , flossing for  $E_f = 500$  epochs on  $k = 75$  Lyapunov exponents before training. Task: binary delayed XOR, delay  $d = 70$ , loss: cross entropy  $(y, \hat{y})$ .
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 \in \{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  $\left|\frac{d\mathcal{L}}{dh_0}\right|$  is initially boosted by gradient flossing, but stays close to indistinguishable once the gradient norm  $\left|\frac{d\mathcal{L}}{dh_0}\right|$  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  $\phi$ , where  $\phi'(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
Figure 9: Increase of gradient norm precedes epoch when accuracy exceeds chance level A) Gradient norm of  $\left| \frac{d\mathcal{L}}{dh_0} \right|$  as a function of training epochs for 20 network realizations with flossing during training. Epochs are aligned to last epoch where accuracy is  $\leq 50\%$ . B) Same task and simulations as in A, but here accuracy as a function of epoch, for 20 network realizations with flossing during training. Epochs are aligned to the last epoch with  $\left| \frac{d\mathcal{L}}{dh_0} \right| < 0.001$ . Different colors are different network realizations. Parameters:  $g = 1$ , batch size  $b = 16$ ,  $N = 80$ , epochs  $= 10^4$ ,  $T = 300$ , flossing for  $E_f = 500$  epochs on  $k = 75$  Lyapunov exponents early during training at training epochs  $e \in \{0, 100, 200, 300, 400\}$ . Task: binary delayed XOR, delay  $d = 70$ , loss: cross entropy  $(y, \hat{y})$ .
gradient flossing a systematic relationship between the first Lyapunov exponent and the delay, that can be fitted by approximately by  $\lambda_1(d) = -0.2\exp(-0.03\text{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 \times N$  Jacobian of the memory state. From a dynamical systems standpoint, a  $2N \times 2N$  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 that 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  $\mathcal{O}\left(N^2 b\right)$  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  $\mathcal{O}\left(N^2 k\right)$  per time step, where  $k$  is the number of flossed Lyapunov exponents. Third, the QR decomposition, which scales with  $\mathcal{O}\left(N k^2\right)$ , where  $k$  is the number of Lyapunov exponents considered.
Figure 10: Gradient flossing decreases condition number of recurrent weight error gradient A) Singular values of recurrent weight gradient  $\sigma_{i}\left(\frac{d\mathcal{L}}{dW}\right)$  as a function of training epochs for singular values  $i\in 1,20,40$  for networks without gradient flossing (blue) and early training gradient flossing (green). At epochs of gradient flossing, the subordinate singular value  $\sigma_{20}$  and  $\sigma_{40}$  are peaked, while the first singular value  $\sigma_{1}$  has only a slight peak. This indicates that gradient flossing increases the effective rank of the recurrent weight gradient  $\frac{d\mathcal{L}}{dW}$ . B) Accuracy as a function of training epochs. Note that accuracy of networks with gradient flossing grows beyond chance level approximately when the subordinate singular values singular value  $\sigma_{20}$  and  $\sigma_{40}$  are peaked increase, which enables high-dimensional parameters updates. Parameters:  $g = 1$ , batch size  $b = 16$ ,  $N = 80$ , epochs  $= 10^{3}$ ,  $T = 300$ , gradient flossing for  $E_{f} = 500$  epochs on  $k = 75$  Lyapunov exponents before training. Task: binary delayed XOR, delay  $d = 70$ , loss: cross entropy  $(y,\hat{y})$ .
Together, this results in a total amortized cost of  $\mathcal{O}\left(N^{2}bT\right)$  per training epoch, where  $T$  is the number of training time steps and a total amortized costs per flossing epoch of  $\mathcal{O}\left(N^{2}T_{f}(1 + k / t_{\mathrm{ONS}} + k)\right)$  where  $T_{f}$  is the number of flossing time steps.
In case of preflossing, thus, the total computation cost scale with  $\mathcal{O}\left(N^2 [EbT + E_pT_f(1 + k / t_{\mathrm{ONS}} + k)]\right)$ , 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  $\mathcal{O}\left(N^2 [EbT + E_pT_p + E_fT_f(1 + k / t_{\mathrm{ONS}} + k)]\right)$ , 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  $\sim 37$  seconds, while the overall training on 10000 training epochs with batch size  $b = 16$  took  $\sim 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.
Figure 11: Gradient flossing decreases condition number of recurrent weight error gradient Same as Fig 10 for different network realization.
G Additional controls