SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Nicolas Zucchet*, Robert Meier*, Simon Schug*  
Asier Mujika, João Sacramento
Department of Computer Science, ETH Zürich {nzucchet,romeier,sschug,asiemrjoo}@ethz.ch
Online learning holds the promise of enabling efficient long-term credit assignment in recurrent neural networks. However, current algorithms fall short of offline backpropagation by either not being scalable or failing to learn long-range dependencies. Here we present a high-performance online learning algorithm that merely doubles the memory and computational requirements of a single inference pass. We achieve this by leveraging independent recurrent modules in multi-layer networks, an architectural motif that has recently been shown to be particularly powerful. Experiments on synthetic memory problems and on the challenging long-range arena benchmark suite reveal that our algorithm performs competitively, establishing a new standard for what can be achieved through online learning. This ability to learn long-range dependencies offers a new perspective on learning in the brain and opens a promising avenue in neuromorphic computing.
How can the connections between neurons in a neural network be adjusted to improve behavior? This question, known as the credit assignment problem, is central in both neuroscience [1] and machine learning [2], owing to its fundamental importance for elucidating learning mechanisms in the brain and constructing intelligent artificial systems. However, the complex and nonlinear nature of neural network processing makes the precise allocation of credit an intricate task.
Deep learning provides a compelling solution to the credit assignment problem via gradient descent, which refines network parameters along the locally most promising direction. For networks processing temporal sequences, gradient computation is made possible by backpropagation-through-time [BPTT; 2-4]. BPTT stores and revisits neural activity in reverse-time order to understand how infinitesimal changes to neural activity, and thus to network parameters, would have impacted the objective function. One important drawback of this algorithm is its requirement to store the entire activity trajectory in memory, which constrains the sequence length for exact gradient computation, impairing the learning of long-term interactions. This constraint becomes a critical bottleneck when working within memory-limited systems, such as neuromorphic hardware [5] and presumably the brain [6].
Alternatives to BPTT for gradient computation do exist. One such approach, forward-mode differentiation [7, 8], involves computing gradients online as the input sequence is processed, by keeping track of the sensitivity of neural activity with respect to each of the network parameters. This marks a qualitative departure from BPTT, as it prepares for all potential future trajectories simultaneously; by contrast, BPTT focuses on improving the activity of a past trajectory. Importantly, the memory footprint of this approach does not depend on sequence length. Still, it remains intractable for real-world applications and likely infeasible in the brain due to its cubic memory scaling and quartic computational complexity in the number of neurons. Recent research focused on approximation
strategies to make online gradient estimation more tractable for general-purpose recurrent networks [9-18]. Our work takes a fundamentally different approach: instead of tailoring the learning algorithm to the neural network architecture, we fix the learning algorithm and seek an architecture that makes it tractable.
We build upon recent advances in linear state space models, a class of recurrent neural networks (RNNs) employing linear recurrent blocks [19-23]. These blocks are stacked and interconnected through nonlinear networks. The key insights from this line of research are that linear recurrent connections simplify temporal credit assignment and enable parallel temporal processing, while nonlinearities between recurrent blocks ensure that network expressiveness remains comparable to that of densely-connected nonlinear RNNs. Much, if not all, of the state-of-the-art performance of those models on long-range temporal tasks [24] can be maintained by transitioning from real-valued to complex-valued neural activities [21-23], and restricting the recurrent connectivity matrix to be diagonal. Recurrent neurons within a given layer are now independent of each other. This greatly improves the tractability of online gradient estimation, as the recurrent parameters of a given neuron do not impact other neurons. We leverage this property to achieve exact online gradient computation within a single layer with as little as twice the memory and compute requirements needed for inference. Further, we demonstrate how this leads to improved gradient estimation compared to existing online learning algorithms when recurrent layers are stacked.
This paper is organized as follows. We start by briefly reviewing existing gradient-based online learning methods in Section 2.1. Next, we introduce the concept of independent recurrent modules, showing how some of the recent high-performance models mentioned above fit in this framework in Section 2.2. Deriving our learning rule requires complex differentiation; we give a concise overview of those tools in Section 2.3. In Section 3, we detail our online learning algorithm that combines exact differentiation within a layer of recurrent independent modules with spatial backpropagation across layers. Finally, in Section 4, we analyze our algorithm and relevant baselines on a synthetic copy task and show that it can learn sequential tasks with sequence lengths up to over 4000 steps.
2.1 Online gradient-based RNN learning
We study gradient-based learning of recurrent neural networks, which process input data  $x_{1},\ldots ,x_{T}$  sequentially while maintaining an internal (hidden) state  $h_t$ . The objective of learning is to minimize a cumulative loss  $L(\theta) = \sum_{t=1}^{T} L_t(\theta)$  which measures performance on a task at hand as a function of network parameters  $\theta$ . The standard algorithm for computing the gradient  $\nabla L(\theta)$  is the offline backpropagation-through-time method, which requires storing the entire input  $x_{1:T}$ , loss  $L_{1:T}$  and internal activity  $h_{1:T}$  sequences, and then revisiting them proceeding backwards in time. Here, we focus on online algorithms which carry the information needed to compute or estimate  $\nabla L_t(\theta)$  forward in time. This enables simultaneous processing of inputs and learning for RNNs, without storing past data and network states. In principle, online algorithms can learn arbitrarily long temporal dependencies as well as seamlessly handle sequences of arbitrary length  $T$ .
The classical alternative to BPTT for forward-in-time gradient computation is known as real-time recurrent learning [RTRL; 7] in the context of RNNs², a method which has its roots in control theory [26]. While RTRL enables online gradient-based learning, it requires storing  $\mathrm{d}_{\theta}h_{t}$  in memory and updating it as the RNN processes its inputs. The size of this auxiliary variable is  $\mathcal{O}(n^3)$ , where  $n = |h|$  is the number of hidden units in the RNN. For comparison, the memory requirements of BPTT are  $\mathcal{O}(nT)$ . In practice, this precludes the usage of RTRL for all but the smallest of models.
There has been much effort in developing memory-efficient alternatives to RTRL. We now briefly discuss these prior efforts while referring to a recent review by Marschall et al. [14] for a more detailed treatment. We divide prior work into three broad categories. One class of algorithms relies on neglecting terms in  $\mathrm{d}_{\theta}h_{t}$  to reduce its size, thereby creating a biased gradient estimator. Typically this requires introducing crude approximations, which allow going from  $\mathcal{O}(n^3)$  to a tractable  $\mathcal{O}(n^2)$  size. Despite such approximations, in many cases, performance still holds in non-trivial tasks [13, 15, 16]. At the end of this spectrum sits instantaneous (spatial) backpropagation, which neglects all temporal
2More generally, BPTT can be seen as a special case of reverse-mode automatic differentiation, and RTRL of forward-mode automatic differentiation [25], applied to the problem of RNN learning.
dependencies in the hidden state when approximating the gradient. A second class of algorithms relies on stochastic estimation; this allows retaining unbiased gradient estimates, at the expense of introducing variance [9-11, 17, 18]. Finally, a third class of methods introduces gradient models [critics; 27-29] to produce gradient estimates online. The critics themselves are then either trained separately offline, making such methods hybrid on/offline; or fully online, using temporal difference techniques [30]. We note that despite their widespread use in reinforcement learning, it is not yet well understood whether temporal difference methods can reliably and efficiently improve the performance of a gradient critic on real-world RNN learning problems.
As noted by Irie et al. [31], the RTRL literature mostly focuses on single-layer recurrent networks and remains scarce for deeper networks. Recently, Javed et al. [32] developed a greedy learning algorithm where a growing network is progressively frozen and trained one layer at a time. Existing approximations such as [13, 15] do not prescribe how to learn the parameters of remote layers, and the multi-layer case is not considered in the respective papers. Introducing a powerful RTRL algorithm that scales to networks of multiple layers is the main algorithmic contribution of this paper. This property is of great empirical relevance given the power of depth to learn the temporal structure [e.g. 19, 23].
2.2 Linear recurrent units and independent recurrent modules
Instead of developing approximate, general-purpose forward-mode differentiation methods, our goal shifts towards seeking an expressive architecture allowing exact, tractable online gradient calculation. We propose that networks with linear recurrent units [LRU; 23] and, more generally networks with independent recurrent modules, are particularly well-suited for this purpose.
A linear recurrent unit, depicted in Figure 1, is defined as
$$
h _ {t + 1} = \lambda \odot h _ {t} + B x _ {t + 1}, \quad y _ {t} = \operatorname {R e} \left[ C h _ {t} \right] + D x _ {t}, \tag {1}
$$
with  $\odot$  the element-wise product. Here,  $x_{t} \in \mathbb{R}^{H}$  represents the input received by the LRU at time  $t$ ,  $h_t \in \mathbb{C}^N$  denotes its internal state, and  $y_{t} \in \mathbb{R}^{H}$  its output. The parameters of the unit include  $\lambda \in \mathbb{C}^N$ ,  $B \in \mathbb{C}^{N \times H}$ ,  $C \in \mathbb{C}^{H \times N}$  and  $D \in \mathbb{R}^{H \times H}$ . The version of the LRU we use in our experiments includes an element-wise normalization factor for the input  $Bx$  and uses an exponential parametrization of  $\lambda$  for network stability. We omit these details in the main text for conciseness; see Appendix A.1 for more details.
LRUs differ from traditional recurrent layers in deep learning: they have linear neural dynamics, a complex-valued state  $h_t$ , and a diagonal connectivity pattern. Orvieto et al. [23] found that the absence of temporal linearity in networks that stack those units through nonlinear connections (see Fig. 1) does not alter expressivity and eases gradient-based learning, a notoriously difficult process for nonlinear RNNs [33, 34]. In addition, the diagonal structure of the recurrence matrix provides several benefits over fully connected ones. First, it affords an easy way to control the eigenvalues of the Jacobian of the system and thus ensure that neural dynamics remain stable. Due to the linearity, it also enables processing the input sequence in parallel [35], significantly accelerating the training of such models on modern computers [22]. Importantly, despite its diagonal parametrization, the LRU remains functionally equivalent to a linear recurrent layer with dense recurrence matrix  $A$ , as  $A$  can be approximated as accurately as needed by a complex-diagonalizable matrix.
Each complex neuron in an LRU is an independent recurrent module, meaning its current state does not impact the dynamics of other modules. This property greatly simplifies online credit assignment (see Section 3). We focus on this specific architecture due to its simplicity and great empirical performance, but our theoretical insights also apply to networks of independent recurrent modules with low-dimensional state vectors per module.
2.3 A primer on complex differentiation
The use of complex-valued networks, such as the LRU, and hence complex differentiation remains relatively scarce. In the following, we provide a concise review of the tools of complex differentiation integral to the derivation of our online learning rule. We use  $f$  and  $g$  to denote complex-valued functions that take the complex variable  $z$  as input.
Example of IRM: the linear recurrent unit (LRU)
Figure 1: (Left) Overview of the class of neural networks we consider in this paper. We stack layers of independent recurrent modules (IRMs), augmented with layer norm [36] and gated linear units [GLU; 37]. Light red indicates instantaneous spatial processing, dark red temporal processing. When we examine networks with fully connected recurrent layers, only the dark red block is modified. The linear recurrent unit is the instantiation of a layer of IRMs we use in our experiments. (Right) Overview of our learning rule. As an input sequence is processed, hidden states  $h_t$  and their sensitivities  $e_t$  to the parameters are updated. Learning ensues by combining the sensitivities  $e_t$  with spatially backpropagated error signals. No information flows in reverse time; our rule is fully online.
The Wirtinger derivatives of  $f$  are defined through
$$
\frac {\mathrm {d} f}{\mathrm {d} z} := \frac {1}{2} \left(\frac {\mathrm {d} f}{\mathrm {d R e} [ z ]} - i \frac {\mathrm {d} f}{\mathrm {d I m} [ z ]}\right), \frac {\mathrm {d} f}{\mathrm {d} \bar {z}} := \frac {1}{2} \left(\frac {\mathrm {d} f}{\mathrm {d R e} [ z ]} + i \frac {\mathrm {d} f}{\mathrm {d I m} [ z ]}\right). \tag {2}
$$
Using them for complex differentiation allows using similar calculus rules as for real functions. Note that we use the row convention for derivatives, that is  $\mathrm{d}_z f$  is a row vector of size  $|z|$ . The following formula holds in general  $\overline{\mathrm{d}_z f} = \mathrm{d}_{\bar{z}} \bar{f}$ .
The complex derivative of a complex function is similar to a  $2 \times 2$  real-valued matrix as both  $\mathrm{d}_z f \in \mathbb{C}$  and  $\mathrm{d}_{\bar{z}} f \in \mathbb{C}$  are necessary to characterize it. Yet, there exists a subclass of functions, called holomorphic functions for which it can be reduced to a 2-dimensional real-valued vector, leading to a more compact representation of derivatives. A continuous function  $f$  is holomorphic if it satisfies the Cauchy-Riemann equations
$$
\frac {\mathrm {d} \operatorname {R e} [ f ]}{\mathrm {d} \operatorname {R e} [ z ]} = \frac {\mathrm {d} \operatorname {I m} [ f ]}{\mathrm {d} \operatorname {I m} [ z ]} \text {a n d} \frac {\mathrm {d} \operatorname {R e} [ f ]}{\mathrm {d} \operatorname {I m} [ z ]} = - \frac {\mathrm {d} \operatorname {I m} [ f ]}{\mathrm {d} \operatorname {R e} [ z ]}, \text {i . e .} \frac {\mathrm {d} f}{\mathrm {d} \bar {z}} = 0. \tag {3}
$$
Any affine function, as well as the composition of two holomorphic functions, are themselves holomorphic.
The chain rule of complex differentiation, crucial for automatic differentiation, is
$$
\frac {\mathrm {d} (f \circ g)}{\mathrm {d} z} = \frac {\mathrm {d} f}{\mathrm {d} g} \frac {\mathrm {d} g}{\mathrm {d} z} + \frac {\mathrm {d} f}{\mathrm {d} \bar {g}} \frac {\mathrm {d} \bar {g}}{\mathrm {d} z}. \tag {4}
$$
When either  $f$  or  $g$  is holomorphic, the second term vanishes as  $d_{\overline{g}}f = 0$  or  $\mathrm{d}_z\bar{g} = \overline{\mathrm{d}_{\overline{z}}g} = 0$ . When  $f$  is a real-valued function, it can be optimized through gradient descent by iteratively updating its input variable  $z$  through  $\Delta z \propto -\mathrm{d}_{\mathrm{Re}[z]}f^{\top} - i\mathrm{d}_{\mathrm{Im}[z]}f^{\top} = -2\mathrm{d}_{\overline{z}}f^{\top}$ .
3 Online learning of networks of independent recurrent modules
Based on the foundations laid down in the preceding section we now derive our online gradient-based learning algorithm for multi-layer networks of independent recurrent modules. We first focus on a single layer, demonstrating that exact forward-mode differentiation is tractable. This insight then guides the derivation of our rule for multi-layer networks.
3.1 Single-layer networks
We focus on parameters  $\theta$  that influence the hidden states; computing the gradient of any other parameter does not require temporal credit assignment. For the LRU, we have  $\theta = \{\lambda, B\}$ . Recall that  $L(\theta) = \sum_{t=1}^{T} L_t(y_t(\theta))$  denotes the loss function that measures how good the outputs  $y_{1:T}$  of a network parametrized by  $\theta$  are. Its derivative  $\mathrm{d}_{\theta}L(\theta)$  can be calculated using forward-mode (complex) differentiation:
$$
\frac {\mathrm {d} L}{\mathrm {d} \theta} = \sum_ {t = 1} ^ {T} \frac {\partial L}{\partial h _ {t}} \frac {\mathrm {d} h _ {t}}{\mathrm {d} \theta} + \frac {\partial L}{\mathrm {d} \bar {h} _ {t}} \frac {\mathrm {d} \bar {h} _ {t}}{\mathrm {d} \theta} = \sum_ {t = 1} ^ {T} \frac {\partial L _ {t}}{\partial h _ {t}} \frac {\mathrm {d} h _ {t}}{\mathrm {d} \theta}. \tag {5}
$$
As mentioned in Section 2.3, the last equality holds as  $h_{t + 1}$  is a holomorphic function of  $h_t$  and of  $\theta$ , hence  $h_t$  is a holomorphic function of  $\theta$  by recursive composition, and  $L$  only directly depends on  $h_t$  through  $L_t$ . The term  $\delta_t \coloneqq \mathrm{d}_{h_t}L^\top$  that is here equal to  $\mathrm{d}_{h_t}L_t^\top$  can easily be computed by spatial backpropagation, as the output  $y_t$  at time  $t$  only depends on the current hidden state  $h_t$ . We are left with computing the sensitivities  $\mathrm{d}_{\theta}h_t$  of the states to the parameters.
Independent recurrent modules do not influence each other. The parameters  $\theta_{i}$  that directly influence the state  $h_{t,i}$  of module  $i$  never impact the state  $h_{t',j}$  of another module. As a consequence, the number of non-zero entries of the sensitivity  $\mathrm{d}_{\theta}h_{t}$  grows linearly with the size of  $\theta$  whenever the number of recurrent neurons within each module is fixed. Applying this to the LRU,  $\mathrm{d}_{\theta}h_{t}$  is entirely characterized by  $e_t^\lambda \coloneqq (\mathrm{d}_{\lambda_i}h_{t,i})_i$  and  $e_t^B \coloneqq (\mathrm{d}_{B_{ji}}h_{t,j})_{ji}$ . Differentiating Equation 1 using the product rule gives the corresponding updates:
$$
e _ {t + 1} ^ {\lambda} = \lambda \odot e _ {t} ^ {\lambda} + h _ {t}, e _ {t + 1} ^ {B} = \operatorname {d i a g} (\lambda) e _ {t} ^ {B} + 1 x _ {t + 1} ^ {\top}, \tag {6}
$$
with 1 a vector of size  $|h|$  filled with ones. More detail on the derivation of Equation 6 and on how to efficiently simulate this update are given in Appendix A.2. Keeping track of those quantities only requires considering an additional hidden state of size  $|\theta|$ . Finally, the  $\lambda$  and  $B$  updates can be obtained by following the gradient, as calculated in Equation 5:
$$
\Delta \lambda \propto \sum_ {t = 1} ^ {T} \delta_ {t} \odot e _ {t} ^ {\lambda}, \quad \Delta B \propto \sum_ {t = 1} ^ {T} \operatorname {d i a g} (\delta_ {t}) e _ {t} ^ {B}. \tag {7}
$$
Interestingly, all the  $e$ -updates are local to the neuron or synapse in question, and no approximations were required to accomplish this. This feature makes the algorithm particularly promising for neuroscience and neuromorphic engineering, where localized computation is highly desirable. The parameter update for  $\lambda$  and  $B$  is also fully local, as it combines a parameter-specific sensitivity, sometimes considered as an eligibility trace [13], and a postsynaptic error term.
The idea that element-wise recurrence simplifies RTRL precedes our work. It can be found in early work by Mozer [38] and Gori et al. [39], and has been revisited recently [32, 31]. In this paper, we extend this insight to complex numbers and thus do not lose expressivity, unlike previous work. We also note that some approximations to RTRL such as e-prop [13] or SnAp-1 [15] end up being exact when applied to networks with independent recurrent modules.
3.2 Multi-layer networks
The derivation in the last section presumes that the loss  $L$  only directly depends on  $h_t$  through  $L_t$ . This assumption no longer holds when layers are stacked, which is crucial to the expressiveness of the model. In the following, we explain how we can extend our rule to the multilayer case. Let us consider layer  $l$  of the network where we aim to compute the gradient of the loss  $L$  with respect to its parameters  $\theta^l$ . The sensitivity  $\mathrm{d}_{\theta^l}h_t^l$  can be computed as before, assuming independent recurrent modules, as  $\theta^l$  does not influence the behavior of the inputs it receives from previous layers. Hence, we are left with computing  $\delta_t^l = \mathrm{d}_{h_t^l}L^\top$ . The simplification  $\delta_t^l = \mathrm{d}_{h_t^l}L_t^\top$  we made in the previous section still holds for the last layer but is violated for the other layers. This is because  $L$ , taken as function of the hidden states  $h^l$  of layer  $l$ , now has an internal memory through the subsequent recurrent layers. The hidden state  $h_t^l$  at time  $t$  will thus directly affect all future losses  $L_{t'}$  for  $t' \geq t$ .
<sup>3</sup> If  $h_t$  is not a holomorphic function of  $\theta$ , one would need to keep  $2|\theta|$  instead of  $|\theta|$  states.
As a consequence, one has to resort to backpropagation-through-time to compute  $\mathrm{d}_{h_t^l}L^\top$  exactly, which breaks causality and rules out the possibility of learning online. To circumvent this issue, we approximate the error signal each layer receives by  $\delta_t^l\approx \mathrm{d}_{h_t^l}L_t^\top$  so that it can be computed instantaneously with spatial backpropagation. We emphasize that the only source of approximation of this algorithm is the one above. Given that there is no approximation for the last layer, we will always compute the exact gradient for that layer.
We summarize our learning rule in Figure 1. It prescribes augmenting the hidden state of each recurrent layer  $l$  with the sensitivity  $e^l$ . For each input/output sample  $(x_t, y_t^{\mathrm{target}})$ , we first update the full entire hidden state  $(h_t, e_t)$  using the previous one  $(h_{t-1}, e_{t-1})$  and current input  $x_t$ . We then spatially backpropagate the error signal obtained at the last layer by comparing the prediction of the network to its desired value  $y_t^{\mathrm{target}}$ . Finally, we combine  $e_t^l$  and  $\delta_t^l$  available at each layer using Equation 7 to compute the current update. So far, we have only described how to update parameters that directly influence the hidden neurons of the recurrent layer. We update the rest of the parameters with the gradient estimate obtained with spatial backpropagation. Importantly, the size of the extended hidden state is upper bounded by the number of neurons plus the number of parameters and is in practice much smaller.