BRAID: Input-driven Nonlinear Dynamical Modeling of Neural-Behavioral Data

[Q5]: how do you ensure that the model learned all encoded information and the residual is the non-encoded information since the model is nonlinear with high flexibility

Our objective function during training is optimizing for learning encoded information. Specifically, our training objectives are prediction of neural activity and prediction of encoded-behavior (behavior after the preprocessing step, see response to question 2). So if our objective is optimized and converges, then it implies that we have learned most of the encoded information. As with all deep learning methods, we use standard numerical optimization with stochastic gradient descent for learning. We continue the gradient descent up to 2500 epochs or until the loss stops improving for several epochs (i.e., early stopping mechanism). In our results, the early stopping criteria is typically met at a much lower number of epochs suggesting that convergence is relatively fast.

Additionally, during behavior preprocessing, we ensure all encoded behavior information is learned by fitting a high-dimensional preprocessing model (RNN) to neural data which is then used to maintain only the encoded behavioral information while filtering out non-encoded information.

[Q6]: Could you show your learning curve against epochs with zooming in the last epochs if necessary

We have added example learning curves (see Fig. B4 in the temporary supplementary file) to show convergence of our models. We will add these results to the paper.

[Q1] C_z(2) and the difference between X_k(1) and X_k(2):

Thank you for this observation. The reviewer raises a good point, which we clarify here. In Fig. 1a, our aim is to visualize the objectives and supervisions for learning dynamics/states rather than the complete computation graph. While we do learn a mapping $C_z^{(2)}$ from $x^{(2)}$ to $z$, this mapping is learned in a final optimization step, meaning after x1 and x2 are fully learned and fixed through stages 1 and 2. We clarify this point in lines 810-813 and explain why/when it could be beneficial in lines 814-818. $x^{(1)}$ and $x^{(2)}$ are different even if each has a mapping to both behavior ($C_z^{(1)}$, $C_z^{(2)}$) and neural activity ($C_y^{(1)}$, $C_y^{(2)}$). This is because during training, $x^{(1)}$ is learned to optimize the training objective of behavior forecasting in stage 1. Only after $x^{(1)}$ is learned, in stage 2 of training, $x^{(2)}$ is learned to optimize a distinct training objective, which is the forecast of residual neural activity (i.e., neural activity that is not already predicted by $x^{(1)}$). As such, $x^{(1)}$ learns the intrinsic behaviorally relevant neural dynamics and $x^{(2)}$ learns the other/residual intrinsic neural dynamics.

[Q2]: How do you make sure that the behavior prediction is not dominated by the input? Can the model learn any information encoded in neural activity?

We thank the reviewer for this great comment. Addressing the problem raised by the reviewer is precisely the motivation behind the preprocessing stage introduced by BRAID explained under section 3.3. Note that we refer to such non-encoded dynamics as behavior-specific dynamics as well (Fig. 1a, sections 3.3 and 4.1.3). Briefly, the preprocessing stage avoids such non-encoded behavior dynamics by learning an RNN to filter out any behavior dynamics that are predictable by the input but are not encoded in the neural activity. That way, the output of the preprocessing RNN is a preprocessed behavior signal and both stages 1 and 2 of training use this preprocessed behavior rather than the original behavior. Thus, this preprocessing stage restricts the dynamics learned in the forthcoming stages 1, 2 to be encoded in neural activity (see section 3.3). In section 4.1.3, we validate this step in numerical simulations showing that when the preprocessing is included, the learned model has an improved neural prediction matching the one obtained via the true model (Fig A.2). This means that BRAID did indeed learn information encoded in neural activity. Importantly also, in all real data analyses we incorporate this preprocessing RNN stage. We clarify this in line 272 and line 841. Moreover, in real data results, we report neural prediction performance in addition to behavior decoding to highlight the fact that our method not only predicts behavior well, but also outperforms baselines without sensory input in terms of neural prediction. This better neural prediction shows that the model is able to learn dynamics encoded in neural activity rather than being dominated by the non-encoded/behavior-specific input dynamics.

[Q3]: input term in modeling neural activity y, what does this input mean?

The input terms in the decoders ($C_z$/$C_y$) are implemented and included in the formulation for the sake of generality (to support the applications where feedthrough is of interest) but we have not used them in real data analyses as we are interested in the predictability through the latent states. In general, the direct input term to observations (feedthrough term) can sometimes be helpful for discrete-time systems (Åström 2013). Since the dynamic equation takes the input up to the previous time-step to infer the current state and the current observation, in general, a feedthrough term can be helpful for providing concurrent information from the input during the length of the current time-step, which has not already been aggregated into the current latent (is not in the inputs from previous time-steps).

Reference:

Åström, Karl J., and Björn Wittenmark. 2013. Computer-Controlled Systems: Theory and Design, Third Edition. Courier Corporation.

[Q4]: How do you define your initial states x0?

We use zero initialization i.e., set x0 to zero. This is because we process the data in relatively long sequences, which makes the results minimally dependent on the initial state since its effects typically die off within a few samples. We will clarify this in the paper.

[W1] Other modalities:

We thank the reviewer for their suggestion, and we agree that applying our method to other neural modalities would be valuable. Here we focused on analyzing spiking activity as it is one of the most widely used modalities in neuroscience and neurotechnology. In addition to spiking activity, we have extensive synthetic data using which we validated BRAID. Nevertheless, we applied BRAID to another neural modality, that is local field potential (LFP) (see Fig. B3 in the temporary supplementary file), and found results consistent with those obtained on spiking activity. We will include these new results in the paper.

[W2] Unmeasured inputs:

We appreciate the reviewer’s great insight on this point, and acknowledge the value of the complementary approach of input inference in lines 508-513. We should clarify that, while input inference is not the goal of BRAID, similar to other latent models, BRAID also has the capacity to implicitly learn unknown inputs as part of the latent states, as long as these unmeasured inputs encode neural and behavioral dynamics as the reviewer notes. This is because BRAID’s training objective is neural and behavioral prediction. As such, BRAID can automatically dedicate part of its latent states to learning unknown inputs if that benefits neural-behavioral prediction. In that respect, BRAID is not fundamentally less capable than input-inference approaches. The main difference is that in BRAID we do not attempt to explicitly disentangle unmeasured inputs or put any constraints such as smoothness on them.

Adding potential assumptions on the nature of unmeasured inputs as done in input-inference approaches can be an interesting future direction for BRAID to further disentangle unmeasured inputs. We will clarify this by revising our discussion.

[Q1]: How does your predictor relate to sequential posterior in variational inference?

Please refer to the response to major [W1]. VI learns the posterior by optimizing ELBO loss, which is a combination of a KL divergence w.r.t. a fixed prior and LL term. BRAID learns the posterior, i.e., the latent states of the predictor model, only based on past observations (m-step or more in the past) by directly optimizing the m-step LL.

[Q2]: Monkey experiment: What do x^1 and x^2 look like in this task? Is there any non-encoded activity x^3?

We have added a new figure as explained under major [W2,3] to visualize the latent trajectories for the monkey experiment (see Fig. B1 in the temporary supplementary file). BRAID trajectories (Fig. B.1b) correspond to the behaviorally relevant states i.e., $x^{(1)}$ while trajectories identified by the unsupervised U-BRAID (Fig. B1d) correspond to the irrelevant states i.e., $x^{(2)}$. We do not learn any behavior-specific states ($x^{(3)}$) in any of our real data analyses since our primary goal is to learn dynamics encoded in neural activity. As expected, $x^{(1)}$ are much more well-separated for different behavior conditions compared to $x^{(2)}$ which shows successful learning of the behaviorally relevant dynamics by BRAID.

[Q3]: Could you provide more details on the "automatic" selection (L313)? Is it simply picking the best performing?

We thank the reviewer for this question. Yes, this automatic selection procedure selects the best-performing model with an inner-cross validation within the training set. Here is how it works: We explore all possible nonlinearity configurations resulting in $2^4$ cases (each parameter can be either linear or nonlinear). We split the training data into two sections, train models for each case on the first section of training data and evaluate them in the held-out section of training data. We then select the configuration yielding the best behavior prediction on the held-out section of training data. Among options with the same behavior prediction, we select the one with the best neural prediction. Once the best performing nonlinearity is selected, we retrain a model with that configuration using the entire training data and report its performance on the unseen test data. Currently, we describe the automatic nonlinearity selection procedure in lines 314-316 and provide further detail in lines 884-890. We will expand this section to clarify this process as described above.

Notation error:

We thank the reviewer for their correction and will fix the place-holder notation throughout the paper.

Dependency on x^1 on Stage 2 could be made more explicit in the text:

We thank the reviewer for this remark. We now add a clarification on this in the text/caption. We also show this in Fig. 1a by an arrow from x1 to x2.

Bold entries per metric would be more representative:

Our emphasis was on highlighting cases with neural-behavior prediction being simultaneously accurate, to make it clear that other rows do either behavior or neural prediction well, but not both at the same time. We will add a notation (underline) to also highlight models that perform well for an individual metric.

[W3]: We thank the reviewer for this note and will revise the manuscript to make sure all important definitions are provided in the main text.

[W4]: We agree with the reviewer that identifiability is a general challenge for nonlinear latent state models including BRAID and beyond. In practice, different nonlinearities can be useful and easier to learn depending on the specific dataset. For example, while making all model elements nonlinear may also be a solution, making only some parameters nonlinear may be equally accurate, but more parsimonious and thus may result in a better fit for a given finite number of training samples. While our results suggest that in our dataset having nonlinear decoders provides superior performance compared to other nonlinearities, this might not be the case in another dataset. In real data, therefore, one can search for optimal nonlinearity based on the desired performance. Indeed, BRAID enables this search through an automatic selection so that users won’t need to manually check different options for nonlinearity. We show that the source of nonlinearity in the ground-truth simulated data can be identified by the automatic nonlinearity selection in BRAID, which selects the best performing configuration (section 4.1.1, line 314, and Table 2). We will add this point to the discussion.

We thank the reviewer for this important suggestion. We now include two categories of alternative results (see Figs. B1,2 in the temporary supplementary file): 1) predicted behavior trajectories i.e., position and velocity of the arm in the non-human primate dataset, 2) predicted low-dimensional posterior trajectories (i.e., latent state trajectories) for visualization. We can see that BRAID is accurate in predicting behavior. BRAID also has the most well-separated latent trajectories for different reach directions in the task. U-BRAID’s trajectories are the least separated, showing that BRAID is more successful in extracting the behaviorally-relevant dynamics that are more congruent with behavior. Also, DPAD’s trajectories are also not as well-separated as BRAID’s, though they are more separated than the unsupervised version as expected.

We thank the reviewer for their positive feedback on our work and appreciate their openness to raise their score. We now address their weaknesses/questions below.

The reviewer raises an interesting point regarding the analogy between BRAID’s predictor-generative structure and the encoder-decoder architecture of the VAEs. Below are several key distinctions:

Causality and real-time inference:

First, we apologize for the use of the word “real-time”, which is unclear. By real-time, we mean that inference is done causally in time; in particular, inferring the state at time $k$ only uses neural observations up to time $k$ but none of the future neural observations. Furthermore, this inference is also recursive and thus computationally efficient. Recursive means that BRAID can infer the state at the next time step based on the inferred state in the present time-step (without redoing any computation on past neural observations and by just updating the present inferred state based on the present neural observation). BRAID can perform inference causally in time and recursively because its predictor and generative models are both recursive unidirectional RNNs and generate predictions given the past neural observations. In other words, BRAID’s RNNs all run only forward in time. This is in contrast with SOTA sequential VAEs such as TNDM and LFADS, which require neural observations for an entire trial/segment to approximate the posterior distribution at the initial point and subsequently any time-step in the trial/segment using a bidirectional RNN (i.e., thus their inference is done non-causally in time). The causal (unidirectional) and recursive features of inference in BRAID can be of utmost interest for real-time decoding applications such as brain-computer interfaces (BCIs). We illustrate the causality of BRAID through its computation graph in Fig. 1b. We will change the word real-time to explain that we mean causal (unidirectional) and recursive.

Objective:

In general, a fundamental distinction of BRAID with variational methods is that the latter have the statistical parameters of a latent distribution (e.g., mean and variance of a Gaussian) as the ultimate output of the encoder, from which the decoder samples a point to initiate decoding. As such, variational methods also often include a KL-divergence term in the loss to enforce the distribution of this variational latent variable. In contrast, BRAID does not have a KL-divergence term, nor does it have such a sampling mechanism between the encoder and decoder, and is thus not a variational approach. BRAID learns all parameters (predictor and generative) through direct optimization of m-step ahead prediction (i.e., forecasting), which is quantified by the m-step-ahead LL. We will note both the analogy and the distinction in objective in the manuscript.

[Q6]: Modeling multiple brain regions:

We appreciate this comment and agree with the reviewer that studying multiple brain regions and their shared vs private dynamics is an interesting topic in computational neuroscience. Indeed BRAID can facilitate such multi-region investigations because it can model the activity of an upstream brain region as a measured input into the downstream region. Exploring the use of BRAID for this application would be a great future direction.

[Q7]: Typographical errors:

We thank the reviewer for the corrections and will fix them in the manuscript.

[Q1]: Could the authors attempt to make the writing clearer with regards to explaining the model and the various training stages? Perhaps a summary as a table might help:

We thank the reviewer for this excellent suggestion. We will create a summary table, showing how the 2 representations of dynamics are learned by two RNNs (RNN and RNN_fw) and how the states in each RNN are divided into 3 subsets that are learned by a sequential multi-stage optimization. This reviewer’s explanation of these RNNs and stages was quite accurate and succinct, which we will use for inspiration to make the writing more clear as explained under W1.

[Q2]: Results for ablating second and third stages, and comment on end-to-end training:

We apologize for the unclarity. Our main results ablate for second and third stages of training and only use stage 1 alone (Fig. 3 and Table 3 with nx=16). We indicate this by setting n1 (dimension of stage 1) equal to nx (full dimension), which indicates that only stage 1 is used in building the model. We refer to this stage-one-only setting as a low-dimensional regime in our real data analyses and explain this in section 4.2 lines 431. We additionally add stage 2 for the high-dimensional regime with nx=64 in Table 3 (that is n1=16 and n2=48). We explain what these settings mean in section 4.2, lines 431-432 but will make this much more clear: “In addition to these results in the low-dimensional …”. As Table 3 shows, adding stage 2 improves neural forecasting while leaving behavior forecasting essentially unchanged by adding the neural-specific latents. The third (post-hoc) stage is not used in real data analyses for the reason explained in response to Weakness 1. We have however used and validated the third stage in our simulated analyses provided in Fig A.2 and denoted that in the legend as “with non-encoded dynamics”. As this figure shows, ablating for stage 3 (denoted “BRAID (with preprocessing)”, green) results in lower behavior decoding compared with having stage 3 (denoted “BRAID (with preprocessing and non-encoded dynamics)”, red) because stage 3 learns behavior-specific dynamics.

End-to-end training without stages: As the reviewer notes, one can train all components simultaneously instead of the multi-stage approach proposed by BRAID (as done by TNDM and mmPLRNN baselines). We want to highlight that the multi-stage training is an important aspect for prioritization of the intrinsic behaviorally relevant neural dynamics in learning. Having separate stages allows for the intrinsic behaviorally relevant dynamics to be learned first by optimizing a behavioral loss, so that these dynamics are not mixed with other neural dynamics, which are learned in stage 2 while fixing RNN1’s parameters. On the other hand, in an end-to-end training, one needs to simultaneously optimize a combined neural-behavioral prediction loss and as such may learn latent states that mix up behaviorally relevant and other neural dynamics, thus not prioritizing the former and not learning the former as accurately. Indeed TNDM and mmPLRNN both use end-to-end training with neural and behavioral terms, and so the superior behavior decoding of BRAID compared to TNDM and mmPLRNN suggests the benefit of multi-stage learning.

[Q3]: Scaling performance experiments and comparisons to LFADS (and CEBRA) if possible:

We thank the reviewer for this comment. As with any model, we believe BRAID will improve forecasting with more neurons. To show this, we are now working on running BRAID with different numbers of neurons and will report the results by the end of the discussion period. Please see response to [W4] above regarding CEBRA and LFADS.

[Q4]: Comment on multi-session models and generalization:

Please refer to the response to W3.

[Q5]: Use of similarity metric to compared learned and true dynamics:

We thank the reviewer for this question. We agree this is an important step for validations and have done so in our simulated datasets by comparing the learned intrinsic dynamical system to the ground truth system. The results are shown in Fig. 2c,g and Fig. A1d. As intrinsic dynamics in a system can be quantified by the eigenvalues of the state transition matrix, we quantify the similarity between dynamics via the accuracy of learning intrinsic eigenvalues. As the results suggest, the eigenvalues of the intrinsic dynamical system are learned much more accurately using BRAID compared to DPAD and linear baseline (IPSID/linear BRAID). As the reviewer suggests this kind of similarity analysis is possible in the synthetic data where the ground truth is known, but as far as we know is not feasible in real data due to the lack of ground truth. As such, neural/behavior decoding are the measures used in real data, which reflect how accurately the underlying neural-behavioral dynamics were learned.

We thank the reviewer for this comment.. We would like to first note a few points. First, we currently do ablation for explicit input modeling by comparing with DPAD, which does not take into account inputs and is actually closer to our architecture than LFADS. Second, DPAD was compared with LFADS in the DPAD paper, suggesting that it improves behavior decoding compared to LFADS. Thus, we expect BRAID to do better than LFADS as it does better than DPAD. Third, regarding CEBRA, the DPAD paper also does comparisons of DPAD with CEBRA, showing that the dynamical modeling in DPAD leads to better neural-behavioral prediction. As BRAID improves upon DPAD, we expect it to also improve upon CEBRA. Nevertheless, we will do our best to run the comparisons with at least one of CEBRA and/or LFADS by the end of the discussion period and report back.

We appreciate this thoughtful comment from the reviewer. With regard to performance on the latter section of a session when trained on the initial parts, this is effectively what happens in the last fold of our 5-fold cross validation, where we train the model on the first 80% of the session and evaluate it on the last 20% of the session. To answer the reviewer’s question, we can compare this final fold with fold 3, where the first and last 40% of the session are used for training while the middle 20% is used for evaluation. As expected, the last fold has slightly worse performance than the middle fold, which can be explained by non-stationarities throughout one session:

Regarding cross-session generalizability, BRAID as presented in our work is designed for the single-session setting. We will add a discussion on how similar approaches taken elsewhere in the literature for other methods could be taken to add cross-session decoding support to BRAID. One key approach is to add session-specific encoding and decoding matrices between the neural data and the input and outputs of the model in each session, while keeping all other model parameters the same across sessions. This will allow data from multiple sessions to be used for training the model and when a new session is given, one can apply the learned model to it by simply using very minimal data to learn just the simple session-specific encoding/decoding matrices for that new session. Beyond cross-session generalizability of the same task, BRAID may also better generalize across different variations of tasks with different sensory instructions (e.g., center-out targets vs. random targets in a grid); as BRAID can account for these differences, it can help with aligning the intrinsic dynamics across tasks. Extending BRAID by these techniques to account for variability across sessions/datasets is indeed an interesting future direction. We will add a discussion in the manuscript.

We thank the reviewer. We will perform an analysis on this and report back.

We thank the reviewer for bringing up the clarity issue for use cases of RNN2 and RNN3. The reviewer’s understanding of our method architecture (2 RNNs, 3 subsets of states in each, preprocessing, etc.) is exactly correct, which we truly appreciate. Our primary focus is to learn intrinsic behaviorally relevant neural dynamics with RNN1 and with priority. We will revise the manuscript to carefully explain the use cases for the optional neural-specific (RNN2) and behavior-specific (RNN3) components, which we describe below in detail:

RNN2: One would use this RNN if, beyond intrinsic behaviorally relevant neural dynamics, they are also interested in learning the other neural dynamics, which are not shared with behavior (i.e., residual neural dynamics). Since all neural dynamics may not be relevant to the measured behavior, learning RNN2 helps explain the neural-specific dynamics not already captured in RNN1, which learns the shared neural-behavioral dynamics. Note that BRAID allows the two subtypes of neural dynamics (behaviorally relevant vs. other) to be dissociated, as one subtype is learned in RNN1 and one is learned in RNN2. So a user can actually study these subtypes separately. We have demonstrated this use case in a high-dimensional regime where we fit both RNN1 and RNN2 to the real datasets, leading to improved neural predictions due to learning both subtypes of neural dynamics (Table 3, nx=64).

RNN3: One would use this RNN3 if they are interested in studying the dynamics in behavior that are not encoded in the recorded neural activity (i.e., the behavior-specific dynamics). RNN3 disentangles behavior-specific dynamics from the two subtypes of intrinsic neural dynamics learned in RNN1 and RNN2. Furthermore, one would use RNN3 if they are interested in improved behavior decoding by using behavior-specific latents. We show this in simulated datasets (Fig. A2). Since our main interest lies in identifying neural dynamics, we do not employ this RNN3 in our real data analysis and always decode behavior purely from neural activity alone. RNN3 is meant to complement the preprocessing stage where the behavior-specific dynamics are filtered out via a separate preprocessing RNN so that RNN1 and RNN2 can exclude those dynamics. We also explain this use-case in section 3.3, lines 273-278 as well as section 4.1.3, lines 375-377.

I would like to thank the authors for their response and appreciate their efforts in addressing my concerns. I am mostly satisfied and would be open to increasing my score, pending the results associated with W2.

We are glad the reviewer is satisfied with our revisions and really appreciate their openness to raise their score pending our results on [W2]. We have now completed this analysis and provide the results in the table below. We will also add this to the appendix. As the table shows, behavior forecasting improves with more neurons and neural forecasting remains largely stable. This suggests that our method can aggregate behaviorally relevant information across a larger population of neurons, while still being able to model this higher-dimensional population activity well.

Thank you for the update. In light of the authors' rebuttal I am inclined to increase my score. Once again I thank the authors and appreciate their efforts in addressing the reviewers' concerns, and I look forward to seeing the updated manuscript.

We are grateful that the reviewer is inclined to increase their score in light of our rebuttal and new analyses. We have now updated the manuscript by incorporating the temporary supplementary file with all the new analyses into it. The updated manuscript now includes all suggestions from the reviewer, including the new results for the effect of the number of neurons on our method's performance that the reviewer had asked for. Furthermore, we have now completed the comparisons to two additional baselines suggested by the reviewer, i.e., LFADS (with controller which infers input) and CEBRA (supervised by behavior and with sensory inputs (w.s.i.)). We have now also added these results to the uploaded revised manuscript and copy them below for convenience.

As the results suggest, BRAID outperformed both new baselines in neural-behavioral prediction, except for low-dimensional models where the neural prediction is on par with LFADS. Note that this is expected as LFADS is an unsupervised model that purely optimizes neural reconstruction while low-dimensional BRAID optimizes behavior prediction to prioritize the intrinsic behaviorally relevant dynamics.

We again sincerely thank the reviewer for their valuable insights and discussions, which we believe significantly improved our manuscript. We are also grateful that the reviewer is inclined to raise their score.

This is a great question and lack of clarity on this may have also led to [W1] from the reviewer.

We need multiple RNNs and a multi-stage learning algorithm because the goal here is to disentangle multiple subtypes of neural-behavioral dynamics. First, we have 3 sections/stages to learn and separate 3 types of neural-behavioral dynamics:

1. intrinsic behaviorally relevant neural dynamics (RNN1/RNN1_fw),
2. other intrinsic neural dynamics (RNN2/RNN2_fw),
3. behavior-specific dynamics (RNN3/RNN3_fw).

Second, for each of the 3 groups, we learn both predictor (RNN) and generative (RNN_fw) representations of dynamics, because in general one representation cannot be deduced from the other one, so we need to learn both from data (see section 3.1).

Our approach of having a separate RNN for each type of neural-behavioral dynamics and formulating the losses and stages of learning such that each RNN learns exactly the type and representation of dynamics that it is meant to is very beneficial. First, this way, each RNN is dedicated to a specific subtype of neural-behavioral dynamics, effectively disentangling the subtypes. Second, as we describe in methods section 3.2, the multi-stage learning of these 3 subtypes of dynamics allows us to prioritize the learning of intrinsic behaviorally relevant neural dynamics over other subtypes, such that we can learn them more accurately without them being confounded by other subtypes. Indeed, the primary focus here is the learning of this first subtype of dynamics. We have provided results showing utility of all these subtypes of dynamics (see Table 3, nx=16 and Fig. 3, for RNN1, Table 3, nx=64 for RNN2 and Fig. A2 for RNN3).

The reviewer also points out a great question of why a separate RNN (RNN1_fw) is needed for forecasting. We apologize that this was not clear in our write up and will update the manuscript to clarify. While the predictor RNN is indeed a dynamical model, it learns a different representation of the dynamics compared with the generative RNN_fw:

• RNN performs a recurrent transformation that aggregates past input and neural observations to predict the latent state, whereas
• RNN_fw generates future latent states only by aggregating inputs (and without relying on neural observations).

These are two conceptually distinct processes and even their inputs are different so neither of them is able to replace the other even in terms of dimensionality of their inputs. Moreover, the recursive computation in these two cases are mathematically distinct (i.e., $A$ is not the same as $A_{fw}$).

The distinction between predictor ($A$) and generative ($A_{fw}$) forms of dynamics is very fundamental and has parallels in linear system theory, where their exact relationship can be proven (section 3.1). Briefly, in a linear dynamical system model with state transition matrix $A_{fw}$, the generative form representation is how the states evolves autonomously with the recursion $A_{fw}$, while the predictor form representation shows how observations can be integrated to update state estimates via a Kalman filter (Van Overschee & De Moor, 1996). Similarly, in the nonlinear case, these two representations are not the same and furthermore their relationship is not known (unlike the linear case); therefore, it is necessary to directly learn generative models ($A_{fw}$) from the data, which BRAID does with its RNN_fw and forecasting objective. To make this more clear, in simulations, we now compare the recurrence learned by the predictor form RNN (instead of the generative form RNN_fw) with the true intrinsic dynamics and show that without learning a separate generative recursion, the predictor form RNN alone cannot learn the intrinsic dynamics (row 2 in table below). Moreover, we show that even if we extend DPAD to incorporate BRAID’s forecasting loss and inputs into DPAD (rows 3, 4 in table below), the single RNN in DPAD, which is in predictor form, still cannot find the intrinsic eigenvalues accurately. We will add the new results below for the ablation study inspired by the reviewer comment to the paper.

Eigenvalue errors:

Simulation with spiral manifold (Fig. 2a):

References: Peter Van Overschee and Bart De Moor. Subspace Identification for Linear Systems. Springer US, Boston, MA, 1996. ISBN 978-1-4613-8061-0. doi: 10.1007/978-1-4613-0465-4

We appreciate this important comment by the reviewer. We show that BRAID can effectively disentangle input dynamics from intrinsic dynamics in our extensive simulations in which the ground-truth intrinsic dynamics, quantified by the eigenvalues, are known (Fig. 2c,g and Fig. A1d). Our results show that BRAID more accurately recovers the intrinsic eigenvalues by explicitly modeling external inputs and learning the generative form of dynamics, compared to DPAD, which does not account for the external input. In real data, since ground-truth intrinsic eigenvalues are not known, we use the forecasting metrics that show the forward model has successfully learned and integrated both intrinsic and input-driven contributions as BRAID outperforms alternatives in forward prediction of both neural and behavioral data.

We thank the reviewer for their comment. Based on this comment and the reviewer’s comment 3 about “why is it necessary to have a separate RNN for forecasting”, it seems that the goal and conceptual advance of BRAID were not clearly communicated in our text. The primary goal of BRAID is to disentangle intrinsic dynamics from input dynamics, which is not even conceptually formulated in DPAD. Furthermore, we now show with new analyses expanded under W3 that doing so is not achieved by straightforward additions (e.g., simply adding input) to DPAD. Specifically, to do this, BRAID introduces two major conceptual advances in addition to multiple technical advances over DPAD.

Conceptual advances:

1. We show that such disentanglement requires two RNNs learning fundamentally distinct dynamics and cannot be as accurately achieved with the single RNN used in DPAD: we need a separate generative RNN (RNN_fw) to learn the intrinsic dynamics in addition to another predictive RNN used in DPAD that integrates the neural observations for immediate state prediction. We now show that even if we add BRAID’s forecasting loss and/or inputs to the single RNN in DPAD, DPAD still learns the intrinsic dynamics inaccurately, as quantified by eigenvalues . See details and additional results below in response to [W3]. The distinction between predictor (RNN) and generative (RNN_fw) forms of dynamics is very fundamental and has parallels in linear dynamics systems theory, which we lay out in section 3.1 (see also the response to [W3]). DPAD only learns the former, while BRAID adds a distinct RNN to also learn the latter representations.

2. We show that incorporation of measured inputs necessitates a new computational step that uses a separate RNN to exclude behavior-specific dynamics before training, without which the learned behaviorally relevant dynamics may not even be neural, i.e., may not be present in neural activity (preprocessing stage, section 3.3). Indeed, if one trivially just adds input to DPAD as the reviewer notes for example by concatenating inputs to neural activity, DPAD will learn dynamics (i.e., find eigenvalues) that are not in neural activity and are just in behavior, because some dynamics in input just drive behavior and DPAD will have no way of knowing that. Note that this BRAID preprocessing step has no equivalent element in DPAD.

In addition to the two major conceptual advances above, below are the technical advances over DPAD:

Other technical advances:

3. The incorporation of a forecasting loss and showing that it is necessary to learn intrinsic dynamics. Note that having forecasting loss in the optimization is not at all supported in DPAD. As an additional baseline, we expand DPAD to support this and show that it is not enough to achieve the goals of BRAID.

4. Enabling the extraction of private dynamics in behavior (behavior-specific dynamics) by introducing a third set of RNNs (RNN3 and RNN3_fw) with the proper loss and proper links to the first two RNNs (see section 3.3). Again, this third set of RNNs has no equivalent in DPAD.

In addition to the above, it is also worth noting that even adding the generative RNN_fw in two sections (stages 1 and 2, see section 3.2), with proper connections between the two sections of the predictor RNN was not straightforward and required new custom RNN cells with support for our orthogonal two RNN formulation (visualized in Fig. 1b) and was not possible as a simple extension of any existing method.

Finally, at the end of the day, we use DPAD as a major baseline in both simulations and data and extensively show the benefits of BRAID over DPAD, with extended analyses as expanded on under W3 below. It is also worth noting that we did these extensive comparisons with DPAD even though according to ICLR policy, DPAD is considered concurrent work (because it was published at a peer reviewed venue on Sep. 6, 2024, which is 2 months after the July 1, 2024 cut-off for concurrency) and thus comparisons to it are not required by ICLR.

We thank the reviewer for this comment. As explained in response to [W2], as an attempt at interpreting intrinsic versus input dynamics, we have shown that BRAID successfully disentangles intrinsic and input dynamics in our simulations (eigenvalue errors e.g., in Fig. 2c.g). In interpreting the disengagement of dynamics, we focus on simulations because in simulations the ground truth for the sources of dynamics is known. When we show small errors in learning a specific type of dynamics, this is a direct interpretation of the learned dynamics by pointing at exactly what aspect (e.g., eigenvalues) in the simulated models were learned by BRAID. However, we agree with the reviewer that further analysis of the individual intrinsic vs input-driven components enabled by BRAID would be very interesting. We hope that our validation in simulations will encourage others to use our method to interpret their own datasets.

We thank the reviewer for this suggestion. We will add R^2 values (at least for our main results) to the paper. We observed similar trends in terms of both R^2 and CC metrics, with BRAID’s R^2 values being roughly the square of the CC values, which suggests that no major scale or baseline mismatches existed for BRAID. In addition to R^2, we will also keep the CC results in the interest of being more easily comparable with those papers in the field that report CC. Please see the R^2 version of our main results in the temporary supplementary file (Table B.2).

We are not sure what the reviewer means by “optimize separately for behavior reconstruction”. First, if they mean the optional behavior-specific state ($x^{(3)}$ learned by RNN3) for behavior decoding, we should clarify that we never use these states in the decoding results as our aim is to decode behavior solely using neural activity and input (i.e., via the identified latent states in neural activity). Thus, the improvements in behavior decoding achieved by BRAID are not due to behavior-specific dynamics and the comparisons with TNDM and DPAD are fair. Second, if the reviewer means the use of behavior reconstruction during learning, both TNDM and DPAD, similar to BRAID, optimize for behavior reconstruction already using their behaviorally relevant latent states as one of their main objectives, with DPAD doing so in a separate optimization similarly to BRAID; so there is no need to add this to them. Third, to further ensure a fair comparison with TNDM, we also modified TNDM to accept sensory inputs alongside neural activity for inference of latent states. Even with this modification, BRAID outperformed the modified TNDM even though they both use the exact same signals for decoding (Table A3).

The reviewer asks:

How would this be different from just training a model to decode the behavior 8 time steps ahead directly from the neural data

Our new ablation studies in new table A.3 show that simply training a model (e.g., DPAD) to decode behavior multiple steps ahead does not learn the correct intrinsic neural dynamics, which is a primary goal here. In this analysis, we added multi-step ahead loss to DPAD, and showed in simulations where the ground truth is known that it still does not learn the intrinsic dynamics. BRAID in contrast accurately learns the intrinsic dynamics using its distinct predictor-generator architecture (figure 1b).

Next, we clarify why the disentanglement of subtypes of dynamics is of interest and what interpretability it provides. The reviewer asks:

what is the additional utility or interpretability that one gets out of inferring an underlying latent state?

First, one of the key goals of learning disentangled latent representations, or disentangled representation learning (Wang et al 2024), is to make the learned embedding interpretable. The fact that we know which latent state dimensions are associated with each subtype of dynamics makes the model more interpretable. One of the ways this interpretability can be useful is now shown in new Fig. A4. In this figure, we have visualized the trajectories (temporal evolution) of the latent states that BRAID learns for the first subtype of dynamics in our experimental data. These latent state trajectories are more congruent with reach directions compared with those extracted from other methods, suggesting that BRAID has been more successful in learning and disentangling the intrinsic behaviorally relevant neural dynamics of motor cortex (M1). Learning more interpretable and behavior related latent representation has been a key goal for many impactful methods in neuroscience (Churchland et al 2012, Kobak et al 2016, Sani et al 2021, etc).

Second, the latent states in BRAID model the temporal evolution of a neural population as a dynamical system, for example the eigenvalues associated with each state specify the decay and oscillations in the temporal evolution. In the “computation through neural population dynamics” view (Vyas et al, 2020), the neural population performs computations through these dynamics to generate behavior, e.g., movement. As such, the study of these latent states in dynamical system models (e.g., BRAID) can be used to study neural computations (Remington et al., 2018; Shenoy et al., 2013). The contribution of our work to this framework is to help address the major problem of disentangling the dynamics that are intrinsic to the population from those that are due to external sensory inputs. Doing so is important, for example, to identify cortical dynamics that are more autonomous/intrinsic (Seely et al. 2016; Shenoy et al., 2021; Sauerbrei et al., 2020; Vyas et al 2020; Vahidi et al 2024). Indeed, our work allows not only for such intrinsic vs. input disentanglement, but also for dissociating the 3 subtypes of intrinsic neural-behavioral dynamics as stated in our response above. As demonstrated in our simulations (Figs. 2, A1), our method enables accurate learning of the intrinsic dynamics as quantified by the eigenvalues. Also, as explained in our response above, the new ablation analysis (new table A.3) shows that the BRAID architecture is critical in learning such accurate intrinsic dynamics.

The above descriptions, references and results clarify why in BRAID we aim to dissociate the three subtypes of intrinsic dynamics in neural-behavioral data, and to further disentangle input dynamics from them. These also explain the interpretability offered by doing so. Please let us know if you have any other questions or would like further clarification on any point. We again thank the reviewer for all their important comments and are encouraged that the reviewer finds our new analyses "very helpful" and "will increase" their score.

Please see the next comment for the references.

References:

Mark M Churchland, John P Cunningham, Matthew T Kaufman, Justin D Foster, Paul Nuyujukian, Stephen I Ryu, and Krishna V Shenoy. Neural population dynamics during reaching. Nature, 487 (7405):51–56, 2012.

Dmitry Kobak, Wieland Brendel, Christos Constantinidis, Claudia E Feierstein, Adam Kepecs, Zachary F Mainen, Xue-Lian Qi, Ranulfo Romo, Naoshige Uchida, and Christian K Machens. Demixed principal component analysis of neural population data. elife, 5:e10989, 2016.

Evan D Remington, Seth W Egger, Devika Narain, Jing Wang, and Mehrdad Jazayeri. A dynamical systems perspective on flexible motor timing. Trends in cognitive sciences, 22(10):938–952, 2018. Omid G Sani, Hamidreza Abbaspourazad, Yan T Wong, Bijan Pesaran, and Maryam M Shanechi. Modeling behaviorally relevant neural dynamics enabled by preferential subspace identification. Nature Neuroscience, 24(1):140–149, 2021.

Britton A Sauerbrei, Jian-Zhong Guo, Jeremy D Cohen, Matteo Mischiati, Wendy Guo, Mayank Kabra, Nakul Verma, Brett Mensh, Kristin Branson, and Adam W Hantman. Cortical pattern generation during dexterous movement is input-driven. Nature, 577(7790):386–391, 2020.

Jeffrey S Seely, Matthew T Kaufman, Stephen I Ryu, Krishna V Shenoy, John P Cunningham, and Mark M Churchland. Tensor analysis reveals distinct population structure that parallels the different computational roles of areas m1 and v1. PLoS computational biology, 12(11):e1005164, 2016.

Krishna V Shenoy and Jonathan C Kao. Measurement, manipulation and modeling of brain-wide neural population dynamics. Nature communications, 12(1):633, 2021. Krishna V Shenoy, Maneesh Sahani, and Mark M Churchland. Cortical control of arm movements: a dynamical systems perspective. Annual review of neuroscience, 36(1):337–359, 2013.

Parsa Vahidi, Omid G Sani, and Maryam M Shanechi. Modeling and dissociation of intrinsic and input-driven neural population dynamics underlying behavior. Proceedings of the National Academy of Sciences, 121(7):e2212887121, 2024.

Saurabh Vyas, Matthew D Golub, David Sussillo, and Krishna V Shenoy. Computation through neural population dynamics. Annual review of neuroscience, 43(1):249–275, 2020.

Xin Wang, Hong Chen, Si’ao Tang, Zihao Wu, and Wenwu Zhu. Disentangled Representation Learning. IEEE Transactions on Pattern Analysis and Machine Intelligence, 46(12):9677–9696, 2024.

We have now updated the manuscript by incorporating all the new analyses and discussions from the temporary supplementary file into it. The uploaded updated manuscript now addresses all comments from all reviewers. We once again thank all reviewers for the valuable and constructive feedback, which we believe has significantly improved our manuscript.