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