SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Wayne W.M. Soo
wmws2@cam.ac.uk
Vishwa Goudar
Center for Neural Science
vishwa.goudar@nyu.edu
Xiao-Jing Wang
xjwang@nyu.edu
Training recurrent neural networks (RNNs) has become a go-to approach for generating and evaluating mechanistic neural hypotheses for cognition. The ease and efficiency of training RNNs with backpropagation through time and the availability of robustly supported deep learning libraries has made RNN modeling more approachable and accessible to neuroscience. Yet, a major technical hindrance remains. Cognitive processes such as working memory and decision making involve neural population dynamics over a long period of time within a behavioral trial and across trials. It is difficult to train RNNs to accomplish tasks where neural representations and dynamics have long temporal dependencies without gating mechanisms such as LSTMs or GRUs which currently lack experimental support and prohibit direct comparison between RNNs and biological neural circuits. We tackled this problem based on the idea of specialized skip-connections through time to support the emergence of task-relevant dynamics, and subsequently reinstitute biological plausibility by reverting to the original architecture. We show that this approach enables RNNs to successfully learn cognitive tasks that prove impractical if not impossible to learn using conventional methods. Over numerous tasks considered here, we achieve less training steps and shorter wall-clock times, particularly in tasks that require learning long-term dependencies via temporal integration over long timescales or maintaining a memory of past events in hidden-states. Our methods expand the range of experimental tasks that biologically plausible RNN models can learn, thereby supporting the development of theory for the emergent neural mechanisms of computations involving long-term dependencies.
The adoption of deep learning methods in neuroscience research has fostered entirely new approaches to modeling neural systems and developing scientific hypotheses [1-4]. Moreover, rapid advances in network initialization and optimization techniques has reduced the threshold knowledge for designing and training neural networks and improved the accessibility of model development in neuroscience. Consequently, it is now far more feasible to formulate theories and even take a theory-first approach to neuroscience. Technological advances in experimental neuroscience are also making it easier to study increasingly complex neural computations, escalating the demand for models that support much longer time-scale dependencies than are conventionally modeled with biologically plausible RNNs (see Appendix A). However, fundamental physiological and anatomical properties of the brain constrain the neural network architectures that are permissible for study in neuroscience. First, the brain's connectivity is highly recurrent at multiple spatial scales, making recurrent neural networks (RNNs) [5] the modeling tool of choice in neuroscience [2, 4, 6]. The training and analysis of RNNs
has advanced theory on the neural basis of perception [7-9], motor behavior [10-13], cognition [14-27], memory and navigation [28, 29]. Second, physiological properties of neurons places a greater demand on the emergent computational properties of neural populations. Specifically, in the absence of inputs, a neuron's response rapidly decays away due to its passive membrane properties. Consequently, "leaky" RNN models rely on the reverberation of activity throughout a population of units to integrate inputs over time and to maintain this information in the absence of ongoing inputs. This has even contributed to the population doctrine gaining traction in neuroscience - that neural populations rather than single neurons form fundamental computation units [30, 31].
1.1 Gradient instability in recurrent neural networks
The stability challenges of training biologically-plausible (leaky) RNNs with backpropagation through time are widely understood - the larger the number of time steps the network must compute over, the more likely it is that the gradients become catastrophically large or small which arrests further learning. Long-short-term memory (LSTM) [32] and Gated Recurrent Units (GRU) [33] based RNNs were developed precisely to mitigate this adversity posed by exploding or vanishing gradients, doing so by storing memory in self-sustaining cells rather than leaky population activity. Unfortunately, there is a dearth of evidence in neuroscience in support of the gating mechanisms that help LSTMs and GRUs overcome gradient stability issues. There exist multiple biological mechanisms for gating in the brain [34], but none of them is known to turn on and off the neuronal leak, input and output flows in such unconstrained ways as presumed by LSTM or GRU. Moreover, memory function through emergent population dynamics (rather than via a population of self-sustaining units) are a central hypothesis in neuroscience. An extended review of other aspects of biological implausibility of gating mechanisms can be found in Appendix B, which highlights issues with gating properties, ease of training, network dynamics as well as mechanisms underlying memory storage. Other methods that help RNNs learn long-term dependencies have been proposed in the past that do not involve gating mechanisms. A straightforward way would be to simply initialize the recurrent weight matrix as an identity matrix [35]. When applied to ReLU RNNs, these networks were competitive with LSTMs on speech and language benchmarks. Similarly, RNNs whose recurrent weight matrices are constrained to have absolute eigenvalues of 1 have been proposed [36, 37], at the cost of lower convergence rates [38]. At their cores, both methods alleviate the problem by reducing the base multiplier in which gradients explode or vanish. Auxiliary losses targeting intermediate time steps [39] have also been proposed.
1.2 Skip connections through time
It is difficult to justify any architectural modifications to biologically plausible RNNs, especially when it involves artificial elements. The alternative, which we propose, is to introduce surrogate elements which improve convergence rates during training and subsequently remove them after training to reinstitute biological plausibility at test time. The ease of adding (and eventually removing) skip connections makes it an ideal candidate for this objective. We present and compare a set of engineering solutions using specialized skip connections for training biologically plausible RNNs on an array of commonly studied cognitive tasks [16]. The idea of using skip connections to improve gradient-based model training gained mainstream popularity when its implementation in deep convolutional neural networks (CNNs) achieved state-of-the-art performance in image classification [40]. Since then, skip connections have become a staple feature of newer architectures in machine learning, such as transformers [41] and MLP-mixers [42]. Skip connections of this nature are also referred to as residual connections. In this context, skip connections solve the degradation problem, where increasing the depth of the model reduces the maximum attainable training performance by direct backpropagation [40]. A conceptually similar idea of skip connections through time involves adding a direct connection between two non-consecutive states in sequential models, thereby facilitating a shortened path during backpropagation and alleviating gradient stability problems. One of the earliest formulations of this idea was proposed more than two decades ago in the training of NARX networks [43]. The architecture of NARX networks naturally contains the mathematical equivalent of skip connections through time, which were referred to as "jump-ahead connections". Network architectures and investigations based on skip connections through time continue to surface even in recent years [44, 45]. For example, dilated skip connections reduce the computational complexity of the RNN, in contrast to regular skip connections which requires additional computations [45]. LSTMs and GRUs are not mutually exclusive with skip connections [46, 47].
B
Figure 1: Limitations of training biologically plausible RNNs for tasks in neuroscience. A. Overview of RNN model. B. RNN performance on three cognitive tasks against the number of training steps. RNNs require more training steps in order to successfully perform tasks that require the learning of long-term dependencies. C. Number of training steps required to train biologically plausible RNNs to attain 0.99 performance on 16 standard tasks [16] in neuroscience (Appendix D).
2.1 Biologically plausible recurrent neural networks
The continuous-time dynamics of biologically plausible leaky ("vanilla") RNNs are described by:
$$
\mathbf {T} \frac {d \mathbf {r}}{d t} = - \mathbf {r} + f \left(\mathbf {W} _ {\text {r e c}} \mathbf {r} + \mathbf {b} + \mathbf {h} _ {\text {e x t}} + \boldsymbol {\eta}\right) \tag {1}
$$
where  $\mathbf{T}$  is a diagonal matrix whose elements are the time constants of each individual neuron (typically constant across all neurons), and  $\mathbf{r}$  is the firing rate vector (the hidden-state of the RNN),  $\mathbf{W}_{\mathrm{rec}}$  is the recurrent synaptic weight matrix,  $\mathbf{b}$  is the bias,  $\mathbf{h}_{\mathrm{ext}}$  is the external input vector and  $\pmb{\eta}$  represents additive noise present in the system.  $f$  is the activation function (or input-output function), which we choose to be ReLU for our simulations. Alternative functions include the hyperbolic tangent (tanh) and softmax. The output of the network is a straightforward readout layer given by
$$
\mathbf {y} _ {t} = g \left(\mathbf {W} _ {\text {o u t}} \mathbf {r} _ {t}\right) \tag {2}
$$
for some activation function  $g$ , which we choose to be the sigmoid function. Other common choices include softmax, softplus or linear, depending on the decoding objective. These models reflect the diversity in function across the brain areas they model, such perceptual functions in early sensory areas such as the visual cortex [9], maintenance of items in working memory, decision-making, context-dependent or rule-dependent responses in higher-order cortical areas such as the prefrontal cortex [14, 16, 19], and motor control of eye saccades or arm movements in motor cortical areas [10-13]. Consequently, the task structure, requisite computations, RNN inputs and outputs vary widely according to the function. We review the values of some of these parameters from previously-published models [9, 11-25, 29, 48-60] in Appendix C and Table S1 in the supplementary.
2.2 Difficulty in learning tasks with long-term dependencies
The suite of cognitive tasks used in our benchmark tests have been adapted from [16], which includes straightforward tasks involving responding after stimulus presentation (go, anti, decision making, decision making with distractors, multi-sensory decision making), reaction time-based tasks (react go/anti), as well as more complex tasks that involve storing information in memory (delayed match/non-match to sample/category, delayed go/anti, delayed decision making). Exact
task descriptions and setup are explained in Appendix D. It has been previously shown that RNNs can learn these tasks without inter-neural interactions, where recurrent connections are restricted to only self-coupling of each neuron [61]. To prevent this class of semi-degenerate and biologically unrealistic solutions, we designated distinct populations of neurons to interface with the input and output units respectively, with the rest of the unassigned neurons designated as auxiliary neurons to aid computations (Figure 1A). This ensures that, at a minimum, there must be connections between the input and output populations, leading to the emergence of meaningful solutions found by the trained networks. For all tasks, the primary inputs and outputs have a ring structure of 32 units, with each unit  $i$  having a preferred direction  $\psi_{i}$ . An input stimulus is characterized by a modality, a magnitude  $\gamma$  and a direction  $\psi$  such that:
$$
h _ {i} = \gamma \cdot 0. 8 \exp \left[ - \frac {1}{2} \left(\frac {8 | \psi - \psi_ {i} |}{\pi}\right) ^ {2} \right] \tag {3}
$$
The input consists of two separate rings, representing two different stimulus modalities. The output is a single ring characterized by a direction  $\psi$ , representing a saccadic response,
$$
y _ {i} = 0. 8 \exp \left[ - \frac {1}{2} \left(\frac {8 | \psi - \psi_ {i} |}{\pi}\right) ^ {2} \right] + 0. 0 5 \tag {4}
$$
When cued to "fixate" by a binary fixation input unit, the RNN must respond by maintaining the value of its fixation output unit, constrained to lie between 0 and 1, above 0.5. The desired input-output mapping at other times in a trial are task dependent. The loss function to be minimized is the mean-squared error between the network output and the target output. The network is optimized with respect to the input weights  $\mathbf{W}_{\mathrm{inp}}$ , output weights  $\mathbf{W}_{\mathrm{out}}$ , recurrent weights  $\mathbf{W}_{\mathrm{rec}}$  and network bias  $\mathbf{b}$ . For each task, we train 100 models in parallel with time steps ranging from 1500 to 2500 depending on the task, time constant of  $100\mathrm{ms}$  and  $dt = 5\mathrm{ms}$ , which together represents a higher temporal resolution than typical models (see Table S2 for detailed model parameters). We train each set of 100 models until the average performance of the networks in the  $10^{th} - 90^{th}$  percentile reaches 0.99 or until 40,000 training steps have been reached (Figure 1B-C). The number of training steps required for the tasks span more than an order of magnitude (Table 1), thus depicting their varying difficulties. Importantly, we find that tasks that involve some form of long-term dependency typically require more training steps and are harder to train in general. In fact, RNNs trained on two tasks that require long-term dependencies were not able to attain the required threshold performance within 40,000 training steps, but instead maintained a low but constant performance at the later stages of training (Figure 1B).
The most common way of training biologically plausible RNNs involves simulating its dynamics using Euler's method for the forward pass (Figure 2A):
$$
\mathbf {r} _ {t + \Delta t} \leftarrow \mathbf {r} _ {t} + \frac {d \mathbf {r} _ {t}}{d t} \Delta t \tag {5}
$$
Here,  $\Delta t$  is the discretization time step that plays a central role in performing backpropagation through time on the simulated dynamics. Consequently, iterating this over  $\theta$  time steps results in:
$$
\mathbf {r} _ {t + \theta \Delta t} ^ {\text {b a s e}} \leftarrow \mathbf {r} _ {t} + \sum_ {k = 0} ^ {\theta - 1} \frac {d \mathbf {r} _ {t + k \Delta t}}{d t} \Delta t \tag {6}
$$
which we will denote as the dynamics of the base model  $\mathbf{r}_{t + \theta \Delta t}^{\mathrm{base}}$ . We explore several ways in which this process can be modified, only during training, in order to help with the learning of long-term dependencies.
3.1 Coarsened discretization (CD)
CD begins by training the network with a large step size to support stable gradients while learning long-term dependencies (Figure 2B). Here, (5) is modified to:
$$
\mathbf {r} _ {t + \theta \Delta t} ^ {\mathrm {C D}} \leftarrow \mathbf {r} _ {t} + \frac {d \mathbf {r} _ {t}}{d t} \Delta t \times \theta \tag {7}
$$
for some discretization factor  $\theta > 1$ . Thereafter,  $\theta$  is gradually reduced over the course of training until  $\theta = 1$  at the end of training, resulting in a reversion to the base model described by (5).
D
Figure 2: Proposed algorithms for training RNNs to perform tasks requiring long-term dependencies. A. Base model: simulating the RNN using Euler's method and training by backpropagation through time. B. Coarsened discretization (CD): training in the same way as the base model, except with a larger discretized time step (by a factor of  $\theta$ ) and then gradually reducing  $\theta$  as training progresses until  $\theta = 1$ . C. Skip connections through time (SCTT): implementing a skip connection every  $\theta$  time steps which takes a weighted average of the hidden state on the two ends of the skip connection with weights  $\beta$  and  $1 - \beta$ , and then gradually increasing  $\beta$  until  $\beta = 1$ . D. Dynamics-aligned skip connections (DASC): SCTTs but with a CD step before taking the weighted average.
3.2 Skip connections through time (SCTT)
SCTT similarly confines the well-established idea of skip connections through time to training only (Figure 2C). For some mixing ratio  $\beta$ , a skip connection between time steps  $t$  and  $t + \theta$  alters the hidden state as:
$$
\mathbf {r} _ {t + \theta \Delta t} ^ {\mathrm {S C T T}} \leftarrow (1 - \beta) \mathbf {r} _ {t} + \beta \mathbf {r} _ {t + \theta \Delta t} ^ {\mathrm {b a s e}} \tag {8}
$$
The skip length  $\theta$  remains unchanged throughout training, while the mixing ratio  $0 < \beta \leq 1$  starts at a small value. This provides shortcuts for gradient backpropagation, thus reducing the risk of gradient instability. Similar to CD, the mixing is gradually reduced (i.e.  $\beta$  is gradually increased to 1) until we are left with a trained network without skip connections (5).
3.3 Dynamics-aligned skip connections (DASC)
We propose an additional method, DASC, which combines the essence of CD with SCTT, in order to respect dynamical properties of the system (Figure 2D). DASC achieves this by simulating large time steps through its skip connections. The skip connections between  $t$  and  $t + \theta$  alters the hidden state as:
$$
\mathbf {r} _ {t + \theta \Delta t} ^ {\mathrm {D A S C}} \leftarrow (1 - \beta) \mathbf {r} _ {t + \theta \Delta t} ^ {\mathrm {C D}} + \beta \mathbf {r} _ {t + \theta \Delta t} ^ {\mathrm {b a s e}} \tag {9}
$$
DASC aims to combine the benefits of CD and SCTT while limiting the risk of gradient instability. Unlike CD, DASC simultaneously trains the model with large and appropriately small time steps; unlike SCTT, it mitigates misalignment in the network's dynamics when it mixes hidden state estimates through the paths with and without skip connections.
Figure 3: Effects of proposed algorithms on network dynamics. A. Diagram of proposed dynamics in activity subspace over time. Coarsened discretization (CD, red) typically results in a large discretization error. Skip connections through time (SCTT, green) result in an attenuation in dynamics. Dynamics-aligned skip connections (DASC, blue) make its hidden-states more aligned with the true dynamics compared to CD. B. Normalized mean-squared error between true and modified dynamics for  $\beta = 0.5$ ,  $\theta \in \{5,10,20,30,40\}$ . Shaded area ( $\theta > 20$ ) implies a time discretization greater than the time constant. C. Difference in task performance computed between modified and true dynamics.
3.4 Network dynamics
We define the neural dynamics of an RNN in the absence of skip connections or time-step over-discretization as its true dynamics; RNNs simulated with any of the methods under study (CD, SCTT or DASC) will thus result in modified dynamics. We provide an intuition of the resultant modified dynamics underlying each method with reference to the true dynamics (Figure 3A). Under reasonable circumstances, CD typically provides a relatively poor estimation of the true network activity as a result of extrapolating over a longer time interval. DASC bridges this discrepancy by taking a weighted average between the true neural activity and neural activity produced by CD (Figure 3A, blue), and is therefore more aligned to the true dynamics than CD by design. We refer to this property as being "dynamics-aligned" (hence its name). SCTT, on the other hand, takes a weighted average between the true network state at two different time points. From the perspective of the later time point, the dynamics can be said to be attenuated (Figure 3A, green).
In order to quantify the effects of modified dynamics on trained networks, we train 100 networks on the aforementioned 16 standard tasks [4] and evaluated the networks with performance in the  $10^{th} - 90^{th}$  percentile over 400 training steps using all three proposed methods. For SCTT and DASC, we keep the mixing ratio  $\beta$  constant at 0.5, while varying the skip length  $\theta \in \{5, 10, 20, 30, 40\}$ , where the time constant is 20 times the base discretization step. For CD, we vary the discretization factor, also referred to as  $\theta$ , using the same 5 values. For all three methods,  $\theta$  represents the number of skipped time steps which effectively shortens the minimum number of steps over which backpropagation through time is performed to alleviate gradient stability problems. Therefore, they are plotted on the same axis (Figure 3B-C). We find that it is difficult to train networks with CD using  $\theta$  values greater than the time constant. We provide an empirical rationale for this in Appendix E. In contrast, SCTT and DASC remained stable even with skip lengths equal to twice the time constant, which highlights a key advantage in using skip connections. To compare the accuracy of the methods, we compute the mean-squared distance in neural activity between each of their true and modified dynamics for each network and average this distance over all networks and across all tasks (Figure 3B). SCTT's discrepancy is consistently large due to the fact that it offers a zeroth-order approximation of DASC, as shown in the first terms of the RHS of (8) and (9). Conversely, DASC uses a first-order approximation to improve numerical accuracy along the skip connections. These results are mirrored by the discrepancy in output performance, averaged here in the same manner as the mean-squared distance (Figure 3C).
Table 1: Training biologically plausible RNNs on cognitive tasks requiring long-term dependencies. For each task, 100 models were simultaneously trained by backpropagation through time until the average performance of the networks in the  $10^{th} - 90^{th}$  percentile reaches 0.99 or when 40,000 training steps have been reached. Control represents the base model in Figure 2A. The other columns are methods introduced in Figure 2B-D.
Tasks that require long-term dependencies
Task	Control	CD10	SCTT10	SCTT40	DASC10	DASC40
Training steps (×104)
dmc	0.196	0.251	1.089	0.907	0.244	0.343
dnmc	0.203	0.282	1.089	0.916	0.252	0.558
dnms	1.167	1.311	1.269	1.299	0.993	1.065
dms	1.324	1.079	1.293	1.282	1.469	1.041
ddmd	1.900	1.465	1.270	1.018	1.757	1.059
ddm	2.413	2.030	1.283	1.088	2.129	1.055
msddm	3.434	2.096	1.269	1.071	2.196	1.121
danti	failed	failed	1.268	1.088	failed	1.005
dgo	failed	failed	1.106	1.268	failed	1.003
Wall-clock training time (×104s)
dmc	0.085	0.016	0.191	0.156	0.119	0.167
dnmc	0.088	0.018	0.171	0.172	0.125	0.273
dnms	0.501	0.237	0.498	0.554	0.491	0.517
dms	0.565	0.105	0.515	0.499	0.723	0.508
ddmd	0.820	0.326	0.484	0.595	0.863	0.518
ddm	1.030	0.632	0.548	0.664	1.023	0.523
msddm	1.464	0.683	0.634	0.595	1.079	0.543
danti	failed	failed	0.402	1.345	failed	0.409
dgo	failed	failed	0.325	0.261	failed	0.412
Floating point operations (×1013)
dmc	0.114	0.080	0.229	0.187	0.144	0.201
dnmc	0.119	0.090	0.207	0.207	0.148	0.327
dnms	0.682	0.416	0.599	0.685	0.585	0.623
dms	0.773	0.343	0.625	0.609	0.865	0.609
ddmd	1.110	0.465	0.582	0.706	1.034	0.620
ddm	1.409	0.645	0.667	0.804	1.254	0.617
msddm	2.005	0.666	0.768	0.727	1.293	0.656
danti	failed	failed	0.575	2.062	failed	0.588
dgo	failed	failed	0.465	0.386	failed	0.587
We train networks on the same set 16 cognitive tasks [16] as in the previous section. For each method, we train 100 separate networks across different hyperparameter values (see Appendix F) until the average performance of the networks in the  $10^{th} - 90^{th}$  percentile reaches 0.99, or until a maximum of 40,000 training steps. We then verify that the performance of each network that failed to attain the required performance threshold has converged and remains constant, albeit at a lower level of performance (for example, Figure 1B, blue). The hyperparameters are: