Meta-Dynamical State Space Models for Integrative Neural Data Analysis
Authors: Ayesha Vermani
Josue Nassar
Hyungju Jeon
Matthew Dowling
Il Memming Park
Venue: ICLR 2025 Spotlight (2025)
Keywords: neural dynamics
state-space model
meta learning
Primary area: applications to neuroscience & cognitive science

Abstract: Learning shared structure across environments facilitates rapid learning and adaptive behavior in neural systems. This has been widely demonstrated and applied in machine learning to train models that are capable of generalizing to novel settings. However, there has been limited work exploiting the shared structure in neural activity during similar tasks for learning latent dynamics from neural recordings.
Existing approaches are designed to infer dynamics from a single dataset and cannot be readily adapted to account for statistical heterogeneities across recordings. In this work, we hypothesize that similar tasks admit a corresponding family of
related solutions and propose a novel approach for meta-learning this solution space from task-related neural activity of trained animals. Specifically, we capture the variabilities across recordings on a low-dimensional manifold which concisely parametrizes this family of dynamics, thereby facilitating rapid learning of latent dynamics given new recordings. We demonstrate the efficacy of our approach on
few-shot reconstruction and forecasting of synthetic dynamical systems, and neural recordings from the motor cortex during different arm reaching tasks.

Published as a conference paper at ICLR 2025
META-DYNAMICA STATE SPACE MODELS FOR INTE- GRATIVE NEURAL DATA ANALYSIS
Ayesha Vermani', Josue Nassar?. Hyungju Jeon', Matthew Dowling's. 1l Memming Park'
1 Champalimaud Centre for the Unknown, Champalimaud Foundation, Portugal
2 Ryvivyk, USA
(ayesha vermani, memming, park @eeeeerch Freseaachh.fcammppp aechampalimaud... org
ABSTRACT
Learning shared structure across environments facilitates rapid learning and adap- tive behavior neural systems. This has been widely demonstrated and applied machine learning train models that capable of generalizing to novel settings. However, there has been limited work exploiting the shared structure in neural activity during similar tasks for learning latent dynamics from neural recordings. Existing approaches are designed to infer dynamics from single dataset and can- not be readily adapted to account for statistical heterogeneities across recordings. In this work, we hypothesize that similar tasks admit a corresponding family of related solutions and propose a novel approach for meta-learning this solution space from task-related neural activity of trained animals. Specifically, we capture the variabilities across recordings on a low-dimensional manifold which concisely parametrizes this family of dynamics, thereby facilitating rapid learning of latent dynamics given new recordings. We demonstrate the efficacy of our approach on few-shot reconstruction and forecasting of synthetic dynamical systems, and neural recordings from the motor cortex during different arm reaching tasks.
INTRODUCTION
Latent variable models are widely used in neuroscience extract dynamical structure underlying high-dimensional neural activity (Pandarinath et al., 2018 Schimel et al. 2022 Dowling et al.) 2024). While latent dynamics provide valuable insights into behavior and generate testable hypotheses of neural computation (Luo et al. 2023 Nair et al., 2023). they are typically inferred from a single recording session. As a result, these models are sensitive to small variations the underlying dynamics and exhibit limited generalization capabilities. In parallel, a large body of work in machine learning has focused on training models from diverse datasets that can rapidly adapt to novel settings. However, there has been limited work inferring generalizable dynamical systems from data, with existing approaches mainly applied to settings with known low-dimensional dynamics (Yin et al. 2021 Kirchmeyer et al. 2022).
Integrating noisy neural recordings from different animals and/or tasks for learning the underlying dynamics presents a unique set of challenges. This partly due heterogeneities recordings across sessions such as the number and tuning properties of recorded neurons, as well as different stimulus statistics and behavioral modalities across cognitive tasks. This challenge further compounded by the lack of inductive biases for disentangling the variabilities across dynamics into shared and dataset- specific components. Recent evidence suggests that learned latent dynamics underlying activity of task-trained biological and artificial neural networks demonstrate similarities when engaged in related
tasks (Gallego al. 20088] Maheswaranathan et al. 2019 [Safaie et al. (023). related line of work, neural networks trained to perform multiple cognitive tasks with overlapping cognitive components learn to reuse dynamical motifs, thereby facilitating few-shot adaptation on novel tasks (Turner &
Barak 2023
Driscoll et al. 2024).
Motivated by these observations, we propose a novel framework for meta-learning latent dynamics from neural recordings (Vermani al. 2024a). Our approach is to encode the variations in the latent dynamical structure present across neural recordings in a low-dimensional vector, e E Rd: which refer the dynamical embedding. During training, the model learns adapt common

---

Published as a conference paper at ICLR 2025
A.
B.
Subject C
Embedding Manifold
Generative Model Overview
Embedding
Hypemetwork
Lateni State
Subject Je
T0
Position
voweresr onuet
31k
h1-
rater
Wi
Subject Ni
Observation
Shared weights
Figure A. Neural recordings display heterogeneities in the number and tuning properties of recorded neurons and reflect diverse behavioral responses. The low-dimensional embedding manifold captures this diversity in dynamics. Our method learns adapt a common latent dynamics conditioned on the embedding via low-rank changes to the model parameters.
latent dynamical system model conditioned on the dynamical embedding. We learn the dynamical embedding manifold from diverse collection of neural recordings, allowing rapid learning of latent dynamics in the analysis of data-limited regime commonly encountered in neuroscience experiments.
Our contributions can be summarized as follows:
1. We propose novel parameterization of latent dynamics that facilitates integration and learning of meta-structure over diverse neural recordings.
2. We develop an inference scheme to jointly infer the embedding and latent state trajectory, as well the corresponding dynamics model directly from data.
3. We demonstrate the efficacy of our method on few-shot reconstruction and forecasting for synthetic datasets and motor cortex recordings obtained during different reaching tasks.
CHALLENGES WITH JOINTLY LEARNING DYNAMICS ACROSS DATASETS
Neurons from different sessions and/or subjects partially observed, non-overlapping and ex- hibit diverse response properties. Even chronic recordings from single subject exhibit drift in neural tuning over time (Driscoll et al. 2017). Moreover, non-simultaneously recorded neural activity lack pairwise correspondence between single trials. This makes joint inference of la- tent states and learning the corresponding latent dynamics by integrating different recordings ill- posed and highly non-trivial.
A.
Example
M
B.
Mf
M
M-2
20
To 075
k-step
Figure 2: A. Three different example neural record- ings, where the speed of the latent dynamics varies across them. B. One generative model is trained on M = 2 or M 20 datasets. While increasing the number of datasets allows the model to learn limit cycle, it is unable capture the different speeds leading to poor forecasting performance.
As an illustrative example, let's consider a case where these recordings exhibit oscillatory latent dynamics with variable velocities (Fig. (A) One possible strategy for jointly inferring the dynam- ies from these recordings is learning a shared dy- namics model, along with dataset-specific like- lihood functions that map these dynamics to individual recordings (Pandarinath al.) (018).. However, without additional inductive biases, this strategy does not generally perform well when there are variabilities in the underlying dynamics. Specifically, when learning dynamics from two example datasets (M 2), we observed that model
2

---

Published as a conference paper at ICLR 2025
with shared dynamics either learned separate solutions overfit to dataset, obscuring global structure across recordings (Fig. 2 When we increased the diversity of training data (M 20). the dynamics exhibited a more coherent global structure, albeit with an overlapping solution space frig.23)) As a result, this model had poor forecasting performance of neural activity in both cases, which evident in the k-step ,² (Fig. 2 While have priori knowledge of the source of variations in dynamics for this example, this is typically not the case with real neural recordings. Therefore, we develop an approach for inferring the variation across recordings and use define solution space of related dynamical systems (Fig. TA).
3 INTEGRATING NEURAL RECORDINGS FOR META-LEARNING DYNAMICS
Let 1:M denote neural time series datasets of length T, with v\ E R collected from M different sessions and/or subjects performing related tasks. We are interested in learning generative model that can jointly describe the evolution of the latent states across these datasets and rapidly adapt novel datasets from limited trajectories. In this work, we focus on nonlinear state-space models (SSM), a powerful class of generative models for spatio-tempord datasets. An SSM is described via the following pair of equations (we drop the superscript for ease of presentation),
24 2 ~ Po (ze 21=,,,
(1)
y/ 2 ~ (ye =t),
(2)
where 2, E IRd the latent state at time 1, po - is the dynamics model and po 2x) the likelihood function that maps the latent state to observed data.
We parametrize the dynamics as a Gaussian distribution po ( 22-1)) (= Q), where the mean is modeled by deep neural network (DNN) fa and Q is the covariance matrix] As previous work has shown that highly expressive likelihood and dynamics can cause optimization issues (Bowman et al., 2015), we model the mean of the likelihood as affine function of 2,. For instance, the likelihood for realevalued observations is defined as p (8 2) = N 1Czi + DD R).
3.1 HIERARCHICAL STATE-SPACE MODEL FOR MULTIPLE DATASETS
We introduce a hierarchical structure in the latent dynamical system model capture variations across datasets and jointly describe the spatiotemporal evolution across M neural recordings in a unified SSM. A natural choice for learning this generative model is a fully Bayesian approach, where each dataset would have its own latent dynamics, parameterized by 8', and hierarchical prior would tie these dataset-specific parameters to shared parameters, 0 ~ p(0) (Linderman al. 2019), leading to the following SSM,
0¹ p(0' 1()))))
2 o 0¹ N(2 for (x
(4)
si Pos 1z))
(5)
where dataset specific likelihoods, ps: (v/ =t), are used to account for different dimensionality and/or recording modality. If we assume p00' (0)
)) is Gaussian, i.c., p(0' 10) = NN00000 10, E), we can equivalently express the dynamics for the hierarchical generative model as,
N(E=(0,,))
(6)
+11=-110000 0. N (ote)
where the dataset-specific dynamics parameter, 0', is expressed as sum of the shared parame- ters, 0, and dataset-specific term, e². While this formulation intuitive, the latent dynamics are approximated using a DNN, which introduces substantial number of parameters and constrains scal- ability. To address these limitations, we propose a modified hierarchical framework that significantly improves both scalability and parameter efficiency, making it suitable for large-scale settings.
Specifically, we introduce a low-dimensional latent variable, E Rdde Rdd < Rd. R4ich we refer as the dynamical embedding that encodes dynamical variations across datasets (Rusu et al.
We that Q can also he parameterized via neural network well.
3

---

Published as a conference paper at ICLR 2025
2019). This dataset-specific dynamical embedding subsequently maps to the parameter space of the latent dynamics function via hypernetwork (Ha al. 2016). ho _ R Apart from improving scalability, this formulation also facilitates efficient few-shot learning since it requires simply inferring the embedding given trials from novel recordings. The generative model for this hierarchical SSM is then described as,
e¹ N(0,1), 0. = 0 hhhee)))
(8)
(9)
+ 2 1 ~ (z
(10)
v Poi zt),
(11)
where we drop the prior over the shared dynamics parameter, 0, significantly reducing the dimension- ality of the inference problem. Similar to the hierarchical Bayesian model, all datasets share the same latent dynamies, with the dataset-specifie variation captured by the dynamical embedding, e,
We encourage learning of shared dynamical structure and further improve parameter efficiency by constraining ha make low-rank changes to the parameters of fo (Fig. ( For example, if we parameterize fo as a 2-layer fully-connected neural network and constrain the hypernetwork to only make rank d, changes to the hidden weights, then for would be expressed as,
(o) = W. (( + ho(e')) of (Win 2())
(12)
embedding modification
W a W + uo(e") vo(e')
W Win =i))
(13)
Rd R xdy
Rd
Raixds
where (.) denotes point-nonlinearity, and the two functions ve(e") : R + Rdixa, up(e") Ry + Rd map the embedding representation to form the low-rank perturbations. Both ug and Vy are parametrized via neural network.
3.2 INFERENCE AND LEARNING
Given 1:M we want to infer both the latent states, 21-7" and the dynamical embeddings, el:M
"e....,,eM as well as learn the parameters of the generative model, θ {0, 0 o' M }. Exact
inference and learning requires computing the posterior, pe (2 el:d I.MM ), and log marginal likelihood, log (
1:A1
). which are both intractable.
In this paper, we use a sequential variational autoencoder an extension of variational autoencoders for state-space models speelsssspiiiilllyy..... the Deep Kalman Filter (DKF) (Krishnan al., (015)., to circumvent this issue. In order to learn the generative model, we maximize a lower-bound to the log marginal likelihood (commonly referred to the ELBO). The ELBO for 1:M is defined follows (trial indices are omitted for ease of notation).
C(yi:M)
Σ Bag. [logPs (vil [()] 1,
(14)
- -Eys [DK (g Pier et) po. po.o (e. - DxL (9a ( l|p(et)
where Ga and q8 encoders that approximate the posterior distributions the dynamical embedding and latent state for dataset respectively, and the joint expectation factorizes as Ey = As described in Sec. 2 one of the challenges with integrating recordings in common latent space is different dimensionalittis (number of recorded neurons) well the dependence of neural activity on the shared latent space. We address this by training additional
Ego(s; "e.
read-in networks 22 R + ddds for each dataset that map y an intermediate vector, which denote by & E Rd This read-in network ensures that the latent states and dynamical-emi -emnicallembeddinng inferred from each dataset are aligned to live in the same space (Vermani et al. 2024b).
4

---

Published as a conference paper at ICLR 2025
While there are many choices for parameterizing the encoders, follow the parameterization in (Krishnan et al. 2015) for simplicity defined as follows,
Farr = 52' (v).
(15)
la(e' I.. (et, agglua (Vh, 1.)) agglo? 1r),.))))
(16)
T
c Birr et)
II Mx lug ==1
e(.), (concat( e()),,
(17)
where y denotes a randomly sampled mini-batch of trials b from dataset i, concat is the concatenation operation, and agg is an aggregation operation. We aggregate the dynamical embedding over trials in mini-batee batch that belong to the same dataset since we are interested in capturing inter-dataset, rather than intra-dataset variations, in the underlying dynamical systems. In practice, we parameterize paa (-), of (.) by a bidirectional recurrent neural network, and (*)()), o.(() by a regular recurrent network, and agg corresponds simple averaging function. We emphasize that Pax oa pa and of are shared across all datasets (See Fig. 14 for details on inference).
3.3 PROOF OF CONCEPT
A.
M-20
B. Inferred embedding
As a proof of concept, we revisit the motivating example presented in Section 2 as means to validate the efficacy of our approach and inves- tigate how it unifies dynamics across datasets. For both A = 2 and M = 20 datasets, we used an embedding dimensionality of and al- lowed the network to make a rank- change to the dynamics parameters.
jel
Velocity
C.
I-step
Prediction
After training, we observed that the shared dy- namics (when e 0) converged to limit ey- cle with a slow velocity (Fig. ( capturnnn the global topology that is shared across all datasets-aad -nd the model learned to modulate the velocity of the dynamics conditioned on the dynamical embedding which strongly correlated with the dataset specific velocity (Fig 3B). This demonstrated that the proposed approach is able to capture dataset-specific variability. Lastly, Fig. BC demonstrates that the proposed approach is able to forecast well for both M = 2 and M 20 datasets. We include further validation experiments when there is no model mismatch as well the generalization of the trained model to new data in Appendix B We additionally include results on these recordings from multi-session CEBRA (Schneider et al. 2023) in Appendix
0.9
Central Dynamics Learnt by the Model
M-2
Figure 3: A. Mean dynamical system correspond- ing to the slowest velocity recording learned by the proposed approach when trained with M = 20 datasets. B. Samples from the inferred dynamical embedding for each dataset (see eq. [6]. C. Fore- casting 22 at (A: 5te-step for models trained with M = 2 or M = 20 datasets.
RELATED WORKS
Multi-Dataset Training in Neuroscience. Previous work has explored multi-dataset training for extracting latent representations in neuroscience, especially across datasets recorded during the same behavioral tasks. LFADS (Pandarinath al. (2018), a variant of the seq VAE framework, used session-stitching with datasct-specific likelihood functions, but focused on single-animal record- Linderman et al. (009) used hierarchical Bayesian state-space model with switching linear dynamical systems, while Herrero- Vidal et al. (2021) developed a joint model with shared linear dynamies and dataset-specific likelihoods. In contrast to these approaches, we incorporate a more
evaluate alternative inference and learning formulations Appendix
Note that we plot the absolute embedding samples since the likelihood function can introduce arbitrary invariance such direction flipping, rotation, and on.
5

---

Published as a conference paper at ICLR 2025
expressive function to approximate the underlying family of dynamical systems which can disentan- gle variabilities across recordings. CEBRA (Schneider et al. (023) and CS- VAE (Yi (003) have been developed for extracting latent representations by integrating multiple datasets. The multi-session training objective in CEBRA promotes invariant feature learning across datasets, while CS-VAE partitions the latent space to learn structured features from behavioral videos. In frame- work, we jointly infer latent trajectories from multi-session recordings and learn a unified generative model that captures variations in dynamical systems from the inferred trajectories. Recently, there has been growing interest in using diverse neural recordings for training large-scale foundation models in neuroscience (Ye et al. 2023 Zhang et al. 2023] Caro al. 2024) (Azabou al. 2024] Vermani et al.) 2024a). These models leverage transformer-basee architectures which lack recurrent hidden states and only incorporate temporal information indirectly via positional encoding. While our approach shares the same broad goal of pretraining a single generative model for rapid learning on downstream recordings, the focus of our work is on learning a family of dynamical systems underlying recordings.
Neural Network Models in Neuroscience Integrative modeling of dynamical behaviors has also been explored in RNN models of neural systems (Yang et al.) 2019 Driscoll et al. 2024). In Driscoll et al. (2024), the authors trained an RNN to perform multiple cognitive tasks and observed motifs corresponding to distinct dynamical behaviors. The broad idea of dynamical structure re-use is similar our work but there are subtle differences-ww are interested in capturing both topological and geometrical differences, and the "context" is learned from data. The motifs in Driscoll et al.) (2024) corresponded to a distinct fixed point structure or topology with a pre-specified context input that could push the dynamics to task-relevant regions in the state space. The embedding analysis in Cotler et al. 2023). where meta-model was trained capture the activity of multiple trained RNNs is quite similar to our main idea since they observed similar dynamical properties in models that were close in the embedding space. Recent work on modeling motor adaptation (Pellegrino et al.) 2023) by low-tensor rank learning in RNNs is broadly similar to our work since the authors adapt the weights low-rank RNN to capture variations in dynamics across trials. contrast, we are interested in modeling dynamical variations across recording sessions and/or tasks.
Recurrent
Additional related works can be found in Appendix A
5 EXPERIMENTS
We first validate the proposed method on synthetic data and then test method on neural recordings from the primary motor and premotor cortex. We compare the proposed approach against the following baselines for all experiments.
We train separate Single Session model using the seq VAE framework on each dataset. Given sufficient training data, this should result in the best performance, but will fail in trial-limited regimes. We consider multi-session Shared Dynamics model with dppaset- likelihoods (Pandarinath al. 2018 Herrero-Vide et al. 2021). We also compare against a baseline where the embedding provided as an additional input to the dynamics model (Embedding-Inpell
similar formulation
to CAVIA (Concat) (Zinggraf et al.) 2019) and DYNAMO (Cotler et al., 2023). We also test the hypernetwork parametrization proposed in CoDA (Kirchmeyer et al. (0022). where the hypernetwork adapts all parameters as a linear function of the dynamical embedding (Linear-Adapter),
We include additional baselines for the motor cortex experiment; we evaluate single session LFADS (Pandarinath et al. (018)) with the controller as an alternative generative model comparison.
We also consider different methods for learning and inference. Specifically, we include single-session models as well as our proposed generative model trained using Variational Sequential Monte Carlo (VSMC) (Naesseth et al. 2018). and the Deep Variational Bayes Filter (DVBF) (Karl et al. 2016).
For all experiments, we split each of the M datasets into a training and test set and report recon- struction and forecasting metrics the test set. To measure the generalization performance, also report these metrics on held-out datasets. Further details on training and evaluation metrics can be found in Appendix G

---

Published as a conference paper at ICLR 2025
5.1 BIFURCATING SYSTEMS
In these experiments, we test whether our method could capture variations across multiple datasets, particularly in the presence of sigmificant dynamical shifts, such as bifurcations commonly observed in real neural populations. We selected two parametric classes of dynamical systems for testing: i) system undergoing Hopf bifurication and, ii) the unforced Duffing system. We include the results for training on datasets generated only from the Hopf system in Appendix C and we discuss the results from joint training on both systems here. We briefty outline the data generation process for the Duffing system (details of the data generation for the Hopf system can be found in Appendix E.2)
The latent trajectories for the Duffing system were generated from a family of stochastic differential equations,
1 = 22 + 5dW,
2 = z1 2b' 00 + czi) ) + 55dW,
(18)
with e = 0.1, a. E R, and dW, denoting the Wiener process. In Fig. 4A we visualize how the dynamical system changes as and b vary. We chose M 20 pairs of (a', b') values (Fig 13
A.
Example True Dynamics bm.01
Leamnd Dynamics
B.
Generalization on Training Tasks
Reconstrction
Method
Single
Shoved Dynumics
Emtbedding Input
0.8
C.
Prodiction P.
0.5
25 k-step
varying
Figure 4: A. (Left) True underlying dynamics from some example datasets used for pretraining as function of parameters and and (Right) the embedding conditioned dynamics learnt by model. B., C. Mean reconstruction and forecasting 2² of the observations for all datasets used for pretraining on test trials.
and generated latent trajectories of length T = 300. Observations were generated according to Vt N/(C= =4, 0.01T) with the dimensionality the observations varying between and 100. In addition to these 20 Duffing system datasets, we included 11 datasets from the Hopf system (Appendix C). and used 128 trajectories from each these 31 datasets for training all methods. We report performance 64 test trajectories from each dataset. We used de for all embedding- conditioned approaches and constrained the hypemnework to make rank d, = 1 changes for our approach.
Our approach learned a good approximation to the ground truth dynamics of the Duffing oscillator system, successfully disentangling different dynamical regimes (Fig. 4 B). Apart from learning the underlying topology of dynamics, also better captured the geometrical properties compared to other cmbedding-conditton the Hopf system-while our approach approximated the ground-truth system well, the Embedding- Input baseline displayed interference between dynamies and the Linear-Adapter learned poor approximation to the ground-trun system (Fig. 16.. Consequently, our approach outperformed other methods on forecasting observations with all methods having comparable reconstruction performance (Fig. 43, C). Notably, apart from the de, we used the same architecture as when training only the Hopf datasets, and did not observe any drop in performance for our approach, contrast to baselines (Fig.) ITC (Bottom), Fig. 4C).
baselines (Fig. [5 We observed similar results for datasets from
Next, we tested the few-shot performance of all methods on new datasets, two generated from the Duffing oscillator system and one from the Hopf system, function of the number of trials used for learning the dataset specific read-in network, Ω² and likelihood. Our approach and the Linear-Adapter demonstrated comparable forecasting performance when using n, 1 and n, training trajectories. However, with n, = 16 training trials, unlike other methods, our approach continued to improved and outperformed them (Table This could be explained by looking the
7

---

Published as a conference paper at ICLR 2025
f1
f1
Our Method
0.69 +: 0.072
0.78 +: 0.051
0.87 + 0.037
Linear-Are Addater
0.68 = 0.08
0.79 + 0.026
0.74 - 0.039
Single Session
0.47 = 0.119
0.79 + 0.014
0.79 0.047
Shared Dynamies
-0.31 0.103
-0.34 0.086
-0.13 0.065
Embedding-Inge
0.59 + 0.084
0.77 0.04
0.74 + 0.039
Table 1:
Few shot forecasting performance (k = 30-step) on 3 held-out datasets as a function of n.. number of trials used learn dataset specific read-in network and likelihood. (#: s.e.m)
inferred embedding on held-out datasets-ar as we increased the number of training trajectories, the model was able consistently align to the "correct" embedding (Fig. 17
5.2 MOTOR CORTEX RECORDINGS
Next, tested the applicability of the proposed approach neural data. We used single and multi- unit neural population recordings from the motor and premotor cortex during two behavioral tasks-the Centre-Our (CO) and Maze reaching tasks (Perich et al. 2018 Gallego al.) 2020 Churchland to reach one of eight target
iet al. 2012). In the CO task, subjects are trained to use manipulandun touch screen to reach target location, while potentially avoiding obstacles. These recordings spanned different sessions, animals, and labs, and involved different behavioral modalities, while still having related behavioral components, making them good testbed for evaluating various methods. For training, we used 40 sessions from the CO task, from subjects M and C, and sessions from the Maze task from subjects and Ni. We set the dimensionality of latent dynamics to d2 = 30, and used dynamies models. For our approach, we constrain the hypernetwork
locations screen. In the Maze task, subjects
embedding dimensionality of d, = 2,
for all embedding-conditionn to make rank d, 6 changes, although we verified that the performance was not sensitive to d, (Fig 18). As proxy for how well the various approaches learned the underlying dynamics, we report metrics on inferring the hand velocity using reconstructed and forecasted neural data from the models. Note that we align all recordings to the movement onset (details in Appendix
Inferred embedding
Condition-nveraged Dynamics
Position
Monkey Predicted
Maze session Sub Ni
CustaoOu Seb
CO session Sub C
Cestoo-Ca
Muse Seb Jo
P= 0.32
Figure 5: Visualizing the embedding manifold. (Left) Each point corresponds to a sample from the inferred embedding distribution (see eq. [6] corresponding to each recording. (Right) The condition-averagel latent dynamics for session from Maze (Sub Ni) (Top) and CO Session (Bottom) generated by the model, along with the corresponding real and forecasted behavior.
The inferred dynamical embedding displayed distinct structures across behavioral tasks and subjects (Fig. [5] Left). While the CO task involves more stereotyped straight reaching behavior with the same stimulus conditions across datasets, the Maze task has more complex stimulus statistics which vary across sessions. The family learned dynamics reflected this heterogeneity across recordings. We visualize these learned dynamical systems for two example sessions, one from each task, in Fig 5 (Right). Specifically, we used the trained encoders, qa and qa to estimate the latent state and embedding the beginning of movement onset. We subsequently generate the latent dynamics from that state using fe.e till the end of the movement onset. The condition condition-averaged principal components (PCs) of these generated latents shown in the figure.
8

---

Published as a conference paper at ICLR 2025
We observed that most of the approaches had adequate performance reconstructing velocity from neural recordings, with our method and Linear- Andaterr outperforming single session reconstruction performance on the CO task (Fig. [A,, top). Multi-Session CEBRA was not able to adequately capture the variability in the Maze sessions and had low reconstruction r². In terms of forecasting, the single-session model trained using the seq VAE framework had the best performance. Notably, our approach managed to balance learning both the CO and Maze tasks relative to other multi-session baselines, with all performing better on the CO task than the Maze (Fig. ØA. bottom). The generative model learned from CEBRA had poor forecasting performance which resulted in a negative r value (not plotted). Next, we tested we can transfer these learned dynamics to new recordings as we varied n, from 8 to 64 trials for learning the read-in network and likelihood. We used trials from 2 held-out sessions from Sub C and M, as well as sessions from a new subject (Sub T) for evaluating all methods. We observed that our approach consistently performed well on both reconstruction and forecasting for held-out sessions from previously seen subjects, and reached good performance on sessions from Sub T as we increased the training trials (Fig. 6B, C (n, = 32)). Moreover, our method outperformed all other baselines on forecasting, especially in very low-sample regimes, while having comparable reconstruction performance (Fig. 19
A.
Training sessions
B.
Few Shot Performace
C. Hand Velocity Trajectories
(n. -2)
Reconstruction
0.7
True Recomiincted Prrdicted
2 0.5
0.5
AAA
Held-out Session Sub
Held scasion
New subject
Prodiction
0.6
0.
0.4
0.0
Medion
Laoer Adayeer
Sagle Scario
Sharod Divs
brest
LEADS
MS Gabe
0.2
Session from Sub T
16
32
64
Training Samples (n)
trials
Figure 6: A. (top) 2 for hand velocity decoding from reconstructed and (bottom) forecasted neural observations for Maze and Centre-Out sessions. B. Behavior reconstruction (top) and forecasting (bottom) performance on held-out sessions and sessions from a new subject as a function of the number of training samples. C. Hand velocity trajectories (400 ms after movement onset) predicted by our approach on 17 test trials from held-out session (top) and 13 test trials from a session on a new subject (bottom), after using n, = 32 trials for aligning to the pre-trained model.
Next, we evaluated the impact of the inference framework on effective learning and few-shot per- formance. We specifically tested single session models as well as our proposed generative model trained after performing inference using VSMC and DVBF (Details in Appendix DD. In both cases, we observed that the inferred embedding distribution learned the underlying dynamical structure across datasets (Fig. A)). Morcover, we were able to similarly exploit this learned structure for few-shot forecasting
novel recording sessions (Fig. B2 We additionally investigated the effect of large-scale training for sample-efficient transfer on downstream tasks by only pretraining the model on 128 trials from 4 sessions spanning different tasks and subjects. Even in this case, the embedding distribution displayed clear clustering based on the task and subject. Moreover, the model performed comparably the Single-Session model on reconstruction, while outperforming it on prediction for both tasks (Fig. 20 A. B). However, it demonstrated poor performance on new sessions given limited trials for learning the read-in and likelihood parameters (Fig. 20 C), underscoring the importance of large-scale training for generalizing to novel settings.
Finally, we probed the differences in the latent state evolution given the same initial condition while interpolating across the learned embedding. In order to do this, we chose an example session from the Maze and CO datasets and obtained their corresponding dynamical embedding from the model, shown as solid blue and green circle in Fig. 7 (middle), respectively. A grid of points was sampled around each of these inferred embeddings (shown as shaded squares in Fig. 7 middle), and for each point we obtained the corresponding low-rank parameter changes to generate the latent trajectories. We observed that the embedding space learned a continuous representation of dynamics, which was reflected in similar predicted behaviors close to the original learned embedding (Fig 7). Interestingly,

---

Published as a conference paper at ICLR 2025
Interpolating across embedding space
Predicted behavior (Maze)
Predicted behavior (CO)
The predicted behavior for a Maze (Sub Je) session and CO (Sub C) session at 9 grid points around the original inferred embedding. The point closest to the original embedding is highlighted in blue and green respectively.
Figure 7:
when we interpolated through the entire embedding space. the predicted behavior and corresponding dynamics continuously varied as well. Specifically, the predicted behavior and dynamics trajectories on the CO session demonstrated similarities over a large portion of the embedding space, with the trajectories shifting to more curved reaches further from the original embedding (Fig. (2 On the Maze task, the trajectories demonstrated more heterogeneity in responses, and decayed to a fixed point further away from the original embedding (Fig. 22
6 DISCUSSION
We present a novel framework for jointly inferring and learning latent dynamics from heteroge- neous neural recordings across sessions/subjects during related behavioral tasks. To the best of our knowledge, this is the first approach that facilitates learning family of dynamical systems from heterogeneous recordings in a unified latent space, while providing concise, interpretable manifold over dynamical systems. Our meta meta-learning approach mitigates the challenges of statistical inference from limited data, common issue arising from the high flexibility models used to approximate latent dynamics. Empirical evaluations demonstrate that the learned embedding manifold provides useful inductive bias for learning from limited samples, with our proposed parametrization offering greater flexibility in capturing diverse dynamics while minimizing interference. We demonstrate that the few-shot performance of our proposed generative model is largely agnostic to the inference method. We observe that the generalization of our model depends on the amount of training data- when trained on smaller datasets, the model learns specialized solutions, whereas more data allows learn shared dynamical structures. This work enhances our capability integrate, analyze, and interpret complex neural dynamics across diverse experimental conditions, broadening the scope of scientific inquiries possible neuroscience.
LIMITATIONS AND FUTURE WORK
Our current framework uses event aligned neural observations; in the future, it would be usefuu to incorporate task-related events, to broaden its applicability to complex, unstructured tasks. Further, the model's generalization to novel settings depends accurate embedding inference, challenge noted in previous works that disentangle task inference and representation learning (Hummos et al., 2024). However, we observe consistent improvement in embedding inference with increase in the number of training samples from novel recordings. Our empirical observations demonstrate that using a hypernetwork improves the expressivity of the dynamical systems model relative other parametrizations. It would be interesting to investigate the theoretical basis of this observation in the future. While our latent dynamics parametrization expressive, assumes shared structure across related tasks. Future work could extend the model to accommodate recordings without expected shared structures (for instance, by adding explicit modularity (Márton et al., 2021)). Investigating the performance of bebeedddnn--onditiionn another avenue for future research. Finally, the embedding manifold provides a map for interpolating across different dynamics. While we focus rapid learning in this paper, our framework could have interesting applications for studying inter-subject variability, learning-indeen inducedd changes in dynamics, or changes in dynamics across tasks in the future.
low-rank adaptation RNN-based architectures presents
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
We would like to thank the anonymous reviewers for providing helpful feedback on the manuscript. This work was supported by NIH RF1-DA056404 and the Portuguese Recovery and Resilience Plan (PPR), through project number 62, Center for Responsible AI, and the Portuguese national funds, through FCT- Fundação para a Ciência e Tecnologia in the context of the project UIDB/0444//220.
REFERENCES
Mehdi Azabou, Vinam Arora, Venkataramana Ganesh, Ximeng Mao, Santosh Nachimuthu, Michael Mendelson, Blake Richards, Matthew Perich, Guillaume Lajoie, and Eva Dyer. A unified, scalable framework for neural population decoding. Advances in Neural Information Processing Systems, 36, 2024.
Samuel R Bowman, Luke Vilnis, Oriol Vinyals, Andrew M Dai, Rafal Jozefowicz, and Samy Bengio. Generating sentences from continuous space. arXiv preprint arXiv:1511.06349 2015.
Josue Ortega Caro, Antonio
Fonseca, Christopher Averill, Syed A. Rizvi, Matteo Rosati, James L. Cross, Prateek Mittal, Emanuele Zappala, Daniel Levine, Rahul M. Dhodapkar, Insu Han, Amin Karbasi. Chadi G. Abdallah, and David van Dijk. Brainlm: A foundation model for brain activity recordings. bioRxiv, 2024.
Mark M Churchland, John Cunningham, Matthew T Kaufman, Justin D Foster, Paul Nuyujukian, Stephen Ryu, and Krishna V Shenoy. Neural population dynamics during reaching. Nature, 487 487 (7405):51-56, 2012.
Jordan Cotler, Kai Sheng Tai, Felipe Hernández, Blake Elias, and David Sussillo. Analyzing populations of neural networks via dynamical model embedding. arXiv [cs.LG], February 2023.
Arnaud Doucet, Nando De Freitas, and Neil Gordon. An introduction to sequential monte carlo methods. Sequential Monte Carlo methods in practice pp. 14, 2001.
Matthew Dowling, Yuan Zhao, and Memming Park. eXponential FAmily dynamical systems (XFADS): Large-scale nonlinear gaussian state-space modeling. In Advances in Neural Information Processsigg Systems (NeurIPS), December 2024. URL https: http////pp////eeee//e/e/ee net forum?
id Ln8oggh22225
Laura N Driscoll, Noah Pettit, Matthias Minderer, Selmaan Chettih, and Christopher D Harvey. Dynamic reorganization of neuronal activity patterns in parietal cortex. Cell, 170(5)9986-99)
2017.
Laura N. Driscoll, Krishna Shenoy, and David Sussillo. Flexible multitask computation in recurrent networks utilizes shared dynamical motifs. Nature Neuroscience. 2777)1349-1133, 2024.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of
deep networks. In International conference on machine learning, pp. 1235.133. PMLR, 2017.
Juan A. Gallego, Matthew G. Perich, Stephanie N. Naufel, Christian Ethier, Sara A. Solla, and Lee Miller. Cortical population activity within a preserved neural manifold underlies multiple motor behaviors. Nature Communications, 91)):2233, 2018.
Juan A Gallego, Matthew G Perich, Raeed H Chowdhury, Sara A Solla, and Lee E Miller. Long-term stability of cortical population dynamics underlying consistent behavior. Nature neuroscience, 23 (2):260-270, 2020.
David Ha, Andrew M Dai, and Quoc Le. Hypernetworks. In International Conference Learning Representations, 2016.
Pedro Herrero- Vidal, Dmitry Rinberg, and Cristina Savin. Across-animas odor decoding by proba- bilistic manifold alignment. Advances Neural Information Processing Systems, 34:20360-222 20372, 2021.
11

---

Published as a conference paper at ICLR 2025
Ali Hummos, Felipe del Río, Brabeeba Mien Wang, Julio Hurtado, Cristian B. Calderon, and Guangyu Robert Yang. Gradient-basee inference of abstract task representations for generalization in neural networks, 2024.
Maximilian Karl, Maximilian Soelch, Justin Bayer, and Patrick Van der Smagt. Deep varia- tional bayes filters: Unsupervised learning of state space models from raw data. arXiv preprint arXiv: 06432, 2016.
Matthieu Kirchmeyer. Yuan Yin, Jérémie Donå, Nicolas Baskiotis, Alain Rakotomamonjy, and Patrick Gallinari. Generalizing to new physical systems via context-informed dynamics model. In International Conference on Machine Learning, pp- 1128-11301. PMLR, 2022.
Rahul G Krishnan, Uri Shalit, and David Sontag.
Deep kalman filters.
arXiv preprint
arXiv:1511.0512,, 2015.
Zhenguo Li, Fengwei Zhou, Fei Chen, and Hang Li. Meta-sgd: Learning to learn quickly for few-shot learning. arXiv preprint arXiv: vrXXiv:707.00900 0000.77775..... 2017.
Scott Linderman, Annika Nichols, David Blei, Manuel Zimmer, and Liam Paninski. Hierarchical recurrent state space models reveal discrete and continuous dynamics of neural activity in c. elegans. BioRxiv, pp. 621540, 2019.
Thomas Zhihao Luo, Timothy Doyeon Kim, Diksha Gupta, Adrian Bondy, Charles D Kopec, Verity A Elliot, Brian DePasquale, and Carlos D Brody. Transitions in dynamical regime and neural mode underlie perceptual decision-making. bioRxiv, pp. 2023-10, 2023.
Niru Maheswaranathan, Alex Williams, Matthew Golub, Surya Ganguli, and David Sussillo. Uni- versality and individuality in neural dynamics across large populations of recurrent networks. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
Christian David Márton, Léo Gagnon, Guillaume Lajoie, and Kanaka Rajan. Efficient and robust multi-task learning in the brain with modular latent primitives. arXiv preprint arXiv:2 14108, 2021.
Christian Naesseth, Scott Linderman, Rajesh Ranganath, and David Blei. Variational sequential monte carlo. In International conference on artificial intelligence and statistics, pp. 968-977.
PMLR, 2018.
Aditya Nair, Tomomi Karigo, Bin Yang. Surya Ganguli, Mark J. Schnitzer, Scott W. Linderman, David J. Anderson, and Ann Kennedy. An approximate line attractor in the hypothalamus encodes an aggressive state. Cell, 186(1):178-193.15 2023.
Alex Nichol and John Schulman. Reptile: a scalable metalearning algorithm. arXiv preprint
18K3::8803002000 2(3):4, 2018.
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, Jaimie M. Henderson, Krishna V. Shenoy, L. F. Abbott, and David Sussillo. Inferring single-trial neural population dynamics using sequential auto-encoders. Nature Methods, 15(10):8058815 2018.
Arthur Pellegrino, N Alex Cayco Gajie, and Angus Chadwick. Low tensor rank learning of neural dynamics. Advances in Neural Information Processing Systems, 36: 1 161141177722 2023.
Matthew G Perich, Juan A Gallego, and Lee E Miller. A neural population mechanism for rapid learning. Neuron, 100(4):964-97, 2018
Andrei A. Rusu, Dushyant Rao, Jakub Sygnowski, Oriol Vinyals, Razvan Pascanu, Simon Osindero, and Raia Hadsell. Meta-learning with latent embedding optimization, 2019.
Mostafa Safaie, Joanna C. Chang, Junchol Park, Lee E. Miller, Joshua T. Dudman, Matthew G. Perich, and Juan A. Gallego. Preserved neural dynamics across animals performing similar behaviour. Nature, 623(7988) 2023.
12

---

Published as a conference paper at ICLR 2025
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. In International Confer- ence on Learning Representations, 2022.
Steffen Schneider, Jin Hwa Lee, and Mackenzie Weygandt Mathis. Learnable latent embeddings for joint behavioural and neural analysis. Nature, 617(7960):3603366 2023.
Andrew R Sedler and Chethan Pandarinath. Ifads-torch: A modular and extensible implementation of latent factor analysis via dynamical systems. arXiv preprint arXiv:2309.01230, 2023.
Elia Turner and Omri Barak. The simplicity bias multi-task RNNs: Shared attractors. reuse of dynamics, and geometric representation. In A. Oh, T. Naumann, A. Globerson, K. Saenko, Hardt, and S. Levine (eds.), Advances in Neural Information Processing Systems, volume 36, pp. 25495 25507. Curran Associates, Inc., 2023.
Ayesha Vermani, Matthew Dowling, Hyungju Jeon, Ian Jordan, Josue Nassar, Yves Bernaerts, Yuan Zhao, Steven Van Vaerenbergh, and Memming Park. Real-time machine learning strategies for new kind of neuroscience experiments. In European Signal Processing Conference, 2024a.
Ayesha Vermani, Il Memming Park, and Josue Nassar. Leveraging generative models for unsupervised alignment of neural time series data. In International Conference on Learning Representations (ICLR), May 2024b.
wwwp::/w.eeeeeeeeeeeeeeee wtt.:/rrrmmmimddh/yyy//y /ttttps://openeeeeee htt.ee/rreeeeee..eeeeoorrrrrrroo
Rui Wang, Robin Walters, and Rose Yu. Meta-learning dynamics forecasting using task inference. Advances in Neural Information Processing Systems, 35:2 55:21640-21653, 2022.
Guangyu Robert Yang, Madhura R Joglekar, H Francis Song, William T Newsome, and Xiao-Jing Wang. Task representations in neural networks trained perform many cognitive tasks. Nature neuroxcience, 22(2):297-306 2019.
Joel Ye, Jennifer I.. Collinger, Leila Wehbe, and Robert Gaunt. Neural data transformer 2: Multi- context pretraining for neural spiking activity. bioRxiv, 2023.
Daiyao Yi, Simon Musall, Anne Churchland, Nancy Indilla-Coreanl and Shreya Saxena. Disentan- gled multi-subject and social behavioral representations through constrained subspace variational autoencoder (cs-vae). eLife, 12, 2023.
Yuan Yin, Ibrahim Ayed, Emmanuel de Bézenac, Nicolas Baskiotis, and Patrick Gallinari. Leads: Learning dynamical systems that generalize across environments. Advances in Neural Information Processing Systems, 34:7561-7573, 2021.
Daoze Zhang, Zhizhang Yuan, YANG YANG, Junru Chen, Jingjing Wang, and Yafeng Li. Brant: Foundation model for intracranial neural signal. In A. Oh, T. Naumann, Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Advances in Neural Information Processing Systems, volume 36, pp. 26304-26321. Curran Associates, Inc., 2023.
Luisa Zintgraf, Kyriacos Shiarli, Vitaly Kurin, Katja Hofmann, and Shimon Whiteson. Fast context adaptation via meta-learning. In International Conference on Machine Learning, pp. 7693-7700 PMLR, 2019.
CONTENTS
Additional Related Works
13
Proof-of-Concept Experiment
14
C
Hopf Bifurcation Systems
14
13

---

Published as a conference paper at ICLR 2025
D Alternative Inference and Learning Approaches
15
D.1 Variational SMC
16
D.2 Deep Variational Bayes Filter
16
E Data Generation Details
17
E.1
Limit Cycle
17
E,2
Hopf Bifurcation
17
E.3 Duffing System
18
E,4 Motor Cortex Recordings
18
Additional Figures
18
G Experiment Details
22
G.I Training
22
G.2 Figure Generation
22
G.3 Metries
22
G.4 Hyperparameters
23
A ADDITIONAL RELATED WORKS
Several meta-learning approaches have been developed for few-shot adaptation, including gradient- based methods (Finn et al. 2017 Li et al. 2017] Nichol & Schulman, [2018] Zintgraf et al. (2019). Amongst these, LEO Rusu et al. 2019) shares the same idea of meta-learning in low-dimensional space of parameter embeddings. However, test-time, and have had limited success for meta-learning dynamies. Similar to work (Cotler ct al. 2023) also learns an embedding space of dynamics learned from trained RNNs, however, we interested in learning dynamics directly from data. Some methods for learning generalizable dynamics been previously
mudient-based approaches require fine-tuning during
Beroosd-DyAAl (Wang et al., 2022) adapts across environments by neural style transfer, however it operates on images of dynamical systems, LEADS (Yin al. (021) learns a constrained dynamics function that is directly added to some base dynamics function, and CoDA (Kirhhmeyer al. 2022) which learns task-specific parameter changes conditioned low- dimensional context similar to our approach. However, these approaches were applied in supervised settings on low-dimensional systems whereas we operate in an unsupervised setting.
B PROOF-OF-CONCEPT EXPERIMENT
1-shot performance
1-shot Performance, We evaluated the generalization performance of our approach on new dataset with M+++1 = 4.1, which was not included the training set, by using training trajectory to
True Observations
train a new read-in network, 2 and likelihood / +1 After
training, the model displayed similar prediction performance on the new dataset (r2 = 0.94 + 0.001) (Fig. 8
rocomsinacted
predicted
No Model Mismatch. Here, we investigated the performance of our approach when there is no mismatch between the proposed generative model and the true system. For this experiment, we generated synthetic data from the model trained on M 20 datasets. We used the observations from validation trials till = 100 to infer the embedding, e¹, and latent state, 2/. We subsequently used the dynamics model to generate latent trajectories of length 250, 2t 2551: and mapped them back the observations via the learned
100
200
time
Figure 8: (Top) True observa- tions on new data with w 4.1 and (Bottom) the corre- sponding reconstructed and predicted observations after aligning to the trained model.
14

---

Published as a conference paper at ICLR 2025
likelihood function. We re- re-trrined model with the same architecture while keeping the likelihood readout and read-in parameters fixed since the likelihood could arbitrarily flip the direction of the flow field.
Similar to the ground truth generative model, the inferred embedding co co-varied with the different velocities (Fig. g left). Further, the model recovered the correct topology of the ground truth dynamics (Fig. 98,, reflected in the forecasting performance on held out test trials (Fig, Big..A right).
A.
Ground Truth
B.
0.8
Velocity
k-step
Ground Truth
Recovered
Figure 9: A. (Left) Samples from the inferred embedding (see eq. 16) after training on trajectories generated from model overlaid the ground truth embeddings. (Right) Forecasting performance of the trained and ground truth model on held out trials. B. Ground truth dynamics generated from an embedding sample and the corresponding recovered dynamics.
Multi-Session CEBRA. We evaluate the performance of multi-session CEBRA, an approach for inferring latents by integrating datasets. This variant of CEBRA is designed to learn invariant latent features across datasets, and has not been evaluated on recordings with variations in underlying dynamical features. In this experiment, we fit M 20 datasets using CEBRA and post-hoc trained a generative model with shared dynamics and dataset-specific likelihood functions, since does not learn generative model. After training CEBRA on these datasets, we observed that the model recovered oscillatory latent trajectories; however, these trajectories were jagged and did not capture the characteristics of the true latents (Fig. [OAA). Next, trained generative model using these latent trajectories. We observed that the learned dynamical system managed to capture a global limit-cycle like structure (Fig. [O8, left). However, this limit cycle also contained fixed-point like structure causing rapid slow-down or noise-induced oscillations, capturing the characteristics of the latent trajectories inferred by CEBRA. Due to this behavior, we observed poor forecasting and reconstruction of observations (Fig. 10B, right).
A.
Example Latent Trajectories
B.
Generative Model Learning
True
Prediction
055
reconstnuctdd
producned
Multi-Session CEBRA
30
k-step
100
100
time
time
Figure 10: A. Example latent trajectories from two datasets (Top) and inferred latent trajectories from Multi-Session CEBRA trained on M = 20 datasets. B. The dynamics learned by a generative model trained on the latent trajectories (Left). Example reconstructed and predicted observations on the two datasets by using the learned generative model (Middle). The k-step prediction 7d of the forecasted observations for M = 20 datasets.
C HOPF BIFURCATION SYSTEMS
For all embedding conditioned approaches, we set de = and learned a rank-1 change to the dynamics for our approach. Our approach successfully learned both dynamical regimes present the datasets and the embedding distribution encoded differences in these dynamics with high certainty given limited time bins on test trials(Fig. ITA, B). While all approaches performed well on reconstructing observations these datasets, our approach and the Embedding- Input outperformed other multi- session baselines on forecasting (Fig (TC). We also evaluated the generalization performance of all methods the 2 held-out datasets function of training data used for training the read-in network
15

---

Published as a conference paper at ICLR 2025
and observed similar trends demonstrated by the reconstruction and k-step 20 r² on test trials from these datasets, shown in Table
A.
Example
C.
Performance on Training Tasks
Observations
True
0.75
15 0
0.5
60
B.
time
Lincir Adupler
Shared Dynamics
Single
1.5
Emnbesling Input
time
time
Figure 11: Example observations along with the reconstructed and predicted observations from approach for the fixed point (Top) and limit cycle (Bottom) dynamical regimes. B. Samples from the embedding as a function the number of time bins in the test trials for 4 example datasets. C. Reconstruction (Top) and forecasting (Bottom) performance of all approaches on the datasets used for pretraining.
Reconstruction n =
Forecasting ng =
12 = 8
1 = 8
Ours
0,85 + 0.054
0.89 + 0.04
0.64 + 0.1
0.69 + 0.07
Linear-Adapte
0.84 + 0.059
0.89 + 0.04
-0.1 + 0.34
0.55 + 0.08
Single Scssion
0.8 + 0.054
0.88 + 0.044
0.27 + 0.08
0.77 + 0.03
Shared Dynamics
0.83 + 0.068
0.89 + 0.04
0.32 + 0.08
0.32 + 0.04
Humbedding-Input
0.86 + 0.049
0.89 + 0.04
0.61 + 0.09
0.56 + 0.11
Table Few-shot reconstruction and forecasting performance (k-step=20) for held-out datasets in C where is the number of trials used for learning the dataset specific read-in network and likelihood.
D ALTERNATIVE INFERENCE AND LEARNING APPROACHES
While we focus on the DKF method for performing inference in the main text, we evaluate the efficacy of our proposed generative model with alternative inference and learning schemes.
D.I VARIATIONAL SMC
We considered the variational sequential monte carlo (VSMC) framework proposed by Naesseth et al. (2018) for learning and inference. VSMC is combination of variational inference and sequential Monte Carlo (SMC) (Doucet al. 2001), allowing for optimization of the parameters of the proposal; for simplicitiy, unlike traditional SMC, no resampling was performed after every time step, t. Given N samples from the encoder, i.e., 21. N VSMC optimizes the following lower bound the log marginal likelihood,
/ II p(za 2-1)p(11 ( 21)dZı Liomm t=1
1 N
logp(vir) log
Σ 1=1
Ey(z,)
log
==1
where,
(19)
p(yel 2 2))p(=
w As the proposed formulation requires inference of the dynamical embedding, e-which is constant over time have to modify VSMC non-trivial to infer constants in state-space models using
(20)
q (1.T)
16

---

Published as a conference paper at ICLR 2025
SMC (Doucet et al. 2001). We can express the log marginal likelihood of the proposed generative
model as
log p(vn.r) log
p(e) II 2 e)p( 1e)d1rdde,
(21)
log / p(e)de / T
(22)
/ p(e)p(U1: le)de,
log
(23)
/ ((e) p(e)
log
P le)de,
(24)
g(e)
p(e)
log By(e)
( I le)
(25)
g(e)
Applying Jensen's inequality to 225
log p(e)
log p(ph.)))) 2 Ey(())
g(e) p (e)
(26)
logr()) 2 By(()) p(yi( | + (gte) (log((e) Byte) [log ((e)).
(27)
As expectations respect inequalities, we can lower bound Eg [log p(y1.r le)-rr/)))))))))) using the VSMC lower- bound (9), leading to the following lower-bound that optimize
logp(yi:)) Egge) [Luomo]
Eggte [logp(e]] Eg(e) [log (()),, (28)
N
w((e) )
logr())r)) 2
log
1
Euele [log Ro(e) (e)), (29)
N
where
wi(e) = plye 21)p(2( 1.
(30)
{ V1.T.C)
A
B.
0.6
06
0.5
VSMC
Revonstruction
0.3
Prediction
0.6
0.6
0,3
dVBF
0.3
32
Truining Sumples
Training Samples
Centre-Cue
Mize Suk
Sub
Our Method
Single Sessinn
Figure 12: Samples from the learned embedding distribution using VSMC (Top) and dVBF (Bottom). B. Behavior decoding performance from reconstructed (Lefft and forecasted trajectories (Right) using the two inference methods with single session generative model vs aligning to our pretrained generative model.
D.2 DEEP VARIATIONAL BAYES FILTER
We additionally considered the DVBF framework proposed in (Karl et al., 2016) for performing learning and inference. This framework explicitly ties the inference network the generative model by forcing the samples from the inference network through the dynamical systems model. In our implementation, we defined the inference network as follows.
qa (ut W et) =N(u ps(concat ioooonattlyb..r +t)), of (coneat( e())))
17

---

Published as a conference paper at ICLR 2025
where 9s encoded the observations backward in time and parametrized by an RNN to infer the parameters of the Gaussian distribution, uy ~ ((u)). The latent trajectory for each dataset was subsequently sampled as,
= ( (2/_1) Q/?
Parameters of the generative model and inference networks were learned jointly by optimizing the following ELBO,
ΣΣ
Ego pogg =[]] - Eg (Dx.(g(u))lllllele - ("xx(ua(())lllll (31)
where p(u) ~ ~N(O.I).
E DATA GENERATION DETAILS
E. LIMIT CYCLE
For the experiments in Sections
3.3 we simulated data from the following system of equations,
r(1 -r)?
w.
21 = r cos 0 + 5dW, 23 = rin 0 + 5dW,
where w², e {1...., .M) } is the dataset specific velocity and dW, is the Wiener process. Specifically, for the experiment with M = datasets, we set = and " = 5; for the experiment with M = 20 datasets, we uniformly sampled 20 values for w between 0.25 and 5. For each value of w', we generated 128 latent trajectories for training, 64 for validation and 64 for testing, each of length 300, where we used Guler-Mayurama with At of 0.04 for integration. Observations were generated according to yi (C'' z4, Rt) where R = 0.01 * / and the elements of the readout matrix O were sampled from AO, d=); the dimensionality of the observations varied between 30 and 100. with R ~ (0, 0.01), where the dimensionality of the observations varied between 30 and 100.
E.2 HOPF BIFURCATION
Latent trajectories were generated from the following wo-dimensional dynamical system,
Z2 5 dW,
22 = -1 (n² =]) + 5dW,
(32)
where the parameter nı controls the topology of the dynamics. Specifically, when j² < 0, this system follows fixed point dynamics, and when > 0 the system bifurcates to a limit cycle.
We uniformly sampled M 21 values of 2² between -1.5 and 1.5, and for each n², we generated 128 trajectories for training and 64 for testing, each trajectory was T = 350. Observations were generated according to vt ~ NCC zi; R) where /t = 0.01 / and the elements of the readout matrix C were sampled from (0, T/ 1/vdd); the dimensionality of the observations varied between 30 and 100. For evaluating few-shot performance, we generated two additional novel datasets where j. -0.675 and p M++ 1.125 (both values were not included in the training set).
F.3 DUFFING SYSTEM
For the Duffing system described in 18 set 10 and varied values and b (shown in blue, Fig. for generating 20 datasets. We additionally used 11 datasets from obtained by uniformly sampling y between -1.5 and 1.5. For few shot evaluation various approaches, used two held-out datasets from the Duffing system (shown in orange, Fig. 13 as well as a dataset from the Hopf example generating by setting μ 1.8. All empirical evaluation was performed on 64 test trials from each dataset.
00.1
b
-1.5
-0.6
-0.2
a
Training Data Helld-Out Data
Figure 13: (a*, 8²) values used to generate different datasets from the Duffing system.
18

---

Published as a conference paper at ICLR 2025
E.4 MOTOR CORTEX RECORDINGS
We binned the spiking activity in 20ms bins and smoothed with 25ms causal Gaussian filter to obtain the rates for all datasets. We further removed neurons that had firing rate of less than 0.1 and aligned the neural activity to movement onset. We used 512 trials when available or 80 percent of the trials, each of length 36, for training all methods.
F ADDITIONAL FIGURES
A.1
Read-in networks
A.2
F0C
5
ga
Embedding Encoder
Neural Recordings
91
56
g/
A.3
96
Latent State Encoder
Recordings in common space
61
trials from reconding i, F15,1:7
inteneed eubeedding
Figure 14: Inference Overview. A.1. Each recording vi.r is projected into a common space Virr by recording specific read-in networks I2 R - 1R A.2. After being projected in this common space, the recording is processed by an RNN (qa) which infers the distribution over the dynamical embedding. Note that the dynamical embedding is aggregated across trials belonging to the same recording session. A.3. This inferred embedding is concatenated with Øı.r to obtain the latent state trajectories for each recording via the encoder q3, parametrized by bi-directional RNN.
A.
True Dynamics
B.
Learnt Dynamics
Learnt Dynamies (Embedding Input)
(Linear Adapter)
Figure 15: (Left) Dynamics learnt by Linear-Adapel and (Right) Embedding corresponding the example dynamics on the true system shown in Fig. 4A.
19

---

Published as a conference paper at ICLR 2025
A. True Dynamies
B. Our Approach
Embedding Input
Linear Adapter
Figure 16: (Left) True dynamics from example datasets used for pretraining in experiment 5.1
(Right)
Dynamics Learnt by different condiiiooo parametrizations.
Example dataset
Example dataset (Extrapolation)
(Interpolation)
Closest embedding on train tasks
Inferred embedding
16
Figure 17: Samples from the embedding distribution as function of the number of training trajecto- ries seeds held-out datasets from the Duffing system (denoted by orange, Fig. 13)
A.
Inferred embedding
B.
Training sessions
d
d
0.1
Reconstruction
0,1
o.
Masc
4,55
Mare
Prodiction
0.4
Rank d
Figure 18: A. Sample from the inferred embedding for different ranks (best of 3 seeds).B. The behavior reconstruction (Top) and prediction (Bottom) 2 for different ranks averaged over 3 training seeds.
20

---

Published as a conference paper at ICLR 2025
09
0.6
Held-out Sessions
Revonsttiction p²
066
Single Session
Shared Dynamies
Adapter
Prodiction
Embedding lupat
L.FDS Single Session
0.9
Held-out Subject
0.6
16
64
16
Training Samples
Training Sumples
Figure 19: Behavior reconstruction (Left) and forecasting (Right) for all methods function of the number of training samples for held-out sessions from Sub M and C, and two sessions from a held-out Subject.
A.
Inferred embedding
B.
Training sessions
C.
Few Shot Performace
Maze
Reconstructed
New subject
"2 0.4
0.1
4,
0.0
Reconstruztio:
Prediction
Prediction
Single Session
0.0
Costre-uu Sub
Marz Sub le Centre-Dul Maze Sub
16
64
Training Samples (a)
Figure 20: A. Samples from the embedding distribution after training our approach using 2 CO sessions from Sub C and M, and 2 Maze sessions from Sub Je and Ni. B. Reconstruction and forecasting performance of the model on held out test trials relative the Single Session models. C. Few-shot reconstruction and forecasting performance on held out sessions and a new subject (Sub T).
Predicted behavior
Predicted dynamics
Figure 21: (Left) Grid points used for generating the latent dynamical trajectories and with the inferred embedding distribution overlaid. The embedding of the CO session from Sub C used infer the initial condition of the latent state is highlighted in green. (Right) Single trials of the predicted hand position and PC projections of the corresponding latent dynamics trajectories. The initial latent state is denoted by the black points.
21

---

Published as a conference paper at ICLR 2025
Predicted behavior
Predicted dynamics
Figure 22: Same 21 but for Maze session from Sub Je (highlighted in blue).
22

---

Published as a conference paper at ICLR 2025
G EXPERIMENT DETAILS
G.1 TRAINING
Since CAVIA (Concat) (Zintgraf al. 2019), DYNAMO (Cotler et al. (023) and CoDA (Kirchmeyer et al. 2022 have not been developed for joint inference and learning of dynamics, use framework for inference with modifications to the parametrization of the embedding conditioned dynamics function all experiments. We used the official implementation of LFADS in PyTorch to obtain the results reported in the paper (Sedler & Pandarinath 2023)-
For our method and CoDA, restricted parameter changes to the input, Win, and hidden weights, W. We additionally included the Frohenius norm on the uunbedding-conditionn (e OI] along with the ELBO (eq. 14 for both approaches.
weight change
Synthetic Experiments. For the results reported in C and 5.1 we used the Adam optimizer with weight decay and Cosine annealing schedule on the learning rate for pretraining all approaches.
Motor Cortex Recordings. We used the LAMB optimizer for pretraining all multi-session methods on the motor cortex recordings and used Adam with weight decay for the single-session models, with a Cosine annealing schedule on the learning rate in both cases. We also incorporated masking during training for all approaches to encourage learning better latent dynamics. Specifically, we sampled the latent state from the dynamics instead of the state inference network on randomly masked time bins to compute the likelihood.
Aligning New Data. We pretrained all multi-session approaches for 3 seeds and picked the best performing model evaluate few-shot performance. When aligning new data to this pre-trained model, we trained the dataset-specific read-in network, Ω2 and likelihood functions p for the new dataset, in addition the state noise, Q' by optimizing the ELBO for new data (eq. 14). We used Adam with weight decay for aligning, and additionally incorporated masking when aligning held-out motor cortex datasets.
G.2 FIGURE GENERATION
Vector Field. We generated all the vector field plots for synthetic experiments by sampling random points 2-D grid to obtain 2z. We used the mean dynamics learned by the model to estimate the velocity at 2, as, - = (2+) zt. In the embedding-con conditiond methods, we additionally used the inference network, Ja to estimate the dynamical embedding corresponding to each dataset, which was used to conditionally generate the vector field plots.
We additionally align all learned vector field plots the true system for of visual comparison (Fig. A 15 16). We do do this by learning a linear transformation from the true latent trajectories to the latent trajectories learned by the model. Note that we follow the same procedure when visualizing latent trajectories inferred by multi-session CEBRA (Fig. TOAA
Context Interpolation. We fed the neural recordings up till movement onset time to the latent state encoder, qa. to obtain the latent state at movement onset, žt. We sampled points on a 2-D grid and simulated the corresponding samples from the embedding distribution as ~ Ne,,, 0.11). Given the latent state at movement onset for particular recording session, we were able to obtain different dynamical trajectories by giving these embedding samples to the latent dynamics model fe.e (2)) along with We used this procedure to obtain the results in in Fig. and 22
G.3 METRICS
Synthetic Experiments. We report the 72 observation reconstruction for test trials over the entire length of the trial for all approaches. In order to evaluate the forecasting performance, we observations till time f and sample the corresponding latent trajectories, = from the inference network and the corresponding for the embedding-condiee condiitoooo methods. We subsequently use the learned dynamics model generate K steps ahead from 2(-1+ and map these generated
23

---

Published as a conference paper at ICLR 2025
trajectories back to the observations. The k-step r² for each dataset is computed as,
Σ 2im1 (yi - 22)²
1 =
Ei (3/k D)²
where y the mean activity during the trial, and M is the number of test trials.
Motor Cortex Recordings. On the motor cortex experiments, we report behavior decoding from the reconstructed and forecasted observations for all methods. Specifically, for each session, we trained linear behavior decoder from the neural observations to the hand velocity of the subject, assuming uniform delay of 100ms between neural activity and behavior for all sessions.
After training all methods, we use reconstructed observations from the test trials to evaluate the behavior reconstruction r². For evaluating the decoding performance from forecasted observations. we used the first 13 time bins (around time till movement onset) to estimate the latent state and embedding, and subsequently use the trained dynamies model to forecast the next 20 time bins. The observations corresponding tp these forecasted trajectories were used to evaluate the prediction r2.
G.4 HYPERPARAMETERS
Synthetic Examples. We used the following architecture for pretraining all methods, with de = for the Motivating example and Hopf bifurcating system, and de for the combined Duffing and Hopf example.
Inference network
Q MLP 64,
qx: (GRU(16), Linear(16, 2 x de)]
4s: (biGRU(64)) Linear 4)]
Generative model
- So: M2, 32,
hg: MLP 16, 16, (64 + 33) x d,)
Par: [Linear(2, d,,))
Training
Ir: 0.005
weight decay: 0.001
batch size: 8 from each dataset
Motor Cortex Experiment.
Inference network
Ω MLP(dd, 128, Propout(O.6). 64)
qx: (64), Linear(64, (64, 2 de)]
qs: (12G) Linear(128, 60))
Generative model
fø: MLP(30, 128, 128,30 30)
he: MLP 64, 64, (256 + 158) d,)
Post [Linear(30 (30,, d,.)]
Training
Ir: 0.01
weight decay: 0.05
batch size: 64 trials from 20 datasets
24


Reviews:
[{"review_id": "jIXmdwhbc6", "reviewer": "Reviewer_eQP9", "summary": "The proposed meta-learning framework for analyzing neural dynamics across diverse datasets that have shared underlying structure. By encoding dataset-specific variations in a low-dimensional embedding space and conditioning a shared dynamical model on this embedding, it allows for the integration of heterogeneous neural recordings. However, it's essential to consider the model's limitations, such as its dependence on shared dynamics, semi-large-scale pre-training data, and its lack of systematic benchmarking. However, overall, the proposed meta-learning framework presents a promising new direction for neural data analysis.", "questions": "See weaknesses.", "limitations": "", "rating": 6, "confidence": 5, "soundness": 2, "presentation": 1, "contribution": 3, "strengths": "*   **Integration of Multi-Session Neural Recordings:** The proposed framework effectively addresses the challenges of integrating (potentially heterogeneous) neural recordings. It achieves this through a parameterization of latent dynamics using a low-dimensional dynamical embedding to capture variations across datasets. While they are not the first report to do so, it is an important direction/application.\n*   **Learning a Shared Dynamical Structure:** The model learns a shared dynamical structure across diverse datasets, which is adapted based on the specific characteristics of each dataset. \n*   **Rapid Adaptation to New Data:** The proposed method demonstrates good performance in few-shot learning scenarios, where it can rapidly adapt to new recordings with limited training data. This is particularly valuable in neuroscience research, where collecting large amounts of data can be time-consuming and expensive. The ability to learn from small datasets expands the applicability of the framework and makes it suitable for a wider range of experimental paradigms.\n*   **Generalizability:** The framework shows promising results in generalizing to new datasets, as evidenced by its performance on held-out data from both synthetic and motor cortex recordings. By learning a manifold of dynamical embeddings, the model is able to capture the variations in dynamics across different recordings and leverage this knowledge to adapt to new, unseen data.", "weaknesses": "**Missing Benchmarks & Issues in Related Works** \n\n- There are some inaccurate statements in the introduction; for example, \"CEBRA (Schneider et al., 2023) and CS-VAE (Yi et al., 2022) extract latent representations but do not learn underlying dynamics.\" Schneider et al. specifically learn the underlying dynamics rather than a priori prescribing them, as other algorithms do (e.g., SIDS, LFADS). CS-VAE assumes the same shared underlying dynamics and uses witching linear dynamical systems (SLDS) to discover them (and the authors should update their citation to: https://elifesciences.org/reviewed-preprints/88602. \n\n\n- Furthermore, this paper does not provide any benchmarking against models that they themselves acknowledge have the same goal: to create robust encoders across multiple datasets for downstream tasks (e.g., decoding). Therefore, they should report their results in comparison to Ye et al., 2023; Zhang et al., 2023; Caro et al., 2024; Azabou et al., 2024; & Schneider et al., 2023 (which they also fail to cite, but which showed joint training of an encoder rapidly adaptable to unseen data).\n\n\n **Dependence on Large-Scale Pretraining Data, and Only Within One Neural Domain:** \n- The effectiveness of the proposed framework depends on the availability of large-scale pretraining data. The sources show that models trained on smaller datasets tend to learn specialized solutions that may not generalize well to new recordings. Moreover, they only test on limited motor cortex settings, vs. more complex settings such as vision decoding.\n\n **Accurate Embedding Inference:** \n\n- The model's generalization ability hinges on accurate inference of the dynamical embedding for a new dataset. If the embedding is not accurately inferred, the model's performance on downstream tasks can be significantly compromised. Showing how corrupted/noisy data, and in other domains (see above) affects the utility of the algorithm will be important for adoption in neuroscience.\n\n**Assumptions of Shared Structure:** \n- The current framework assumes that there is a shared structure across the related tasks being analyzed. This assumption might not hold true in all cases, particularly when dealing with recordings from very different brain regions or unrelated behaviors. Future work may need to explore ways to incorporate modularity or other mechanisms to accommodate datasets without expected shared structures. The authors should comment on this.", "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": "9gZcwgbBjc", "reviewer": "Reviewer_zU5A", "summary": "This paper tackles the problem of analysis across different animals, tasks, and recordings by introducing an unsupervised and dataset-hierarchical model of neural time series data. Their model class seeks to characterize a unified latent space over dynamical structures, which the authors claim has advantages for both inference efficiency and interpretability.  The authors propose different simplifications or ablations to their model class as baselines, showcasing the usefulness of the full interaction between the components. The authors investigate a synthetic example to show recovery performance and provide intuition for the shortcomings of typical approaches. Finally, they analyze recordings in the motor cortex of monkeys performing reaching tasks, highlighting the common dynamics across recordings.", "questions": "- How does the read-in network $\\Omega$ handle different data dimensionalities?\n- A readout matrix $C^i \\sim \\mathcal{N}(0,0.01)$ is highly contractive. What effect does that have on the model, does it not induce degeneracy?\n- What are the color schemes in Figures 3A and 4A? How are these figures generated? \n- Why did the linear adapter and Embedding-Input methods get worse as $n_s$ increased in Table 1? Is this from a similar intuition as high M in Section 2?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The paper is well-motivated. The introduction frames the paper well by highlighting surrounding literature and results on shared structure across subjects/tasks in task-trained models. Then section 2 provides a compelling example of the limitations of a shared latent dynamics model, with dataset-specificity only in the emission likelihood. \n- The paper is generally well-written. \n- The synthetic results are strong and form a cohesive story with Section 2\n- The motor cortex recording analysis is extensive, analyzing two different tasks, and provides compelling results. In particular, assessing the variability in dynamics across different stimuli and further showing how this variability itself varies between tasks showcases the model's fit. Finally, interpolation in an embedding space (Figure 7) that encompasses both tasks showcases the model's use for interpretability generally.", "weaknesses": "While the authors do explore ablations and specific perspectives on their own model class, a comparison to other models is missing. There are two fronts to this comparison:\n- **Theoretical and modeling connections**, investigating how their SSM relates to known models. Can you cast known models as special cases of yours? For instance the shared dynamics motifs (Driscoll et al, 2024), how would they best fit within this framework? Alternatively, the embedding analysis of (Cotler et al., 2023), how would it fare on SOTA unsupervised models of neural activity compared to your analysis directly on the data?\n- **Empirical comparisons**, implementing the other models mentioned in Section 4 using the way they would typically handle the variability across datasets (such as inputs or indicator variables).\n\nMedium: \n- Figures 3 and 4 lack detail on how they were generated, specifically as to which parameters (/embedding values $e^i$) are being used. \n- Section 3.1. takes unnecessary turns in presenting the hierarchical model in my opinion. It starts with eq (3) with $\\theta^i = \\theta + \\epsilon_i$, and but then discards this model (due to dimensionality of $\\epsilon_i$) in favor of eq. (8) with $e^i \\sim p(e)$, only to come back in eq. (11) to a form $\\theta^i = \\theta + h(e^i)$. I believe a more streamlined exposition could help the reader, going from eq (3) to eq (8) directly in some fashion. \n\nMinor (/typos): \n- L131, \"we model the likelihood as a linear function of $z_t$\": the mean of the likelihood is an affine model (from the $+D$), but the likelihood function/distribution itself is not \"linear\". \n- Eq (5), $\\phi^i$ undefined. I gather from L205 that each $\\phi^i$ is learned independently  -- I would clarify this aspect. \n- L155: I am unsure about the use of the term \"inductive bias\". $\\theta$ and $\\epsilon_i$ capture shared and dataset-specific structure *by construction* -- I see the \"bias\" here as referring to how dynamical structures are similar to a mean structure with $\\theta$.\n- Eq. (14) and surrounding (and further, L347), $d_r$ and $r$ are used interchangeably", "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": "qBOT8a0hi6", "reviewer": "Reviewer_qSiJ", "summary": "The authors present a method for learning the dynamics of neural data across individuals and tasks. Due to the heterogeneity of neural data (e.g. different neurons are recorded across different days, and of course across different subjects) much of the work in inferring dynamics from neural data has focused on single-session models. This limits the ability to train large models with lots of data, and to directly compare dynamics across individuals/sessions. The authors propose a state space model approach that learns a single set of shared latent dynamics, as well as session-level embedding that modifies the shared dynamics. They demonstrate the efficacy of their method on simulated data, and then apply it to real-world neural recordings from monkeys.", "questions": "minor comments:\n- L390, \"Fig\" missing before 17?\n- would be helpful to show the true dynamics in fig 12", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The paper addresses a long-standing problem in the field of neural data analysis, and one that will only become more pronounced as data collection methods improve. The approach is novel and well-motivated, and provides a tool for not only nicely reconstructing neural dynamics but also providing interpretable model components that will be crucial for scientific understanding. Additionally, the paper is well-written.\n\nThe proof of concept presented in figs 2+3 is extremely helpful for introducting the problem, and much appreciated.\n\nThe analysis of the synthetic data in sec 5.1 is also well-designed and instructive.\n\nThe real neural data is a well-chosen dataset that balances real world heterogeneity and simplicity for exploring model performance and behavior.", "weaknesses": "While the authors generally do a great job keeping track of all the notation, I find myself confused about how different trials are handled. It seems that trial averages over, say, certain reach directions are being modeled, rather than individual trials (perhaps I missed this somewhere). But how, then, are the different conditions handled in the notation y_{1:T}^{1:M}? I became confused about this in section 3.2 regarding the read-in network and the dynamical embeddings. Is there a single dynamical embedding per session, or per session-condition? If the former, does this require concatenating data across all conditions? Please clarify. [related, in sec 5.2, are test _trials_ left out, or test _conditions_?]\n\nFig 5B is difficult to interpret; some more annotations or a clearer caption would be appreciated.\n\nThe authors cite large-scale training of models as a motivation for their work, but there are no results that directly demonstrate how adding more data improves performance - the only results here are single session models vs the \"all sessions\" model. It would be interesting to see some of the metrics computed as a function of the number of sessions used for training the model (and then, for example, testing few-shot performance on held-out sessions/subjects).", "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": "R0VV7iUpQS", "reviewer": "Reviewer_6MTE", "summary": "This work proposes a new model class for learning shared sets of latent neural dynamics across heterogeneous data sets. While a number of recent papers have demonstrated empiricially that such population dynamics may be shared across tasks or even species, differences in number of neurons, individual behaviors, etc. present an obstacle to learning a shared model from which inferences might be drawn. The solution in this work involves learning a dataset-specific embedding vector used to parameterize low-rank updates to the weight matrices of state space models. These embedding vectors then parameterize differences across sessions, tasks, and individuals.\n\nThe paper is clearly written, and the problem is a timely one for the field. The experiments are thoughtfully chosen and give good insights, though some aspects (data type, model comparisons) are a bit limited. In all, a good paper that's above the bar for acceptance in my view.", "questions": "- There are some experiments to demonstrate that the parameterized low-rank embeddings outperform linear adapters or using embeddings as inputs, and these are convincing, but the intuition supplied for this is pretty terse. Why should one expect interference in these cases? Why does the rank-1 update protect against this?\n- Figure 1: It would be helpful if panel B (or some other panel) could include $\\Omega$ and the inference networks proposed by the model. For generation, this seems fine, but that's only capturing half the model.\n- Is there a ground truth for Figure 2? It's hard to judge what one _should_ be looking for here.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- Methods for combining insights across datasets are a huge priority in neuroscience. \n- The model builds on several well-understood components and proposes a sensible ansatz for how embeddings alter to underlying system dynamics. \n- Multiple experiments carefully examine several key variables, including data set size, look-ahead prediction, and number of training trials. These are all clearly presented.", "weaknesses": "- The model makes a number of design choices that are unclearly or incompletely motivated. For instance: does the dynamical VAE matter? Should one use the Bowman et al. inference scheme or the SMC-based autoencoder models of [1, 2, 3]? What about the bidirectional RNNs for $\\mu_\\alpha$ and $\\sigma^2_\\alpha$? One can't ablate everything, but it would be good to know how much conclusions depend on these choices.\n- Related but distinct: The experimental comparisons are done for alternative approaches to learning the embeddings, but there aren't really comparisons with other models. At least one of the metrics tabulated is $k$-step ahead $R^2$, which could surely be compared with LFADS on single-session data or POYO [4] on multi-session data. I recognize that this paper has goals in addition to pure prediction, but since the authors do consider that metric, it would be good to see how the model proposed here compares to other predictive models. Likewise, it would be good to know how to assess (other than qualitatively) what constitutes a \"good\" learned latent space in the absence of ground truth.\n\n[1] Variational Sequential Monte Carlo\nChristian Naesseth, Scott Linderman, Rajesh Ranganath, David Blei Proceedings of the Twenty-First International Conference on Artificial Intelligence and Statistics, PMLR 84:968-977, 2018.\n\n[2] C. J. Maddison, D. Lawson, G. Tucker, N. Heess, M. Norouzi, A. Mnih, A. Doucet, and Y. Whye Teh. Filtering variational objectives. In Advances in Neural Information Processing Systems, 2017.\n\n[3] T. A. Le, M. Igl, T. Jin, T. Rainforth, and F. Wood. Auto- Encoding Sequential Monte Carlo. arXiv:1705.10306, May 2017.\n\n[4] A Unified, Scalable Framework for Neural Population Decoding\nMehdi Azabou, Vinam Arora, Venkataramana Ganesh, Ximeng Mao, Santosh Nachimuthu, Michael Mendelson, Blake Richards, Matthew Perich, Guillaume Lajoie, Eva Dyer", "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 (Spotlight)