[论文解读] A differentiable brain simulator bridging brain simulation and brain-inspired computing

Chaoming Wang, Tianqiu Zhang+4 more|arXiv (Cornell University)|Nov 9, 2023

Advanced Memory and Neural Computing被引用 4

一句话总结

本文介绍了 BrainPy，一个基于 JAX 构建的可微分大脑模拟器，通过支持高效、可扩展且灵活的多尺度大脑模型，实现了大脑模拟与类脑计算的统一。该框架利用 JAX 的自动微分和自定义稀疏算子，实现了在训练复杂脉冲神经网络时的最先进效率与可微分性。

ABSTRACT

Brain simulation builds dynamical models to mimic the structure and functions of the brain, while brain-inspired computing (BIC) develops intelligent systems by learning from the structure and functions of the brain. The two fields are intertwined and should share a common programming framework to facilitate each other's development. However, none of the existing software in the fields can achieve this goal, because traditional brain simulators lack differentiability for training, while existing deep learning (DL) frameworks fail to capture the biophysical realism and complexity of brain dynamics. In this paper, we introduce BrainPy, a differentiable brain simulator developed using JAX and XLA, with the aim of bridging the gap between brain simulation and BIC. BrainPy expands upon the functionalities of JAX, a powerful AI framework, by introducing complete capabilities for flexible, efficient, and scalable brain simulation. It offers a range of sparse and event-driven operators for efficient and scalable brain simulation, an abstraction for managing the intricacies of synaptic computations, a modular and flexible interface for constructing multi-scale brain models, and an object-oriented just-in-time compilation approach to handle the memory-intensive nature of brain dynamics. We showcase the efficiency and scalability of BrainPy on benchmark tasks, highlight its differentiable simulation for biologically plausible spiking models, and discuss its potential to support research at the intersection of brain simulation and BIC.

研究动机与目标

• 解决大脑模拟与类脑计算（BIC）领域之间缺乏共享编程框架的问题。
• 克服传统大脑模拟器缺乏可微分性、难以进行基于梯度训练的局限性。
• 弥合大脑模拟中的生物物理真实性与类脑计算中深度学习框架的训练效率之间的差距。
• 开发一个统一、可扩展且可扩展的模拟环境，支持可微分、事件驱动和多尺度大脑模型。

提出的方法

• 利用 JAX 和 XLA 实现自动微分与高性能计算，扩展其在大脑动力学模拟中的能力。
• 引入自定义稀疏与事件驱动算子，高效模拟大规模、稀疏连接的神经网络。
• 采用模块化、可组合的接口，用于构建跨神经元、突触与网络的多尺度大脑模型。
• 提出一种新颖的突触计算抽象，灵活处理多样化的突触投射与动力学。
• 采用面向对象的即时（JIT）编译方法，高效管理内存密集型大脑动力学。
• 通过使用替代梯度和电压正则化，实现脉冲网络中的时间反向传播（BPTT）稳定训练。

Figure 1: The overview of BrainPy architecture.

实验结果

研究问题

• RQ1 能否构建一个统一框架，同时支持生物合理的大脑模拟与类脑计算的可微分训练？
• RQ2 如何高效地将可微分性集成到高保真的大脑模拟器中，而不牺牲模拟性能或生物物理准确性？
• RQ3 基于 JAX 的框架如 BrainPy 在多大规模上能够实现事件驱动、稀疏神经网络模拟与可微分动力学的扩展？
• RQ4 在大脑模拟器中集成自动微分，如何实现脉冲神经网络在认知任务上的端到端训练？
• RQ5 BrainPy 在混合大脑模拟与 BIC 工作负载中，相较于现有大脑模拟器与深度学习框架，具有哪些性能与可扩展性优势？

主要发现

• BrainPy 利用 JAX 的自动微分能力，实现了生物合理脉冲神经网络的可微分模拟，支持通过时间反向传播（BPTT）进行端到端训练。
• 该框架在基准任务中实现了高效率与可扩展性，表现出适合大规模大脑网络模拟的性能。
• 通过使用替代梯度和电压正则化，实现了脉冲网络的稳定训练，并成功在脉冲递归网络中完成延迟匹配任务。
• BrainPy 中的速率编码模型收敛速度较慢，需数十个训练周期才能达到高精度，凸显了速率编码网络的训练挑战。
• BrainPy 的模块化接口与自定义稀疏算子，实现了跨不同生物细节层次的多尺度大脑模型的灵活、可扩展且内存高效的模拟。
• 该框架通过在单一、可扩展且可微分的编程环境中统一模型定义、模拟与训练，成功弥合了大脑模拟与 BIC 之间的鸿沟。

Figure 2: Synaptic projections in BrainPy. (A) The AlignPre and AlignPost projections offer a decoupled interface for managing dynamics and the communication between dynamics. (B) The synaptic communication allows for diverse implementations, including the utilization of DL models.