[論文レビュー] 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

ひとこと要約

この論文では、JAXに基づいて構築された微分可能で生物学的・生物理的に現実的なスパイクニューロンネットワークのシミュレーションを可能にする、微分可能な脳シミュレータBrainPyを紹介する。このフレームワークは、スケーラブルで柔軟なマルチスケール脳モデルのサポートを備え、微分可能で生物学的・生物理的に現実的なスパイクニューロンネットワークのシミュレーションを統合的に実現する。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.

研究の動機と目的

• 脳シミュレーション分野と脳にインspiredされたコンピューティング（BIC）分野の間で共有されるプログラミングフレームワークの欠如に対処すること。
• 従来の脳シミュレータが勾配ベースの学習に必要な微分可能性を欠いているという制限を克服すること。
• 脳シミュレーションにおける生物理的リアリズムとBIC分野のディープラーニングフレームワークのトレーニング効率のギャップを埋めること。
• 微分可能でイベント駆動的かつマルチスケールの脳モデルをサポートする統合的で拡張可能でスケーラブルなシミュレーション環境を開発すること。

提案手法

• JAXとXLAを活用して自動微分と高性能計算を実現し、脳動態シミュレーションの能力を拡張する。
• 大規模で疎結合なニューロンネットワークを効率的にシミュレートするためのカスタムスパースおよびイベント駆動型演算子を導入する。
• ニューロン、シナプス、ネットワークの各レベルでマルチスケール脳モデルを構築可能なモジュラーかつ合成可能なインターフェースを採用する。
• 多様なシナプスプロジェクションとダイナミクスを柔軟に扱えるように、シナプス計算のための新規抽象化を採用する。
• メモリ集約的な脳動態を効率的に管理するため、オブジェクト指向のジャストインタイム（JIT）コンパイルアプローチを適用する。
• サロゲート勾配と電圧正則化を用いて、スパイクネットワークにおける時間遅延を伴うバックプロパゲーション（BPTT）を可能にし、安定した学習を実現する。

Figure 1: The overview of BrainPy architecture.

実験結果

リサーチクエスチョン

• RQ1 生物学的・生物理的に現実的な脳シミュレーションと、脳にインspiredされたコンピューティングのための微分可能トレーニングを両立できる統合フレームワークを構築可能か？
• RQ2 高精度な脳シミュレータに微分可能性を効率的に統合する方法は何か？ その際、シミュレーションのパフォーマンスや生物理的正確性を損なわないようにできるか？
• RQ3 JAXベースのフレームワーク、たとえばBrainPyは、微分可能動的特性を備えた大規模でイベント駆動的かつスパースなニューロンネットワークのシミュレーションにどの程度スケーラブルに拡張可能か？
• RQ4 脳シミュレータに自動微分を統合することで、認知タスクのためのスパイクニューロンネットワークのエンドツーエンドトレーニングがどのように可能になるか？
• RQ5 BrainPyは、既存の脳シミュレータやディープラーニングフレームワークに比べて、ハイブリッド脳シミュレーションとBICワークロードにおいて、どの程度のパフォーマンスとスケーラビリティの利点を提供するか？

主な発見

• BrainPyは、JAXの自動微分を活用して生物学的に妥当なスパイクニューロンネットワークの微分可能シミュレーションを可能にし、時間遅延を伴うバックプロパゲーション（BPTT）によるエンドツーエンドトレーニングを実現する。
• ベンチマークタスクにおいて高い効率性とスケーラビリティを達成し、大規模な脳ネットワークシミュレーションに適したパフォーマンスを示している。
• サロゲート勾配と電圧正則化の活用により、スパイクネットワークの安定したトレーニングが可能となり、スパイク再帰的ネットワークにおける遅延マッチングタスクの成功実装が達成された。
• BrainPyのレートベースモデルでは収束が遅く、高い精度に到達するまでに数十エポックのトレーニングが必要であることが判明し、レートコードドネットワークのトレーニングの課題が浮き彫りになった。
• 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.