SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Here,  $e[t]$  was an  $N_{\mathrm{rec}}$  dimensional vector, rather than a scalar as was the case for the original work [35]. We initialized  $P[t] = \lambda^{-1}I$ , where  $\lambda$  is the regularization parameter.
Short-comings of FORCE for dRNN training
Despite its wide use in the literature for dRNN training [24, 25, 27-29], there are several shortcomings of this way of training RNNs using FORCE that render it incompatible with the realities of experiments:
1. The training procedure requires a particular range of  $\lambda$  to converge (See Fig. S4). Even when FORCE does converge, this is for a very limited range of  $\lambda$  and epochs of training.
2. FORCE is initially designed for theoretical investigations and requires the training dataset to be continuous. Even subtle jumps in the data results in stability issues with FORCE training, as the premise of FORCE approach is to have small error from the start to the end [35].
Point 1 necessitates low-level hyperparameter tuning. The only work around to point 2 is to reinitialize the inverse covariance matrix  $P$  whenever there is a jump in data points. This practically means that if we want to perform multiple training steps on the same data, each epoch simply corresponds to picking a better initialization for the next epoch, whereas all the learned correlation structure, which is stored in the inverse covariance matrix  $P$ , needs to be forgotten. The latter point showcases that FORCE is fundamentally ill-suited to many aspects of experimental data, for example, when multiple imaging sessions need to be combined.
S2.3 Newton's method and gradient descent with single-step prediction error
We implemented Newton's method by using the exact Hessian given in Eq. (S5) for two different scenarios: i) when  $c_{t} = 1$  corresponding to minimizing cross entropy loss, and ii) when  $c_{t} = [1 - d_{t}^{2}]^{-1}$  corresponding to the weighted loss. When implementing ii), we performed the outlier detection not only on the gradient, but also the Hessian level; since the division by  $[1 - d_t^2 ]$  could have destabilized the training. Given that each subproblem had a unique Hessian for the Newton descent, we solved the separate subproblems in parallel using multiprocessing tool in numpy.
To implement the gradient descent, we used a single-layer Pytorch model with a linear  $^+$  tanh structure. We implemented both a CPU and GPU version, and used the ADAM and SGD optimizers. Gradient descent was implemented for two types of loss functions: i) loss function, and ii) cross-entropy with  $c_{t} = 1$
S3 Benchmarking details for generator RNNs
S3.1 Randomly initialized chaotic RNNs
When benchmarking dRNNs on randomly initialized chaotic RNNs, we picked the recurrent weights of the generator randomly from a normal distribution  $W_{ij}^{\mathrm{rec}} \sim \mathcal{N}\left(0,\frac{g^2}{n_{\mathrm{rec}}}\right)$  with zero mean and standard deviation of  $\frac{g}{\sqrt{n_{\mathrm{rec}}}}$ . For  $g > 1$  and sufficiently large  $n_{\mathrm{rec}}$ , these networks can sustain their activities without any external stimulus and exhibit chaotic behavior [60]. For these chaotic generators, we picked  $g = 3$  throughout this work to ensure the networks were deep into the sustainable spontaneous activity regime. We chose the chaotic regime specifically, since previous work established that teacher-forcing when training an RNN to output a function can have potentially harmful consequences to the stability of the network training [35]. We aimed to test whether this generalized to the reproduction task we considered in our work.
We note that in figures with randomly initialized chaotic RNNS, where we compared CORNN to other optimizers, we omitted the zero self-excitation constraint, i.e.,  $W_{ii}^{\mathrm{rec}} = 0$ , to remain conservative in our speed comparisons. This is because 2nd order methods (FORCE and Newton descent) cannot incorporate the zero self-excitation constraint without additional optimization steps since projected 2nd order steps are not necessarily descent directions, whereas projected gradient descent is notoriously slower compared to the unconstrained version. In contrast, CORNN was approximately equally fast with or without the ADMM solver; since both ADMM and HAPE updates converged in more or less the same number of iterations.
S3.2 Training of RNNs for the 3-bit flip flop task
To benchmark CORNN on diverse conditions, on top of the randomly connected RNNs discussed above, we established ground truth data by training RNNs on a 3-bit flip flop task. The architecture of the generator RNN was as follows:
$$
r (t) = (1 - \alpha) r (t - 1) + \alpha \tanh  \left[ W _ {\text {i n}} u (t) + W _ {\text {r e c}} r (t - 1) \right], \tag {S27a}
$$
$$
o (t) = W _ {\text {o u t}} r (t), \tag {S27b}
$$
where  $r(t)$  is the hidden state of the RNN at time  $t$ ,  $\alpha$  is the time-constant associated with the dynamics,  $W_{\mathrm{in}}$  is the weight matrix governing input into the RNN,  $u$  is the input into the RNN,  $W_{\mathrm{rec}}$  are the recurrent weights,  $o$  is the output of the RNN and  $W_{\mathrm{out}}$  are the weights which take the hidden state to the output. Since CORNN does not learn  $W_{\mathrm{out}}$ , to compute the outputs for the CORNN learned networks in Fig. S8, we used the same projection matrix  $W_{\mathrm{out}}$  as the original network. See the accompanying code for further details on network training and benchmarking experiments.
S3.3 Training of RNNs for the timed-response task
To test robustness of CORNN to mismatches between generator and inference network dynamics, we established ground truth data by training leaky current RNNs on a timed-response task with the following architecture:
$$
x (t) = (1 - \alpha) x (t - 1) + \alpha W _ {\text {i n}} u (t) + \alpha W _ {\text {r e c}} r (t - 1) + \alpha \epsilon_ {\text {R N N}}, \tag {S28a}
$$
$$
r (t) = \tanh  (x (t)), \tag {S28b}
$$
$$
o (t) = W _ {\text {o u t}} r (t). \tag {S28c}
$$
Here, unlike the CORNN and 3-bit flip flop equations (See Eq. (S27)), the time-derivative was at the level of currents  $(x)$ , not firing rates  $(r)$ . Upon receiving an input during  $t \in [0,100]ms$ , we trained 10 networks to output a Gaussian pulse centered at  $t = 500ms$  with standard deviation  $20ms$ ; otherwise zero everywhere.
During test trials in Figs. 5, S10, and S11, we introduced a novel distractor input ten times stronger than the cue at  $t \in [400, 410]ms$ . To allow fair comparison between networks trained on different numbers of neurons, we computed the  $R^2$  of the neural activities from the first 100 neurons across three 1s trials and performed a uniform average. For the illustrative Fig. 4, to allow visualization of finer details in time-activities after distraction, we instead plotted a scenario with an earlier and weaker (half strength than the cue) distractor at  $t \in [200, 210]ms$ . To obtain the correlated noise in Fig. S11, we first sampled random noise from a normal distribution, then convolved the i.i.d. noise with a 2D Gaussian kernel across time points (s.d. of  $3ms$ ) and neurons (s.d. of 5 neurons), and re-scaled by the targeted standard deviation. See accompanying code for further details.
S3.4 The mismatch between leaky-firing and leaky-current RNNs
In this section, we briefly discuss the extent of generator mismatch between leaky-firing and leaky-current networks. To see this, we consider the dynamics of the currents  $(z)$  under the leaky-firing rate RNNs:
$$
\tau \frac {\mathrm {d} z (t)}{\mathrm {d} t} = \tau W ^ {\text {r e c}} \frac {\mathrm {d} r (t)}{\mathrm {d} t} + \tau W ^ {\text {i n}} \frac {\mathrm {d} u (t)}{\mathrm {d} t}, \tag {S29}
$$
$$
\tau \frac {\mathrm {d} z (t)}{\mathrm {d} t} = - z (t) + W ^ {\mathrm {r e c}} f (z (t)) + W ^ {\mathrm {i n}} u (t) + \tau W ^ {\mathrm {i n}} \frac {\mathrm {d} u (t)}{\mathrm {d} t}.
$$
Although the current dynamics look similar to the leaky-current RNNs, there are several key differences. Firstly, the term  $\tau W^{\mathrm{in}}\frac{\mathrm{d}u(t)}{\mathrm{d}t}$  is added to the leaky-current equation, which can act as an unobserved current when the input stimuli change. Second, the definition of the firing rates are different, with the firing-rate RNNs performing a low-pass filter in time on the linear  $+$  non-linear term,  $f(z)$ , (hence the name, leaky firing-rate RNN). Finally, under non-homogeneous time-scales for the leaky firing-rate RNN,  $\tau$  would promote to a diagonal matrix and would not necessarily commute with  $W^{\mathrm{rec}}$ , so that the Eq. (S29) no longer holds, leading to a stronger mismatch.
S4 Benchmarking CORNN's speed and scalability
To obtain the CORNN's solver in Algorithm 1, we took five distinct steps:
1. Single-step prediction error paradigm, where we teacher-forced  $\hat{r}_{t,i}$ , the RNN unit activities at time  $t$ , to the observed neural activities,  $r_{t,i}$ , to predict  $r_{t+1,i}$ , with a single-step backpropagation.
2. Convexification of the loss function, where we replaced the  $\mathcal{L}_2$  loss function with a cross-entropy loss.
3. Hessian alignment to neural activity correlations, where we weighted individual samples to align subproblem Hessians to a pre-computable quantity, i.e., the correlation matrix of observed neural activities.
4. Initialization to a fixed-point, where we used the approximate least square solution to initialize the networks.
5. Using ADMM approach to enforce constraints without increasing algorithmic runtimes
In this work, we performed several benchmarking experiments to delineate the contribution of each step we take towards the final convex solver to the scalability and speed of CORNN. We proved, with targeted experiments, the need for each of these steps, as we discuss below.
S4.1 Single-step prediction error paradigm
As discussed in Section 2.2, the first step we took was to restrict the family of interest for the inference models to the models that can perform single-step predictions well. This simplifying assumption mapped the recurrent estimation problem to a single layer feed-forward network with a sigmoid/tanh non-linearity. Here, the inputs of the network are  $u_{t,i}$  and  $r_{t,i}$ , whereas the outputs are  $r_{t+1,i}$ . Since this step eliminated the need for the back-propagation through time, both the training times and memory requirements decreased significantly (Compare Figs. S1 and S2).
However, a priori, it was not clear whether the restrictive assumption of single-step prediction error minimization would lead to decreased performance. To test this, we performed the experiments in Fig. S1 and observed that teacher forcing increases the efficiency of the training, and was not deleterious. A potential explanation may lie in the observation that, unlike the traditional scenario where a sparse error signal needs to be propagated back in time to perform credit assignment, network reproduction training had access to rich variety of error signals at each time step and for each neuron.
S4.2 Convexification of the loss function
As discussed in Section 2.3, the convexification of the loss function plays an important role in speeding up the learning process. However, by itself alone, convexification is not the reason of the speed improvement. For example, in Fig. 3, gradient descent on  $\mathcal{L}_2$  loss converges faster than the gradient descent on the logistic loss. On the other hand, for smaller networks with 200 neurons (Figs. S2, S3, S5), convex loss function merits the use of Newton descent and leads to faster convergence than gradient descent. Yet, as shown in Figs. S2 and S5, CORNN is at least an order of magnitude faster even for these small networks. Hence, though convexification of the loss function is crucial, it cannot explain the speed or the scalability of CORNN by itself.
S4.3 Hessian alignment
The main contribution of our work is to pick  $c_{t,i}$ , which approximates the problem as a local least-squares as discussed in Supplementary Section S1.3, and aligns the Hessians of each subproblem to a pre-computable quantity, i.e., the correlation matrix of neural activities, as discussed in Section S1. This alignment allows the use of HAPE updates that have complexity of gradient descent but converge in  $O(10)$  iterations thanks to the approximate Hessian. Figs. S2 and S5 show that Newton's descent with exact Hessian computations on the weighted CORNN loss is an order of magnitude slower than CORNN. Hence, the Hessian alignment was a crucial step to speed up the learning, but most importantly allow scaling to large networks due to its gradient-level complexity.
S4.4 Fixed-point initialization
The fixed-point initialization relies on an approximate solution when the final error is low (See Supplementary Sections S1.2 and S1.3). This initialization is a good first guess, since, if there was no conversion noise at all, one could have simply performed least-squares on the currents estimated via  $z = \operatorname{arctanh}(r)$ . This is akin to regressing to the inverse of the output before the non-linearity in a single layer network. Though such an approach can provide a good first guess, it is not robust to conversion noise. This was shown in Figs. 3 and S3, where non-Gaussian noise components lead to low accuracy for the initialization. Finally, the least-squares solution does not satisfy any additional
constraints on the connectivity matrix, thus is not suitable as a valid solution itself. Yet, since the least-squares solution is computationally negligible, we used it as an initialization for CORNN.
S4.5 Alternating direction method of multipliers
We chose to use ADMM approach to enforce the no self-excitation constraints and present a versatile base model that can incorporate additional design choices such as L1 and/or nuclear norm regularization. In our comparison experiments, we omitted any constraints since other solvers were considerably slower, or outright incompatible. However, in our speed benchmarking results in Figs. 5, S8, and S9, we used the full CORNN solver with the ADMM described in Algorithm 1. The choice of ADMM for enforcing constraints is rooted in our observation that both the ADMM steps and the original HAPE updates take  $O(10)$  iterations to converge. Moreover, other parts of the ADMM updates are negligible compared to the primal problem, where the HAPE updates take place. In other words, the fact that HAPE updates and ADMM converge in similar steps creates a mutual synergy that increases accuracy and promotes versatility in our work without hurting convergence times.
S5 Reproducibility
We provided the code to reproduce each figure in our Github. All experiments, except for those in Figures 3, 5, and S4, were run on a standard computer with Geforce RTX 3090 Ti GPU and Intel Core i9-9900X Skylake X 10-Core processor. Experiments in Fig. 5 were run on a computer with two Geforce RTX 3080 Ti GPUs (though using only one) and Intel Core i9-10980XE 18-Core processor. Experiments in Figure S4 were divided to run on both of the aforementioned computers. Experiments in Figure 3 were run on a computer with two NVIDIA Geforce RTX 4090 (though using only one) and an AMD Ryzen Threadripper PRO 5975WX 32-Core processor (though using only 5 cores due to memory limitations). Illustrative figures are plotted on an Apple MacBook Air with M1 Chip. See the accompanying code provided in the Github repo<sup>2</sup> for the specific hyperparameters used for the experiments.
$^{2}$ Reproduction code is available at https://github.com/schnitzer-lab/CORNN-public
Figure S1: Backpropagation through time (BPTT) on an  $\mathcal{L}_2$  loss function is suboptimal, whereas teacher-forcing can increase the reconstruction accuracy. To assess the optimality of BPTT on an  $\mathcal{L}_2$  loss function and investigate the potential drawbacks of teacher-forcing, we conducted reproduction experiments on chaotic networks with low (A) and high (B) input to conversion noise ratio. We varied the levels of teacher forcing applied to BPTT and employed two distinct loss functions: cross-entropy and  $\mathcal{L}_2$  loss function. Surprisingly, our results challenge traditional beliefs, since cross-entropy with teacher forcing emerged as the superior method compared to BPTT. This superiority may be attributed to the rich information content available during network reproduction training. Parameters: A, B:  $\alpha = 0.9$ ,  $n_{\mathrm{rec}} = 100$ ,  $T = 100$ ,  $n_{\mathrm{epoch}} = 3000$ .  $\forall t = 1, \dots, T$ ,  $u(t) = 0.1$  with  $W^{\mathrm{in}} \sim \mathcal{N}(0, g^2)$  and  $g = 3$ . A:  $\epsilon^{\mathrm{conv}} \sim \mathcal{N}(0, 10^{-8})$ ,  $\epsilon^{\mathrm{input}} \sim \mathcal{N}(0, 10^{-4})$ . B:  $\epsilon^{\mathrm{conv}} \sim \mathcal{N}(0, 10^{-10})$ ,  $\epsilon^{\mathrm{input}} \sim \mathcal{N}(0, 10^{-2})$ . Box plots are obtained from 10 random generator networks shared across optimization algorithms. Orange lines denote median values; boxes span the 25th to 75th percentiles; outliers are not shown.
Figure S2: CORNN achieves a better combination of speed and accuracy compared to alternatives, dramatically accelerating convergence. The plots illustrate the relationship between accuracy (Pearson's correlation coefficient between ground truth and inferred weights) and training time, measured in seconds on a log scale, with low (A) and high (B) input to conversion noise ratio. The blue line represents the performance of our newly developed method, which combines fixed-point initialization and CORNN, and achieves high accuracy while significantly accelerating convergence compared to other methods. Parameters: A-B:  $\alpha = 0.1$ ,  $n_{\mathrm{rec}} = 200$ ,  $T = 3000$ . No input. A:  $\epsilon^{\mathrm{conv}} \sim \mathcal{N}(0,10^{-8})$ ,  $\epsilon^{\mathrm{input}} \sim \mathcal{N}(0,10^{-4})$ . B:  $\epsilon^{\mathrm{conv}} \sim \mathcal{N}(0,10^{-10})$ ,  $\epsilon^{\mathrm{input}} \sim \mathcal{N}(0,10^{-2})$ . In all plots, data points indicate median values, and error bars denote standard error of the mean (SEM) over 100 runs of the simulation.
Figure S3: CORNN is robust to non-Gaussian noise injections. Same as in Figure S2, but with parameters:  $\alpha = 0.1$ ,  $n_{\mathrm{rec}} = 200$ ,  $T = 1500$ , No input.  $\epsilon^{\mathrm{conv}} \sim$  Poisson  $(10^{-1})$ ,  $\epsilon^{\mathrm{input}} \sim$  Poisson  $(10^{-2})$ .
Figure S4: The traditional FORCE approach is severely suboptimal, underperforms compared to fixed point initialization, and is hyper-sensitive to the regularization level. This plot describes the relationship between the reconstruction accuracy (correlation between ground truth and inferred weights) and the strength of regularization applied to the weights during training. The results show that while the FORCE algorithm, as previously applied to a similar task [24, 25, 27-29], is highly sensitive to the choice of hyperparameters, CORNN, whose accuracy is practically bounded below by the initialization, is not. Moreover, in majority of the cases, FORCE cannot train beyond the fixed-point initialization we developed in this work. A with fixed point initialization and B with random initialization. Parameters are the same as in Fig. S2A. In all plots, data points indicate mean values, and error bars denote SEM over 10 runs of the simulation.
Figure S5: CORNN can be dramatically accelerated with on a standard GPU. The same as in Fig. S2, but with GPU acceleration for CORNN and gradient descent algorithms.
Figure S6: Fixed-point initialization speeds up the convergence. A-B Same as in Figure S1, but with fixed-point initialization. C We conducted an ablation study to investigate the effectiveness of our proposed fixed-point initialization method. The results show that the CORNN algorithm with fixed-point initialization (purple) outperforms all other methods, including the vanilla CORNN without this fixed-point initialization technique. Parameters:  $\alpha = 0.1$ ,  $n_{\mathrm{rec}} = 500$ ,  $T = 10000$ ,  $\epsilon^{\mathrm{conv}} \sim \mathcal{N}(0,10^{-10})$ ,  $\epsilon^{\mathrm{input}} \sim \mathcal{N}(0,10^{-2})$ . All algorithms are ran on CPU. For A,B, the boxes cover from the lower quartile to the upper quartile, with an orange line at the median. For C, in all plots, data points indicate median values, and error bars denote SEM over 100 runs of the simulation.
CORNN
Newton descent - weighted
cross entropy
Newton descent - cross entropy
Gradient descent (ADAM) -
cross entropy
- - - - - - - - - - - - - - - - - -
CORNN (FP initialized)
Figure S7: CORNN can reproduce an RNN performing 3-bit flip flop task through dRNN training. We use CORNN to reproduce the behavior of a synthetic RNN trained on a 3-bit flip-flop task [12], where the network is given a short input pulse at a random time. If the input pulse is  $+1$ , the output of the network must transition to  $+1$  or stay at  $+1$ . If the input pulse is  $-1$ , the network's output must transition to  $-1$  or stay at  $-1$ , ignoring input pulses that do not match. Through dRNN training, CORNN is able to accurately reproduce, left, the inputs and outputs of the task-trained RNN, middle, spatiotemporal pattern of the neural population dynamics, and, right, the inputs and (partially) recurrent weights of the synthetic RNN. Parameters:  $\alpha = 0.9$ ,  $n_{\mathrm{rec}} = 1000$ ,  $T = 200$  trials (100 data points per trial),  $\epsilon^{\mathrm{conv}} \sim \mathcal{N}(0, 10^{-6})$ ,  $\epsilon^{\mathrm{input}} \sim \mathcal{N}(0, 10^{-4})$ .
Figure S8: CORNN runtimes scale linearly with increasing trial data, which can facilitate learning from larger networks. Top, reconstruction accuracy of network weights, measured as the correlation between recurrent weights, are plotted as a function of dataset size (number of trials), middle, reconstruction accuracy of network outputs, measured as the  $\mathbf{R}^2$  between the output units, and, bottom, the linear scaling of algorithmic wallclock runtimes with near-unity slope vs increasing number of trials on a log-log plot. The different colors in the plot correspond to different number of neurons in the generator network. Missing data points correspond to the limits of the GPU memory, which can be mitigated by computing prediction errors in temporal chunks in real-data applications, though not implemented for this case study. Parameters:  $\alpha = 0.5$ ,  $\epsilon^{\mathrm{conv}} \sim \mathcal{N}(0,10^{-6})$ ,  $\epsilon^{\mathrm{input}} \sim \mathcal{N}(0,10^{-4})$ . In all plots, data points indicate mean values, and error bars denote SEM over 20 runs of the simulation.
Figure S9: CORNN scales polynomially to large network sizes. Using the networks with 3000 neurons trained for the 3-bit flip flop task in Figure S7, we subsampled neural populations and collected CORNN runtimes for CPU and GPU implementations. Parameters:  $\alpha = 0.5$ ,  $\epsilon^{\mathrm{conv}} \sim \mathcal{N}(0,10^{-6})$ ,  $\epsilon^{\mathrm{input}} \sim \mathcal{N}(0,10^{-4})$ . A with 200 trials (20,000 data points) and B with a single trial (100 data points).
Figure S10: CORNN is robust to variations in the time-scales compared to the ground truth. Same as in Figure 5, but with 500 observed neurons and varying levels of uncertainties in time-scales.
Figure S11: Correlated noise presents a challenge for CORNN, which can be mitigated by increasing the data size. Same as in Figure 5, but with the fully observed network and varying levels of injected random and correlated noise (for details, see Section S3.3).