SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Fatih Dinc*
Department of Applied Physics
Adam Shai*
CNC Program
Mark J. Schnitzer†
Howard Hughes Medical Institute
Hidenori Tanaka†
Physics & Informatics Laboratories, NTT Research, Inc.
Sunnyvale, CA 94085
Center for Brain Science, Harvard University
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)<sup>1</sup>. 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.
Understanding the relationship between neural dynamics and computational function is fundamental to neuroscience research [1]. To infer computational structure from neural population dynamics,
These authors co-supervised this work.
$^{1}$ CORNN software and reproduction code are available at https://github.com/schnitzer-lab/CORNN-public
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.
Figure 1: Using data-constrained recurrent neural networks for the interpretation and manipulation of brain dynamics within a potential experimental pipeline. This approach centers around online modeling of network dynamics, which can enhance hypothesis testing at the single-cell level and support advancements in brain-machine interface research. The training process is motivated by three objectives: (i) predicting the patterns of neural populations, (ii) revealing inherent attractor structures, and (iii) formulating optimal control strategies for subsequent interventions.
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 microcircuity.
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
Figure 2: CORNN: Convex and scalable solver for dRNNs via ADMM [51]. The CORNN algorithm optimizes the parameters of a recurrent neural network so that activity in hidden units align with activities measured from a ground-truth system, we refer to as a generator. The choice of objective function results in a convex loss landscape. The algorithm starts by finding a fixed point for initialization, where the Hessians of all subproblems are aligned to the correlation matrix of neural activities, which can be pre-computed and pre-inverted. The optimization loop then iterates through predicting the neural activities, computing the prediction error, updating the parameter values, and checking for convergence. The final output is the set of optimized parameters, which represents the functional connections between the neurons in the data-constrained recurrent neural network.
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]:
$$
\tau \frac {\mathrm {d} r _ {i}}{\mathrm {d} t} = - r _ {i} + \tanh  (z _ {i}) + \epsilon_ {i} ^ {\text {c o n}}, \quad z _ {i} = \sum_ {j = 1} ^ {N _ {\text {r e c}}} W _ {i j} ^ {\text {r e c}} r _ {j} + \sum_ {j = 1} ^ {N _ {\text {i n}}} W _ {i j} ^ {\text {i n}} u _ {j} + \epsilon_ {i} ^ {\text {i n p u t}}, \tag {1}
$$
where  $\tau$  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_{ij}^{\mathrm{rec}}$  is the weight matrix for the recurrent connections,  $W_{ij}^{\mathrm{in}}$  is the weight matrix from the input to the recurrent units,  $u_{j}$  is the input vector, and  $\epsilon_{i}^{\mathrm{input}}$  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  $\epsilon_{i}^{\mathrm{con}}$  that accounts for errors and mismatches in the non-linearity. Additionally, the inference model has the constraint  $W_{ii}^{\mathrm{rec}} = 0$  preventing any self-excitation of neurons.
Suitable estimators for reconstructing such a model can be obtained (recursively) by discretizing the differential equation:
$$
\hat {r} _ {t + 1, i} = (1 - \alpha) \hat {r} _ {t, i} + \alpha f (\hat {z} _ {t, i}), \tag {2}
$$
where we define the time scale parameter  $\alpha = \frac{\Delta t}{\tau}$ . Here, the model parameters influence  $\hat{z}_{i,t}$  explicitly linearly, and implicitly through previous time activities, such that  $\hat{z}_{t,i} = \sum_{j}\hat{x}_{t,j}\hat{\theta}_{ji}$ . We define the short-hand notations  $x = [r;u]$  and  $\hat{\theta} = [W^{\mathrm{rec}};W^{\mathrm{in}}]^{T}$  as concatenated matrices. The goal of the estimation problem is to find  $\hat{\theta}$  minimizing a loss function.
A reasonable choice for the loss function of this regression problem is the traditional  $\mathcal{L}_2$  loss:
$$
\mathcal {L} _ {2} (\hat {\theta}) = \sum_ {i = 1} ^ {n _ {\mathrm {r e c}}} \sum_ {t = 1} ^ {T} \left(\hat {r} _ {t, i} - r _ {t, i}\right) ^ {2}. \tag {3}
$$
However, due to the time recurrent definition of the general estimator  $\hat{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  $\hat{r}_{t,i}$  is teacher-forced to  $r_{t,i}$  when computing  $\hat{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  $\mathcal{L}_2$  loss function transforms to a simpler form:
$$
\mathcal {L} _ {2} (\hat {\theta}) = \alpha^ {2} \sum_ {i = 1} ^ {n _ {\text {r e c}}} \sum_ {t = 1} ^ {T} \left(\hat {d} _ {t, i} - d _ {t, i}\right) ^ {2}, \tag {4}
$$
where we define  $d_{t,i} = \frac{r_{t + 1,i} - (1 - \alpha)r_{t,i}}{\alpha}$  and correspondingly  $\hat{d}_{t,i} = \tanh (\hat{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  $\mathcal{L}_2$  loss function, we observe that  $\frac{1 + d_{t,i}}{2}$  follows a linear  $^+$  sigmoid form. As a result, we replace the  $\mathcal{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}\to \pm 1$ , the problem reduces to logistic regression. Instead of a naive replacement, we use a weighted loss of the form:
$$
\mathcal {L} _ {\text {w e i g h t e d}} (\hat {\theta}) = \sum_ {i = 1} ^ {n _ {\text {r e c}}} \sum_ {t = 1} ^ {T} c _ {t, i} \mathrm {C E} \left(\frac {1 + \hat {d} _ {t , i}}{2}, \frac {1 + d _ {t , i}}{2}\right). \tag {5}
$$
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_{t,i}^2]^{-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:
$$
\mathcal {L} _ {\mathrm {C O R N N}} (\hat {\theta}) = \frac {1}{T} \mathcal {L} _ {\text {w e i g h t e d}} (\hat {\theta}) + \frac {\lambda}{2} \| \hat {\theta} \| _ {F} ^ {2}, \tag {6}
$$
where  $||\hat{\theta} ||_F$  is the Frobenius norm of  $\hat{\theta}$  and  $\lambda$  is the regularization parameter.
2.4 Fundamental principle of CORNN: Subproblem Hessian 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  $\mathcal{L}_{\mathrm{CORNN}}(\hat{\theta}) = \sum_{i}\mathcal{L}_{i}(\hat{\theta}_{:,i})$ , where  $\hat{\theta}_{:,i}$  denotes the  $i$ th column of the parameter matrix  $\hat{\theta}$  (See Eq. S1). This means that the original problem can be divided into  $n_{\mathrm{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_{t,i}^2 ]^{-1}$ , leads to a shared (approximate) Hessian for all subproblems (See Eq. (S4)):
$$
H \approx \frac {1}{T} X ^ {T} X + \lambda I, \tag {7}
$$
where  $X$  stands for  $x_{t,i}$  in the matrix form,  $I$  is the identity matrix, and  $\lambda$  is the regularization parameter. We note that the approximation is exact in the limit  $\hat{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
$$
\nabla_ {\hat {\theta}} \mathcal {L} _ {\text {C O R N N}} (\hat {\theta}) = - \frac {1}{T} X ^ {T} E + \lambda \hat {\theta}, \tag {8}
$$
where we define the prediction error matrix  $E_{t,i} = \frac{[d_{t,i} - \bar{d}_{t,i}]}{1 - d_{t,i}^2}$ . 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
$$
A ^ {+} = \left[ X ^ {T} X + T \lambda I \right] ^ {- 1}. \tag {9}
$$
we obtain the following update rule (See Eq. (S20))
$$
\theta^ {k + 1} = A ^ {+} X ^ {T} X \theta^ {k} + A ^ {+} X ^ {T} E ^ {k}, \tag {10}
$$
where we highlight all quantities that can be pre-computed with blue. In essence, the update rule, which we call "Hessian aligned prediction error (HAPE) update," consists of computing the prediction error followed by matrix multiplications; thus is well suited to be accelerated on a GPU if available, but still can be efficiently computed on CPU cores.
The alignment of subproblem Hessians relies on the assumption that we can find a "close-enough" starting position, called the fixed-point, for the initial parameters  $\hat{\theta}^{(0)}$ . As we show in Supplementary Section S1.2, a suitable candidate is the approximate least-squares solution:
$$
\theta_ {\mathrm {l s}} := A ^ {+} X ^ {T} Z. \tag {11}
$$
Then, the initial set of predictions becomes  $\hat{d}_{t,i} \coloneqq \tanh (\sum_{j} x_{t,j} (\theta_{\mathrm{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.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_{ii}^{\mathrm{rec}} = 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} \approx r_{t+1,i}$  for small  $\alpha$ ; 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