計(jì)算神經(jīng)科學(xué)作為腦科學(xué)與類腦智能之間的連接橋梁,其目標(biāo)是使用數(shù)學(xué)建模和仿真模擬的方法來(lái)闡明大腦的工作原理,同時(shí)為人工智能的發(fā)展提供新的模型和新的思想。任何一個(gè)學(xué)科的發(fā)展都離不開(kāi)工具的進(jìn)步。以人工智能為例,這一輪人工智能的快速發(fā)展很大程度上得益于深度學(xué)習(xí)框架(比如TensorFlow和PyTorch)的普及。深度學(xué)習(xí)框架目前已成為人工智能研究的基礎(chǔ)設(shè)施,因?yàn)樗鼈兲峁┝艘粋€(gè)通用的編程接口,能支持研究人員在各個(gè)領(lǐng)域靈活高效地定義各種AI模型。然而,由于計(jì)算神經(jīng)科學(xué)中的腦動(dòng)力學(xué)模型(Brain Dynamics Models)與現(xiàn)有AI模型的顯著性區(qū)別,該領(lǐng)域長(zhǎng)期缺乏一款類似于PyTorch和TensorFlow的、真正允許用戶自主靈活編程的、簡(jiǎn)單易用但靈活高效、通用的腦動(dòng)力學(xué)編程框架(General-purpose Brain Dynamics Programming Framework)。特別地,隨著越來(lái)越多海量神經(jīng)數(shù)據(jù)的產(chǎn)生、模型仿真復(fù)雜度的日益增長(zhǎng)、建模手段、方法及目標(biāo)的日趨多樣化,我們比以往任何時(shí)候都更迫切地需要開(kāi)發(fā)通用的建模工具,以幫助我們輕松構(gòu)建、模擬、訓(xùn)練和分析多尺度及大尺度的大腦動(dòng)力學(xué)模型。
然而,構(gòu)建這樣的編程框架面臨諸多挑戰(zhàn):
- 首先,大腦的復(fù)雜性使得通用框架必須能夠支持不同層次(如離子通道、神經(jīng)元、神經(jīng)網(wǎng)絡(luò))的模型構(gòu)建和跨尺度的模型組合。現(xiàn)有的大腦仿真平臺(tái)往往只關(guān)注單一或少數(shù)幾個(gè)尺度,而我們需要一個(gè)能夠全面涵蓋各種尺度的框架。
- 其次,為了全面理解大腦功能,我們不僅需要模擬神經(jīng)活動(dòng),還需要分析其背后的機(jī)制,甚至根據(jù)數(shù)據(jù)或任務(wù)訓(xùn)練模型。這就需要一個(gè)通用的編程框架能夠整合多種建模需求。然而,目前的大腦模擬器主要集中在模擬方面,而忽略了訓(xùn)練和分析的需求。
- 第三個(gè)挑戰(zhàn)是在保持編程便捷性的同時(shí)實(shí)現(xiàn)高運(yùn)行性能,這一點(diǎn)對(duì)于腦動(dòng)力學(xué)建模尤為重要,因?yàn)槠涮匦允蛊浜茈y在便捷的 Python 界面中進(jìn)行高效仿真。
- 最后,腦動(dòng)力學(xué)建模是一個(gè)快速發(fā)展的領(lǐng)域,新的概念、模型和數(shù)學(xué)方法不斷涌現(xiàn)。這就要求通用編程框架必須具備可擴(kuò)展性,以便適應(yīng)這一領(lǐng)域的發(fā)展。
在2023年12月22日,北京大學(xué)心理與認(rèn)知科學(xué)學(xué)院、IDG麥戈文腦科學(xué)研究所、定量生物學(xué)中心、以及北大-清華生命科學(xué)聯(lián)合中心的吳思研究團(tuán)隊(duì)在《eLife》雜志上發(fā)表了一篇題為“BrainPy, a flexible, integrative, efficient, and extensible framework for general-purpose brain dynamics programming” 的文章。他們開(kāi)發(fā)了BrainPy這一先進(jìn)框架,旨在為腦動(dòng)力學(xué)編程提供一套定制化的基礎(chǔ)設(shè)施。BrainPy框架具備高度的靈活性,允許用戶自由定義腦動(dòng)力學(xué)模型。它集成了事件驅(qū)動(dòng)算子、微分方程求解器和通用模型構(gòu)建接口等工具,使用戶能夠根據(jù)需要靈活調(diào)整模型。這一綜合性的基礎(chǔ)設(shè)施為構(gòu)建全面而強(qiáng)大的腦動(dòng)力學(xué)建??蚣艿於藞?jiān)實(shí)基礎(chǔ)。BrainPy不僅提供了一個(gè)綜合性的研究平臺(tái),還支持模擬、訓(xùn)練和分析等多種應(yīng)用場(chǎng)景。通過(guò)一個(gè)模型,用戶可以進(jìn)行模擬、離線學(xué)習(xí)、在線學(xué)習(xí)或反向傳播訓(xùn)練,并進(jìn)一步進(jìn)行低維分岔分析或高維慢點(diǎn)分析等分析工作。在性能方面,BrainPy通過(guò)面向?qū)ο蟮募磿r(shí)編譯(JIT)和針對(duì)腦動(dòng)力學(xué)特性的專用算子實(shí)現(xiàn)了顯著的提升。這使得代碼執(zhí)行更為高效,進(jìn)一步提高了模擬和分析的準(zhǔn)確性和速度。此外,BrainPy還具備良好的可擴(kuò)展性。新的功能和擴(kuò)展可以通過(guò)插件模塊輕松實(shí)現(xiàn),而底層算子甚至可以在用戶級(jí)Python接口中進(jìn)行擴(kuò)展。這使得BrainPy能夠適應(yīng)不斷發(fā)展的腦動(dòng)力學(xué)領(lǐng)域,滿足不斷變化的科研需求??偠灾?,BrainPy是一個(gè)功能強(qiáng)大、靈活且可擴(kuò)展的腦動(dòng)力學(xué)通用編程框架。它為用戶提供了定制化的基礎(chǔ)設(shè)施,支持模擬、訓(xùn)練和分析等多種應(yīng)用場(chǎng)景,并實(shí)現(xiàn)了高效的代碼執(zhí)行。BrainPy的推出將極大地推動(dòng)腦動(dòng)力學(xué)領(lǐng)域的發(fā)展,為科研人員提供強(qiáng)有力的工具,以深入探索大腦的奧秘。
eLife編委和評(píng)委對(duì)BrainPy的貢獻(xiàn)給予了高度評(píng)價(jià)(其中評(píng)委Marcel Stimberg是BrainPy的最大競(jìng)爭(zhēng)對(duì)手Brian2的開(kāi)發(fā)者):
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編程系統(tǒng)
以下是對(duì)BrainPy的簡(jiǎn)要介紹。
圖1:BrainPy框架概覽圖
首先,為支持其成為通用編程框架的目標(biāo),BrainPy提供了腦動(dòng)力學(xué)建模所必需的基礎(chǔ)設(shè)施(圖1A)。這些基礎(chǔ)設(shè)施是一系列實(shí)用工具的集合,旨在提供基本服務(wù),使用戶能夠輕松、靈活和高效地進(jìn)行各種類型的腦動(dòng)力學(xué)建模。具體而言,BrainPy實(shí)現(xiàn)了:1)基于密集矩陣的常規(guī)計(jì)算和基于稀疏連接的事件驅(qū)動(dòng)計(jì)算的數(shù)學(xué)算子;2)各種微分方程的數(shù)值積分器,用于動(dòng)力學(xué)系統(tǒng)演化的數(shù)值積分模擬,是神經(jīng)動(dòng)力學(xué)系統(tǒng)的重要計(jì)算模塊;3)用于構(gòu)建多尺度腦動(dòng)力學(xué)模型的通用模型構(gòu)建接口,以及與之相關(guān)的即時(shí)編譯,以實(shí)現(xiàn)這些模型的高效運(yùn)行;4)專門(mén)用于腦動(dòng)力學(xué)建模的工具箱。
其次,腦動(dòng)力學(xué)模型具有模塊化、多尺度和層次化的關(guān)鍵特性,BrainPy設(shè)計(jì)了一種模塊化、可組合、靈活的編程范式以匹配這些特性。這一范式通過(guò)內(nèi)部的DynamicalSystem接口實(shí)現(xiàn),DynamicalSystem支持在任何組織水平上定義腦動(dòng)力學(xué)模型。給定一個(gè)動(dòng)力學(xué)系統(tǒng),無(wú)論其復(fù)雜性如何,用戶都可以將其實(shí)現(xiàn)為DynamicalSystem類。對(duì)于復(fù)雜的動(dòng)力學(xué)模型,如霍奇金-赫胥黎(HH)類型神經(jīng)元模型到大規(guī)模皮層網(wǎng)絡(luò),它們的模型定義可以通過(guò)子組件的組合來(lái)實(shí)現(xiàn)。DynamicalSystem支持層次化的可組合編程,使得由底層組件組成的模型可以層次化地作為新組件構(gòu)成更上層的模型。這一屬性對(duì)于構(gòu)建多尺度腦模型非常有用。圖2展示了一個(gè)從通道(圖2A)到神經(jīng)元(圖2B)到網(wǎng)絡(luò)(圖2C)再到系統(tǒng)(圖2D)的遞歸組合模型的例子。此外,為了方便用戶,BrainPy提供了許多常用模型,包括離子通道、神經(jīng)元、突觸、群體和網(wǎng)絡(luò),作為構(gòu)建大規(guī)模模型的構(gòu)建塊。
圖2:使用BrainPy模塊化和層級(jí)化地構(gòu)建腦動(dòng)力學(xué)模型
接著,BrainPy提供了一個(gè)集成平臺(tái),可以全面地執(zhí)行腦動(dòng)力學(xué)模型的模擬、訓(xùn)練和分析。對(duì)于模擬,BrainPy設(shè)計(jì)了接口brainpy.DSRunner來(lái)模擬腦模型的動(dòng)力學(xué),可用于模擬任何層次的模型。對(duì)于訓(xùn)練,BrainPy提供了brainpy.DSTrainer接口來(lái)支持這一功能。DSTrainer的不同子類提供了不同的訓(xùn)練算法,這些算法可以用來(lái)訓(xùn)練不同類型的模型。分析模型動(dòng)力學(xué)對(duì)于理解模型行為的潛在機(jī)制來(lái)說(shuō),與模型模擬和訓(xùn)練一樣重要。針對(duì)一個(gè)動(dòng)力學(xué)系統(tǒng),BrainPy提供了接口brainpy.DSAnalyzer用于自動(dòng)動(dòng)力學(xué)分析,不同類別的DSAnalyzer實(shí)現(xiàn)了不同的分析方法。
不僅如此,BrainPy有著高效的性能表現(xiàn),這背后得益于即時(shí)編譯(JIT)技術(shù)和基于事件驅(qū)動(dòng)的稀疏計(jì)算算子。JIT編譯在運(yùn)行時(shí)將動(dòng)態(tài)Python代碼轉(zhuǎn)換成靜態(tài)機(jī)器代碼,這可以顯著減少Python解釋的時(shí)間成本。BrainPy還創(chuàng)新性地引入了專門(mén)設(shè)計(jì)的底層算子,用于加速稀疏連接網(wǎng)絡(luò)中事件驅(qū)動(dòng)的計(jì)算。這些專用算子包括與突觸前神經(jīng)元、突觸后神經(jīng)元和突觸相關(guān)的變量之間的轉(zhuǎn)換,以及稀疏計(jì)算算子、事件驅(qū)動(dòng)計(jì)算算子和即時(shí)連接算子。通過(guò)采用這些專門(mén)算子,BrainPy顯著減少了突觸計(jì)算所需的時(shí)間。和多個(gè)已有的腦動(dòng)力學(xué)仿真工具進(jìn)行了比較,例如NEURON[7]、Nest[8]、Brian2[9]等,在COBA、COBAHH、Decision making networks等常用的benchmarks,BrainPy都展現(xiàn)出了優(yōu)越性。
圖3:NEURON、Nest、Brian2 和 BrainPy 在不同計(jì)算設(shè)備下的速度比較。
最后,BrainPy還具有非常好的可擴(kuò)展性,實(shí)用工具、功能模塊、機(jī)器學(xué)習(xí)方法等,都可以直接通過(guò)Python接口輕松定制和擴(kuò)展。值得一提的是,BrainPy還為事件驅(qū)動(dòng)算子的自定義提供了Python接口,用戶寫(xiě)Python代碼即可在CPU和GPU上完成事件驅(qū)動(dòng)算子的自定義。這為引入新的方法和理論帶來(lái)極高的便利性,也有利于計(jì)算神經(jīng)科學(xué)更好地與其他學(xué)科進(jìn)行交叉。
近期,吳思團(tuán)隊(duì)對(duì)BrainPy功能做了進(jìn)一步拓展,強(qiáng)調(diào)BrainPy可用于構(gòu)建可微分的大腦模擬,使得神經(jīng)網(wǎng)絡(luò)的生物物理仿真和任務(wù)訓(xùn)練的結(jié)合成為可能,從而可以將深度學(xué)習(xí)與神經(jīng)計(jì)算建模連接起來(lái)。該工作已被今年人工智能的頂級(jí)會(huì)議ICLR 2024接受。
BrainPy的生態(tài)建設(shè)
一個(gè)軟件平臺(tái)的成功最重要也是唯一的評(píng)判標(biāo)準(zhǔn)是這個(gè)軟件被廣大用戶所選用,為此BrainPy也在積極建設(shè)生態(tài)環(huán)境。BrainPy已在GitHub和OpenI上開(kāi)源,建立了brainpy-examples,brainpy-datasets等倉(cāng)庫(kù)為用戶提供支持,累計(jì)下載超過(guò)15萬(wàn)次。開(kāi)源地址為:
- https://github.com/brainpy/BrainPy
- https://openi.pcl.ac.cn/OpenI/BrainPy
BrainPy的配套教材《神經(jīng)計(jì)算建模實(shí)戰(zhàn)》在2023年由電子工業(yè)出版社出版,結(jié)合了神經(jīng)計(jì)算的理論知識(shí)和代碼實(shí)戰(zhàn),目前被部分高校和研究所在教學(xué)中采用,線上培訓(xùn)人數(shù)也超600人。BrainPy 連續(xù)兩年獲得OpenI 新一代人工智能開(kāi)源開(kāi)放平臺(tái)優(yōu)秀項(xiàng)目嘉獎(jiǎng)。此外,BrainPy已經(jīng)適配由廣東省智能科學(xué)與技術(shù)研究院研發(fā)的全球首顆億級(jí)神經(jīng)元規(guī)模的可編程類腦晶圓計(jì)算芯片天琴芯。
作者信息
北京大學(xué)心理與認(rèn)知科學(xué)學(xué)院、麥戈文腦科學(xué)研究所、定量生物學(xué)中心、北大-清華生命科學(xué)聯(lián)合中心的吳思教授為該論文的通訊作者,北京大學(xué)心理與認(rèn)知科學(xué)學(xué)院博雅博士后王超名博士為該論文的第一作者,北京大學(xué)前沿交叉學(xué)科研究院博士生張?zhí)烨?、研究助理陳嘯宇和賀思超為該論文做出重要貢獻(xiàn)。該文章得到了科技部重大項(xiàng)目計(jì)劃的資助。
文章鏈接
https://doi.org/10.7554/eLife.86365
2024-01-19