# JEDI: Jointly Embedded Inference of Neural Dynamics

arXiv: https://arxiv.org/abs/2603.10489 · HTML: https://mlsanity.com/paper/2603.10489/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 26 of its 58 references that are in the map.
- First posted on arXiv 11 Mar 2026.

## Builds on (most cited first, 26 of 26)

- [HyperNetworks](https://mlsanity.com/paper/1609.09106/) (2016)
- [Continual learning with hypernetworks](https://mlsanity.com/paper/1906.00695/) (2019)
- [Spectral Bias and Task-Model Alignment Explain Generalization in Kernel Regression and Infinitely Wide Neural Networks](https://mlsanity.com/paper/2006.13198/) (2020)
- [Learnable latent embeddings for joint behavioral and neural analysis](https://mlsanity.com/paper/2204.00673/) (2022)
- [Coupled Oscillatory Recurrent Neural Network (coRNN): An accurate and (gradient) stable architecture for learning long time dependencies](https://mlsanity.com/paper/2010.00951/) (2020)
- [Generalized Teacher Forcing for Learning Chaotic Dynamics](https://mlsanity.com/paper/2306.04406/) (2023)
- [Linking connectivity, dynamics and computations in low-rank recurrent neural networks](https://mlsanity.com/paper/1711.09672/) (2017)
- [Generalizing to New Physical Systems via Context-Informed Dynamics Model](https://mlsanity.com/paper/2202.01889/) (2022)
- [Lyapunov spectra of chaotic recurrent neural networks](https://mlsanity.com/paper/2006.02427/) (2020)
- [On Lyapunov Exponents for RNNs: Understanding Information Propagation Using Dynamical Systems Tools](https://mlsanity.com/paper/2006.14123/) (2020)
- [Inferring stochastic low-rank recurrent neural networks from neural data](https://mlsanity.com/paper/2406.16749/) (2024)
- [Learning Interpretable Hierarchical Dynamical Systems Models from Time Series Data](https://mlsanity.com/paper/2410.04814/) (2024)
- [Low Tensor Rank Learning of Neural Dynamics](https://mlsanity.com/paper/2308.11567/) (2023)
- [Nuclear Norm Regularization for Deep Learning](https://mlsanity.com/paper/2405.14544/) (2024)
- [Dynamics of random recurrent networks with correlated low-rank structure](https://mlsanity.com/paper/1909.04358/) (2019)
- [Shaping dynamics with multiple populations in low-rank recurrent networks](https://mlsanity.com/paper/2007.02062/) (2020)
- [Aligned and oblique dynamics in recurrent neural networks](https://mlsanity.com/paper/2307.07654/) (2023)
- [Generalizable, real-time neural decoding with hybrid state-space models](https://mlsanity.com/paper/2506.05320/) (2025)
- [Charting and navigating the space of solutions for recurrent neural networks](https://mlsanity.com/paper/2111.09356/) (2021)
- [Meta-Dynamical State Space Models for Integrative Neural Data Analysis](https://mlsanity.com/paper/2410.05454/) (2024)
- [Analyzing Populations of Neural Networks via Dynamical Model Embedding](https://mlsanity.com/paper/2302.14078/) (2023)
- [eXponential FAmily Dynamical Systems (XFADS): Large-scale nonlinear Gaussian state-space modeling](https://mlsanity.com/paper/2403.01371/) (2024)
- [The computational role of structure in neural activity and connectivity](https://mlsanity.com/paper/2308.16772/) (2023)
- [CORNN: Convex optimization of recurrent neural networks for rapid inference of neural dynamics](https://mlsanity.com/paper/2311.10200/) (2023)
- [Weight-Space Linear Recurrent Neural Networks](https://mlsanity.com/paper/2506.01153/) (2025)
- [Expressivity of Neural Networks with Random Weights and Learned Biases](https://mlsanity.com/paper/2407.00957/) (2024)

## Abstract (arXiv)

Animal brains flexibly and efficiently achieve many behavioral tasks with a single neural network. A core goal in modern neuroscience is to map the mechanisms of the brain's flexibility onto the dynamics underlying neural populations. However, identifying task-specific dynamical rules from limited, noisy, and high-dimensional experimental neural recordings remains a major challenge, as experimental data often provide only partial access to brain states and dynamical mechanisms. While recurrent neural networks (RNNs) directly constrained neural data have been effective in inferring underlying dynamical mechanisms, they are typically limited to single-task domains and struggle to generalize across behavioral conditions. Here, we introduce JEDI, a hierarchical model that captures neural dynamics across tasks and contexts by learning a shared embedding space over RNN weights. This model recapitulates individual samples of neural dynamics while scaling to arbitrarily large and complex datasets, uncovering shared structure across conditions in a single, unified model. Using simulated RNN datasets, we demonstrate that JEDI accurately learns robust, generalizable, condition-specific embeddings. By reverse-engineering the weights learned by JEDI, we show that it recovers ground truth fixed point structures and unveils key features of the underlying neural dynamics in the eigenspectra. Finally, we apply JEDI to motor cortex recordings during monkey reaching to extract mechanistic insight into the neural dynamics of motor control. Our work shows that joint learning of contextual embeddings and recurrent weights provides scalable and generalizable inference of brain dynamics from recordings alone.