BRAID: Input-driven Nonlinear Dynamical Modeling of Neural-Behavioral Data
Authors: Parsa Vahidi
Omid G. Sani
Maryam Shanechi
Venue: ICLR 2025 Poster (2025)
Keywords: Deep learning
Dynamic modeling
Sensory stimuli
RNN
Intrinsic
Behavior
Primary area: applications to neuroscience & cognitive science

Abstract: Neural populations exhibit complex recurrent structures that drive behavior, while continuously receiving and integrating external inputs from sensory stimuli, upstream regions, and neurostimulation. However, neural populations are often modeled as autonomous dynamical systems, with little consideration given to the influence of external inputs that shape the population activity and behavioral outcomes. Here, we introduce BRAID, a deep learning framework that models nonlinear neural dynamics underlying behavior while explicitly incorporating any measured external inputs. Our method disentangles intrinsic recurrent neural population dynamics from the effects of inputs by including a forecasting objective within input-driven recurrent neural networks. BRAID further prioritizes the learning of intrinsic dynamics that are related to a behavior of interest by using a multi-stage optimization scheme. We validate BRAID with nonlinear simulations, showing that it can accurately learn the intrinsic dynamics shared between neural and behavioral modalities. We then apply BRAID to motor cortical activity recorded during a motor task and demonstrate that our method more accurately fits the neural-behavioral data by incorporating measured sensory stimuli into the model and improves the forecasting of neural-behavioral data compared with various baseline methods, whether input-driven or not.

Published as a conference paper at ICLR 2025
BRAID: INPUT-DRIVEN NONLINEAR DYNAMICAL MODELING OF NEURAL-BEHAVIORAL DATA
Parsa Vahidi', Omid G. Sani', Maryam M. Shanechi
"Electrical and Computer Engincering, CComputer Science, 'Biomedical Engineering Viterbi School of Engineering, University of Southern California (USC), Los Angeles, CA (pvahidi, omid. ghasemsan,, shanechi | edu, "Corresponding author
ABSTRACT
Neural populations exhibit complex recurrent structures that drive behavior, while continuously receiving and integrating external inputs from sensory stimuli, up- stream regions, and neurostimulation. However, neural populations are often mod- eled autonomous dynamical systems, with little consideration given the influ- ence of external inputs that shape the population activity and behavioral outcomes. Here, we introduce BRAID, deep learning framework that models nonlinear neural dynamics underlying behavior while explicitly incorporating any measured external inputs. Our method disentangles intrinsic recurrent neural population dynamics from the effects of inputs by including forecasting objective within input-driven recurrent neural networks. BRAID further prioritizes the learning of intrinsic dynamics that are related to behavior of interest by using multi-stage optimization scheme. We validate BRAID with nonlinear simulations, showing that it can accurately learn the intrinsic dynamics shared between neural and be- havioral modalities. We then apply BRAID to motor cortical activity recorded during motor task and demonstrate that our method more accurately fits the neural-behavioral data by incorporating measured sensory stimuli into the model and improves the forecasting of neural-behavioral data compared with various baseline methods, whether input-driven or not.
INTRODUCTION
Understanding the relationship between neural activity and behavior is critical goal in neuroscience and neurotechnology. Neural activity and its temporal structure, or "dymamics" during behavior are formed by the interplay between (1) the recurrent networks within a brain area, i.e., intrinsic dynamics, and the (2) temporally-sstrucuree inputs it receives during the behavior (Remington al.) 2018 Vyas et al., 2020). A neural population may receive inputs from measurable sources such as sensory stimuli, cloctricalooptogenetie neurostimulation [Buonomano & Maass 2009 Seely al. 2016 Susilaradeya et al.) 22119; Sauerbrei et al.) 2020): Shenoy &
Yang et al. 202 Vahidi et al.] 2024), well as from other upstream brain areas (Sauerbrei et al.; 2020 Shenoy & Kao 2021), which could be included in multi-regional recordings (Jun et al. [017]] Steinmetz et al., 2019). However, even easily measurable external inputs (e.g., sensory stimuli) are often not explic- itly considered when modeling neural-behavine activity, which can lead to a conflation of intrinsic and input-driven contributions, creating challenges for interpretation (Seely et al. 2016 Sauerbrei et al. 20200: Vahidi et al. 2024). Beyond disentangling intrinsic dynamics from input dynamics, incorporating measured inputs into models can also enhance the behavior decoding performance in neurotechnologies such as stimulation-base closed-loop controllers (Yang et al. 2018 2021b).
Another challenge to disentangle neural dynamics that are relevant to a specific behavior from other neural dynamics, and to prioritize the former. This is critical as the majority of neural variance
may not be relevant to the behavior of interest (Churchland et al.) 2012 Mante et al. 20013 Kobak
et al.
2016 Allen et al. 2019 Engel & Steinmetz] 2019 Stringer al. 2019 Sani et al. 2021
While most prior works use unsupervised methods when modeling neural activity as latent variable dynamical systems (Aghagolzadeh & Truccolo 2015; Gao al. 2016 Wu al. 2017 Pandarinath jct al. 20188 Hernandez et al. [2000] Rutten et al. 20200 Kim et al. 2021), recent works have shown
improved learning of behaviorally relevant neural dynamics by using behavior data during learning

---

Published as a conference paper at ICLR 2025
in supervised manner (Sani et al., 2021 Hurwitz et al. 2021 Kramer et al.
2022; Gondur et al.
2024 Vahidi et al., 2024; Abbaspourazad et al. 2024 Sani et al. 2024 Oganesian et al. 2024).
Yet another challenge is posed by the nonlinearities in behualbbaaavorall data. While linear models
have been extremely effective in approximating neural dynamics (Hastie et al. 2009; Churchland
Kobak et al. 2016 Abbaspourazad et al. [021] Sani et al.] (021), , they may require higher dimensional latent repre-
et al., 2012 Mante et al. 2013) Cunningham & Yu 2014
Kao et al. 2015
sentations compared to nonlinear models (Yang et al. 2019 Nozari et al. 2024 Sani et al. 2024), and do not provide interpretability for nonlinear dynamical phenomena such as multi-stable fixed points and limit cycles (Kim al., Durstewwit et al., 2023). Moreover, unlike linear models, for nonlinear models the relationship between the intrinsic dynamics and an inference model that is fitted to estimate the latent states from observations is not analytically known, posing a challenge for studying intrinsic dynamics (see section 3.1.
Here, we address all aforementioned challenges by introducing Behaviorally Relevant Analysis of Intrinsic Dynamics (BRAID), a new method with the following key contributions. First, BRAID captures complex nonlinear structures in ceural-behavioral--nnur data, offering greater expressiv- ity than linear methods. Second, by optimizing multi-step-ahead forecasts of neural-behavior data, BRAID simultancously learns two representations for neural dynamics: the predictor and the gener- ative form representations (see section (..l). the latter of which describes intrinsic dynamics. Third, by explicitly modeling the influence of measured inputs, BRAID disentangles their dynamics from intrinsic dynamics more closely reflect the neuronal networks within the recorded brain region. Fourth, we introduce a multi-stage learning framework that dissociates and prioritizes the learning of intrinsic behaviorally relevant neural dynamics, while considering measured inputs (see section 3.2 Fifth, we introduce additional preprocessing and post-hoc learning stages that allow behavior- specific dynamies be dissociated from behaviorally relevant neural dynamics (see section 3.3)
We validate BRAID in multiple simulated datasets with distinct nonlinear structures and show its capability to accurately learn the underlying nonlinear model, resulting an interpretable represen- tation of intrinsic dynamics. We then apply our method to electrophysiologica data recorded from a non-human-prin -armaat (NHP) performing sequential reaches (O'Doherty et al. (017)) Our results indi- cate that accounting for both nonlinearity and sensory inputs improves behavioraa prediction, suggesting a more accurate representation of intrinsic behaviorally relevant neural dynamics.
2 RELATED WORK
Our work addresses multiple problems simultaneously, which makes it related to various methods that tackle a subset of these problems. A summary of related methods is provided in table [
First, a key ability of BRAID is to incorporate measured inputs to disentangle intrinsic dynamics from input dynamics. Other nonlinear modeling methods (Gao et al.] 2016, (Sussillo et al. 2016
Wu et al. 2017: Pandarinath et al. 2018 Rutten al., 2020; Hurwitz et al.) 202 Kim et al. Abbaspourazad et al.] 2024 Sani et al. 2024)) have not addressed modeling measured external inputs and their impact on neural-behavion
2021
data. As demonstrated by Vahidi et al. (2024). not consider- ing external inputs can lead to the dynamics of these inputs being misinterpreted as intrinsic neural dynamics. To overcome this challenge, |ahidi et al.) (2024) introduce linear dynamical model- ing method, termed IPSID, which explicitly incorporates measured external inputs into the model. However, IPSID is an analytical, and strictly linear method that cannot capture nonlinearities. By incorporating the strengths of this linear modeling work into BRAID, we can account for measured external inputs and dissociate their dynamies while allowing every model element to be nonlinear. We use IPSID as a key baseline to show the benefit of enabling nonlinearity in our method (see appendix (..2 for details). We also show the results for the special case of setting all model elements as linear in our method (referred to lincar BRAID). which fits in a linear model similar to IPSID.
Second, a key capability of BRAID is that it dissociates behaviorally relevant neural dynamics into a distinct part of the latent states and prioritizes their learning, while also being able to learn neural- specific and shhaiior-specife dynamics using additional latent states. Among prior works, two recent nonlinear methods termed DPAD (Sani et al. 2024) and TNDM (Hurwitz et al. 2021) aim to dissociate behaviorally relevant dynamics from other neural dynamics, but neither method dissoci- ates the third category of dynamics, i.e., the behavior-specific dynamics. More importantly, neither
2

---

Published as a conference paper at ICLR 2025
DPAD nor TNDM incorporates external inputs into the model to dissociate intrinsic dynamics from input dynamics. Finally. DPAD learns models based on 1-step-ahead prediction of eeurra-behaviona data and does not explicitly learn the intrinsic dynamies, whereas BRAID adds m-step-ahead pre- dictions into the loss to optimize forecasting and also explicitly learns a generative representation of intrinsic dynamics. TNDM on the other hand is sequential autoencoder (similar to L.ADDS, Pandarinath et al. 2018), i.c., it optimizes reconstruction of a window of data after ingesting the entire window as input. We compared our results with both DPAD and TNDM, although DPAD's architecture is closer to ours. In fact, the comparisons with DPAD can also be thought of as ablation studies that show the benefit of incorporating external inputs and forecasting in our method.
Third, we learn behaviorally relevant neural dynamics, or in other words the shared neural- behavioral dynamics, in an initial optimization focused on learning these dynamies, while leaving the learning of other neural dynamics to a separate subsequent optimization. This approach, priori- tizes behaviorally relevant neural dynamics in the sense that we can fit models with low dimensional latent states that are purely focused on these dynamics (Sani al.
2021 Besides DPAD (Sani al., 2024), a few other works, including TNDM, propose nonlinear approaches for learning dynamics shared between two modalities (Hurwitz et al., 2021 Kramer al. 2022 Gondur al. 2024). How- ever, these works use a combined loss to optimize the reconstruction of both modalities in the same optimization. While this approach can capture the dynamics shared between modalities, it does not prioritize them over dynamics specific to either modality (Sani et al., 2024). Moreover, most multi- modal approaches do not model the effect of external inputs (Hurwitz 2024). One multi-modal model, termed mmPLRNN (Kramer et al. (002)). which models dynamics of two modalities with a piecewise- linear RNN (see appendix A.. for details), supports modeling the effect of external inputs, although this capability was not demonstrated in Kramer et al. (2022). Nevertheless, we include comparisons with mmPLRNN with input as one baseline. Finally, as an- other ablation study to assess the importance of prioritizing behaviorally relevant dynamics, we also implement an unsupervised version BRAID, termed U-BRAID, that removes the behaviorally relevant optimization step and instead learns all neural dynamics in optimization step while still incorporating external inputs into the model (see section and appendix AA.2 for details).
2021] Gondur et al.,
Most other prior nonlinear methods only consider neural signals during modeling without consider- ing behavior or external inputs (Gao et al. 2016: Pandarinath et al., 2018 Hernandez et al. 2020
Rutten et al. 2020 Kim et al. [222] do not dynamic models (Zhou & Wei [220 Schneider et al. (023) and thus are vastly different from our method. Nevertheless, we include comparisons with LFADS (Pandarinath et al., (218)) and CEBRA (Schneider et al., (2023) as additional baselines.
We list the differences of some of these methods with our method in table
Table 1: Related works (see Discussion). ELBO: evidence lower bound, LL: log- log-likelihood.
Prioritize behaviorally relevant
Dissociate non-neural
Dissociate intrinsic
Method
Nonlinear
Training objective
IPSID
Projection-based Multi-modal ELBO
TNDM
xx
X
LFADS
X
ELBO
mmPLRNN
V
Multi- ELBO
DPAD
X
X
I-step-ahead LL
CEBRA
X
X
Contrastive
BRAID
V
m-step-ahead L.I.
3 METHODS
3.1
BRAID MODEL.
We model the neural activity (y& E IR'v) and behavior (zy E IR²') as observations of nonlinear dynamical system that is driven by measured (uę E R=' and/or unmeasured inputs (W; E R²)) )
x A (x + K (ux) We
Yk =
Cy Us) + Vb
(1)
Zk -x(x us + Ek
3

---

Published as a conference paper at ICLR 2025
Here, x3 E I== represent the latent states of the system and evolve according to intrinsic dynamics Afw. V& E 1 and EK E R" represent observation noises. Given this dynamical system, one can recursively infer the latent state from neural observations Yı using RNN follows
A(x + K (yk, Up)
(2)
where Xs+1|k (or simply x is defined as the inferred latent state based on {yi,,..,, Yk} and {ui,..., ug.. Given the latent nature of the states, even when inference is optimal (e.g.. in a Kalman
filter), the inferred states will not be equal to the internal states x in equation T (Katayama 2006 which is why we use different notations for the states in equations [ and 2 More importantly.
note that A and K in equation 2 are distinct from Af and K fw in equation I This is because A and K represent the "predictor form" representation of dynamics, describing how the inferred latent state recursively evolves over time as samples of Yı and u₁ are observed, whereas f and K represent the "generative form" representation of dynamics that describe how the latent states themselves evolve, purely based on their intrinsic dynamics - so f what ultimately describes the intrinsic dynamics. For linear systems, there is an analytical bidirectional relationship between predictor and generative form representations (defined by the Kalman filter, see [Katayama 2006), whereas for nonlinear systems in general, this relationship is not known. Thus, we devise an approach that allows us to learn both representations of dynamics from data.
Critically, to predict the latent state (or neural-behavioral data) multiple (m >> 1) steps into the future without new neural observations and using only new inputs ug, we need to propagate the latent state ahead according its intrinsic dynamics, i.., the "generative form" representation of dynamics, as
Af 11k) + f (ux+m-1)
X =
where XR+m|jk denotes the latent state at time step A: + m, generated given {y,, y& and {Ug,... Uc+m-1)- Note that for m = 2. the right hand side of equation 3 would have Ijki which is given by equation 2 Thus, m-step-ahead inference of the latent state engages both the pre- dictor and generative form representations of the dynamics via equations [ and [] respectively. As an alternative interpretation, the m-step-ahead prediction of the latent state (or neural or behavioral data), for m > 1, involves two RNNs operating in complementary fashion (figurep):
The first RNN (RNN, parameterized by A and K) takes in neural and input time series and recursively estimates the 1-step-ahead prediction (Xxjk-1).
2. The second RNN (RNN far, parameterized by Af and K fww) takes in the 1-step-ahead predicted state from the first RNN and propagates it m mn-1a 1 additional steps ahead according to the intrinsic latent dynamics of the model, to get the m-step-ahead predictions (X for m >> 1).
Overall, the BRAID model is comprised of six distinct transformations: A(-), (-) K()), K f.)( C.()), and (.). A/Afu describe predictor/generativy form recursions of the latent state. K/K fur describe predictor/generatiiv form encoders. C2 and Cy describe behavior and neural de- coders. We implement these six transformations as multi-layer perceptrons (MIPs) with arbitrary user-specified number of units and hidden layers. As special case, any (or all) of these mappings can be replaced by a linear mapping (i.c., an MLP with no hidden layer and a linear activation).
We learn the parameters specifying all six transformations of the model by optimizing a weighted sum of m-step-ahead behavor-a prediction errors (for m E [m1, m2...., m.....,m_) as our losses
L2 1==3a Σ/= a.mm MSE(Z)+m, C2(Xx+ (ky Uk+m.)) Ly = D
1 O'ymm MSE(yk+m. Cy [X++m//k: Uk+m,))
where MSE(.) indicates the mean-squared error loss, L. denotes the number of steps ahead simulta- neously included the loss, and a'am; and 0'ym; denote the weights used in the sum. In this work, we always set Osm, and Oyym, to 1. Moreover, although the decoders in BRAID can optionally take both the latent state and the external input uy (lines 2-3 of equation D in our real data analyses we do not provide u₂ to decoders and generate predictions only based the latent states.
3.2 PRIORITIZATION OF BEHAVIORALLY RELEVANT OVER OTHER NEURAL DYNAMICS
To dissociate behaviorally relevant neural dynamics from other neural dynamics and prioritize the former, we break the latent state X into sections (x(²)
(1) and (2)) and learn these sections
4

---

Published as a conference paper at ICLR 2025
(a)
Stage 2
(b)
RNN1
Neural activity
y
Neural
decoder
1-dlep-cheed
Encode
y Neural activity
Stage
Input
AAe
ut
K
C.
Exrtue-1
Behavior decoder
2-step-ahead
2g Behavior
8
Encoder
Apri
K
Mess
3-step-shead
Apey
Behavior-specifi (non-encoded)
RNN1
Tw
Figure 1: BRAID model architecture. (a) BRAID dissociates the dynamics of neural- neurao-behavioral data into three latent states x x and X (3). the dynamics shared between neural and behav- ioral modalities (learned in stage 1 by RNNI and RNNIJW), 2) any remaining dynamics private to neural activity (learned in stage 2 by RNN2 and RNN2 ww). and 3) Input-driven, behavior-specifee dy- namics not encoded in neural activity (learned per section 3.3 by RNN3 and RNN3 (b) For each latent state, we simultaneously learn predictor and generative form representation of the dynamics (denoted by A and A fw), by optimizing m-step-ahead prediction of interconnected two-RNN system is visualized for RNNI and RNNI fw in this computation graph. The superscript..... (1) indicating that parameters are for RNNI and RNNI fw is omitted for simplicity.
data. This
in two learning stages. We denote the model parameters associated with each model section using (1) or (2) superscript, e.g., A1)) and (2)) We provide the full two-section formulation for the model in appendix [.] and the optimization details in appendix A.1.22 Briefly, each of the two learning stages consist of optimizations, follows:
Stage 1: Learning RNN/ and RNNI
(1) Ia Learn A(1, A fwo ((1), K fwr and C( and extract latent states x and x (1)
[4mjk (for 1)
by minimizing the behavior prediction loss Lg from equation
Learn c by predicting neural data from xx tion loss Ly from equation 4
(1) and x k+| while minimizing the neural predic-
(1)
Stage 2: Learning RNN2 and RNN2
(2) 2a Learn 12), A fw (2), K fw. and Cy and extract latent states x and X+ (for m 1)
(2)
by minimizing the neural loss Ly. while including outputs of stage 1b part of the predictions.
2b Learn ((22) by predicting behavior from x and (2) k+mlkk> while minimizing the behavior loss Lg from equation 4 and including outputs of stage la as part of the predictions.
The explicit dissociation of the relevant dynamics and the above two-stage optimization allow us to first preferentially learn the now-dimensional) shared dynamics between the two observations, i.e.,
the behaviorally relevant neural dynamics (x (1)), ) in stage 1. Then in stage 2, we learn any residual neural dynamics (x(2)), which, depicted in figure 1 can depend the behaviorally relevant dynamic (see appendix [ for details). The optional stage 2 is of interest for explaining neural dynamics beyond the ones related behavior. This multi-stage approach has similarities Sani et al. (2024), but here we have: 1) additional signals (uę). 2) different losses, 3) a forecasting RNN within each model section (figure Tp), and additional steps that are discussed in the next section.
3.3 DISSOCIATION OF DEHAVIOR-SPECTIE DYNAMICS
Optimizing behavior prediction (stage la) given neural activity ys and input uş can lead to learning behavior dynamics that are predictable from the input but not encoded in the recorded neural activity. Although learning such rehavior-specifie dynamics enhances behavior decoding, it poses an interpretation challenge for neuroscience applications because would not know what part of
5

---

Published as a conference paper at ICLR 2025
the learned dynamics are represented in the recorded brain regions. As shown in Vahidi al. (2024), this may lead to a misinterpretaatonn of input-driven behavior-specific dynamics as intrinsic dynamics of the recorded brain region. To mitigate this possibility, we develop two additional steps in our method, that can 1) exclude such behavior-specific dynamics from x and 2) learn them separately
distinct latent state x (figure [p.. Details are provided in appendix A.1.3 Briefly, first, to exclude behavior-specific dynamics, we introduce an optional preprocessing stage that predicts behavior from neural data, and passes this neurally-predicted behavior be used in stages la and
2b. This preprocessing step ensures that the behaviorally relevant states learned in stage 1 (xx (x encoded in recorded neural activity yk. which can be crucial for interpretability in neuroscience studies. In our analyses of the real datasets (section 4.2), we always include this preprocessing step. Second, to still be able learn behavior-specifie dynamics, we add an optional post-hoc learning step (i.e., stage 3) that fits RNN3 and RNN3 fw to any unexplained behavior and learns these input-
driven behavior-specift dynamics as a distinct latent state x(3) As we show in simulations (see section I and figure f.gure the preprocessing step can exclude non-encoded behavior dynamics. When desired in an application, the optional stage can learn such dynamics to offset any behavior decoding loss incurred due to the preprocessing step, while still maintaining the interpretability of the model. We did not apply this post-hoc step in our real data analyses (section 4.2)
3.4 INFERENCE AND EVALUATION METRICS
After learning BRAID's parameters, we can readily use the learned mappings A, K (and A and
KJ fww) to infer the I-nand multii-step-ahead predicted states X5 (and X using equations 2 (and
3) for the held-out test data. Predicted neural activity and behavior are obtained by applying their
corresponding decoders Cy and C2 these inferred states. We also use the term "decoding" for behavior predictions because our model predicts behavior only using neural data and inputs, and never using behavior itself. To evaluate the performance of our models, we perform and 2-fold valssvvaliattion, for real data and simulation analyses, respectively. We report Pearson's 's Correlation Coefficient (CC) and in some cases (see table (. also the coefficient of determination (dR²) between the predicted and actual observation, averaged over dimensions. We further report the m-step-ahead prediction accuracy, which reflects how well the intrinsic dynamics Af w are learned. For simulation analyses with linear recursions (Afw), we additionally evaluate the learned intrinsic dynamics by comparing the eigenvalues of f between the true and learned model (see appendix A.1.)).
4
EXPERIMENTAL RESULTS
4.1 SIMULATION EXPERIMENTS
We validated BRAID three simulations with different nonlinear ceuraalbehavioral-ae structures
to show that it can learn intrinsic behaviorally relevant neural dynamics in presence of inputs.
4.1.1 BRAID ACIIEVEE NEAR OPTIMAL NEURAL-BEHAVIOOR PREDICTIVE ACCURACY IN NONLINEAR INPUT-DRES DRIVEN SIMULATIONS
First, we considered an input-driven dynamical system as in equation | but with only the behavior mapping C2 being nonlinear with the mapping Jc, (v) == sin(v) + as detailed in section A.4.2 (figure b). We generated 10 random parameter sets as our true models and generated data from them. First, we implemented an automatic selection of nonlincarity for BRAID by setting each of A. K. Cy, or C₂ to linear or nonlinear, resulting in 24 different BRAID models, and finding the model with the best behavior decoding in the training data. Across all realizations and 2 cross- validated folds, setting the behavior decoder C₂ to be nonlinear was correctly identified as the best performing nonlinearity in 100% of the cases. Additionally, we evaluated BRAID with nonlinearity only in one of A, K, Cy. or Cz. The model with nonlinear behavior decoder C₂ outperformed other nonlinearity choices as well as linear models (i.e., IPSID and the fully linear BRAID) in behavior decoding and neural prediction. BRAID further outperformed DPAD, which is nonlinear but does not account for the input 4g in neural-beite behavioral prediction. In fact, both BRAID with nonlinear C2 and BRAID with automatic nonlinearity selection achieved almost the same neural- heural-behavioral

---

Published as a conference paper at ICLR 2025
prediction the true simulated models, demonstrating BRAID's success in accurately learning the nonlinear input-driven dynamical system (table
Table 2: I-step-ahead prediction results for nonlinear simulations with sinusoidal behavior mapping.
Mean + s.e.m. is across 20 runs (10 datasets, folds). State dimension is always set to ground truth.
True model's outcome indicates the "Ideal" accuracy. Bold: within 1 s.e.m. of ideal.
Method
Behavior decoding CC
Neural prediction CC
IPSID
0.4567 = 0.0527
0.8901 = 0.0304
linear BRAID
0.4558 + 0.0528
0.8893 + 0.0304
DPAD Nonlin C₂
0.4958 + 0.0620
0.3767 0.0680
BRAID Nonlin
0.4735 + 0.0572
0.8127 + 0.0528
BRAID Nonlin K
0.7680 = 0.0467
0.6983 0.0473
BRAID Nonlin Cy
0,4558 + 0.0528
0.8887 = 0.0305
BRAID Nonlin C₁
0.8696 + 0.0487
0.8913 0.0305
BRAID Auto Nonlin
0.8693 + 0.0487
0.8913 = 0.0305
True model (ideal)
0.8737 + 0.0486
0.8921 = 0.0306
4.1.2 BRAID DISSOCIATES INTRINSIC DYNAMICS FROM INPUT DYNAMICS IN SIMULATIONS
Next, we sought to validate BRAID's ability to disentangle intrinsic and input-driven contributions neural dynamics. BRAID simultaneously learns a predictor form and a generative form represen- tation of the dynamics, the latter of which directly describes the intrinsic dynamics in terms of the mapping A f (-) (section 3.1 figure [p.. In this simulation, we kept the ground truth state transi- tions linear so that we could precisely quantify the intrinsic dynamics and their learning error via the fo (appendix [A.1.5]. We analyzed data from three sets of simulated dynamical systems with distinct nonlinear structures (see appendix A.4 for details): (1) Spiral behavior manifold (figure 2b-d), (2) trigonometric behavior manifold (also explained in section 4.1.1 figure h2- and (3) trigonometric input-encoder (figure A.I). For each simulation, we generated realizations from 10 different systems with randomly generated sets of parameters. Across all three simulations, BRAID accurately learned the intrinsic dynamics, resulting in smaller error in eigenvalues of the transition matrix compared with DPAD, which is nonlinear but does not consider the input, and compared with linear BRAID and IPSID, which consider input but are linear (figures 2.g If). We demonstrate in an ablation study that learning a separate generative model for forecasting crucial for correct learning of the intrinsic dynamics (table A.3) These more ac- curate intrinsic dynamics coupled with the input also resulted in BRAID achieving better behavior decoding (figures 2p.f [A.Ip) and neural prediction (figure (A.Ip) compared baselines. These re- sults suggest that failing to account for either nonlinearity or input may lead to less accurate models and misinterpretation of the intrinsic dynamics in nonlinear data.
eigenvalues of the state transition matrix f
Another metric for how well intrinsic dynamics are learned is forecasting, where behavior is pre- dicted multiple steps into the future, without observing new neural data and only by observing the future input (section 3.1.. Forecasting evol ves the state dynamics according to the learned intrinsic dynamics (equation 3) and thus validates their accurate learning. We performed forecasting up to 32 steps ahead and found that the nonlinear models with input consistently outperformed linear models as well as the nonlincar DPAD, which does not consider input (figures 21.h A.I.f.
4.2 NON-HUMAN PRIMATE MOTOR CORTICAL ACTIVITY DURING REACHING
We applied our method to a publicly available dataset recorded from a non-human primate (NHP) performing reaching movements O O'Doherty et al. (217)) (figure 3p.. We took either the smoothed spike counts or raw LFP from primary motor cortex (M1) as neural time-series yı. fingertip's po- sition and velocity as behavior Zk. and sensory task instructions (target location) the input u, (appendix A.3). Sensory inputs can have their own dynamics, which are distinct from the intrinsic dynamics of the motor cortex. Our goal is to learn the intrinsic dynamics in M1 related to movement while disentangling them from the dynamics of sensory input and also from any behavior behavior-specific dynamics. Therefore, we importantly include BRAID's behavior preprocessing stage (section 3.3)
7

---

Published as a conference paper at ICLR 2025
(a)
Spiral behavior manifold
log, normalized EV error
0.5
0.5
Decoding CO
0.4
Time
0.3
-0.5
Decoding
0.3
-1
10%
0.2
0.1
0.1
-1.5
1,
# steps ahead
(e)
(
(9)
(h)
Trigonometric behavior manifold
log, nomalized EV error
0.8
-0.5
Decoding CC
Decoding C8
0.6
10%
0.6
0.4
-1.5
0.4
0.2
1%
0.2
0
BRAID Nonlinear Cy (steps-ahead-(1.22 (stepe-ahead-(1.25)
DPAD Nonlinear Cy
Linear BRAID
IPSID
# steps ahead
(steps-ahead=11123..
Figure 2: BRAID better learns the intrinsic shared dynamics by simultaneously modeling in- put and nonlinearity, and by optimizing forecasting. (a-d) Results for simulations with spiral behavioral manifold.
(b) I-step-ahead behavior decoding for nonlinear BRAID, nonlincar DPAD, linear BRAID, and IPSID (c). Error in identifying intrinsic dynamics the true model quantified by the error in learning the eigenvalues of fw- (d) Behavior decoding forecasts for to 32 steps ahead, enabled by learning the intrinsic dynamics (A) with predictions optimized for [1, 2, 3, 4, 5|-steps- ahead (section) 3 (e-h) Same (a-d) for simulations with trigonometrie behavior mapping.
We fitted BRAID with different nonlinearity choices as in our simulations: (1) nonlinear recursion A(-)AAp fm (.) (2) nonlinear encoder K (-))/jw((.), (3) nonlinear decoders C-) and Cz(),, and a fully linear variant, linear BRAID. We included m = [1. 2, 4, -s--es--ahhadd predictions in the BRAID loss, which for m engage the intrinsic behaviorally relevant dynamies (A fw) and allow their learning. To evaluate the learned intrinsic dynamics, we report neural-behavioraa forecasting accuracy for different step-ahead horizons (figures 3p-c, and A.3 for LFP),, while tabulating the 4- steps-ahead (i.e., 200ms) results to highlight BRAID's advantage in forecasting (tables 3 and A.4
In addition to these results in the liw-dimensional regime (stage only, ng n1 16), we also report the results in the high-dimensionall regime (both stages, 71 64,
16) (table 3 figures
A.5 and A.6). Among nonlinearity configurations, BRAID with nonlinear decoders provided the best fit to ceural-behavioral data (table Moreover, BRAID's behavior forecasting performance improved as more neurons were included (table A.7).
Next, compared BRAID's neural-hehavioral forecasting several ablation baselines (table and figures and A.5). First, BRAID outperformed linear BRAID, i.e., a similar but fully linear model. Second, BRAID outperformed DPAD (Sani al. (2024), which can have decoder nonlin- carities but does not consider inputs. BRAID's advantage shows the importance of considering the effects of sensory inputs neural-behavioral dynamics. Third, we compared U-BRAID, which removes the first stage of BRAID and thus loses prioritization. BRAID consistently outperformed U-BRAID in behavioral forecasting, but did not match the neural forecasting of U-BRAID unless 64). which is expected given U-BRAID's
was given enough latent state dimensions (table 3 singular objective being neural prediction. This comparison shows the benefit prioritization for learning low-dimensional representations of intrinsic behaviorally relevant dynamics, and confirms that with its stage 2, BRAID can capture any remaining -on-behavioral neural dynamics. Finally, BRAID's low dimensional latent state trajectories were better separated for different movement di- rections compared to those of U-BRAID and DPAD, suggesting that BRAID's 's latent states are more congruent with behavior, i.e., more behaviorally relevant (figure A.4)
We also compared BRAID's s neural- soural-behavioral forecasting performance to mmPLRNN which mod- els multi-modal data using piecewise-linear RNNs, and has the option to model inputs although prior
8

---

Published as a conference paper at ICLR 2025
Table 3: Forecasting performance (4-step-ahead) compared to basclines in NHP dataset for models
with low (ng = 16) and high-dimensional (z... -0,==6 64) latent states, 71 = for BRAID, linear BRAID,
and DPAD. R results were similar (table A.8 Tables A.4 A.5 and A.6 have additional results.
Behavior forecasting CC
Neural forecasting CC
Method
72 16
f2 = 64
T2
16
fg = 64
linear BRAID
0.7453 + 0.0066
0.7409 = 0.0059
0.1767 + 0.0054
0.3784 + 0.0078
DPAD
0.6706 + 0.0096
0.7352 0.0079
0.2067 0.0062
0.3611 = 0.0080
U-BRAID
0.7663 + 0.0069
0.8049 = 0.0068
0.4089 + 0.0076
0.4185 = 0.0074
mmPLRNN
0.6851 + 0.0143
0.7328 + 0.00361
0.3162 0.0107
0.3570 + 0.0223
BRAID (ours)
0.8042 + 0.0085
0.7970 = 0.0086
0.3274 0.0078
0.4123 + 0.0077
work had not explored this input option. BRAID outperformed input-driven mmPLRNN networks in forecasting both behavior and neural activity, across all forecasting horizons, with the exception of 2 steps-ahead neural prediction (table figure B) Note that BRAID's better decoding is achieved despite the fact that mmPLRNN, by design, incorporates behavior as an input during inference, whereas BRAID and DPAD do not. Also, we compared BRAID to TNDM (Hurwitz et al. 2021), nonlinear sequential autoencoder that models neural-behavioral dynamics, but does not include the effect of external inputs. We extended TNDM beyond the original work to create a version that also adds sensory inputs as an additional input besides neural activity (appendix(A.22 (.2.66.. We analyzed non-smoothed spike counts from the same dataset and found that BRAID significantly outperformed TNDM in behavior decoding while achieving comparable neural prediction (table A.5). Extending TNDM to include inputs significantly improved its behavior decoding, but it still did not reach that of BRAID (table A.5). Finally, we compared BRAID with CEBRA (Schneider et al. (2023). which non-dynamic convolutional encoder with contrastive loss behavior (section A.2.8), and with LFADS (Pandarinath et al.] (018), which is an unsupervised sequential autoencoder (section A.2.7
BRAID outperformed both these baselines in ceural-behaviorll prediction (table A.60)...6)...
(a)
Smoothed spike counts
Target position u,
Cursor position and velocity 8,
2s
(b)
(c)
0.8
Neural N prediction CC
0.4
0.6
Decoding
0.4
0.2
0.2
2
8
16
32
8
16
32
steps ahead
steps ahead
BRAID Nonlinear decoders (steps-ahead-(1.2444
DPAD Nonlinear decoders (steps-ahead(112243))
Linear BRAID (steps-ahead=(112 (1,244,8))
mmPLRNN Nonlinear decoders
Figure 3: BRAID outperforms baselines in neural-behavioral forecasting. (a) Dataset and task visualization. (b, c) Behavior and neural activity forecasting correlation coefficient (CC) for BRAID, linear BRAID, DPAD, and mmPLRNN at low state dimension regime (n,,=16). Shaded areas show the s.e.m. across the 7 recording sessions and cross-validation folds.
DISCUSSION
We introduced BRAID, method for input-driven nonlinear dynamical modeling that disentangles intrinsic shared dynamics between two observation modalities from the effect of imput. Here, we assume external inputs are measured (e.g., from sensory stimuli other brain regions) and

---

Published as a conference paper at ICLR 2025
available for modeling. This approach is distinct from the input-inference approach, where unmeasured inputs are inferred from measured neural activity (Pandarinath et al. 2018; Schimel et al. 2022). These two approaches are in sense complementary. In our approach, any measured inputs can be explicitly incorporated into the model to dissociate their dynamics from the intrinsic data. Practically, one cannot measure all inputs
dynamics of the measured neural-behavions given brain area, so our approach does not rule out the influence of unmeasured inputs on the learned intrinsic dynamics. One could thus use the input-inference approach to infer such unmeasured inputs from all measured signals. Note, however, that inferred inputs are ultimately function of measured signals and thus do not add any new information (unlike measured inputs): rather they can be thought of as a decomposition of the measured signals based on certain assumptions (e.g., smoothness).
While BRAID's learning is supervised by behavior, it only uses neural activity and input (not be- havior) during inference. This supervision allows stage 1 to extract behaviorally relevant intrinsic dynamics with priority, while later stages learn neural-specific or behavior- sppriifo---pefiiccceeeee.ooeeeee dynamics in in- dependent optimizations. This multi-stage learning has similarities to some prior works (Sani al.
2021 Vahidi et al.) 2024, Sani et al. 2024] Oganesian et al. 2024), but is fundamentally different from other works that use single
optimization loss (Kramer et al. 2022 Gondur et al., 2024 Hurwitz et al., (02T). which may miss prioritization of shared dynamics over unshared dy- namies (Sani et | 2024). Some of the latter group are indeed focused on multi-modal inference and fuse all shared and unshared information into the same latent space (Kramer et al.) [022 Gondur ct al., 2024). Multi-stage methods avoid this fusion by focusing on shared dynamics during a first optimization stage with only cross-modality prediction (e.g., behavior decoding) as the objective.
To disentangle input dynamics from intrinsic dynamic, BRAID consists of three stages, each learn- ing a predictor and a generator model. In each stage, BRAID's predictor model in that stage (e.g., RNNZ) infers the latent states, which are subsequently employed as the input to compute m-step- ahead predictions via the associated generative model (e.g.. RNNI fw.) BRAID's predictor and gen- erative models are learned jointly to maximize the m-step-ahead log-likelihood. This has analogies to the mecder-decoder architectures such as those commonly used in variational inference (Kingma) & Welling, 2013). Specifically, the predictor and generator RNNs in BRAID have roles similar to those of the encoder and decoder in variational inference, respectively. However, in variational in- ference, the posterior distribution of the unobserved variables given the data is parametrized and a part of the optimization loss aims to enforce that distribution
the inferred latent variables (Chung jet al. 2013; Krishnan et al., 2015) Fraccaro et al., 2016 Luk et al.] (2024). In contrast, we do not impose such parametrization the latent states in BRAID. Developing variational methods with the same multi-section architecture and multi-stage learning as BRAID is an interesting future direction.
Similar to many prior works including some in neuroscience (Abbaspourazad et al. 2024| Sani al., 2024 Chung et al.) (2155] Krishnan et al. 20155] Fraccaro et al. |016 Luk et al.) (024). BRAID has a causal formulation and performs inference by recursively inferring the next latent state after each
new observation. This is distinct from some nonlinear methods in neuroscience (Gao et al. 2016
Pandarinath et al.) 2018: Hernandez et al.] 2020 Hurwitz et al.] 2021 Keshtkaran et al. 2022
Gondur et al. 2024 Karniol-Tambour et al. 2024 that perform inference non-causally in time.
Thus, BRAID may be useful for real-time decoding in brain-computer interfaces (Shanechi 2019).
BRAID, as presented, is designed for single-session settings. Extending it for cross-session gener- alization is an interesting future direction that can follow approaches used in the literature, such as aligning neural manifolds across sessions (Farshchian et al. 2019), adaptive modeling approaches (Ahmadipour et al., I Yang et al., 202 or incorporating session-specific read-in and read- mappings while keeping shared model parameters fixed (Pandarinath et al. 2018). The latter allows training across multiple sessions and adapting to new sessions with minimal data by learning session-specific matrices. Furthermore, BRAID's s ability to disentangle intrinsic dynamics may also facilitate generalization across tasks with different sensory instructions.
A fundamental challenge for nonlinear latent state models is that many alternative models may plain the data equally well. Thus, evaluating learned dynamics relies on computable quantities from measured signals, primarily m--tee--ahead beurai-behavioral predictions. While our results suggest that in our dataset having nonlinear decoders provides higher performance than other nonlinearities, this might not be case in another dataset. In practice, the optimal nonlinearity can be selected based on the desired metric as done by BRAID, For linear models, all correct latent models have the same eigenvalues (Katayama, 2006), allowing direct evaluation of learned dynamics (figure
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
This work was supported, in part, by the following organizations and grants: National Institutes of Health (NIH) R01MH123770, NIH BRAIN Initiative R61MH135407, NIH RFIDA056402, NIH Director's New Innovator Award DP2-MH126378, Foundation for OCD Research (FFOR), the Of- fice of Naval Research (ONR) Young Investigator Program under contract 100014-19-1-211 and the Army Research Office (ARO) under contract 1911N--16-1-08
as part of the collaboration between the US DOD, the UK MOD and the UK Engineering and Physical Research Council (EP- SRC) under the Multidisciplinary University Research Initiative (MURI), We sincerely thank the Sabes lab at the University of California San Francisco for making the NHP datasets that we used here publicly available.
AUTHOR CONTRIBUTIONS
P.V., O.G.S., and M.M.S. developed the new BRAID method and wrote the manuscript. P.V. and O.G.S. implemented the code for the method and the analyses. P.V. performed the analyses. M.M.S. supervised the work.
8 REPRODUCIBILITY STATEMENT
To ensure the reproducibility of our work, we are sharing the code for BRAID along with Python notebook demonstrating its usage at wtfttps://giiie provide model architecture and training details in appendix A.1.4. Finally, the dataset used OODDoherty et al. 2017) is publicly available for anyone interested in reproducing the results re- ported in section 4.2
com/ /aa.eeehnaaabbaAAAD We also
REFERENCES
Hamidreza Abbaspourazad, Mahdi Choudhury, Yan T Wong, Bijan Pesaran, and Maryam M Shanechi. Multiscale low-dimensional motor cortical state dynamics predict naturalistic reach- and-grasp behavior. Nature communications, 12(1):607, 2021.
Hamidreza Abbaspourazad, Eray Erturk, Bijan Pesaran, and Maryam M Shanechi. Dynamical flex- ible inference nonlinear latent factors and structures in neural population activity. Nature Biomedical Engineering, 8(1):85-108, 2024.
Mehdi Aghagolzadeh and Wilson Truccolo. Inference and decoding of motor cortex low- dimensional dynamies via latent state-space models. IEEE Transactions Neural Systems and Rehabilitation Engineering, 24(2):272-282, 2015.
Parima Ahmadipour, Yuxiao Yang, Edward F Chang, and Maryam M Shanechi. Adaptive tracking of human ecog network dynamics. Journal of Neural Engineering, 18())016011, 2021. doi: 10. abae42
1088/1741-2552/abaed URLpttps:
doi. oww..o.r111... 10w.11108111111111
William E Allen, Michael Z Chen, Nandini Pichamoorthy, Rebecca H Tien, Marius Pachitariu, Liqun Luo, and Karl Deisseroth. Thirst regulates motivated behavior through modulation of brainwide neural population dynamics. Science, 364(643)):eavv339..
2019. 2019.
Dean V Buonomano and Wolfgang Maass. State-dependent computations: spatiotemporal process- ing in cortical networks. Nature Reviews Nearoscience, 10(2):113-122 2009.
Junyoung Chung, Kyle Kastner, Laurent Dinh, Kratarth Goel, Aaron Courville, and Yoshua Ben- gio. A recurrent latent variable model for sequential data. Advances in neural information pro- cessing systems, 28, 2015.
Mark M Churchland, John P Cunningham, Matthew T Kaufman, Justin D Foster, Paul Nuyujukian, Stephen Ryu, and Krishna V Shenoy. Neural population dynamics during reaching. Nature, 487 (7405):51-56 2012.
11

---

Published as a conference paper at ICLR 2025
John P Cunningham and Byron M Yu. Dimensionality reduction for large-scale neural recordings. Nature neuroscience, 17(11):500-1159 2014.
Daniel Durstewitz. A state space approach for piecewise-linear recurrent neural networks for iden- tifying computational dynamics from neural measurements. PLoS computational biology, 13(6): c1005542,, 2017.
Daniel Durstewitz, Georgia Koppe, and Max Ingo Thurm. Reconstructing computational system dynamics from neural data with recurrent neural networks. Nature Reviews Neuroscience, 24 (11):693-710, 2023.
Tatiana A Engel and Nicholas A Steinmetz. New perspectives dimensionality and variability from large-scale cortical dynamics. Current opinion in neurobiology. 58:181-190, 2019.
Ali Farshchian, Juan A. Gallego, Joseph P. Cohen, Yoshua Bengio, Lee E. Miller, and Sara A.
Solla. ADVERSARIAL DOMAIN ADAPTATION FOR STABLE BRAIN-MACHNNE IN- TERFACES. In International Conference Learning Representations, 2019. URL https: /
Marco Fraccaro, Søren Kaae Sønderby, UIrich Paguet, and Ole Winther. Sequential neural models with stochastic layers. Advances in neural information processing systems, 29, 2016.
Yuanjun Gao, Evan W Archer, Liam Paninski, and John P Cunningham. Linear dynamical neural population models through nonlinear embeddings. Advances neural information processing systems, 29, 2016.
Rabia Gondur, Usama Bin Sikandar, Evan Schaffer, Mikio Christian Aoi, and Stephen L Keeley. Multi-modal gaussian process variational autoencoders for neural and behavioral data. In Inter- national Conference on Learning Representations, 2024.
Trevor Hastie, Robert Tibshirani, Jerome H Friedman, and Jerome H Friedman. The elements of statistical learning: data mining, inference, and prediction, volume 2. Springer, 2009.
Daniel Hernandez, Antonio Khalil Moretti, Ziqiang Wei, Shreya Saxena, John Cunningham, and Liam Paninski. Nonlinear evolution via spatially-deeendeet linear dynamics for electrophysiol- ogy and calcium data, 2020. URL htt.:httss///eee httttps:/aaxxiv.oot
Cole Hurwitz, Akash Srivastava, Kai Xu, Justin Jude, Matthew Perich, Lee Miller, and Matthias Hennig. Targeted neural dynamical modeling. Advances in Neural Information Processing Sys- tems, 34:29379-29392. 2021.
James Jun, Nicholas A Steinmetz, Joshua H Siegle, Daniel Denman, Marius Bauza, Brian Barbar- its, Albert Lee, Costas A Anastassiou, Alexandru Andrei, Çağatay Aydın, et al. Fully integrated silicon probes for high-density recording of neural activity. Nature, 55 551(7679)
2017.
Jonathan C Kao, Paul Nuyujukian, Stephen
Ryu, Mark M Churchland, John Cunningham, and Krishna V Shenoy. Single-trial dynamics of motor cortex and their applications to brain-machine interfaces. Nature communications, 6(1):7759, 2015.
Orren Karniol-Tambour, David M Zoltowski, E Mika Diamanti, Lucas Pinto, Carlos D Brody, David W Tank, and Jonathan W Pillow. Modeling state-dependent communication between brain regions with switching nonlinear dynamical systems. In The Twelfth International Conference Learning Representations, 2024.
Tohru Katayama. Subspace Methods for System Identification. Springer Science & Business Me- dia, 2006. ISBN 978- 1-8-188-4888--1 158-7. URL htttp:://iiii springer. homm/book/C... 1007 - 158-X
Mohammad Reza Keshtkaran, Andrew R Sedler, Raeed H Chowdhury, Raghav Tandon, Diya Bas- rai, Sarah L Nguyen, Hansem Sohn, Mehrdad Jazayeri, Lee E Miller, and Chethan Pandarinath. A large-scale neural network training framework for generalized estimation of single-trial popu- lation dynamics. Nature Methods, 19(12):152-11577, 2022.
12

---

Published as a conference paper at ICLR 2025
Timothy D Kim, Thomas Z Luo, Jonathan W Pillow, and Carlos D Brody. Inferring latent dy- namics underlying neural population activity via neural differential equations. In International Conference on Machine Learning, pp. 5551-5561. PMLR, 2021.
Diederik P. Kingma and Jimmy Ba. Adam: A Method for Stochastic Optimization. 2017. URL
httt:://axiin org/ abs / / 1412. 6980
Diederik P Kingma and Max Welling. Auto-encoding variational bayes, 2013. URL https:
6114
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.
Daniel Kramer, Philine Bommer, Carlo Tombolini, Georgia Koppe, and Daniel Durstewitz. Re- constructing nonlinear dynamical systems from multi-modal time series. In Proceedings of the 39th International Conference on Machine Learning, volume 162, pp. 11613-11633, PMLR, 17- 23 Jul 2022.
Rahul G. Krishnan, Uri Shalit, and David Sontag. Deep kalman filters, 2015. URL https:
Enoch Luk, Eviatar Bach, Ricardo Baptista, and Andrew Stuart. Learning optimal filters using variational inference, 2024. URL /tttpss::///rxx wtgg/bb///0.......////// 18066
Valerio Mante, David Sussillo, Krishna Shenoy, and William T Newsome. Context-dependent computation by recurrent dynamics in prefrontal cortex. nature, 503(77744:77-84, 2013.
Erfan Nozari, Maxwell A Bertolero, Jennifer Stiso, Lorenzo Caciagli, Eli Combblath, Xiaosong He, Arun S Mahadevan, George Pappas, and Dani S Bassett. Macroscopic resting-state brain dynamics are best described by linear models. Nature biomedical engineering, 8(1):68-84, 2024.
Joseph E. O'Doherty, Mariana M. B. Cardoso, Joseph G. Makin, and Philip N. Sabes. Nonhuman Primate Reaching with Multichannel Sensorimotor Cortex Electrophysiology. May 2017. URL
https://dof orooorrgglo. /tlgg/00. 2tt////eeoo...o 58333 583331
Lucine I. Oganesian, Omid G Sani, and Maryam Shanechi. Spectral learning of shared dynamics between generalized-lines processes. Advances in Neural Information Processing Systems, 37: 89150-89183, 2024.
Dan O'Shea and Chethan Pandarinath. Lfads run manager, December 2021. httsss//:
2tt.100.52811eeoodo.o 5790161
Chethan Pandarinath, Daniel J O'Shea, Jasmine Collins, Rafal Jozefowicz, Sergey D Stavisky, Jonathan C Kao, Eric M Trautmann, Matthew T Kaufman, Stephen I Ryu, Leigh R Hochberg, et al. Inferring single-trial neural population dynamics using sequential auto-encoders. Nature methods, 15(10):805-88 -15, 2018.
F. Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Pretten- hofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, and E. Duchesnay. Seikit-learn: Machine learning in Python. Journal of Machine Learning Research, 12:2825-2830, 2011.
Evan D Remington, Seth W Egger, Devika Narain, Jing Wang, and Mehrdad Jazayeri. A dynamical systems perspective on flexible motor timing. Trends cognitive sciences, 22(10):938-92) 95288 2018.
Virginia Rutten, Alberto Bernacchia, Maneesh Sahani, and Guillaume Hennequin. Non-reversible gaussian processes for identifying latent dynamical structure in neural data. Advances in neural information processing systems, 33:9622-9632 2020.
Omid G Sani, Hamidreza Abbaspourazad, Yan T Wong, Bijan Pesaran, and Maryam M Shanechi. Modeling behaviorally relevant neural dynamics enabled by preferential subspace identification. 2021.
Nature Neuroscience,
13

---

Published as a conference paper at ICLR 2025
Omid G Sani, Bijan Pesaran, and Maryam M Shanechi. Dissociative and prioritized modeling of behaviorally relevant neural dynamics using recurrent neural networks. Nature Neuroscience, pp- 1-13, 2024.
Britton A Sauerbrei, Jian Shan-Zhon 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):386339, 2020.
Marine Schimel, Ta-Chu Kao, Kristopher T Jensen, and Guillaume Hennequin. ilqr-vae: control- based learning of input-driven dynamics with applications to neural data. biorxiv. In International Conference Learning Representations, 2022.
Steffen Schneider, Jin Hwa Lee, and Mackenzie Weygandt Mathis. Learnable latent embeddings for joint behavioural and neural analysis. Nature, 617(7960):360-366) 2023.
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 ml and vl. PLoS computational biology, 12011),4005514, 2016.
Maryam M Shanechi. Brain- interfaces from motor to mood. Nature neuroscience, 22(10):
1554-1564, 2019.
Krishna V Shenoy and Jonathan C Kao. Measurement, manipulation and modeling of brain-wide neural population dynamies. Nature communications, 12(1):633, 2021.
Nicholas A Steinmetz, Peter Zatka-Haas, Matteo Carandini, and Kenneth D Harris. Distributed coding of choice, action and engagement across the mouse brain. Nature, 576(7786):266-773, 2019.
Carsen Stringer, Marius Pachitariu, Nicholas Steinmetz, Charu Bai Reddy, Matteo Carandini, and Kenneth D Harris. Spontaneous behaviors drive multidimensional, brainwide activity. Science, 364+64377)eaav7893, 2019.
Damar Susilaradeya, Wei Xu, Thomas M Hall, Ferran Galán, Kai Alter, and Andrew Jackson. Ex- trinsic and intrinsic dynamics movement intermittency. Elife, 8:e40145, 2019.
David Sussillo, Sergey D Stavisky, Jonathan C Kao, Stephen Ryu, and Krishna V Shenoy. Making brain machine interfaces robust to future neural variability. Nature communications, 711):33749 2016.
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 Acudemy of Sciences, 121(7))e212288712) 2024.
Peter Van Overschee and Bart De Moor. Subspace Identification for Linear Systems. Springer US, Boston, MA, 1996. ISBN 978-1-4613-80614 doi: 10.1007/978- 1-4613-0465-4,
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.
Anqi Wu, Nicholas A Roy, Stephen Keeley, and Jonathan W Pillow. Gaussian process based non- linear latent structure discovery in multivariate spike train data. Advances in neural information processing systems, 30, 2017.
Yuxiao Yang, Allison Connolly, and Maryam M Shanechi. control-theoretic system iden- tification framework and real-time closed-loop clinical simulation testbed for electrical brain stimulation. Journal of neural engineering. 1566006, 2018.
Yuxiao Yang, Omid G Sani, Edward F Chang, and Maryam M Shanechi. Dynamic network model- ing and dimensionality reduction for human ecog activity. Journal of neural engineering, 16(5): 056014, 2019.
14

---

Published as a conference paper at ICLR 2025
Yuxiao Yang, Parima Ahmadipour, and Maryam M Shanechi. Adaptive latent state modeling of brain network dynamics with real-time learning rate optimization. Journal of Neural Engineering, 18(3):330133, 2021a. doi: (0.1088//741125522aaeee URL https:/ // dx doi. org/10.
1088/1741-2552/abbe
Yuxiao Yang, Shaoyu Qiao, Omid G Sani, J Isaac Sedillo, Breonna Ferrentino, Bijan Pesaran, and Maryam M Shanechi. Modelling and prediction of the dynamic responses of large-scale brain net- works during direct electrical stimulation. Nature biomedical engineering, 5(4):324-345 2021b.
Ding Zhou and Xue-Xin Wei. Learning identifiable and interpretable latent models of high- dimensional neural activity using pi-vae. Advances Neural Information Processing Systems, 33:7234-7247, 2020.
15

---

Published as a conference paper at ICLR 2025
A APPENDIX
A.1 METHOD DETAILS
A.1.1 TWO-SECTION FORMULATION
In equations 2 and we combined both latent state sections of our model (xg and x for simpler exposition. Here, we present the complete two-section formulation. The predictor form part of the model (equation can be written as follows:
A (xI--))
(y up
(x k(1-11)
+
K ( Uhs x
x [|-
(A.1)
cy (1) (1)
(2) (2)
9RRK-1
(X Up
Cy x1
8 o (1) (1)
(x u&)
Cz
(2)
[ 4&
(2)
where we have also included the two-section formulation for the prediction of observations as lines 2-3 of the equation. As before, Y& E "R" and Zy =k= are the observed high-dimensional neural activity and behavior respectively while Ug E R" represents the measured inputs to the dynami- cal system. Here, the overall latent state, X₃ R which describes the dynamics underlying the data, is constructed such that the behaviorally relevant neural dynamics, repre-
(2) R"= sented by x E " are dissociated from the the irrelevant ones, represented by X
(1)
The predictor form RNNs in equation A. (i.e., RNNI and RNN2) are complemented by another set of RNNs (i.e., RNNI for and RNN2 fw) that constitute the generative form part of the model (equation 3). and enable m-step-ahead (for m > 1) prediction of latent states and neurva-behavioral data. The generative RNNs again shown with combined latent state in equation for simpler exposition. The following equations show the complete two-section formulation:
x x
A (1)
fwr X
k 112)
Fwu (Ux+
x X
fw [x k+mm 112
(2)
K fw
(U++m k+m
(A.2)
Fktm/k
C/ (x (1)
(x k++m Uk+m)
(2)
(x (2)
k+m|h | us+m)
c2 (x()((((1) [m)k Uk+m) (2)
mk=
(x mm|k u
where m > 1, and X and x x++1|k (i.c., m = 1) are taken from equation A.1 A visualization of how the formulations in equations A. and A.2 are connected is provided in figure I In equa- tion A.2] we have also included the two-section formulation for the prediction neural hehaviobeaavioral
(1)
observations as lines 2-3 of the equation, showing that applying the same decoders
/C(2) and
cf (C( as in equation A.1 the m-step-ahead predicted latents x (1) (2)
*k-m|/ k4 gives the m-
step-ahead predictions the neural-behavione behavioral data (f |s and 2 A.).
Equations A. and A.2 together constitute the two-section formulation of the BRAID model, which consists of 12 transformations in total: A(.). (-)( K(-), f( Cz(+), and Cy(,), each having
two sections denoted with the (1) and (2) superscripts.
A.1.2 LEARNING ALGORITIM STEPS
In sections A.1.2 A. 1.3
we provide detailed formulations of the optimization stages used BRAID during learning. For simplicity, we explain the optimizations in terms of 1-step ahead predictions, which involve predictor form parameters of the model. Formulations for m-step-ahead predictions, which constitute additional terms in the overall loss (equations 4), are analogous to those provided here, but instead of the predictor form parameters (equation A.I they engage the generative form parameters (equation A.2).
16

---

Published as a conference paper at ICLR 2025
Note that regardless of what step-ahead predictions included during training, the learned model can be used to predict the latent state and neural-bchaviora data at m-steps ahead for any desired using equations 2 and 3 (or A.1 and (A.2) together. For example, in figure [. only [1, 2, 4, 8] step ahead predictions are incl luedd in the optimization loss (equation 4 but we evaluate the learned models with predictions up 32-steps ahead.
We develop two-stage optimization algorithm for learning parameters of the two sections of the BRAID model (equation A.I). We note that the following stages are sequential. This means that parameters associated with behaviorally relevant states (x: ( are fully learned with RNNI, then
needed, the remaining parameters corresponding to non-relevant dynamics (xx²)) can be learned via RNN2. In all the optimizations described below, we use the rrean-square-erre (MSE) of predicting observations as the loss function, but we note that the MSE is proportional the negative log- likelihood (NLL) for isotropic Gaussian-distribuue optimizations that are performed the two learning stages of BRAID.
data. Below we provide the details for the 4
Stage 1:
1a First, BRAID learns a recurrent neural network (RNNI) with n4 states, to minimize behavior
prediction MSE given past neural data and inputs (equation A.3). This ensures that RNNI only learns neural dynamics that are relevant to (i.e., predictive of) behavior. The states of RNNI
constitute the first set of latent states in the BRAID model: x (1) This optimization step can be formulated as:
(1)
(x
(1) + (y us) (x .uk
X Xk+1
Zk =
(A.3)
loss
MSE(ZE. C. 11) (x .x))
1b Next, in a second optimization, learn a transformation Cy that maps x while minimizing neural prediction MSE (equation A.4)
(1) neural activity
{
Yk
(1) = c
(xx up)
(A.4)
loss
MSE(Y Cy (x+ ,,:)
The above steps conclude stage 1 of learning, i.e., learning intrinsic behaviorally relevant dynamics x
(1) Next explain the (optional) remaining stage 2, which can learn any remaining dynamics in neural activity x
Stage 2:
2a We Iearn a second recurrent neural network (RNN2) with n2 = n, states, to minimize the MSE loss of predicting the residual neural activity, i.e., y/. yk
c (xg .u), given
past neural activity and inputs (equation A.5) States of RNN2, i.c. x
(2) together with x from
stage constitute the full neural dynamics, i.e., X be formulated as:
(2)
This optimization step can
x
x
x k+1 = (x (2) + K (yk Ub, x
y/ ( (x 44
(A.5)
loss
MSE Cy (x uu))
(2)
2b Finally, another readout, c can be learned to map x to the residual behavior, i.e., z/. = Zk -. (x ux), to minimizing the overall behavioral loss (equation A.6
loss : MSE(Z', (2) (x Up) {
z = -122 up
(A.6)
17

---

Published as a conference paper at ICLR 2025
(2) Note that the optimization in stage 2b does not change RNN2 or Xx that were learned in stage 2a. So although stage 2b is supervised by behavior, the second set of states x are still learned unsupervised with respect to behavior.
(2)
In case a very low state dimension is specified by the user for stage 1 (n1 lower than the ground truth shared dimensionality), RNN/ would not have enough capacity to learn all behaviorally relevant dynamics. In that case, some behaviorally relevant neural dynamics will be left for RNN2 in stage (2)
2 to learn. This is why the c²²²) transformation from Xx to behavior included in the model, to
(2) allow such behaviorally relevant information in X to be utilized to improve behavior decoding.
1.1.3 NON-ENCODED DEHAVIOR-SPECIIEE DYNAMICS
To remove the non-encoded behavior-specific dynamics, as preprocessing step, we fit high- dimensional (nz 150 in all real data analyses) unsupervised RNN extract neural dynamics alone by minimizing neural prediction MSE (equation A.7)
(0)
A
(x (0)
+ K (()) (y)) up)
Xk+1
Yk = c (xx uu
(0)
(A.7)
loss
MSF(Yy, Cy (xx ..))
Then readout C2 is trained to map these neurally relevant states x3
(0) to behavior:
Zk
C
(x
(0)
(A.8)
loss
MSF(ZZ C2 (0) (xx (0)
After parameters of the above are learned, behavior as output of the preprocessing RNN model (first line in equation A.8) and subsequently use in place of the original behavior in BRAID (in equations A. validate that this additional stage can successfully remove any input-driven behavior dynamics not encoded in the neural recordings (figure A.2 We include this preprocessing step in all reported real data analyses with BRAID.
inference the training data to obtain the filtered
A.3 A.6) In simulations, we
Stage 3: As mentioned in 3.3 if desired, BRAID can also learn behavior-specifie dynamics sep- arate dissociated latent states using post-hoe learning step. This step is performed after BRAID's main learning is done and meant to be used in conjunction with BRAID's s preprocessing stage ex- plained above. In this post-hoc learning step (stage 3), we first infer the behavior using the originally learned BRAID model. We then obtain the residual behavior (z"²) by subtracting the inferred behav- ior from the measured behavior. We then learn third RNN (RNN3) that optimizes the prediction of the residual behavior using only the external inputs. The following equations summarize this step:
Zか [CL
(x uu) - c22 (x uR)]
z ==
(3)
A (3)
(3)
K (ux)
X/+1
(A.9)
z C (3) up
loss
MSF C (x
(3)
u.))
We summarize BRAID's three stages and their use-case in table A.I
A.1.4 MODEL ARCHITECTURE DETAILS AND HYPERPARAMETERS
Throughout the manuscript, model nonlinearities within any of the transformations i.e., A()) A (.). (.), K (-). C.). and C.).) we use a multi-layer perceptron (MLP), also known as feedforward neural network, with single hidden layer, 64 units in the hidden layer, and ReLU nonlinearity as activation function. Otherwise, to keep a transformation linear, we replace the MLP with linear mapping implementing matrix multiplication, which is special case of an MLP with
18

---

dedicated the irrelevant states Published as a conference paper at ICLR 2025
Summary of the three stages of BRAID and their use-case in learning various dynamics. Each stage has two RNNs: predictor form and generator form RNN, denoted without and with fw subscript, respectively.
Table A.I:
Stage
Models
Functionality
RNNI, RNNI fwr
Intrinsic behaviorally relevant neural dynamics
2
RNN2, RNN2 /w RNN3, RNN3
Residual neural-specific intrinsic dynamics Residual Sohavior-specifiie dynamics not encoded in neural activity
3
no hidden layers and linear activation function. For example, linear BRAID is BRAID model with all mappings being linear whereas BRAID Nonlinear C₂ has a nonlinear MLP as the behavior decoder (both 2 and c while all of its other transformations linear.
We use an Adam optimizer (Kingma & Ba 2017) in all BRAID optimizations. We train models up maximum number of 2500 epochs ensure convergence, while employing early stopping to avoid overfitting, We provide example learning curves in figure A.7 Details of the hyperparameters used for BRAID provided in table A.2
Table A.2: BRAID hyperparameters used in real data experiments and simulations
Hyperparameter
Value
Number of hidden layers in nonlinear maps
Number of hidden units in nonlinear maps
64
Nonlinear activation
ReLU
Learning rate
0.001
Batch size
32
Sequence length Optimizer
128
Adam
We also show that BRAID can locate the correct structure of nonlinearity within all possible com- binations in our simulations (table 2). Here, we set each of the following four groups of transfor- mations i..., A(- y- K() (.)/k (-), (.).. C-(.). as linear or nonlinear, resulting in a total of 24 cases. To select one final configuration for the nonlinearity, we follow an automatic nonlinearity selection procedure for given dataset. In this procedure, within the training data, we perform 2-fold inner cooss-validation in which we fit BRAID models with all 24 nonlinearity configurations and then pick the nonlinearity structure with the best volsdaaldd behavior decoding the held- out section of the training data. Then we retrain a BRAID model with that selected structure on the entire training data to get final model. Finally. evaluate that final model the unseen test data. We refer to this approach as automatic nonlinearity selection (table
In simulations, state dimensions are data. In real data analyses, to investigate the effect of nz. we vary the state dimension nu E [1, 2, 4, 8, 16, 32, 64] and report the results (figure A.5) For BRAID and DPAD models, we always learn the first 16 dimensions via stage i.e., X E Ih with n1 min(16, nz), and if there is any more capacity left (i.e., if n₂ ny n2 is positive),
be the
that of the true model underlying the
x E Rm=- and is learned using stage 2. We pick 16 as the dimensionality of the behaviorally relevant states as in our experiments, because BRAID reached close to its peak behavior decoding at this dimension. We refer to models with state dimensions n = 16 and 64 as low and high- dimensional regimes, respectively.
(2)
A.1.5 INTRINSIC DYNAMICS AND FIGENVALUES
To evaluate how well the intrinsic behaviorally relevant neural dynamics are learned, in simulations
(figures 2 A.I. we assess the eigenvalues of the generative transition A which characterize the intrinsic dynamics. Note that the ground truth and learned models in these simulations both have a linear intrinsic state transition A fw and the nonlinearity either lies in the transformation from latent
19

---

Published as a conference paper at ICLR 2025
states to behavior (figure transformation from external inputs to the latent space (figure A.1 We learn the forward recursion parameters f and K fw of equations and as we optimize multi-step-aheatt predictions. We take the eigenvalues of the intrinsic transition A fwr and compare them to that of the ground truth model. For the ground truth (A₂) and identified (A:) eigenvalues we first pair them as {A1, Ag...... ........... An, } and {Å1, Âg...., An, } such that the sum of squared distances of pairs is minimized. We then calculate the normalized eigenvalue error as:
Σ
a
(A.10)
Σ |1:11
A.2 BASELINES
First, assess the impact of the nonlinearities learned by BRAID, we compare it against two fully linear dynamical methods: IPSID (Vahidi et al. 2024) and linear BRAID. Second, to highlight the significance of modeling the effect of measured inputs the huuaa--behavooal dynamies, take an autonomous dynamical model, DPAD (Sani et al. 2024). as another baseline. Another important aspect of BRAID is supervision of the behaviorally relevant dynamies in presence of inputs in its first stage. To assess prioritization of the behaviorally relevant dynamics due to this supervision, we take an equivalent unsupervised baseline termed U-BRAID detailed below. We also compare BRAID against multi-modal nonlinear method that allows accounting for external inputs termed mmPLRNN (Kramer et al., 2022). Finally, we compare BRAID to TNDM (Hurwitz et al. 2021). second autonomous method for modeling neural-behaving dynamics.
A.2.1 IPSID
IPSID, similar to BRAID, models the effect of measured external inputs on neural-behavione dy- namics but operates under a fully linear framework. fits the parameters a linear version of equation A via projection-based analytical algorithm called subspace identification (Van Over- schee& & De Moor 1996).
A.2.2 LINEAR BRAID
Linear BRAID serves as another linear baseline, retaining the same architecture and learning stages as BRAID but with all transformations replaced by linear mappings. In essence, linear BRAID and IPSID have a similar model that are learned differently. Linear BRAID uses the same numerical optimization used in BRAID. BRAID reduces to linear BRAID by removing all hidden layers within model transformations and setting all activation functions to lincar. In simulations, we find that linear BRAID and IPSID perform similarly expected (figure 2
A.2.3 DPAD
Dissociative Prioritized Analysis of Dynamics (DPAD) (Sani et al. 2024), learns nonlinear model that dissociates and prioritizes dynamics shared between neural activity and behavior, but it impor- tantly does not account for the external inputs. Originally, DPAD also does not allow for multi- step-ahead optimization and thus does not learn a generative form representation of the dynamics, which is in contrast to the learning of A fw in BRAID. We extend DPAD to add optimization of multi-step-ahead predictions into the DPAD framework for a more fair comparison to BRAID in terms of forecasting.
A.2.4 U-BRAID
U-BRAID is unsupervised method, which only performs stage 2 of the BRAID learning proce- dure. As such, U-BRAID learns all neural dynamics irrespective of their relevance to the behavior, but still while considering inputs (equation A.5) U-BRAID does utilize behavior information in
20

---

Published as a conference paper at ICLR 2025
learning dynamics, and the extracted latent states are later mapped to the behavior data via down- stream decoder (equation A.6 but without x.)).. fact, U-BRAID is special case of BRAID with 1 = 0 and f2 R2.
A.2.5 MMPLRNN
Multi-modal sereewsse-lineal RNN (mmPLRNN) is a method previously introduced for multi- modal dynamical modeling with piecewise- RNNs (Kramer et al. (022). This method allows for modeling external inputs, although this aspect of it has not been investigated in any prior work. Nevertheless, we compare BRAID to an input-driven mmPl PLRN to further assess its performance. mmPLRNN builds on a prior work, PLRNN (Durstewitz) (017), by fusing information from two modalities (e.g., neural activity and behavior). By design, mmPLRNN utilizes both modalities (and input) during inference, which is in contrast to BRAID that only uses neural activity (and input) during inference. Although this provides the benefit of using more data for behavior decoding to mmPLRNN and confounds the comparison with BRAID, still include the mmPLRNN results. We train mmPLRNN models with nonlinear readouts comparable to BRAID and compare their neural-behavioral forecasting. mmPLRNN is generative model whose parameters learned via variational inference. We used the recommended hyperparameters from the original work !!
A.2.6 TNDM
Targeted Neural Dynamical Modeling (TNDM) (Hurwitz et al. 202T) is method based on sequen- tial autoencoders that learns sets of dynamies: one contributing to both neural and behavioral data, and the other only contributing to neural data. TNDM uses a non-causal, bidirectional RNN as the encoder to infer the initial conditions for its relevant and irrelevant generator/decoder RNNs. The dynamical model is learned via variational inference with both neural and behavioral reconstructions optimized simultancously with a combined loss. Unlike BRAID, TNDM does not account for mod- cling the effect of external inputs on neural-behavioral dynamics. Therefore, for comparisons, we also implement an extension of TNDM to allow the inclusion of external inputs. In this version, we provide the external inputs to TNDM model as input by concatenating them with the neural activity as the input to the model. Importantly, we do not add reconstruction of inputs as part of the loss to keep the loss the same as that of the original TNDM and keep the learned model focused on beural-behavioral reconstruction.
We compare BRAID to TNDM (with and without addition of sensory stimuli as external input) in our real data experiments. Unlike BRAID that models the neural observations with a Gaussian distri- bution, TNDM uses Poisson observation model for the neural data. Therefore, in our comparisons to TNDM, we analyze non-smoothed spike counts in 50ms bins (for both TNDM and BRAID). We use the default hyperparameters from the original work for TNDM.
A.2.7 LFADS
Latent Factor Analysis via Dynamical Systems (LFADS) (Pandarinath al. 2018 is unsuper- vised method that combines nonlinear dynamical modeling with sequential autoencoders. Similar to TNDM, LFADS uses non-causal, bidirectional RNN as the encoder to infer the initial conditions for a generator/decoder RNN. The dynamical model is learned via variational inference for unsu- pervised neural reconstruction. LFADS has version that uses additional RNNs called controller networks to infer unmeasured inputs from the neural data and use those inferred inputs to drive its generator RNN, We use this version of LFADS (O'Shea & Pandarinath 2021 to serve as addi- tional benchmark in comparison to BRAID's 's modeling of measured inputs. For a full visualization of the LFADS model including the controller networks see Supplementary Fig. 12 in Pandarinath et al. 2018 Briefly, without controller network, all the information about the complete trial has to be encoded into the initial state of the LFADS generator RNN, which then autonomously evolves to extract states and factors over the course of the trial. In contrast, controller network acts regular non-autonomous RNN for LFADS and allows its generator RNN to take corrections from the neural data throughout the trial, hence improving its capacity to accommodate mooth changes
We use the implementation provided in htt.tp/s//ggaaab.ccoo///tsttttta LRNN
A We use the implementation provided in
http:////eeennlllaaaataa.aa ql.ab/tndm
We use the implementation provided in
manaaaagrr
21

---

Published as a conference paper at ICLR 2025
in the middle of trials (Pandarinath et al., 2018). h.. The controller network itself does not directly take neural data as input, rather an additional bidirectional RNN called the controller-encodee operates on the neural data of each trial, and the states of this controller-encoder are passed as input to the controller network. As a result, the 'inferred inputs' in LFADS are also non-causally inferred. To compare results with BRAID, the controller-encoder networks to be the same as BRAID. We also pass the same smoothed neural data to LFADS and use the Gaussian neural loss option accordingly. Other hyperparameters were set as those in the original work (Pandarinath et al. 2018
set the number of factors, and the latent states of the generator and
A.2.8 CEBRA
CEBRA (Schneider et al. 2023) is recent method proposed for extracting latent embeddings from neural data in a way that can be guided by behavior. CEBRA uses a 1-dimensional convolutional net- work to extract embeddings from small windows of neural data and guides the extraction of latents via a contrastive loss. The supervised version of CEBRA (i.e., TEBRA-BBharye loss on behavior to learn embeddings that dissociate samples with different behavior data, and are thus behaviorally relevant. Due to its use of convolutional network to extract embeddings, CEBRA does not learn explicit recursive dynamics and also extracts embedding from finite window of neu- ral data at time. CEBRA also does not learn any models to decode behavior neural data form the extracted embeddings. Thus, as done in the original work (Schneider et al. 2023), we fit k-NN regression" models decode heural-behavioral data from the extracted CEBRA A embeddings. For a fair comparison, we also provided the measured sensory inputs time-series of the task to CEBRA by concatenating them with neural activity as the input. We use the default hyperparameters from the original work for CEBRA.
uses a contrastive
A.3 HON-HUMAN PRIMATE ELECTROPHYSSOLOGIIN RECORDINGS FROM
We analyzed publicly available dataset O'Dooeeryy et al. 2017) in which macaque (monkey I) performs a motor task. Spiking activity was recorded from primary motor cortex (M1), while the subject controlled a 2D cursor reach targets that appeared on random locations on grid within a virtual reality environment. Targets appeared back to back, without any time gaps. We took the subject's 2D fingertip position and velocity as the behavior time-series Zy, and the sensory input, taken as 2D location of the current target, as the input signal ug. We analyzed the first spike dimension available for each channel-resulting recording sessions and randomly selected half of these units to model as our neural activity. For neural modality, we use spike counts within 50 ms non non-overlapping windows. Finally, we smoothed the spike counts by Gaussian kernel with 50 ms s.d. (except for in table A.5 and took that the neural seme-seriss Yı- We report the mean and standard error of the mean (s.c.m.) computed across 7 sessions and cross-validated folds.
89 to 92 units from the first 7 available
As a second neural modality, we also model the raw local field potential (LFP) activity recorded from the same monkey during the same task. As the only preprocessing, we apply 10 Hz anti- aliasing filter to be able to downsample the raw LFP to the sampling rate of behavior (i.e., 20 Hz). We refer to this modality with those obtained for spiking activity (figure
LFP data. Results for raw LFP data (figure A.3) were consistent
A.4 SIMULATION DETAILS
We analyze three simulated datasets based dynamical systems. In all the three, generate 10 different sets of random linear matrices for equation A. then generate the ground truth latent states Xk. neural activity y& and behavior observations zg- In equation A.IT] Wk. V&, and e zero-mean white Gaussian noises accounting for unmeasured excitations, neural observation noise, and behavioral observation noise respectively. fc, (-) and fø (-) are nonlinear functions, as described below for each simulation.
"sklearn..neighee hagaaammnniiooooooigngnnnneee KNeissobosReerresoo (Pedregosa et al.
2011
We use the implementation provided in https: // thub. wwm hamm//dapt/// hallaapptivveottoooooooooooo l.ab/
CEBRA
22

---

Published as a conference paper at ICLR 2025
X A + f +W
Yk Cyxy + Dyu + Vi.
(A.11)
Z
(C2x1) + Dzu; + Ck
To generate temporally structured input in all cases, we simulate a separate random linear state space model according to equation A.12 and take its output as the external input u to the main model of equation A.11
k+1
A Ju t C C v
(A.12)
A.4.1 SIMULATION 1: SPIRAL BEHAVIOR MANIFOLD
For the first simulation, we generate data with spiral behavior manifold. To do so, we apply y sin(z)] the readout from the latent states (third line in equation A.I1I. The bar over the function input indicates a scaling factor that normalizes it before nonlinear function is applied. See figure 2h for a visualization of the nonlincarity. We take f-) identity, set dimensions to ny R2 nu n2 n2 = 2, and take Dy = D2 for this simulation.
pointwise nonlinear mapping Sc. (v) = T
2 cos 2
A.4.2 SIMULATION TRIGONOMETRIC BEHAVIOR MANIFOLD
For the second simulation with trigonometric behavior map. we apply another nonlinearity, point- wise sinusoidal nonlinear function, fc, (v) = sin bv, the latent states to generate the behavior Zy: See figure 2e for a visualization of the nonlinearity. In this simulation f'((.) is taken as identity function and we fly T12 fy fl 71 1.
A.4.3 SIMULATION 3: TRIGONOMETRIC INPUT-ENCODER
For the third simulation, iterate the state equation (first line in equation A.11 a pointwise sinusoidal nonlinear function, (v) = a sin(() + bv, to the input. See figure A. for visualization of the nonlinearity. Here fc, (-) is taken be identity function. In this simulation, we set n2 =
apply
n1 = 1.
We also perform two additional simulations (figure A.2 that are similar in nonlinearity structure to the second and third simulations explained above, but incorporate an additional I I-dimensional latent (3) in figure a, representing input-driven behavior-specific dynamics encoded in the state X neural activity, as follows:
x
[ (1) (3) (3)
x
+ fB (B) B(3)
=
+ Wk
x
(A.13)
cy) X + Dy +VE
(1)
Yk
fc, (Ct) (1) + C² x
(3)
+Du Ek
A.5 SUPPLEMENTARY RESULTS
A.5.1 BRAID CAN EXCLUDE NON-ENDDDED SEHAVIOR-SPECIHC DYNAMICS
Here, demonstrate that the optional preprocessing step in BRAID (detailed in section 3.3 dissociate behavior-specific dynamics (i.e., those that are not encoded the neural activity) during learning and make sure they are not conflated with intrinsic neural dynamics and are not mixed into the neural states (x( and x (2)). We conducted two additional simulations similar in structure to the second and third simulations explained in section 4.1.2 However, here, for all simulated models, we added input-driven dynamics that influenced behavior but were not encoded neural activity (denoted as x
The preprocessing step is intentionally expected to yield latent states that are potentially less predictive of behavior, but are encoded in neural activity. When desired, BRAID provides the option further learn behavior-specifie dynamics post-hoc
(3) in figure 1 and
23

---

Published as a conference paper at ICLR 2025
with separate latent state (Xg (x²)). The preprocessing and post-hoc learning steps allow BRAID to avoid conflation of non-encoded behavior dynamics with others, while also being able to learn these dynamics and thus not incurring any overall reduction in behavior decoding.
We fitted BRAID models with the preprocessing, and both with and without post-hoe learning of behavior-specific dynamics. With the preprocessing, BRAID reached the neural prediction perfor- mance of the ground truth model indicating correct removal of behavior-specific dynamics (figure A.2 Moreover, the optional learning of behavior-speciff dynamics led reaching the behavior decoding performance of the ground truth model (figure A.2). suggesting that one could option- ally learn these dynamics as well within BRAID to gain interpretability (by learning disentangled model) without compromising decoding performance.
A.5.2 ADDITIONAL SUPPLEMENTARY TABLES
In this section we include additional supplementary tables that further support the results from the main text. The caption for each supplementary table includes all the details, but here we provide list of these tables:
Table A.3
Ablation analysis showing the importance of the RNN generative model in BRAID
for learning intrinsic dynamics.
Table A.4
BRAID results for different nonlinearity configurations in real NHP data.
Table A.5
Comparison with TNDM in real NHP data.
Table A.6
Comparison with LFADS and CEBRA in real NHP data.
Table A.7
BRAID results for modeling different number of neurons real NHP data.
Table A.8
Same as table
3 shown in terms of the R metric.
Table A.3: A separate generative model is essential to learn the intrinsic dynamics accurately with BRAID.
Row 1: BRAID, when learning separate generative model (RNN fw) more accurately learns the intrinsic dynamics as quantified by the error identifying eigenvalues of the ground truth intrinsic dynamics in the simulated dataset with spiral manifold in figure 2
Row The error when ablating the forward RNN from BRAID.
Rows 3-4: DPAD, even when optimized with m-step-ahead prediction loss and/or with input (uk),
does not learn the intrinsic dynamics accurately.
Method
log1o normalized eigenvalue error
BRAID
-1.3963 + 0.2551
BRAID without RNN
0.0635 + 0.2212
DPAD + m-step loss
0.0357 + 0.1587 -0.6512 0.0354
DPAD + input (ug) + m-step loss
Table A.4: Comparison of BRAID model's nonlinearity configurations in the NHP dataset (nz
M1 = 16, --step-ahead)
Model nonlincarity
Behavior forecasting CC
Neural forecasting CC
Linear
0.7453 + 0.0066
0.1767 + 0.0054
Recursion (A, fa)
0.7121 + 0.0059
0.2719 + 0.0061
Encoder K, K fw))
0.7181 + 0.0078
0.1646 + 0.0049
Decoder (Cz, Cy)
0.8042 + 0.0085
0.3274 + 0.0078
24

---

Published as a conference paper at ICLR 2025
Table A.5: Comparison to TNDM, both when sensory input is additionally provided to the TNDM model and when it is not (see appendix A.2.6). All models are learned in low-dimensional regime, i.e., BRAID with n2 = R- = 16, and TNDM with 16 relevant factors only. We used non-smoothed spike counts the neural signals in this analysis. BRAID performances are for causal I-step-ahead prediction, whereas the TNDM performances are non-causal smoothing performances, which are the only option for TNDM since it is sequential autoencoder.
Method
Behavior decoding CC
Neural prediction CC
TNDM
0.3752 + 0.0170
0.3021 + 0.0051
TNDM with sensory input
0.6219 = 0.0103
0.3075 = 0.0050
BRAID (ours)
0.7841 0.0079
0.2935 + 0.0053
Table A.6: Comparison to LFADS (with controller), and CEBRA (CEBRA-behavior). LFADS inferred the external input the dynamical system (appendix A.2.7 For CEBRA, w.s.i. (with sensory input) indicates that sensory input is additionally provided to the model to obtain the em- beddings (appendix 6.2.6). All models are learned with both low-dimensional (nz = 16) and high- dimensional (n, 64) latent states (n, 16 for BRAID). BRAID performances are for causal I-step-ahead prediction, whereas LFADS and CEBRA performances are non-causal smoothing, and O-step-ahead reconstruction respectively.
Behavior decoding CC
Neural prediction CC
Method
f1g = 16
1hr = 64
fg = 16
f2 = 64
LFADS
0.4714 + 0.0192
0.5891 = 0.0135
0.5615 + 0.0086
0.5920 + 0,0072
CEBRA w.s.i.
0.7544 + 0.0073
0.7514 = 0.0072
0.5070
0.0053
0.5693 + 0.0041
BRAID (ours)
0,8109 + 0.0074
0.8085 + 0.0076
0.5571
0.0051
0.8401 + 0.0061
Table A.7: Effect of the number of neurons on BRAID's performance. Results of the BRAID modeling when different numbers of neurons included as the neural signal. We included neurons from the same channel sets across different sessions and the ranges indicate the number of neurons available from those channels across different sessions. all 3 rows, neural predictions are eval- uated on the same common set neurons (i.e., the smallest set shown in row 1) to make the neural forecasting results comparable across rows. Behavior forecasting improved with more neurons and neural forecasting remained largely stable. These results suggest that BRAID aggregate behav- iorally relevant information across larger populations of neurons, while still being able to model this higher-dimensional population activity well.
Behavior forecasting CC
Neural forecasting CC
Scale
f1y = 16
nr = 64
= 16
64
20-21 neurons
0.7727 + 0.0091
0.7709 = 0.0089
0.3206 + 0.0081
0.4115 + 0.0094
41-43 neurons
0,8042 + 0.0085
0.7970 = 0.0086
0.3220 1 0.0088
0.4202 + 0.0091
89-92 neurons
0.8337 + 0.0061
0.8302 = 0.0072
0.3329 + 0.0082
0.4195 + 0.0088
25

---

Published as a conference paper at ICLR 2025
Table A.8: R results for the same analyses provided in table 3 Forecasting performance (4-step- ahead) compared to baselines NHP dataset for models with low-dimensional (n, = 16) and high- dimensional (n, = 64) latent states. sessions were excluded from this table due to a very negative outlier in the mmPLRNN results.
16 for BRAID, linear BRAID, and DPAD. I of the
Behavior forecasting R
Neural forecasting R
Method
f1g 16
fl 64
f2g 16
n2 64
linear BRAID
0.5821 + 0.0056
0.5763
0.0084
0.0239
0.0036
0.1519 0.0067
DPAD
0.4561 + 0.0141
0.5497 = 0.0101
0.0321 = 0.0042
0.1344 + 0.0069
U-BRAID
0,6047 + 0.0097
0.6684 = 0.0088
0.1768
0.0068
0.1860 0.0064
mmPLRNN
0.4737 + 0.0118
0.6127 = 0.0369
0.0734 + 0.0079
0.0563 + 0.0670
BRAID (ours)
0.6680 + 0.0118
0.6578 + 0.0129
0.1083
0.0060
0.1792 + 0.0064
26

---

Published as a conference paper at ICLR 2025
A.5.3 ADDITIONAL SUPPLEMENTARY FIGURES
In this section include supplementary figures that further support the results from the main The caption for each supplementary figure include all the details, but here we provide a list of these supplementary figures:
Figure
A.1
Simulation results for simulation 3 with trigonometric nonlinear input-encoder.
Figure
Validation of the optional third stage of learning in BRAID for learning behavior- specific input driven dynamics.
A.2
Figure
A.3
Results of modeling the raw LFP modality in the real NHP data.
Figure
A.4
Average low-dimensional latent state trajectory extracted from real NHP data.
Figure
A.5
Results in real NHP data for different latent state dimensions.
Figure
A.6
Example BRAID decoded time series real NHP data.
Figure
Example plot showing loss versus epochs for BRAID.
(b)
log. normalized EV error
BRAID Nonlinear X
(1epp-ahead=(112.3/44
0.8
-0.5
(Steps-ahead=11.233...5
DPAD Nonlinear
Docoding
0.6
Neural prediction
0.6
-1
10%
0.4
-1.5
0.4
BRAID
1%
(sleps (he.saaaadd-11..4.4...,
0.2
-2255
IPSID
0
(a)
(e)
(9)
Trigonometric input-encoder
0.8
Neural prediction CC
0.8
0.6
Decoding
0.4
0.2
0.2
8
32
I1 steps ahead
# steps ahead
Figure A.1: BRAID results, optimized for forecasting, in simulation with trigonometric input- encoder. (a) Visualization of example nonlinearity in the simulation. (b-c) 1-step-ahead behavior decoding and neural prediction for nonlinear BRAID, nonlinear DPAD, linear BRAID, and IPSID. (d) Error in identifying intrinsic dynamics of the true model, quantified by the eigenvalues of the state transition matrix Afw- (c-f) Behavior and neural forecasting accuracy for to 32 steps ahead, enabled by learning the intrinsic dynamics (A ww), with predictions optimized for [1, 3, 4, -tepss ahead (section3.1)
27

---

Published as a conference paper at ICLR 2025
(8)
(b)
(c)
Trigonomelric
Neural prediction CC
0.8
behavior manifold
Decoding CC
0.8
0.6
0.8
0.4
0.4
0.2
0.2
True model (Ideal)
0
BRAID (with preprocessing)
BRAID (without preprocessing
(d)
(e)
nan-encaded dynamics)
Trigonometric input-encoder
BRAID (with preprocessing and non-encoded dynamics)
0.8
Decading cc
Neural prediction
0.75
0.6
0.4
0.2
0.25
Figure A.2: Behavior preprocessing successfully excludes non-encoded, behavior-specific dy- namics in simulations. I-step-ahead behavior decoding and neural predictions for simulations with (a-c) trigonometric behavior decoder, and (d-f) input-encoder. Note that the true model includes behavior-specific dynamics, so here we expect that decoding of BRAID with preprocessing but without the post-hoc learning of behavior-speciffc dynamics (shown as green), to worse than that of true model. Once the post-hoc learning step is also performed (shown red), BRAID reaches ideal performance, but importantly does so while these behavior-specific dynamics are dissociated into separate latent state C/ In contrast, without the preprocessing step (shown as blue), BRAID reaches ideal decoding performance, but does so without having dissociated behavior specific dy- namics to not be included in ((!) See section A.1.3 for details.
LFP modality
(a)
(b)
0.8
Neural prediction CC
Decoding CC
0.6
0.4
0.4
0.2
0.2
0
0
4
2
16
32
steps ahead
# steps ahead
Nonlinear decoders (sleps-ahead=11122445
DPAD Nonlinear decoders (1sps-aaeea=11.2244
Linear BRAID (steps-ahead:(1122..]))
mmPLRNN Nonlinear decoders
Figure A.3: Analysis of Local Field Potential (LFP) neural modality. BRAID outperforms baselines in neural-behavioral forecasting. We used L.FP neural data as second different modal- ity and performed analysis similar the one in figure 3 for smoothed spike counts. (a) Behavior and, (b) neural activity forecasting correlation cocfficient (CC) for BRAID. linear BRAID, DPAD, and mmPLRNN for n.=16. Shaded areas show the s.e.m,, acToss the 3 recording sessions and cross- vross-validation folds. NHP results all other figures and tables are for spiking data.
28

---

Published as a conference paper at ICLR 2025
(a)
Position
Velocity
(b)
BRAID
(c)
DPAD
(d)
U-BRAID
Figure A.4: Latent states trajectories revealed non-human primate reaching dataset. We divide reach trials to 8 conditions based on reach direction (shown by colors) and find the condition- averaged (4-stepaahead) latent states trajectories for 2 dimensional models. (a) Condition Condition-averaged behavior i.e., movement position and velocity. (b-d) Condition-averagee latent state trajectories for (b) BRAID, (c) DPAD and (d) U-BRAID. BRAID learns the most well-separated posterior (latent) trajectories for the 8 conditions in the task. U-BRAID's 's trajectories are the least separated, showing that BRAID more successful in extracting the behaviorally relevant intrinsic dynamies that more congruent with behavior. Also, DPAD's trajectories are not as well-separated as BRAID's, although they are more separated than the unsupervised version (U-BRAID) as expected.
(a)
(b)
0.8
BRAID Nonlinear decoders (aops-aahed-11122...3
Neural prediction CC
0.4
Decoding DC
Linear BRAID (steps-ahead=(11.244 (1224,))))))))
0.3
0.4
0.2
DPAD Nonlinear decoders (steps-ahead=111 2,4.8)
0.2
0.1
U-BRAID Nonlinear decoders (steps-ahead=11122..8)
10
20
60
0
10
20
30
40
50
60
R8
Figure A.5: Forecasting (4-step-ahead predictions) correlation coefficient across latent dimen- sions. BRAID predicts both (a) behavior and (b) neural activity more accurately than DPAD due to modeling input, and than linear BRAID due to modeling nonlinearity. As state dimension increases beyond 16 (dedicated to behaviorally relevant dynamics, i.e., ny 16), BRAID
stage 2 to
leamn neural specific dynamics, thus reaching the unsupervised baseline. i.e., U-BRAID, in neural prediction. For BRAID and DPAD, the first 16 state dimensions dedicated to the behaviorally relevant neural dynamics (i.e., ni = 16) while any remaining dimensions (n,, > 16) are dedicated to the residual non shared neural dynamics. 4-steps-ahead in this dataset corresponds 200ms ahead.
29

---

Published as a conference paper at ICLR 2025
Position X
Aptte My m h
wom comeu
Position -Y
Predicted behavior
Velociy X
True behavior
Velocity Y
160
180
200
220
240
260
Time (6)
Figure A.6: Predicted behavior visualization. True behavior versus BRAID's (4-step-ahead) pre-
dicted behavior for representative session corresponding to the results table (n, = 64).
(a) Example 1
(b) Example
Stage leamiing
Stage 1 learning curve
2.5
Loss
2.0
2.0
1.5
50
100
150
200
250
300
150
250
300
350
Training epochs
Training Validation
Training epochs
Stage 1b learning curve
Stage 1b learning curve
1.4
1.2
Loss
1.0
1.0
0.8
0.8
Training epochs
Training epochs
Figure A.7: Learning curve examples for BRAID model. Loss values as a function of number of epochs for training and validation datasets for example runs of BRAID. Top row: learning curve examples for the RNNI, RNN/ fu. i.e., stage 1a, in section 33.2 Bottom row: learning curves examples for stage in section 3.2 These examples are taken from the same models whose per- formances are reported in figure and table 3
30


Reviews:
[{"review_id": "cO7SFDav2p", "reviewer": "Reviewer_QzSK", "summary": "BRAID aims to distinguish the effect of measured inputs towards neural dynamics, while dissociating shared neural-behavioral dynamics, neural-only dynamics, and behavior-only dynamics. It does so by training a series of recurrent neural networks: (1) Stage 1 trains a shared recurrent neural network that outputs both neural activity and behavior for (a) 1-step ahead and (b) multi-step ahead prediction, (2) Stage 2 does similarly but for neural dynamics only, and (3) Stage 3 does so for behavioral dynamics only. The authors show better correlation coefficients with the neural and behavioral activity in validation datasets as compared to some baseline comparisons.", "questions": "None noted / see above.", "limitations": "", "rating": 3, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The method dissociates neural-only, vs. behavior-only, vs. shared neural-behavioral spaces.\n- BRAID has the ability to take in measured inputs. However, this is not providing any conceptual advance from existing methods, since they are simply fed into the RNNs with an additional input transformation in the form of K.", "weaknesses": "- BRAID is very similar to DPAD (Sani et al., 2024), but with the addition of measured inputs, behavior-only dynamics, and an additional network that predicts m-steps ahead. None of these additions are conceptual advances, and are extremely straightforward additions to an existing method.\n- The paper claims that \"BRAID disentangles [input] dynamics from intrinsic dynamics\", however the contribution of the inputs is not analyzed further at all - can one effectively disentangle input dynamics from intrinsic dynamics via this approach?\n- The modeling strategy is multi-stage and quite involved, with multiple RNNs being trained without fully going into the utility of each one. The 'RNN_{fw}' models seem to be forecasting, but why is it necessary to have a separate RNN for forecasting when the 'RNN_1' model is a dynamical model that should be in theory capable of predicting m-steps forward in time?\n- The R^2 should be reported throughout the paper instead of Pearson's; the R^2 is more standard in this field, and takes into account the predictability using the mean value of the signal.\n- There is no attempt at interpretability of the underlying dynamics and the contribution from different sources as identified by this method.\n- While the authors show that BRAID performs with higher behavior reconstructions than TNDM as shown in the Appendix, this is very much to be expected since TNDM does not optimize separately for behavior reconstruction, as BRAID does. Similarly, DPAD does not either (and does not take in inputs). However, these are very simple to add to both of these methods, and thus do not provide fair comparisons in their existing form.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Jy093v9jS6", "reviewer": "Reviewer_FaeG", "summary": "This paper proposes a new deep learning method to jointly model neural and behavioural data by explicitly modelling task inputs, in order to better model and disentangle input effects and intrinsic neural dynamics that are predictive of behaviour. The learning of the intrinsic dynamics is enabled by the use of a forecasting objective, i.e., $m$-step-ahead neural and behavioural prediction. The method involves 2 (or 3) RNN models, and optimisation is done in multiple stages: a pre-processing stage to filter out behaviours that are not relevant to recorded neural dynamics (used instead of actual behaviour for training), a stage to learn the neural dynamics predictive of behaviour, and a stage to learn any residual neural dynamics. The proposed method outperforms several baselines at predicting neural activity and decoding behaviour, and is also amenable to real-time predictions due to its causal formulation.", "questions": "* 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.\n* Is it possible to show results when ablating for the second and third (post-hoc) stages of training? Apologies if I've missed this. Also, could the authors comment on what happens when just training the model end-to-end?\n* Could the authors consider experiments that show the scaling performance of the model vs number of neurons, and attempt to compare with LFADS (and CEBRA) if possible?\n* Could the authors comment on multi-session models, generalisation to unseen sessions?\n* Apart from neural predictivity and behaviour decoding, could the authors perhaps use a similarity metric to explicitly compare the learned and true dynamical systems (e.g. using [Dynamical Similarity Analysis](https://openreview.net/forum?id=7blSUMwe7R))? This should be fairly easy for the synthetic tasks, and should be possible in case of the neural data as well.\n* It would be interesting to see the model's performance on a dataset involving multiple brain regions, but I understand that this would take more time to run and depend on the availability of a dataset.\n* I spotted a potential typo on Line 073 (\"prepossessing\" -> \"preprocessing\"?), and another minor issue on Line 111 (\"intrinsic representation of dynamic.\" -> \"dynamics\"?).", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "* The method is well-motivated given prior work, and the causal formulation makes it amenable to real-time inference, which is a strong asset.\n* The figures are neat and the experiments are comprehensive. The writing is mostly clear, but I have some comments on explaining the method in a slightly clearer manner (see Weaknesses).\n* The method performs well not only on synthetic tasks designed to show its efficacy over several baselines, but also on decoding real neural data.", "weaknesses": "* The model architecture and the learning stages are quite complicated, and the writing makes this hard to understand in some places. If I've understood this correctly, there are mainly two RNNs: one to predict the next timestep neural activity given current activity and observed inputs, and another to predict activity $m$ timesteps into the future using just observed inputs. Each of these RNNs' parameters are split into up to 3 subsets, which are optimised sequentially: one to learn behaviourally relevant neural dynamics, one to learn any residual dynamics, and another to learn behaviour that is not encoded in the neural dynamics. This is in addition to an RNN that preprocesses inputs.\n\n  While this preprocessing RNN has been ablated for, can the authors comment on how one identifies in practice whether or not to use the second and third stages of training (RNN2 and RNN3), which were mentioned to be optional? In general, some additional clarity in the writing here would be appreciated.\n\n* It would be important to see how the method scales with the number of neurons – based on the details in the appendix, it seems like the maximum dimensionality explored here (in the neural data case) is around 45. Perhaps the authors could run an experiment comparing decoding performance and neural predicitivity for differing numbers of neurons from the same dataset to show this (and also comment on the time taken to train BRAID).\n\n* From the experiments it seems that BRAID is mainly a single-session model. It is well-known that there can be a lot of variability in neural activity across sessions as animals learn to perform the task better or due to some representational drift. This does not seem to be addressed in the paper as far as I can see, but could the authors comment on how BRAID generalises to unseen sessions, and also specifically for the later parts of a session when training on the initial parts (one of the 5 folds)?\n\n* While the experiments are comprehensive and baselines have been compared against, I think comparisons with LFADS (and if possible, CEBRA) could be useful here – both these approaches seem slightly less involved in terms of training complexity, but would represent baselines when ablating for explicit input modelling (LFADS) or explicit dynamics modelling (CEBRA). The idea here is also that LFADS is designed to infer inputs to the dynamical system, so it might perform better than the extended TNDM baseline included currently.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "EafmMt93ga", "reviewer": "Reviewer_XENV", "summary": "The authors tackle the problem of disentangling intrinsic from input-driven neural dynamics underlying behavioral measurements. They do so through a novel nonlinear model called BRAID. Their approach is to conceptualize the intrinsic as the generative (i.e. forward) dynamics, and the input-driven as the (posterior-) predictor dynamics. The authors model each of these components with nonlinear DNNs for flexibility, allowing each to be learned. Importantly, they devise a multi-stage training procedure that prioritizes learning behavior-relevant dynamics, with neural reconstruction placed second.  They show how this approach can help disentangle the neural dynamics directly relevant to behavior, in both synthetic experiments and monkey motor reaching neural activity.", "questions": "Some of the following questions relate directly to the weaknesses raised above. \n- How does your predictor relate to sequential posterior in variational inference?\n- Monkey experiment: What do x^1 and x^2 look like in this task? Is there any non-encoded activity x^3?\n- Could you provide more details on the \"automatic\" selection (L313)? Is it simply picking the best performing?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- The work tackles a very relevant question in the field of statistical neuroscience of disentangling the relationships between behavior, neural activity and inputs, and shows originality in doing so with careful modeling. \n- The modeling is well-formulated in the main text, and made clear for the reader in the appendix and through provided code.\n- The decomposition into multiple stages and sub-components to encourage learning behavior-relevant activity is interesting and novel (to my knowledge).\n- The authors showcase their model on synthetic and real experimental data, showing strong performance in both. \n- The authors help support the significance of their results by (1) comparing them against many baselines and extensive ablations of their model, and (2) performing many runs, providing error bars for all numerical results. \n- The metrics are meaningful and bypass common non-identifiability problems (such as considering the eigenvalues of A)\n- The figures are clean and easy to understand.", "weaknesses": "I am putting a score of 6 (marginal accept) but I would be willing to increase it if the authors can help address my weaknesses/questions, and in particular the \"major\" ones below. \n\nEDIT: The authors have sufficiently addressed the weaknesses below -- I have updated my score accordingly. \n\nMajor:\n- I am not convinced by the claims of real-time inference in the discussion, and the lack of surrounding literature on the predictor components of the model. Variational inference approaches for sequential models similarly model the (sequential, i.e. filtering) posterior and optimize it jointly with the generative model in the ELBO. I believe the similarities could be actually quite exact, in which case a fairer depiction of the relationship with VI is necessary. On that note, Table 1 refers to your method as pure LL optimization, but perhaps it could be more easily interpreted through an ELBO objective.\n- Most of the results are numerical, and the paper lacks a bit in alternate results such as posterior trajectories or visual reconstruction.  \n- In line with the previous comment, the monkey experiment results are numerical and would benefit from some post-training analysis if the goal is to show the model as a modeling and analysis tool. \n\nMedium:\n- Appendix should be for additional but not necessary details to understand the paper. One example of this is the simulations in section 4.1. The notation from L312 follows the appendix where you introduce $f_{C_z}$ and $\\nu, \\bar\\nu$ -- I would make it consistent with eq. (1) or add eq (A.11) to the main text.\n- Identifiability: the authors discuss behavioral- vs neural-relevant activity, but having nonlinear A, Cs are another type of inter-dependence, which could be discussed further.\n- Monkey experiment: numerical performance results are good but sometimes lack transparency in their presentation. For instance, U-BRAID does do better on neural, but that is by construction and does not put BRAID in any lesser light. The corresponding paragraph (3 of section 4.2) however does not acknowledge this better performance. Similarly, the authors refer to mmPLRNN as having an \"unfair advantage\" (L451). I would remove this.\n- The relationship with probabilistic formulations is skimmed but still alluded to with the ELBO/LL in Table 1. I would expand further on these ideas. \n\nMinor:\n- Using $\\cdot$ (\\cdot) instead of $.$ (dot) is more standard for place-holder variables in functions (L191, L312, L1015)\n- Dependency on x^1 on Stage 2 could be made more explicit in the text\n- Table 2: bold entries per metric would be more representative", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "gcb8tdROf3", "reviewer": "Reviewer_kFWw", "summary": "This paper introduced a novel framework, BRAID, modeling nonlinear neural and behavioral dynamics with external input. This model dissociates behaviorally relevant neural dynamics, neural specific and behavioral specific dynamics, and outperforms baselines.", "questions": "1.\tIn Figure 1, should there be a connection between stage 2 and behavior Z_k if you are learning C_z(2) in Stage 2 2b? If so, what is the difference between X_k(1) and X_k(2) since both are connected to neural activity via [C_y(1), C_y(2)] and to behavior via [C_z(1), C_z(2)]?\n2.\tIn equation 1, the model can learn behavior dynamics that are predictable from the input but are not encoded in the recorded neural activity, but how do you make sure that the behavior prediction is not dominated by the input? In Figure 3a, there is a high correlation between your input target position and your behavior cursor position and velocity. If so, can the model learn any information encoded in neural activity?\n3.\tThe model also has an input term in modeling neural activity y, what does this input mean? Because you have an input in latent space that encodes intrinsic contribution, then you also have the same input to encode input-driven contribution. \n4.\tHow do you define your initial states x0?\n5.\tI am curious, 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? \n6.\tFollowing Q5, you mentioned you trained the model until convergence, could you show your learning curve against epochs with zooming in the last epochs if necessary?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "BRAID considers input-driven neural and behavioral dynamics and dissociates them. The authors performed ablation studies to shown which parts of BRAID contributed to the neural and behavior forecasting.", "weaknesses": "1. Although the model has been shown to be highly efficient in modeling neural spiking data, it has not been tested on other modalities, such as widefield calcium imaging.\n2. Although the authors mentioned that their model is not to infer unmeasured input, I still think this may be a weakness of this model, because (for example) the neural and behavioral dynamics can be encoded by unmeasured input.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)