计算神经科学作为脑科学与类脑智能之间的连接桥梁,其目标是使用数学建模和仿真模拟的方法来阐明大脑的工作原理,同时为人工智能的发展提供新的模型和新的思想。任何一个学科的发展都离不开工具的进步。以人工智能为例,这一轮人工智能的快速发展很大程度上得益于深度学习框架(比如TensorFlow和PyTorch)的普及。深度学习框架目前已成为人工智能研究的基础设施,因为它们提供了一个通用的编程接口,能支持研究人员在各个领域灵活高效地定义各种AI模型。然而,由于计算神经科学中的脑动力学模型(Brain Dynamics Models)与现有AI模型的显著性区别,该领域长期缺乏一款类似于PyTorch和TensorFlow的、真正允许用户自主灵活编程的、简单易用但灵活高效、通用的脑动力学编程框架(General-purpose Brain Dynamics Programming Framework)。特别地,随着越来越多海量神经数据的产生、模型仿真复杂度的日益增长、建模手段、方法及目标的日趋多样化,我们比以往任何时候都更迫切地需要开发通用的建模工具,以帮助我们轻松构建、模拟、训练和分析多尺度及大尺度的大脑动力学模型。

然而,构建这样的编程框架面临诸多挑战:

  • 首先,大脑的复杂性使得通用框架必须能够支持不同层次(如离子通道、神经元、神经网络)的模型构建和跨尺度的模型组合。现有的大脑仿真平台往往只关注单一或少数几个尺度,而我们需要一个能够全面涵盖各种尺度的框架。
  • 其次,为了全面理解大脑功能,我们不仅需要模拟神经活动,还需要分析其背后的机制,甚至根据数据或任务训练模型。这就需要一个通用的编程框架能够整合多种建模需求。然而,目前的大脑模拟器主要集中在模拟方面,而忽略了训练和分析的需求。
  • 第三个挑战是在保持编程便捷性的同时实现高运行性能,这一点对于脑动力学建模尤为重要,因为其特性使其很难在便捷的 Python 界面中进行高效仿真。
  • 最后,脑动力学建模是一个快速发展的领域,新的概念、模型和数学方法不断涌现。这就要求通用编程框架必须具备可扩展性,以便适应这一领域的发展。

在2023年12月22日,北京大学心理与认知科学学院、IDG麦戈文脑科学研究所、定量生物学中心、以及北大-清华生命科学联合中心的吴思研究团队在《eLife》杂志上发表了一篇题为“BrainPy, a flexible, integrative, efficient, and extensible framework for general-purpose brain dynamics programming” 的文章。他们开发了BrainPy这一先进框架,旨在为脑动力学编程提供一套定制化的基础设施。BrainPy框架具备高度的灵活性,允许用户自由定义脑动力学模型。它集成了事件驱动算子、微分方程求解器和通用模型构建接口等工具,使用户能够根据需要灵活调整模型。这一综合性的基础设施为构建全面而强大的脑动力学建模框架奠定了坚实基础。BrainPy不仅提供了一个综合性的研究平台,还支持模拟、训练和分析等多种应用场景。通过一个模型,用户可以进行模拟、离线学习、在线学习或反向传播训练,并进一步进行低维分岔分析或高维慢点分析等分析工作。在性能方面,BrainPy通过面向对象的即时编译(JIT)和针对脑动力学特性的专用算子实现了显著的提升。这使得代码执行更为高效,进一步提高了模拟和分析的准确性和速度。此外,BrainPy还具备良好的可扩展性。新的功能和扩展可以通过插件模块轻松实现,而底层算子甚至可以在用户级Python接口中进行扩展。这使得BrainPy能够适应不断发展的脑动力学领域,满足不断变化的科研需求。总而言之,BrainPy是一个功能强大、灵活且可扩展的脑动力学通用编程框架。它为用户提供了定制化的基础设施,支持模拟、训练和分析等多种应用场景,并实现了高效的代码执行。BrainPy的推出将极大地推动脑动力学领域的发展,为科研人员提供强有力的工具,以深入探索大脑的奥秘。

eLife编委和评委对BrainPy的贡献给予了高度评价(其中评委Marcel Stimberg是BrainPy的最大竞争对手Brian2的开发者):

Editor's evaluation:

“ The paper introduces a new, important framework for neural modelling that promises to offer efficient simulation and analysis tools for a wide range of biologically-realistic neural networks. It provides convincing support for the ease of use, flexibility, and performance of the framework, and features a solid comparison to existing solutions in terms of accuracy. The work is of potential interest to a wide range of computational neuroscientists and researchers working on biologically inspired machine learning applications.”

BrainPy编程系统

以下是对BrainPy的简要介绍。


图1:BrainPy框架概览图

首先,为支持其成为通用编程框架的目标,BrainPy提供了脑动力学建模所必需的基础设施(图1A)。这些基础设施是一系列实用工具的集合,旨在提供基本服务,使用户能够轻松、灵活和高效地进行各种类型的脑动力学建模。具体而言,BrainPy实现了:1)基于密集矩阵的常规计算和基于稀疏连接的事件驱动计算的数学算子;2)各种微分方程的数值积分器,用于动力学系统演化的数值积分模拟,是神经动力学系统的重要计算模块;3)用于构建多尺度脑动力学模型的通用模型构建接口,以及与之相关的即时编译,以实现这些模型的高效运行;4)专门用于脑动力学建模的工具箱。

其次,脑动力学模型具有模块化、多尺度和层次化的关键特性,BrainPy设计了一种模块化、可组合、灵活的编程范式以匹配这些特性。这一范式通过内部的DynamicalSystem接口实现,DynamicalSystem支持在任何组织水平上定义脑动力学模型。给定一个动力学系统,无论其复杂性如何,用户都可以将其实现为DynamicalSystem类。对于复杂的动力学模型,如霍奇金-赫胥黎(HH)类型神经元模型到大规模皮层网络,它们的模型定义可以通过子组件的组合来实现。DynamicalSystem支持层次化的可组合编程,使得由底层组件组成的模型可以层次化地作为新组件构成更上层的模型。这一属性对于构建多尺度脑模型非常有用。图2展示了一个从通道(图2A)到神经元(图2B)到网络(图2C)再到系统(图2D)的递归组合模型的例子。此外,为了方便用户,BrainPy提供了许多常用模型,包括离子通道、神经元、突触、群体和网络,作为构建大规模模型的构建块。


图2:使用BrainPy模块化和层级化地构建脑动力学模型

接着,BrainPy提供了一个集成平台,可以全面地执行脑动力学模型的模拟、训练和分析。对于模拟,BrainPy设计了接口brainpy.DSRunner来模拟脑模型的动力学,可用于模拟任何层次的模型。对于训练,BrainPy提供了brainpy.DSTrainer接口来支持这一功能。DSTrainer的不同子类提供了不同的训练算法,这些算法可以用来训练不同类型的模型。分析模型动力学对于理解模型行为的潜在机制来说,与模型模拟和训练一样重要。针对一个动力学系统,BrainPy提供了接口brainpy.DSAnalyzer用于自动动力学分析,不同类别的DSAnalyzer实现了不同的分析方法。

不仅如此,BrainPy有着高效的性能表现,这背后得益于即时编译(JIT)技术和基于事件驱动的稀疏计算算子。JIT编译在运行时将动态Python代码转换成静态机器代码,这可以显著减少Python解释的时间成本。BrainPy还创新性地引入了专门设计的底层算子,用于加速稀疏连接网络中事件驱动的计算。这些专用算子包括与突触前神经元、突触后神经元和突触相关的变量之间的转换,以及稀疏计算算子、事件驱动计算算子和即时连接算子。通过采用这些专门算子,BrainPy显著减少了突触计算所需的时间。和多个已有的脑动力学仿真工具进行了比较,例如NEURON[7]、Nest[8]、Brian2[9]等,在COBA、COBAHH、Decision making networks等常用的benchmarks,BrainPy都展现出了优越性。


图3:NEURON、Nest、Brian2 和 BrainPy 在不同计算设备下的速度比较。

最后,BrainPy还具有非常好的可扩展性,实用工具、功能模块、机器学习方法等,都可以直接通过Python接口轻松定制和扩展。值得一提的是,BrainPy还为事件驱动算子的自定义提供了Python接口,用户写Python代码即可在CPU和GPU上完成事件驱动算子的自定义。这为引入新的方法和理论带来极高的便利性,也有利于计算神经科学更好地与其他学科进行交叉。

近期,吴思团队对BrainPy功能做了进一步拓展,强调BrainPy可用于构建可微分的大脑模拟,使得神经网络的生物物理仿真和任务训练的结合成为可能,从而可以将深度学习与神经计算建模连接起来。该工作已被今年人工智能的顶级会议ICLR 2024接受。

BrainPy的生态建设

一个软件平台的成功最重要也是唯一的评判标准是这个软件被广大用户所选用,为此BrainPy也在积极建设生态环境。BrainPy已在GitHub和OpenI上开源,建立了brainpy-examples,brainpy-datasets等仓库为用户提供支持,累计下载超过15万次。开源地址为:

  • https://github.com/brainpy/BrainPy
  • https://openi.pcl.ac.cn/OpenI/BrainPy

BrainPy的配套教材《神经计算建模实战》在2023年由电子工业出版社出版,结合了神经计算的理论知识和代码实战,目前被部分高校和研究所在教学中采用,线上培训人数也超600人。BrainPy 连续两年获得OpenI 新一代人工智能开源开放平台优秀项目嘉奖。此外,BrainPy已经适配由广东省智能科学与技术研究院研发的全球首颗亿级神经元规模的可编程类脑晶圆计算芯片天琴芯。

作者信息

北京大学心理与认知科学学院、麦戈文脑科学研究所、定量生物学中心、北大-清华生命科学联合中心的吴思教授为该论文的通讯作者,北京大学心理与认知科学学院博雅博士后王超名博士为该论文的第一作者,北京大学前沿交叉学科研究院博士生张天秋、研究助理陈啸宇和贺思超为该论文做出重要贡献。该文章得到了科技部重大项目计划的资助。

文章链接

https://doi.org/10.7554/eLife.86365

 


2024-01-19