SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We also investigate the effects of gradient flossing during the training with orthogonal weight initializations and confirm our finding that gradient flossing improves trainability on tasks that have long time horizons to bridge. Moreover, we find that gradient flossing during training can further improve trainability. We replicated the two more challenging tasks from the main paper (Fig 4) for orthogonal initialization with variable temporal complexity and performed gradient flossing either both during and before training, only before training, or not at all.
Fig 14A shows the test accuracy for Vanilla RNNs with orthogonal initialization trained on the delayed temporal XOR task  $y_{t} = x_{t - d / 2} \oplus x_{t - d}$  with random Bernoulli process  $x \in \{0,1\}$ . The accuracy of orthogonal Vanilla RNNs falls to chance level for  $d \geq 40$  (Fig 14C). With gradient flossing before training, the trainability can be improved, but still falls close to chance level for  $d = 70$ . In contrast, for initially orthogonal networks with gradient flossing during training, the accuracy is improved to  $>80\%$  at  $d = 70$ . In this case, we preflossed for 500 epochs before task training and again after 500 epochs of training on the task. In Fig 14B, D the networks have to perform the nonlinear XOR operation  $y_{t} = x_{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 identical to Fig 4 in the main text. Again, we observe similar to networks with Gaussian initialization that flossing before training improves the performance compared to baseline, but starts failing for long delays  $d > 60$ . In contrast, orthogonal networks that are also flossed during training can solve even more difficult tasks (Fig 14D). We note that for Fig 14B and D, we trained the network only on 5000 epochs, compared to 10000 epochs in networks with random Gaussian initialization because for 10000 epochs, both networks with gradient flossing only before training and with gradient flossing before and during training were able to bridge  $d = 70$ . These results suggest that orthogonal initialization does seem to slightly improve performance for tasks with long time horizons to bridge and gradient flossing and additionally boost the performance. Thus orthogonal initialization and gradient flossing seems to go well together. It would be interesting to study if orthogonal initialization also reduces the number of gradient flossing steps necessary to improve performance.
H Additional Details on Training Tasks
In this section, we provide a more rigorous definition of the tasks used for training, as discussed in Section 3:
C
Figure 12: Gradient flossing throughout training can be detrimental to learning A) Accuracy as a function of training epochs for binary temporal delayed XOR task for gradient flossing throughout training every 100 training epochs (red). Accuracy drops down close to chance level every time after gradient flossing but recovers quickly between. Same for only 5 episodes of gradient flossing at epochs  $e \in \{0, 100, 200, 300, 400\}$  (green) and no flossing at all (blue). B) Test error as a function of training epochs. C) Gradient norm of  $\left| \frac{d\mathcal{L}}{dh_0} \right|$  as a function of training epochs for networks without gradient flossing (blue) and networks with flossing throughout training (red) and early training gradient flossing (green). Error gradient norm is boosted after each gradient flossing. In networks without gradient flossing, the gradient norms  $\left| \frac{d\mathcal{L}}{dh_0} \right|$  are much smaller overall. Parameters:  $g = 1$ , batch size  $b = 16$ ,  $N = 80$ , epochs  $= 10^4$ ,  $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})$ .
H.1 Copy task
For the copy task, the target network readout at time  $t$  is  $y_{t} = x_{t - d}$ , where  $d$  denotes the delay. We chose the input to be sampled i.i.d. from a uniform distribution between 0 and 1.
H.2 Temporal XOR task
The temporal XOR task requires the target network readout  $y_{t}$  at time  $t$  to be computed as follows:
$$
y _ {t} = \left| x _ {t - d / 2} - x _ {t - d} \right| \tag {20}
$$
where again  $d$  denotes a time delay of  $d$  time steps. In the case of  $x \in \{0,1\}^2$  and  $y \in \{0,1\}$ , the output  $y_{t}$  follows the truth table of the XOR digital logic gate (Table 2). Thus, the function  $f(x_{a}, x_{b}) = |x_{a} - x_{b}|$  can be seen as an analytical representation of the XOR gate. It is important to note that  $f(x,0) = x$  only for  $x \geq 0$ , and that this task requires a nonlinearity. The implementation can easily be constructed analytically, for example, using two rectified linear units  $\phi(x) = \max(x,0)$  the output can be constructed by
$$
f \left(x _ {a}, x _ {b}\right) = \left| x _ {a} - x _ {b} \right| = \phi \left(x _ {a} - x _ {b}\right) + \phi \left(x _ {b} - x _ {a}\right). \tag {21}
$$
Together with a delay line to transmit the signal  $x_{t - d}$  over time, this can solve the task.
Figure 13: Lyapunov exponents of trained networks with and without gradient flossing A) First Lyapunov exponents  $\lambda_{1}$  for Vanilla networks trained on spatial delayed XOR task as a function of the delay with no gradient flossing. Colored-coded is test accuracy at the end of training where red corresponds to  $100\%$  accuracy and blue to chance level  $(50\%)$ . B) Same as A for networks with gradient flossing during training. Black dashed line shows that Lyapunov exponents of successfully trained networks can be approximated by the empirical fit  $\lambda_{1}(d) = -0.2\exp(-0.03\text{delay})$ . (Protocol for gradient flossing during training same as main text Fig 4B). C) Same as A for tenth Lyapunov exponents  $\lambda_{10}$ . D) Same as B for tenth Lyapunov exponents  $\lambda_{10}$ . Same fit as in B also describes  $\lambda_{10}$ . Parameters:  $g = 1$ , batch size  $b = 16$ ,  $N = 80$ , epochs  $= 10^{4}$ ,  $T = 300$ , gradient flossing for  $E_{f} = 500$  epochs on  $k = 75$  Lyapunov exponents before training. Task: binary spatial delayed XOR, loss: cross entropy  $(y, \hat{y})$ .
I Additional Background on Lyapunov Exponents of RNNs
An autonomous dynamical system is usually defined by a set of ordinary differential equations  $\mathrm{d}\mathbf{h} / \mathrm{d}t = \mathbf{F}(\mathbf{h})$ ,  $\mathbf{h} \in \mathbb{R}^N$  in the case of continuous-time dynamics, or as a map  $\mathbf{h}_{s+1} = \mathbf{f}(\mathbf{h}_s)$  in the case of discrete-time dynamics. In the following, the theory is presented for discrete-time dynamical systems for ease of notation, but everything directly extends to continuous-time systems [49]. Together with an initial condition  $\mathbf{h}_0$ , the map forms a trajectory. As a natural extension of linear stability analysis, one can ask how an infinitesimal perturbation  $\mathbf{h}_0' = \mathbf{h}_0 + \epsilon \mathbf{u}_0$  evolves in time. Chaotic systems are sensitive to initial conditions; almost all infinitesimal perturbations  $\epsilon \mathbf{u}_0$  of the initial condition grow exponentially with time  $|\epsilon \mathbf{u}_t| \approx \exp(\lambda_1 t) |\epsilon \mathbf{u}_0|$ . Finite-size perturbations, therefore, may lead to a drastically different subsequent behavior. The largest Lyapunov exponent  $\lambda_1$  measures the average rate of exponential divergence or convergence of nearby initial conditions:
$$
\lambda_ {1} \left(\mathbf {h} _ {0}\right) = \lim  _ {t \rightarrow \infty} \frac {1}{t} \lim  _ {\epsilon \rightarrow 0} \log \frac {| | \boldsymbol {\epsilon} \mathbf {u} _ {t} | |}{| | \boldsymbol {\epsilon} \mathbf {u} _ {0} | |} \tag {22}
$$
In dynamical systems that are ergodic on the attractor, the Lyapunov exponents do not depend on the initial conditions as long as the initial conditions are in the basins of attraction of the attractor. Note that it is crucial to first take the limit  $\epsilon \to 0$  and then  $t\to \infty$ , as  $\lambda_{1}(\mathbf{h}_{0})$  would be trivially zero for a bounded attractor if the limits are exchanged, as  $\lim_{t\to \infty}\log \frac{||\epsilon\mathbf{u}_t||}{||\epsilon\mathbf{u}_0||}$  is bounded for finite perturbations even if the system is chaotic. To measure  $k$  Lyapunov exponents, one has to study the evolution of  $k$  independent infinitesimal perturbations  $\mathbf{u}_s$  spanning the tangent space:
$$
\mathbf {u} _ {s + 1} = \mathbf {D} _ {s} \mathbf {u} _ {s} \tag {23}
$$
forward pass	backward pass
RNN dynamics	O (N2b)	"
Jacobian step	O (N2k)	"
QR step	O (N k2)	"
total amortized costs per training epoch	O (N2bT)	"
total amortized costs per gradient flossing epoch	O (N2Tf(1+k/tONS+k))	"
total amortized costs of preflossing	O (N2[EbT+EpTf(1+k/tONS+k)])	"
total amortized costs flossing during training	O (N2[EbT+EpTp+EfTf(1+k/tONS+k)])	"
Table 1: Computational cost for gradient flossing and training of RNNs
$N$  denotes number of neurons,  $b$  is the batch size,  $T$  is the number of time steps in forward pass of training,  $T_{f}$  is the number of time steps in forward pass of flossing,  $t_\mathrm{ONS}$  is the reorthonormalization interval,  $k$  is the number of flossed Lyapunov exponents,  $E$  is the number of training epochs,  $E_{p}$  is the number of preflossing epochs,  $E_{f}$  is the number of flossing epochs during training. Empirically, we find that the necessary number of preflossing epochs  $E_{p}$  and flossing episodes  $E_{f}$  is much smaller than both the total number of training epochs  $E$ . Moreover,  $T_{p}$  can be smaller than  $T$ .
Table 2: XOR
input xt-d	input xt-2d	target output yt
0	0	0
0	1	1
1	0	1
1	1	0
where the  $N\times N$  Jacobian  $\mathbf{D}_s(\mathbf{h_s}) = \mathrm{df}(\mathbf{h_s}) / \mathrm{dh}$  characterizes the evolution of generic infinitesimal perturbations during one step. Note that this Jacobian along the trajectory is equivalent to a stability matrix only at a fixed point, i.e., when  $\mathbf{h}_{s + 1} = \mathbf{f}(\mathbf{h}_s) = \mathbf{h}_s$ .
We are interested in the asymptotic behavior, and therefore we study the long-term Jacobian
$$
\mathbf {T} _ {t} \left(\mathbf {h} _ {0}\right) = \mathbf {D} _ {t - 1} \left(\mathbf {h} _ {t - 1}\right) \dots \mathbf {D} _ {1} \left(\mathbf {h} _ {1}\right) \mathbf {D} _ {0} \left(\mathbf {h} _ {0}\right). \tag {24}
$$
Note that  $\mathbf{T}_t(\mathbf{h}_0)$  is a product of generally noncommuting matrices. The Lyapunov exponents  $\lambda_1 \geq \lambda_2 \dots \geq \lambda_N$  are defined as the logarithms of the eigenvalues of the Oseledets matrix
$$
\boldsymbol {\Lambda} \left(\mathbf {h} _ {0}\right) = \lim  _ {t \rightarrow \infty} \left[ \mathbf {T} _ {t} \left(\mathbf {h} _ {0}\right) ^ {\top} \mathbf {T} _ {t} \left(\mathbf {h} _ {0}\right)\right] ^ {\frac {1}{2 t}}, \tag {25}
$$
where  $\top$  denotes the transpose operation. The expression inside the brackets is the Gram matrix of the long-term Jacobian  $\mathbf{T}_t(\mathbf{h}_0)$ . Geometrically, the determinant of the Gram matrix is the squared volume of the parallelotope spanned by the columns of  $\mathbf{T}_{\mathbf{t}}(\mathbf{h}_{\mathbf{0}})$ . Thus, the exponential volume growth rate is given by the sum of the logarithms of its first  $k$  (sorted) eigenvalues. Oseledets' multiplicative ergodic theorem guarantees the existence of the Oseledets matrix  $\Lambda (\mathbf{h}_0)$  for almost all initial conditions  $\mathbf{h}_0$  [48]. In ergodic systems, the Lyapunov exponents  $\lambda_{i}$  do not depend on the initial condition  $\mathbf{h}_0$ . However, for a numerical calculation of the Lyapunov spectrum, Eq 25 cannot be used directly because the long-term Jacobian  $T_{t}(\mathbf{h}_{0})$  quickly becomes ill-conditioned, i.e., the ratio between its largest and smallest singular value diverges exponentially with time.
J Algorithm for Calculating Lyapunov Spectrum of Rate Networks
For calculating the first  $k$  Lyapunov exponents, we exploit the fact that the growth rate of a  $k$ -dimensional infinitesimal volume element is given by  $\lambda^{(m)} = \sum_{i=1}^{m}\lambda_{i}$ . Therefore,  $\lambda_{1} = \lambda^{(1)}, \lambda_{2} = \lambda^{(2)} - \lambda_{1}, \lambda_{3} = \lambda^{(3)} - \lambda_{1} - \lambda_{2}, \ldots$  [50]. The volume growth rates can be obtained via QR-decomposition.
orthogonal weight initialization
Figure 14: Gradient flossing before and during training improves trainability for orthogonal nets
A) Test accuracy for orthogonally initialized 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 (orange), and with no gradient flossing (blue) for  $d = 70$ . Solid lines are mean, transparent thin lines are individual network realizations B) Same as A for delayed spatial XOR task with  $y_{t} = x_{t - d}^{1}\oplus x_{t - d}^{2}\oplus x_{t - d}^{3}$ . C) Test accuracy as a function of task difficulty (delay  $d$ ) for delayed temporal XOR task. D) Test accuracy as a function of task difficulty (delay  $d$ ) for delayed spatial XOR task. Parameters:  $g = 1$ , batch size  $b = 16$ ,  $N = 80$ , epochs  $= 10^4$  for delayed temporal XOR, epochs  $= 5000$  for delayed spatial XOR,  $T = 300$ , flossing for  $E_{f} = 500$  epochs on  $k = 75$  Lyapunov exponents before training and during training for green lines, and only before training for orange lines. Shaded areas are  $25\%$  and  $75\%$  percentiles, solid lines are means, transparent dots are individual simulations, task loss is cross-entropy between  $y, \hat{y}$ .
First, we evolve an initially orthonormal system  $\mathbf{Q}_s = [\mathbf{q}_s^1,\mathbf{q}_s^2,\dots \mathbf{q}_s^m ]$  in the tangent space along the trajectory using the Jacobian  $\mathbf{D}_s$ :
$$
\widetilde {\mathbf {Q}} _ {s + 1} = \mathbf {D} _ {s} \mathbf {Q} _ {s} \tag {26}
$$
A continuous system can be transformed into a discrete system by considering a stroboscopic representation, where the trajectory is only considered at certain discrete time points. We use here the notation of discrete dynamical systems where this corresponds to performing the product of Jacobians along the trajectory  $\tilde{\mathbf{Q}}_{s + 1} = \mathbf{D}_s\mathbf{Q}_s$ . We study the discrete network dynamics in the limit of small time step  $\Delta t\to 0$  and for discrete time  $\Delta t = 1$ . The notation can be readily extended to continuous systems [49].
Second, we extract the exponential growth rates using the QR-decomposition,
which uniquely decomposes  $\widetilde{\mathbf{Q}}_{s+1}$  into an orthonormal matrix  $\mathbf{Q}_{s+1}$  of size  $N \times k$  so  $\mathbf{Q}_{s+1}^{\top} \mathbf{Q}_{s+1} = \mathbb{1}_{m \times m}$  and to an upper triangular matrix  $\mathbf{R}^{s+1}$  of size  $k \times k$  with positive diagonal elements. Geometrically,  $\mathbf{Q}_{s+1}$  describes the rotation of  $\mathbf{Q}_s$  caused by  $\mathbf{D}_s$  and the diagonal entries of  $\mathbf{R}^{s+1}$  describe the stretching and shrinking of the columns of  $\mathbf{Q}_s$ , while the off-diagonal elements represent the shearing. Fig 15 visualizes  $\mathbf{D}_s$  and the QR-decomposition for  $k = 2$ .
The Lyapunov exponents are given by time-averaged logarithms of the diagonal elements of  $\mathbf{R}^s$ :
$$
\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 {27}
$$
Figure 15: Geometric illustration of Lyapunov spectrum calculation. An orthonormal matrix  $\mathbf{Q}_s = [\mathbf{q}_s^1, \mathbf{q}_s^2, \dots, \mathbf{q}_s^m]$ , whose columns are the axes of an  $k$ -dimensional cube, is rotated and distorted by the Jacobian  $\mathbf{D}_s$  into an  $k$ -dimensional parallelotope  $\widetilde{\mathbf{Q}}_{s+1} = \mathbf{D}_s \mathbf{Q}_s$  embedded in  $\mathbb{R}^N$ . The figure illustrates this for  $k = 2$ , in which case the columns of  $\widetilde{\mathbf{Q}}_{s+1}$  span a parallelogram, which can be divided into a right triangle and a trapezoid and rearranged into a rectangle. Thus, the area of the gray parallelogram is the same as that of the orange rectangle. The QR-decomposition reorthonormalizes  $\widetilde{\mathbf{Q}}_{s+1}$  by decomposing it into the product of an orthonormal matrix  $\mathbf{Q}_{s+1} = [\mathbf{q}_{s+1}^1, \mathbf{q}_{s+1}^2, \dots, \mathbf{q}_{s+1}^m]$  and the upper-triangular matrix  $\mathbf{R}^{s+1}$ .  $\mathbf{Q}_{s+1}$  describes the rotation of  $\mathbf{Q}_s$  caused by  $\mathbf{D}_s$ . The diagonal entries of  $\mathbf{R}^{s+1}$  gives the stretching/shrinking along the columns of  $\mathbf{Q}_{s+1}$ , thus the volume of the parallelotope formed by the first  $k$  columns of  $\widetilde{\mathbf{Q}}_{s+1}$  is given by  $V_m = \prod_{i=1}^m \mathbf{R}_{ii}^{s+1}$ . The time-averaged logarithms of the diagonal elements of  $\mathbf{R}^s$  give the Lyapunov spectrum:  $\lambda_i = \lim_{t_{\mathrm{sim}} \to \infty} \frac{1}{t_{\mathrm{sim}}} \log \prod_{s=1}^t \mathbf{R}_{ii}^s = \lim_{t_{\mathrm{sim}} \to \infty} \frac{1}{t} \sum_{s=1}^t \log \mathbf{R}_{ii}^s$ .
Note that the QR-decomposition does not need to be performed at every simulation step, just sufficiently often, i.e., once every  $s_{\mathrm{ONS}}$  steps such that  $\widetilde{\mathbf{Q}}_{s + s_{\mathrm{ONS}}} = \mathbf{D}_{s + s_{\mathrm{ONS}} - 1} \cdot \mathbf{D}_{s + s_{\mathrm{ONS}} - 2} \dots \mathbf{D}_s \cdot \mathbf{Q}_s$  remains well-conditioned [50]. An appropriate reorthonormalization interval  $s_{\mathrm{ONS}} = t_{\mathrm{ONS}} / \Delta t$  thus depends on the condition number, the ratio of the smallest and largest singular value:
$$
\kappa_ {2} \left(\widetilde {\mathbf {Q}} _ {s + s _ {\mathrm {O N S}}}\right) = \kappa_ {2} \left(\mathbf {R} ^ {s + s _ {\mathrm {O N S}}}\right) = \frac {\sigma_ {1} \left(\mathbf {R} ^ {s + s _ {\mathrm {O N S}}}\right)}{\sigma_ {m} \left(\mathbf {R} ^ {s + s _ {\mathrm {O N S}}}\right)} = \frac {\mathbf {R} _ {1 1} ^ {s + s _ {\mathrm {O N S}}}}{\mathbf {R} _ {m m} ^ {s + s _ {\mathrm {O N S}}}}. \tag {28}
$$
An initial transient should be disregarded in the calculation of the Lyapunov spectrum because  $\mathbf{h}$  first has to converge towards the attractor and  $\tilde{\mathbf{Q}}$  has to converge to the unique eigenvectors of the Oseledets matrix (Eq 25) [60]. A simple example of this algorithm in pseudocode is:
Algorithm 2 Jacobian-based algorithm for Lyapunov spectrum
initialize h, Q  
evolve h until it is on attractor (avoid initial transient)  
evolve Q until it converges to the eigenvectors of the backward Oseledets matrix set  $\gamma_{i} = 0$   
for  $t = 1 \rightarrow T$  do  
 $\mathbf{h} \gets \mathbf{f}(\mathbf{h})$ $\mathbf{D} \gets \frac{\mathrm{df}}{\mathrm{dh}}$ $\mathbf{Q} \gets \mathbf{D} \cdot \mathbf{Q}$   
if  $s \equiv 0 \pmod{s_{\mathrm{ONS}}}$  then  
 $\mathbf{Q}, \mathbf{R} \gets \operatorname{qr}(\mathbf{Q})$ $\gamma_{i} += \log(R_{ii})$   
end if  
end for  
 $\lambda_{i} = \gamma_{i}/T$
It is guaranteed that under general conditions initially random orthonormal systems will exponentially converge towards a unique basis that is given by the eigenvectors of the Oseledets matrix Eq 25 [60]. A minimal example of this algorithm in pseudocode is shown in appendix 3. A feasible strategy to determine the reorthonormalization time interval  $t_{\mathrm{ONS}}$  is to get first a rough estimate of the Lyapunov spectrum using a short simulation time  $t_{\mathrm{sim}}$  and a small  $t_{\mathrm{ONS}}$  and repeat with a longer simulation time and a  $t_{\mathrm{ONS}}$  based on the Lyapunov spectrum of the rough estimate of the Lyapunov spectrum. Another strategy is, to first iteratively adapt  $t_{\mathrm{ONS}}$  on a short simulation run to get an acceptable condition number. It should be noted that there exists a diversity of other methods to estimate the Lyapunov spectrum [14, 49, 79, 80].
K Convergence of Lyapunov Exponents of RNNs
In Fig. 16, we demonstrate the convergence of the Lyapunov exponents. We show the estimate of the Lyapunov exponents  $\lambda_{i}$  for  $i = 1,20,60,80$  for different initial conditions but identical network realization.
Figure 16: Convergence of Lyapunov exponents Convergence of selected Lyapunov exponents  $\lambda_{i}$  for ten identical network realizations with different initial conditions with simulation time ( $i = 1, 20, 60, 80$ ) for  $\sigma = 1$  and  $g = 1$ . (Other parameters:  $N = 80$ ,  $t_{\mathrm{sim}} = 100$  steps,  $t_{\mathrm{ONS}} = 1$ ).