JEDI: Jointly Embedded Inference of Neural Dynamics

Anirudh Jamkhandi; Ali Korojy; Olivier Codol; Guillaume Lajoie; Matthew G. Perich

q-bio.NCMar 11, 2026

JEDI: Jointly Embedded Inference of Neural Dynamics

Authors: Anirudh Jamkhandi, Ali Korojy, Olivier Codol, Guillaume Lajoie, Matthew G. Perich

Organizations: Mila - Quebec AI Institute · Université de Montréal

Abstract

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.

Figures & tables

Figure 1: JEDI leverages a hypernetwork-based framework to flexibly generate RNN weights based on contextual inputs, all learned directly from a loss computed against the RNN output and time series neural recordings.

Figure 2: Quantifying the quality of the embeddings. A) Synthetic multi-task data generation setup. B) 2D PCA visualization of context embeddings. Each point corresponds to a trial, color-coded by input signal type. C) Training reconstruction accuracy $R^{2}$ across different methods. D) Accuracy of task classification from the learned embeddings. E) Generalization accuracy $R^{2}$ from the center of the learned embedding for each task F) Confusion matrix of generalization $R^{2}$ scores, expanding the results in Panel E. Rows represent training tasks, and columns indicate test tasks.

Figure 3: Impact of embedding noise on model performance, comparing reconstruction

R^{2}

as increasing noise is applied to the embeddings. The added noise was scaled according to the standard deviations of embeddings for each model

Figure 4: A) We drove the chaotic Teacher RNN with sinusoidal inputs at different frequencies (1–10 Hz). B) The eigen spectra of the weights inferred with JEDI exhibit a characteristic expansion along the imaginary axis as input frequency increases.

Figure 5: JEDI identifies fixed point structure in task-trained networks. A) We fit JEDI to a network trained to perform the MemoryPro task. This task has four contextual trial phases. B) Context embedding learned by JEDI colored by trial phase (left) and response direction (right). C) Fixed point structure for the four trial phases for the ground-truth task-trained network. D) Fixed point structure inferred by JEDI.

Figure 6: JEDI applied to monkey motor cortex data during reaching. A) We apply JEDI to recordings of the motor and premotor cortex during a center-out reaching task. Schematic adapted from Gallego et al. (2020) . B) 2D PCA projections of the learned embeddings for each model. Each point corresponds to a single trial, color-coded by reach direction. C) Reconstruction $R^{2}$ across all trials for each model. D) Classification accuracy for reach direction from each learned embedding. E) Generalization $R^{2}$ using the center of each embedding cluster for the different models. F) Confusion matrix for generalization across different contexts for each model, further visualizing the results in Panel E. G) Eigenspectra of learned JEDI weights during preparation (top) and movement (bottom) colored by the 8 reach directions. During movement, eigenvalues clustered close to the unit circle, indicating dynamics approaching the edge of stability. H) Quantification of eigenspectra changes by the magnitude of each eigenvalue for preparation (gray) and execution (green). I) Quantification of eigenspectra changes by the phase of each eigenvalue. J) Lyapunov exponents across all learned weights for preparation and execution. K) Distributions of maximum Lyapunov exponents across all learned trials for preparation and execution. L) Stable fixed points during execution inferred by JEDI colored by reach directions. Single-trial neural trajectories for each corresponding reach is also plotted.

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Figure S1 : Reconstructions of neural trajectories from various trained models. A) Reconstructions of Synthetic Data. B) Reconstructions of the Monkey cortex data.

Figure S2 : Jedi performance on the synthetic data generated by the teacher RNN. The heatmaps show accuracy of reconstruction as measured by the coefficient of determination $R^{2}$ score. The red circular dot highlights the hyperparameter that was chosen for experiments. A) $R^{2}$ heatmap for various model rank vs Hypernetwork hidden layer size. B) $R^{2}$ heatmap for various Model rank vs Embedding size of the hypernetworks and c) $R^{2}$ heatmap for various embedding size vs Hidden size of the hypernetwork.

Figure S3 : Performance for Jedi on the multi-task synthetic data generated by the teacher RNN with different rank of connectivity matrix

J

. The bar plot shows accuracy of reconstruction as measured by coefficient of determination

R^{2}

score (with 1 s.e.m). Jedi was able to fit well on varied data ranks with the generated weights set to rank=5.

Figure S4 : Eigenvalue spectra of learned recurrent weights on multi-task synthetic data. JEDI learned consistent structure across tasks, while JEDI-full exhibited a dense, isotropic cloud without task separation.

Figure S5 : Eigenvalue spectra of Jedi-full weights on monkey motor cortical activity. JEDI-full exhibited a dense, isotropic cloud without task/direction separation.

Figure S6 : Eigenvalue spectra of learned recurrent weights on multiple frequency dataset. The legend indicates the specific sine frequencies used to train the models, JEDI-full consistently exhibits a dense, isotropic spectral cloud. This distribution lacks any discernible structure or cluster formation corresponding to the underlying task frequencies.

Figure S7 : Assessing embedding quality on monkey reaching data during preparation. a) 2D PCA visualization of embeddings from different models. Each point represents a sample trial, color-coded by reach direction. b) Confusion matrices of direction classification accuracies. Rows correspond to the true direction, and columns to the predicted direction. c) Confusion matrices showing generalization performance across reach directions. Each cell indicates the $R^{2}$ score when decoding one direction (column) using mean embeddings of the other(row).

Figure S8 : Assessing embedding quality on monkey reaching data during preparation. a) 2D PCA visualization of embeddings from different models. Each point represents a sample trial, color-coded by reach direction. b) Confusion matrices of direction classification accuracies. Rows correspond to the true direction, and columns to the predicted direction. c) Confusion matrices showing generalization performance across reach directions. Each cell indicates the $R^{2}$ score when decoding one direction (column) using mean embeddings of the other(row).

Figure S9 : Stable fixed points during preparation inferred by JEDI colored by reach directions. Single-trial neural trajectories for each corresponding reach is also plotted.

Figure S10 : 3d PCA projection of fixed points from Jedi across 4 periods on MemoryPro task. The fixed points structure closely resembled the structure of task-trained RNN.

Figure S11 : When the context embedding remains frozen during training, the eigenvalue spectra of the learned recurrent weights fail to capture task-specific signatures. Consequently, joint optimization of both contexts and weights is essential to develop meaningful, trial-specific representations.

Explore similar work

Mar 27, 2026

Identifying Connectivity Distributions from Neural Dynamics Using Flows

Connectivity structure shapes neural computation, but inferring this structure from population recordings is degenerate: multiple connectivity structures can generate identical dynamics. Recent work uses low-rank recurrent neural networks (lrRNNs) to infer low-dimensional latent dynamics and connectivity from observed activity, enabling a mechanistic interpretation of the dynamics. However, standard approaches for training lrRNNs can recover spurious structures irrelevant to the underlying dynamics. We first characterize the identifiability of connectivity structures in lrRNNs and determine conditions under which a unique solution exists. To find such solutions, we develop an inference framework based on maximum entropy and continuous normalizing flows (CNFs), trained via flow matching. Instead of estimating a single connectivity matrix, our method learns a distribution over connection weights that is maximally unbiased over unidentifiable components while matching the observed dynamics. This approach captures complex yet necessary distributions such as heavy-tailed connectivity found in empirical data. We validate our method on synthetic datasets with connectivity structures that generate multistable attractors, limit cycles, and ring attractors, and demonstrate its applicability in recordings from rat frontal cortex during decision-making. Our framework shifts circuit inference from recovering connectivity to identifying which connectivity structures are computationally required, and which are artifacts of underconstrained inference.

Timothy Doyeon Kim, Ulises Pereira-Obilinovic, Yiliu Wang +2

Allen Institute, Seattle, WA, USA · University of Washington, Seattle, WA, USA

Recurrent Neural Networks Functional Connectivity

Ask PDF

Sep 29, 2026

NeuroDyn-EEG: An Interpretable Pre-trained Model for EEG Based on Neural Dynamics

Clinical scalp electroencephalography (EEG) offers a noninvasive window into neural dynamics of neuropsychiatric disorders. However, discriminative deep models often lack anatomically indexed physiological interpretability. We propose NeuroDyn-EEG, a pretraining framework integrating generative priors from neural dynamics. It couples an extended Jansen-Rit neural mass model, leadfield-based source projection, and simulation-based parameter inversion. Trained on synthetic parameter-EEG pairs within physiological ranges, NeuroDyn-EEG estimates 11 regional parameter families across 90 AAL regions plus one global parameter from standard 19-channel EEG, using only ~2.43M trainable parameters. We evaluate the framework across three levels. First, controlled simulations demonstrate robust parameter recovery under diverse noise conditions, while real resting-state EEG evaluations confirm spectral and phase consistency in an inverse-forward closed loop. Second, on four clinical benchmarks (AD65, PD31, Figshare MDD, and TUAB), NeuroDyn-EEG achieves competitive classification performance, securing the highest BACC, AUROC, and AUCPR on PD31 and MDD, and highest BACC on AD65. Third, post hoc regional analyses reveal disease-specific alterations: local synaptic connectivity C_1 involves the most altered regions in AD65, whereas the firing threshold theta ranks first in MDD, offering testable mechanistic hypotheses. Overall, NeuroDyn-EEG maps scalp EEG to anatomically indexed dynamical parameters, bridging representation learning and mechanistic neurophysiology. Code: https://github.com/Gnosis-Neurodynamics/NeuroDyn-EEG.

Yi Cui, Tong Zhao, Jiaxin Lei +5

Gnosis Neurodynamics Co. Ltd, Room 202, West Area, Building North 1, No. 9 Yard, Shengmingyuan Road, Changping District, Beijing, China · School of Biomedical Engineering, Tsinghua Medicine, Tsinghua University, Beijing, China · University of California, Davis, department of psychology +2

Electroencephalography Neural Posterior Estimation

Ask PDF

May 19, 2026

BrainDyn: A Sheaf Neural ODE for Forecasting and Generating Brain Dynamics

Brain activity unfolds as coordinated patterns across regions over time, and modeling these dynamics is a central goal in computational neuroscience. Because activity is distributed over connected brain regions, it is naturally represented on a graph whose nodes are brain regions or recording electrodes and whose edges capture relationships between these nodes, with message-passing modeling regional communication. However, standard message-passing networks place every node in a single shared feature space and typically aggregate messages from neighbors by scalar summation, which does not align with neurobiological evidence of brain communication. In order to model the heterogeneous nature of the brain, we introduce BrainDyn, a graph ordinary differential equation (ODE) network built on a cellular sheaf with dynamic restriction maps, recomputed from the forecasting context window rather than fixed once, mediating region-to-region communication through learned linear transformations that drive the time evolution of neural activity. We evaluate BrainDyn on resting-state fMRI, EEG, and simulated activity from the NEST simulator. Across these modalities, we test multi-region forecasting and autoregressive generation of brain transients, and, using NEST, prediction of altered transients following controlled perturbations. BrainDyn outperforms competing models across modalities, while analysis of the learned features shows that they are expressive, interpretable, and biologically aligned.

Siddharth Viswanath, Panayiotis Ketonis, Cassidy All +10

Yale University · Boise State University · University of Wisconsin–Madison

Graph Neural Networks Sheaf Neural Networks Neuroimaging Neural ODEs Spatiotemporal GNNs Brain Network Analysis

Ask PDF