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.
Abstract (translated)
动物的大脑能够灵活而高效地利用单一神经网络完成多种行为任务。现代神经科学的一个核心目标是将大脑的灵活性映射到神经群体的动力学机制上。然而，从有限、嘈杂且高维度的实验神经记录中识别特定任务的动力学规则仍然是一个重大挑战，因为实验数据往往只能提供对脑状态和动力学机制的部分访问。尽管直接受约束于神经数据的递归神经网络（RNNs）在推断底层动态机制方面已经十分有效，但它们通常仅限于单一任务领域，并且难以跨行为条件进行泛化。 在这里，我们介绍了一种名为JEDI的分层模型，该模型通过学习共享权重的嵌入空间来捕捉跨任务和情境的神经动力学。此模型能够在单个统一模型中重现个体样本的神经动态的同时扩展到任意大小和复杂度的数据集，并揭示条件间共有的结构。使用模拟RNN数据集，我们展示了JEDI能够准确地学习出强大的、具有泛化能力的特定条件下的嵌入结构。通过逆向工程由JEDI学到的权重，我们证明它恢复了真实的固定点结构，并在特征谱中揭露了神经动态的关键特性。 最后，我们将JEDI应用于猴子伸手运动过程中记录到的运动皮层数据集，从而提取出对运动控制神经动力学机制的理解。我们的工作表明，同时学习情境嵌入和递归权重为从记录本身推断大脑动力学提供了可扩展且泛化的推理方式。
URL
https://arxiv.org/abs/2603.10489
