Meta-Dynamical State Space Models for Integrative Neural Data Analysis

## Abstract

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.

## 1 Introduction

Latent variable models are widely used in neuroscience to extract dynamical structure underlying high-dimensional neural activity (Pandarinath et al., ; Schimel et al., ; Dowling et al., ).
While latent dynamics provide valuable insights into behavior and generate testable hypotheses of neural computation (Luo et al., ; Nair et al., ), they are typically inferred from a single recording session.
As a result, these models are sensitive to small variations in 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 on inferring generalizable dynamical systems from data, with existing approaches mainly applied to settings with known low-dimensional dynamics (Yin et al., ; Kirchmeyer et al., ).

Integrating noisy neural recordings from different animals and/or tasks for learning the underlying dynamics presents a unique set of challenges.
This is partly due to heterogeneities in 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 is 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 et al., ; Maheswaranathan et al., ; Safaie et al., ).
In a 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, ; Driscoll et al., ).

Motivated by these observations, we propose a novel framework for meta-learning latent dynamics from neural recordings.
Our approach is to encode the variations in the latent dynamical structure present across neural recordings in a low-dimensional vector,
$e\in\mathbb{R}^{d_{e}}$,
which we refer to as the dynamical embedding.
During training, the model learns to adapt a common latent dynamical system model conditioned on the dynamical embedding.
We learn the dynamical embedding manifold from a 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 a 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 as 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.

Figure: Figure 1: 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. B. Our method learns to adapt a common latent dynamics conditioned on the embedding via low-rank changes to the model parameters.
Refer to caption: x1.png

## 2 Challenges with Jointly Learning Dynamics across Datasets

Figure: Figure 2: A. Three different example neural recordings, where the speed of the latent dynamics varies across them. B. One generative model is trained on $M=2$ and $M=20$ datasets. While increasing the number of datasets allows the model to learn limit cycle, it is unable to capture the different speeds leading to poor forecasting performance.
Refer to caption: x2.png

Neurons from different sessions and/or subjects are partially observed, non-overlapping and exhibit diverse response properties.
Even chronic recordings from a single subject exhibit drift in neural tuning over time (Driscoll et al., ).
Moreover, non-simultaneously recorded neural activity lack pairwise correspondence between single trials.
This makes joint inference of latent states and learning the corresponding latent dynamics by integrating different recordings ill-posed and highly non-trivial.

As an illustrative example, let’s consider a case where these recordings exhibit oscillatory latent dynamics with variable velocities (Fig. 2A).
One possible strategy for jointly inferring the dynamics from these recordings is learning a shared dynamics model, along with dataset-specific likelihood functions that map these dynamics to individual recordings.
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 a model with shared dynamics either learned separate solutions or overfit to one dataset, obscuring global structure across recordings (Fig. 2A).
When we increased the diversity of training data $(M=20)$, the dynamics exhibited a more coherent global structure, albeit with an overlapping solution space (Fig. 2B).
As a result, this model had poor forecasting performance of neural activity in both cases, which is evident in the k-step $r^{2}$ (Fig. 2B). While we have a-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 it to define a solution space of related dynamical systems (Fig. 1A).

## 3 Integrating Neural Recordings for Meta-Learning Dynamics

Let $y_{1:T}^{1:M}$ denote neural time series datasets of length $T$, with $y_{t}^{i}\in\mathbb{R}^{d_{y}^{i}}$, collected from $M$ different sessions and/or subjects during related tasks.
We are interested in learning a generative model that can jointly describe the evolution of the latent states across these datasets and rapidly adapt to novel datasets from limited trajectories.
In this work, we focus on state-space models (SSM), a powerful class of generative models for spatio-temporal datasets.
An SSM is described via the following pair of equations (we drop the superscript for ease of presentation),

$$ $\displaystyle z_{t}\mathchoice{\!\mid\!}{\!\mid\!}{\mid}{\mid}z_{t-1}$ $\displaystyle\sim p_{\theta}(z_{t}\mathchoice{\!\mid\!}{\!\mid\!}{\mid}{\mid}z _{t-1}),$ (1) $\displaystyle y_{t}\mathchoice{\!\mid\!}{\!\mid\!}{\mid}{\mid}z_{t}$ $\displaystyle\sim p_{\phi}(y_{t}\mathchoice{\!\mid\!}{\!\mid\!}{\mid}{\mid}z_{ t}),$ (2) $$

where $z_{t}\in\mathbb{R}^{d_{z}}$ is the latent state at time $t$, $p_{\theta}(z_{t}\mathchoice{\!\mid\!}{\!\mid\!}{\mid}{\mid}z_{t-1})$ is the dynamics model
and $p_{\phi}(y_{t}\mathchoice{\!\mid\!}{\!\mid\!}{\mid}{\mid}z_{t})$ is the likelihood function that maps the latent state to observed data.

Following standard practice, we parameterize the dynamics as $p_{\theta}(z_{t}\mathchoice{\!\mid\!}{\!\mid\!}{\mid}{\mid}z_{t-1})=\mathcal{N
}(z_{t}\mathchoice{\!\mid\!}{\!\mid\!}{\mid}{\mid}f_{\theta}(z_{t-1}),Q)$, where $f_{\theta}$ is a deep neural network (DNN) and $Q$ is a covariance matrix(^1^11We note that $Q$ can also be parameterized via a neural network as well.).
As previous work has shown that making both the likelihood and dynamics highly expressive can lead to optimization issues (Bowman et al., ), we model the likelihood as a linear function of $z_{t}$.
For instance; for real-valued observations, it is defined as $p_{\phi}(y_{t}\mathchoice{\!\mid\!}{\!\mid\!}{\mid}{\mid}z_{t})=\mathcal{N}(y_
{t}\mathchoice{\!\mid\!}{\!\mid\!}{\mid}{\mid}Cz_{t}+D,R)$.

### 3.1 Hierarchical State-Space Model for Multiple Datasets

To learn an SSM that can jointly describe the spatio-temporal evolution across $M$ neural recordings, we introduce
a hierarchical structure in the latent dynamics model that captures the variations in dynamics across datasets.
A natural choice for learning such a generative model is
a fully Bayesian approach, where each dataset would have its own latent dynamics, parameterized by $\theta^{i}$, and a hierarchical prior would be used to tie the dataset specific parameters to the shared parameters, $\theta\sim p(\theta)$ (Linderman et al., ),
leading to the following generative model

$$ $\displaystyle\theta^{i}\mathchoice{\!\mid\!}{\!\mid\!}{\mid}{\mid}\theta$ $\displaystyle\sim p(\theta^{i}\mathchoice{\!\mid\!}{\!\mid\!}{\mid}{\mid} \theta),$ (3) $\displaystyle z_{t}^{i}\mathchoice{\!\mid\!}{\!\mid\!}{\mid}{\mid}z_{t-1}^{i}, \theta^{i}$ $\displaystyle\sim\mathcal{N}(z_{t}^{i}\mathchoice{\!\mid\!}{\!\mid\!}{\mid}{ \mid}f_{\theta^{i}}(z_{t-1}^{i}),Q^{i}),$ (4) $\displaystyle y_{t}^{i}\mathchoice{\!\mid\!}{\!\mid\!}{\mid}{\mid}z_{t}^{i}$ $\displaystyle\sim p_{\phi^{i}}(y_{t}^{i}\mathchoice{\!\mid\!}{\!\mid\!}{\mid}{ \mid}z_{t}^{i}),$ (5) $$

where dataset specific likelihoods, $p_{\phi^{i}}(y_{t}^{i}\mathchoice{\!\mid\!}{\!\mid\!}{\mid}{\mid}z_{t}^{i})$, are used to account for different dimensionality and/or modality of recordings.

If we assume that $p(\theta^{i}\mathchoice{\!\mid\!}{\!\mid\!}{\mid}{\mid}\theta)$ is Gaussian, i.e., $p(\theta^{i}\mathchoice{\!\mid\!}{\!\mid\!}{\mid}{\mid}\theta)=\mathcal{N}(
\theta^{i}\mathchoice{\!\mid\!}{\!\mid\!}{\mid}{\mid}\theta,\Sigma)$, we can equivalently express dynamics for the hierarchical generative model as,

$$ $\displaystyle\varepsilon^{i}$ $\displaystyle\sim\mathcal{N}(\varepsilon^{i}\mathchoice{\!\mid\!}{\!\mid\!}{ \mid}{\mid}0,\Sigma),$ (6) $\displaystyle z_{t}^{i}\mathchoice{\!\mid\!}{\!\mid\!}{\mid}{\mid}z_{t-1}^{i}, \theta,\varepsilon^{i}$ $\displaystyle\sim\mathcal{N}\left(z_{t}^{i}\mathchoice{\!\mid\!}{\!\mid\!}{ \mid}{\mid}f_{\theta+\varepsilon^{i}}(z_{t-1}^{i}),Q^{i}\right),$ (7) $$

where the dataset-specific dynamics parameter, $\theta^{i}$, is expressed as a sum of the shared parameters, $\theta$, and a dataset-specific term, $\varepsilon^{i}$. This decomposition of latent dynamics explicitly introduces an inductive bias that encourages $\theta$ and $\varepsilon^{i}$ to capture the shared and dataset-specific dynamical structure respectively.
While this formulation is appealing, the equal dimensionality of $\varepsilon^{i}$ and $\theta$ could lead the model to overly rely on dataset-specific terms. Additionally, using DNNs for modeling latent dynamics introduces a large number of parameters, thereby limiting the scalability of this approach.

Drawing inspiration from the hierarchical Bayesian model, we develop an alternative hierarchical framework that both improves scalability and more effectively incorporates the desired inductive bias.
Specifically, we introduce a low-dimensional latent variable, $e^{i}\in\mathbb{R}^{d_{e}},\mathbb{R}^{d_{e}}\ll\mathbb{R}^{d_{\varepsilon}}$—which we refer to as the dynamical embedding—that maps to the space of dynamics parameters and encodes dataset-specific variations (Rusu et al., ).
This approach enhances scalability and enables efficient few-shot learning, as it requires simply inferring the embedding given new data.

The generative model for this hierarchical SSM is then described as,

$$ $\displaystyle e^{i}$ $\displaystyle\sim p(e),$ (8) $\displaystyle z^{i}_{t}\mathchoice{\!\mid\!}{\!\mid\!}{\mid}{\mid}z^{i}_{t-1}, e^{i}$ $\displaystyle\sim\mathcal{N}\left(z_{t}^{i}\mathchoice{\!\mid\!}{\!\mid\!}{ \mid}{\mid}f_{\theta}(z_{t-1}^{i},e^{i}),Q^{i}\right),$ (9) $\displaystyle y^{i}_{t}\mathchoice{\!\mid\!}{\!\mid\!}{\mid}{\mid}z^{i}_{t}$ $\displaystyle\sim p_{\phi^{i}}(y^{i}_{t}\mathchoice{\!\mid\!}{\!\mid\!}{\mid}{ \mid}z^{i}_{t}),$ (10) $$

where we drop the prior over the shared dynamics parameter, $\theta$, significantly reducing the dimensionality of the inference problem.
Similar to the hierarchical Bayesian model, all datasets share the same latent dynamics, $\theta$, with the dataset-specific variation captured by the dynamical embedding, $e_{i}$.

Next, we expand on the functional form of the latent dynamics $f_{\theta}$ conditioned on $e^{i}$.
Our goal is to design $f_{\theta}$ in a way that can effectively capture complex changes and disentangle underlying variability in dynamics across recordings, while simultaneously incentivizing the extraction of shared dynamical structure.
While the simplest formulation would be using the embedding as an input to the dynamics model (Zintgraf et al., ; Cotler et al., ), we observed that this led to interference between dynamics as we increased the diversity of dynamical regimes (for example, Fig. 13).
Instead, we use the dynamical embeddings to directly output changes to the parameters of the dynamics model via a hypernetwork, $h_{\vartheta}$, (Ha et al., ), which facilitates increased expressivity when modeling conditional functions (Galanti & Wolf, ),

$$ $\displaystyle\theta^{i}=\theta+h_{\vartheta}(e^{i}),$ (11) $\displaystyle z_{t}^{i}\mathchoice{\!\mid\!}{\!\mid\!}{\mid}{\mid}z_{t-1}^{i}, e^{i}\sim\mathcal{N}(z_{t}^{i}\mathchoice{\!\mid\!}{\!\mid\!}{\mid}{\mid}f_{ \theta^{i}}(z_{t-1}^{i}),Q^{i}).$ (12) $$

We encourage learning of shared dynamical structure by constraining $h_{\vartheta}$ to make low-rank changes to the parameters of $f_{\theta}$ (Fig. 1B).
For example, if we parameterize $f_{\theta}$ as a 2-layer fully-connected network and constrain the hypernetwork to only make rank $d^{r}=r$ changes to the hidden weights, then $f_{\theta^{i}}$
would be expressed as,

$$ $\displaystyle f_{\theta^{i}}(z_{t}^{i})$ $\displaystyle=\,\,\,\vW_{o}\,\,\,\sigma\bigl{(}\{\vW_{hh}\,\,\,+\,\,\, \underbrace{h_{\vartheta}(e^{i})}_{\mathclap{\text{embedding modification}}}\} \,\sigma(\vW_{\text{in}}\,\,\,z^{i}_{t})\bigr{)}$ (13) $\displaystyle=\underbrace{\vW_{o}}_{\mathclap{\mathbb{R}^{d_{z}\times d_{2}}}} \sigma\bigl{(}\{\underbrace{\vW_{hh}}_{\mathclap{\mathbb{R}^{d_{2}\times d_{1} }}}+\underbrace{\vu_{\vartheta}(e^{i})}_{\mathclap{\mathbb{R}^{d_{2}\times d_{ r}}}}\cdot\underbrace{\vv_{\vartheta}(e^{i})^{\top}}_{\mathclap{\mathbb{R}^{d_ {r}\times d_{1}}}}\}\,\sigma(\underbrace{\vW_{\text{in}}}_{\mathclap{\mathbb{R }^{d_{1}\times d_{z}}}}z^{i}_{t})\bigr{)}$ (14) $$

where $\sigma(\cdot)$ denotes a point-nonlinearity,
and the two functions
$\vv_{\vartheta}(e^{i}):\mathbb{R}^{d}_{e}\to\mathbb{R}^{d_{1}\times r}$, $\vu_{\vartheta}(e^{i}):\mathbb{R}^{d}_{e}\to\mathbb{R}^{d_{2}\times r}$
map the embedding representation to form the low-rank perturbations, and both $\vu_{\vartheta}$ and $\vv_{\vartheta}$ are parameterized by a neural network.

### 3.2 Inference and Learning

Given $y_{1:T}^{1:M}$, we want to infer both the latent states, $z_{1:T}^{1:M}$ and the dynamical embeddings, $e^{1:M}=[e^{1},\ldots,e^{M}]$ as well as learn the parameters of the generative model,
$\Theta=\{\theta,\vartheta,\phi^{1},\ldots,\phi^{M}\}$. Exact inference and learning requires computing the posterior, $p_{\Theta}(z_{1:T}^{1:M},e^{1:M}\mathchoice{\!\mid\!}{\!\mid\!}{\mid}{\mid}y_{
1:T}^{1:M})$, and log marginal likelihood, $\log p_{\Theta}(y_{1:T}^{1:M})$, which are both intractable. To circumvent this issue, we use a sequential variational autoencoder (seqVAE) (Krishnan et al., )—an extension of variational autoencoders for state-space models—to learn the generative model by optimizing a lower-bound to the log marginal likelihood (commonly referred to as the ELBO).
The ELBO for $y^{1:M}_{1:T}$ is defined as,

$$ $\begin{split}\mathcal{L}(y_{1:T}^{1:M})&=\sum_{t,m}\mathbb{E}_{q_{\alpha,\beta }}\left[\log p_{\phi^{i}}(y^{i}_{t}\mathchoice{\!\mid\!}{\!\mid\!}{\mid}{\mid} z^{i}_{t})\right]\\ &\quad-\mathbb{E}_{q_{\beta}}\left[\mathbb{D}_{KL}(q_{\beta}(z^{i}_{t} \mathchoice{\!\mid\!}{\!\mid\!}{\mid}{\mid}\bar{y}^{i}_{1:T},e^{i})||p_{\theta ,\vartheta}(z^{i}_{t}\mathchoice{\!\mid\!}{\!\mid\!}{\mid}{\mid}z^{i}_{t-1},e^ {i}))\right]-\mathbb{D}_{KL}(q_{\alpha}(e^{i}\mathchoice{\!\mid\!}{\!\mid\!}{ \mid}{\mid}\bar{y}_{1:T}^{i})||p(e^{i}))\end{split}$ (15) $$

where $q_{\alpha}(e^{i}\mathchoice{\!\mid\!}{\!\mid\!}{\mid}{\mid}\bar{y}_{1:T}^{i})$ and $q_{\beta}(z_{t}\mathchoice{\!\mid\!}{\!\mid\!}{\mid}{\mid}\bar{y}_{1:T}^{i},e^
{i})$ are the encoders for the dynamical embedding and latent state, respectively, and $\mathbb{E}_{q_{\alpha,\beta}}\equiv\mathbb{E}_{q_{\beta}(z_{t}\mathchoice{\!
\mid\!}{\!\mid\!}{\mid}{\mid}\bar{y}_{1:T}^{i},e^{i})q_{\alpha}(e^{i}
\mathchoice{\!\mid\!}{\!\mid\!}{\mid}{\mid}\bar{y}_{1:T}^{i})}$. We additionally train a per-dataset read-in network $\Omega_{i}:\mathbb{R}^{d_{y}^{i}}\to\mathbb{R}^{d_{\bar{y}}}$ that maps $y^{i}_{t}$
to a vector $\bar{y}^{i}_{t}\in\mathbb{R}^{d_{\bar{y}}}$. As described in Sec. 2, one of the challenges with integrating recordings in a common latent space is different dimensionalities (number of recorded neurons) and
the dependence of neural activity on the shared latent space. The read-in network ensures that the latent states and dynamical-embeddings inferred from each dataset live in the same space.

We use the following parametrization of the inference networks (Krishnan et al., ),

$$ $\displaystyle\bar{y}_{t}^{i}=\Omega^{i}(y_{t}^{i}),$ (16) $\displaystyle q_{\alpha}(e^{i}\mathchoice{\!\mid\!}{\!\mid\!}{\mid}{\mid}\bar{ y}_{1:T}^{i})=\mathcal{N}(e^{i}\mathchoice{\!\mid\!}{\!\mid\!}{\mid}{\mid}\mu_ {\alpha}(\bar{y}_{t}^{i}),\sigma^{2}_{\alpha}(\bar{y}_{t}^{i})),$ (17) $\displaystyle q_{\beta}(z_{1:T}^{i}\mathchoice{\!\mid\!}{\!\mid\!}{\mid}{\mid} \bar{y}_{1:T}^{i},e^{i})=\prod_{t=1}^{T}\mathcal{N}(z_{t}^{i}\mathchoice{\! \mid\!}{\!\mid\!}{\mid}{\mid}\mu_{\beta}(\textrm{concat}(\bar{y}_{t}^{i},e^{i} )),\sigma^{2}_{\beta}(\textrm{concat}(\bar{y}_{t}^{i},e^{i}))),$ (18) $$

where concat is the concatenation operation. We parameterize $\mu_{\alpha}$, $\sigma^{2}_{\alpha}$
by a bidirectional recurrent neural network, and $\mu_{\beta}$, $\sigma^{2}_{\beta}$
by a regular recurrent network.
We emphasize that $\mu_{\alpha}$, $\sigma^{2}_{\alpha}$, $\mu_{\beta}$, and $\sigma^{2}_{\beta}$ are shared across all datasets.

### 3.3 Proof of Concept

Figure: Figure 3: A. Example dynamics learned by the proposed approach when trained with $M=20$ datasets. B. Samples from the inferred dynamical embedding for each dataset (see eq. 17). C. Forecasting $r^{2}$ at k-step=50 for models trained with $M=2$ and $M=20$ datasets.
Refer to caption: x3.png

As a proof of concept,
we revisit the motivating example presented in Section 2 as a means to validate the efficacy of our approach and investigate how it unifies dynamics across datasets.
For both $M=2$ and $M=20$ datasets, we used an embedding dimensionality of 1 and allowed the network to make a
rank-1 change to the dynamics parameters.

After training, we observed that the shared dynamics (when $e=0$) converged to a limit cycle with a slow velocity (Fig. 3A)—capturing the global topology that is shared across all datasets—and the model learned to modulate the velocity of the dynamics conditioned on the dynamical embedding which strongly correlated with the dataset specific velocity(^2^22Note that we plot the absolute embedding samples since the likelihood function can introduce arbitrary invariance such as direction flipping, rotation, and so on.) (Fig 3B).
This demonstrated that the proposed approach is able to capture dataset-specific variability.
Lastly, Fig. 3C 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 as the generalization of the trained model to new data in Appendix A.

## 4 Related Works

Multi-Dataset Training in Neuroscience. Previous work has explored multi-dataset training for extracting latent representations in neuroscience. LFADS (Pandarinath et al., ) used session-stitching, but focused on single-animal recordings. Linderman et al. () used hierarchical Bayesian state-space models with switching linear dynamical systems, while Herrero-Vidal et al. () developed a joint model with shared linear dynamics but dataset-specific likelihoods. In contrast to these approaches, we incorporate a more expressive dynamics function that disentangles variabilities across recording sessions. CEBRA (Schneider et al., ) and CS-VAE (Yi et al., ) extract latent representations but do not learn underlying dynamics. Recently, there has been growing interest in using diverse neural recordings for training large-scale foundation models in neuroscience (Ye et al., ; Zhang et al., ; Caro et al., ; Azabou et al., ). Our approach shares the same broad goal of pretraining a single model for rapid adaptation on downstream recordings; however, we focus on learning a dynamics model across recordings while these methods leverage transformer-based architectures that encode temporal information indirectly via positional embeddings.

Meta-Learning.
Several meta-learning approaches have been developed for few-shot adaptation, including gradient-based methods (Finn et al., ; Li et al., ; Nichol & Schulman, ; Zintgraf et al., ). Amongst these, LEO (Rusu et al., ) shares the same idea of meta-learning in low-dimensional space of parameter embeddings. However, gradient-based approaches require fine-tuning during test-time, and have had limited success for meta-learning dynamics.
Similar to our work (Cotler et al., ) also learns an embedding space of dynamics learned from trained RNNs, however, we are interested in learning dynamics directly from data. Some methods for learning generalizable dynamics been previously proposed—DyAD (Wang et al., ) adapts across environments by neural style transfer, however it operates on images of dynamical systems, LEADS (Yin et al., ) learns a constrained dynamics function that is directly added to some base dynamics function, and CoDA (Kirchmeyer et al., ) which learns task-specific parameter changes conditioned on a 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.

## 5 Experiments

We first validate the proposed method on synthetic data and then test our method on neural recordings from the primary motor and premotor cortex. We compare the proposed approach against the following baselines:

Single Session Model.
We train a separate model on each dataset.
Given sufficient training data, this should result in the best performance, but will fail in trial-limited regimes.

Multi-Session Shared Dynamics.
We also consider a model that assumes the same dynamics with dataset-specific likelihoods (Pandarinath et al., ; Herrero-Vidal et al., ).

Embedding-Input. In this baseline, the embedding is provided as an additional input to the dynamics model. This is a similar formulation to CAVIA (Concat) (Zintgraf et al., ) and DYNAMO (Cotler et al., ).

Linear Adapter.
We additionally test the hypernetwork parametrization proposed in CoDA (Kirchmeyer et al., ), where the hypernetwork adapts all parameters as a linear function of the dynamical embedding.

For each experiment, we split each of the $M$ datasets into a training and test set and report reconstruction and forecasting metrics on the test set.
To measure the generalization performance, we also report these metrics on held-out datasets.
Further details on training and evaluation metrics can be found in Appendix D.

### 5.1 Bifurcating Systems

In these experiments, we test whether our method could capture variations across multiple datasets, particularly in the presence of significant dynamical shifts, such as bifurcations commonly observed in real neural populations.
To test this, we chose two parametric classes of dynamical systems: i) a system undergoing a Hopf bifurication and, ii) the unforced Duffing system. We include the results for training on datasets generated only from the Hopf system in Appendix B.1 and discuss the results of jointly training on both systems here. We briefly outline the data generation process for the Duffing system (with data generation for Hopf included in Appendix B.1).

The latent trajectories for the Duffing system were generated from a family of stochastic differential equations,

$$ $\dot{z_{1}}=z_{2}+5\,\mathrm{d}W_{t},\hskip 30.0pt\dot{z_{2}}=a^{i}z_{2}-z_{1} (b^{i}+cz_{1}^{2})+5\,\mathrm{d}W_{t}$ (19) $$

with $c=0.1$, $a,b\in\mathbb{R}$, and $\mathrm{d}W_{t}$ denoting the Wiener process.
In Fig. 4A, we visualize how the dynamical system changes
as $a$ and $b$ vary.

Figure: Figure 4: A. (Left) True underlying dynamics from some example datasets used for pretraining as a function of parameters $a$ and $b$ and (Right) the embedding conditioned dynamics learnt by our model. B., C. Mean reconstruction and forecasting $r^{2}$ of the observations for all datasets used for pretraining on test trials.
Refer to caption: x4.png

We chose $M=20$ pairs of $(a^{i},b^{i})$ values (Fig 11), and generated latent trajectories of length $T=300$.
Observations were generated according to $y_{t}^{i}\sim\mathcal{N}(C^{i}z_{t}^{i},0.01\mathbb{I})$;
the dimensionality of the observations varied between 30 and 100.
In addition to these 20 datasets, we included 11 datasets from the Hopf system (Appendix B.2), and used 128 trajectories from each of these 31 datasets for training all methods. We report performance on 64 test trajectories from each dataset.
We used $d_{e}=2$ for all embedding-conditioned approaches and constrained the hypernetwork to make rank $d^{r}=1$ changes for our approach.

Our approach was able to learn a good approximation to the ground-truth dynamics of the Duffing oscillator system while disentangling different dynamical regimes (Fig. 4 A).
Apart from learning the underlying topology of dynamics, it also better captured the geometrical properties compared to other embedding-conditioned baselines (Fig. 12).
We observed similar results for datasets from the Hopf system–while our approach approximated the ground-truth system well, the Embedding-Input baseline displayed interference between dynamics and the Linear-Adapter learned a poor approximation to the ground-truth system (Fig. 13).
Consequently, our approach outperformed other methods on forecasting observations with all methods having comparable reconstruction performance (Fig. 4B, C). Notably, apart from the $d^{e}$, we used the same architecture as when training on only the Hopf datasets, and did not observe any drop in performance for our approach, in contrast to baselines (Fig. 10C (Bottom), Fig. 4C).

**Table 1: Few shot forecasting performance ($k=30$-step) on 3 held-out datasets as a function of $n_{s}$, the number of trials used to learn dataset specific read-in network and likelihood. ($\pm$ 1 s.e.m)**
|  | $n_{s}=1$ | $n_{s}=8$ | $n_{s}=16$ |
| --- | --- | --- | --- |
| Our Method | 0.69 $\pm$ 0.072 | 0.78 $\pm$ 0.051 | 0.87 $\pm$ 0.037 |
| Linear-Adapter | 0.68 $\pm$ 0.08 | 0.79 $\pm$ 0.026 | 0.74 $\pm$ 0.039 |
| Single Session | 0.47 $\pm$ 0.119 | 0.79 $\pm$ 0.014 | 0.79 $\pm$ 0.047 |
| Shared Dynamics | -0.31 $\pm$ 0.103 | -0.34 $\pm$ 0.086 | -0.13 $\pm$ 0.065 |
| Embedding-Input | 0.59 $\pm$ 0.084 | 0.77 $\pm$ 0.04 | 0.74 $\pm$ 0.039 |

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, as a function of $n_{s}$, the number of trials used for learning the dataset specific read-in network, $\Omega^{i}$ and likelihood.
Our approach and the Linear-Adapter demonstrated comparable forecasting performance when using $n_{s}=1$ and $n_{s}=8$ training trajectories.
However, with $n_{s}=16$ training trials, unlike other methods, our approach continued to improved and outperformed them (Table 1).
This could be explained by looking at the inferred embedding on held-out datasets—as we increased the number of training trajectories, the model was able to consistently align to the “correct” embedding (Fig. 14).

### 5.2 Motor Cortex Recordings

Figure: Figure 5: Visualizing the embedding manifold. A. Each point corresponds to a sample from the inferred embedding distribution (see eq. 17) corresponding to each recording. B. The condition-averaged of latent dynamics projections generated by the model, along with the corresponding real and forecasted behavior.
Refer to caption: x5.png

We move on to testing the applicability of the proposed approach on neural data.
We use
single and multi-unit neural population recordings from the motor and premotor cortex during two behavioral tasks–the Centre-Out (CO) and Maze reaching tasks (Perich et al., ; Gallego et al., ; Churchland et al., ).
In the CO task, subjects are trained to use a manipulandum to reach one of eight target locations on a screen. In the Maze task, subjects use a touch screen to reach a 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 a good testbed for evaluating various methods.
For training, we used $40$ sessions from the CO task, from subjects M and C, and $4$ sessions from the Maze task from subjects Je and Ni. We set the dimensionality of latent dynamics to $d_{z}=30$, and used an embedding dimensionality of $d_{e}=2$, for all embedding-conditioned dynamics models.
For our approach, we constrain the hypernetwork to make rank $d^{r}=6$ changes, although we verified that the performance was not sensitive to $d^{r}$ (Fig 15).
As a 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 D).

The inferred dynamical embedding displayed distinct structures across behavioral tasks and subjects (Fig. 5A).
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 of learned dynamics reflected this heterogeneity across recordings, as visualized by the condition-averaged principal components (PCs) of the latent trajectories generated by the model (Fig. 5B).

Many of the approaches had adequate performance
on reconstructing velocity from neural recordings, with our method and Linear Adapter outperforming single session reconstruction performance on the CO task (Fig. 6A, top). However, the single-session models had the best performance on forecasting.
Since we are using the same base architecture for all methods, this is expected given enough training data.
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. 6A, bottom). Next, we tested if we can transfer these learned dynamics to new recordings as we varied $n_{s}$ 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 2 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_{s}=32$)).
Moreover, our method outperformed all other baselines on forecasting, especially in very low-sample regimes, while having comparable reconstruction performance (Fig. 16).

Figure: Figure 6: A. (top) $r^{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_{s}=32$ trials for aligning to the pre-trained model.
Refer to caption: x6.png

We further 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 to the Single-Session model on reconstruction, while outperforming it on prediction for both tasks (Fig. 17 A, B). However, it demonstrated poor performance on new sessions given limited trials for learning the read-in and likelihood parameters (Fig. 17 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 a 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, 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. 18).
On the Maze task, the trajectories demonstrated more heterogeneity in responses, and decayed to a fixed point further away from the original embedding(Fig. 19).

Figure: Figure 7: 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.
Refer to caption: x7.png

## 6 Discussion

We present a novel framework for jointly inferring and learning latent dynamics from heterogeneous neural recordings across sessions/subjects during related behavioral tasks. To the best of our knowledge, this is the first approach that facilitates learning a family of dynamical systems from heterogeneous recordings in a unified latent space, while providing a concise, interpretable manifold over dynamical systems.
Our meta-learning approach mitigates the challenges of statistical inference from limited data, a common issue arising from the high flexibility of models used to approximate latent dynamics.
Empirical evaluations demonstrate that the learned embedding manifold provides a useful inductive bias for learning from limited samples, with our proposed parametrization offering greater flexibility in capturing diverse dynamics while minimizing interference.
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 it to learn shared dynamical structures.
This work enhances our capability to integrate, analyze, and interpret complex neural dynamics across diverse experimental conditions, broadening the scope of scientific inquiries possible in neuroscience.

### Limitations and Future Work

Our current framework uses event aligned neural observations; in the future, it would be useful to incorporate task-related events, to broaden its applicability to complex, unstructured tasks. Further, the model’s generalization to novel settings depends on accurate embedding inference, a challenge noted in previous works that disentangle task inference and representation learning (Hummos et al., ).
While our latent dynamics parametrization is expressive, it assumes shared structure across related tasks. Future work could extend the model to accommodate recordings without expected shared structures (for instance, by adding modularity (Márton et al., )). Investigating the performance of embedding-conditioned low-rank adaptation on RNN-based architectures presents another avenue for future research. Finally, the embedding manifold provides a map for interpolating across different dynamics. While we focus on rapid learning in this paper, our framework could have interesting applications for studying inter-subject variability, learning-induced changes in dynamics, or changes in dynamics across tasks in the future.

## Appendix A Details of Motivating Experiment from Sections 2 and 3.3

Data Generation.
The latent dynamics follow a two-dimensional limit cycle which is defined by the following system of equations

$$ $\displaystyle\dot{r}=r(1-r)^{2},$ $\displaystyle\dot{\theta}=\omega^{i},$ $\displaystyle z_{1}^{i}=r\cos\theta+5\,\mathrm{d}W_{t},\quad z_{2}^{i}=r\sin \theta+5\,\mathrm{d}W_{t},$ $$

where $\omega^{i}$, $i\in\{1,\cdots,M\}$ is the dataset specific velocity and $\mathrm{d}W_{t}$ is the Wiener process.
Specifically, for the experiment with $M=2$ datasets, we set $\omega^{1}=2$ and $\omega^{2}=5$;
for the experiment with $M=20$ datasets, we uniformly sampled 20 values for $\omega^{i}$ between 0.25 and 5.
For each value of $\omega^{i}$, we generated 128 latent trajectories for training, 64 for validation and 64 for testing, each of length $T=300$, where we used Euler-Mayurama with $\Delta t$ of 0.04 for integration.
Observations were generated according to $y_{t}^{i}\sim\mathcal{N}(C^{i}z_{t}^{i},R)$ where $R=0.01*I$ and the elements of the readout matrix $C^{i}$ were sampled from $\mathcal{N}(0,I/\sqrt{d^{z}})$; the dimensionality of the observations varied between 30 and 100.

1-shot Performance. We evaluated the generalization performance of our approach on a new dataset with $\omega^{M+1}=4.1$, which was not included in the training set, by using 1 training trajectory to train a new read-in network, $\Omega^{M+1}$ and likelihood $p_{\phi}^{M+1}$.
After training, the model displayed similar prediction performance on the new dataset ($r^{2}_{k=50}=0.94\pm 0.001$) (Fig. 8).

Figure: Figure 8: (Top) True observations on new data with $\omega=4.1$ and (Bottom) the corresponding reconstructed and predicted observations after aligning to the trained model.
Refer to caption: x8.png

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 $t=100$ to infer the embedding, $e^{i}$, and latent state, $z^{i}_{t}$.
We subsequently used the dynamics model to generate latent trajectories of length 250, $z^{i}_{t+1:t+250}$ and mapped them back to the observations via the learned likelihood function. We re-trained a 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-varied with the different velocities (Fig. 9A, left).
Further, the model recovered the correct topology of the ground truth dynamics (Fig. 9B), reflected in the forecasting performance on held out test trials (Fig. 9A, right).

Figure: Figure 9: A. (Left) Samples from the inferred embedding (see eq. 17) after training on trajectories generated from our model overlaid on 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.
Refer to caption: x9.png

## Appendix B Synthetic Bifurcating Systems

### B.1 Hopf Bifurcation

Data Generation.
Latent trajectories were generated from the following two-dimensional dynamical system,

$$ $\dot{z_{1}}=z_{2}+5\,\mathrm{d}W_{t},\hskip 30.0pt\dot{z_{2}}=-z_{1}+(\mu^{i}- z_{1}^{2})z_{2}+5\,\mathrm{d}W_{t},$ (20) $$

where the parameter $\mu^{i}$ controls the topology of the dynamics.
Specifically, when $\mu^{i}<0$, this system follows fixed point dynamics, and when $\mu>0$ the system bifurcates to a limit cycle.

We uniformly sampled $M=21$ values of $\mu^{i}$ between -1.5 and 1.5, and for each $\mu^{i}$, we generated 128 trajectories for training and 64 for testing, each trajectory was $T=350$.
Observations were generated according to $y_{t}^{i}\sim\mathcal{N}(C^{i}z_{t}^{i},R)$ where $R=0.01*I$ and the elements of the readout matrix $C^{i}$ were sampled from $\mathcal{N}(0,I/\sqrt{d^{z}})$; the dimensionality of the observations varied between 30 and 100.
For evaluating few-shot performance, we generated two additional novel datasets where $\mu^{M+1}=-0.675$ and $\mu^{M+2}=1.125$ (both values were not included in the training set).

Results.
For all embedding conditioned approaches, we set $d_{e}=1$ and learned a rank-1 change to the dynamics for our approach.
Our approach successfully learned both dynamical regimes present in the datasets and the embedding distribution encoded differences in these dynamics with high certainty given limited time bins on test trials(Fig. 10A, B).
While all approaches performed well on reconstructing observations on these datasets, our approach and the Embedding-Input outperformed other multi-session baselines on forecasting (Fig 10C).
We also evaluated the generalization performance of all methods on the 2 held-out datasets as a function of training data used for training the read-in network and observed similar trends as demonstrated by the reconstruction and k-step = 20 $r^{2}$ on test trials from these datasets, shown in Table 2.

Figure: Figure 10: A. Example observations along with the reconstructed and predicted observations from our approach for the fixed point (Top) and limit cycle (Bottom) dynamical regimes. B. Samples from the embedding as a function of 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.
Refer to caption: x10.png

**Table 2: Few-shot reconstruction and forecasting performance (k-step=20) for held-out datasets in B.1 where $n_{s}$ is the number of trials used for learning the dataset specific read-in network and likelihood.**
|  | Reconstruction |  | Forecasting |  |
| --- | --- | --- | --- | --- |
|  | $n_{s}=1$ | $n_{s}=8$ | $n_{s}=1$ | $n_{s}=8$ |
| Ours | 0.85 $\pm$ 0.054 | 0.89 $\pm$ 0.04 | 0.64 $\pm$ 0.1 | 0.69 $\pm$ 0.07 |
| Linear-Adapter | 0.84 $\pm$ 0.059 | 0.89 $\pm$ 0.04 | -0.1 $\pm$ 0.34 | 0.55 $\pm$ 0.08 |
| Single Session | 0.8 $\pm$ 0.054 | 0.88 $\pm$ 0.044 | 0.27 $\pm$ 0.08 | 0.77 $\pm$ 0.03 |
| Shared Dynamics | 0.83 $\pm$ 0.068 | 0.89 $\pm$ 0.04 | 0.32 $\pm$ 0.08 | 0.32 $\pm$ 0.04 |
| Embedding-Input | 0.86 $\pm$ 0.049 | 0.89 $\pm$ 0.04 | 0.61 $\pm$ 0.09 | 0.56 $\pm$ 0.11 |

### B.2 Duffing System

Figure: Figure 11: $(a^{i},b^{i})$ values used to generate different datasets from the Duffing system.
Refer to caption: x11.png

Data Generation. For the Duffing system described in 19, we set $c=\frac{1}{10}$ and varied values of $a$ and $b$ (shown in blue, Fig. 11) for generating 20 datasets. We additionally used 11 datasets from B.1 obtained by uniformly sampling $\mu$ between -1.5 and 1.5. For few shot evaluation of various approaches, we used two held-out datasets from the Duffing system (shown in orange, Fig. 11), as well as a dataset from the Hopf example generating by setting $\mu=-1.8$. All empirical evaluation was performed on 64 test trials from each dataset.

## Appendix C Additional Figures

Figure: Figure 12: (Left) Dynamics learnt by Linear-Adapter and (Right) Embedding-Input corresponding to the example dynamics on the true system shown in Fig. 4A.
Refer to caption: x12.png

Figure: Figure 13: (Left) True dynamics from example datasets used for pretraining in experiment 5.1. (Right) Dynamics Learnt by different embedding-conditioned parametrizations.
Refer to caption: x13.png

Figure: Figure 14: Samples from the embedding distribution as a function of the number of training trajectories on 3 seeds on held-out datasets from the Duffing system (denoted by orange, Fig. 11)
Refer to caption: x14.png

Figure: Figure 15: A. Sample from the inferred embedding for different ranks (best of 3 seeds).B. The behavior reconstruction (Top) and prediction (Bottom) $r^{2}$ for different ranks averaged over 3 training seeds.
Refer to caption: x15.png

Figure: Figure 16: Behavior reconstruction (Left) and forecasting (Right) for all methods as a function of the number of training samples for held-out sessions from Sub M and C, and two sessions from a held-out Subject.
Refer to caption: x16.png

Figure: Figure 17: 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 to the Single Session models. C. Few-shot reconstruction and forecasting performance on held out sessions and a new subject (Sub T).
Refer to caption: x17.png

Figure: Figure 18: (Left) Grid points used for generating the latent trajectories and behavior prediction with the embedding distribution overlaid. The true embedding of the CO session from Sub C used to inferring the initial condition is highlighted in green. (Right) Single trials of the predicted hand position and PCs of latent dynamics trajectories.
Refer to caption: x18.png

Figure: Figure 19: Same as 18 but for a Maze session from Sub Je (highlighted in blue).
Refer to caption: x19.png

## Appendix D Experiment Details

### D.1 Training

Since CAVIA (Concat) (Zintgraf et al., ), DYNAMO (Cotler et al., ) and CoDA (Kirchmeyer et al., ) have not been developed for joint inference and learning of dynamics, we use our framework for inference with modifications to the parametrization of the embedding conditioned dynamics function on all experiments. For our method and CoDA, we restricted parameter changes to the input, $\vW_{in}$, and hidden weights, $\vW_{hh}$. We additionally included the Frobenius norm on the embedding-conditioned weight change $\lVert h_{\varphi}(e^{i})\rVert_{P}$ along with the ELBO (eq. 15) for both approaches.

Synthetic Experiments. For the results reported in B.1 and 5.1, we used the Adam optimizer with weight decay and a Cosine annealing schedule on the learning rate for pretraining all approaches.

Motor Cortex Recordings. We binned the spiking activity in 20ms bins and smoothed it with a 25ms Gaussian filter to obtain the rates for all datasets. We further removed neurons that had a firing rate of less than 0.1Hz and aligned the neural activity to movement onset. We used 512 trials when available or 80 percent of the trialsm, each of length 36, for training all methods.

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 to evaluate few-shot performance. When aligning new data to this pre-trained model, we trained the dataset-specific read-in network, $\Omega^{i}$ and likelihood functions $p_{\phi}^{i}$ for the new dataset, in addition to the state noise, $Q^{i}$, by optimizing the ELBO for new data (eq. 15). We used Adam with weight decay for aligning, and additionally incorporated masking when aligning held-out motor cortex datasets.

### D.2 Metrics

Synthetic Experiments. We report the $r^{2}$ on observation reconstruction for test trials over the entire length of the trial for all approaches. In order to evaluate the forecasting performance, we use observations till time $t$ and sample the corresponding latent trajectories, $z^{i}_{0:t}$, from the inference network and the corresponding $e^{i}$ for the embedding-conditioned methods. We subsequently use the learned dynamics model to generate K steps ahead from $z^{i}_{t+1:t+K}$ and map these generated trajectories back to the observations. The k-step $r^{2}$ for each dataset is computed as,

$$ $r^{2}_{k}=1-\frac{\sum_{i=1}^{M}(y_{k}-\hat{y}_{k})^{2}}{\sum_{i=1}^{M}(y_{k}- \bar{y})^{2}}$ $$

where $\bar{y}$ is 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 a linear behavior decoder from the neural observations to the hand velocity of the subject, assuming a 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^{2}$. 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 dynamics model to forecast the next 20 time bins. The observations corresponding tp these forecasted trajectories were used to evaluate the prediction $r^{2}$.

### D.3 Hyperparameters

Synthetic Examples. We used the following architecture for pretraining all methods, with $d^{e}=1$ for the Motivating example and Hopf bifurcating system, and $d^{e}=2$ for the combined Duffing and Hopf example.

- •
Inference network
–
$\Omega^{i}$: $\mathrm{MLP}(d^{y^{i}},64,8)$
–
$q_{\alpha}$: [$\mathrm{GRU}(16),\mathrm{Linear}(16,2\times d^{e})]$
–
$q_{\beta}$: $[\mathrm{biGRU}(64),\mathrm{Linear}(128,4)]$
- •
Generative model
–
$f_{\theta}$: $\mathrm{MLP}(2,32,32,2)$
–
$h_{\vartheta}$: $\mathrm{MLP}(d^{e},16,16,(64+33)\times d^{r})$
–
$p_{\phi^{i}}$: $[\mathrm{Linear}(2,d^{y_{i}})]$
- •
Training
–
lr: 0.005
–
weight decay: 0.001
–
batch size: 8 from each dataset

Motor Cortex Experiment.

- •
Inference network
–
$\Omega^{i}$: $\mathrm{MLP}(d^{y^{i}},128,\mathrm{Dropout}(0.6),64)$
–
$q_{\alpha}$: [$\mathrm{GRU}(64),\mathrm{Linear}(64,2\times d^{e})]$
–
$q_{\beta}$: $[\mathrm{biGRU}(128),\mathrm{Linear}(128,60)]$
- •
Generative model
–
$f_{\theta}$: $\mathrm{MLP}(30,128,128,30)$
–
$h_{\vartheta}$: $\mathrm{MLP}(d^{e},64,64,(256+158)\times d^{r})$
–
$p_{\phi^{i}}$: $[\mathrm{Linear}(30,d^{y_{i}})]$
- •
Training
–
lr: 0.01
–
weight decay: 0.05
–
batch size: 64 trials from 20 datasets

## Peer Reviews

confidence: 5
contribution: 3
initial_rating: 6
presentation: 1
questions: See weaknesses.
soundness: 2
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.
*   **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. 
*   **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.
*   **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.
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.
weaknesses: **Missing Benchmarks & Issues in Related Works** 

- 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. 


- 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).


 **Dependence on Large-Scale Pretraining Data, and Only Within One Neural Domain:** 
- 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.

 **Accurate Embedding Inference:** 

- 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.

**Assumptions of Shared Structure:** 
- 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.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 4
questions: - How does the read-in network $\Omega$ handle different data dimensionalities?
- 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?
- What are the color schemes in Figures 3A and 4A? How are these figures generated? 
- 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?
soundness: 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. 
- The paper is generally well-written. 
- The synthetic results are strong and form a cohesive story with Section 2
- 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.
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.
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:
- **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?
- **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).

Medium: 
- Figures 3 and 4 lack detail on how they were generated, specifically as to which parameters (/embedding values $e^i$) are being used. 
- 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. 

Minor (/typos): 
- 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". 
- Eq (5), $\phi^i$ undefined. I gather from L205 that each $\phi^i$ is learned independently  -- I would clarify this aspect. 
- 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$.
- Eq. (14) and surrounding (and further, L347), $d_r$ and $r$ are used interchangeably

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 4
questions: minor comments:
- L390, "Fig" missing before 17?
- would be helpful to show the true dynamics in fig 12
soundness: 4
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.

The proof of concept presented in figs 2+3 is extremely helpful for introducting the problem, and much appreciated.

The analysis of the synthetic data in sec 5.1 is also well-designed and instructive.

The real neural data is a well-chosen dataset that balances real world heterogeneity and simplicity for exploring model performance and behavior.
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.
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_?]

Fig 5B is difficult to interpret; some more annotations or a clearer caption would be appreciated.

The 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).

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
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?
- 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.
- Is there a ground truth for Figure 2? It's hard to judge what one _should_ be looking for here.
soundness: 4
strengths: - Methods for combining insights across datasets are a huge priority in neuroscience. 
- The model builds on several well-understood components and proposes a sensible ansatz for how embeddings alter to underlying system dynamics. 
- Multiple experiments carefully examine several key variables, including data set size, look-ahead prediction, and number of training trials. These are all clearly presented.
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.

The 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.
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.
- 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.

[1] Variational Sequential Monte Carlo
Christian Naesseth, Scott Linderman, Rajesh Ranganath, David Blei Proceedings of the Twenty-First International Conference on Artificial Intelligence and Statistics, PMLR 84:968-977, 2018.

[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.

[3] T. A. Le, M. Igl, T. Jin, T. Rainforth, and F. Wood. Auto- Encoding Sequential Monte Carlo. arXiv:1705.10306, May 2017.

[4] A Unified, Scalable Framework for Neural Population Decoding
Mehdi Azabou, Vinam Arora, Venkataramana Ganesh, Ximeng Mao, Santosh Nachimuthu, Michael Mendelson, Blake Richards, Matthew Perich, Guillaume Lajoie, Eva Dyer

Meta-review: The authors present a meta-learning approach to learning dynamical systems models of neural activity, which can allow for rapid learning of new dynamical systems by training on a large set of other systems. The reviewers were unanimous that this is good work, with many giving it high marks, and so I recommend acceptance.

Decision: Accept (Spotlight)