Meta-Dynamical State Space Models for Integrative Neural Data Analysis

TL;DR

This work addresses learning generalizable latent neural dynamics across heterogeneous recordings by meta-learning a family of dynamical systems. It introduces a low-dimensional dynamical embedding $e \in \mathbb{R}^{d_e}$ that conditions a shared latent dynamics via a hypernetwork, with a low-rank constraint to keep parameters tractable, and uses dataset read-in networks to align recordings into a common latent space. Inference is performed with a sequential variational scheme (DKF) that jointly estimates $e$ and latent trajectories while learning dataset-specific readouts, demonstrated on synthetic bifurcations and motor cortex data, including few-shot transfers. The approach yields an interpretable embedding manifold over dynamics, enabling rapid generalization to new recordings and tasks, with potential impact on integrative neuroscience analyses and foundation-model–style pretraining of neural dynamics.

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.

Figures (22)

• 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.
• 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$ or $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.
• Figure 3: A. Mean dynamical system corresponding 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. \ref{['eq:embedding_inference']}). C. Forecasting $r^2$ at $(k=50)$-step for models trained with $M=2$ or $M=20$ datasets.
• 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.
• Figure 5: Visualizing the embedding manifold. (Left) Each point corresponds to a sample from the inferred embedding distribution (see eq. \ref{['eq:embedding_inference']}) corresponding to each recording. (Right) The condition-averaged latent dynamics for a session from Maze (Sub Ni) (Top) and a CO Session (Bottom) generated by the model, along with the corresponding real and forecasted behavior.