[논문 리뷰] A differentiable brain simulator bridging brain simulation and brain-inspired computing

Chaoming Wang, Tianqiu Zhang외 4명|arXiv (Cornell University)|2023. 11. 09.

Advanced Memory and Neural Computing인용 수 4

한 줄 요약

이 논문은 JAX 기반으로 구축된 미분 가능 뇌 시뮬레이터인 BrainPy를 소개한다. 이는 뇌 시뮬레이션과 뇌에 영감을 받은 계산을 통합하여, 효율적이고 확장 가능하며 유연한 방식으로 다중 척도 뇌 모델을 지원하는 생물학적으로 현실적인 스파iking 신경망 시뮬레이션을 가능하게 한다. 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) 분야 간의 공통 프로그래밍 프레임워크 부족 문제를 해결하기 위해.
• 기존 뇌 시뮬레이터가 기울기 기반 학습을 위한 미분 가능성이 부족한 한계를 극복하기 위해.
• 뇌 시뮬레이션의 생물학적 현실성과 BIC 분야의 딥러닝 프레임워크의 학습 효율성 간 격차를 메우기 위해.
• 다양한 스케일의 뇌 모델을 지원하는 통합적이고 확장 가능하며 확장 가능한 시뮬레이션 환경을 개발하기 위해.

제안 방법

• JAX와 XLA를 활용하여 자동 미분 및 고성능 계산을 수행하며, 뇌 동역학 시뮬레이션에 대한 기능을 확장한다.
• 대규모이고 희소하게 연결된 신경망을 효율적으로 시뮬레이션하기 위해 고유한 희소 및 이벤트 기반 연산자를 도입한다.
• 신경세포, 시냅스, 네트워크 수준에서 다양한 척도의 뇌 모델을 구성할 수 있도록 모듈러하고 조합 가능한 인터페이스를 구현한다.
• 다양한 시냅스 투영과 동역학을 다루기 위한 새로운 추상화를 도입하여 탄력적인 방식으로 시냅스 계산을 처리한다.
• 메모리 집약적인 뇌 동역학을 효율적으로 관리하기 위해 객체 지향적 즉시 컴파일(JIT) 컴파일 방식을 적용한다.
• 서브티튜트 기울기와 전압 정규화를 사용하여 스파킹 네트워크에서 시간을 거슬러 올라가는 backpropagation(BPTT)을 가능하게 한다.

Figure 1: The overview of BrainPy architecture.

실험 결과

연구 질문

• RQ1 생물학적으로 현실적인 뇌 시뮬레이션과 뇌에 영감을 받은 계산을 위한 미분 가능 학습을 동시에 지원하는 통합 프레임워크를 구축할 수 있는가?
• RQ2 고성능 뇌 시뮬레이터에서 성능이나 생물학적 정확도를 희생시키지 않고도 효율적으로 미분 가능성을 통합할 수 있는가?
• RQ3 JAX 기반 프레임워크인 BrainPy가 다양한 스케일의 이벤트 기반 및 희소 신경망 시뮬레이션에 대해 미분 가능한 동역학을 얼마나 잘 확장할 수 있는가?
• RQ4 자기 미분 기능을 뇌 시뮬레이터에 통합함으로써 스파킹 신경망의 엔드 투 엔드 학습이 어떻게 가능해지는가?
• RQ5 기존의 뇌 시뮬레이터와 딥러닝 프레임워크에 비해 BrainPy는 하이브리드 뇌 시뮬레이션 및 BIC 워크로드에서 어떤 성능 및 확장성 이점을 제공하는가?

주요 결과

• BrainPy는 JAX의 자동 미분 기능을 활용해 생물학적으로 타당한 스파킹 신경망의 미분 가능 시뮬레이션을 가능하게 하며, 시간을 거슬러 올라가는 backpropagation(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.