SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Figure 3: CORNN reduces training times by several orders of magnitude. The plots illustrate the relationship between reconstruction accuracy (Pearson's correlation coefficient between ground truth and inferred weights) and training time, measured in seconds on a log scale. The FORCE approach (here, on firing rates) is the default method in neuroscience literature for dRNN training [24, 25, 28, 29]. Parameters:  $\alpha = 0.1$ ,  $n_{\mathrm{rec}} = 5000$ ,  $T = 30000$ . No input.  $\epsilon^{\mathrm{conv}} \sim \mathrm{Poisson}(10^{-3})$ ,  $\epsilon^{\mathrm{input}} \sim \mathcal{N}(0,10^{-4})$ . Lines: median. Error bars: s.d. over 7 networks.
use of fast HAPE updates in the computationally extensive first subproblem, where Hessians can be aligned.
While we only considered the equality constraint in this work, the ADMM framework can be used to add additional L1 and/or nuclear norm regularizations to the learned connectivity matrix [51, 54], and utilize biological priors for further inductive biases. Finally, we note that given the highly irregular nature of the prediction error when  $d_{t,i} \approx \pm 1$ , we supplied the fast solver with a simple automated outlier detection step explained in Supplementary Section S1.5.
3.2 Backpropagation through time on the  $\mathcal{L}_2$  loss is suboptimal
In a traditional setting, in which neural activations are hidden and need to be learned through training to provide a correct output, backpropagation through time (BPTT) is necessary to learn long-term dependencies in the data. However, in the network reproduction where neural activations are no longer hidden, it is not clear whether BPTT would be beneficial for the learning, since long-term dependencies are inherently present in the state space of the dynamical system, whose equations are given in Eq. (1). Specifically, given that Eq. (1) has only a single time derivative, the time evolution is performed in a Markovian manner, i.e., given the current state, the next state can be computed without need for additional information (up to a corruption by random noise).
To test whether BPTT is a suitable training algorithm for network reproduction, we examined randomly connected chaotic recurrent neural networks as a synthetic ground-truth benchmark (See Section S3.1 for details). As shown in Fig. S1, contrary to expectation, teacher-forcing enhanced the convergence speed of the BPTT. Moreover, we also observed higher reproduction accuracy with the cross-entropy loss compared to the traditional  $\mathcal{L}_2$  loss, inline with our observation from Fig. 2 that  $\mathcal{L}_2$  loss leads to vanishing gradients in large error regime. As expected, BPTT took around a minute even in the toy example of Fig. S1 with 100 time points and 100 neurons, approximately two orders of magnitude smaller than a standard calcium imaging session in both dimensions.
3.3 CORNN decreases training times several orders of magnitude
Having shown that teacher-forcing can speed-up convergence in chaotic networks, we next validated that chaotic dynamics can be reconstructed via a single step prediction error paradigm that CORNN utilizes. Specifically, we initialized randomly connected recurrent neural networks as generators, extracted neural activities for a fixed duration of length  $T$ , and trained the dRNNs using CORNN, Newton's solver, gradient descent, and FORCE on varying number of iterations to benchmark their speed and accuracy (See Supplementary Sections S2 and S3.1 for details). All algorithms were able to train, some achieving near perfect accuracies in reproducing the internal connectivity of the generator network (See Figs. 3, S2, and S3). However, FORCE, the current default method in neuroscience literature for dRNN training [24, 25, 28, 29], was 4 orders of magnitude slower (Fig. 3) and needed fine-tuned hyperparameters to barely outperform the fixed-point initialization (Fig. S4).
Whether minimizing the weighted or logistic loss led to more accurate reconstruction depended on the underlying noise distribution (we tested two noise distributions: Poisson(mean) and  $\mathcal{N}$  (mean, variance)) of the data as shown in Figs. 3, S2, and S3. However, CORNN outperformed all other algorithms in speed in both cases with several orders of magnitude and converged
in sub-second times even with the CPU implementation, whereas others took up to several days for training (Fig. 3). We also observed increased efficiency for CORNN on a GPU (Fig. S5) and that fixed-point initialization could stabilize other algorithms (Figs. S4 and S6).
3.4 CORNN runtimes scale linearly with data size and polynomially with the network size
To this point, we provided validity evidence for the speed and accuracy of CORNN, or single-step prediction error paradigm in general, via the randomly connected chaotic RNN generators. Next, we focused on RNNs with more structured connectivity matrices. Specifically, we trained a model RNN using BPTT to perform a 3-bit flip flop task following [12] as shown in Fig. S7 (See Supplementary Section S3.2 for details). Fig. S7A shows the output of an example test trial, not used for the generator RNN training or the CORNN reproduction, with 200 time points followed by 100 time points of inter trial interval. We observed that CORNN reproduced network performed the task the same way the original network did; potentially making similar mistakes. Thus, as a first step, we confirmed that CORNN can learn to output the same outputs as the original network.
Next, we looked at underlying neural activities of an example (original) generator and reproduced (CORNN learned) network in Fig. S7B. The reproduction matched not only the output of the network, but also the internal firing dynamics of individual units. Yet, as shown in Fig. S7C, the synaptic connections were not perfectly learned even in such a simple scenario with matching models, which required more trials and potentially interventional data [15]. Thus, we next tested whether increasing the number of trials could lead to better reproduction of the original network and whether CORNN can scale well to fit the increased data and network size.
Fig. S8 shows that accurately reproducing the large networks required significantly more trials. Moreover, the data size requirement for the estimation process increased with the number of units in the generator network, which emphasizes the importance of an optimization algorithm that can scale well to multiple day recordings [55]. Fortunately, CORNN had a near-linear scaling of training times vs. the increasing number of trials and hence can handle large datasets (See Fig. S8). In the other direction, we observed that the training times scaled polynomially (Fig. S9), following  $n_{\mathrm{rec}}^{\beta}$
Generator network dynamics
D
Figure 4: CORNN can reproduce underlying neural dynamics and attractor structures despite several non-idealities in the timed-response task. A. To assess the robustness of CORNN against time-scale mismatches, assumed inference models, and unobserved influences from subsampled neural populations, we trained 10 networks using a different generator model than CORNN for the timed-response task. B. We conducted a set of novel test trials to evaluate the reproduction accuracy of CORNN when the network was disturbed by a small distractor, which was not present during the learning of the timed-response task. C. Despite observing only 500 out of 5000 neurons and with an average  $10\%$  mismatch between  $\alpha_{G}$  and  $\alpha_{I}$ , the CORNN-learned network reproduced both the neural dynamics and the output. D. When reproducing the original network without subsampling or time-scale mismatches but with generator-inference mismatches, the learned weights exhibited imperfect, but non-zero, correlation with the ground truth. E. Similar to D, we reproduced the network shown in C with subsampling and time-scale mismatches. Parameters:  $n_{G} = 5000$ ,  $\epsilon^{\mathrm{RNN}} \sim \mathcal{N}(0,10^{-2})$  for 100 training trials, one example network,  $\alpha_{G} = \frac{\Delta t}{\tau_{G}} = 0.1$ , and  $\Delta t = 1ms$ . B, D:  $n_{O} = 5000$ ,  $\alpha_{I} = 0.1$ . C, E:  $n_{O} = 500$ ,  $\alpha_{I} \sim \mathcal{N}(0.1,10^{-4})$ .
with  $\beta \in [1.3,2.2]$ , vs. the increasing number of neurons. Bringing both observations together, we concluded that there is an  $O(n_{\mathrm{rec}}^{\beta}T)$  scaling of the empirical training times.
3.5 CORNN is robust to a diverse set of non-ideal conditions
To this point, we had considered ideal scenarios in which the generator network shared the same dynamical system equations as the inference network, all neurons were observed, and the time-scales of the generator and inference RNN units matched perfectly. Next, we considered cases in which these assumptions are violated (See Fig. 4A). Specifically, we trained a generator RNN with different governing equations (See Supplementary Section S3.3) to perform a timed-response task (See Fig. 4B), which put the network in a limit-cycle like dynamical attractor state initiated by an input. We tested the existence of the underlying attractor through a novel distractor input, which was not present during the training of the task (See Fig. 4B, C). If CORNN was able to learn the underlying state-space geometry, the reproduced network activities should be able to return to their attractive trajectory under novel perturbations.
When we subsampled the generator networks by  $10\%$  and introduced jitter in the time scales, we observed that the inference model reproduced the neural dynamics in the observed population in a novel perturbation trial (See Fig. 4C), even though the sub-connectivity matrix was not well reproduced (See Fig. 4D, E). On the one hand, this observation reinforces the findings of previous literature that the connectivity matrices learned by dRNNs should be interpreted as functional connections, accounting for the behavior of networks at the level of neural dynamics, and not synaptic connections [15, 25]. On the other hand, the ability of the mismatched network to reproduce neural activities in a novel perturbation trial provided evidence that the effects of the underlying attractor structures can be robustly reproduced with CORNN despite severe experimental non-idealities. We quantified these robustness aspects in Figs. 5A,B (subsampling), S10 (time-scale mismatch), and S11 (correlated noise). Moreover, we observed the  $n_{\mathrm{rec}}^{\beta}$  scaling for the training times with increased number of network size (Fig. 5C), inline with the results of Fig. S9.
To place in perspective the experimental paradigm enabled by CORNN introduced in this work, we now return to the experimental scenario in Fig. 1. Imagine an experiment in which mice perform a predefined task several times, e.g., for an half hour, with imaging of 3000-4000 neurons. This experimental scenario yields roughly ten million parameters to be trained in the dRNN. The current study showcased CORNN's efficacy in training networks of thousands of neurons, requiring just  $O(10)$  iterations, each comparable in complexity to gradient computation and taking seconds. Consequently, training such a network from the initial imaging session would take less than a minute. Once trained, the network can enable real-time planning of experimental interventions, testing
Figure 5: CORNN runtimes scale polynomially with increasing number of neurons, whereas reconstruction accuracies remain robust to subsampling. We use the timed-response task as a testbed for quantifying the robustness of CORNN to unobserved influences due to subsampling in neural populations. A Reconstruction accuracy of input weights are plotted as a function of subsampling ratio (fraction of observed neurons), B reconstruction accuracy of neural activities, measured as the  $\mathbb{R}^2$  between the ground truth and predicted activations, and C the linear scaling of the training times with asymptotic slopes  $\in [1.2,1.3]$  vs increasing number of neurons on a log-log plot, hinting at polynomial scaling. The different colors in the plot correspond to different number of training trials. Parameters:  $\alpha = 0.1$ ,  $\epsilon^{\mathrm{RNN}}\sim \mathcal{N}(0,10^{-2})$  for the training trials. See Supplementary Section S3.3 for further details. Data points: mean, error bars: s.e.m. over 10 networks.
multiple scenarios in parallel to identify several optimal targeting strategies, a template of potential interventions, for neuron combinations driving desired animal behavior. Moreover, with CORNN's rapid convergence and typical inter-trial intervals of a few seconds in behavioral experiments, the network and the interventional strategy can be refined between trials with incoming data streams, facilitating real-time learning.
To understand the timescales for real-time inference with a dRNN for within-trial intervention experiments using calcium imaging, we can look at typical values. Acquiring each image frame takes about  $30~\mathrm{ms}$  [2]. While new frames continue to be acquired, motion correction and neural activity extraction can happen in  $\sim 5$  ms per frame [56]. In addition to these ongoing imaging steps, the simulation and decision-making with the CORNN-fitted dRNN needs to be performed. The dRNN starts with the current observed activity and can quickly simulate multiple future responses under different input scenarios, i.e., potential interventions. Since the simulation involves only iterative matrix multiplications and point-wise nonlinearities, it finishes very quickly—in just a few ms. For instance, running a 1000neuron RNN forward for 10 timesteps ( $\sim$  a second), for 100 different initial conditions, takes  $< 6$  ms on a GPU, and  $< 25$  ms on a CPU, in our hands. Once a desired intervention is identified from the simulations, the phase mask on a spatial light modulator can be updated in  $\sim 10$  ms to optically stimulate the chosen neurons [56]. Putting all the steps together, in under  $100~\mathrm{ms}$  one could capture brain activity, simulate future responses, decide on an intervention, and update the optical stimulation parameters. This is fast enough for real-time closed-loop applications. By streamlining network training, CORNN provides dRNNs for these types of experiments.
An important final point in considering the use of CORNN in experimental settings comes from the fact that neural networks are, in general, non-identifiable [57]. That is, for any given settings of the parameters, there are other settings which give the same input-output function. This means that CORNN does not aim to infer the true underlying synaptic connectivity matrix from a neural activity dataset. Instead, the main use of CORNN is to infer an RNN model which recapitulates the dynamical trajectories in a neural population. CORNN may also capture the underlying attractor structure of a system. However, we caution that any claim having to do with attractor structures must be experimentally validated with perturbation experiments that directly test for attractors. In our work, we simulated such experimental validation (Figs. 4 and 5) and found that in the setting tested, CORNN was able to predict the dynamical effects of perturbations on the neural population.
In this work, we introduced a fast and scalable convex solver (CORNN) for rapid training of data-constrained recurrent neural networks. We showed that CORNN is as accurate yet faster than existing approaches by orders of magnitude and can easily scale to the large datasets of today, training in seconds to a few minutes depending on the data, network size and the availability of a GPU. CORNN, as a base model, lays the groundwork for future developments and can be improved by integrating experimentally relevant assumptions and regularizations within the ADMM framework.
When applied to simulated data from structured networks, CORNN picked up the underlying attractor structure despite non-idealities (Figs. 4, 5, S10, and S11). Inspired by this observation, further development of data-constrained RNNs can support systems neuroscience research aiming to understand inter-area interactions and can complement or augment studies of canonical or pairwise correlations toward characterizations of functional connectivity [33]. Moreover, reverse engineering the learned network may provide a deeper understanding of how neuronal populations contribute to emergent computation, cognition, and memory [12, 30]. Finally, with the advent of interventional methods such as single-cell targeted optogenetics [10], dRNNs can be trained to test causal connections computationally and supply theoretical predictions at a scale previously unachievable for single-cell targeted experimentation [8].
This work constitutes a first step towards the application of CORNN to experimental data. However, several steps remain to apply dRNNs in real-time to interventional experiments. Some example steps may include the transformation of calcium traces or spike trains into traces of firing rates normalized within  $[-1, 1]$ , applying the CORNN solver developed in this work into first offline and then online experimental scenarios, estimation of neuronal time-scales from the experimental data instead of tuning them as hyperparameters, and perhaps implementing a low-rank regularization approach that opens the door to interpreting the observed dynamics in terms of latent variables [26, 52].
We would like to thank Liam Storan, Udith Haputhanthri, Parth Nobel, Dr. Itamar Landau, Dr. Yoshihisa Yamamoto, Dr. Surya Ganguli, Dr. Jay Mclelland, Dr. John Duchi, and Dr. Stephen Boyd for valuable feedback and insightful discussions. FD receives funding from Stanford University's Mind, Brain, Computation and Technology program, which is supported by the Stanford Wu Tsai Neuroscience Institute. MJS gratefully acknowledges funding from the Simons Collaboration on the Global Brain and the Vannevar Bush Faculty Fellowship Program of the U.S. Department of Defense. FD expresses gratitude for the valuable mentorship he received at PHI Lab during his internship at NTT Research.
[1] Christopher Langdon, Mikhail Genkin, and Tatiana A Engel. A unifying perspective on neural manifolds and circuits for cognition. Nature Reviews Neuroscience, pages 1-15, 2023.
[2] Tony Hyun Kim and Mark J Schnitzer. Fluorescence imaging of large-scale neural ensemble dynamics. Cell, 185(1):9-41, 2022.
[3] Anne E Urai, Brent Doiron, Andrew M Leifer, and Anne K Churchland. Large-scale neural recordings call for new insights to link brain and behavior. Nature neuroscience, 25(1):11-19, 2022.
[4] Nicholas A Steinmetz, Christof Koch, Kenneth D Harris, and Matteo Carandini. Challenges and opportunities for large-scale electrophysiology with neuropixels probes. Current opinion in neurobiology, 50:92-100, 2018.
[5] Nicholas James Sofroniew, Daniel Flickinger, Jonathan King, and Karel Svoboda. A large field of view two-photon mesoscope with subcellular resolution for in vivo imaging. *elife*, 5:e14472, 2016.
[6] Sadegh Ebrahimi, Jérôme Lecoq, Oleg Rumyantsev, Tugce Tasci, Yanping Zhang, Cristina Irimia, Jane Li, Surya Ganguli, and Mark J Schnitzer. Emergent reliability in sensory cortical coding and inter-area communication. Nature, 605(7911):713-721, 2022.
[7] Nicholas A Steinmetz, Cagatay Aydin, Anna Lebedeva, Michael Okun, Marius Pachitariu, Marius Bauza, Maxime Beau, Jai Bhagat, Claudia Böhm, Martijn Broux, et al. Neuropixels 2.0: A miniaturized high-density probe for stable, long-term brain recordings. Science, 372(6539): eabf4588, 2021.
[8] Matthew G Perich and Kanaka Rajan. Rethinking brain-wide interactions through multi-region 'network of networks' models. Current opinion in neurobiology, 65:146-151, 2020.
[9] Shenqin Yao, Peng Yuan, Ben Ouellette, Thomas Zhou, Marty Mortrud, Pooja Balaram, Soumya Chatterjee, Yun Wang, Tanya L Daigle, Bosiljka Tasic, et al. Recv recombinase system for in vivo targeted optogenomic modifications of single cells or cell populations. Nature methods, 17 (4):422-429, 2020.
[10] Kayvon Daie, Karel Svoboda, and Shaul Druckmann. Targeted photostimulation uncovers circuit motifs supporting short-term memory. Nature Neuroscience, 24(2):259-265, 2021.
[11] David Sussillo, Mark M Churchland, Matthew T Kaufman, and Krishna V Shenoy. A neural network that finds a naturalistic solution for the production of muscle activity. Nature neuroscience, 18(7):1025-1033, 2015.
[12] David Sussillo and Omri Barak. Opening the black box: low-dimensional dynamics in high-dimensional recurrent neural networks. Neural computation, 25(3):626-649, 2013.
[13] John J Hopfield. Neural networks and physical systems with emergent collective computational abilities. Proceedings of the national academy of sciences, 79(8):2554-2558, 1982.
[14] Sung Soo Kim, Hervé Rouault, Shaul Druckmann, and Vivek Jayaraman. Ring attractor dynamics in the drosophila central brain. Science, 356(6340):849-853, 2017.
[49] Ibrahim Ayed, Emmanuel de Bézenac, Arthur Pajot, Julien Brajard, and Patrick Gallinari. Learning dynamical systems from partial observations. arXiv preprint arXiv:1902.11136, 2019.
[50] Daniel Durstewitz, Georgia Koppe, and Max Ingo Thurm. Reconstructing computational system dynamics from neural data with recurrent neural networks. Nature Reviews Neuroscience, pages 1-18, 2023.
[51] Stephen Boyd, Neal Parikh, Eric Chu, Borja Peleato, Jonathan Eckstein, et al. Distributed optimization and statistical learning via the alternating direction method of multipliers. Foundations and Trends® in Machine learning, 3(1):1-122, 2011.
[52] Manuel Beiran, Alexis Dubreuil, Adrian Valente, Francesca Mastrogiuseppe, and Srdjan Ostoic. Shaping dynamics with multiple populations in low-rank recurrent networks. Neural Computation, 33(6):1572-1615, 2021.
[53] Jorge Nocedal and Stephen J Wright. Numerical optimization. Springer, 1999.
[54] Jian Yang, Lei Luo, Jianjun Qian, Ying Tai, Fanlong Zhang, and Yong Xu. Nuclear norm based matrix regression with applications to face recognition with occlusion and illumination changes. IEEE transactions on pattern analysis and machine intelligence, 39(1):156-171, 2016.
[55] Walter G Gonzalez, Hanwen Zhang, Anna Harutyunyan, and Carlos Lois. Persistence of neuronal representations through time and damage in the hippocampus. Science, 365(6455): 821-825, 2019.
[56] Zihui Zhang, Lloyd E Russell, Adam M Packer, Oliver M Gauld, and Michael Häusser. Closed-loop all-optical interrogation of neural circuits in vivo. Nature methods, 15(12):1037-1040, 2018.
[57] A Emin Orhan and Xaq Pitkow. Skip connections eliminate singularities. arXiv preprint arXiv:1701.09175, 2017.
[58] Peter J Huber. Robust estimation of a location parameter. In Breakthroughs in statistics, pages 492-518. Springer, 1992.
[59] Brian DePasquale, Christopher J Cueva, Kanaka Rajan, G Sean Escola, and LF Abbott. full-force: A target-based method for training recurrent networks. *PloS one*, 13(2):e0191527, 2018.
[60] Haim Sompolinsky, Andrea Crisanti, and Hans-Jurgen Sommers. Chaos in random neural networks. Physical review letters, 61(3):259, 1988.
S1 Details on CORNN derivations and implementation
S1.1 Separation into subproblems and derivation of the gradient and Hessian
The loss function given in Eq. (6) can be written in details as:
$$
\mathcal {L} (\hat {\theta}) = \sum_ {i = 1} ^ {n _ {\mathrm {r e c}}} \underbrace {\left[ \frac {1}{T} \sum_ {t = 1} ^ {T} c _ {t , i} \mathrm {C E} \left(\frac {1 + \hat {d} _ {t , i}}{2} , \frac {1 + d _ {t , i}}{2}\right) + \sum_ {j = 1} ^ {n _ {\mathrm {t o t}}} \frac {\lambda}{2} \hat {\theta} _ {j i} ^ {2} \right]} _ {\mathcal {L} _ {i} (\hat {\beta} ^ {(i)})}, \tag {S1}
$$