JEDI: Jointly Embedded Inference of Neural Dynamics

Animal brains flexibly and efficiently achieve many behaviors 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 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 by 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 sharing statistical strength across trials and conditions, uncovering shared structure across conditions in a single, unified model. Using simulated RNN datasets, we demonstrate JEDI accurately learns robust, generalizable, condition-specific embeddings, and recovers the fixed point structure of a task-trained network. 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 neural recordings.