Title: CORNN: Convex optimization of recurrent neural networks for rapid inference of neural dynamics
Venue: NeurIPS 2023
Accepted: yes
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

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) 1 . In studies of simulated recordings, 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.

1 Introduction

Understanding the relationship between neural dynamics and computational function is fundamental to neuroscience research [1]. To infer computational structure from neural population dynamics,

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 a steady increase in the numbers of cells that can be recorded and manipulated [2-7]. 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 transformed into a computational model through network reconstruction, provides a promising way of bridging experiments and theory [8]. This approach not only complements recent experimental breakthroughs in targeted stimulation of individual neurons ([9, 10]) for hypothesis testing and brain-machine interfaces but also facilitates the distillation of fundamental computational principles from large amounts of experimental data [1, 11, 12]. Historically, especially for recordings of individual or small numbers of cells, hand-crafted network models were used [13-20]. While valuable, traditional hand-crafted models often struggle to handle the large datasets generated in contemporary neuroscience.

To bridge the gap, a 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-29]. These models, part of a broader suite of computational strategies [11, 21, 30-32], align the dynamics 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 [33]-still a mainstay in systems neuroscience [6]-and provides a viable path to distill underlying computations by extracting the state-space representations of the learned dRNNs [12]. Additionally, dRNNs may open the door to computational control for causal experimentation [9, 10], serving as a valuable computational adjunct to widely used techniques like optogenetics [8, 34].

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 might have been a barrier to the widespread adoption of dRNN approaches. Moreover, 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. 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.

In this work, we present CORNN, a 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-propagation through time (BPTT), which have been employed in previous seminal works utilizing dRNNs [24-29]. Our main contributions include the development of CORNN (Fig. 2), the introduction of an initialization strategy to accelerate convergence (Figs. 2, S4, and S6), and the demonstration of CORNN's scalable speed (Figs. 3, 5, S2, S3, S5, S8, and S9), which enables rapid training on the activity patterns of thousands of neurons (Figs. 4, 5, S7, S8, S9, S10, and S11). CORNN's 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. S1, S6), making it a user-friendly technology for biologists. Lastly, we highlight CORNN's robustness against non-idealities such as mismatches in the assumed dynamical system equations (Figs. 4, 5, S10, and S11), subsampling of neural populations (Fig. 5), differences in dynamical time-scales (Fig. S10), and existence of non-Gaussian (Fig. S3) or correlated noise (Fig. S11).

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 experimental/computational technology. While this work focuses on introducing and validating the fast solver on simulated benchmarks, future work with CORNN should focus on addressing the challenges that arise when analyzing large-scale neural recordings from experimental neuroscience.

RNN

Prepare for

interventions

Single cell interventions

Behavioral experiment

Neural recordings

Real-time extraction

of neural signal

Dynamics reconstruction

with a dRNN model

Simulate single-cell targeting

Desired behavior

This work

Real-time interventions to neural populations

Predicting neural dynamics

Firing rates

3

Time

Revealing attractorstructures

M

Neuralpopulationcontrol

a

2 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 computational research [8]. To date, however, the use of dRNNs has remained offline. Past studies extracted computational principles from neural recordings by fitting dRNNs after data collection [24-29]. Moreover, the use of dRNNs has generally required the expertise of computational neuroscientists, 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 brain-machine interface research is likely a forthcoming application [36-38]. In such use cases, neurobiologists might wish to train dRNNs via a user-friendly approach. Hence, a fast and straightforward dRNN training procedure would enable a new breed of interventional experiments to dissect the brain's microcircuitry.

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 at the single-cell level (Fig. 1). In this scenario, neural activities from mice performing behavioral tasks are captured and extracted in real-time using advanced imaging or electrophysiological technologies and pre-processing algorithms [2, 3, 7, 39-41]. 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 a tight feedback loop between incoming measured data and optimal experimental design, in order to refine the inferred model. Techniques similar to adversarial attacks [42] might be used to devise optimal cell targeting strategies from these dRNNs [8], allowing one to test hypotheses about the computational roles of individual neurons or to identify optimal neurons for use within brain-machine interfaces. Once a 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, a 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 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

L2

Random initialization

Dissimilar

local structure

Similar

local structure

Fixed point initialization &

Hessian inverse computation

Compute error

Update

parameters

Prediction

Generator

RNN

Activities

Optimization loop

from brain-imaging movies [43] (though see [44] for a promising development in smaller scale 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, as 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-50], which chiefly aims to maximize reconstruction accuracy of dynamical systems without regard to biological interpretability, training speed, and scalability.

2.2 Inference model and estimator family of interest

Recurrent neural networks are universal approximators for dynamical systems, whose internal computations can be reverse engineered and interpreted [12, 52]. Thus, we choose leaky firing-rate RNNs as the inference models of CORNN, which follow the dynamical system equations [22]:

where τ is the neural decay time, z i is the total input to the recurrent unit i , r i is the firing rate of the unit i , and tanh( . ) is the pre-defined non-linearity, W rec ij is the weight matrix for the recurrent connections, W in ij is the weight matrix from the input to the recurrent units, u j is the input vector, and ϵ input i is the input noise to the recurrent units. Unlike previous literature, given our goal of reconstruction, we also consider the existence of a conversion noise ϵ con i that accounts for errors and mismatches in the non-linearity. Additionally, the inference model has the constraint W rec ii = 0 preventing any self-excitation of neurons.

Suitable estimators for reconstructing such a model can be obtained (recursively) by discretizing the differential equation:

where we define the time scale parameter α = ∆ t τ . Here, the model parameters influence ˆ z i,t explicitly linearly, and implicitly through previous time activities, such that ˆ z t,i = ∑ j ˆ x t,j ˆ θ ji . We define the short-hand notations x = [ r ; u ] and ˆ θ = [ W rec ; W in ] T as concatenated matrices. The goal of the estimation problem is to find ˆ θ minimizing a loss function.

A reasonable choice for the loss function of this regression problem is the traditional L 2 loss:

However, due to the time recurrent definition of the general estimator ˆ r t,i , 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 of the network. We discuss this in the results below.

Instead of using BPTT, we consider a smaller subset of estimators during training where ˆ r t,i is teacher-forced to r t,i when computing ˆ r t +1 ,i . This turns the potentially infinite time problem into a single time-step one, which we call 'single step prediction error paradigm.' Under this condition, the L 2 loss function transforms to a simpler form:

where we define d t,i = r t +1 ,i -(1 -α ) r t,i α and correspondingly ˆ d t,i = tanh(ˆ z t,i ) becomes an estimator consisting of a linear term through a single non-linearity. However, this simplified version of the loss is not convex and leads to a counter-intuitive global loss function with vanishing gradients (Fig. 2).

2.3 Convexification of the loss function

Rather than minimizing the naive L 2 loss function, we observe that 1+ ˆ d t,i 2 follows a linear + sigmoid form. As a result, we replace the L 2 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 →± 1 , the problem reduces to logistic regression. Instead of a naive replacement, we use a weighted loss of the form:

Here, we use CE to denote the cross-entropy loss function and c t,i are non-negative weighting factors chosen as c t,i = [1 -d 2 t,i ] -1 for the CORNN solver. The specific choice of c t,i , which resembles (but is not) preconditioning [53], is motivated from a theoretical view in Supplementary Section S1.3. 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:

where || ˆ θ || F 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 L CORNN ( ˆ θ ) = ∑ i L i ( ˆ θ : ,i ) , where ˆ θ : ,i denotes the i th column of the parameter matrix ˆ θ (See Eq. S1). This means that the original problem can be divided into n rec sub-problems and solved independently. However, inspired by the fact that solving least-squares for a vector or a matrix target has the same complexity, we can devise a faster method. Specifically, the choice of the specific weighting factors, c t,i = [1 -d 2 t,i ] -1 , leads to a shared (approximate) Hessian for all subproblems (See Eq. (S4)):

where X stands for x t,i in the matrix form, I is the identity matrix, and λ is the regularization parameter. We note that the approximation is exact in the limit ˆ d t,i = d t,i . In other words, we align the Hessian of each subproblem to the correlation matrix of neural activities. Unlike for

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. (7) shows similarities with quasi-Newton methods [53] but is distinct in the pre-computed nature of the Hessian.

Following straightforward algebra, we obtain the gradient as

where we define the prediction error matrix E t,i = [ d t,i -ˆ d t,i ] 1 -d 2 t,i . It is worth noting that each column of the gradient matrix contains the gradient for the scalar subproblem corresponding to the i th output. Then, defining the inverse (fixed-point) Hessian matrix as

we obtain the following update rule (See Eq. (S20))

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 is the approximate least-squares solution:

Then, the initial set of predictions becomes ˆ d t,i := tanh( ∑ j x t,j ( θ ls ) j,i ) , which is subsequently refined through the HAPE iterations. This step is 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. (1). However, CORNN can be applied to the other widely used variant of RNNs, i.e. , the leaky current RNN described in Eq. (S28) and regularly employed in neuroscience literature [12, 26, 52]. 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, i.e. , r = tanh( z ) and subsequently replacing d t,i with r t,i 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 a versatile base model

While Hessian alignment, and subsequent HAPE updates, are the core principles driving the fast solver of CORNN, a versatile base model should be able to incorporate various regularization schemes and sparsity constraints. For example, the addition of constraints W rec ii = 0 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. (1), r t,i ≈ r t +1 ,i for small α ; making the unconstrained problem ill-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 a solver using the alternating direction method of multipliers (ADMM) [51] (summarized in Fig. 2).

In the ADMM solver, the primary problem is divided into two subproblems that are solved individually during subsequent iterations, yet linked through a consensus variable that ensures agreement 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/regularization (See Supplementary Section S1 for derivations). This division allows the

Reconstruction

accuracy (r) of W

rec

0

1

10

10 3

10 5

Training time (s)

1 min

1 hr

1 day

CORNN

GD on logistic loss

GD on L2 loss

FORCE

CPU

GPU

use of fast HAPE updates in the computationally extensive first subproblem, where Hessians can be aligned.

While we only considered the equality constraint in this work, the ADMM framework can be used to add additional L1 and/or nuclear norm regularizations to the learned connectivity matrix [51, 54], and utilize biological priors for further inductive biases. Finally, we note that given the highly irregular nature of the prediction error when d t,i ≈ ± 1 , we supplied the fast solver with a simple automated outlier detection step explained in Supplementary Section S1.5.

3.2 Backpropagation through time on the L 2 loss is suboptimal

In a traditional setting, in which neural activations are hidden and need to be learned through training to provide a correct output, backpropagation through time (BPTT) is necessary to learn long-term dependencies in the data. However, in the network reproduction where neural activations are no longer hidden, it is not clear whether BPTT would be beneficial for the learning, since long-term dependencies are inherently present in the state space of the dynamical system, whose equations are given in Eq. (1). Specifically, given that Eq. (1) has only a single time derivative, the time evolution is performed in a Markovian manner, i.e. , given the current state, the next state can be computed without need for additional information (up to a corruption by random noise).

To test whether BPTT is a suitable training algorithm for network reproduction, we examined randomly connected chaotic recurrent neural networks as a synthetic ground-truth benchmark (See Section S3.1 for details). As shown in Fig. S1, contrary to expectation, teacher-forcing enhanced the convergence speed of the BPTT. Moreover, we also observed higher reproduction accuracy with the cross-entropy loss compared to the traditional L 2 loss, inline with our observation from Fig. 2 that L 2 loss leads to vanishing gradients in large error regime. As expected, BPTT took around a minute even in the toy example of Fig. S1 with 100 time points and 100 neurons, approximately two orders of magnitude smaller than a standard calcium imaging session in both dimensions.

3.3 CORNN decreases training times several orders of magnitude

Having shown that teacher-forcing can speed-up convergence in chaotic networks, we next validated that chaotic dynamics can be reconstructed via a single step prediction error paradigm that CORNN utilizes. Specifically, we initialized randomly connected recurrent neural networks as generators, extracted neural activities for a fixed duration of length T , and trained the dRNNs using CORNN, Newton's solver, gradient descent, and FORCE on varying number of iterations to benchmark their speed and accuracy (See Supplementary Sections S2 and S3.1 for details). All algorithms were able to train, some achieving near perfect accuracies in reproducing the internal connectivity of the generator network (See Figs. 3, S2, and S3). However, FORCE, the current default method in neuroscience literature for dRNN training [24, 25, 28, 29], was 4 orders of magnitude slower (Fig. 3) and needed fine-tuned hyperparameters to barely outperform the fixed-point initialization (Fig. S4).

Whether minimizing the weighted or logistic loss led to more accurate reconstruction depended on the underlying noise distribution (we tested two noise distributions: Poisson (mean) and N (mean , variance) ) of the data as shown in Figs. 3, S2, and S3. However, CORNN outperformed all other algorithms in speed in both cases with several orders of magnitude and converged

in sub-second times even with the CPU implementation, whereas others took up to several days for training (Fig. 3). We also observed increased efficiency for CORNN on a GPU (Fig. S5) and that fixed-point initialization could stabilize other algorithms (Figs. S4 and S6).

3.4 CORNN runtimes scale linearly with data size and polynomially with the network size

To this point, we provided validity evidence for the speed and accuracy of CORNN, or single-step prediction error paradigm in general, via the randomly connected chaotic RNN generators. Next, we focused on RNNs with more structured connectivity matrices. Specifically, we trained a model RNN using BPTT to perform a 3-bit flip flop task following [12] as shown in Fig. S7 (See Supplementary Section S3.2 for details). Fig. S7A shows the output of an example test trial, not used for the generator RNN training or the CORNN reproduction, with 200 time points followed by 100 time points of inter trial interval. We observed that CORNN reproduced network performed the task the same way the original network did; potentially making similar mistakes. Thus, as a first step, we confirmed that CORNN can learn to output the same outputs as the original network.

Next, we looked at underlying neural activities of an example (original) generator and reproduced (CORNN learned) network in Fig. S7B. The reproduction matched not only the output of the network, but also the internal firing dynamics of individual units. Yet, as shown in Fig. S7C, the synaptic connections were not perfectly learned even in such a simple scenario with matching models, which required more trials and potentially interventional data [15]. Thus, we next tested whether increasing the number of trials could lead to better reproduction of the original network and whether CORNN can scale well to fit the increased data and network size.

Fig. S8 shows that accurately reproducing the large networks required significantly more trials. Moreover, the data size requirement for the estimation process increased with the number of units in the generator network, which emphasizes the importance of an optimization algorithm that can scale well to multiple day recordings [55]. Fortunately, CORNN had a near-linear scaling of training times vs . the increasing number of trials and hence can handle large datasets (See Fig. S8). In the other direction, we observed that the training times scaled polynomially (Fig. S9), following n β rec

A

ℳ

G

τ

G

1.

2.

3.

ℳ

G

≠ ℳ

I

τ

G

≠

τ

I

NG

≫

NO

=

NI

: Generator,

: Observed,  : Inference

G

O

I

Observed subpopulation

Inference Model (     )

ℳ

I

τ

I

Neuron ID

100

500

0

300

500

Input window

Output window

0.4

-0.4

Input

Network Output

Target Output

Time (ms)

500

0

Input/Output

(a.u.)

B

Distractor

C

D

CORNN

predictions

Input

weights

-0.05

0.05

-0.025

0.025

r = 0.89

-0.1

0.1

-0.02

0.02

Recurrent

weights

r = 0.45

-0.04

0.04

-0.025

0.025

Output

weights

r = 0.95

Generator network dynamics

CORNN learned 10% subsampled,

time-scale and model mismatched

network dynamics

-0.02

0.02

-0.1

0.1

-0.025

0.025

-0.05

0.05

-0.02

0.02

-0.1

0.1

E

r = 0.82

r = 0.37

r = 0.31

Neuron ID

100

500

0

300

500

Input window

Output window

0.4

-0.4

Distractor

Input

Target Output

Time (ms)

500

0

Input/Output

(a.u.)

Input

weights

Recurrent

weights

Output

weights

with β ∈ [1 . 3 , 2 . 2] , vs. the increasing number of neurons. Bringing both observations together, we concluded that there is an O ( n β rec T ) scaling of the empirical training times.

3.5 CORNN is robust to a diverse set of non-ideal conditions

To this point, we had considered ideal scenarios in which the generator network shared the same dynamical system equations as the inference network, all neurons were observed, and the time-scales of the generator and inference RNN units matched perfectly. Next, we considered cases in which these assumptions are violated (See Fig. 4A). Specifically, we trained a generator RNN with different governing equations (See Supplementary Section S3.3) to perform a timed-response task (See Fig. 4B), which put the network in a limit-cycle like dynamical attractor state initiated by an input. We tested the existence of the underlying attractor through a novel distractor input, which was not present during the training of the task (See Fig. 4B, C). If CORNN was able to learn the underlying state-space geometry, the reproduced network activities should be able to return to their attractive trajectory under novel perturbations.

When we subsampled the generator networks by 10% and introduced jitter in the time scales, we observed that the inference model reproduced the neural dynamics in the observed population in a novel perturbation trial (See Fig. 4C), even though the sub-connectivity matrix was not well reproduced (See Fig. 4D, E). On the one hand, this observation reinforces the findings of previous literature that the connectivity matrices learned by dRNNs should be interpreted as functional connections, accounting for the behavior of networks at the level of neural dynamics, and not synaptic connections [15, 25]. On the other hand, the ability of the mismatched network to reproduce neural activities in a novel perturbation trial provided evidence that the effects of the underlying attractor structures can be robustly reproduced with CORNN despite severe experimental non-idealities. We quantified these robustness aspects in Figs. 5A,B (subsampling), S10 (time-scale mismatch), and S11 (correlated noise). Moreover, we observed the n β rec scaling for the training times with increased number of network size (Fig. 5C), inline with the results of Fig. S9.

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. 1. Imagine an experiment in which mice perform a predefined task several times, e.g ., for an half hour, with imaging of 3000-4000 neurons. This experimental scenario yields roughly ten million parameters to be trained in the dRNN. The current study showcased CORNN's efficacy in training networks of thousands of neurons, requiring just O (10) iterations, each comparable in complexity to gradient computation and taking seconds. Consequently, training such a network from the initial imaging session would take less than a minute. Once trained, the network can enable real-time planning of experimental interventions, testing

A

B

C

Test accuracy (R 2 ) of

reconstructed neural activity

0.6

0.7

0.8

0.9

Subsampling percentage

out of 5000 neurons

1 %

10 %

100 %

Number of

trials

30

50

100

Reconstruction

accuracy (r) of W

in

0.6

0.7

0.8

0.9

Subsampling percentage

out of 5000 neurons

1 %

10 %

100 %

Number of neurons

100

1000

5000

Training time (min)

0.1

1

10

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 decision-making with the CORNN-fitted dRNN needs to be performed. The dRNN starts with the current observed activity and can quickly simulate multiple future responses under different input scenarios, i.e. , potential interventions. Since the simulation involves only iterative matrix multiplications and point-wise nonlinearities, it finishes very quickly-in just a few ms. For instance, running a 1000neuron RNN forward for 10 timesteps ( ∼ a second), for 100 different initial conditions, takes < 6 ms on a GPU, and < 25 ms on a CPU, in our hands. Once a desired intervention is identified from the simulations, the phase mask on a spatial light modulator can be updated in ∼ 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 a system. However, we caution that any claim having to do with attractor structures must be experimentally validated with perturbation experiments that directly test for attractors. In our work, we simulated such experimental validation (Figs. 4 and 5) and found that in the setting tested, CORNN was able to predict the dynamical effects of perturbations on the neural population.

5 Conclusion

In this work, we introduced a fast and scalable convex solver (CORNN) for rapid training of dataconstrained recurrent neural networks. We showed that CORNN is as accurate yet faster than existing approaches by orders of magnitude and can easily scale to the large datasets of today, training in seconds to a few minutes depending on the data, network size and the availability of a GPU. CORNN, as a base model, lays the groundwork for future developments and can be improved by integrating experimentally relevant assumptions and regularizations within the ADMM framework.

When applied to simulated data from structured networks, CORNN picked up the underlying attractor structure despite non-idealities (Figs. 4, 5, S10, and S11). Inspired by this observation, further development of data-constrained RNNs can support systems neuroscience research aiming to understand inter-area interactions and can complement or augment studies of canonical or pairwise correlations toward characterizations of functional connectivity [33]. Moreover, reverse engineering the learned network may provide a deeper understanding of how neuronal populations contribute to emergent computation, cognition, and memory [12, 30]. Finally, with the advent of interventional methods such as single-cell targeted optogenetics [10], dRNNs can be trained to test causal connections computationally and supply theoretical predictions at a scale previously unachievable for single-cell targeted experimentation [8].

This work constitutes a first step towards the application of CORNN to experimental data. However, several steps remain to apply dRNNs in real-time to interventional experiments. Some example steps may include the transformation of calcium traces or spike trains into traces of firing rates normalized within [ -1 , 1] , applying the CORNN solver developed in this work into first offline and then online experimental scenarios, estimation of neuronal time-scales from the experimental data instead of tuning them as hyperparameters, and perhaps implementing a low-rank regularization approach that opens the door to interpreting the observed dynamics in terms of latent variables [26, 52].

References

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.

Tony Hyun Kim and Mark J Schnitzer. Fluorescence imaging of large-scale neural ensemble dynamics. Cell , 185(1):9-41, 2022.

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.

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.

Nicholas James Sofroniew, Daniel Flickinger, Jonathan King, and Karel Svoboda. A large field of view two-photon mesoscope with subcellular resolution for in vivo imaging. elife , 5:e14472, 2016.

Sadegh Ebrahimi, J´ erˆ ome Lecoq, Oleg Rumyantsev, Tugce Tasci, Yanping Zhang, Cristina Irimia, Jane Li, Surya Ganguli, and Mark J Schnitzer. Emergent reliability in sensory cortical coding and inter-area communication. Nature , 605(7911):713-721, 2022.

Nicholas A Steinmetz, Cagatay Aydin, Anna Lebedeva, Michael Okun, Marius Pachitariu, Marius Bauza, Maxime Beau, Jai Bhagat, Claudia B¨ ohm, Martijn Broux, et al. Neuropixels 2.0: A miniaturized high-density probe for stable, long-term brain recordings. Science , 372(6539): eabf4588, 2021.

Matthew G Perich and Kanaka Rajan. Rethinking brain-wide interactions through multi-region 'network of networks' models. Current opinion in neurobiology , 65:146-151, 2020.

Shenqin Yao, Peng Yuan, Ben Ouellette, Thomas Zhou, Marty Mortrud, Pooja Balaram, Soumya Chatterjee, Yun Wang, Tanya L Daigle, Bosiljka Tasic, et al. Recv recombinase system for in vivo targeted optogenomic modifications of single cells or cell populations. Nature methods , 17 (4):422-429, 2020.

Kayvon Daie, Karel Svoboda, and Shaul Druckmann. Targeted photostimulation uncovers circuit motifs supporting short-term memory. Nature Neuroscience , 24(2):259-265, 2021.

David Sussillo, Mark M Churchland, Matthew T Kaufman, and Krishna V Shenoy. A neural network that finds a naturalistic solution for the production of muscle activity. Nature neuroscience , 18(7):1025-1033, 2015.

David Sussillo and Omri Barak. Opening the black box: low-dimensional dynamics in highdimensional recurrent neural networks. Neural computation , 25(3):626-649, 2013.

John J Hopfield. Neural networks and physical systems with emergent collective computational abilities. Proceedings of the national academy of sciences , 79(8):2554-2558, 1982.

Sung Soo Kim, Herv´ e Rouault, Shaul Druckmann, and Vivek Jayaraman. Ring attractor dynamics in the drosophila central brain. Science , 356(6340):849-853, 2017.

Abhranil Das and Ila R Fiete. Systematic errors in connectivity inferred from activity in strongly recurrent networks. Nature Neuroscience , 23(10):1286-1296, 2020.

Mikail Khona and Ila R Fiete. Attractor and integrator networks in the brain. Nature Reviews Neuroscience , pages 1-23, 2022.

Sukbin Lim and Mark S Goldman. Balanced cortical microcircuitry for maintaining information in working memory. Nature neuroscience , 16(9):1306-1314, 2013.

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.

Xiao-Jing Wang. Probabilistic decision making by slow reverberation in cortical circuits. Neuron , 36(5):955-968, 2002.

Hidehiko 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-271, 2022.

Daniel LK Yamins, Ha Hong, Charles F Cadieu, Ethan A Solomon, Darren Seibert, and James J DiCarlo. Performance-optimized hierarchical models predict neural responses in higher visual cortex. Proceedings of the national academy of sciences , 111(23):8619-8624, 2014.

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 neuroscience , 22(7):1159-1167, 2019.

Guangyu Robert Yang, Madhura R Joglekar, H Francis Song, William T Newsome, and XiaoJing Wang. Task representations in neural networks trained to perform many cognitive tasks. Nature neuroscience , 22(2):297-306, 2019.

Kanaka Rajan, Christopher D Harvey, and David W Tank. Recurrent network models of sequence generation and memory. Neuron , 90(1):128-142, 2016.

Matthew G Perich, Charlotte Arlt, Sofia Soares, Megan E Young, Clayton P Mosher, Juri Minxha, Eugene Carter, Ueli Rutishauser, Peter H Rudebeck, Christopher D Harvey, et al. Inferring brain-wide interactions using data-constrained recurrent neural network models. bioRxiv , pages 2020-12, 2021.

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.

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.

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.

Arseny Finkelstein, Lorenzo Fontolan, Michael N Economo, Nuo Li, Sandro Romani, and Karel Svoboda. Attractor dynamics gate cortical information flow during decision-making. Nature Neuroscience , 24(6):843-850, 2021.

Valerio Mante, David Sussillo, Krishna V Shenoy, and William T Newsome. Context-dependent computation by recurrent dynamics in prefrontal cortex. nature , 503(7474):78-84, 2013.

Chethan Pandarinath, Daniel J O'Shea, Jasmine Collins, Rafal Jozefowicz, Sergey D Stavisky, Jonathan C Kao, Eric M Trautmann, Matthew T Kaufman, Stephen I Ryu, Leigh R Hochberg, et al. Inferring single-trial neural population dynamics using sequential auto-encoders. Nature methods , 15(10):805-815, 2018.

Ruxandra Barbulescu, Gonc ¸alo Mestre, Arlindo L Oliveira, and Lu´ ıs Miguel Silveira. Learning the dynamics of realistic models of c. elegans nervous system with recurrent neural networks. Scientific Reports , 13(1):467, 2023.

Ed Bullmore and Olaf Sporns. Complex brain networks: graph theoretical analysis of structural and functional systems. Nature reviews neuroscience , 10(3):186, 2009.

Karl Deisseroth. Optogenetics. Nature methods , 8(1):26-29, 2011.

David Sussillo and Larry F Abbott. Generating coherent patterns of activity from chaotic neural networks. Neuron , 63(4):544-557, 2009.

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.

Riichiro Hira, Fuki Ohkubo, Yoshito Masamizu, Masamichi Ohkura, Junichi Nakai, Takashi Okada, and Masanori Matsuzaki. Reward-timing-dependent bidirectional modulation of cortical microcircuits during optical single-neuron operant conditioning. Nature communications , 5(1): 1-12, 2014.

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):331335, 2012.

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. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems , volume 30. Curran Associates, Inc., 2017. URL https://proceedings.neurips.cc/paper/ 2017/file/4edaa105d5f53590338791951e38c3ad-Paper.pdf .

Andrea Giovannucci, Johannes Friedrich, Pat Gunn, Jeremie Kalfon, Brandon L Brown, Sue Ann Koay, Jiannis Taxidis, Farzaneh Najafi, Jeffrey L Gauthier, Pengcheng Zhou, Baljit S 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.

Hakan Inan, Murat A Erdogdu, and Mark Schnitzer. Robust estimation of neural signals in calcium imaging. Advances in neural information processing systems , 30, 2017.

Aleksander Madry, Aleksandar Makelov, Ludwig Schmidt, Dimitris Tsipras, and Adrian Vladu. Towards deep learning models resistant to adversarial attacks. arXiv preprint arXiv:1706.06083 , 2017.

Hakan Inan, Claudia Schmuckermair, Tugce Tasci, Biafra O Ahanonu, Oscar Hernandez, J´ er ˆ ome Lecoq, Fatih Dinc ¸, 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.

Zhe Chen, Garrett J Blair, Changliang Guo, Jim Zhou, Juan-Luis Romero-Sosa, 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.

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 , 113(15):3932-3937, 2016.

Kadierdan Kaheman, J Nathan Kutz, and Steven L Brunton. Sindy-pi: a robust algorithm for parallel implicit sparse identification of nonlinear dynamics. Proceedings of the Royal Society A , 476(2242):20200279, 2020.

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:1910.03471 , 2019.

Georgia Koppe, Hazem Toutounji, Peter Kirsch, Stefanie Lis, and Daniel Durstewitz. Identifying nonlinear dynamical systems via generative recurrent neural networks with applications to fmri. PLoS computational biology , 15(8):e1007263, 2019.

Ibrahim Ayed, Emmanuel de B´ ezenac, Arthur Pajot, Julien Brajard, and Patrick Gallinari. Learning dynamical systems from partial observations. arXiv preprint arXiv:1902.11136 , 2019.

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.

Stephen Boyd, Neal Parikh, Eric Chu, Borja Peleato, Jonathan Eckstein, et al. Distributed optimization and statistical learning via the alternating direction method of multipliers. Foundations and Trends® in Machine learning , 3(1):1-122, 2011.

Manuel Beiran, Alexis Dubreuil, Adrian Valente, Francesca Mastrogiuseppe, and Srdjan Ostojic. Shaping dynamics with multiple populations in low-rank recurrent networks. Neural Computation , 33(6):1572-1615, 2021.

Jorge Nocedal and Stephen J Wright. Numerical optimization . Springer, 1999.

Jian Yang, Lei Luo, Jianjun Qian, Ying Tai, Fanlong Zhang, and Yong Xu. Nuclear norm based matrix regression with applications to face recognition with occlusion and illumination changes. IEEE transactions on pattern analysis and machine intelligence , 39(1):156-171, 2016.

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.

Zihui Zhang, Lloyd E Russell, Adam M Packer, Oliver M Gauld, and Michael H¨ ausser. Closedloop all-optical interrogation of neural circuits in vivo. Nature methods , 15(12):1037-1040, 2018.

A Emin Orhan and Xaq Pitkow. Skip connections eliminate singularities. arXiv preprint arXiv:1701.09175 , 2017.

Peter J Huber. Robust estimation of a location parameter. In Breakthroughs in statistics , pages 492-518. Springer, 1992.

Brian DePasquale, Christopher J Cueva, Kanaka Rajan, G Sean Escola, and LF Abbott. fullforce: A target-based method for training recurrent networks. PloS one , 13(2):e0191527, 2018.

Haim Sompolinsky, Andrea Crisanti, and Hans-Jurgen Sommers. Chaos in random neural networks. Physical review letters , 61(3):259, 1988.

S1 Details on CORNN derivations and implementation

S1.1 Separation into subproblems and derivation of the gradient and Hessian

The loss function given in Eq. (6) can be written in details as:

where n tot = n rec + n in . We note that for a given i , the functions L i depend only on ˆ β ( i ) := ˆ θ : ,i , i.e. , the i th column of the ˆ θ matrix. In other words, the lack of cross-talk between the columns of ˆ θ in the loss function means each output can be regressed independently with respect to the corresponding part of the loss function, i.e. , L i ( ˆ θ i ) . Thus, for the rest of this section, without loss of generality, we simply considered one instance of the regression, i.e. , fix i , derived the learning rule, and later generalized to the full problem.

The loss function of interest for the subproblem is:

The gradient can be computed as:

The Hessian follows directly from the gradient and the fact that tanh ′ ( x ) = (1 -tanh 2 ( x )) :

At this point, it is worth emphasizing that while x t,i are shared across all subproblems, d t are not. In fact, for any neuron i , a different set of d t,i would contribute to the Hessian of the corresponding subproblem. Thus, in its current form, using the exact Hessian would require repeating matrix inversions for each of the N rec neurons. We performed benchmarking with the exact Newton method, See Supplementary Section S2, though the repeated matrix inversion was no longer feasible for large networks (Fig. 3). To mitigate this, we took a 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 t = tanh( ˆ β T fp x t ) , such that ¯ d t ≈ d t . We discussed in Supplementary Section S1.2 on how to estimate ˆ β fp . Then, we replaced the Hessian with a local approximation ∇ 2 ˆ β L ( ˆ β ) ≈ ∇ 2 ˆ β L ( ˆ β fp ) within the vicinity of ¯ d t ≈ d t , leading to the the local approximate Hessian:

From this equation, we can infer why the choice of c t in the main text was scalable. The local Hessian depends only on the inputs, which are shared across all the n rec independent subproblems, and not on

the predicted or true outputs ( ˆ d t or d t ). 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:

Even though ∇ ˆ β L ( ˆ β ) 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 ∇ ˆ θ L ( ˆ θ ) to solve all subproblems at once. Thus, one can vectorize the computation for each subproblem, i.e. , minimization of L i ( ˆ β ( i ) ) for i = 1 , . . . , N rec , without the need for an external for loop.

At this point, we note that since H fp ⪰ 0 is positive semi-definite (PSD), the direction of the parameter updates is guaranteed to be a descent direction. This is because the update direction has a negative angle with the gradient following ∆ T fp ∇ ˆ θ L ( ˆ θ ) = -∇ ˆ θ L ( ˆ θ ) T H -1 fp ∇ ˆ θ L ( θ ) ≤ 0 by positive semi-definiteness. Thus, even if the error between ¯ d t and d t 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-point updates can be written as:

where γ k is the user-controlled learning parameter chosen either as 1 (to bypass the line search if fixed-point was close enough) or through back-tracking line search. For this work, we simply picked γ k = 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. (11), we performed a Taylor approximation of the gradient and carried out a single update near the vicinity of the putative fixed-point:

where a single-step converges to ˆ β ls = ( X T X + TλI ) -1 X T Z , i.e. , the solution to the least squares problem on the currents Z . Thus, ˆ β ls is expected to be near to both the solution β ∗ and the regime described by the fixed-point Hessian H fp ; 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 T X + TλI ) -1 X T 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 ≈ p :

where -H ( p ) is a constant independent of the optimization variables, and thus can be ignored. Then, the low-error limit of the CORNN loss is:

Here, in the second row, we performed a Taylor approximation around the optimal value ˆ d t ≈ d t + ( ˆ d ′ t ) | ˆ d t = d t (ˆ z -f -1 ( d t )) . In this low error limit, minimizing the scalable weighted loss approximately accounts to minimizing the L 2 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 z = f -1 ( d ) . The observation that low-error limit of the CORNN loss approximates the least-squares provided additional support to the claim that least-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 a general scenario, it is unexpected that all neurons would be connected to each other. Rather, in many cases, one might be interested in introducing a trainable parameter mask such that ∀ ( i, j ) ∈ B ( i, j ) W ij = 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, we will first setup the full problem of interest with constraints, and then review the ADMMequations [51]. To start with, we considered a problem of the following form:

subject to ∀ k ∈ B ( k ) , β = 0

We transformed this problem by adding a new variable ˆ χ , which was constrained to be equal to ˆ β :

subject to ∀ k ∈ B ( k ) , ˆ χ k = 0 , ˆ χ = ˆ β.

Defining the indicator function 1( x ) as

The ADMM loss function became:

minimize

subject to ˆ χ = ˆ β.

(S13)

ADMMLagrangian has two separable loss functions of independent variables, connected through a linear equality constraint, which can be augmented to obtain:

where κ is the dual variable. The ADMM steps (after re-defining κ → (1 /ρ ) κ ) became (See Eqs. (3.5-3.7) in [51]):

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:

where we recall the prediction error matrix:

Then, a single fixed-point update step became:

with ˜ λ = Tλ and ˜ ρ = Tρ . Here, we defined the short-hand notation for the inverse matrix:

Consequently, the update rule for γ k = 1 became:

where we highlighted all quantities that can be pre-computed with blue. For ρ = 0 , this reproduces the update rule for the unconstrained problem in Eq. (10). As before, we promoted the vector ˆ β to a matrix ˆ θ , hence vectorized across the full problem.

S1.4.3 Solving second ADMM iteration via projection

For this step, the goal is to minimize

Fortunately, this problem was perfectly separable across the vector entries of ˆ χ j . For cases where χ k = 0 was constrained, this was the only feasible point and thus the answer to that particular problem. For when j / ∈ B ( j ) , we obtained the minimum loss value by simply picking ˆ χ j = ˆ β k +1 j + κ k j . In short notation, the update rule for this subproblem became:

where we defined Π B as the projection operator to the constraint satisfying subspace ( i.e. , ∀ k ∈ B ( k ) , ˆ χ k = 0 ). Similar to before, we promoted the vector ˆ χ to a matrix to vectorize the full problem.

S1.4.4 Variable initialization and choice of ρ

A priori, it was not clear how ρ should be chosen. However, looking at the Hessian inverse A + , we observed that ρ ≥ λ was an appropriate choice for the step size ρ to have a reasonable impact on the optimization problem. Empirically, we observed that ρ ≫ λ lead to faster convergence.

We initialized the primal and dual variables with χ 0 = ˆ β 0 = ˆ β fp , i.e. the pre-computed least-squares solution, ˆ β fp = A + X T Z , without additional computational complexity. Moreover, given that ADMMensured convergence ˆ χ = ˆ θ , the dual variable κ was bound to converge to non-zero values 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:

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 d t ≈ 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 | E t,i | > v = ⇒ E t,i = 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 error leads to a 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 a 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 | d t | ≥ | 1 -δ | , we projected d t back to ± (1 -δ ) to prevent numerical overflow errors. For experiments in Figs. S1 and S6 with BPTT, we projected the firing rates | r t | ≥ | 1 -δ | to the boundary r t = ± (1 -δ ) for the same reasons. Empirically, this stabilized all solvers and was performed to ensure a fair comparison.

Bringing all together, we state the full CORNN solver in Algorithm 1. We note that for the figures in this work, we simply used a maximum number of iterations (30-100 depending on experiment) instead of monitoring the convergence of CORNN; as CORNN runtimes remained sub-minute and a few more iterations than ideal were perfectly tolerable.

S2 Implementation details for other solvers

S2.1 Back-propagation through time

We performed back-propagation through time for the experiments in Figs. S1 and S6. In our experiments, the teacher forcing means that the neural activities at time t +1 is computed via Eq. (2), but we replace ˆ x t,i in ˆ z t,i = ∑ j ˆ θ t,j ˆ x t,j with the data x t,i , and ˆ r t,i with r t,i . Initially, ˆ x t,i depends on ˆ θ 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. Specifically, we use the following forward-propagation equations:

Algorithm 1: Convex and scalable solver for CORNN via ADMM framework

CORNN.fit ( X,D,λ, nIter = 30 , ρ in = 100) ;

%

X

contains the aggregated input, e.g. firing rates and input units. It is of size

% D contains the targets. It is of size [ T × N rec ] ;

%To start, rescale the regularization parameter λ := Tλ

ρ := ρ in λ % ρ is picked to be comparable to entries of X T X and rescaled λ ;

%Pre-compute the inverses once during initialization;

X + := A + X T ;

Z := f -1 ( D ) %Compute the currents from the targets ;

θ fp := X + Z %Initialize the fixed-point to the ls solution;

%Initialize the primary and dual variables;

θ := θ fp %Initialize to the least-squares solution;

χ := θ fp %Initialize to the least-squares solution;

u := 0 ∗ θ fp %Initialize the dual variable;

for i = 1 : nIter;

Predict the targets ˆ D = f ( Xθ k ) ;

Compute the prediction error E k = ( D -ˆ D ) ⊘ (1 -D ⊙ D ) and perform outlier detection % ⊘ / ⊙ stand for element-wise division/product ;

Update first primal variable θ k +1 := X -θ k + X + E k + ρA + ( χ k -κ k ) ;

Update second primal variable χ k +1 := Π B ( θ k +1 + u k ) ;

Update the dual variable u k +1 := κ k + θ k +1 -χ k +1 ;

Probe convergence by checking || χ k +1 -θ k +1 || , continue if not converged; end return χ

where u is a uniform random number generated between 0 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 [25] for training the recurrent weights to reproduce target neural activities via FORCE. For both network types, we defined the prediction errors as:

Here, ˆ r [ t ] is the predicted firing rate at time point t , vice versa for ˆ x [ t ] . We no longer performed any teacher forcing, thus r [ t -1] and x [ t -1] needed to predict ˆ r [ t ] 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:

Here, e [ t ] was an N rec dimensional vector, rather than a scalar as was the case for the original work [35]. We initialized P [ t ] = λ -1 I , 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 shortcomings of this way of training RNNs using FORCE that render it incompatible with the realities of experiments:

[

T

×

N

tot

]

;

The training procedure requires a particular range of λ to converge (See Fig. S4). Even when FORCE does converge, this is for a very limited range of λ and epochs of training.

FORCE is initially designed for theoretical investigations and requires the training dataset to be continuous. Even subtle jumps in the data results in stability issues with FORCE training, as the premise of FORCE approach is to have small error from the start to the end [35].

Point 1 necessitates low-level hyperparameter tuning. The only work around to point 2 is to reinitialize the inverse covariance matrix P whenever there is a jump in data points. This practically means that if we want to perform multiple training steps on the same data, each epoch simply corresponds to picking a better initialization for the next epoch, whereas all the learned correlation structure, which is stored in the inverse covariance matrix P , needs to be forgotten. The latter point showcases that FORCE is fundamentally ill-suited to many aspects of experimental data, for example, when multiple imaging sessions need to be combined.

S2.3 Newton's method and gradient descent with single-step prediction error

We implemented Newton's method by using the exact Hessian given in Eq. (S5) for two different scenarios: i) when c t = 1 corresponding to minimizing cross entropy loss, and ii) when c t = [1 -d 2 t ] -1 corresponding to the weighted loss. When implementing ii), we performed the outlier detection not only on the gradient, but also the Hessian level; since the division by [1 -d 2 t ] could have destabilized the training. Given that each subproblem had a unique Hessian for the Newton descent, we solved the separate subproblems in parallel using multiprocessing tool in numpy.

To implement the gradient descent, we used a single-layer Pytorch model with a linear + tanh structure. We implemented both a CPU and GPU version, and used the ADAM and SGD optimizers. Gradient descent was implemented for two types of loss functions: i) loss function, and ii) cross-entropy with c t = 1 .

S3 Benchmarking details for generator RNNs

S3.1 Randomly initialized chaotic RNNs

When benchmarking dRNNs on randomly initialized chaotic RNNs, we picked the recurrent weights of the generator randomly from a normal distribution W rec ij ∼ N ( 0 , g 2 n rec ) with zero mean and standard deviation of g √ n rec . For g > 1 and sufficiently large n rec , these networks can sustain their activities without any external stimulus and exhibit chaotic behavior [60]. For these chaotic generators, we picked g = 3 throughout this work to ensure the networks were deep into the sustainable spontaneous activity regime. We chose the chaotic regime specifically, since previous work established that teacher-forcing when training an RNN to output a function can have potentially harmful consequences to the stability of the network training [35]. We aimed to test whether this generalized to the reproduction task we considered in our work.

We note that in figures with randomly initialized chaotic RNNS, where we compared CORNN to other optimizers, we omitted the zero self-excitation constraint, i.e. , W rec ii = 0 , to remain conservative in our speed comparisons. This is because 2nd order methods (FORCE and Newton descent) cannot incorporate the zero self-excitation constraint without additional optimization steps since projected 2nd order steps are not necessarily descent directions, whereas projected gradient descent is notoriously slower compared to the unconstrained version. In contrast, CORNN was approximately equally fast with or without the ADMM solver; since both ADMM and HAPE updates converged in more or less the same number of iterations.

S3.2 Training of RNNs for the 3-bit flip flop task

To benchmark CORNN on diverse conditions, on top of the randomly connected RNNs discussed above, we established ground truth data by training RNNs on a 3-bit flip flop task. The architecture of the generator RNN was as follows:

where r ( t ) is the hidden state of the RNN at time t , α is the time-constant associated with the dynamics, W in is the weight matrix governing input into the RNN, u is the input into the RNN, W rec are the recurrent weights, o is the output of the RNN and W out are the weights which take the hidden state to the output. Since CORNN does not learn W out , to compute the outputs for the CORNN learned networks in Fig. S8, we used the same projection matrix W out as the original network. See the accompanying code for further details on network training and benchmarking experiments.

S3.3 Training of RNNs for the timed-response task

To test robustness of CORNN to mismatches between generator and inference network dynamics, we established ground truth data by training leaky current RNNs on a timed-response task with the following architecture:

Here, unlike the CORNN and 3-bit flip flop equations (See Eq. (S27)), the time-derivative was at the level of currents ( x ), not firing rates ( r ). Upon receiving an input during t ∈ [0 , 100] ms , we trained 10 networks to output a Gaussian pulse centered at t = 500 ms with standard deviation 20 ms ; otherwise zero everywhere.

During test trials in Figs. 5, S10, and S11, we introduced a novel distractor input ten times stronger than the cue at t ∈ [400 , 410] ms . To allow fair comparison between networks trained on different numbers of neurons, we computed the R 2 of the neural activities from the first 100 neurons across three 1 s trials and performed a uniform average. For the illustrative Fig. 4, to allow visualization of finer details in time-activities after distraction, we instead plotted a scenario with an earlier and weaker (half strength than the cue) distractor at t ∈ [200 , 210] ms . To obtain the correlated noise in Fig. S11, we first sampled random noise from a normal distribution, then convolved the i.i.d. noise with a 2D Gaussian kernel across time points (s.d. of 3 ms ) and neurons (s.d. of 5 neurons), and re-scaled by the targeted standard deviation. See accompanying code for further details.

S3.4 The mismatch between leaky-firing and leaky-current RNNs

In this section, we briefly discuss the extent of generator mismatch between leaky-firing and leakycurrent networks. To see this, we consider the dynamics of the currents ( z ) under the leaky-firing rate RNNs:

Although the current dynamics look similar to the leaky-current RNNs, there are several key differences. Firstly, the term τW in d u ( t ) d t is added to the leaky-current equation, which can act as an unobserved current when the input stimuli change. Second, the definition of the firing rates are different, with the firing-rate RNNs performing a low-pass filter in time on the linear + non-linear term, f ( z ) , (hence the name, leaky firing-rate RNN). Finally, under non-homogeneous time-scales for the leaky firing-rate RNN, τ would promote to a diagonal matrix and would not necessarily commute with W rec , so that the Eq. (S29) no longer holds, leading to a stronger mismatch.

S4 Benchmarking CORNN's speed and scalability

To obtain the CORNN's solver in Algorithm 1, we took five distinct steps:

Single-step prediction error paradigm , where we teacher-forced ˆ r t,i , the RNN unit activities at time t , to the observed neural activities, r t,i , to predict r t +1 ,i , with a single-step backpropagation.

Convexification of the loss function , where we replaced the L 2 loss function with a crossentropy loss.

Hessian alignment to neural activity correlations , where we weighted individual samples to align subproblem Hessians to a pre-computable quantity, i.e. , the correlation matrix of observed neural activities.

Initialization to a fixed-point , where we used the approximate least square solution to initialize the networks.

Using ADMM approach to enforce constraints without increasing algorithmic runtimes.

In this work, we performed several benchmarking experiments to delineate the contribution of each step we take towards the final convex solver to the scalability and speed of CORNN. We proved, with targeted experiments, the need for each of these steps, as we discuss below.

S4.1 Single-step prediction error paradigm

As discussed in Section 2.2, the first step we took was to restrict the family of interest for the inference models to the models that can perform single-step predictions well. This simplifying assumption mapped the recurrent estimation problem to a single layer feed-forward network with a sigmoid/tanh non-linearity. Here, the inputs of the network are u t,i and r t,i , whereas the outputs are r t +1 ,i . Since this step eliminated the need for the back-propagation through time, both the training times and memory requirements decreased significantly (Compare Figs. S1 and S2).

However, a priori, it was not clear whether the restrictive assumption of single-step prediction error minimization would lead to decreased performance. To test this, we performed the experiments in Fig. S1 and observed that teacher forcing increases the efficiency of the training, and was not deleterious. A potential explanation may lie in the observation that, unlike the traditional scenario where a sparse error signal needs to be propagated back in time to perform credit assignment, network reproduction training had access to rich variety of error signals at each time step and for each neuron.

S4.2 Convexification of the loss function

As discussed in Section 2.3, the convexification of the loss function plays an important role in speeding up the learning process. However, by itself alone, convexification is not the reason of the speed improvement. For example, in Fig. 3, gradient descent on L 2 loss converges faster than the gradient descent on the logistic loss. On the other hand, for smaller networks with 200 neurons (Figs. S2, S3, S5), convex loss function merits the use of Newton descent and leads to faster convergence than gradient descent. Yet, as shown in Figs. S2 and S5, CORNN is at least an order of magnitude faster even for these small networks. Hence, though convexification of the loss function is crucial, it cannot explain the speed or the scalability of CORNN by itself.

S4.3 Hessian alignment

The main contribution of our work is to pick c t,i , which approximates the problem as a local leastsquares as discussed in Supplementary Section S1.3, and aligns the Hessians of each subproblem to a pre-computable quantity, i.e. , the correlation matrix of neural activities, as discussed in Section S1. This alignment allows the use of HAPE updates that have complexity of gradient descent but converge in O (10) iterations thanks to the approximate Hessian. Figs. S2 and S5 show that Newton's descent with exact Hessian computations on the weighted CORNN loss is an order of magnitude slower than CORNN. Hence, the Hessian alignment was a cruical step to speed up the learning, but most importantly allow scaling to large networks due to its gradient-level complexity.

S4.4 Fixed-point initialization

The fixed-point initialization relies on an approximate solution when the final error is low (See Supplementary Sections S1.2 and S1.3). This initialization is a good first guess, since, if there was no conversion noise at all, one could have simply performed least-squares on the currents estimated via z = arctanh ( r ) . This is akin to regressing to the inverse of the output before the non-linearity in a single layer network. Though such an approach can provide a good first guess, it is not robust to conversion noise. This was shown in Figs. 3 and S3, where non-Gaussian noise components lead to low accuracy for the initialization. Finally, the least-squares solution does not satisfy any additional

constraints on the connectivity matrix, thus is not suitable as a valid solution itself. Yet, since the least-squares solution is computationally negligible, we used it as an initialization for CORNN.

S4.5 Alternating direction method of multipliers

We chose to use ADMM approach to enforce the no self-excitation constraints and present a versatile base model that can incorporate additional design choices such as L1 and/or nuclear norm regularization. In our comparison experiments, we omitted any constraints since other solvers were considerably slower, or outright incompatible. However, in our speed benchmarking results in Figs. 5, S8, and S9, we used the full CORNN solver with the ADMM described in Algorithm 1. The choice of ADMM for enforcing constraints is rooted in our observation that both the ADMM steps and the original HAPE updates take O (10) iterations to converge. Moreover, other parts of the ADMM updates are negligible compared to the primal problem, where the HAPE updates take place. In other words, the fact that HAPE updates and ADMM converge in similar steps creates a mutual synergy that increases accuracy and promotes versatility in our work without hurting convergence times.

S5 Reproducibility

We provided the code to reproduce each figure in our Github. All experiments, except for those in Figures 3, 5, and S4, were run on a standard computer with Geforce RTX 3090 Ti GPU and Intel Core i9-9900X Skylake X 10-Core processor. Experiments in Fig. 5 were run on a computer with two Geforce RTX 3080 Ti GPUs (though using only one) and Intel Core i9-10980XE 18-Core processor. Experiments in Figure S4 were divided to run on both of the aforementioned computers. Experiments in Figure 3 were run on a computer with two NVIDIA Geforce RTX 4090 (though using only one) and an AMD Ryzen Threadripper PRO 5975WX 32-Core processor (though using only 5 cores due to memory limitations). Illustrative figures are plotted on an Apple Macbook Air with M1 Chip. See the accompanying code provided in the Github repo 2 for the specific hyperparameters used for the experiments.

A

Cross entropy loss

L2 loss

Reconstruction

accuracy (r) of W

rec

0.2

0.8

Fixed point

initialization

No teacher

forcing

25% teacher

forcing

50% teacher

forcing

75% teacher

forcing

100% teacher

forcing

Runtime per experiment is 47 ± 1 seconds

B

Reconstruction

accuracy (r) of W

rec

0.2

0.8

Runtime per experiment is 47 ± 1 seconds

Cross entropy loss

L2 loss

Fixed point

initialization

No teacher

forcing

25% teacher

forcing

50% teacher

forcing

75% teacher

forcing

100% teacher

forcing

Training time (s)

0.1

1

10

0.1

1

10

1

0.9

Reconstruction

accuracy (r) of W

rec

Reconstruction

accuracy (r) of W

rec

A

B

0.97

0.93

CPU implementations

CORNN

Newton descent - weighted

cross entropy

Newton descent - cross entropy

Gradient descent (ADAM) -

cross entropy

Original FORCE applied

to currents (z)

Modified FORCE applied

to firing rates (r)

Training time (s)

Training time (s)

Reconstruction

accuracy (r) of W

rec

CPU implementations

CORNN

Newton descent - weighted

cross entropy

Newton descent - cross entropy

Gradient descent (ADAM) -

cross entropy

Original FORCE applied

to currents (z)

Modified FORCE applied

to firing rates (r)

1

100

0.01

1

0

Figure S3: CORNN is robust to non-Gaussian noise injections. Same as in Figure S2, but with parameters: α = 0 . 1 , n rec = 200 , T = 1500 , No input. ϵ conv ∼ Poisson (10 -1 ) , ϵ input ∼ Poisson (10 -2 ) .

Reconstruction

accuracy (r) of W

rec

Reconstruction

accuracy (r) of W

rec

A

B

Fixed point initialization only

Original FORCE applied

to currents (z)

Modified FORCE applied

to firing rates (r)

0

0.5

10 -4

10 -2

1

100

Regularization level (a.u.)

Regularization level (a.u.)

10 -4

10 -2

1

100

CORNN

1

0

0.5

1

Training time (s)

0.1

1

10

Reconstruction

accuracy (r) of W

rec

Reconstruction

accuracy (r) of W

rec

A

B

0.97

0.93

GPU/CPU implementations

CORNN (GPU)

Newton descent - weighted

cross entropy (CPU)

Gradient descent (ADAM) -

L2 (GPU)

Gradient descent (ADAM) -

cross entropy (GPU)

Original FORCE applied

to currents (z)

Modified FORCE applied

to firing rates (r)

0.9

1

0.1

1

10

Training time (s)

A

Reconstruction

accuracy (r) of W

rec

Fixed point

initialization

No teacher

forcing

25% teacher

forcing

50% teacher

forcing

75% teacher

forcing

100% teacher

forcing

Runtime per experiment is 56 ± 2 seconds

B

Reconstruction

accuracy (r) of W

rec

Fixed point

initialization

No teacher

forcing

25% teacher

forcing

50% teacher

forcing

75% teacher

forcing

100% teacher

forcing

Cross entropy loss

L2 loss

0.8

0.7

0.6

0.3

Runtime per experiment is 56 ± 2 seconds

C

CORNN

Newton descent - weighted

cross entropy

Newton descent - cross entropy

Gradient descent (ADAM) -

cross entropy

CORNN (FP initialized)

Random initialization

Training time (s)

0.1

1

10

Reconstruction

accuracy (r) of W

rec

100

0

1

Input and output for the first trial

fi

Trace of neural activations

Neurons

Time step

Original

CORNN learned

Weights

Input

Recurrent

Original

Learned

Learned

Time step

CORNN

network

Original

network

D

0

0.5

10

10

0.0

20

20

0.5

0.5

0.0

0.5

30

30

0.1

0.0

40

40

0

100

200

30

0

100

200

300

0.1

0.0

0.1

input

output

25

150

175

200

Network Output

Correct Output

200

250

300

Network Output

Correct Output

200

250

300

Number of trials

100

1000

One

day

One

month

Test accuracy (R 2 ) of

reconstructed output

1

0.5

Reconstruction

accuracy (r) of W

rec

0.2

1

Training time (min)

100

1

0.01

Number of

neurons

100

1000

3000

CPU implementation

1

0.5

GPU implementation (64 bit)

100

1

10 -4

Number of trials

100

1000

GPU implementation (32bit)

1

0.5

0.2

1

0.2

1

100

1

0.01

10 -4

Number of trials

100

1000

One

day

One

month

slopes

0.98

0.96

0.92

slopes

0.94

0.92

0.89

slopes

0.49

0.64

0.58

0.01

Number of neurons

100

1000

100

1

10 -4

0.01

Training time (min)

3000

A

B

Number of neurons

100

1000

100

1

10 -4

0.01

Training time (min)

3000

Implementation

CPU

GPU (32 bit)

GPU (64 bit)

Slope

1.36

1.44

1.96

Slope

2.17

1.85

2.07

Uncertainty (s.d.) in time-scales (α)

A

B

1

0

-1

-2

-3

Test accuracy (R 2 ) of

reconstructed neural activity

0.01

0.1

1

Number of

trials

30

50

100

α ~ N(0.1,[s.d.]

2

)

Uncertainty (s.d.) in time-scales (α)

Reconstruction

accuracy (r) of W

in

0.01

0.1

1

0.2

0.8

C

Uncertainty (s.d.) in time-scales (α)

Reconstruction

accuracy (r) of W

rec

0.01

0.1

1

0.2

0.4

Test accuracy (R 2 ) of

reconstructed neural activity

Standard deviation

of the injected noise

10 -3

10 -2

Number of

trials

30

50

100

Independent noise

Correlated noise

Type of

noise

10 -1

1

0

Unsuccessful reproduction level

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\n- 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.\n\nLIMITATIONS:\n1. 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-27"}

comments: SUMMARY:
The paper introduces a new method for training RNNs that allows fast convergence which could make realtime parameter update possible.

STRENGTHS:
- The method conceptually is simple and intuitive.
- It tackles an important problem which is the time consuming nature and computing intensivity of ML models
- It achieves impressive results through simple hacks
- One possible use of this is also quick experimentation with RNN models
- 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). 
- 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.
- I would have appreciated results experimenting on real data and comparing neural dynamics computed via this method compared to other methods

QUESTIONS:
NA

LIMITATIONS:
NA
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-03

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n* 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.\n\nWEAKNESSES:\nBelow 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.\n\nQUESTIONS:\n* 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?\n\nLIMITATIONS:\nSee weaknesses section", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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\nSTRENGTHS:\nThe 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\nWEAKNESSES:\nThe reasoning about the enforced use of CPU does not sound convincing to me. See the question part.\n\nQUESTIONS:\non 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.\n\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-07"}