SlowGuess/NeurIPS_2023_OCR huggingface dataset:
- Step count. CD reverts to the base model when  $\theta = 1$ , while SCTT and DASC revert to the base model when  $\beta = 1$ . Training starts at some initial value  $\theta_0 > 1$  (CD) or  $\beta_0 < 1$  (SCTT or DASC). We progressively update the value of  $\theta$  or  $\beta$  up to a value of 1 (Figure 2, right). The number of steps over which we carry out this progressive update through the course of training is defined as the step count. We consider three values of step counts:  $\{1,5,10\}$ . This annealing process is performed over the first 20,000 training steps in equal intervals.
- **Skip length.** As described in the previous section, all methods have some way of shortening the number of backpropagation steps during training. For SCTT and DASC, this is facilitated by the skip connections. The skip length is therefore defined as the temporal length of the skip connection, which takes three values  $\theta \in \{10, 20, 40\}$ , representing half, one and twice the time constant respectively. For CD, the skip length is defined as initial discretization factor (i.e. prior to progressive updates),  $\theta_0$ , which can only take the value of 10 but not 20 and 40 due to network instability problems, as explained in Section 3.4 and Appendix E.
- Initial ratio. For SCTT and DASC, we also vary the initial weighting ratio  $\beta_0 \in \{0.2, 0.5, 0.8\}$ . There is no equivalent term for CD.
As an example, suppose we train using DASC with step size 5 and initial ratio 0.5. Then,  $\beta$  would start at 0.5 and then progressively take the values of 0.5, 0.6, 0.7, 0.8, 0.9 and finally 1.0 in fixed intervals over the course of training. In summary, we trained 4,800 networks using CD and 43,200 networks each using SCTT and DASC, totalling over 90,000 trained networks (Table S3). For each training session, we recorded the number of training steps (either due to attaining 0.99 performance or hitting the training step limit), the wall-clock time spent in training, and the number of floating point operations involved during training. We highlight the advantage of using skip connections over timescales longer than the time constant by comparing six training configurations: the base model (control), CD with discretization factor 10 ( $\mathrm{CD}_{10}$ ), SCTT with skip lengths 10 and 40 ( $\mathrm{SCTT}_{10}$ ,  $\mathrm{SCTT}_{40}$ ) and DASC with skip lengths 10 and 40 ( $\mathrm{DASC}_{10}$ ,  $\mathrm{DASC}_{40}$ ) on tasks that require long-term dependencies (9 out of 16 tasks) in Table 1. Detailed results can be found in Appendix F.
In the following analyses, we compare the training performance of methods we introduce here, namely SCTT and DASC, to previously known methods, namely CD and direct training (although no explicit literature was found for CD with the best of our efforts). Skip connections alleviate gradient stability problems during training which improves training efficiency by facilitating a shorter path for backpropagating through time. The most suitable metric to evaluate this objective is the number of training steps needed to learn the tasks (Table 1, top). RNNs trained with SCTT or DASC required fewer training steps than control or CD networks on 7 of the 9 tasks that require long-term dependencies. Moreover, control and CD networks were unable to learn 2 of these 9 tasks. This presents a clear advantage of our proposed methods. However, in general, when evaluating training algorithms, other metrics such as wall-clock training time and total floating point operations are more practical and may be more emphasized (Table 1, bottom). In most tasks, while CD requires more training steps, the wall-clock time and floating point operations required for each step is significantly less (determined by the discretization factor). However, CD networks are not always successful at learning long-term dependency tasks. Also, SCTT and DASC achieved shorter training times and required fewer floating point operations than contol networks on 7 of the 9 tasks. Taken together, we believe that these considerations tilt the advantage in favor of SCTT and DASC. Note that the results in Table 1 are computed using optimal hyperparameter configurations. Hyperparameter-specific analyses can be found in Figure S3 and Appendix G.
4.1 Strategic placement of skip connections
Another feature of skip connections is that they do not need to be regular or periodic, as implemented in the networks explored thus far. Skip connections can be placed anywhere over the task duration, where they may enhance training performance or inadvertently disrupt the true network dynamics. However, poorly placed skip connections can produce dynamical instabilities during training. Conversely, we show that by introducing additional yet strategically-placed skip connections, we can exploit its properties to improve training efficiency (Figure 4D). To best illustrate this, we train biologically plausible RNNs on a rule-based visuomotor association task [62, 63]. This task is characterized by the need to maintain long-term dependencies across trials so that it may maintain the current rule in memory across trials and detect implicit changes in the task rule to update this memory. Therefore, in contrast to the tasks studied thus far where networks are trained on single task trials, for this task networks must be trained on sequences of consecutive trials. Consequently, these tasks pose a higher level of difficulty for training biologically plausible RNNs due to the sheer number of time steps involved.
The network architecture is different compared to the previous models (Figure 4A) in that there is no separation of input and output populations, the input structure is changed, and the output of the
Figure 4: Strategic placement of skip connections in multi-trial tasks. A. Overview of RNN model. B. Two hidden task rules. C. Example of a trial where uncued rule switching occurs. D. The task is trained by implementing one of our proposed methods (CD, SCTT, DASC, marked with colored dots) with additional strategically-placed skip connections. In tasks that encompass multiple trials, adding skip connections between the start of every trial provides minimal disruption to network dynamics.
network now passes through a softmax in order to read out a decision:
$$
\mathbf {y} = \operatorname {s o f t m a x} \left(\mathbf {W} _ {\text {o u t}} \mathbf {r}\right) \tag {10}
$$
The network receives a 2-dimensional one-hot vector representing one of 2 stimuli as inputs. The target response is dictated by one of two hidden rules. Under Rule 1, the network must associate stimulus 1 with response 1 and stimulus 2 with response 2, and under rule 2 these associations are reversed (Figure 4B). The network may infer the current rule based on the feedback it receives in the form another 2-dimensional one-hot vector that indicates positive or negative feedback following the network's response. Each rule persists across multiple trials, and then switches to the other rule in an uncued manner. Similar to the monkeys in [62], our trained models detect and adapt their behavior within a single trial after the rule switch (Figure 4C). We trained RNNs on 5 consecutive trials, which necessitates learning dependencies that are  $5\mathrm{x}$  longer than the previous tasks. In light of this challenge, we define a new metric for performance: the number of networks out of 100 that successfully learn to perform the task at more than 0.99 accuracy. By introducing strategic skip connections between the start of consecutive trials, the proportion of trained networks increased across all methods (Table 2), thereby demonstrating the effectiveness and flexibility of skip connections.
Table 2: Proportion of networks successfully trained to perform a rule reversal task. 100 models were trained independently with their respective best hyperparameter configuration for each of the proposed methods and the proportion of successfully trained models are presented as baseline results. This is then repeated with additional strategically-placed skip connections (Figure 4).
Proportion of trained networks
Control	CD10	SCTT40	DASC40
Baseline	0.67 (±0.02)	0.19 (±0.02)	0.71 (±0.02)	0.71 (±0.02)
+ Strategic SC	0.76 (±0.02)	0.25 (±0.02)	0.81 (±0.02)	0.78 (±0.02)
It can be challenging to train biologically plausible RNNs on cognitive tasks requiring long-term dependencies without the aid of artificial gating elements. In a "vanilla" RNN, an intuitive way to alleviate this difficulty is to coarsen the time discretization (CD) in order to reduce the number of time steps of the overall simulation, thereby shortening the path for backpropagation. However, we show that over-discretization can produce poor approximations of the true continuous-time dynamics (red curves in Figure 3B,C), thus hindering model reliability. While our implementation of CD overcomes this by progressively refining the discretization, our results suggest that it does not offer an improvement in training efficiency (Table 1), even if some time savings is achieved. We also find that its benefits are constrained by the single unit time constant, beyond which the network quickly destabilizes during training (Figure 3B). Finally even with CD, networks may fail to learn challenging long-term dependency tasks (Tables 1, 2). In contrast, SCTT and DASC offer significant improvements in training efficiency and reliability. However, the two approaches accomplish this in different ways. While it provides a shorter path to backpropagate gradients, SCTT inadvertently attenuates the network's dynamics (Figure 3A). This significantly misaligns the network's dynamics in the presence versus absence of the skip connections (Figure 3B,C), which compels more learning as the network must change as the skip connections are annealed away. Instead, DASC emphasizes an alignment of the network's dynamics with (modified) and without (true) skip connections (Figure 3) while still supporting shorter paths for backpropagation.
Experimental advances in neuroscience are facilitating the study of tasks involving increasingly complex and lengthy temporal dependencies. Concurrently, technological advances are enabling simultaneous recordings from increasingly large neural populations. This underscores the need for complementary theory through high-fidelity biologically plausible models that perform challenging computations over lengthy timescales. By introducing methods that enable the development of such models without sacrificing biological plausibility, our work will serve to broaden the applicability of deep learning methods in systems neuroscience. This includes the development of model for evidence accumulation [64], spatial navigation [65], scene processing [66], rule-guided behavior [67, 68] complex rule learning [69], few-shot learning and adaptive behavior [70]. An important reason for using large time steps when training biologically plausible RNNs is that it decreases the computational resources or wall-clock time required for training. While true for simpler tasks, we have shown that there are harder tasks that cannot be learned with conventional methods regardless of time step size.
It is noteworthy that we focus on RNNs built to be directly comparable with neural circuits in the brain, rather than biological plausibility of a training algorithm itself; a neural circuit in the brain is sculpted by not only learning but also development and evolution. In general, training biologically plausible RNNs is a resource-intensive undertaking, and while our methods improve training efficiency and model fidelity, they do not alleviate the resource load per-training step. Additionally, at the present time approaching human-level computational competency requires architectures such as transformers unknown in neuroscience. More work will be necessary to develop biologically plausible yet computationally powerful mechanisms for learning long-term dependencies.
We have proposed and evaluated a set of methods to learn long-term dependencies in biologically plausible leaky RNNs. These methods improve model fidelity while offering substantial learning performance improvements over conventional training approaches, and even facilitate learning of tasks that are difficult to learn with conventional methods.
W.S. developed the algorithm. W.S. and V.G. performed the experiments. All authors designed the study, took part in discussions, interpreted the results, and wrote the paper.
This work was partly supported by the NIH grant R01MH062349 and ONR grant N00014-23-1-2040.
[60] Kalidindi, H. T. et al. Rotational dynamics in motor cortex are consistent with a feedback controller. eLife 10 (2021).
[61] Khona, M., Chandra, S., Ma, J. J. & Fiete, I. Winning the lottery with neurobiology: faster learning on many cognitive tasks with fixed sparse rnns. arXiv 10.48550, 2207.03523 (2022).
[62] Asaad, W. F., Rainer, G. & Miller, E. K. Neural activity in the primate prefrontal cortex during associative learning. Neuron 21, 1399-1407 (1998).
[63] Fusi, S., Asaad, W. F., Miller, E. K. & Wang, X.-J. A neural circuit model of flexible sensorimotor mapping: learning and forgetting on multiple timescales. Neuron 54, 319-333 (2007).
[64] Morcos, A. S. & Harvey, C. D. History-dependent variability in population dynamics during evidence accumulation in cortex. Nature neuroscience 19, 1672-1681 (2016).
[65] Terry Jr, A. V. Spatial navigation (water maze) tasks. Methods of behavior analysis in neuroscience 2, 153-166 (2009).
[66] Epstein, R. A. & Baker, C. I. Scene perception in the human brain. Annual review of vision science 5, 373-397 (2019).
[67] Bernardi, S. et al. The geometry of abstraction in the hippocampus and prefrontal cortex. Cell 183, 954–967 (2020).
[68] Kaefer, K., Nardin, M., Blahna, K. & Csicsvari, J. Replay of behavioral sequences in the medial prefrontal cortex during rule switching. Neuron 106, 154–165 (2020).
[69] Grant, D. A. & Berg, E. A behavioral analysis of degree of reinforcement and ease of shifting to new responses in a weigl-type card-sorting problem. Journal of experimental psychology 38, 404 (1948).
[70] Wang, J. X. et al. Prefrontal cortex as a meta-reinforcement learning system. Nature neuroscience 21, 860–868 (2018).