CORNN: Convex optimization of recurrent neural networks for rapid inference of neural dynamics
Authors: Fatih Dinc
Adam Shai
Mark Schnitzer
Hidenori Tanaka
Venue: NeurIPS 2023 poster (2023)
Keywords: brain-machine interfaces
recurrent neural networks
convex optimization
computational neuroscience

Abstract: Advances in optical and electrophysiological recording technologies have made it possible to record the dynamics of thousands of neurons, opening up new possibilities for interpreting and controlling large neural populations in behaving animals. A promising way to extract computational principles from these large datasets is to train data-constrained recurrent neural networks (dRNNs). Performing this training in real-time could open doors for research techniques and medical applications to model and control interventions at single-cell resolution and drive desired forms of animal behavior. However, existing training algorithms for dRNNs are inefficient and have limited scalability, making it a challenge to analyze large neural recordings even in offline scenarios. To address these issues, we introduce a training method termed Convex Optimization of Recurrent Neural Networks (CORNN). In studies of simulated recordings, CORNN attained training speeds $\sim$100-fold faster than traditional optimization approaches while maintaining or enhancing modeling accuracy. We further validated CORNN on simulations with thousands of cells that performed simple computations such as those of a 3-bit flip-flop or the execution of a timed response. Finally, we showed that CORNN can robustly reproduce network dynamics and underlying attractor structures despite mismatches between generator and inference models, severe subsampling of observed neurons, or mismatches in neural time-scales. Overall, by training dRNNs with millions of parameters in subminute processing times on a standard computer, CORNN constitutes a first step towards real-time network reproduction constrained on large-scale neural recordings and a powerful computational tool for advancing the understanding of neural computation.

CORNN: Convex optimization of recurrent neural
networks for rapid inference of neural dynamics
Fatih Dine Department of Applied Physics Stanford University Stanford, CA 94305
Adam Shai" CNC Program Stanford University Stanford, CA 94305
Mark J. Schnitzer'
Howard Hughes Medical Institute
CNC Program
Stanford University
Stanford, CA 94305
Hidenori Tanaka'
Physics & Informatics Laboratories, NTT Research, Inc.
Sunnyvale, CA 94085
Center for Brain Science, Harvard University Cambridge, MA 02138
Abstract
Advances in optical and elottrophysiological recording technologies have made it possible to record the dynamics of thousands of neurons, opening up new possibili- ties for interpreting and controlling large neural populations in behaving animals. A promising way to extract computational principles from these large datasets is to train cant-coont recurrent neural networks (dRNNs). Performing this training in real-time could open doors for research techniques and medical applications to model and control interventions at single-cell resolution and drive desired forms animal behavior. However, existing training algorithms for dRNNs are inefficient and have limited scalability, making a challenge to analyze large neural record- ings even in offline scenarios. To address these issues, we introduce a training method termed Convex Optimization of Recurrent Neural Networks (CORNN) studies of simulated recordings, CORNN attained training speeds ~ 100-ffold faster than traditional optimization approaches while maintaining or enhancing modeling accuracy. We further validated CORNN simulations with thousands of cells that performed simple computations such as those a 3-bit flip-flop or the execution timed response. Finally, we showed that CORNN can robustly reproduce network dynamics and underlying attractor structures despite mismatches between generator and inference models, severe subsampling of observed neurons, or mismatches in neural time-scales. Overall, by training dRNNs with millions of parameters subminute processing times standard computer, CORNN constitutes first step towards real-time network reproduction constrained on large-scale neural recordings and powerful computational tool for advancing the understanding of neural computation.
Introduction
Understanding the relationship between neural dynamics and computational function fundamental to neuroscience research []]. To infer computational structure from neural population dynamics,
'These authors contributed equally to this work.
These authors co-supervised this work.
'CORNN software and reproduction code are available at http://sstlalcoom/ccccccceelhhhh
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

neuroscientists regularly collect and analyze large-scale neural recordings using electrophysiological or optical imaging recording methods. Both recording modalities have undergone rapid progress in recent years, yielding steady increase in the numbers of cells that can be recorded and manipulated (277). Owing to this progress, an urgent need has emerged for developing and refining theoretical and computational tools suited for the analysis of large recording datasets.
Data-driven modeling of neural network dynamics, in which recordings of neural activity are trans- formed into computational model through network reconstruction, provides a promising way of bridging experiments and theory [ This approach not only complements recent experimental breakthroughs in targeted stimulation of individual neurons (19, [O)) for hypothesis testing and interfaces but also facilitates the distillation of fundamental computational principles
from large amounts of experimental data [ [2]. Historically, especially for recordings of indi- vidual or small numbers of cells, hand-crafted network models were used (0) While valuable, traditional hand-crafted models often struggle to handle the large datasets generated in contemporary neuroscience.
To bridge the gap. new approach has recently emerged in which artificial networks are trained on tasks mimicking those performed in the laboratory by animal models [21-23). However, this approach lacks the ability to transform observed patterns of neural activity into data-driven, mechanistic computational models. A promising solution to these issues, which retains the advantages of neural network- based modeling, is the training of dRNNs [24 These models, part of a broader suite of computational strategies [T] 21 30 32], align the dynamies of the neural units within the RNN to the activities of recorded neurons. This data-constrained methodology for RNN training has the potential to modernize estimations of functional connectivity
[3 -till mainstay in systems neuroscience [6] and provides a viable path to distill underlying computations by extracting the state-space representations of the learned dRNNs [[2.
Additionally, dRNNs may open the door to
computational control for causal experimentation 9 10 serving as valuable computational adjunct
to widely used techniques like optogenetics
[8]
Despite recent advances in the dRNN framework. how to perform fast and scalable reconstructions of neural activity traces from large- scale empirical recordings has remained unclear. Notably, the slowness of existing dRNN optimization algorithms may often necessitate the use of high- biigh-performance clusters and several days of computation large-scale experiments. These limitations might have been barrier to the widespread adoption of dRNN approaches. Moreover, to fully harness the potential of dRNNs, it essential to extend their range of applicability from offline analyses to on-the-fly applications within individual recording sessions. An approach to training dRNNs in a nearly immediate manner would be a key advancement that would enable novel experimental approaches, such theory-driven real-time interventions targeting individual cells with specific computational or functional roles (Fig. D. However, realizing these benefits requires having an optimization routine that is fast, robust, and scalable to large networks.
In this work, we present CORNN, convex solver designed for fast and scalable network reproduction. We demonstrate its accuracy and efficiency compared to alternative network training methods such as FORCE 35 and back- propaatopnn through time (BPTT), which have been employed in previous seminal works utilizing dRNNs (9). Our main contributions include the development of CORNN (Fig. [ the introduction of an initialization strategy to accelerate convergence (Figs. 2 S4
and
S6). and the demonstration of CORNN's s scalable speed (Figs. 3 S
S2
S3
S5 S8 and S9) which
enables rapid training on the activity patterns of thousands of neurons (Figs.
S7
S8
S1O
(ST]. CORNN's 'ss performance can be further enhanced by 1-2 orders of magnitude through the use of a standard graphical processing units (GPU) on a desktop computer (Figs. S5 and [S8. Unlike BPTT and FORCE, CORNN does not require fine-tuning (Figs. SS S6). making it a user-frrendde technology for biologists. Lastly, we highlight CORNN'S robustness against non-idcalities such as mismatches in the assumed dynamical system equations (Figs.
4 STO and S11 . subsampling
of neural populations (Fig. 5). differences in dynamical time-scales (Fig. $10 and existence of non-Gaussian (Fig. S33 or correlated noise (Fig. (SII).
By enabling user-friendly, subminute training on standard computers, CORNN represents a first necessary step in transforming data-constrained recurrent neural networks from a theoretical concept into an experimmental/compttaionn technology. While this work focuses on introducing and validating the fast solver on simulated benchmarks, future work with CORNN should focus addressing the challenges that arise when analyzing large-scale neural recordings from experimental neuroscience.
2

---

Neural recordings
Real-time interventions to neural populations
Real-time extraction of neural signal
Behavioral experiment
Predicting neural dynamic
Fang www
This work
Dynamics reconstruction with model
Revealing
Prepare for interventions
Single cell interventions
Neural population control
44
Desired behavior
Simulate single-cell targeting
Using tata-constrained recurrent neural networks for the interpretation and manipulation of brain dynamics within potential experimental pipeline. This approach centers around online modeling of network dynamics, which can enhance hypothesis testing at the single-cell level and support advancements in machine interface research. The training process is motivated by three objectives: (i) predicting the patterns of neural populations, (ii) revealing inherent attractor structures, and (iii) formulating optimal control strategies for subsequent interventions.
Figure 1:
Approach
2.1 Experimental setup for real-time interventions using dRNNs
Previous work suggested the use of dRNNs for real-time feedback between experimental and com- putational research B To date, however, the use of dRNNs has remained offline. Past studies extracted computational principles from neural recordings by fitting dRNNs after data collection 29 Moreover, the use of dRNNs has generally required the expertise of computational neu- roscientists, likely due to the complications associated with the optimization and training of neural network models. However, with the advent of large-scale recordings and data-driven approaches, real-time wrain-machine interface research is likely a forthcoming application (36-38]. In such use cases, neurobiologists might wish to train dRNNs via sser-friendly approach. Hence, fast and straightforward dRNN training procedure would enable new breed of interventional experiments to dissect the brain's microcireuitry.
To appreciate the advantages of fast data-driven reconstruction of neural dynamics, consider a hypothetical scenario in which dRNNs, facilitated by CORNN, enable real-time interventions the single-cell level (Fig. [ In this scenario. neural activities from mice performing behavioral tasks are captured and extracted in real-time using advanced imaging or electrophysiologice technologies and pre-processing algorithms [2] 3 39H As the experiment progresses, neural activity traces from each mouse are reproduced by training dRNNs, which are then reverse-engineered to reveal underlying attractor structures. Rapid dRNN training allows for tight feedback loop between incoming measured data and optimal experimental design, in order refine the inferred model. Techniques similar to adversarial attacks [42] might be used to devise optimal cell targeting strategies from these dRNNs R allowing one to test hypotheses about the computational roles of individual neurons or to identify optimal neurons for use within brain-machine interfaces. Once perturbation strategy is determined, it can be tested on subsequent trials of the experiment, using the same animal whose recorded neural dynamics led to the trained dRNN. Thus, fast dRNN training algorithm could allow for better fitting of a dynamical system model to the experimental data and provide a natural testbed to probe hypotheses about the computational structure of the biological neural circuitry.
Motivated by the goals described above, our paper focuses on training dRNNs accurately and fast as possible. However, experimental concerns and the need for biological interpretability lead to several constraints. Firstly, to ensure real-time communication with the experimental apparatus, we require that the training process takes place on standard lab computers, not clusters. Secondly, given the computational complexity of real-time processing in large scale recordings, especially with thousands of neurons, expect that at least one of the GPUs is reserved for extracting neural activities
3

---

Fixed point rntiiiizaaiion Hession inverse compulstion
Activtios
Optimization loop
Mm
Compun ents
Convex loes
Convex
Iocal eirustume
RNN
Random intialization
Prodictio
Dioskrder local eiruttum
Figure 2: CORNN: Convex and scalable solver for dRNNS via ADMM [51). The CORNN algorithm optimizes the parameters of a recurrent neural network so that activity in hidden units align with activities measured from ground-truth system, refer to generator. The choice of objective function results in a convex loss landscape. The algorithm starts by finding a fixed point for initialization, where the Hessians of all subproblems activities, which can be pre-computed and pre-inverted. The optimization loop then iterates through predicting the neural activities, computing the prediction error, updating the parameter values, and checking for convergence. The final output is the set of optimized parameters, which represents the functional connections between the neurons in the data-constrained recurrent neural network.
aligned to the correlation matrix of neural
from irainnimaging movies 43 experiments), or spikes from electrophysiological recordings [7], leaving the central processing unit (CPU), or perhaps a second GPU, for training dRNNs. Finally, maintaining biological interpretability via 1-1 matching of observed and modeled neurons is crucial, arbitrariness brought by hidden units can explain away existing functional connectivity and/or contributions of observed neurons to underlying attractor structures. These concerns make many traditional machine learning architectures and optimization routines, typically trained on large clusters with billions of parameters, unsuitable for our purposes. This includes recent work on the dynamical system reconstruction paradigm 31 45|550 50 which chiefly aims to maximize reconstruction accuracy of dynamical systems without regard biological interpretability, training speed, and scalability.
(though see for a promising development in smaller scale
2.2 Inference model and estimator family of interest
Recurrent neural networks universal approximators for dynamical systems, whose internal computations can reverse engineered and interpreted [2). Thus, we choose leaky firing-rate RNNs the inference models of CORNN, which follow the dynamical system equations [22]:
N Fee
N
dr: dt
-r; tanh(2;) + con 21
Wree
Tj
w W
uy +
E,
imput
(I)
where 7 is the neural decay time, 2; the total input to the recurrent unit ry is the firing rate of the unit i, and tanh(.) is the pre-defined non-linearity, W is the weight matrix for the recurrent connections, Win is the weight matrix from the input the recurrent units, uj is the input vector, and input is the input noise to the recurrent units. Unlike previous literature, given our goal of reconstruction, we also consider the existence of conversion noise eçoa that accounts for and mismatches in the non-linearity. Additionally, the inference model has the constraint Wree 0 preventing any self-excitation of neurons.
Suitable estimators for reconstructing such a model can be obtained (recursively) by discretizing the differential equation:
Pt.1. = (1 a1-a) + a(
(2)
where we define the time scale parameter a = At Here, the model parameters influence É,, explicitly linearly, and implicitly through previous time activities, such that â,,; = =E, Ez We define the short-hand notations æ [r: u| and Ô = [Wree; Winj" as concatenated matrices. The goal of the estimation problem is to find 0 minimizing loss function.

---

A reasonable choice for the loss function of this regression problem is the traditional Ca loss:
Pree
(2(
fr. re.i)²
However, due to the time recurrent definition of the general estimator ét,s, this minimization problem would require backpropagation through time (BPTT). It is not apriori clear if BPTT is the right optimization method for network reproduction, for which one has access to high information content regarding the internal dynamics the network. We discuss this the results below.
Instead of using BPTT, we consider smaller subset of estimators during training where fr.a is teacher-forced to re, when computing ft- This turns the potentially infinite time problem into single time-step one, which we call "single step prediction error paradigm." Under this condition, the L2 loss function transforms to a simpler form:
Theeo T
C2(0) = a
E(de, (de,s dr,)² f=1 1=1
where we define de,; (1 and correspondingly d = (a,,())).)) becomes an estimator
consisting of linear term through single non-linearity. However, this simplified version of the loss is not convex and leads to a counter-intuine global loss function with vanishing gradients (Fig. 2).
Convexification of the loss function
Rather than minimizing the naive L2 loss function, we observe that 1+de.s follows a linear sigmoid form. As result, we replace the C2 loss function with cross-entropy loss function that is well known to be convex under this estimator. In fact, in the limit of de,; + =1, the problem reduces to weighted loss of the form:
logistic regression. Instead of a naive replacement,
1 de, 1 di,s
Lweighter (0)
Cr. CE
Here, we use CE to denote the cross-entropy loss function and Cr2 non-negative weighting factors chosen Ct2i [1 df 1 for the CORNN solver. The specific choice of Ct,i. which resembles (but is not) preconditioning (53), is motivated from theoretical view in Supplementary Section S1.33 The reader can verify that the Hessian of this loss, derived in Eq. (S4). is positive semi-definite. For now, we write down the regularized loss function for CORNN as:
LCORNN (0) -
1
fweighted (0) 1011
(6)
where |0]|| is the Frobenius norm of ê and is the regularization parameter.
2.4 Fundamental principle of CORNN: Subproblem Hessians align to the correlation matrix of neural activity traces
Having convexified the loss function, we next observe that the loss function itself is perfectly separable such that LCORNN(0) E. (.(...). where 0,, denotes the ith column of the parameter matrix (See Eq. SI). This means that the original problem can divided into Nree sub-problems and solved independently. However, inspired by the fact that solving least-squares for vector or matrix target has the same complexity, we can devise faster method. Specifically, the choice of the specific weighting factors, Ct = dZ,,1 1- leads to shared (approximate) Hessian for all subproblems (See Eq. (S4):
1 x
H
where X stands for Ztz in the matrix form, I is the identity matrix, and À is the regularization parameter. We note that the approximation exact in the limit de, de,. In other words, we align the Hessian of each subproblem to the correlation matrix of neural activities. Unlike for
5

---

least-squares minimization, this is a local, not global, Hessian; but as long as we can initialize the network sufficiently close to the optimal solution, we expect that the descent direction with the approximate Hessian converges quickly. In this sense, the approximation in Eq. 77 shows similarities with quasi-Newton methods [53] but distinct in the pre-computed nature of the Hessian.
Following straightforward algebra, we obtain the gradient as
V&LCORNN (0) -0)-000)--- -- E+
1
(8)
where we define the prediction error matrix Ex,; = [de,.- dd... 1-
It is worth noting that each column of the gradient matrix contains the gradient for the scalar subproblem corresponding to the ith output. Then, defining the inverse (fixed-poin) Hessian matrix as
A = [
(9)
we obtain the following update rule (See Eq.
(20
gh+l = A - XO* + A+
(10)
where we highlight all quantities that can be pre-computed with blue. In essence, the update rule, which we call "Hessian aligned prediction error (HAPE) update," consists of computing the prediction
error followed by matrix multiplications; thus is well suited to be accelerated on a GPU if available, but still can be efficiently computed on CPU cores.
The alignment of subproblem Hessians relies on the assumption that we can find a "close-enough" starting position, called the fixed-point, for the initial parameters ((0), As we show in Supplementary Section S1.2 a suitable candidate the approximate least-squares solution:
O = AT
(11)
Then, the initial set of predictions becomes de, :- tanh æz. (O,)),,), which is subsequently refined through the HAPE iterations. This step computationally negligible, as the blue colored matrix is already pre-computed for the iterations that follow.
In this section, we focused on the leaky firing rate RNNs as described in Eq. ① However, CORNN can be applied to the other widely used variant of RNNs, Le.. the leaky current RNN described in Eq. (S28) and regularly employed in neuroscience literature [12 26 [2]. The reproduction of the leaky current RNNs can be made convex by observing the direct link between the firing rates, r, and the currents, z, via linear plus non-linear relationship, i.e., r tanh(z) and subsequently replacing di,; with Tt,; in Hessian and gradient calculations. To prevent introducing unnecessary complexity. we focus on leaky firing rate RNNs in this work.
3 Results
3.1 CORNN as versatile base model
While Hessian alignment, and subsequent HAPE updates, are the core principles driving the fast solver of CORNN, versatile base model should be able to incorporate various regularization schemes and sparsity constraints. For example, the addition of constraints W or a potential low-rank regularization would prevent the Hessian alignment and subsequently the use of fast HAPE updates. However, these constraints could be biologically relevant and perhaps vital to eliminate overfitting in the reproduced network. In fact, for the inference model given in Eq. [T). Tt, Ttt1.s for small a; making the unconstrained problem iol-posed since self excitation would explain a good amount of the activity of each neuron. Thus, to incorporate the equality constraint and many others, we developed solver using the alternating direction method of multipliers (ADMM) [1 (summarized in Fig. 2).
In the ADMM solver, the primary problem is divided into two subproblems that are solved indi- vidually during subsequent iterations, yet linked through a consensus variable that ensures agree- ment between the solutions of both subproblems upon convergence. In this framework, the first subproblem solves the unconstrained optimization problem, whereas the second one enforces the constraints/reguulriiatioo (See Supplementary Section SI for derivations). This division allows the

---

Figure 3: CORNN reduces training times by several orders of magnitude. The plots il- lustrate the relationship between reconstruction accuracy (Pearson's correlation coefficient be- tween 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 lit- erature for dRNN training (24 25 28 29]. Pa- rameters: = 0.1, free 5000, T = 30000. No input. coonv Poisson (10son10010 300 cinput (0, 10-4), Lines: median. Error bars: s.d. over networks.
CORNN
GD on loss
logistic loss
FORCE
CPU
GPU
min
hr
1 day
accuracy (1) of W...
Reconstr cuctuccccioonn
10
10
10
Training time (8)
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 LI and/or nuclear norm regularizations to the learned connectivity matrix [4], and utilize biological priors for further inductive biases. Finally, we note that given the highly irregular nature of the prediction error when de,; & =1, we supplied the fast solver with a simple automated outlier detection step explained in Supplementary Section $1.5
3.2 Backpropagation through time the L2 loss is suboptimal
In traditional setting, in which neural activations are hidden and need to be learned through training to provide 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. (T). Specifically, given that Eq. has only 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 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.T for details). As shown in Fig. ST contrary to expectation, feacher-forcing enhanced the convergence speed of the BPTT. Moreover, we also observed higher reproduction accuracy with the cross-entropy loss compared to the traditional La loss, inline with our observation from Fig. that Le loss leads to vanishing gradients in large error regime. As expected, BPTT took around a minute even in the toy example of Fig. ST with 100 time points and 100 neurons, approximately two orders of magnitude smaller than standard calcium imaging session in both dimensions.
3.3 CORNN decreases training times several orders of magnitude
Having shown that teacher-forcin 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 fixed duration of length T. and trained the dRNNs using CORNN, Newton's '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 28 29, was 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
(mman, variance)) of the data as shown in Figs. 3 S2 and S3 However, CORNN outper-
formed all other algorithms in speed in both cases with several orders of magnitude and converged
7

---

in sub-second times even with the CPU implementation. whereas others took up to several days for training (Fig. E]. We also observed increased efficiency for CORNN on a GPU (Fig. (S5) and that fixed-point initialization could stabilize other algorithms (Figs. S4 and S6p.
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 model RNN using BPTT to perform a 3-bit flip flop task following [12] as shown in Fig. [S7 (See Supplementary Section SSction(S3.2jeee for details). Fig. SS7 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, 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 example (original) generator and reproduced (CORNN learned) network in Fig. [S7 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, next tested whether increasing the number of trials could lead to better reproduction of the original network and whether CORNN can scale well fit the increased data and network size.
Fig. | 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 linear scaling of training times vs. the increasing number of trials and hence can handle large datasets (See Fig. (S)) In the other direction, observed that the training times scaled polynomially (Fig. S9). following n
CORNN learned 10% subsampled, time-scale and model mismatched
A
Generator network dynamics
C
Generstor Model (4(g)
natwork dynamics
Imput Window Distractor
Outzut window
Input window Distractor Oytou window
5000
p
Nowrn
500
Inferwron Model Le)
100
A.4
InpurCerput
500
[au1]
Turget Output
TSrget Output
Observed Subeopttetion
CORNN
predidione
0
0:0
Time (ms)
0:45
0.24, 1- 0.06
500
0.08
0
Time (mrs)
D.1
0.37
550
0.31
G: Canerrtrr, O2. Closerond, Inference
80.05
-0.1
04
0,255 Input
0.025
0.00
0.825 Output
Recurrent
Input weights
Recumant
Outpul
weights
weights
waights
weights
weights
Figure 4: CORNN can reproduce underlying neural dynamics and attractor structures despite several non-idealities in the timed -reponse 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 ag and a,, the CORNN-learner 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: ng 5000, R (O, 10 10-2) for 100 training trials. example network, ag At 0.1, and At Ims. B, D: no 5000, To
or = 0.1. C, E: no 500, a₂ (O.1,, 10.1,11-110 10-4).
8

---

with 3 E [1.3, 2.2], vs. the increasing number of neurons. Bringing both observations together, we concluded that there is an Of 7) 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, considered cases in which these assumptions are violated (See Fig. A)). Specifically, we trained generator RNN with different governing equations (See Supplementary Section S3.3) perform timed-response task (See Fig. 43), 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 novel distractor input, which present during the training of the task (See Fig. 43. C). If CORNN was able to learn the underlying state-space geometry, the reproduced network activities should be able return their attractive trajectory under novel perturbations.
not
When we subsampled the generator networks by 10% and introduced jitter in the time scales, observed that the inference model reproduced the neural dynamics in the observed population in novel perturbation trial (See Fig. 4C), even though the sub-connectivity matrix was not well reproduced (See Fig. D. E). On the one hand, this observation reinforces the findings of previous literature that the connectivity matrices learned by dRNNs should be interpreted functional connections, accounting for the behavior of networks at the level of neural dynamics, and not synaptic On the other hand, the ability of the mismatched network 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, (subsampling), [ST0 (time-scale mismatch), and S11 (correlated noise). Moreover, we observed the n scaling for the training times with increased number of network size (Fig. (C), inline with the results of Fig. S9
connections 15 25
4
Discussion
To place in perspective the experimental paradigm enabled by CORNN introduced in this work, we now return to the experimental scenario in Fig. T Imagine an experiment in which mice perform 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 ((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
A
Test acmuracy (R) of recoonsstun reconssruced set.uueedeedddde neural activily
accuracyy (1) W,
(rairr time (min)
Reconstruction
0.8
neurall
10
0.7
0.7.
0.
0.1
10% %
"10 Subsampling percentage out of 5000 neurons
100
1300
Subsampling percentage out of 5200 neurons
Number of neurons
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 R² between the ground truth and predicted activations, and C the linear scaling of the training times with asymptotic slopes [1.2, 1.3] vs increasing number of neurons on log-log plot, hinting at polynomial scaling. The different colors in the plot correspond to different number of training trials. Parameters: = 0.1, a.l,RRM RNN ~ (0, 10 10-2) for the training trials. See Supplementary Section S3.3 for further details. Data points: mean, error bars: 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 ms [2]. While new frames continue to be acquired, motion correction and neural activity extraction can happen in ~-5 ms per frame [56]. In addition to these ongoing imaging steps, the simulation and gpeision-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, L.e, potential interventions. Since the simulation involves only iterative matrix multiplications and point-wise nonlinearities, it finishes very quickly-in just few ms. For instance, running a 1000neuron RNN forward for 10 timesteps (~~ a second), for 100 different initial conditions, takes < 6 ms on a GPU, and < 25 ms 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 ~ 10 ms to optically stimulate the chosen neurons [56). Putting all the steps together, in under 100 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 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, simulated such experimental validation (Figs. 4 aad 5 and found that in the setting tested, CORNN was able to predict the dynamical effects of perturbations on the neural population.
Conclusion
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. 5 STO and[S1T) Inspired by this observation, further devel- opment 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 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 optogenctics [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).
[33] Moreover, reverse engineering the learned
This work constitutes 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).
10

---

Acknowledgements
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. MIS 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 PHI Lab during his internship at NTT Research.
References
[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 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 vivo imaging. elife, 5te 5:e14472, 2016.
[6] Sadegh Ebrahimi, Jérôme Lecoq, Oleg Rumyantsev, Tugce Tasci, Yanping Zhang, Cristina Irimia, Jane Li, Surya Ganguli, and Mark Schnitzer. Emergent reliability in sensory cortical coding and inter-area communication. Nature, (05(7911):713-77)) 13-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 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 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-25) 226.2 2021.
[11] David Sussillo, Mark M Churchland, Matthew Kaufman, and Krishna Shenoy. neural net- work that finds naturalistic solution for the production of muscle activity. Nature neuroscience, 1807):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-644 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-255 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.
11

---

[15] Abhranil Das and lla R Fiete. Systematic enrors in connectivity inferred from activity in strongly recurrent networks. Nature Neuroscience, 23(10): 2020.
[16] Mikail Khona and Ila R Fiete. Attractor and integrator networks in the brain. Nature Reviews Neuroscience, pages 1-23, 2022.
[17] Sukbin Lim and Mark S Goldman. Balanced cortical microcircuitry for maintaining information working memory. Nature neuroscience, (6(9):330--1144 2013.
[18] Rodrigo Laje and Dean V Buonomano. Robust timing and motor patterns by taming chaos in recurrent neural networks. Nature neuroscience, 16(7):925-933, 2013.
[19] Xiao-Jing Wang. Probabilistic decision making by slow reverberation in cortical circuits. Neuron, 36(5):555 2002.
[20] Hidchiko K Inagaki, Susu Chen, Kayvon Daie, Arseny Finkelstein, Lorenzo Fontolan, Sandro Romani, and Karel Svoboda. Neural algorithms and circuits for motor planning. Annual Review of Neuroscience, 45:249-274 2022.
[21] Daniel LK Yamins, Ha Hong, Charles F Cadieu, Ethan A Solomon, Darren Seibert, and James J DiCarlo. Performance-opimmized hierarchical models predict neural responses higher visual cortex. Proceedings of the national academy of sciences, 111(23):8619 8624, 2014.
[22] Nicolas Y Masse, Guangyu R Yang. H Francis Song, Xiao-Jing Wang, and David J Freedman. Circuit mechanisms for the maintenance and manipulation of information in working memory. Nature nearoscience, (177:1159-116,,, 2019.
[23] Guangyu Robert Yang. Madhura R Joglekar, H Francis Song, William T Newsome, and Xiao- Jing Wang. Task representations in neural networks trained to perform many cognitive tasks. Nature neuroscience, 22(2):297-366, 2019.
[24] Kanaka Rajan, Christopher D Harvey, and David W Tank. Recurrent network models of sequence generation and memory. Neuron, 90(1):128-142, 2016.
[25] Matthew G Perich, Charlotte Arlt, Sofia Soares, Megan E Young, Clayton Mosher, Juri Minxha, Eugene Carter, Ueli Rutishauser, Peter H Rudebeck, Christopher D Harvey, et al. Inferring train-wide interactions using data-constrained recurrent neural network models. bioRxiv, pages 2020-12, 12, 2021.
[26] Adrian Valente, Jonathan W Pillow, and Srdjan Ostojic. Extracting computational mechanisms from neural data using low-rank rnns. Advances in Neural Information Processing Systems, 35: 24072-24086. 2022.
[27] Lea Duncker and Maneesh Sahani. Dynamics on the manifold: Identifying computational dynamical activity from neural population recordings. Current opinion in neurobiology, 70: 163-170, 2021.
[28] Zach Cohen, Brian DePasquale, Mikio C Aoi, and Jonathan W Pillow. Recurrent dynamics of prefrontal cortex during context-dependent decision- making.. bioRxiv, pages 2020-11, 2020.
[29] Arseny Finkelstein, Lorenzo Fontolan, Michael N Economo, Nuo Li, Sandro Romani, and Karel Svoboda. Attractor dynamics gate cortical information flow during decision decision-making. Nature Neuroscience, 24(6):843 840-8500 2021.
[30] Valerio Mante, David Sussillo, Krishna V Shenoy, and William T Newsome. Context-dependent computation by recurrent dynamics in prefrontal cortex. nature, 503(7474):78-884, 2013.
[31] Chethan Pandarinath, Daniel J O' Shea, Jasmine Collins, Rafal Jozefowicz, Sergey Stavisky, Jonathan C Kao, Eric M Trautmann, Matthew T Kaufman, Stephen Ryu, Leigh R Hochberg, al. Inferring single-trial neural population dynamics using sequential auto-encoders. Nature methods, 15(10):805-815 2018.
[32] Ruxandra Barbulescu, Gonçalo Mestre, Arlindo L. Oliveira, and Luis Miguel Silveira. Learning the dynamics of realistic models of e. elegans nervous system with recurrent neural networks. Scientific Reports, 13(1):467, 2023.
12

---

[33] Ed Bullmore and Olaf Sporns. Complex brain networks: graph theoretical analysis of structural and functional systems. Nature reviews neuroscience, 10(3):186, 2009.
[34] Karl Deisseroth. Optogenetics. Nature methods, 8(1):26-29) 29, 2011.
[35] David Sussillo and Larry Abbott. Generating coherent patterns of activity from chaotic neural networks. Neuron, 63(4):544-55) 557, 2009.
[36] Kelly B Clancy, Aaron C Koralek, Rui M Costa, Daniel E Feldman, and Jose M Carmena. Volitional modulation of optically recorded calcium signals during neuroprosthetic learning. Nature neuroscience, 17(6):807-809 2014.
[37] Ricchiro Hira, Fuki Ohkubo, Yoshito Masamizu, Masamichi Ohkura, Junichi Nakai, Takashi Okada, and Masanori Matsuzaki. seward-timing-dependeee bidirectional modulation of cortical microcircuits during optical single-neuron operant conditioning. Nature communications, 5(1): 1-12, 2014.
[38] Aaron C Koralek, Xin Jin, John D Long II, Rui M Costa, and Jose M Carmena. Corticostriatal plasticity is necessary for learning intentional neuroprosthetic skills. Nature, 483(7389):331- 335, 2012.
[39] Andrea Giovannucci, Johannes Friedrich, Matt Kaufman, Anne Churchland, Dmitri Chklovskii, Liam Paninski, and Eftychios A Pnevmatikakis. Onacid: Online analysis of calcium imaging data in real time. In I. Guyon, U. Von Luxburg. S. Bengio, H. Wallach, R. Fergus, S. Vish- wanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, vol- wwwriis..cpppappeeepeeeeee
ume 30. Curran Associates, Inc., 2017. URL httpttttt////ooeeeeteeee pdf 0017//il//ecaa116fff555333333333
[40] Andrea Giovannucci, Johaauecce,ooooeessse Johannes Friedrich, Pat Gunn, Jeremie Kalfon, Brandon L Brown, Sue Ann Koay, Jiannis Taxidis, Farzaneh Najafi, Jeffrey L Gauthier, Pengcheng Zhou, Baljit Khakh, David W Tank, Dmitri B Chklovskii, and Eftychios A Pnevmatikakis. Caiman: An open source tool for scalable calcium imaging data analysis. eLife, 8:e38173, 2019.
[41] Hakan Inan, Murat A Erdogdu, and Mark Schnitzer. Robust estimation of neural signals in calcium imaging. Advances in neural information processing systems, 30, 2017.
[42] Aleksander Madry, Aleksandar Makelov, Ludwig Schmidt, Dimitris Tsipras, and Adrian Vladu. Towards deep learning models resistant to adversarial attacks. arXiv preprint arXiv: /00.:/700660833 2017.
[43] Hakan Inan, Claudia Schmuckermair, Tugce Tasci, Biafra O Ahanonu, Oscar Hernandez, Jérôme Lecoq. Fatih Dinç, Mark J Wagner, Murat A Erdogdu, and Mark J Schnitzer. Fast and statistically robust cell extraction from large-scale neural calcium imaging datasets. bioRxiv. 2021.
[44] Zhe Chen, Garrett J Blair, Changliang Guo, Jim Zhou, Juan- Romero- SSsao, Alicia Izquierdo, Peyman Golshani, Jason Cong, Daniel Aharoni, and Hugh T Blair. A hardware system for real-time decoding of in vivo calcium imaging data. Elife, 12:e78344, 2023.
[45] Steven L Brunton, Joshua L Proctor, and J Nathan Kutz. Discovering governing equations from data by sparse identification of nonlinear dynamical systems. Proceedings of the national academy of sciences, 13((5))33332-330 3937, 2016.
[46] Kadierdan Kaheman, J Nathan Kutz, and Steven L Brunton. Sindy-pi: robust algorithm for parallel implicit sparse identification of nonlinear dynamics. Proceedings of the Royal Society A, 47662224))202200279, 2020.
[47] Dominik Schmidt, Georgia Koppe, Zahra Monfared, Max Beutelspacher, and Daniel Durstewitz. Identifying nonlinear dynamical systems with multiple time scales and long-range dependencies. arXiv preprint arXiv: 19000037,, 2019.
[48] Georgia Koppe, Hazem Toutounji, Peter Kirsch, Stefanie Lis, and Daniel Durstewitz. Identifying nonlinaar dynamical systems via generative recurrent neural networks with applications to fmri. PLoS computational biology, 15(8))e 2019.
13

---

[49] Ibrahim Ayed, Emmanuel de Bézenac, Arthur Pajot. Julien Brajard, and Patrick Gallinari. Learning dynamical systems from partial observations. arXiv preprint arXiv:1902. 190.::9992.11166, 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 opti- mization and statistical learning via the alternating direction method of multipliers. Foundations and Trends@ in Machine learning, 3(1): 2011.
[52] Manuel Beiran, Alexis Dubreuil, Adrian Valente, Francesca Mastrogiuseppe, and Srdjan Os- tojie. Shaping dynamics with multiple populations in low-rank recurrent networks. Neural Computation, 33(6):1572-161: 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 pattern analysis and machine intelligence, 39(1): (09())1566-71)) 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 17110:170..09175, 2017.
[58] Peter Huber. Robust estimation of location parameter. In Breakthroughs in statistics, pages 492-518. Springer, 1992.
[59] Brian DePasquale, Christopher Cueva, Kanaka Rajan, G Sean Escola, and LF Abbott. full- force: A target-based method for training recurrent networks. PloS one, 13())=e0111527, 2018.
[60] Haim Sompolinsky, Andrea Crisanti, and Hans-Jurgen Sommers. Chaos in random neural networks. Physical review letters, 61(3):259, 1988.
14

---

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:
Ryee
T + dr.
1 dei
nttt
(
=
C
(S1)
T 1=1
2
2
j=1
(((80)
where Rtot. Thyor + Win. We note that for given i, the functions C; depend only on 3(0) :- 0..5. i.e. the ith column of the 0 matrix. In other words, the lack of cross-talk between the columns of 0 in the loss function means each output can be regressed independently with respect to the corresponding part of the loss function, i.e., L, (f.). Thus, for the rest of this section, without loss of generality, simply considered one instance of the regression, iee. fix derived the learning rule, and later generalized to the full problem.
The loss function of interest for the subproblem is:
1 1 de
1 1+ de
Rool
(3,}
cCE
+ A
(S2)
T
2
The gradient can be computed as:
1 de
1 de
(V (VB((0))
Ct V
+BBrr
T
=
T
Σ
T
Ct
2
(1
d:)(1 + +d+)
[
d
+ dr.
+ de
Ft, + ABk.
(S3)
1=1
1 T
- Σ
cel (oldd--e)staa + XBE
The Hessian follows directly from the gradient and the fact that tanh'(x) (x) (1 sanh²(2))
T Σ a(1 (e(1-))a Dee, + Adgt
(v3c(8)) 1
(S4)
At this point, is worth emphasizing that while Ft, shared across all subproblems, de are not. In fact, for any neuron i, different set of de,; would contribute to the Hessian of the corresponding subproblem. Thus, its current form, using the exact Hessian would require repeating matrix inversions for each of the Nrec neurons. We performed benchmarking with the exact Newton method, See Supplementary Section SS2] though the repeated matrix inversion was no longer feasible for large networks (Fig. To mitigate this, we took fixed-point solver approach in CORNN, described below.
First, we assumed that one can initialize the full problem with an initial point, call it d, = tanh (aan(335 (8).)). such that d, x d,. We discussed Supplementary Section S1.2 on how to estimate 2pp- Then, we replaced the Hessian with a local approximation CC ~ V2 ((2p)) within the vicinity of d, R de. leading to the the local approximate Hessian:
ΣΣ
1ー11
T
1- d2
1 x Hip Ad 2 T
(S5)
From this equation, we can infer why the choice of e the main text was scalable. The local Hessian depends only the inputs, which are shared across all the Rree independent subproblems, and not
15

---

the predicted or true outputs (de or de). Moreover, it needs to be computed once, and at the beginning of the optimization only.
Using the approximate Hessian and gradient, the parameter update rule simply follows:
Ap iD V C(3).
(S6)
Even though V 3 (3) was a vector for the subproblem, given that the (approximate) Hessians for all subproblems are aligned to the correlation matrix, we can promote the gradient to a matrix V²((0) to solve all subproblems once. Thus, one vectorize the computation for each subproblem, i.e., minimization of (3(0)) for ....., Nzec. without the need for external for loop.
At this point, we note that since Htp = 0 is positive semi-definite (PSD), the direction of the parameter updates is guaranteed to be descent direction. This is because the update direction has negative angle with the gradient following A V)[(() semi-definiteness. Thus, even if the error between de and de is large for an instance, e.g., fixed-point is incorrectly chosen, it can only hurt the speed of the optimization. Bringing all together, fixed fxxedppoint updates can be written as:
V V, V
V (f((() I by positive
Woll g(k-1) = (()) wh VoL(O())) (VCC
(S7)
where YY is the user-controlled learning parameter chosen either as (to bypass the line search fixed-point close enough) or through back-tracking line search. For this work, we simply picked Th 1 without any line search as we did not observe any convergence issues. We will simplify Eq.
(S7) to its final form when discussing the ADMM solver in Supplementary Section S1.4
S1.2 Fixed point initialization
To derive the fixed point update in Eq. I1 we performed Taylor approximation of the gradient and carried out a single update near the vicinity of the putative fixed-point:
g(++1) -((00 V ((30),
8(0) (+XXx+N)) * 300) -
(S8)
=3(0 -3(0 (+x**+N)
T x
where single-step converges to 3/. - (XTx x TAT) X Z, iee., the solution to the least squares problem the currents Z. Thus, Brs is expected to be near to both the solution 8 and the regime described by the txiet-piitt Hessian Hip: thus a suitable candidate to be the fixed-point itself. As an added bonus, this choice did not increase the complexity of the solver; since (X (X"X+TAX) -1 x has already been pre-computed for the iterative updates.
S1.3 Interpretation of the weighted convex loss
To interpret CORNN's loss function, we consider the low error limit of the cross entropy CE(p.p) around p 2 p:
CB(p.)) -plog(0) (1 p) log(1 p),
-H(pp) (plog(0) + (1 log(1 -D]]
(p p)
dp
p=p
d²
2 (plog(p) + (1 - p) log(1 - p)]
(p- p)
kmp
(p - 2)
-p)(pp+
+ O )))²)
2p(1 p)
16

---

where - (p) is a constant independent of the optimization variables, and thus can be ignored. Then, the low-error limit of the CORNN loss
Σ CCE
1 1 de
1 1+
de
1
(d; - dd
T
2
2
T
2(1 dlt2²
t=1
1
(de de (1 d,) [2-
(de))²
(S9)
2(1 de)
T 1
(d.))²
T
2
Here, the second row, we performed a Taylor approximation around the optimal value de R de (Z)lāzzd f-
this low error limit, minimizing the scalable weighted loss approximately accounts to minimizing the L2 loss on the currents. The main difference is, since the original weighted loss is on the firing rates and not currents, the weighted loss can account for the conversion noise, or non-linearity mismatches, that would induce bias in the explicit conversion & = f (d). The observation that low-error limit of the CORNN loss approximates the least- squsrss provided additional support to the claim that lqast-squares solution serves as a suitable initialization.
S1.4 Derivation of the ADMM updates
So far, we discussed the fixed-point solver for minimizing the unconstrained CORNN loss with L2 regularization. In this section, we discuss how to add constraints on weights. Specifically, in general scenario, it unexpected that all neurons would be connected each other. Rather, in many cases, one might be interested in introducing a trainable parameter mask such that v(i.j) E B(i,j) Wi 0. To account for this scenario, we incorporated our fast solver into the ADMM framework.
S1.4.1 Setting up the full problem
In this section, will first setup the full problem of interest with constraints, and then review the ADMM equations 51 To start with, we considered a problem of the following form:
T Σ
1 1
d
1
+ de
That (B,)
minimize
C(B)
T GCE 1=1
2
2
2
(S10)
subject to Vk E B()), Br 0.
We transformed this problem by adding a new variable X. which was constrained to be equal to 9:
1 T
1+ d 1 de
nttt
(3,)
minimize
OCE
1
2
(S11)
/=1
subject to wk B(k), Xk
8
Defining the indicator function 1 (xx) as
0
if x is true,
1(x)
(S12)
x
if false
The ADMM loss function became:
1 + de
1 + de
Near (3
minimize (ADMM((,8)
Σ
cCE
Σ
1(xx 0)
T
2
2
E=1
J=L
KEB(K)
(3)
Econstriaa (8)
subject to x= 3.
(S13)
17

---

ADMM Lagrangian has two separable loss functions of independent variables, connected through linear equality constraint, which can be augmented to obtain:
(p 8 k) = LADMM (3, X) + x x= x)
P 113 2
x112-
(S14)
2
where i is the dual variable. The ADMM steps (after re-defining ni - (1/p)rt) became (See Eqs. (3.5-3.7) in (1):
getz arg min (3)oRNN(3)+ + (p/2)]13 x + x²11))
(S15a)
B
x arg min (Coonstraine (8) + (p/2) 8 + x x1))
(S15b)
k+1 1 + gh++1 ght11
(S15c)
S1.4.2 Solving first ADMM iteration via fixed-point updates
Similar to the unconstrained case, we started by finding the gradient and the fixed-point Hessian:
CCORNN (8) + (p/2)118 x n 1
X E (A + p)B ()+ p(k* x*), (S16a)
[cORNN(3) + (p/2)113 x* + n² 11
~ T
x + (A+
(S16b)
Atp
where we recall the prediction error matrix:
d,,i dr,s
E
(S17)
dl.
Then, single fxed-point update step became:
App + M)T
+ (A+
1 x F. + (A + p)B p(* x5) T
X
+[**+T(>++
+ T(>+ |)] X E T() + +)
Tpx* +
-xt)]
(S18)
= -AtX*E (A+ 3 pAA
(x* -
with Ā = T' and أ = Tp. Here, we defined the short-hand notation for the inverse matrix:
At = x + (Ā p)1
(S19)
Consequently, the update rule for 7 became:
ge+z = 8h + x M)+ 3 px* (x=- + - px* )*)),
(S20)
where we highlighted all quantities that can be pre-computed with blue. For O, this reproduces the update rule for the unconstrained problem in Eq. (10). As before, we promoted the vector 3 to a matrix 0, hence vectorized across the full problem.
$1.4.3 Solving second ADMM iteration via projection
For this step, the goal is to minimize
Reao
minimize > (p/2) 82 + #,)g
(S21)
subject to 4k: E B(k). =0
Fortunately, this problem was perfectly separable across the vector entries of Xz. For cases where X 0 was constrained, this was the only feasible point and thus the answer that particular problem. For when & B(j), we obtained the minimum loss value by simply picking X; = gh+1 nk In short notation, the update rule for this subproblem became:
*+1 = arg min (Coonstraine (x) + (p/ (0/2)]1 x + n² (11) = IT (3et1 nt), (S22)
&
where we defined Bh as the projection operator to the constraint satisfying subspace (i.e, Vk B(k), X 0). Similar to before, promoted the vector & to a matrix to vectorize the full problem.
18

---

$1.4.4 Variable initialization and choice of
priori, it was not clear how p should be chosen. However, looking at the Hessian inverse A+, we observed that p 2 À was an appropriate choice for the step size p have a reasonable impact on the optimization problem. Empirically, we observed that p > À lead to faster convergence.
We initialized the primal and dual variables with x 3 Bep i.e. the pre-computed least-squares solution, Spp = A+ x+x,,, without additional computational complexity. Moreover, given that 0, the dual variable was bound to converge to non-zero values
ADMM ensured convergence only at the diagonal such that the 2nd ADMM step would converge.
To sum up, without any additional computational cost, we defined the initial conditions for all three variables as:
8 x Z. 8 = A+X
(S23a)
(S23b)
N =.
(S23c)
S1.5 Automated outlier detection and the pseudo-code for the full solver
Upon close inspection of the prediction error, we observed that the extremely large error values stemming from de se 1 would amplify the conversion noise and carry little-to-no learning signal. Thus, we implemented a simple outlier detection by ignoring errors larger than some value such that [Eel > = E 0. In the CPU solver, we scaled the gradient to account for the missing data points. Since these data points constituted small percentages, mostly < 1% of the full dataset, we did not perform the scaling for the GPU based solver to prevent unnecessary overhead.
It is worth emphasizing that this choice of outlier detection takes away the convexity of the problem since zeroing out part of the prediction emror leads to non-convex global loss. However, the outliers tended to be ignored throughout the optimization process, thus the problem remained convex in remaining data points. Future work can consider Huber-like (58) global convex robust loss that would automatically detect outliers. For the purpose of this work, we observed that the version of outlier detection we employed here performed well on empirical benchmarks.
For all solvers, we also defined a working precision (= 10-6) and whenever [dg] 2 81 we projected de back to =(1 - 6) to prevent numerical overflow errors. For experiments in Figs. S1 and S6 with BPTT, we projected the firing rates |rx 2 |1 - 8) to the boundary r. = = - 5) for the same reasons. Empirically, this stabilized all solvers and was performed ensure fair comparison.
Bringing all together, we state the full CORNN solver in Algorithm We note that for the figures in this work, we simply used a maximum number of iterations (30-100 100 depending on experiment) instead of monitoring the convergence of CORNN; as CORNN runtimes remained sub-min minute and a few more iterations than ideal were perfectly tolerable.
S2
Implementation details for other solvers
S2.1 Back- proppapopan through time
We performed back-propagation through time for the experiments in Figs. SI and S6 In our experiments, the teacher forcing means that the neural activities time ( + 1 is computed via Eq. (2) but we replace d..; in 2,, E, ft. with the data Zt,i, and fr,s with rt, Initially, ât,; depends on i implicitly, this replacement cuts off the time-dependent computation graph for the gradient. The probability of teacher forcing is the probability with which the replacement is performed at every single time point during the forward propagation of the network. Specifically, use the following borward-propggatio equations:
(1 a)re, + af(a))) p.
(S24)
(1 affe + af()))) u u2p p.
19

---

Algorithm 1: Convex and scalable solver for CORNN via ADMM framework
CORNN.fit (X; D, A, nlter = 30, Pin 100);
% X contains the aggregated input, e.g. firing rates and input units. It of size [T x Ntot];
% D contains the targets.
of size [T Nreel:
% To start, rescale the regularization parameter À = TX
p :- Pin %p is picked to be comparable to entries of X X and rescaled A;
% Pre-compute the inverses once during initialization;
At (x"x X + (p + 1
X A xx;T;
X AxxxX=XX
Z = F(D) &Compute the currents from the targets;
Pp X+Z%Initializz the fixed-point to the Is solution;
% Initialize the primary and dual variables;
= Bip &Initialize the least-squares solution;
== 0ip %Initialize to the least-squares solution;
u * 0p &Initialize the dual variable;
for = 1 nlter;
Predict the targets ò = f(Xe"):
Compute the prediction error E* = DD-
D) and perform outlier detection
2/0 stand for element-wise division/product:
Update first primal variable gh+1 ;= X 0h + X E* pA*(x** x²));
Update second primal variable x ++1 = II (Ok+1 + w²):
Update the dual variable w=*1 == n* + gk+1 x
Probe convergence by checking ||x gl+11| continue if not converged;
end return
where u is uniform random number generated between and 1. It is worth emphasizing that the teacher forcing in this context is distinct from [59], where the teacher signal is acquired not from the real-data but from a second network, though the overall idea is similar.
S2.2 Current and firing rate based FORCE
We followed the approach by 225 for training the recurrent weights to reproduce target neural activities via FORCE. For both network types, we defined the prediction errors as:
elt] = (f(t) - rrtt)/a, or elt] = =(t) - 2(t].
(S25)
Here, i((] the predicted firing rate time point I, vice versa for a({). We no longer performed any teacher forcing, thus r(t 1] and z/{ - 1] needed to predict +[4] came from the previous prediction step of the network, and were not the ground truth values.
The recurrent weights and the corresponding inverse covariance matrix P[t] were updated by the following rule:
P[[] P(t P(t 1)f(t 1 It-
1| - 1] 1| 10-1177 1
(S26a)
W ree [t] W rt 1] elt (|(
(S26b)
Here, e[(] was an Nrve dimensional vector, rather than scalar the case for the original work 35 We initialized P[[] - where 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 short-
comings this way of training RNNs using FORCE that render incompatible with the realities of experiments:
20

---

1. The training procedure requires particular range of À to converge (See Fig, S4 Even when FORCE does converge, this is for very limited range of À 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, the premise of FORCE approach is to have small error from the start to the end [35].
Point necessitates low-level hyperparameter tuning. The only work around point is to re- initialize the inverse covariance matrix P whenever there is 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 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. (S3) for two different scenarios: i) when e
corresponding to minimizing cross entropy loss, and ii) when ea [1 - --dd)-- 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?j 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 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 = 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 Wiee ~ /V (0. g² with zero mean and
Tiace
standard deviation of Vitee" For g > and sufficiently large Trress these networks can sustain their activities without any external stimulus and exhibit chaotic behavior [60]. For these chaotic generators, we picked g throughout this work ensure the networks were deep into the sustainable spontaneous activity regime. We chose the chaotic regime specifically, since previous work established that reacher-foreing when training an RNN to output 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., Wrree O. 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, established ground truth data by training RNNs on a 3-bit flip flop task. The architecture of the generator RNN was as follows:
r(I)= (1 a)r(1 - 1) + a tanh (Wn + W Wreer(1- (((1-))) 1)},
(S27a)
ott) W (t),
(S27b)
21

---

where rtt) is the hidden state of the RNN at time t, a is the time- conseent associated with the dynamics, Wn is the weight matrix governing imput into the RNN, are the recurrett weights, o is the output of the RNN and oon are the weights which take the hidden state to the output. Since CORNN does not learn Wout: to compute the outputs for the CORNN learned networks in Fig. [S8] we used the same projection matrix Wout as the original network. See the accompanying code for further details on network training and benchmarking experiments.
the input into the RNN, Wren
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 timed-response task with the following architecture:
x(t) (1 a)x(t-)(((t) - + aWinu(t) + aaWoor 1) + CERNN,
(S28a)
r(t) tanh (t)),
(S28b)
o(l) rl).
(S28c)
Here, unlike the CORNN and 3-bit flip flop equations (See Eq. SS77 the time-derivative was at the level of currents (I)), not firing rates (r). Upon receiving an input during [0, 100 we trained 10 networks to output a Gaussian pulse centered at t = 500ms with standard deviation 20mss otherwise zero everywhere.
we introduced novel distractor input ten times stronger than the cue at & [400, 410 ms. To allow fair comparison between networks trained on different numbers of neurons, we computed the R² of the neural activities from the first 100 neurons across three 1s trials and performed a uniform average. For the illustrative Fig. of finer details in time-activities after distraction, we instead plotted a scenario with an earlier and weaker (half strength than the cue) distractor i [200, 210 ms. To obtain the correlated noise in Fig. SII, we first sampled random noise from normal distribution, then convolved the i.i.d. noise with 2D Gaussian kernel across time points (s.d. of 3ms) and neurons (s.d. of neurons), and re-scaled by the targeted standard deviation. See accompanying code for further details.
During test trials in Figs. S10 and S11
allow visualization
S3.4 The mismatch between leaky-firing and leaky-current RNNs
In this section, we briefly discuss the extent of generator mismatch between Yeaky-firing and leaky- current networks. To see this, we consider the dynamics of the currents (z) under the leaky-firing rate RNNs:
da(t) dt = T rree dr(t) di
-W du(t)
dt
(S29)
dz(()
dt =()) + Whee f(=(t) wh(( + -Win du(() dt
Although the current dynamics look similar to the leaky-current RNNs, there are several key dif- ferences. Firstly, the term TWin du(t) is added to the leaky-current equation, which 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 the linear + non-lime linelr term, f((z,, (hence the name, leaky firing-rate RNN). Finally, under non-homogeneous time-scales for the leaky firing-rate RNN, would promote diagonal matrix and would not necessarily commute with Wrree 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 |
took five distinct steps:
1. Single-step prediction error paradigm, where we teacher-forced År,5, the RNN unit activities at time i, to the observed neural activities, Tt,i, to predict rt+1. with single-step back- propagation.
2. Convexification of the loss function, where we replaced the La loss function with cross- entropy loss.
22

---

3. Hessian alignment to neural activity correlations, where we weighted individual samples align subproblem Hessians to pre-computable quantity, i.e. the correlation matrix of observed neural activities.
4. Initialization 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 single layer feed-forward network with sigmoid/tanh non-linearity. Here, the inputs of the network are de, and rt, whereas the outputs are 1,+,, Since this step eliminated the need for the proppaapaoaanoonn through time, both the training times and memory requirements decreased significantly (Compare Figs. ST 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. and observed that teacher forcing increases the efficiency of the training, and was not deleterious. 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, Fig. 3 gradient descent on La loss converges faster than the gradient descent 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, shown in Figs. [2|and 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, cannot explain the speed or the scalability of CORNN by itself.
S4.3 Hessian alignment
The main contribution of our work is to pick Cr,ú, which approximates the problem local least- squares as discussed in Supplementary Section [S1.3] and aligns the Hessians of each subproblem pre-computable quantity, i.e., the correlation matrix of neural activities, as discussed in Section This alignment allows the use HAPE updates that have complexity of gradient descent but converge in (100) iterations thanks the approximate Hessian. Figs. S2 and S5 show that Newton's '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 cruical step to speed up the learning, but most importantly allow scaling to large networks due to its gradient-level complexity.
ST
S4.4 Fixed-point initialization
The fixed-point initialization relies Supplementary Sections S1.2 and (1.3). This initialization is 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 = arctanh(r). This is akin to regressing to the inverse of the output before the non-linearity in single layer network. Though such an approach can provide good first guess, it is not robust to conversion noise. This was shown Figs. 3 BgandSS [53] where non-Gaussian noise components lead to low accuracy for the initialization. Finally, the least-squares solution does not satisfy any additional
approximate solution when the final error is low (See
23

---

constraints on the connectivity matrix, thus is not suitable as valid solution itself. Yet, since the least-squares solution is computationally negligible, we used 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 ver- satile base model that incorporate additional design choices such LI 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 SS,, we used the full CORNN solver with the ADMM described in Algorithm| I The choice of ADMM for enforeing constraints is rooted in our observation that both the ADMM steps and the original HAPE updates take (100) iterations to converge. Moreover, other parts of the ADMM updates are negligible compared the primal problem, where the HAPE updates take place. In other words, the fact that HAPE updates and ADMM converge in similar steps creates 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 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 computer with two Geforce RTX 3080 Ti GPUs (though using only one) and Intel Core 19-10980XE 18-Core processor. Experiments in Figure $4 were divided run 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 cores due to memory limitations). Illustrative figures are plotted on an Apple Macbook Air with MI Chip. See the accompanying code provided in the Github repo for the specific hyperparameters used for the experiments.
Core 19. 9900X Skylake X 10-Core processor. Experiments in Fig. were
"Reproduction code is available at http://ghhalccooo/ssmiirrr/ra/rp
24

---

A
B
accurto of) W
Runtime per experiment is 47 seconds
Reconstruction accuracy (1) of
Runtime per experiment 47 i seconds
Reconstruction
0.8
Cross entrapy loss
Gross entropy loss
L2 loss
loss
tH
0.2
Fixed point initatization
No tsachar
25% teacher
50% tsacher
75% teasher
100% teacher foroing
Fised point
Intialization
No teucher
25% teacher
50% teacher
78% teacher fording
100% teacher fording
forcing
lorcing
furcing
forcing
forcing
forcing
forsing
Figure S1: Backpropagation through time (BPTT) on an L2 loss function is suboptimal, whereas teacher-forcing can increase the reconstruction accuracy. To assess the optimality of BPTT on an L2 loss function and investigate the potential drawbacks of teacher-forcing, we conducted reproduction experiments on chaotic networks with low (A) and high (B) input conversion noise ratio. We varied the levels of teacher forcing applied to BPTT and employed two distinct loss functions: cross-entropy and C2 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: a = 0.9, Tirec- 100, T 100, Repoch 3000. Vt = 1, 1..... T. u(t) 0.1 with Win ~ (O. g²) and g 3. A: <oonv (0, 10- cimput ~ NOO, 10- ). B: coou ~ (O, 10-10), eimput (0, 10 10. 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.
CPU implementations
A
Reconstruction
accuracy of W
Newton descent- weighted
CORNN
Reconsiruction
accuracy (r) of W
entropy
0.97
Newton descent- cross entropy
Gradient descent (ADAM)-
0.8-
cross entropy
Original FORCE applied
0.93
0.1
10
to currents (2)
0.1
10
Training time (s)
Modified FORCE applied to fring rates
Training time (s)
Figure S2: CORNN achieves a better combination of speed and accuracy compared to al- ternatives, dramatically accelerating convergence. The plots illustrate the relationship between accuracy (Pearson's '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: a 0.1, Pree 200, T 3000. No input. A: connv (0, 10- 8), cinput (O, 10 10-4). B: conv (0, 10 10), imppu AO 10 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.
CPU implementations
Figure S3: CORNN is robust to non-Gaussian noise injections. Same as in Figure S2 but with parameters: a Prec = 200, T = 1500, No input. conv
CORNN
Reconstruction accuracy (r) of W
Newton descent weighted
cross entropy
Newton descent- - cross entropy
Gredient descent (ADAM)-
0.1.
cross entropy
Original FORCE applied to currents (x)
Poisson(10-1 (100)) input ~ Poisson (10-2)
0.01
100
Training time (s)
Modified FORCE applied to firing rates (f)
25

---

B
Reconstruction accuracy (r) of W.
Reconstruction accuracy (f) W
CORNN
point inttalization only
0.5
Original FORCE applied to currents
Modified FORCE spplied firing
10
10
Regularization level (a.u.)
Regularization level (a.u.)
Figure S4: The traditional FORCE approach 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 similar task [24 25 [2]. is highly sensitive to the choice of hyperparameters, CORNN, whose accuracy is practically bounded below by the initialization, is not. Moreover, majority 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. (S2 In all plots, data points indicate mean values, and error bars denote SEM over 10 runs of the simulation.
GPU/CPU implementations
B
accuraoy (r) W...
CORNN (GPU)
accuracy (1) W
Reconstruction
Newton descent weighted
entropy (CPU)
Reconstruction
0.97
Gradient descent (ADAM)-
L2 (GPU)
0.9
Gradient descont (ADAM)
cross entropy (GPU)
Original FORCE applied to currents Modified FORCE applied to fring rates (r)
0.93
Training time (8)
Training time (8)
Figure S5: CORNN be drammatically accelerated with standard GPU. The same in Fig. S2 but with GPU acceleration for CORNN and gradient descent algorithms.
26

---

A
B
Runtime per experiment is 56 t 2 seconds
Runtime per experiment is 56 2 seconds
accuracy (r) W
Reconstruction
0.8
Reconstruction
80.6
accuracy (1)
0.7
Cross entropy loss L2 loss
Fored point intitaiization
No taacthar fording
26% teacher larding
50% teacher fording
75% teacher fording
100% teacher farcing
Fiee paint Intialization zalizzzzion
No macchar forcing
25% teacher fording
50% teacher forcing
75% teacher fording
100% teacher forning
Random initialization
Reconstruchin accuracy (r) of W
CORNN
Newton descent- weighted
entropy
Newton descent cross entropy
Gradient descent (ADAM)
cross entropy
0.1
10 100
CORNN initielized)
Training time (s)
Figure S6: Fixed-point initialization speeds up the convergence. A-B Same as in Figure ST but with fixed-point initialization. C We conducted ablation study 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: a 0.1, Tree = 500, T = 10000, conv (0, 10 10-10), imput (O, 10-2). All algorithms
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.
Input and output for the first trial
Trace of neural activations
Weights
Original
CORNN learned
801
Input
km1
Learnd
6m0
CORNN Original
network network
4m2
Neurons
-05
03
Recurrent
100
291
259
Learned
6.1
4m2
am0
-81
200
200
2.0
-6.
00
01
Time step
Time step
Original
Figure S7: CORNN can reproduce an RNN performing 3-bit flip flop task through dRNN training. We CORNN reproduce the behavior of synthetic RNN trained on 3-bit flip-flop task (12), where the network is given a short input pulse at random time. If the input pulse is +1, the output of the network must transition to output must transition to -1 or stay at -1, ignoring input pulses that do not match. Through dRNN training. CORNN able 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: = 0.9, Tree = 1000, 200 trials (100 data points per trial), conv ~ (0, 10 10-6), input ~ N(0, 10 10-4).
stay at +1. If the input pulse is -1, the network's
27

---

implementation
GPU implementation (64 bit)
implementation (32bit)
accuracy (r) W
Reconstruction
Number neurons @ 1.00
@ 3030
0.2
0.2
reconstructed output
Test accuracy (R*)
0.5.
0.5
100
slopes 0.98
100
slopes
100
Training time (min)
0.49
0.64
0.58
0.01
0.01
10
One
1000
Onee
100
1000
100
1000
One
day
month
day
month
Number of trials
Number trials
Number of trials
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 R² between the output units, and, bottom, the linear scaling of algorithmic wallclock runtimes with near-unity slope vs increasing number trials on log-log plot. The different colors the plot correspond 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 real-data applications, though not implemented for this case study. Parameters: = 0.5, conv ~ N 10-9), cnpput (0, 10 10-4). all plots, data points indicate mean values, and error bars denote SEM over 20 runs of the simulation.
A
B
100
Slope
100
Training time (min)
1.36
Training time (min)
1.44
Slope
1:96
2.17
Implementation CPU
2.07
0.01
bit)
0.01
100
1000
3000
100
1000
3000
Number of neurons
Number of neurons
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 SS7 subsampled neural populations and collected CORNN runtimes for CPU and GPU implementations. Parameters: a 0.5, conv ~ (O 10 "), inppu NOO 10 4). A with 200 trials (20.000 data points) and with single trial (100 data points).
28

---

Tost accuracy (R*) of > reconstructed neural activity
Nanteral
22
Reconstruction socurady of W.
Reconstruction accuracy (r of
0.2
-2-
N(O..[/sd.)))
Uncertainty (s.d.) in time-scales (a)
Uncertainty (s.d.) in time-scales (a)
Uncertainty (s.d.) in time-scales (a)
Figure S10: CORNN is robust to variations in the time-scales compared the ground truth.
Same as in Figure 5 but with 500 observed neurons and varying levels of uncertainties in time-scales.
Text rocuracy (R+ reconstructed noural activity
Unsuccessful reproduction level
Number
Typa
Figure S11: Correlated noise presents chal- lenge for CORNN, which can be mitigated by increasing the data size. Same in Figure [ but with the fully observed network and varying levels of injected random and correlated noise (for details, see Section S3.3).
intependent noise
Corre Correltedd ltte noise
10
Standard deviation injected
29


Reviews:
[{"review_id": "hbqHEiHCaQ", "reviewer": "Reviewer_1Kip", "summary": "The authors introduce a method to produce a digital twin of a biological neural network that can be inferred from measurements of neuronal activity.  The major advantage of the model is that it can be trained in real time to gain the basic dynamics of neural population. In the future this could allow advance HRI applications, like linking brain activity to controlling prosthesis.\n\nThe digital twin is trained using a randomised neural network that represents a neural population. It is used as a teacher network for student network that has a specific structure, designed in a way this is allows approximate convexification of the optimisation problem, enabling a ridiculously easy training of the network that can scale to large neural populations.  The computational cost of performing the analysis of the biological neural network scales linearly with its size.\n\nThe system is evaluated by a test, where a random neural network is initialised repeatedly with Gaussian random weights in a range that can produce also chaotic self-sustained signalling. The is assumed to provide a large enough variation for a neural signal generator\n", "questions": "on line 81: Why would the GPU use of real time acquisition of thousands of cell signals would force the dRNN calculation to CPUs? As these are tasks that can be pipelined,  the obvious answer is to use more GPUs .  Also, on can expect that signal detection part will be moved to special purpose chips tuned for inference as the signal dimension grows. \n\nThere has been a recent article (https://doi.org/10.1038/s41598-022-25421-w, published Jan 2023 ) on data-driven modelling of C. elegans brain with different recurrent neural networks architectures that found that 4-hidden layer GRU is able to accurately reproduce the stimuli. As this is based on comparing measured data to artificial neural network behaviour like in the current manuscript, without the real time acceleration, but using  back propagation training, it would be a very good test case for the advantages of using the CoRNN.\nPerhaps the Authors could add this to their submission.", "limitations": "NA", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "The real time, fast learning capability of the suggested architecture. The convexificaiton, as a technology. has definitely uses and application potential beyond the domain of the paper.\n\nExcellent figures to clarify the process.\n", "weaknesses": "The reasoning about the enforced use of CPU does not sound convincing to me. See the question part.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the Reviewer for their comments and endorsement of our work. We carefully reviewed each comment and have provided below our responses detailing how we changed the manuscript to incorporate the Referees’ suggestions. We would also greatly appreciate it if the Reviewer would like to propose any additional research areas upon which  the specific convexification trick we developed might have an impact, for potential mention in the Discussion.\n\n- Why would the GPU use of real time acquisition of thousands of cell signals would force the dRNN calculation to CPUs? \n\nWe thank the reviewer for this comment. The reviewer is totally correct - distinct analyses and algorithms can run in parallel on different GPUs if the computer has multiple GPUs, or analysis can be put into FPGAs (for example see https://elifesciences.org/articles/78344), or other computational resources. In general, it is the case that at least one GPU will be unavailable during an online neuroscience imaging experiment for things like online motion correction and cell extraction, and so we wanted an algorithm that works well with limited computational resources. That being said, given that a typical contemporary neuroscience experiment is offline, we designed CORNN to benefit from the availability of GPUs, as shown in Fig. S2, S4, and S5. We have changed the manuscript with these issues in mind in the following ways. \n\nFirst, we showed that in a large-scale scenario, GPU accelerated gradient descent on Pytorch falls behind CORNN by 1-2 orders of magnitude (experiment still running, see Fig. N2 in the appended PDF for the current stage of the experiment). Hence, CORNN is valuable even when GPUs are available. We plan to replace current Figure 4, which will become a supplementary figure, with this experiment if our work is accepted for publication.\n\nSecond, we revised lines 79-82 of the manuscript to clearly state that GPUs may be available for the real-time training as well and incorporated the discussion above regarding the on-chip processing:\n\n“Motivated by the goals described above, our paper focuses on training dRNNs accurately and as fast as possible. However, experimental concerns and the need for biological interpretability lead to several constraints. Firstly, to ensure real-time communication with the experimental apparatus, we require that the training process takes place on standard lab computers, not clusters. Secondly, given the computational complexity of real-time processing in large scale recordings, especially with thousands of neurons, we expect that at least one of the GPUs is reserved for extracting neural activities from brain-imaging movies (though see Chen et al. for a promising development in smaller scale experiments), or spikes from electrophysiological recordings, leaving the central processing unit (CPU), or perhaps a second GPU, for training dRNNs…”\n\n- There has been a recent article on data-driven modelling of C. elegans brain (...). (Request by the reviewer for real data application)\n\nWe thank the reviewer for bringing this relevant work to our attention. We now cited this work in our manuscript, where we discuss data-constrained applications of neural networks (line 39). \n\nWe perfectly agree that real-data applications would be very helpful to illustrate the usefulness of our approach. However, as it stands, the manuscript is 11 pages long after revisions, which we need to trim if accepted. The previous work (see line 55) aimed to apply existing optimization methods to train these dRNNs on real-data and already demonstrated the validity of this paradigm, though via an inefficient training paradigm that takes days to converge in our new experiment (Fig. N2). \n\nWe wanted our work to convey one powerful idea and in an accessible manner for readers from multiple fields. Our work is a necessary step towards the development of dRNNs as an experimental technology, especially for real-time interventions, but further steps should be taken. To address this, we added the following paragraph to the discussion:\n\n“All in all, 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 transformation of calcium traces, or electrophysiological spikes, into firing rates normalized within $[-1,1]$, applying CORNN solver developed in this work into first offline than online experimental scenarios, estimation of neuronal time-scales from the experimental data instead of tuning them as hyperparameters (See Table 2 in \\cite{perich2021inferring} for tuned hyperparameters) and perhaps developing a low-rank regularization approach into CORNN that opens the door to interpreting observed dynamics in terms of latent variables \\cite{beiran2021shaping,valente2022extracting}.“\n\nTo some extent, that previous work applying dRNN training to real data necessitated the fine-tuning of several hyperparameters [See Table 2 in \\cite{perich2021inferring} for tuned hyperparameters] constitutes evidence that applications to real-data are non-trivial and will require further work. For now, we address these issues as limitations in the Discussion section, and we lay out a feasible path for future research.\n"}, {"review_id": "cLqeuXxmTb", "reviewer": "Reviewer_GD24", "summary": "The paper proposes a convex loss for training data constrained recurrent neural networks (dRNNs). This is achieved by performing the optimization in the $d_{t,i}$ space instead of $r_{t,i}$ space (as defined in the paper) and using cross-entropy instead of $l_2$ loss. The authors discuss that this choice makes the optimization convex, and using a weighting scheme they show that a pre-computed inverse Hessian (using the data correlation) makes the optimization much faster than original BPTT or FORCE algorithms. Results are shown on synthetic chaotic rate networks under model misspecification as well as noise and subsampling. Better accuracies and faster runtimes suggest that CoRNN has desirable properties compared to baselines.", "questions": "* How are the $d_t$'s initialized to be close to the true ones?\n\n* Since the Hessian is not the exact Hessian, does the optimization still guaranteed to find the (local) minimum? Or is it biased? Can the authors use debiasing strategies to account for the Hessian mismatch?\n\n> 128 $\\mathcal{L}_{\\text{CoRNN}} = \\sum_i \\mathcal{L}_i(\\hat{\\theta}_i)$.\n\n* Is this a trivial fact? Is $\\hat{\\theta}_{:,i}$ a specific row/column of the weight matrices? This basically means the optimization problem is separable over different rows/columns. This intuitively sounds unrealistic to me unless again I'm missing something important here.\n\n> 111 [...] This turns the potentially infinite time problem into a single time-step one.\n\n* How is this done? The parameters of the model are still $\\theta$. If the loss sums over all $t,i$ then the gradients with respect to $\\theta$ are still computed in a similar fashion to before unless I'm missing something important here.\n\n> 122 $c_{t,i}=[1-d_{t,i}^2]^{-1}$.\n\n* How is this chosen? This means that some neurons/time points are contributing more to the loss and they are more important to be captured properly, right? Why is this a good choice as opposed to treating all neurons/time points the same?", "limitations": "See weaknesses section", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "* Developing tools for online experimentation and targeted manipulation of brain circuits is a significant problem. Given the current experimental advances in the neuroscience field these tools are in crucial demand to help us understand geometrical and computational properties of the brain networks. This makes the problem considered by this paper timely and significant.\n\n* Arguably one of the main bottlenecks of this online experimentation is fitting/inference time complexity. This paper takes a step towards speeding up the optimization by convexifying the loss and pre-computing objects to avoid additional computational burden.\n\n* The motivation is clear and the presentation of the paper is logical, understandable, and easy to follow.", "weaknesses": "Below I've included a list of questions and concerns. In summary, I believe that the experimental part of the paper can be improved by a large margin. Specifically since the method is pitched as a general purpose optimization framework it's important to present results on more complex synthetic datasets as well as real neuroscience dataset where the performance of the model is well investigated under various parameter changes. In addition, I think the discussion section can be further improved by being more upfront about the limitations of the work and paths for future work.\n\n> 81 extraction, especially with thousands of neurons, requires GPUs, leaving only the CPU for dRNNs.\n\n* I'm a little confused by this repeating statement, can we not use multiple GPUs? Is this a constraint for the specific hardware you're using or a more fundamental constraint?\n\n* Eq. 1: correct me if I'm wrong but most of the theoretical neuroscience work is done on the equation where the nonlinearity is applied pre-summation $\\tau \\frac{dr_i}{dt} = -r_i + \\sum w_{ij} \\phi(r_j)$. Although from a model capacity perspective, the two models are shown to have similar capacities the one I included here is more biologically relevant. Can this be implemented in your approach and can you include some results on this alternative model? \n\n* Related to this, the framework seems to heavily rely on the specific instantiation of the network dynamics. For example, given a different form, we cannot immediately consider the loss form of $d_{t,i}$ and use its linear+sigmoid form. This seems to be an important limitation and some discussion on how to extend this to other forms of ODE would be helpful.\n\n\n> 114 $d_{t,i} = \\tanh(\\hat{z})$\n\n* Should it be $\\hat{z}_{t,i}$?\n\n\n* Why is (5) convex with respect to $\\theta$? Is this a trivial fact that the authors leave it unexplained?\n\n* Equation 7 is playing a crucial role in the optimization, can the authors bring the proof to the main text instead of SM?\n\n* Based on (S4), we can in principle use other nonlinearities and a corresponding weighting function $c_{t,i}$ right?\n\n\n**Figures**\n\n* Fig. 1 seems to be over-promising. Although in principle dRNNs can lead to the applications that are discussed but those are not shown in the current work. Unless those applications are not shown in real data/simulations I recommend rewording/having a more concise Fig. 1 that focuses on the contributions. Usually, Fig. 1 is a schematic of what's done in the paper, not an idealistic scenario that can be done in future studies.\n\n* Fig. 3: The colors are very hard to see, specifically because the error bars are too small, can the authors change the formatting to make it easier to read? After reading this paper two times I'm still unsure what \"teacher enforcing\" means in this context and how it's achieved in various percentages. Instead of these two parameter configurations, can the authors show these as line plots when changing the parameters more continuously on a certain range that includes the performance extremes? Related to this, are the claims true about this specific instantiation of the network or does it hold if we change the noise model? For example, intuitively, cross-entropy loss corresponds to the log probability of a categorical distribution, and as the authors show it provides a better surrogate for the data generated from normal distribution. What if we change the noise model to something like a Poisson or some noise model that's not independent across different time points?\n\n* Does the result hold if we change the dynamical regime of the network? In neuroscience, previous work has hypothesized various network mechanisms such as limit cycles, multi-stability, input-driven networks, transitions between multiple attractors, etc. Is it possible to show results on these other regimes to ensure that the presented optimization framework is not beneficial for a particular regime? Can the authors include experiments with varying the dimension of the chaotic attractor or a more complex data-generating model as opposed to a simple chaotic (low-d) rate network? \n\n* Fig. 4B: The result is a little surprising for me, how can the accuracy go down for other methods? Aren't all of these methods guaranteed to converge to a local minimum? Are the authors controlling for the initialization of the parameters?\n\n* Are the times reported in Fig. 4 CPU times or GPU times? \n\n* Fig. 5: In (A) there are these weird bumps/spikes in the network output. It looks like the errors are happening in a space that's not within the null space of the output. Whereas my experience with the flip-flop task is that BPTT errors don't look weird like this. Is this due to the specific architecture chosen or is this a result of the convex loss/optimization framework?\n\n* Fig. 7: miss-matched -> mismatched\n\n* I'm again surprised that the network can account for the common inputs/unobserved nodes. I think it's an established fact that by just using observational data no method can generically recover the weights in the presence of common inputs and unobserved nodes. This might be the result of the specific data-generating model that's used here. A discussion on this would be helpful to point to the limitations of the work.\n\n* Although the work is motivated by neuroscience data and experiments, no experimental data is included in the paper and that to me is the biggest weakness of the paper. Brain data is much more complex and heterogenous and the noise model, nonlinearity, dynamical regimes, and other factors make it harder to build a robust estimator that generalizes to test trials/test time points. I would assume that various loss functions have different properties that make them more suitable to specific datasets and it's hard for me to believe that a single loss can improve loss on all datasets. Unless results on more complex datasets and simulations are shown it's hard to decide whether the presented framework is preferred to the traditional ones.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the time the Reviewer put into this detailed review. As they noted, developing tools for real-time experiments and targeted intervention of neural activity is a significant and timely problem. The reviewer also commented that the motivation was clear, logical, and easy to follow. In our revisions, we maintained these strengths while addressing the concerns raised regarding additional experiments and discussion of limitations. We believe the feedback provided has improved the quality of the manuscript substantially. Below we describe the specific changes we made.\n\n- (Request for real data applications and complex synthetic datasets, discussion of limitations)\n\nThe reviewer raised important points regarding applying CORNN to real data and testing on more complex datasets. In the revised manuscript, we newly added three more experiments further testing more complex noise conditions in addition to the existing benchmarks. In particular, correlated noise, in both the space and time dimensions, and poisson noise (Figs N1-N3). While application to alternate models like LSTMs and real data remains for future work, the manuscript does include a number of non-random networks as synthetic benchmarks, in Figs. 5 and 7. \n\nIn addition, we have expanded the discussion of limitations and future directions in the manuscript, including considerations like nonlinearity mismatch and applications to real data. Please also see our general comments to all Reviewers above about our proposed changes to the abstract. Additionally, in our response to Reviewer 1Kip, we described limitations regarding real data applications and plans to address those in future work.\n\n- (Clarification on CPU usage)\n\nPlease see our response to a similar question by the Reviewer 1Kip.\n\n- (Question  on Eq. (1))\n\nPlease see our response to a similar question by the Reviewer mPZo.\n\n- (Proof of convexity)\n\nWe addressed this comment by making the following change before Eq. (5):\n\n“​​(...) As a result, we replace the L2 loss function with a cross-entropy loss function that is well known to be convex under this estimator. In fact, in the limit of $d_{t,i} \\to \\pm 1$, the problem reduces to logistic regression. The reader can also verify that the Hessian of this loss, derived in Eq. (S4), is positive semi-definite…”\n\n- Fig. 1 seems to be over-promising. (...)\n\nWe made substantial changes to clarify the contribution of our work and how Fig. 1 directly relates to our findings, please see our general comments to all reviewers. Please also see our response to the reviewer Gp7K’s second question.\n\n- (Clarification on teacher forcing) \n\nThe teacher forcing in this scenario means that the time activities at time $t$ is computed via Eq. (2), but we replace \\hat x_{t,i} in \\hat z_{t,i} = \\sum_j \\theta_{t,j} \\hat x_{t,j} with the given data x_{t,i}. \\hat x_{t,i} depends on \\theta implicitly, so this replacement cuts off the time-dependent computation graph for the gradient. The probability of teacher forcing is the probability with which the replacement is performed at every single time point during the forward propagation of the network. While we currently discuss the approach in supplementary material section S2 and provide code, we will add further implementation detail if this work is accepted. For now, due to time constraints, we focussed on performing additional experiments requested by the reviewers.\n\n- Are the claims true about this specific instantiation of the network or does it hold if we change the noise model? \n\nPlease see our general comments regarding the newly added figures.\n\n- Does the result hold if we change the dynamical regime of the network? \n\nYes, the results hold when networks have specific attractor structures. Please see Figs. 5 and 7 and our proposed changes to the abstract.\n\n- (Surprising Fig. 4B results)\n\nThe accuracy in these graphs are measured through the correlations between learned and original networks. The models are minimizing errors defined on the predictions. It is not possible to define a loss function on the ground truth weights, nor is it experimentally relevant.\n\n- (Unexpected spikes in Fig 5A)\n\nTo clarify, Fig 5 A is an illustrative example of a network that arises when the specific architecture we describe in Eq. (1) is trained with BPTT, not CORNN, to perform a 3-bit flip flop task (See also the spikes in Fig. 2 of Sussillo and Barak, 2013). Therefore, this figure or any error the network makes has no relationship with CORNN or the convex loss. Based on the reviewer’s feedback, we have updated lines 217-218 to explicitly state that we used BPTT for training networks to perform the 3-bit flip flop task in Fig. 5(A).\n\n- (observational data vs synaptic connectivity)\n\nWe are not interested in reproducing synaptic connectivity. We only used it as a surrogate to test the reproduction accuracy in cases where generator and inference networks followed the same dynamical system equations. To clear any confusion regarding the learned connectivity matrix, we added a paragraph to the discussion, though omitted here due to space limitations.\n\n- (no real data + how to decide if our work is preferred to traditional ones)\n\nPlease see our general response, as well as response to Reviewer 1Kip for the changes we made to address this comment.\n\n- (d_t initialization)\n\nThey can be initialized using theta_ls as defined in Eq. (11). We added the following sentence right after Eq. (11) to clarify this:\n“Then, the initial set of predictions becomes $\\hat d_{t,i} := \\tanh(\\sum_j x_{t,i} \\theta_{\\rm ls})$, which is subsequently refined through the HAPE iterations.”\n\n- (Inexact Hessian and convergence)\n\nSince the approximate Hessian is positive semi-definite, the parameter update direction is guaranteed to be a descent direction. We updated the relevant part in supplementary text.\n\n- Question on lines 128, 111 and 122\n\nWe noted the relevant supplementary sections in the main text so that the reader can find them easily.\n"}, {"review_id": "eY3QjSfX9e", "reviewer": "Reviewer_Gp7K", "summary": "The paper introduces a new method for training RNNs that allows fast convergence which could make realtime parameter update possible.", "questions": "NA", "limitations": "NA", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "- The method conceptually is simple and intuitive.\n- It tackles an important problem which is the time consuming nature and computing intensivity of ML models\n- It achieves impressive results through simple hacks\n- One possible use of this is also quick experimentation with RNN models\n- Outlier detection is included", "weaknesses": "- The methods section is easy to understand until the end of setion 2.3 where there are many equations with variables that are not defined either in the main text or in the SM (any equation in the main text has to be defined very explicitly in the main text). \n- To evaluate the value of this implementation on the proposed application, I would appreciate a comparison of the speed of computation to the sampling rates of data collection and intervention delivery.\n- I would have appreciated results experimenting on real data and comparing neural dynamics computed via this method compared to other methods", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the Reviewer for their careful reading of our manuscript. The suggestion to better tie our results back to Figure 1 is an excellent one, which we implemented. You can find below our detailed, point-by-point responses to each question raised.\n\n- any equation in the main text has to be defined very explicitly in the main text \n\nWe thank the Reviewer for this comment. The Reviewer is absolutely correct and we apologize for the inconvenience. To address this, we refined the presentation of our manuscript. We ensured that every math variable is defined in the main text, obviating the need to navigate to the supplementary material.\n\n- To evaluate the value of this implementation on the proposed application, I would appreciate a comparison of the speed of computation to the sampling rates of data collection and intervention delivery.\n\nExcellent point. Indeed, we could have done a better job tying our results back to the experimental paradigm introduced in Fig. 1. Specifically, we should have discussed how our results advance experimental neuroscience research in practice and can be readily applied to real-time interventional experiments. To address this aspect, we added the following paragraphs to the discussion section:\n\n“To place in perspective the experimental paradigm enabled by CORNN introduced in this work, we now return to the experimental scenario in Figure 1. Imagine an experiment, where mice perform a predefined task several times, say for an half hour, with 3000-4000 observed neurons. This 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 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. \n\nTo 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 ms (Kim and Schnitzer, 2022). While new frames continue to be acquired, motion correction and neural activity extraction can happen in \\sim 5 ms per frame (Zhang et al., 2018). In addition to these ongoing imaging steps, the simulation and decision-making with the CORNN-fitted dRNN is 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 1000 neuron RNN forward for 10 timesteps (~ a second), for 100 different initial conditions, takes <6 ms on GPU, and <25 ms on CPU, in our hands. Once a desired intervention is identified from the simulations, the phase mask on the spatial light modulator can be updated in \\sim 10 ms to stimulate the chosen neurons (Zhang et al., 2018). Putting all the steps together, in under 100 ms, one could capture brain activity, simulate future responses, decide on an intervention, and update the optical stimulation - fast enough for real-time closed-loop application. By streamlining network training, CORNN enables dRNNs for these types of experiments.”\n\n- I would have appreciated results experimenting on real data (...)\n\nThank you for this comment. Please see our response to the reviewer 1Kip on a similar question.\n"}, {"review_id": "44hB7JpfRA", "reviewer": "Reviewer_mPZo", "summary": "The motivation behind this paper is to replace FORCE/RLS fitting of biologically-flavored RNNs to data by a faster and more robust alternative training scheme. This is accomplished by 1) altering the training objective to a cross-entropy loss; 2) using an approximate pre-computable Hessian to facilitate the use of Newton steps; 3) use of a suitable fixed point initialization; 4) using teacher forcing (training only one-step losses); 5) employing ADMM to enforce constraints (e.g., no self-excitation). Experiments show improvements to various metrics, including accuracy of recovered weights, fits to observed activity, and runtime. \n\nThe stated goal of the project is to produce a method fast enough for real-time analysis, though this work focuses only on weight inference and fitting accuracy of the firing rates. Still these models are increasingly used in the field to fit neural data, so better algorithms for doing so are needed.", "questions": "- It is unclear to me what the difference is between panels A and B of Figure 3. Just different parameters? It would help if this were reflected in the titles.\n- Figure 4: I realize that the authors would like to zoom in to show differences between models, but the choice of y axis range obscures the fact that these are modest to minor differences between the models.\n- The authors use correlation between the true and inferred weights as their accuracy measure, but this (so far as I can see) first noted in the caption to Figure 4, even though it occurs in Figure 3.", "limitations": "1. Equation (1) is not the usual form for the equations defining these networks (e.g., refs 22, 24). Typically, the leak term is on the voltage variable, not the firing rate, and the nonlinear transfer is from voltage to firing rate, not input current to effective drive. It looks to me as if this choice is essential to make the cross-entropy convex in $\\theta$. If so, this should be discussed, since it's a qualitatively minor change that has large implications for the feasibility of the approach.\n\n2. As noted above, it seems to me there are five distinct technical pieces, and it is not always clear what is doing the heavy lifting. For instance, Figure 3 (and S3) seems to indicate that initialization in these models is doing a huge amount of work, which is unsurprising but not well emphasized in the text. Better delineation in the main text of which choices are most important to the results would be helpful.", "rating": 7, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Good algorithm performance as assessed along multiple dimensions (speed, inference accuracy, reconstruction accuracy).\n- Thoughtful comparisons to other baseline models.\n- Emphasis on real-time and streaming settings.\n- Experiments targeting model mismatch.\n- ADMM approach allows for many different types of potential constraints.", "weaknesses": "- Addresses some key sources of potential model error but not others (e.g., correlated noise, wrong nonlinearity, mismatch between network timescale and measurement timescale for some data types).\n- No experiments on real data. Even if accuracy of weight recovery is impossible to assess, reconstruction accuracy would be.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the Reviewer for their expert review, careful reading, and identification of typos, which we have now fixed. We especially appreciated the expert parsing of our approach into 5 distinct components; this showcased how well the Reviewer understood and appreciated our research. With the Reviewer’s permission, we would like to use similar language to describe the contributions of each individual component of the fast CORNN solver, within a new supplementary section that we plan to add. Our detailed responses to each question are below.\n\n- (Request for more simulations and real data applications)\n\nTo address the request by the Reviewer for incorporating more non-ideal conditions, we performed new experiments with correlated (Fig N1) and non-Gaussian noise (Fig. N2-3). The reviewer also raised a need to check non-linearity mismatch. We believe the current manuscript addresses these points in two ways. In particular, we studied: \n\ni) effect of conversion noise, which creates an inevitable mismatch between non-linearities. Instead of tanh(.) non-linearity, the conversion noise uses a time and neuron dependent non-linearity of the form “tanh(.) + \\epsilon.”\n\nii) a leaky-current RNN in Fig. 7, whose equations do not match those of the leaky firing rate RNNs. We presume that using a different data generator is more stringent than having a different non-linearity on the same model. \n\nYet, we agree that processing datasets where firing rates are not bounded within $[-1,1]$, which might stem from other non-linearities, will require some form of data transformation. This is part of the future work that needs to be addressed to apply CORNN to real data, though previous work successfully addressed similar concerns before (See Rajan et al. 2016, Perich et al. 2021). Please also see our response to the final question of reviewer 1Kip, for a related discussion of real data applications and how we changed the manuscript accordingly.\n\n- Figure 4: (...) minor differences between the model accuracies.\n\nWe thank the Reviewer for this comment. Indeed, our primary claim about CORNN is that it enables faster training, rather than increased accuracy. On the other hand, with the revised Fig. 4 (See Fig. N2 in the shared PDF), the differences in accuracy are far more clear between models. This is partially because we tested the optimizers on a large-scale scenario with 5000 neurons, as opposed to a few hundred neurons in the prior Fig. 4.\n\n- Equation (1) is not the usual form (...)\n\nBoth leaky firing rate and leaky current RNNs are regularly used in the field of neuroscience. After receiving the reviews we realized that our arbitrary choice to focus on one architecture made it seem like our approach was not applicable to the other. However, CORNN can be trivially implemented for both types of RNNs. To address the concern, we added the following paragraph to the end of the Approach section:\n\n“In this section, we focused on the leaky firing rate RNNs as described in Eq. (1). However, CORNN can be applied to the other widely used variant of RNNs, i.e., the leaky current RNN described in Eq. (S27) and regularly employed in neuroscience literature [citation]. The reproduction of the leaky current RNNs can be made convex by observing the direct link between the firing rates, $r$, and the currents, $z$, via a linear plus non-linear relationship, \\emph{i.e.}, $r = \\tanh(z)$. To prevent introducing unnecessary complexity, we focus on leaky firing rate RNNs in this work, though the generalization to leaky current RNNs can be trivially performed by replacing $d_{t,i}$ with $r_{t,i}$ in Hessian and gradient calculations.”\n\n- Better delineation (...) of which choices are most important (...)\n\nWe thank the Reviewer for this extremely helpful comment. The Reviewer is right, we should have made this point clear. Moreover, this comment made us realize that we dissected all aspects but point 3) regarding the effect of initialization in our experiments. To address this, we started a new large-scale experiment with non-Gaussian noise (Fig. N2) and added a new experiment with Poisson noise where the initial guess with least-squares is quite suboptimal (Fig. N3).\n\nTo discuss the effect of each step in our convexification procedure, we plan to add a new section to the supplementary, breaking down the steps as suggested by the Reviewer and explaining how they are tested throughout the paper. Here are bullet points describing how we tested each of the following aspects:\n\n1) “altering the training objective to a cross-entropy loss:” In Fig. 4, and the new Fig. N2 that will replace Fig. 4, we are testing the convergence accuracy and speed of algorithms trained on L2 loss vs cross-entropy loss vs weighted cross-entropy (CORNN) loss. \n\n2) “using an approximate pre-computable Hessian to facilitate the use of Newton steps:” In Fig. 4, which will now become a supplementary figure, we tested Newton’s descent with exact Hessian computations on the weighted CORNN loss. \n\n3) “use of a suitable fixed point initialization:” Our new experiments (Figs. N2-N3) will address this point. Though, as noted in the main text, the least-squares solution used for the initialization does not satisfy any additional constraints on the connectivity matrix, thus is not suitable as a valid solution itself.\n\n4) “using teacher forcing (training only one-step losses):” Teacher forcing is at the heart of our work, without which training a network with 10 million parameters under a minute would not have been possible. Fig. 3 shows that BPTT without teacher-forcing is not necessarily better than teacher-forced approach, whereas (old) Fig. 4 shows that FORCE without teacher-forcing does not necessarily have better accuracy.\n\n5) “employing ADMM to enforce constraints (e.g., no self-excitation):” ADMM approach is not related to the speed of the solver, rather the ability to incorporate additional constraints, though this comes with no added complexity.\n"}]

Author rebuttal:
We thank the AC for considering our work and all reviewers for their expert reports, which we used to refine the presentation of our manuscript and add new experiments. We note that two reviewers thought the paper had excellent impact in its area and one thought there were no major concerns. The comments from the fourth reviewer made us realize that the presentation of the manuscript needed to be improved. We acted on all suggestions by making appropriate changes in the manuscript as detailed below.

We are sharing a PDF file with three new figures (Figs. N1-3) as part of the revision. Please note that we have fewer samples than desired in the current version of these figures due to time limitation. If our work is accepted, we will increase the number of samples in all figures. The figure descriptions are provided within the PDF. Below, we list the points we aimed to address with each new figure:

- Fig. N1: A new experiment with correlated noise, for the networks considered in Fig. 7. As requested by Reviewers mPZo and GD24, we wanted to test whether correlated noise would prevent CORNN from training. Though CORNN trains more accurately under comparable independent noise, increasing the number of trials helped mitigating correlated noise.

- Fig. N2: A new experiment with a chaotic network of 5000 neurons. So far, our simulations tested networks with hundreds of neurons due to inefficiency of other algorithms. With this new experiment, we showed how previous algorithms failed to scale.

- Fig. N3: A similar experiment to (old) Figure 4, but with independent Poisson noise. As requested by Reviewer GD24, we tested the accuracy of CORNN under asymmetric non-Gaussian noise distribution.

Moreover, we refined the presentation of the manuscript significantly by incorporating feedback from the reviewers. Specifically, we made the following changes: 

- We restructured the Introduction and Discussion to explain clearly how the rapid inference made possible via CORNN enables real-time interventional experiments, as well as the next steps needed to actualize this experimental scenario. We provided detailed descriptions of these changes in our point-by-point response to the Reviewers wherever applicable.

- We updated the Introduction to emphasize why CORNN is novel and timely. Specifically, we re-wrote the following paragraph:

“Despite recent advances in the dRNN framework, how to perform fast and scalable reconstructions of neural activity traces from large-scale empirical recordings has remained unclear. Notably, the slowness of existing dRNN optimization algorithms may often necessitate the use of high-performance clusters and several days of computation in large-scale experiments. These limitations have been a barrier to the widespread adoption of dRNN approaches. To fully harness the potential of  dRNNs, it is essential to extend their range of applicability from offline analyses to on-the-fly applications within individual recording sessions. Therefore, an approach to training dRNNs in a nearly immediate manner would be a key advancement that would enable novel experimental approaches, such as theory-driven real-time interventions targeting individual cells with specific computational or functional roles (Fig. 1). However, realizing these benefits requires having an optimization routine that is fast, robust, and scalable to large networks.”

Finally, inspired by the reviewer reports, we propose to make the following changes to the title and abstract to refine the presentation: 

Proposed Title: CORNN: Convex optimization of recurrent neural networks for rapid inference of neural dynamics

Proposed Abstract:
Advances in optical and electrophysiological recording technologies have made it possible to record the dynamics of thousands of neurons, opening up new possibilities for interpreting and controlling large neural populations. A promising way to extract computational principles from these large datasets is to train data-constrained recurrent neural networks (dRNNs). Making this training real-time could open doors for research techniques and medical applications to model and control interventions at single-cell resolution and drive desired behavior. However, existing training algorithms for dRNNs are inefficient and have limited scalability, making it a challenge to analyze large neural recordings even in offline scenarios. To address these issues, we introduce a training method termed Convex Optimization of Recurrent Neural Networks (CORNN). In studies of simulated recordings of hundreds of cells, CORNN attained training speeds ~ 100-fold faster than traditional optimization approaches while maintaining or enhancing modeling accuracy. We further validated CORNN on simulations with thousands of cells that performed simple computations such as those of a 3-bit flip-flop or the execution of a timed response. Finally, we showed that CORNN can robustly reproduce network dynamics and underlying attractor structures despite mismatches between generator and inference models, severe subsampling of observed neurons, or mismatches in neural time-scales. Overall, by training dRNNs with millions of parameters in subminute processing times on a standard computer, CORNN constitutes a first step towards real-time network reproduction constrained on large-scale neural recordings and a powerful computational tool for advancing the understanding of neural computation.

The first change in title is emphasizing that it is optimization, not the RNN, that is convex. The second change is regarding the “unit resolution,” which was unnecessarily restrictive since dRNNs could be applied to non-unit resolution activity. Inspired by the comments of Reviewers GD24 and mPZo, the final change emphasizes what exactly we achieved towards a real-time application, i.e. that our work allows rapid inference. Abstract changes are cosmetic in nature to refine the presentation. We would appreciate reviewers' feedback.

Meta-review: The heart of this work is to provide a convex approximation to RNN learning where the learning is reduced in complexity to the point where dynamical systems can be learned online. The primary points are to 1) replace the full cost with a "step-by-step" cost wherein the dynamics are learned between every two pairs of measurements instead of learning a full sequence that matches the data at all points and 2) replace the standard cost with a cross entropy cost. Given these changes the authors show that the dynamics can be learned efficiently even at scale.

The work has many merits in terms of the goals and methods. The main detriments came in the form of confusion brought on by the terminology and presentation, that the authors were able to respond to in the response period. To add to this discussion from the AC point, the "teacher" terminology is confusing since full-FORCE (DePasquale et al. 2018 PLOS ONE) uses a "teacher" network that serves a distinctly different role that also minimizes instability during learning (this work should also probably be cited). There is also some commentary about vanilla RNNs being interpretable, which is a point still up for debate. I believe with the appropriate clarification changes promised by the authors, this work can be a nice addition to the NeurIPS conference and I thus recommend this work be accepted.
Decision: Accept (poster)