可微分生物物理模拟JAXLEY：实现大规模神经动力学模型的梯度优化新范式

《Nature Methods》：Jaxley : differentiable simulation enables large-scale training of detailed biophysical models of neural dynamics

本刊推荐：研究人员针对生物物理神经元模型参数优化困难的问题，开发了基于自动微分和GPU加速的可微分模拟器JAXLEY。该工具通过反向传播实现梯度下降优化，能够高效拟合电压/钙成像数据，并训练包含10万个参数的形态详细神经元网络完成计算机视觉任务。这项技术为跨尺度神经计算机制研究提供了新范式。

　　

在神经科学领域，生物物理神经元模型犹如一台精密的"细胞计算机"，能够通过离子通道、突触传导等机制揭示神经活动的细胞基础。然而，科学家们长期面临一个核心挑战：如何为这些包含大量微分方程的复杂模型找到合适的参数，使其既能匹配实验观测数据，又能执行特定的计算任务？传统方法如同"盲人摸象"，难以高效优化成千上万的参数。

目前主流的NEURON等模拟器缺乏自动微分功能，无法利用深度学习中的梯度下降方法。参数优化往往依赖遗传算法等无梯度方法，计算成本随参数数量急剧增加。对于具有复杂形态的神经元模型，参数空间更是浩瀚如海，使得大规模生物物理网络的训练几乎成为不可能完成的任务。

来自德国蒂宾根大学的研究团队在《Nature Methods》发表了创新性研究成果，他们开发的JAXLEY框架将可微分模拟引入神经科学领域。这一工具融合了自动微分(automatic differentiation)和GPU加速技术，使研究人员能够像训练深度神经网络一样，使用梯度下降法优化大规模生物物理模型。

关键技术方法包括：基于JAX框架构建的可微分微分方程求解器，支持多级检查点(multilevel checkpointing)的内存优化策略，Polyak梯度下降优化算法，以及面向电压/钙成像数据的动态时间规整(dynamic time warping, DTW)损失函数。研究利用艾伦细胞类型数据库(Allen Cell Types Database)的公开数据和新合成的计算数据集，通过哈密顿蒙特卡洛(Hamiltonian Monte Carlo)进行贝叶斯推断验证。

JAXLEY的准确性、速度与可扩展性

研究人员首先验证了JAXLEY的数值准确性。通过构建大鼠海马CA1锥体细胞和鼠标视觉皮层第5层神经元的生物物理模型，JAXLEY在亚毫秒和亚毫伏分辨率上与NEURON模拟器结果高度一致。在计算效率方面，GPU上的JAXLEY比CPU上的NEURON快两个数量级，能够并行模拟多达100万个单室神经元。

拟合单神经元模型到细胞内记录

研究团队将JAXLEY应用于拟合第5层锥体细胞(L5PC)模型。仅需9步梯度下降迭代，JAXLEY就能找到与合成体细胞电压记录高度匹配的参数集，比遗传算法节省近90%的模拟次数。在对艾伦细胞类型数据库中真实膜片钳记录的拟合中，通过并行化1000次梯度下降运行，获得的电压轨迹与实验记录高度相似。

拟合多参数单神经元模型

JAXLEY成功优化了包含1390个参数的L5PC模型，这些参数控制着树突树上每个分支的离子电导分布。梯度下降找到了与整个树突树电压响应高度匹配的参数集，而遗传算法在500次迭代后损失值仍高出两个数量级。通过贝叶斯推断获得的参数后验分布，揭示了数据对不同树突区域电导分布的约束程度差异。

非线性单神经元计算

研究展示了一个具有双树突结构的神经元模型如何通过学习离子通道密度和形态学参数，实现非线性模式分离任务。该神经元学会仅在其中一个树突受到强刺激时产生动作电位，形成了非线性的决策边界。

视网膜树突钙信号的混合模型

结合视网膜神经节细胞(RGC)的钙成像数据，研究团队构建了包含光感受器、双极细胞和生物物理RGC模型的混合架构。JAXLEY同时优化了287个突触参数和320个细胞参数，使模型能够预测15,000对噪声刺激-钙响应数据。训练后的模型再现了实验观察到的钙信号区室化响应特性，其感受野集中在记录位点周围。

生物物理RNN解决工作记忆任务

研究实现了由Hodgkin-Huxley型神经元组成的递归神经网络(RNN)，并训练其完成证据整合和延迟匹配样本任务。通过优化输入权重、递归权重和输出权重(109个参数)，网络能够在500毫秒时间尺度上以99.9%的准确率区分正负刺激。随着突触增益的增加，网络动力学从稳定状态过渡到混沌状态，最大李雅普诺夫指数(Lyapunov exponent)定量刻画了这一转变。

训练10万个参数的生物物理网络

最终，研究团队构建了一个包含28×28输入神经元、64个形态详细的隐藏神经元和10个输出神经元的前馈网络，共优化55,000个膜电导参数和51,000个突触参数。该网络在MNIST手写数字识别任务中达到94.2%的准确率，显著高于线性分类器。对抗性攻击分析表明，生物物理网络与传统人工神经网络具有相似的脆弱性。

JAXLEY的出现标志着生物物理建模领域的重要转折点。正如深度学习框架革命了人工神经网络的训练，可微分模拟将为神经科学研究开启新篇章。这一技术不仅使大规模生物物理网络的优化成为可能，还将促进梯度引导的贝叶斯推断、动力学稳定性分析、对抗鲁棒性评估等新型分析方法的开展。

随着连接组学、形态学和转录组学等跨尺度数据的爆炸式增长，JAXLEY提供了一个强大框架，能够将这些多模态测量整合到统一的生物物理模拟中。正如深度学习网络可以通过端到端训练从数据中自动学习特征，JAXLEY使研究人员能够以类似方式"训练"生物物理网络，从而在保持生物学合理性的同时，发现实现神经计算的最佳参数配置。

虽然挑战依然存在——如毫秒级生物物理过程与秒级行为时间尺度的衔接、非凸损失面上的局部极小值问题等——但JAXLEY为探索这些基本问题提供了前所未有的工具。这项研究不仅推动了计算神经科学的方法学创新，更为理解脑功能的多尺度机制开辟了新途径。