Modeling Neural Activity with Conditionally Linear Dynamical Systems

We are very grateful to the reviewer for their detailed and insightful comments, and overall positive review. You and other reviewers found merit in the conceptual advance, found the paper well written, and suggested better outlining and testing the scopes of applicability. This provides valuable feedback towards improving our contribution.

Thank you for pointing out the scope of the experiments conducted so far in the paper, which we hope to highlight model recovery and inference on real-world experiments. Your feedback on more challenging experiments targeting the model assumptions is echoed by reviewers and is a valuable concern. To address this overall weakness, we introduce below a new experiment that readily fits within the experiments considered so far while targeting the idealized conditions of low dimensionality and model alignment. Furthermore, we point you to our answer to Question 4 of Reviewer pLGj, which details new numerical calculations of the empirical covariance in $x$ given $u$. As per Appendix A.2., this provides a diagnostic feature and proxy for our approximation error to any true nonlinear dynamics in the system.

Mechanistic model of ring-attractor dynamics

Inspired by reviewer feedback, we provide an alternate synthetic ring attractor experiment based on the neuroscience literature—see review by Hulse and Jayaraman (2020, “Mechanisms Underlying the Neural Computation of Head Direction”). The intention with this additional task is to test the CLDS and its inference in a synthetic data modality, where we know the underlying computation, but which, importantly, is not generated from a CLDS model, to test the model in mispecified settings.

Model description: We write a model of continuous ring attractor dynamics with bumps of activity integrating angular velocity (Zhang, 1996, “Representation of spatial orientation by the intrinsic dynamics of the head-direction cell ensemble: a theory”); a model and computational account of spatial orientation shared across many species. In this model, attractor dynamics are implemented through head-direction preferential units arranged in a “ring”, with short-range excitation and long-range inhibition forming a bump of activity for a specific unit in a specific head-direction, with a velocity integrator causing a shift in the location of the bump as the head-direction changes. Concretely, the dynamics of the new generative model follow

$\frac{1}{\tau} \dot u(\phi, t) = - u + f(w * u) + v \frac{\partial}{\partial \phi} u(\phi, t).$

Here, $u(\phi, t)$ represents the network activity of cells with preferred head-direction $\phi$ at time $t$, $f$ is a (ReLU) nonlinearity, $w$ is a “Mexican-hat” convolution filter, and $v_t = \dot \theta_t$ is the instantaneous difference in head-direction $\theta_t$. For implementation purposes, $u$ is discretized to $u_t \in \mathbb{R}^N$, $t \in \{1, \dots, T\}$, with each unit $u_i$ having preferred directions $\phi_i$ regularly spanning the interval $[0, 2\pi)$, and $w * u$ is implemented with a circulant matrix $Wu_t$. We further include additive Gaussian random noise to make the process stochastic. We use this activity $u$ as our observations, of dimension $N=32$. The head direction $\theta_t$ still acts as our conditions, sampled the same way as in our previous synthetic HD model.

Results: We know the underlying model captures high-dimensional (N=32 here) ring attractor dynamics. When fitting a low-dimensional CLDS model, we expect the latent dynamics to capture the ring-attractor structure and the nonlinearity of the dynamics, with a linear emission model bringing it back to the high-dimensional activity. This is precisely what we found when fitting the CLDS—all necessary figures to relay the results will be included in the final revision. After a search over kernel hyperparameters, we found the best-performing model to encode a ring of fixed points that closely resembles the synthetic HD fixed points, and with eigenvalues closer to the ones observed from the CLDS fits on the mouse ADn recordings. The results make us confident that the CLDS can capture this nonlinear structure even under model mismatch.

Response to individual weaknesses

(1) Spike data: We agree that the extension to additional noise models (e.g. Poisson) is of interest, and we intend to follow up soon with a full implementation of this. There are no technical hurdles in principle—we just have not prioritized this particular direction. In our experience, simple linear Gaussian methods perform well and are popular among practitioners in neuroscience. For example, Gaussian Process Factor Analysis (GPFA) uses a Gaussian noise model.

We agree with the reviewer that this extension would enhance the paper. However, we think the core conceptual advances of the paper do not rely on the noise model. For example, we think most neuroscientists would regard the head direction system as an inherently nonlinear system. Here, we show that one can model it using linear methods (assuming that head direction is simultaneously measured with neural activity).

(2) EM Scaling: For now, we are focused on obtaining point estimates of model parameters — i.e. the functions $A(u), b(u), C(u), d(u)$. We mentioned MCMC and variational inference as possible extensions for obtaining an approximate posteriors over these parameters/functions. In contrast, EM only provides point estimates (the maximum of the posterior). In other words, this is an apples-to-oranges comparison. Relative to EM, we would expect that MCMC to be much more expensive (since it involves sampling) and variational inference to be moderately more expensive (since it involves an optimization problem with additional parameters).

We think that EM is scalable for problems of interest in neuroscience. Extremely high-dimensional latent spaces would be susceptible to overfitting anyways. Furthermore, it is not clear that alternative methods are more scalable. In particular, any method that computes the marginal log likelihood of $y_1, \dots, y_T$ integrating over $x_1, \dots, x_T$ will perform the Kalman filtering step anyways, which accounts for a large chunk of the computation time in EM. In particular, the M-step is very cheap. In principle, it is possible to do gradient ascent on the marginal log likelihood, but then one must tune learning rates and in general take much smaller steps in parameter space.

In practice, we believe the biggest challenge for scaling will be datasets of very long duration (large $T$). Here, it is possible to leverage parallel computation on a GPU to perform the Kalman filtering, reducing the scaling behavior from $T$ to $\log T$. See Sarkka & Garcia-Fernandez (2020, “Temporal parallelization of Bayesian smoothers”)

(3) Kernel function: Draws from a GP with an RBF kernel are infinitely differentiable. A Matern kernel would be a weaker prior (leading to less smooth functions). This would introduce additional hyperparameters and lead to a complex model selection process. For example, the kernel hyperparameters are known to be unidentifiable in a simple regression setting. See Zhang (2004, “Inconsistent estimation and asymptotically equal interpolations in model-based geostatistics”).

A CLDS model can account for abrupt task switches or transitions if they are associated with a measurable behavioral event. For example, we show this in our analysis of the nonhuman primate reaching task by selecting $u_1, \dots, u_T$ to be a step function centered at the time of movement onset.

If abrupt task switches or transitions are unobserved (i.e. latent variables), then one could consider a switching linear dynamical system. We discuss the complementary challenges and advantages of this alternative model in the paper. In particular, inference in this model is less straightforward because Kalman filtering and smoothing is not possible. Speaking from personal experience, we found that open-source implementations of these models were quite difficult to tune in practice.

(4) Input-driven DS: The most relevant prior work we are aware of is by Foerster et al. (2017, “Input Switched Affine Networks: An RNN Architecture Designed for Interpretability”) and more recently the SSM models like Mamba (Gu & Dao, 2024, “Mamba: Linear-Time Sequence Modeling with Selective State Spaces”). Neither of these papers explicitly mention attractor dynamics or ring attractors specifically.

The main difference is that these models are applied to natural language processing tasks and therefore take in discrete tokenized inputs. The CLDS model we propose takes inputs over a continuous space and furthermore incorporates a smoothness prior over this space. Thus, in the head direction example, our model supports an exact ring attractor. We think that an input-switched affine network could in theory learn to approximate a ring attractor if the head direction variable were binned onto a discrete number of values. However, (1) it does not incorporate noise in the latent space, a crucial component of SSMs and probabilistic modeling, and (2) the model would have a hard time learning this from a random initialization because it has no smoothness prior and therefore no inductive bias towards learning similar representations for nearby heading directions.

Thank you again for your time; we look forward to your final assessment.

We are very grateful to the reviewer for their time and insightful questions. We are pleased to see your overall positive assessment of our work!

Weaknesses

The point on modern state-space models is a great one. We had mentioned input-switched affine networks in the first submission, but missed Mamba and its related literature. Our revision will include a longer discussion emphasizing the following:

• Presence of noise: SSM models developed in the ML community do not model noise in the latent space. This is thought to be important when analyzing neural data due to limited dataset sizes and noisy observations. Latent noise complicates the inference process and common strategies for inference include EM and marginal optimization for tractable models and variational inference or MCMC for intractable models.
• Structure of the state-space: SSM models developed in the ML community are often applied to natural language processing tasks and therefore take in discrete tokenized inputs. The CLDS model we propose takes inputs over a continuous space and furthermore incorporates a smoothness prior over this space, which is an important inductive bias. Thus, in the head direction example, our model supports an exact ring attractor. Furthermore, our work shows how input-selective models can be used to model continuous attractors (e.g. ring attractors for heading direction). Given the fundamental importance of attractor dynamics to modern theories of neural circuit dynamics, we think this is an important conceptual advance which cannot be found in existing ML literature.
• Models of nonlinearity: Finally, we can comment on the GP v.s. MLP modeling. GPs importantly carry a prior over functions, which is of great help in inferring the non-parametric forms of the coefficients in the low-data regimes often seen in neuroscience (as mentioned above). Furthermore, this same prior makes for the potential to obtain full posterior estimates over these coefficients modeled with GPs—we discussed in the paper how this could be done in Lines 127-128.

Questions

Addressing your questions one-by-one:

• Dimensionality and GP trade-off: An important aspect of this tradeoff is that the CLDS is linear in dynamics ($D$-dimensional) and is nonlinear over conditions (with degree of nonlinearity enforced by $L$). Hence, while both trade off in final expressivity, they do so in complementary ways that model predictions (in time v.s. over different conditions) can help further parse out. Second, on the topic of overfitting, we typically treat $L$ as fixed to a high-enough value to allow sufficient expressivity by the GP prior, as opposed to an optimized hyperparameter—see also our related comments to Reviewer pLGj on this topic. The dimension $D$ is typically variable and searched over, but selected through cross-validation and restricted to be low-dimensional, both avoiding overfitting.

• Nonlinear biases: We agree with your intuition here. It is worth noting that the specific model mentioned (standard LDS + nonlinear $\mathbf{b}(\boldsymbol{u}_t)$) is a subcase of the more general CLDS with both $\mathbf{A}(\cdot)$ and $\mathbf{b}(\cdot)$ nonlinear: while the GP prior implicity favors against that, one could learn a degenerate set of weights emphasizing the constant function for $\mathbf{A}(\cdot)$, making it constant over $\boldsymbol{u}$. Still on the topic of biases, one can bypass the dynamics entirely by using $\mathbf{d}(\boldsymbol{u}_t)$ to essentially fit a nonlinear tuning curve or PSTH to each neuron (disregarding how single-trial dynamics vary around these trial-average firing rates). We think it is feasible and interesting to see how all of these “ablations” of the CLDS model perform on various experimental datasets, in particular the nonhuman primate reaching dataset.

Thank you again for your time; we look forward to your final assessment.

We are grateful to the reviewer for their time and constructive comments. We are pleased to see your overall positive assessment of our work! Below are the answers to your questions.

Weaknesses

The CLDS model is a statistical description of condition-dependent high-dimensional time series with latent structure, primarily motivated by and dedicated to the analysis of neural dynamics. We would not consider the CLDS to be “brain-inspired” in the sense of a biophysically realistic circuit model.

Regarding the soundness of the approach besides experimental results, we would point to the fact that the CLDS model is a natural extension of a lot of existing, influential models. In particular, we would point to past work using linear network models of neural circuits (see citations [4-8] in the paper) and more recent work that develops nonlinear dynamical systems models (e.g. citations [9-14] in the paper). Our central argument is that CLDS models provide some of the nice benefits of both.

In practice, neuroscientists are often interested in finding the relationship between experimental task variables or observed behavioral covariates with neural data. Common models (e.g. LFADS) are fit either in an entirely unsupervised manner or incorporate this task and behavioral covariates in a highly complex, nonlinear manner. CLDS models strike a nice balance between interpretability and expressivity.

Questions

Answering your questions in order

1. There is a very broad variety of statistical models for neural data analysis that are targeted at different specific circumstances. While it is not easy to fit all of these methods into a table, we think the most important advance of our paper is to show that simple and interpretable models (i.e. linear) can be tweaked to model complex (i.e. nonlinear) neural circuit dynamics.
2. $\mathbf{A}(\boldsymbol{u}_t)$, $\mathbf{b}(\boldsymbol{u}_t)$, $\mathbf{C}(\boldsymbol{u}_t)$, and $\mathbf{d}(\boldsymbol{u}_t)$ are all nonlinear functions of $\boldsymbol{u}_t$ and belong to the nonparametric family of functions determined by a Gaussian Process (GP). To your question, with the coefficients framed as GPs, they allow for a nonlinear dependency on the conditions, making the system linear in dynamics and nonlinear in conditions. The amount of nonlinearity introduced depends on the kernel and its parameters. For example, if we use RBF kernels, the length scale determines the probability of how much high-frequency content (or wiggliness) is in the function. This provides an appropriate mechanism for controlling the level of smoothness either using some prior knowledge or in a data-driven manner.
3. Certainly! We comment on how this can be done on lines 163-170 in the submitted manuscript.

Finally, thank you for providing a pointer to this missed reference; we will make sure to incorporate it.

Thank you again for your time; we look forward to your final assessment.

We are very grateful to the reviewer for their detailed and insightful comments. We address your questions point by point below, and point to our reply to Reviewer hqD4 (due to space constraints) which details a "Mechanistic model of ring-attractor dynamics" inspired by Question #5 in your review.

Responses to questions

(1) Hyperparameters: Tuning hyperparameters is a common challenge across many models, with cross-validation being the de facto approach. Although CLDS models are a relatively simple class of models (e.g., in comparison to any deep network based approach), a comprehensive grid search of all potential hyperparameter settings is still computationally infeasible. Here, we stick to a relatively simple choice of an approximate GP kernel with a single hyperparameter governing smoothness (i.e. the lengthscale). We tune this by cross-validated co-smoothing.

While it is possible to treat $L$ as a hyperparameter tuned to improve performance, we propose to instead set $L$ to a large enough value that results in negligible error in the GP kernel approximation. As $L \to \infty$, we pay a larger computational price in fitting the model (because we have more basis functions), but the statistical properties of the model are essentially unchanged because high-frequency basis functions have nearly zero amplitude. This idea is well established in GP regression literature; see for example (Greengard et al. 2025; "Equispaced Fourier representations for efficient Gaussian process regression from a billion data points."). Hopefully, these clarifications make it clear that $L$ was not selected in an “ad hoc” manner, but set to a “large enough” value. We will explain this in more detail in the revision, as well as numerical error metrics.

Regarding dimensionality, we are unaware of conclusive evidence that 10-15 dimensions is optimal to model nonhuman primate motor cortex. Could the reviewer give us a more concrete citation? In our hands, we find that we get substantially diminishing returns after fitting a ~5D model whether the model is a classic LDS or a CLDS. We report in Table 1 below co-smoothing results as we take $D \in \\{3, 5, 10, 15\\}$. In our revision, we will include a more comprehensive sweep over this hyperparameter as a supplementary figure.

Table 1: Co-smoothing reconstruction of single held-out neurons from the test-set, over varying latent dimensionality $D$, averaged over two random seeds.

(2) Identifiability: Our revision will clarify the set of equivalent CLDS model solutions, known from the literature on LDS models (Glover & Willems, 1974, “Parametrizations of linear dynamical systems: Canonical forms and identifiability”). We note that this form of non-identifiability is relatively mild compared with many alternative approaches. Anything nonlinear dynamical systems model or deep neural network approach (e.g. LFADS) will have much more extreme forms of non-identifiability.

Although the latent space of CLDS is only identifiable up to these transforms, other quantities are uniquely identified. First, eigenvalues of the dynamics matrices are also identifiable (and thus the dynamical regimes), which we make sure to convey and asses recovery on the two head-direction tasks. Second, for a fixed set of inputs $u_1, \dots , u_T$ and a fixed estimate of CLDS model parameters the conditional distribution of observations $y_1, \dots, y_T$ is an $NT$-dimensional Gaussian with identifiable mean and covariance. The parameters of this Gaussian are efficiently computable via Kalman smoothing. Furthermore, these parameters have scientific value. For example, the covariance characterizes across-neuron and across-time noise correlations. See, for example, Panzeri et al. (2022, "The structures and functions of correlations in neural population codes.").

(3) Nonlinear experiment: Thank you for this insightful perspective, it highlights an opportunity for improving the soundness of our approach. Our motivation behind Figure 2 was to show how a canonical nonlinear neural circuit (a ring attractor) could be instantiated by a conditionally linear model. We show (not surprisingly) that we can recover the parameters of this model when fit to noisy simulated observations. Now, as prompted by you and Reviewer hqD4, we sought to challenge the inference and devised a new task, specifically chosen to reflect well-established high-dimensional and non-normal (chain-like) dynamics of interest in the neuroscience community. The ground-truth model is not a CLDS, providing a model-mismatch experiment. We found that the CLDS was able to reliably uncover the ring-attractor structure of the true model, further attesting to its ability to recover the nonlinear structure we've seen in the mouse head-direction data.

Finally, like any statistical model, CLDS models will fail to provide good fits in regimes where the data is too noisy or scarce. Data requirements will grow with the complexity and dimensionality of the underlying system. Unlike nonlinear/deep network alternatives, linear systems are amenable to a theoretical analysis (although this remains to be an open area of research). See, for example, Hardt et al. (2018) "Gradient descent learns linear dynamical systems", which shows that no local minima exist under certain conditions. Some sample complexity results are also available. See, Tsiamis et al. (2023) “Statistical learning theory for control: A finite-sample perspective”.

(4) u-x diagnostic metric: Thank you for this important question; it is challenging but essential to get a good grasp of this relationship, given the importance of inputs in our modeling approach. Already present in the manuscript, we’ve derived in Appendix A.2. bounds to our approximation error for both the CLDS approximation to a “true” nonlinear system (eq. 23) and the composite dynamics to a “true” autonomous system (eq. 42). Both of these bounds revolve around the second moments of the conditions distributions $p(x | u)$ and $p(u | x)$, namely $\xi :=\mathbb{E}_{u}[\mathrm{Tr}[\mathrm{Cov}[x | u_t]]$. We are pleased to provide numerical estimates of this quantity in Table 2 below, computed empirically from posterior estimates in a manner akin to the composite dynamics detailed in the text. A control is obtained by performing the same procedure on conditions shuffled over trials/batches.

Table 2: empirical estimates of $\xi$. The more co-dependent these variables are, the smaller $\xi$ will be, and the better our CLDS approximation is to a ground-truth nonlinear system.

Their relative value and difference to the control help provide this diagnostic feature. These quantities will be included in the updated manuscript, and the methods along with the open-source code. Finally, regarding the SLDS example, we refer to points (1) and (7) on general model selection.

(5) CLDS decompose nonlinear dynamics It is worth noting that Fig. 2b represents the ground truth composite dynamics, not the ground-truth model itself. Nonetheless, and directly to your point, we have developed above an additional example where we fit a CLDS to a mechanistic ring-attractor model (see response to hqD4). We hope that our description and the pending figure will satisfy your questions.

Part of our motivation here was to show that a CLDS model can, in fact, be hand-constructed to produce simulated ring attractor dynamics. We believe this approach of constructing nonlinear dynamics from linear dynamics per condition is a conceptual point worth making to the neural modeling community.

(6) Regarding the head direction system, we will include a supplemental figure showing all neural tuning curves (all of which are well-fit by the model). Regarding condition-averaged trajectories, thank you for pointing this out, our revision will include them—they look similar for the LFADS and CLDS models.

Importantly, both the tuning curves and condition-averaged trajectories are relatively simple to fit. The hard problem is predicting the statistics of single-trial dynamics and an underlying flow field that is capable of forecasting or predicting heldout activity. This is quantified by the co-smoothing analysis.

(7) Choice of inputs: Firing rates of motor cortical neurons rapidly rise at movement onset, see Kaufman et al. (2016, “The Largest Response Component in the Motor Cortex Reflects Movement Timing but Not Movement Type”). In principle, we could have set the step function to coincide with the movement onset of the hand, a standard way to align spike times in reaching tasks, but unfortunately, the data we’re working with didn’t have this information.

The CLDS model does involve a degree of “hand-engineering.” As we explain in the paper, we view this as both a strength and a limitation of the model. It is a limitation because we might miss out on important latent variables by taking a more supervised approach. But, on the other hand, purely unsupervised methods are difficult to fit. Our view is that there is a lack of tools that are easy-to-use and allow users to incorporate domain knowledge into their model. CLDS models fill this gap in the literature.

It is true that the inputs to the model can be engineered in a variety of ways, and therefore fair to ask, “How should one choose $u$”? We can’t give a general answer to this question— the answer will depend on the specifics of each experiment. If a practitioner identifies several possible choices for a particular experiment, we suggest that cross-validation could be used to select between these alternatives.

Thank you again for your time; we look forward to your final assessment.