A differentiable brain simulator bridging brain simulation and brain-inspired computing
Authors: Chaoming Wang
Tianqiu Zhang
Sichao He
Hongyaoxing Gu
Shangyang Li
Si Wu
Venue: ICLR 2024 poster (2024)
Keywords: brain simulator
brain simulation
computational neuroscience
brain-inspired computing
Primary area: applications to neuroscience & cognitive science
TL;DR: We developed BrainPy, a differentiable brain simulator, to help bridge the gap between brain simulation and brain-inspired computing.

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, and highlight its differentiable simulation for biologically plausible spiking models.

Published as a conference paper at ICLR 2024
A DIFFERENTIABLE BRAIN SIMULATOR BRIDGING BRAIN SIMULATION AND BRAIN-INSPIRED COMPUTING
Chaoming Wang', Tianqiu Z Sichao He², Hongyaoxing Gu², Soangyangg Li', Si Wult School of Psychological and Cognitive Sciences, IDG/MMCiovern Institute for Brain Research, Bejing Key Laboratory of Behavior and Mental Health. Beking-Tsinghun Center for Life Sciences, Center of Quantitative Biology, Academy for Advanced Interdisciplinary Studies, Peking University "School of Software Engineering, Beijing Jiaotong University
Institute of Software, Chinese Academy of Sciences
"The authors contributed equally (wancchaoming, syli, s siwu @ppu.. edu. en, tianqiuakita@stu. pku. edu on
Corresponding author
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) frameworkss fail to capture the biophysical realism and complexity of brain dynamics. In this paper, we introduce BrainPy, 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. 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, 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, and highlight its differentiable simulation for biologically plausible spiking models.
INTRODUCTION
Brain simulation aims to elucidate brain functions by building dynamical models that mimic the structure and dynamics the brain (Gerstner et al.) 2014), while brain- irsinninspired computing aims to develop intelligent systems by learning from the structure and computational principles of the brain (Mehonic & Kenyon, 202T). The two fields are intertwined and their developments can facilitate each other. For example, brain simulation can provide BIC with models of neurons, synapses, networks, and inspirational information processing principles; while BIC can provide brain simulation with efficient algorithms for optimizing model parameters, simulation tools for running large-scale networks, and testing bed for validating hypothesized neural mechanisms. Ideally, the two fields should share a common programming framework, so that they can benefit from each other's development by sharing models, mathematical tools, and emerging findings.
However, up to now, none of the existing software in the two fields can fully achieve this goal. Traditional brain simulators, such as NEURON (Hines & Carnevale, 1997). NEST (Gewaltig & Diesmann 2007), and Brian2 (Goodman Brette [2088] Sttmmerg et al./ 2019), are designed for simulating brain dynamics models with high fidelity and accuracy. They rely customized numerical solvers and data structures that are not compatible with automatic differentiation, and hence canmot support training models with standard gradient-based methods. On the other hand, by leveraging the automatic differentiation functionality of deep learning (DL) frameworks like Py PyTorch (Paszke et al. 2019) and TensorFlow (Aaai et al. 2016)), existing BIC libraries, such snn' oonTorch (Eshraghian et al.

---

Published as a conference paper at ICLR 2024
2021), Norse (Pehle & Pedersen 2021), and SpikingJelly (Fang et al. 2020 provide convenient interfaces for building and training spike neural networks (SNNs). They are, however, not designed to capture the unique and important features of brain dynamics, and hence not suitable to simulate large-scale brain networks with realistic biophysical properties.
In this paper, we propose BrainPy brain simulators, BrainPy leverages the power of the JAX (Frostig et al. 2018), allowing seamless integration with AI models. However, BrainPy goes beyond integration and introduces dedicated optimizations that unleash the full potential of flexible, efficient, and scalable brain simulator within the JAX ecosystem. To capture the sparse and event-driven nature of brain computation. BrainPy provides a wide range of customized primitive operators. For enhanced flexibility in model construction across various brain organization scales, BrainPy offers a modular and composable interface. To handle the complexity of synaptic computations, BrainPy introduces novel abstraction for executing diverse synaptic projections. Additionally, to tackle the memory- seetory-itensive demands of brain dynamics, BrainPy employs an object-oriented just-in-time (JIT) compilation approach. Lever- aging the automatic differentiation capabilities of JAX, BrainPy represents a unique differentiable brain simulator that bridges the gap between brain simulation and BIC fields. We demonstrate the efficiency and scalability of BrainPy on several brain simulation and BIC tasks and showcase its ability to train biologically plausible spiking models with differentiability.
innovative solution to bridge this gap. Unlike traditional
2 RELATED WORKS
Brain Simulators. Different brain simulators normally have different focuses. NEURON (Hines & Carnevale 1997) allows users to define detailed biophysical models of neurons and synapses, with complex morphology and ion channels. NEST (Gewaltig & Diesmann) 2007) focuses on large-scale network models of point neurons and synapses, with simplified dynamics and connectivity patterns. Brian2 (Stimberg et al. 2014: 2019) targets being flexible and intuitive, allowing users to easily define dynamical models, environment interactions, and experimental protocols. Currently, the dominant programming approach in brain simulation is descriptive language (Blundell et al.] (2018). by which users can use text (Stimberg al., 2019 Vitay et al. 2015), JSON (Dai et al.] [2020] Dura-Bernal eet al. 2019), or XML (Gleeson et al.) 2010) files to describe the model, and then translate the descriptive information into high-efficient C++ or CUDA code. The main advantage of this approach is the clear decoupling between mathematical description from its implementation details. However, this advantage comes with expensive costs, which include, for instance, the lack of flexibility and generality of defining new models not covered by the predefined constructs and functions of the custom language, the difficulty of integrating and interfacing with other tools and frameworks not using the same format, and the high learning cost for the unfamiliar syntax of the custom language. These limitations prevent the application of existing brain simulators to BIC models.
BIC Libraries. SNNs are the current dominating models in BIC for their advantages in biological interoperability and energy efficiency. A number of programming libraries have been developed for SNNs, such as NengoDL. (Rasmussen] 2018), BindsNet (Hazan et al. 2018) 8), snnTorch (Eshraghian
cct al. 2021). Norse (Pehle & Pedersen 2021). SpikingJelly (Fang et al. 2020). and BrainCog Zeng
ct al. 2022). These librarics utilize DL frameworks, such as Py Torch (aszkke et al.) 2019) and
TensorFlow (Abadi et al. 2016). to achieve the training of SNNs on various tasks that cannot be done in traditional brain simulators. So far, BIC libraries have mainly focused on the combination of spiking neurons with DL models, e.g., spiking convolutional neural networks and spiking recurrent neural networks. However, these libraries fall short of high-fidelity brain simulation. First, DL frameworks lack the dedicated components for sparse, event-driven, and scalable computation required for brain dynamics models. Second, BIC libraries designed for machine learning tasks often lack the essential capabilities to support realistic neuronal and synaptic simulations based on experimental data. The brain encompasses intricate biochemical and biophysical processes that span vast scales in both space and time. Unfortunately, current BIC libraries without these dedicated optimizations face significant challenges in accurately modeling such complex biophysical characteristics.
THE DESIGN OF BRAINPY
BrainPy is designed to take the combined advantages of brain simulators and DL framewokss. This innovative tool is specifically engineered to leverage the strengths of JAX (Frostig et al. 2018),
2

---

Published as a conference paper at ICLR 2024
A
B
JAX integration
Core Designs
Integrated Functions
Object-oriented JIT
Sparse Event-driven
2:ffiesncy
Simulation
Dynamical System Plexfility
Just-in-Time Compilation +Brainny object
Flax
Training
Noural Network
Tranning algondtems
JIT conn & Projection Scalaniity
Analysis
Figure I: The overview of BrainPy architecture.
powerful AI framework, while simultaneously expanding upon it to enable flexible, efficient, and scalable simulation of various brain dynamics models.
By building upon JAX, BrainPy inherits the robust computational capabilities and extensive library support that JAX provides (see Figure [[]. This enables seamless integration of DL. techniques, such as neural network architectures and gradient-based optimization algorithms, into its brain simulation framework. BrainPy facilitates smooth interoperation with existing JAX-based machine learning libraries (Babuschkin et al., 2020). Users can easily transform models from other libraries into BrainPy objects. For example, by utilizing the brainpy. arainpy.dnn... Fromā lax functionality. a Flax (Heek) et al. 2020) model can be treated as a BrainPy module. Conversely, users can convert BrainPy models into formats compatible with other libraries. The brainpy. dnn. ToFlax ax feature, for instance, enables the interpretation a dynamical system in BrainPy as a Flax recurrent cell, allowing brain models developed in BrainPy to be utilized within machine-learning context.
However, BrainPy goes beyond being a mere extension of JAX. introduces novel and fundamental features to empower users to simulate and analyze the intricate dynamics of the brain (refer to F and Section 4. (1) BrainPy provides flexibility in modeling brain dynamics. The brain exhibits unique multi-scale organization characterized by hierarchical modularity. To address this, BrainPy offers a modular and composable interface specifically designed for brain dynamics programming. With this interface, users easily define and customize complex brain models across multiple levels of organization, from low-level ion channels, neurons, and synapses to high-level networks and systems. (2) BrainPy prioritizes efficiency in simulating brain dynamics. Brain dy- namics involve unique properties such as event-driven computation and sparse connections. BrainPy abstracts these operations into primitive operators and gives a full set of primitive operators that achieve major speedups, from two to five orders of magnitude faster than traditional dense operators. (3) BrainPy tackles the scalability challenges of the brain's 's large-scar scale structure. The brain consists of massively interconnected neurons arranged in complex networks and circuits. To handle this, BrainPy offers specialized support, including just-in-time connectivity operators with zero memory footprint, automatic synapse merging for network topology optimization, and parallelization strategies for distributed computing.
Leveraging its distinctive brain simulation capabilities and scamless integration with DL techniques, BrainPy offers an integrated platform for comprehensive simulation, training, and analysis of brain dynamics models (see Figure Firstly, it enables efficient simulation of models various scales of organization. These simulations can be executed in parallel across multiple threads, processors, and devices, facilitating parameter explorations and enhancing performance. Secondly, BrainPy supports the training of diverse model types based on neural data or behavioral tasks. For example, reservoir computing models can be trained using BrainPy's 's online and offline learning algorithms. and detailed spiking neural networks and rate-basce recurrent neural networks can be trained using its back propagation- bbasd algorithms. Thirdly, BrainPy offers automatic analysis interfaces that unravel the underlying mechanisms of simulated or trained models. For instance, low-dimensional analyzers can perform phase plane and bifurcation analysis. On the other hand, high-dimensional analyzers facilitate slow point computation and linearization analysis for high-dimensional systems. Lastly, BrainPy achieves exceptional running speed through the utilization of JIT compilation. Through novel object-oriented JIT transformation. BrainPy enables whole -raph optimization of class- based models into executable processes using JAX and XLA on CPU, GPU, or TPU devices during simulation, training, and analysis tasks (see Figure TD).
3

---

Published as a conference paper at ICLR 2024
4 METHODOLOGY
In this section, we present the specific optimizations implemented in BrainPy for brain modeling, and highlight the advancements achieved the operator, model, and compilation levels to facilitate its goal of unifying models in brain simulation and BIC.
4.1 DEDICATED OPERATORS FOR SPARSE AND EVENT-DRIVED COMPUTATION
Compared to DNN models, brain dynamics models perform computation in a different way. They usually have sparse connections (neurons in network are typically interconnected with a probability smaller than 0.2 (Potjans & Diesmann 2012) ) and perform event-driven computations (synaptic states are updated only when the presynaptic spiking event occurs (Ros et al., (006; Roy al.) 2019)). These unique features greatly hinder the rumning efficiency of brain dynamics models if conventional dense array operators are used (see Section 5.T). Traditional brain simulators utilize specific data structure and accompanying event-driven computation code to address this prohlem (Kunkel et al., 2012 Jordan et al. 2018). Nevertheless, this solution is confined to predetermined simulation scenarios, restricting its applicability to other domains. Moreover, it lacks compatibility with automatic differentiation, a crucial element of the backpropagation algorithm, thereby impeding efficient training of brain dynamics models using gradient- based optimization algorithms.
To unify brain simulation and BIC programming workflows, BrainPy abstracts these specialized func- tions in brain simulation as reusable primitive operators, that they can be flexibly fused, chained, or combined to define any complex computations of brain models as desired by the user. Particularly, BrainPy provides set of sparse and event-driven operators in its breinpy- .math. sparse and brainpy. math. event modules, respectively. Sparse operators can not only store synaptic connec- tivity compactly by avoiding the memory usage of unnecessary zeros, but also compute synaptic currents efficiently by using only nonzero entries. Although modern numerical computing libraries have provided diverse routines for sparse computation, they are prone to encountering problems such as duplicate memory storage of identical synaptic projections and redundant memory access of homogeneous synaptic weights. contrast, BrainPy's sparse operators prioritize the optimization of sparse computations specifically for brain dynamics modeling. As result, they offer superior memory efficiency and significantly faster speeds compared to sparse computation routines designed for general domains. Moreover, event-driven operators in BrainPy can further take advantage of the discontinuous nature of spikes by only performing computations when spiking events arrive. They can lead to significant improvements in efficiency (typically orders of magnitude faster, see Section 5.1. as it does not constantly update the state of the system when no spike arrives.
In comparison to traditional brain simulators, our approach of encapsulating the characteristic brain operations as primitive operators streamlines the support of gradient- based optimization algorithms which are typically used in training BIC models nowadays. Notably, all dedicated operators in BrainPy offer general implementations for both forward- and reverse-mode automatic differentiation, so that brain dynamics models building upon these operators can be differentiable to be used for madient-based optimization tasks (see Section 5.3)
4.2 SCALABLE SIMULATION WITH JIT CONNECTIVITY OPERATORS
Simulating large-scale brain organizations is notoriously challenging since both computing resources and device memory have near quadratic scaling requirementas the number of neurons increases. For the human brain, which consists of approximately 86 billion neurons, storing the Boolean synaptic connectivity pattern would necessitate nearly 100 terabytes of memory storage. This requirement poses a challenge even for modern supercomputer centers. Moreover, supercomputers are expensive, energy-intensive, and less accessible to broad range of researchers. This poses significant obstacles for researchers seeking to engage in large-scale brain modeling endcavors.
To address this challenge, BrainPy introduces a JIT connectivity algorithm with extremely low sampling complexity (see Appendix C). JIT connectivity is method (Lytton et al. 2008 Carvalho et al. 2020 Knight & Nowotny 2020) used for simulating large-scale brain networks with static synaptic parameters. In this method, synaptic connections are determined by fixed connectivity rule, and synaptic weights remain unchanged during simulation. Instead of storing the synaptic connectivity, the JIT connectivity algorithm regenerates connectivity parameters when a presynaptic
4

---

Published as a conference paper at ICLR 2024
neuron fires. BrainPy provides comprehensive range of JIT connectivity operators crafted for performing matrix-matria multiplication and matrix-vector multiplication within the brainpy. matth jiit conn module. Compared to standard dense sparse matrix multiplication operators, BrainPy's JIT operators enable simulations with neural networks two to three orders of magnitude larger on a single device (Section 4.2 paving easy avenue for large-scale brain simulation for researchers.
4.3 INTEGRATING MODELS IN BRAIN SIMULATION AND AT BY DECOUPLING THE DYNAMICS FROM THE COMMUNICATION
Encapsulating characteristic brain operations as primitive operators is an initial step toward integrating brain simulation and AI models. To further this integration, we need unified perspective for linking major DL components with brain simulation models. The main obstacle hindering this perspective is managing the complex nature of synaptic computation.
BrainPy introduces a unique abstraction that effectively captures and simplifies the complexity of synaptic computations. In BrainPy, a synapse projection between two neuron populations is decomposed into four key components, each encompassing various implementation variants. They are (1) synaptic dynamics, such as alpha, exponential, or dual exponential dynamics; (2) synaptic connectivity patterns, including dense or sparse connections; (3) synaptic output models, such as conductance-basett current-based, or magnesium blocking effects; and (4) synaptic plasticity rules, including short-term plasticity or long-term spike timing dependent plasticity. This decomposition enables BrainPy to offer general implementations for executing diverse synaptic computations. Among the various implementations, two special cases called A1 Anig and AlignPost projections offer superior advantages. Both projections assume homogeneous parameters governing synaptic dynamics within projectior is important note that these projections not approximations. They accurately compute the same dynamics as the original projections while providing new benefits.
CSR
proj bp. Ara ( peerpze, delag-delay,
Presynaptic current
con.iinnin pro.num, post num, postsgoss.
Presynaptic potential
V E 3
Synaptic Communications
Dense
Presynaptic
z E 1
by dy.d Pra prespro, del aalayadalag sysnsye, comm-bp Lincur(pre num. past num. g_maz). rustoas, post-jost,
spiiees
Mis M2
Synaptic E 8
Mea Meaz
Convolution
delays
W gt
graf bp. PrajAtsgaPeagge bolay-delay, оya- commebp, ann, ConvaDEs, 10, osteoss, post=poat,
Synaptic
vanables
mm
Figure 2: Synaptic projections in BrainPy. (A) The Al Lignrre and A1 ignPost 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.
Minimal device memory consumption. In general, a projection requires storing and updating pman synaptic variables, where and n are the numbers of pre- and post-synaptic neurons, and p is the connection probability, However, Al AlignPre and Al Aligt projections only require m and n variables, respectively, aligning with the dimensions of the pre- and post-synaptic neuron groups (please refer to Appendix D for the reduction details and Figure S7 for the computing workflow). Another aspect that showcases the memory efficiency of Al AlignPre and Post projections is their ability to automatically merge duplicate synapse variable creation and updating across multiple projections. Al AlingPre is particularly suitable for handling one-to-many connections since it keeps a trace of synaptic dynamics induced by pre-synaptic neurons (Figure S7CC) This trace can be shared and reused across multiple post-synaptic groups if the synaptic parameters
(Figure[S7c).
'We utilize the exponential synapse model to exemplify the homogeneity. The dynamics of the model is described by ÿ -9/T +,, 6(t - )²), where g the conductance, 7 the time constant, f5P the pre-synaptic spike. We consider the projection homogencous if the value of remains consistent across all synapses.
5

---

Published as a conference paper at ICLR 2024
(typically, time constants) are the same across these projections. On the other hand, the AlignPost projection is highly advantageous for many-to-one connections (Figure (S7D). This is because all synaptic interactions with identical time constants can be converged into single trace of variables for modeling the postsynaptic currents. We showcase the advantages of this automatic synaptic merging by constructing large-scale spiking network model consisting of 30 brain areas, inspired by the work of Joglekar et al. (2017). Our findings reveal that this technique not only decreases device memory usage but also significantly reduces compilation and simulation time (see Appendix F.
Decoupling brain dynamics from its communication. Furthermore, the A1 inire and A1 Al gnPost projections offer a novel perspective on the incorporation of conventional DI. com- ponents within brain simulation models. As depicted in Figure 2 there is distinct separation between the dynamics and the communication between these dynamics. The model dynamics align precisely with the dimensional of pre- and post-synaptic neuron groups, exhibiting strong element- wise and memory-intensive properties. The communication between pre- and post-synaptic groups is facilitated by a communication matrix. Standard brain models implement such communication via sparse matrices, while DL models like linear transformations, convolutions, and normalizations can also serve alternative communication mechanisms for propagating brain dynamics (Figure 3).
4.4 FLEXIBLE MODELING WITH A MULTI-SCALE MODEL BUILDING INTERFACE
A
Potassium Channel
HH Neuron
C
E/i Network
D
Visual System
HH
Na
E
IKibp dyp. Io
class IiC Condien
tnst (welf, min T):
def (self, rum):
supert). inst (siec)
Inct (xtow)
supertl inde (:
Jrit
self Laman, solf-e ELeas, soll-p Kerseblei integrat seln.dp, "sipttierr"
setr IK size, 0.
self.N HC nuaLE nu
Dietaara 1289, (93) raagelnusl]
Sclr-Jka ..o- - by dyn THAD011
Esporentiall
Siit-NinnuLLE solt.N.
selr-T-
dye-TLi
in rangelrur))
seit.1 Expotentiali
for in rangelnanl:
def codetelsett, Vh
selt.Winnas_i Winas_Ell selt.N. p=0.02, E--50.)
selrip-value Setograte
connocernis, 11) self. conss appent/pros)
Figure 3: Multi-scale model building interface of BrainPy. Here bp referred to brainpy package.
Brain dynamics models have key features of modularity, multi-scale organization, and hierarchy (Meunier et al. 2010). To match these characteristics, BrainPy implements modular, composable, and flexible programming interface centered around the bra trainpy. Dynami ca Syssteem class. The key idea underlying this multi-scale model-building paradigm is that models any level of resolution can be defined as brainpy. Dynami ca System classes, and higher-level models (e.g.. networks or systems) can be constructed by hierarchically combining lower-level models (e.g., ion channels or neurons). Figure presents an illustrative example of hierarchical model composition. It demonstrates [3]
the recursive construction of a model, progressing from channels (Figure3) to neurons (Figure 33), networks (Figure 3P)) and finally systems (Figure 3P). It is important to note that while tools like NetPyNE (Dura-Bernal et al.] (2019) also enable hierarchical composition, BrainPy clearly distinguishes it through its unique flexibility and customizability. Specifically, BrainPy allows users to flexibly control the depth of composition based on their specific requirements, ensuring seamless alignment with the aforementioned brain characteristics. Additionally. for user convenience, BrainPy offers a wide range of commonly used models such as channels, neurons, synapses, populations, and networks, serving building blocks simplify the construction of large-scale models.
4.5 OBJECT-ORIENTED JIT COMPILATION
Brain dynamics models are intrinsically memory-intensive For example, the computation within the classical leaky inngrate-andffin (LIF) neuron primarily consists of element- wise operations such as addition, multiplication, and division. In contrast to DNN models that are typically filled with computation-intensiv operators, the memory-intessive nature of brain dynamics models poses

---

Published as a conference paper at ICLR 2024
significant challenges for efficient simulation using native Python. To overcome this, BrainPy heavily leverages JAX (Frostig et al.] 2018) and XLA (Artemev et al. 2022) for JIT compilation. JIT compilation executes models outside of the Python interpreter and optimizes cemmory-intensvvy operators by automatic kernel fusion. 'his fusion technique effectively reduces off-chip memory access, kernel launching overhead, and CPU-GPU switching delays, making XIA ideal choice for compiling brain dynamics models. Particularly, BrainPy integrates a collection of object-oriented JIT transformations into its multi-scale model-building interface. These transformations JAX's ' implementation of funetion-oriented JIT interface, as detailed in Appendix G By leveraging these transformations, any multi-scale BrainPy model can be effortlessly converted into an optimized XLA process, compatible with CPU, GPU, and TPU platforms.
built upon
5 DEMONSTRATIONS
In this section, we showcase BrainPy's efficiency and scalability in brain simulation and BIC tasks. We also highlight its differentiable simulation capability by training biologically plausible spiking network model on working memory tasks.
5.1
EFFICIENCY OF BRAINPY
A
Speed comparison on CPU
Speed comparison on GPU
103
10a
-8 JAX dense
Py lorch dense op
10a
10a
PyTorch sparse
Simulation
BreinPy 00/1000
BralnPy os/100 Hz
101
BrainPy oo/10 Hz
100
0.5
1.0
1.5
2.0
2.5
3.0
0.5
1.0
1.5
2.0
2.5
3.0
Matrix size
1o4
Matrix size
1c4
C
COBA LIF Network
D
COBA HH Network
Simulation time [s]
10a
NEST
ANNardhy
10²
CPU
10a
DindsNeet GPU
10
BrainPy CPU
BranPy GPU
10
1.0
neurons
neurans
185
Figure 4: Event-driven operators in BrainPy enable the efficient running of brain simulation models. (A, B) Speed comparison between different operators that perform the matrix-vecte multiplication for synaptic computation on both CPU (A) and GPU (B) devices. (C, D) Speed comparison of different brain simulators when simulating the COBA-LIF (C) and COBA-HH networks (D).
We first showcase the performance of our event-driven operators (refer to Section 4.1 by using brainpy. ..a.h.h..ennm. rmv operator, which is used to compute y Mv, where M mm is the synaptic connectivity, E IR" the presynaptic spikes, and y E 1TMh the postsynaptic currents. Unlike the alternative dense matrix-vector multiplication, it takes advantage of the CSR representation to store M. Different from the corresponding sparse operator, it makes full use of the event property of the and computes only at positions corresponding the spiking event (see Appendix B). Therefore, brainpy. math. event. csrnv is capable of significantly reducing temporal and spatial costs associated with synaptic computations. esrmv with the corresponding sparse and dense operators in JAX and PyTorch. Each operator performs the synaptic computation in (10,000 times with the time step 0.1 ms), where we randomly sample spiking events in so that presynaptic neurons can fire with frequencies at 10 Hz, 100 Hz, and 1000 Hz. The comparison results show that the event-driven operator achieves two to five orders of magnitude faster than other operators both CPU and GPU devices (Figure 4A and Figure 4B). Moreover, with the lower firing frequency, the event-driven operator runs faster. This is in contrast to its counterparts whose computing times are independent of the number of incoming spiking events.
practice, we compare brainpy mathh math. event
7

---

Published as a conference paper at ICLR 2024
To verify the efficiency of BrainPy on realistic brain simulation models, further compare BrainPy with several mainstream brain simulation frameworks, including NEURON (Hines Carnevale 1997), NEST (Gewaltig & Diesmann 2007), Brian2 (Stimberg et al. 2019), ANNArchy (Vitay et al. 2015), and BindsNet (Hazan et al. (018)) We use El balance networks with the LIF neuron (L.e. the COBA-LIF model (Vogels & Abbott 2005)) and the Hodgkin-Huxley neuron (i.e., the COBA-HH model (Brette et al.) (007g) as benchmarks. The E/I balanced network typically exhibits sparse and event-driven properties, and highly suitable to apply brainpy .at.hhevent........mm... . cs rmv for its computation. As shown in Figure # and Figure 4D, BrainPy shows the best performance on the COBA-LIF model, and such speed superiority becomes more pronounced in the COBA-HH network.
5.2
SCALABILITY OF BRAINPY
We evaluate the scalability of BrainPy by focusing on our JIT connectivity operators (refer to Sec- tion 4.2) We investigate the memory usage and execution speed of dense, sparse, and JIT connectivity operator brainpy math. jitconn. mv mv_prrbb,boo_o_o. uni.uniioo form when performing the matrix-vector multi- plication y = Jv. Here, J represents the connection matrix with a connectivity probability of p. and the weights at nonzero positions are sampled from normal distribution N N(x, o²) using the seed J is stored as a dense matrix in the dense operator, a CSR sparse matrix in the sparse operator, and four scalars (p, μl, G, 8) in the JIT connectivity operator. Our experiments reveal that as the matrix size increases, the JIT connectivity operator maintains nearly constant memory consumption (Figure 55A) and executes with one to two orders of magnitude greater speed compared to the corresponding dense and sparse operators (Figure 5B). This emphasizes the potential for performing large-scale brain simulations on limited computing devices.
B
Memory comparison
Speed comparison
El Network
Memory Usage (MD)
10
Sinulitionn time [s]
10
Time Spaco
Sp
0091 op
10
Time 88
M Space [Ca)
10a
10a
Spane B7
10-1
own OP
10a
10a
10
10a
105 Matrix size
10
D
Matrix size
E
F
Nelwork aba
1es
KTH
MNIST
Network speed comparison
0.94
Simusction time [z]
10a
Acorrey
0.92
Accuracy
10
0.90
Dcroc Sparsils dp
0
2
Reservoir size
1e4
Reservoir size
164
Reservoir siza
164
Figure 5: JIT connectivity operators enable large-scale brain dynamics modeling. (A, B) The memory usage (A) and speed (B) comparison between BrainPy's JIT connectivity operator with JAX's sparse and dense operators for performing the matrix multiplication.
(C) Scaling up the COBA-LIF network with the JIT connectivity operator. (D, E) The empirical relationship between the classification performance and the reservoir size using KTH (D) and MNIST (E) datasets. (F) The inference speed
comparison of the reservoir model using the dense, sparse, and JIT connectivity operators.
To demonstrate its utility in realistic brain simulation, we implement a large-scale version of the aforementioned COBA-LIF network (Vogels & Abbott 2005) using the proposed JIT connectivity operator (Appendix N). We have successfully scaled up this El balance network model to over million neurons and 80 incoming synapses per neuron on single GPU device. The memory and computing time scale linearly with network size (Figure CC Figure S16B, and Figure (ST6C).
Furthermore, these large-scale operators can be applied in brain-inspired reservoir models (Lukoše) vičius 2012) as their input and recurrent weights are fixed during training, which aligns well with the capabilities of JIT connectivity operators. To assess its performance, we conducted experiments to scale up the reservoir size (see Appendix H.. Initially, we evaluate the model on the KTH dataset (An) tonik al. 2019), and find that the model's performance improves constantly with an increase of the reservoir size (Figure 5D). In a previous study (Antonik et al., (2019), a reservoir with only 16,384 hidden units achieved testing accuracy of 91.3% In contrast, implementation with JIT
8

---

Published as a conference paper at ICLR 2024
connectivity operators allows easy scaling up to reservoir with 30,000 nodes, resulting in superior accuracy of up to 94.4%. We also verified this scaling experiment using the MNIST dataset. When the reservoir layer size was set to 50,000 nodes, the network achieved an accuracy of 98.9% (Figure 5J,, on par with the state-of-art machine learning algorithms. Additionally, the reservoir model using the JIT connectivity operator is twice fast during inference compared to the sparse implementation. It also performs better than the dense implementation when the reservoir model has over 10,000 nodes (Figure FT
5.3 TRAINING FUNCTIONAL BRAIN DYNAMICS MODELS
A
Fixation
Sample
Delay
Test
E
(mu))
AMPA
0.0
GABA
-0.1
500 ms
500 ms
1000 mc
500 ms
0.0
0.5
1.9
0.0
0.5
1,
B
F
(ms)
(ms)
Neuron Index
100
Matchh sample
100
Non-match sample
50
Loss
Time (ms)
1e5
Time (ms)
ted
Potental (mV)
Adaptive splking
Bursting
Neuron index
-2
2551
5000
(m76)
1000
250
Tlme (mS)
520
1000
Time (ms)
2000 2009
Figure 6: BrainPy helps to train a data-driven, biologically realistic SNN model for performing working memory task. (A) The delayed
working memory task. (B) Two examples that show the match case (the motion orientations are consistent between sample and test periods) and the non-match case. (C, D) Neuron firing patterns in the model. (E) The AMPA and GABA synapse dynamics in the model. (F) The loss and accuracy of the model when training on the DMS task. (G) An example of the spiking dynamics of the network after training.
Finally, we evaluate the differentiable simulation capability of BrainPy for biological brain networks. Particularly, we apply the propapropag algorithm to train data-driven SNN model of the prefrontal cortex (PFC) (refer to Appendix [I) with a working memory task (Figure 6 and Figure We use the generalized leaky integrate-and-fir (GIF) model to fit the spiking patterns of PFC neurons as observed experiments (Mihalas Niebur 2009). After fitting, most PFC neurons exhibit tonic spiking with spike frequency adaptation (Figure 6C). while the remaining have characteristic bursting features (Figure(DD). Furthermore, we consider the exponential synapse model to model AMPA and GABA synapses (Figure (F). We train the network to solve the delayed match-to-sampii (DMS) task, where the network must indicate whether sequentially presented sample and test stimuli (500 ms sample period; 1000 ms delay period) match exactly (Figure(oA). In this task, the network must maintain the stimulus information for an extended time period (> 1000 ms), then compare the held information to the test stimulus make decision. We find that our biologically grounded GIF network can successfully solve this task requiring long-term dependencies. Within a few epochs, the testing accuracy rapidly increases to nearly 100% (Figure (6.). The post-training spiking dynamics of the network also exhibit comparable patterns to the neural activity of PFC neurons recorded from monkeys performing the same DMS task (see Figure 6G and data in Constantinidis al. (0016).
CONCLUSION
In conclusion, BrainPy provides differential brain simulator that serves as a bridge between the worlds of brain simulators and DL frameworks. By leveraging dedicated optimizations, it enables flexible, efficient, and scalable brain simulation capabilities within the JAX framework. Additionally, its seamless integration with the JAX machine learning ecosystem facilitates the incorporation of Al models into brain simulation. Through the unique combination of these two strengths, BrainPy emerges as a powerful tool for exploring the complexities of the brain and a comprehensive platform for interdisciplinary research between brain simulation and BIC.

---

Published as a conference paper at ICLR 2024
REPRODUCIBILITY STATEMENT
We put all the codes in our supplementary materials for reproducing our experiment results which are mainly demonstrated in Section [ Other results in the appendix are also included in our codes. We provide a README file for our code directory structure and running guidance. For the details of Section
(.2oon5.22) 5.2 and Section a complete description of the model and training parameters are given in Appendix and Appendix
ACKNOWLEDGMENTS
This work supported by the Science and Technology Innovation 2030-Brain Science and Brain- inspired Intelligence Project (No. 2021ZD0200204), We thank Yifeng Gong from Beijing Institute of Technology for his valuable contributions during his internship in Si Wu's lab.
REFERENCES
Martín Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, et al. Tensorflow: system for large-scale machine learning. In 12th USENIX symposium on operating systems design and implementation (OSDI 16), pp. 265-283, 2016.
Denis Alevi, Marcel Stimberg. Henning Sprekeler, Klaus Obermayer, and Moritz. Augustin. Brian2cuda: flexible and efficient simulation of spiking neural network models on gpus. Frontiers in Neuroinformatics, 16:883700, 2022.
Piotr Antonik, Nicolas Marsal, Daniel Brunner, and Damien Rontani. Human action recognition with large-scale brain-inspired photonic computer. Nature Machine Intelligence, pp. 1-8, 2019.
Artem Artemev, Tilman Roeder, and Mark van der Wilk. Memory safe computations with xla compiler. ArXiv, abs/2206. 141122206.1444,,.. 2022.
Igor Babuschkin, Kate Baumli, Alison Bell, Surya Bhupatiraju, Jake Bruce, Peter Buchlovsky, David Budden, Trevor Cai, Aidan Clark, Ivo Danihelka, Claudio Fantacci, Jonathan Godwin, Chris Jones, Tom Hennigan, Matteo Hessel, Steven Kapturowski, Thomas Keck, lurii Kemaev, Michael King. Lena Martens, Hamza Merzic, Vladimir Mikulik, Tamara Norman, John Quan, George Papamakarios, Roman Ring. Francisco Ruiz, Alvaro Sanchez, Rosalia Schneider, Eren Sezener, Stephen Spencer, Srivatsan Srinivasan, Wojeiech Stokowiee, and Fabio Viola. The DeepMind JAX Ecosystem, 2020. URL http::tttttttttie thub.b.ttttut.... www.cc///eeep... nd
Inga Blundell, Romain Brette, Thomas A Cleland, Thomas G Close, Daniel Coca, Andrew Davison, Sandra Diaz-Picr Carlos Fernandez Musoles, Padraig Gleeson, Dan FM Goodman, et al. Code generation computational neuroscience: a review of tools and techniques. Frontiers in neuroinformatics, pp. 68, 2018.
Romain Brette, Michelle Rudolph, Ted Carnevale, Michael Hines, David Beeman, James M Bower, Markus Diesmann, Abigail Morrison, Philip H Goodman, Frederick C Harris, al. Simulation of networks of spiking neurons: review of tools and strategies. Journal of computational neuroscience, 23(3):349-398. 2007.
Nathalie Azevedo Carvalho, Sylvain Contassot- Vivier, Laure Buhry, and Dominique Martinez. Simulation of large scale neural models with event-driven connectivity generation. Frontiers in Neuroinformatics, 14, 2020.
Albert Compte, Nicolas J.-B. Brunel, Patricia S. Goldman-Rakin and J. Wang. Synaptic mech- anisms and network dynamics underlying spatial working memory in cortical network model. Cerebral cortex, 10 9:910-23. 2000.
Christos Constantinidis, Xue- Qi. and Travis Meyer. Single-neuron spike train recordings from macaque prefrontal cortex during a visual working memory task before and after training. CRCNS. org, 2016.
10

---

Published as a conference paper at ICLR 2024
Kael Dai, Sergey L Gratiy, Yazan N Billeh, Richard Xu, Binghuang Cai, Nicholas Cain, Atle E Rimehaug, Alexander J Stasik, Gaute Einevoll, Stefan Mihalas, et al. Brain modeling toolkit: An open source software suite for multiscale modeling of brain circuits. PLOS Computational Biology, 16(1 I) ): 008386, 2020.
Salvador Dura-Berma Benjamin A Suter. Padraig Gleeson, Matteo Cantarelli, Adrian Quintana, Facundo Rodriguez, David Kedziora, George Chadderdon, Cliff Kerr, Samuel A Neymotin, et al. Netpyne, a tool for data-driven multiscale modeling of brain circuits. Elife, 8:e44494, 2019.
Jason K Eshraghian, Max Ward, Emre Neftci, Xinxin Wang, Gregor Lenz, Girish Dwivedi, Mo- hammed Bennamoun, Doo Seok Jeong, and Wei D Lu. Training spiking neural networks using lessons from deep learning. arXiv preprint arXiv:2 119...28944. 2021.
Wei Fang, Yanqi Chen, Jianhao Ding, Ding Chen, Zhaofei Yu, Huihui Zhou, Yonghong Tian, and other contributors. Spikingjelly. httpps://iituut wttp:/ccoooaaammmmmm httph//wwwi244////p..i/n.ng 2020. Accessed: YYYY-MM-DD.
Roy Frostig, Matthew Johnson, and Chris Leary. Compiling machine learning programs via high-level tracing. 2018.
Wulfram Gerstner, Werner M Kistler, Richard Naud, and Liam Paninski. Neuronal dynamics: From single neurons to networks and models of cognition. Cambridge University Press, 2014.
Marc-Oliver Gewaltig and Markus Diesmann. Nest (neural simulation tool). Scholarpedia, 2(4): 1430, 2007.
Padraig Gleeson, Sharon Crook, Robert C Cannon, Michael L Hines, Guy O Billings, Matteo Farinella, Thomas M Morse, Andrew P Davison, Subhasis Ray, Upinder Bhalla, al. Neuroml: language for describing data driven models of neurons and networks with a high degree of biological detail. PLoS computational biology, 1600e10, 2010.
Dan FM Goodman and Romain Brette. Brian: simulator for spiking neural networks in python. Frontiers in neuroinformaties, 2:5, 2008
Joachim Hass, Loreen Hertäg, and Daniel Durstewitz. detailed data-driven network model of prefrontal cortex reproduces key features of in vivo activity. PLDS Computational Biology, 12, 2016.
Hananel Hazan, Daniel J. Saunders, Hassaan Khan, Devdhar Patel, Darpan T. Sanghavi, Hava Siegelmann, and Robert Kozma. Bindsnet: A machine carnine- spiking neural networks library in python. Frontiers in Neuroinformatice, 12:89, 2018. ISSN 1662-5196. doi: 10. 3389/miin..01180000 800899 URL /ttp::/tt////.....ooooooa.oooooa
Eninf. 2018.00083
Jonathan Heek, Anselm Levskaya, Avital Oliver, Marvin Ritter, Bertrand Rondepierre, Andreas Steiner, and Mare van Zee. Flax: A neural network library and ecosystem for JAX, 2020. URL
http://g wttpgbtthhuu..... cow www.oo/..oooccccoooooooooo.ooo.o goog leroolll//e flax
Michael Hines and Nicholas T Carnevale. The neuron simulation environment. Neural computation, 9(6):1179-129) 1997.
Madhura R. Joglekar, Jorge F. Mejías, Guangyu Robert Yang, and Xiao-Jing Wang. Inter-areal balanced amplification enhances signal propagation in a large-scale circuit model the primate cortex. Neuron, 18:222-2344.0 2017. URL https: /api www.steeschccaa.......
Jakob Jordan, Tammo Ippen, Moritz Helias, Itaru Kitayama, Mitsuhisa Sato, Jun Igarashi, Markus Diesmann, and Susanne Kunkel. Extremely scalable spiking neuronal network simulation code: From laptops to exascale computers. Frontiers in Neuroinformatics, 12, 2018. (RLhtttss
//api. semant cscholar wrg/Corpus (D:
w
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. CoRR, 0b2/1412,6990, 2014.
11

---

Published as a conference paper at ICLR 2024
James C. Knight and Thomas Nowotny. Larger gpu-accelerated brain simulations with procedural connectivity. Nature Computational Science, 1:136 142, 2020.
Susanne Kunkel, Tobias C. Potjans, Jochen M. Eppler, Hans Ekkehard Plesser, Abigail Morrison, and Markus Diesmann. Meeting the memory challenges of brain-scale network simulation. Frontiers in Neuroinformatiis, 5, 2012. URL RtLhttts::/ /tapps:/aage /pii semant i icscho ar Str.schoo.ar.oog/ooruus......
6261003
Mantas Lukoševičius. practical guide to applying echo state networks. In Neural Networks, 2012.
William W. Lytton, Ahmet Omurtag, Samuel A. Neymotin, and Michael I. Hines. Just-in-time connectivity for large spiking networks. Neural Computation, 20:2745-2755 2008.
Nikola T Markov, Maria M
AR Ribeiro Gomes, Camille Lamy, Loic Magrou, Julien Vezoli, Pierre Misery, Arnaud Falchier, Rene Quilodran, Marie-Alice Gariel, et al. weighted and directed interareal connectivity matrix for macaque cerebral cortex. Cerebral cortex, 24(1):17-36, 2014.
Adnan Mehonic and A. J. Kenyon. Brain-inspired computing needs master plan. Nature, 604:255 260, 2021.
David Meunier, Renaud Lambiotte, and Edward T Bullmore. Modular and hierarchically modular organization of brain networks. Frontiers in neuroscience, 4:200, 2010.
Stefan Mihalas and Ernst Niebur. A generalized linear integrate-and-fing fire neural model produces diverse spiking behaviors. Neural Computation, 21:704-718, 2009.
Emre O. Neftei, Hesham Mostafa, and Friedemann Zenke. Surrogate gradient learning in spiking neural networks: Bringing the power of gradient- based optimization to spiking neural networks. IEEE Signal Processing Magazine, 16:51-63, 2019.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer. James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
Christian Pehle and Jens Egholm Pedersen. Norse A deep learning library for spiking neural networks, January 2021. mentation: https//norse.ai/ddos//
4422025 Docu-
Philipp Plank, A. Rao, Andreas Wild, and Wolfgang Maass. A long short-term memory for ai
applications spike-based neuromorphic hardware. Nature Machine Intelligence, 4:467 479. 2021.
Tobias C. Potjans and Markus Diesmann. The cell-type specific cortical microcircuit: Relating structure and activity in full-scale spiking network model. Cerebral Cortex (New York, NY), 24: 785 - 806, 2012.
Daniel Rasmussen. Nengodl: Combining deep learning and neuromorphic modelling methods. Neuroinformatics, 17:611 1-628, 2018
Eduardo Ros, Richard R. Carrillo, Eva M. Ortigosa, Boris Barbour, and Rodrigo Agís. Event- driven simulation scheme for spiking neural networks using lookup tables characterize neuronal dynamics. Neural Computation, 18:2959-2993, 2006.
Kaushik Roy, Akhilesh R. Jaiswal, and Priyadarshini Panda. Towards spike-based machine intelli- genee with neuromorphic computing. Nature, 575:607 617, 2019.
H. Francis Song, Guangyu Robert Yang, and Xiao-Jing Wang. Training excitatory-inhibitory recurrent neural networks for cognitive tasks: A simple and flexible framework. PLoS Computational Biology, 12, 2016.
Marcel Stimberg, Dan FM Goodman, Victor Benichoux, and Romain Brette. Equation- oriiettd specification of neural models for simulations. Frontiers in neuroinformatics, 8:6, 2014.
12

---

Published as a conference paper at ICLR 2024
Marcel Stimberg, Romain Brette, and Dan FM Goodman, Brian 2, intuitive and efficient neural simulator. Elife, 8:e47314, 2019.
David Sussillo and L. F. Abbott. Generating coherent patterns of activity from chaotic neural networks. Neuron, 63:544-557, 2009.
Harvey A Swadlow. Efferent neurons and suspected interneurons in s-1 forelimb representation of the awake rabbit: receptive fields and axonal properties. Journal of Neurophysiology, 63(6): 1477-1498, 1990.
Julien Vitay, Helge Ulo Dinkelbach, and Fred H Hamker. Annarchy: code generation approach to neural simulations parallel hardware. Frontiers in neuroinformatics, 9:19, 2015.
Tim Vogels and I.. F. Abbott. Signal propagation and logic gating in networks of integrate-and-fire neurons. The Journal of Neuroscience, 25:10786 - 10795, 2005.
Xiao-Jing Wang. Probabilistic decision making by slow reverberation in cortical circuits. Neuron, 36 (5):955-968, 2002.
Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms. ArXiv, 70s.07707,7777 2017.
Mingqing Xiao, Qingyan Meng. Zongpeng Zhang, D.K. He, and Zhouchen Lin. Online training through time for spiking neural networks. ArXiv, abs/2210.04195,, 2022.
Yi Zeng, Dongeheng Zhao, Feifei Zhao, Guobin Shen, Yiting Dong, Enmeng Lu, Qian Zhang, Yinqian Sun, Qian Liang, Yuxuan Zhao, Zhuoya Zhao, Hongjian Fang, Yuwei Wang, Yang Li, Xin Liu, Chen-Yu Yu Du, Qingqun Kong, Zizhe Ruan, and Weida Bi. Braincog: A spiking neural network based brain-inspired cognitive intelligence engine for brapn-inspired ai and brain simulation. ArXiv, abs/2207.08533, 2022.
13

---

Published as a conference paper at ICLR 2024
A ENVIRONMENT SETTINGS
In this paper, all evaluations and benchmarks were conducted in a Python 3.10 environment, which was installed system running Ubuntu 22.04.2 LTS. For CPU experiments, used Intel(R) Xeon(R) W-2255 CPU @ 3.70GHz, 64GB RAM @ 3200MHz. For GPU experiments, we used the NVIDIA RTX A6000 GPU with CUDA 11.7.
For brain simulation, we compared the running performance of BrainPy with state-of-art brain simulators, including NEURON (Hines & Carnevale, 1997) in version 8.2.0, NEST (Gewaltig & Diesmann 2007) at version 3.6, Brian2 (Stimberg et al. 2019) at version 2.5.1, ANNArchy (Vitay al. 2015) in version 4.7.2, and BindsNet (Hazan et 2018 in version 0.3.2,
For brain-inspired computing, we compared BrainPy with several spiking neural networks (SNN)
training packages, including snnTorch (Eshraghian al. 202T) in version 0.6.1. SpikingJelly (Fang et al. 2020 in version 0.0.0.0.14 14, and Norse (Pehle & Pedersen 202T) in version 1.0.0.
When comparing the performance of dedicated operators in BrainPy with conventional operators
in modern deep learning frameworks, we used PyTorch (Paszke et al. 2019) at version 2.0, and JAX (Frostig et al. 2018) at version 0.4.10.
B EVENT-DRIVEN MATRIX-VECTOR MULTIPLICATION
Event-driven matrix- multiplication y Mv in BrainPy is used for synaptic computation, where v is the presynaptic spikes, M the synaptic connection matrix, and y the postsynaptic current. Specifically, it performs matrix- -vector multiplication event property of the input vector v. Instead of multiplying the entire matrix M by the vector v, which can be wasteful if v has many zero elements, event-driven matrix-vector multiplication in BrainPy only performs multiplications for the non-zero elements of the vector, which are called events. This can reduce the number of operations and memory accesses, and improve the running performance of matrix-vector multiplication.
sparse and efficient way by exploiting the
In particular, BrainPy implements bra atainpy. math. event carmv, which the connection matrix M is stored as the compressed sparse row (CSR) sparse matrix. The computation of this operator requires several parameters:
1. data: The weights of M contain the non-zero elements in the row-major order.
2. indices: The array contains the column indices of the non-zero elements in the matrix M.
3. indptr: The array contains the starting index of each row.
4. events: The presynaptic spiking vector v.
5. shape: A tplle with two integers representing the shape of the matrix M.
brainpy. math, event. rermv makes full use of the event property of v, and computes only at positions where the spike in vs True. The pseudo-code of this operator can be written in Listing SI
def esrmv (data, indices, ianptz, events, outs):
2
for i, event in enumerate (events):
if event:
for in range (idn idnptrli 1]):
outs iindices [i]] datali]
Listing S1: The Python pscudo-code of brainpy math, event. carnv, where data, indice ces, and idrptr corresponds to the matrix M, events indicates the vector v, and outs represents the postsynaptic current
C MATRIX-VECTOR MULTIPLICATION WITH THE JUST-IN-TIME CONNECTIVITY
Synaptic connectivity storage imposes a memory bottleneck for large-scale neuronal network simula- tions, it scales quadratically with the number of neurons. Previous studies have demonstrated that
14

---

Published as a conference paper at ICLR 2024
static synaptic parameters in brain modeling be dynamically regenerated during runtime, thereby circumventing the space costs of connectivity Lytton et al. (2008); Carvalho et al. (2020); Knight & Nowotny (2020).
We investigate the dynamic regeneration of synaptic connectivity information in a matrix- product operation y Mv, where M is the synaptic connection matrix to be regenerated, is the presynaptic spike vector, and the postsynaptic current vector.
A common connectivity schema is the fixed probability connection, where each neuron in the presynaptic population connects to each neuron in the postsynaptic population with a fixed probability p. The postsynaptic targets of any presynaptic neuron can thus be drawn from Bernoulli process with success probability
naive way of drawing from the Bernoulli process is to sample repeatedly from the uniform distribution [0, 1] and create a synapse if the sample is smaller than p. However, this is highly inefficient for sparse connectivity (p
A more efficient way of drawing the sampling is using the geometric distribution Geolp] Knight NNwotnyy 2020). Instead of sampling at every possible connection position, we can sample the distance between two consecutive connection positions, avoiding unnecessary sampling at failed positions. The geometric distribution can be sampled in constant time by inverting the cumulative density function of the corresponding continuous distribution to obtain log(U/U.11) To However, this
sampling method is still costly on CPU and GPU devices due to the existence of log operation.
We propose sample the distance between two consecutive connection positions using the uniform 1
random integers from to 1 1]. The expectation of U11, - is p. ensuring that the sampling
has the desired connection probability of p. In practice, sampling from U[1, 1 -J] is one order of magnitude faster than sampling from Geofp). The corresponding Python pseudo-code of our proposed solution shown in Listing
import math, random
def gitconn_prab_haee homo (events, prcb, weight, seed, outs): random. seed (seed) max max_caist caist math. fitth.eeil(e (2/prob 1)
for event in events:
if event:
post_i = random. randint (1, max_cdist)
while post_i
:
10
outs(post.. += weight
11
post_i += random. randint max_cdist)
The Python pseudo-code of just-in-time connectivity operator, where events indicates the presynaptic spikes v. prob, weight, and seed corresponds to the matrix M, and out represents the postsynaptic current y. The seed parameter guarantees the reprodueibility and consistency of matrix regeneration across multiple invocations of this function. Note here, all nonzero elements in the matrix M have the save value weight.
Listing S2:
D SYNAPTIC PROJECTIONS WITH ALIGNPRE AND AL ALIGNPOST
AlignPre holds true because with homogeneous synaptic parameters, spike train coming from the same presynaptic neuron will lead to the same synaptic dynamics. As result, all synapses originating from the same presynaptic neuron can share a single dynamical variable. Morcover, the AlignPre projection is suitable for all types of synapse models.
In contrast, the applicability of Al AlignPost Post is limited to synapse models that exhibit exponential dynamics composed of exponential components. For the exponential synapse, the conductance of a post-synaptic neuron is updated according to g <- + whenever spike arrives and regardless of which presynaptic neuron emitted this spike (refer to Appendix E for details).
AlignPre is capable of modeling scenarios where the pre-synaptic neuron group projects to multiple downstream post-synaptic neuron groups. This mechanism imposes constraint the size of synapse
15

---

Published as a conference paper at ICLR 2024
B
Presymaphic spikes
R
Presynaptic spikes
delay
debay
variables
E R2
Delay spiles
X
R
Mis M12
Wam
Mas
Synaptic weights
Ma: Maz
Mam
Synaptic weights
Mai
1
Mei M42
Pym
W e ue t-0
Synaptic currents
Postsymaptic neurons
W
E mm
Synaptic currents
Synaptic variables
Postsymaptic neurons
D
post1
pre1
pre
post2
prez
post
Synaptic
post3
pre3
Synaptic
variables
variablles
Figure S7: Overview of Synaptic Projections. (A) The workflow of the AlignPre model. (B) The workflow of the AlignFost model. (C) The Al AlignPre model is suitable for one-to-many connection. (D) The AlignFost model is suitable for many-to-one connection.
variables, ensuring alignment with the number of pre-synaptic neuron groups Consequently, if these projections share the same delay, only single synapse variable is stored for all these projections. The operational sequence of A1 i gn unfolds as follows: Initially, it receives spikes from the pre- synaptic neuron group. By considering the delay times, it retrieves spikes that occur after the specified delay and then proceeds to compute synapse variables through synaptic dynamics. Subsequently, the synapse variables traverse the synaptic communication layer and synaptic output layer, where they are transformed into synaptic currents, aligned in accordance with the number of post-synaptic neuron groups. These synaptic currents are then transmitted to the post-synaptic neuron groups for further computations.
AlignPost. on the other hand, is designed to handle scenarios in which the post-synaptic neuron group receives input from multiple upstream pre-synaptic neuron groups. In this case, the size of the synaptic variables adjusted to match the number of post-synaptic neuron groups. The operational sequence of Al AlignPost closely resembles that of AlignFre, re, with the notable difference being that the calculation of synaptic variables occurs after the synaptic communication layer. Specifically, -nssyyaaaicsppiiaa spikes with synaptic weights update synaptic variables. Subsequently, these updated synaptic variables undergo computation through the synaptic output layer, resulting in synaptic currents that are then conveyed to the post-synaptic neurons.
Al AlignPost combines delayed pre-synaptic
E EXPONENTIAL SYNAPSE MODEL ENABLES THE AL GNPOST PROJECTION
For the exponential synapse, the conductance of post-synaptic neuron is governed by
sp
g(()
Σ exp
(I)
where n is the total number of spikes the post-synaptic neuron received at the current time t, 17, the spiking time of the received j-th spike, and the synaptic time constant.
Eq can be rewritten as a differential equation:
g(()
9
5(t
k.k))
(2)
16

---

Published as a conference paper at ICLR 2024
where i is the index of the connected pre-synaptic neuron, and f, is the time of k-th spike of the pre-synaptic neuron i.
In the discrete version, Eq equivalent to the following equations:
g exp(-At/r)g(( - At).
(3)
and whenever pre-synaptic spike arrives, gt)) undergoes update according to:
g(() + g(() +
Therefore, the exponential dynamics property enables us to track and record information using a single variable for each post-synaptic neuron.
F MULTI-AREA SPIKING NEURAL NETWORK
We build spiking network model which inspired by Joglekar al. (2017) with 30 brain areas. Simulations are performed using a network of leaky integrate-and-Hn neurons, with the local circuit and long-range connectivity structure derived from (Markov al.] 2014). Each of the five areas consists of 4000 neurons, with 3200 excitatory and 800 inhibitory neurons.
For each neuron, the membrane potential dynamies are described by the following equations:
dV
dt = (V - Vreat) + RG(()
(5)
when V(t)= > Vea, the neuron generates spike and (t) set to Vreset- V is the membrane potential, Vrest is the resting membrane potential, Vreset is the reset membrane potential, Vo, is the spike threshold, 7 is the time constant, R 12 is the resistance, and G(t) the time-variant synaptic inputs.
For the synapse, we use conductance- based synaptic interactions. Particularly, G(t) is given by:
G(t) G(t)=- = Σ (x (V, -,)
where V; is the membrane potential of neuron i. The synaptic conductance from neuron j to neuron i is represented by (t), while E, signifies the reversal potential of that synapse. For excitatory synapses, E, was set to mV, whereas for inhibitory synapses, was -80 mV.
The dynamics of the synaptic conductance is given by:
dgzs dt
9jz Tdecay
Σ 8(t ((t-tt),
Imax
where th is the spiking time of the presynaptic spike. Whenever a spike occurred in neuron j. the synaptic conductance fji experienced an immediate increase by fixed amount fmax- Subsequently. the conductance Рjz decayed exponentially with time constant of Tderay ms for excitation and Tdecay 10 ms for inhibition.
The connection density is set according to experimental connectivity data Markov et al. 2014). The inter-areal connectivity is measured as weight index, called the extrinsic fraction of labeled neurons (Markov et al. 2014). The intra-area connectivity is set according to the setting in a standard El balance network (Brette et al. 2007).
Moreover, we introduce distance-dependent inter-areal synaptic delays by assuming a conduction velocity of 3.5 m/sec (Swadlow) 1990) and using distance matrix based experimentally measured wiring distances across areas (Markov et al.) 2014).
We build this multi-area model using the synaptic projection with and without AlignPost. Then measure the memory usage after constructing the model, the compilation time during JIT compilation. and the execution time when simulation the model. The experimental data is avaliable in Figure
17

---

Published as a conference paper at ICLR 2024
Memory
With
Merging
Memory
B
Execution Time
With Merging
Without Merging
20
30
Time [s]
Compilation Time
With Merging
Merging
Figure S8: The evaluation of memory usage, model compilation time, and model execution time for the multi-are spiking neural network model using synaptic projection with and without the automatic merging of the AlignPost projection.
G 00 TRANSFORMATION COMPARISON
The JIT compilation interface of JAX/XLA is designed to work with pure Python functions, which is compatible with the modular and composable programming interface in BrainPy for building brain dynamics models (see Section 4.4) To bridge this gap, provide bra brainny.... ieiinnymmm. math. jit for automatic object-oriented JIT transformation.
In this object-oriented transformation, variables that will change during execution should be declared as brainpy .ath. Variable (otherwise, they will be treated as constants and compiled into the transformed models). Then, single line calling model_trans forned brainpy .ath.. jit (model) will transform the model object into a function compatible with JAX's functional JIT interface, before final compilation into optimized XLA process model trans formed. Through this streamlined approach, BrainPy 's object-oriented JIT compilation furnishes a robust infrastructure maximize the performance of diverse brain dynamics models.
18

---

Published as a conference paper at ICLR 2024
Although Flax and Haiku provide functionality for managing state variables, they do not focus specifically on the needs of brain dynamics modeling. In Flax, the initialization of state variables is mixed into the computation, making it hard separate variables from computing logic. In Haiku, state variables have complex syntax for declaration and calculation, which increases the difficulties for users these time-dependent variables in complex brain dynamical systems.
The following codes compares the 00 programming of Haiku, Flax, and BrainPy, and it highlight how straightforward the BrainPy is:
import haiku as hk import jax.numpy numpy as jnp import random jr
def stateful_f()
counter ones)
shape- [], dtype=jnp.inds int32, init=jnp.
ut..et__tate counter
hnt.without_appl__mgg (hk. (ttttsstttt(ttttteool____)
10 params, state (mageir..eeeeee.e.eellll for i in range (3)
state appply (params, state)
13 print After li 11 iterations, State: (state)') Listing S3: The Haiku's object-oriented JIT interface.
from flax Import linen as nn import jax. numpy as jnp
import random as jr
class F (nn, Medule)
compact def __f__cllo ______aaa___________________ (self):
counter self. verlll.ariaaleeeeee stsate",, 'counter', Lambda inp, ones (s, jnp .int32), ())
counter.
10
11 F()
12 variables = inini PRNGKey )1))
for in range (3):
14 variables apply (variables, mutable= state'
15 print After (1 1) iterations, state: (dict (variables ("state")) 10)
Listing $4: The Flax's object-oriented JIT interface.
import brainpy as bp import brainpy.me math as bm
class _ (self):
self.counter = bm. Variable cnes (), dtype=bm, int
9bm. @bm.cls_iie cls_jit
def calllaall_ssll__ (self):
10
self.counter +=
1
12 F()
for in range (3)
14
f()
15
print (E'Aiter li 1) iterations, State: I.. counter)))
Listing S5: The BrainPy's object-oriented JIT interface.
The Pythonic approach of BrainPy makes the object-oriented JIT much more accessible for neurosci- entists without deep MI. expertise and lowers the barriers compared to Flax or Haiku.
19

---

Published as a conference paper at ICLR 2024
H
THE RESERVOIR COMPUTING MODEL
Reservoir model. The dynamics of the reservoir model used here is given by (Lukoševičius 2012):
x(t) (t) (1 a. x(t 1) + a. f(W,, in u(t) + W rec x(t - 1)),
(8)
y()) W out x(t)
where x(t) is the reservoir state, y(t) the readout value, Win and Wrec are input and recurrent connection matrices, respectively, out the readout weight matrix which can be trained by either offline learning or online learning algorithms, a E (0, 1] the leaky rate, u(t) the input at time step t, and f the nonlinear activation function.
Model implementation. Win and Wrec rec are fixed and randomly initialized, and they are highly suitable to be implemented with the just-in-time connectivity operators.
Since Win is usually initialized with the uniform distribution U[-s, s, s], we implement the input oper- ation Wo. (() with brainpy. math. jincoonnnnnn..m monn.my ww hwv_prob_urg uni form (vector, w_low, w w_high conn hann_rrb,, seed), where vector corresponds to the input u(t), 1ow corresponds to w high corresponds to 8, conn prob corresponds to the connection probability of the input matrix W in: and seed the random seed controlling the reproducibility of the input matrix.
W rec is usually initialized with the normal distribution N(0, o) with a desirable spectral radius have the relationship of ity of the matrix, m the matrix size, and p the spectral radius of the recurrent weight W ree- Therefore, implement the recurrent matrix operation W Weee x(1 - 1) with brainpy. math. jit jitconn. mnv_prbb_ormee
P. Usually,
p/ / * p, where p is the connection probabil-
(vector, w_mu, w_sigma, conn_prob, seed), where vector corresponds the reservoir state x(t 1), w mu O, w L.s sigma corresponds o, conn prob corresponds to the connection probability of the recurrent matrix Wrec. and seed the random seed controlling the reproducibility of the recurrent matrix.
Training methods. The training objective of reservoir models is to find the optimal W that minimizes the square error between y() and (1)) The common way to learn the linear output weight W out is using the ridge regression algorithm (Lukoševičius) 2012). However, ridge regression an offline learning method in which the parameters are learned given all samples of data. When training reservoir models with a large amount of samples, ridge regression usually is halted by the limited memory space of devices. Therefore, we use the FORCE learning method (Sussillo & Abbott 2009) to train the model. Contrary ridge regression, FORCE learning is an online learning algorithm that learns the linear readout weight using only local information in time. This learning mechanism can update the parameters of model one sample of data a time, which can significantly reduce computational costs and enable training large-scale reservoir to be possible.
Table S1: The parameter table of the reservoir model on different datasets.
Parameter
KTH dataset
MNIST dataset
W in connection probability
[0.01, 0.005]
0.1
W ree connection probability
[0.001, 0.0002, 0.0001]
[0.1, 0.01]
Distribution of W in
Uniform distribution
Uniform distribution
Distribution of W Wrec
Normal distribution
Normal distribution
Spectral radius p
1.0
1.3
Input scaling
0.1
0.3
Leaky rate a
0.9
0.6
Number of training epoch
10
5
FORCE learning rate
0.1
0.1
Experiment details. We conducted several experiments to investigate the performance of large- scale reservoir models different datasets. First, we evaluate the performance of the model on the KTH dataset (Antonik et al. 2019),, a widely used benchmark dataset for action recognition in computer vision research. This dataset contains spatiotemporal patterns that be captured by
20

---

Published as a conference paper at ICLR 2024
reservoir models. Then evaluate the reservoir model the MNIST dataset, which is static image dataset without temporal information. All parameter selections for classifying two datasets are shown in Table ST
We set the connection probabilities of Win Win and W rec as follows: For the KTH dataset, Win has connection probability of 0.01 for size [2000, 4000, 8000, 10000, 20000] and 0.005 for size [30000]. W ree has a connection probability of0.001 for size E [2000, 4000, 8000], 0.0002 for size [10000], and 0.0001 for size E [20000, 30000]. For the MNIST dataset, W in has connection probability of 0.1 for all reservoir sizes, and W has connection probability of 0.1 for size E [2000, 4000, 8000, 10000], and 0.01 for size E [20000, 30000, 40000, 50000).
RECURRENT NEURAL NETWORKS FOR PERFORMING THE DMS TASK
Network structure. The architecture of recurrent networks used here (both the rate- and spiking- based models) shown in Figure S9 where the recurrent layer consists of recurrent units that receive and process the time-varying inputs from the input layer, and generate the desired time-vayying outputs. The imput layer encodes the sensory information relevant to the task, while the readout layer produces decision in terms of an abstract decision variable.
Input Layer
Recurrent Layer
Readout Layer
Figure The schematic diagram of the recurrent neural network model for implementing working memory task.
The input layer. The input layer consists of N motion direction-tunel input neurons (for the spiking network, N = 100; for the rate model, = 24). The tuning of the motion direction-selective neurons followed von Mises distribution, such that the firing rate activity of the input neuron during the stimulus period was
u A exp cos (f-6%) (pret))
(10)
where 0 the direction of the stimulus, opee is the preferred direction of input neuron i, x was set 2.0, and the maximum firing rate and was 40 Hz. Furthermore, for all periods, input neurons keep a constant background firing rate Tbg. Here, Pog was set to Hz.
In the main text, Figure 3B depicts the spiking activities of the input layer in the spiking- based model. Unlike the spiking encoded input, which produces discrete spikes, the rate-based input layer generates continuous firing rate values as data. Figure STO demonstrates the cases of match (when the stimuli during the sample and test periods are identical) and non-match (when they are different).
The recurrent layer in the rate model. In the rate-based model, the recurrent layer consists of 80 excitatory neurons and inhibitory neurons. The firing activity r of the rate-based recurrent neurons was modeled with the following dynamical equation (Song al. 2016):
dr di
(Wreer win u + rec V2Toroas).
(11)
21

---

Published as a conference paper at ICLR 2024
Match sample
Non-match sample
Neuron index
20
3
2
1
0.0
0.0
0.5
1.0
1.5
2.0
2.5
0.5
1.0
1.5
2.0
2.5
Time (ms)
1e1
Time (ms)
1e1
Figure S10: The rate version of two examples that show the match case and the non-match case. For the spiking version of the input, please refer to Figure 3 in the main text.
where 7 is the neuron time constant, f(*) represents the activation function, Wrec and Win are the recurrent synaptic weights and input synaptic weights, respectively, bree is bias term, Ç is independent Gaussian white noise with zero mean and unit variance applied to all recurrent neurons and oee is the noise intensity. The activation function used in this study is [(x) max(0,
In practice, the simulation of the Eq. 11 was performed with the Euler integration method:
T: = (1 -1- di tf W ri-a + Win us ree
(12)
where a = dt/r,, the numerical integration step, (0, 1) denotes the standard normal distribution.
To maintain the separation of excitatory and inhibitory neurons, decomposed the recurrent weight matrix Wreee as the product between a trainable non-negative matrix Wree, and a fixed diagonal matrix D (Song et al. 2016):
Wree
Wreat D
(13)
1
D =
(14)
The recurrent layer in the spiking model. For the spiking model, the recurrent cell was im- plemented with the spiking neuron model. In this study, the spiking neuron is modified from the generalized integrate-and-ff neuron model (Mihalas & Niebur (2099.. In particular, this model has two internal currents, one fast and one slow. Its dynamic behavior is given by
dl
dt I1, dl -I2.
fast internal current
(15)
slow internal current
(16)
dt
dV
TV dt
V + Vrest + R(I, + I2 + Iext),
membrane potential
(17)
When V of i-th neuron meets VA, the modified GIF model fires:
I + A
(18)
I + I3 +
(19)
V Vrests
(20)
where 7n denotes the time constant of the fast internal current, 7/2 the time constant of the slow internal current, TV the time constant of membrane potential, R the resistance, Text the external input, Vrest the resting potential, and A, and A₂ the milke-triggered currents.
For matching the statistical data of electrophysilogicall experiments in Hass et al, (2016), we assigned 80 excitatory neurons and 20 inhibitory neurons. Inhibitory neurons exhibit bursting firing pattern,
22

---

Published as a conference paper at ICLR 2024
while excitatory neurons show adaptive spikings. We set A1 8.0 for inhibitory neurons, A1 = for excitatory neurons, A2 = -0.6, Th1 10 ms, 712 was sampled from a uniform distribution U[100,3 3000 ms.
The external current lex in the network was modeled through the exponential synapse model, whose dynamics is given by:
d
L Tdecay
+
Σ 2i d'1 Σ W Willll
(21)
where 7decay is the time constant of the synaptic state, t the time of the presynaptic spike, Wree the recurrent weight, Win the input to recurrent weight, and d'd the synaptic delay. Here, d' ms. Tdecay 100 ms was chosen according to the previous literature (Compte et al. 2000).
To inspect the computational graph of the modified GIF network, we also give the discrete description of the model:
[t + At] = where = 1, A,, 07, (1])
(22)
2 At] au [2 + Azz(t]
(23)
Texa It + At] a a Lealt] + Wre + W
(24)
Vit + At] av V(]+ + (Vrest + R(L [t + Ar + 12 Ad) + Text (t))
(25)
z|t + At] = where Vit At] VeL 1,0)
(26)
where 2 is the spiking state, an e 71 -1 at
T/g
01/
Vin
At
av
At
and O' =
Fdecay
Similar to the rate-based model, we decomposed the recurrent weight matrix Wree as the product
between a trainable non-negative matrix Wrree++ and fixed diagonal matrix D (Sog al. 2016):
W
Wree, D
(27)
For the forward spike function, we use the Heaviside function to generate the spike:
spike(x) = HV[[t] - Van) = H(x),
(28)
where r is used to represent V[[]- Veh.
To make the non-differentiable spiking activation compatible with the back-propagation algorithm, we considered a surrogate gradient:
spike' (a) = ReLU (width [x]))
(29)
where width 1.0, and a 0.3. a the parameter that controls the altitude of the gradient, and width is the parameter that controls the width of the gradient. The shape of this function is shown in Figure SIT
0=05, wisth=1
0-0.5, w
175
width=1
1 wisth=2.0
125
0.75
025
Figure S11: The shape of the ReLU gradient function depends on its parameters.
23

---

Published as a conference paper at ICLR 2024
The readout layer. The recurrent neurons project linearly to the output neurons. For the rate-based model, the readout layer is given by:
Wo" r[t] + bout
(30)
where Wrout are the synaptic weights between the recurrent and output neurons, and gont the bias. For the spiking network, the output neuron is the leaky neuron, whose dynamics given by:
dy Tout dt y Wou + boue
(31)
where Tout is the time constant of the output neuron, Wout the synaptic weights between the recurrent and output neurons, and your the bias. In the discrete description, the output dynamics is written as:
yit
At] (o + (Wousz[t]] Att
(32)
where Cout Tout At
Weight initialization. Initial input, recurrent, and readout weights were drawn from a Gaussian distribution W V N(0, 1), where nin is the number of afferent neurons, N (0, 1) is the zero-mean unit-variance Gaussian distribution, and the weight scale. For inhibitory neurons, 8 0.2; for excitatory neurons, = 0.05.
A
Rate-based network accuracy
1.0
Accuracy
0.8
0.8
0
25
50
75
100 Epochs
125
150
175
200
B
Recurrent Activity
Excitatory Inhibitory
30
20
10
500
1000
1500
2000
2500
Time (ms)
Reedout 0 Readout
Output Activity
20
-10
500
1000
1500
2000
2500
Time (ms)
Figure S12: The rate-based recurrent model for implementing a DMS working memory task. (A) The training accuracy the rate-based recurrent neural network for performing DSM working memory task. (B) The recurrent activity during the network receiving a non-match case stimulus. (C) The output activity of the network when receiving a non-match case stimulus.
Training methods. The training performed using the BPTT algorithm. The integration time step At is 1 ms for the spiking neural network, while At is 100 ms for the rate-based model. The Adam optimizer (Kingma& & Ba 2014) was used for computing the gradient- based optimization. The goal of the training was to minimize the cross-entropy between the output activity and the target output during the test period. For the spiking network, add additional voltage regularization loss that
24

---

Published as a conference paper at ICLR 2024
penalizes voltages that fall significantly outside the support of the surrogate gradient function PPlakk ct al. 2021:
Lvoltage = Xo >> (ReL 0.4))² + (ReLU(-VV(d) -2.0))²
(33)
where X is the strength of the voltage regularization.
The performance of the rate-based model. In contrast to the spiking- -base model, the rate model employed here exhibits a slow convergence of training (see Figure (SI2A). It requires dozens of epochs of training achieve high accuracy. Figure SS2 and C also show the recurrent and output activities after the rate-based network receives a non-match case stimulus. Same as in the spiking network, the neural activity of the rate model during the delay period (1000- 2000 ms) maintains high firing rate.
EVALUATION OF BRAINPY ON TRAINING MACHINE-LEARNIA ORIENTED SPIKING NEURAL NETWORKS
To evaluate the performance of BrainPy BIC models, conduct comparative analysis with several popular SNN training frameworks, including snnTorch (Eshraghian et al., 2021). Norse (Pehle & Pedersen 2021), and SpikingJelly (Fang al.] 2020). Note here comparisons only carried out using densely connected spiking neural networks, which are predominantly utilized in existing BIC models.
We first use a simple three-layer SNN model (Neftci et al. 2019), where the input layer delivers synaptic currents the hidden LIF layer with exponential dynamics, then the output layer readouts the hidden dynamics with an exponential synapse. The connections between each layer are dense layers. The training was performed using the BPTT algorithm the Fashion-MNIST dataset (Xiao et al. 2017). The comparison result of the training speed per epoch is presented in Table S22 We find that BrainPy exhibits superior performances on both CPU and GPU devices. Furthermore, to test the performance of BrainPy large-scale SNN networks, we re-implement VGG SNN network (Xiao) (et al. 2022) which was originally coded PyTTorch+SiikingJlll environment. The evaluation result shown in Table S2] Under the same training setting. BrainPy achieves much better running performance, reduces almost half the time for training on GPUs.
Table S2: Comparison of average training speed per epoch for two SNN models using sonTorch (Eshraghian et al. 2021) SpikingJelly (Fang et al. 2020), Norse (Pehle Pedersen 2021), and BrainPy.
Model
Device
sorch
SpikingJelly
Norse
BrainPy
CPU (Intel)
0.33
49.9 + 10s 53.1 1 0.9
52.2 0.2s 49.6 i 0.3
28.6 + 1.0 17.1 + 0.6
Simple SNN
Neftci et al.
2019
GPU (RTX A6000)
1.s
VGG SNN Xiao et al. 2022
GPU (RTX A6000)
104.0 1.0
50.0 + 1.0
K EVALUATION OF COMPILATION TIME BETWEEN BRAINPY AND BRIAN2
Simulating biological spiking neural networks often involves a time-consuming compilation pro- cess. To evaluate the compilation cost of BrainPy, we conducted systematic comparison with Brian2 (Stimberg et al. 2019).
To assess the compilation time, we utilized two different network models: simple El balance network (Vogels & Abbott 2005) and more complex multi-area network model consisting of 30 interconnected brain areas Joglekar et al., (017).
During the initial run, both BrainPy and Brian2 compile the model. Consequently, measured the compilation time by simulating the respective network for single time step (T' At, where T represents the simulation duration, and A/ is the simulation time step).
25

---

Published as a conference paper at ICLR 2024
For Brian2 CPU, measured the compilation time using the cython backend. To measure GPU compilation time, we utilized Brain2CUDA (Alevi et | 2022).
In the case of the EI balance network, varied the network size by increasing the number of neurons. We measured the compilation time across different network sizes and computing platforms. The comparison results are presented in Figure SS13A. [ than ten times faster than Brian2, regardless of the CPU or GPU device used. Brian2CUDA shows gradual increase in compilation cost on GPU devices, which is consistent with the findings in their original paper. In BrainPy, CPU and GPU exhibit comparable compilation speeds.
BrainPy demonstrates compilation speeds more
A
Compliation
B
Compliation MultiArea Network
10
10338.42 42
10
Compitation Time
Compitation Time [3]
10"
10
BrainPy CPU
BrainPy GPU
Brain2 CPU
Brain2CUDA
10
10
5.34
3.75
5000
10000
Nelwork Size
BrainPy CPU
BrainPy GPU
Brian2 CPU
Figure S13: The compilation time comparison between BrainPy and Brian2. (A) The compilation time measured on an El balance network model (Vogels & Abbott, 2005). (B) The compilation time measured on multi-area large-scale network model (Joglekar al. 2017).
With the multi-area network model, BrainPy exhibits significant speedup compared to Brain2 (Figure(S13) B) this network, we employ jax. vmap to merge multiple projections from one brain area, resulting in a much simpler computational graph. In contrast, although the simulation of the network took only around one minute, the compilation process in Brain2 took over an hour. Note that in this model, the Brain2CUDA failed This result unequivocally demonstrates the superior advantage of BrainPy, which leverages modern just-in-time compilation technology, over traditional brain simulators.
ignore the comparison with Brain2's GPU backend.
The compilation time comparisons between BrainPy and Brian2 on both the El balance network and the more complex multi-area network demonstrate clear advantages of BrainPy's JIT compilation approach. BrainPy exhibited over 10x faster compilation times than Brian2 on both CPU and GPU devices, with the advantage being more pronounced larger networks.
The gradual increase in Brian2's compilation time on GPU devices has been reported before (Alevi al. 2022) and is likely due inefficiencies in translating imperative Python code GPU kernels. In contrast, BrainPy's compilation speed was consistent across CPU and GPU. This shows the power of just-in-time compilation and using computational graph optimizations like jax. vmap to simplify projections between areas.
The extremely long (>1 hour) compilation time for Brian2 on the multi-area network compared the short simulation runtime (~1 minute) highlights a key bottleneck. As researchers continue building larger and more complex brain network models, short compilation times are essential to enable rapid testing and iteration. The over three orders of magnitude speedup shown by BrainPy is thus hugely impactful.
Overall, BrainPy's modern compilation approach enables faster development cycles and more complex networks than previously feasible. As brain models continue increasing in scale and detail, optimized just-in-tiin timtttime compilation will become more and more critical. These benchmark results validate BrainPy's 's advantages compilation efficiency today and its ability to scale up to the brain models of tomorrow.
26

---

Published as a conference paper at ICLR 2024
L SUPPORT FOR THE DISTRIBUTED COMPUTATION IN BRAINPY
BrainPy leverages JAX's parallel computation capabilities to enable distributed simulation and training of biological spiking neural network models. Specifically, utilizes JAX's pjit interface which parallelizes execution across devices using XLA's Global Synchronous Parallelism for Multi- Device (GSPMD) protocol. This allows seamless model parallelism in BrainPy, wherein single biological neural network can be partitioned across multiple devices with synchronized updates.
To showcase BrainPy's distributed computing capabilittie,, provide example using decision- making neural network model from Wang (2002). This model has multiple interconnected brain areas that be parallelized across devices. By using pji t, of this complex model by distributing it over multiple GPUs or TPUs. The difference under such parallelization context is that users should specify one sharding parameter when defining the model (for details please see the code in Data availability). Each device simulates a part of the network in syne with the other devices. This enables faster experimentation with larger models, as shown in Figure S14
able to accelerate the simulation
Distributed Computation of the Decision Making Network
102
Compilation Time [s]
10'
TPU 8 cores
TPU cores
TPU 2 cores
TPU core
0.0
0.2
0.4
0.6
0.8
1.0
1.2
1.4
1.6
Network Size
1e7
Figure S14: Distributed computation of a decision-making network (Wang] 2002) using BrainPy on TPUv3 devices. For TPU experiments with and cores, data are truncated due to the out-of-memory error
It's worth noting that this distribution strategy has been successfully used in training large-scale DNN models, but can incur significant communication overhead for spiking neural networks. This is because spike-based models tend to have sparse activations with only small subset of neurons spiking at each timestep. However, GSPMD synchronization requires gathering and broadcasting the full activation state across devices at every step, even though most values are zero.
This overhead is negligible for dense DNN activations but grows substantially with increasing sparsity levels. For example, in extreme cases where less than 1% of neurons in spiking model spike per timestep, over 99% the synchronization communication would transfer redundant zero values. This redundant communication consumes network bandwidth and limits scalability.
To address this inefficiency, we are exploring more optimized synchronization schemes tailored to spiking neural networks. These could leverage compression gather only non-zero spikes to reduce communication costs. Alternately, asynchronous update approaches could be employed that relax strict lockstep synchronization between devices. Such spike-specific optimizations unlock greater scalability for large-scale spiking models in BrainPy while retaining its flexible device distribution capabilities.
27

---

Published as a conference paper at ICLR 2024
M PACKAGE STRUCTURE
BrainPy is comprehensive Python library for modeling and simulating brain dynamics. It provides a flexible and extensible framework for building, simulating, training, and analyzing biological neural networks and brain-inspired algorithms. BrainPy is built on top of efficient Just-In-Tinee (JIT) compilers, enabling high-performance computations for various brain dynamics models.
The BrainPy package is organized into several core modules that encapsulate different aspects of brain dynamics programming:
brainpy math: This module provides collection of mathematical functions and utilities for use in brain dynamics modeling. It includes functions for numerical operations, vector and matrix operations, random number generation, and dedicated event-driven, sparse, and JIT connectivity operators.
bra i ipy. i integrators: This module provides a variety of numerical integration meth- ods for solving diverse differential equations (including ordinary differential equations. stochastic differential equations, fractional differential equations, and delay differential equations) that arise in brain dynamics models. These solvers are based on well-established numerical techniques, such as Euler's method, Runge-Kutt Kutta methods, and adaptive sol vers. The brainpy. integrators module is an essential component of the BrainPy library, providing a powerful and versatile set of tools for integrating differential equations used in brain dynamics models.
bra iny dyn: This module provides fundamental building blocks for biological neural net- works, including ion channels, neurons, synapses, projections, plasticity models, populations, and networks.
a iny.n.iy...... dnr: This module provides high-level API for constructing deep neural net- works. It designed to be simple and easy to use, while still providing the flexibility and power needed to build complex models. The module includes variety of pre-built layers, including convolutional layers, pooling layers, and fully connected layers. The DNN models can also be used as a component in brainpy. dyn models.
brainpy- analysis: The module provides a collection of tools for analyzing the dy- namics of neural networks, including phase plane analysis and bifurcation analysis for low-dimensional dynamical systems, and linearization analysis for nigh-dimensional ones.
brainpy- traint This module provides algorithms for training biological neural networks, including supervised learning, online learning, and offline learning.
brainpylib: This module provides customized operators and utilities for optimizing brain dynamics computations.
In addition to these modules, BrainPy also includes toolbox that extends its core functionality:
brainpy- measure:
The module provides a collection of functions for measuring and evaluating the properties of neural network models, such as firing rates, synchronization, and functional connectivity.
brainpy- inputs: The module provides a collection of functions for defining and handling input data for models in BrainPy.
brainpy connect: The module provides a mechanism for connecting neurons and synapses in BrainPy. It allows users to specify the connectivity pattern between differ- ent populations of neurons and to define the properties of synaptic connections, such as weights.
brainpy.
The module provides a collection of functions for initializing the parameters of models in BrainPy. It includes various initialization schemes, such as random initialization, Xavier initialization, and He initialization. also provides functions for initializing other network components, such as biases and noise sources.
For more details about BrainPy's package structure, please see Figure S15
28

---

Published as a conference paper at ICLR 2024
8 drainPy
Tanalysis
dynanics
cifferential equations.
#highdin
Linearization analysis and fixed/stow point finding.
analysis, bifurcation analysis, fast-slow bifurcation analysis, serialization, save and load the veriables of nodel.
checkpoints
Construct synaptic connection, inclulding built-in self-custonized connectors Construct self-custtenzed connections. Construct randon connections.
cus cusstonocon-yy
candon_conn.pp
regular_conn. cararrryoopppy
Construct regular pattern connections.
Contain Layers in deep learning and interogeration Contain classes and functions describe the dynonics of noural systens.
Flax.
dyn
1ons
Ions include sodium (Nael, potassium (K-), calciun (Ca2+), and chloride (Ct-), Tons channels dynamics intuding nany irplementations. Neural. dynamics including simplified and biophysical neurons.
channels
neurrons
"synapses
Synaplic Dynanics inc Ircluding plennomeoooogcc lugical, biological synapses, gap junctions, Syneptic projections inc mecluding AlienPre, AlignPost and STDP.
projections
Trates
Firing rate nodels including population rate nodels, RNN and reservoir computing Synaptic autputs inc including CUBA, COBA, and MgBlock. Comnon dynamics models including leaky integrator, input//utput groups. Compatible with previaus brain dynamics models interface. Encoding volues
fouts
Bothers
dyno'ld
fencoding
troins.
Initialize
This nodule provides nethods to initiatize weights. Calculate input
inputs
different forats.
integrators
Nunerical integrators inc used uncluning 00ES, SDEs, FOES and DOEs.
losses
learning
math
Provides the nathenatical foundation for brain dynanics programming. Delay varisbles.
della
Tevent
Event-driven operaters.
jätcoonn
JIT connections operators. Sparse operators. object-ariented Transfornations. Operators registrations.
#sparse
#object_transtee
fop_register
ssurrogate
Surrogate gradients for gradient-based learning. Paraltelizate Support for model sharding. Reasurerent for calculating corretations and firing rates. Gradient-based optinizers Parallel simulation techniques with milti-processes and multi-devices. Assenble tests.
#sharding
heasure
optinizers
in ceep learning.
Frunning
Ttests
tools
Acdtitional, tools used for brain dynenics programiing
running training algorithns varicus neural networks. Back-propagation learning nethad orfline training nethods, Online training methods, like RLS, LMS, etc. Nixin objects provide Visuatiization toolkit.
(back_prapagation)
ofntine.
ridge regression, Linear regression,
Lonline. py
mixin
brain dynenics arogramning.
visualization
Figure S15: The structure and components of BrainPy
N IMPLEMENTATION OF EI BALANCE NETWORK USING THE JIT CONNECTIVITY OPERATORS
For the details of the EI balance network, please refer (Vogels Abbout 2005). To implement this balanced network in BrainPy, we leverage the just-in-time (JIT) connectivity operators (see Section 4.2) to dynamically create connections between excitatory and inhibitory neuron populations during the simulation. We first initialize the neuron population specifying the neuron models and population sizes (see Line 12 ListingS66) S6). Then we use JIT connectivity to randomly connect the two layers with a given sparsity and specified weight for and inhibitory-to-excitaae
nhhiitory-to-inniiitoy projections (see Lines 17, 18, 24, and 25 in Listing S6). Note that has implemented event-driven matrix multiplication JIT operator with homogeneous synaptic weight. Specifically, this layer is based on the operator brainpy math. jitcoon.n...... event sireen__vv__ee_ mv__rrbb_____....... homo.
brainpy. ánn.
Key parameters to tune include the maximum connection synapse per neuron and connection weight strengths. Monitoring the relative excitatory and inhibitory firing rates time allows assessment of El balance and guides further parameter adjustments towards stabilizing the network dynamics. The JIT connectivity functionality has the same workflow for simulation BrainPy.
import brainpy as bp import brainpy. math as bo
class KINet (bp. DynSysGroos)
def init (self, scale-t scale-1.)
29

---

Published as a conference paper at ICLR 2024
super init ()
self.num int (4000 scale)
self.num_e
int (3200 scale)
self.num inn = int (800 scalc)
p = 80 / self. num
12
self.N Lyffee. LifRef (self. V_rest=-60., V_th--50.,
13
v_reset=-60., tau=20., tau_ref=5., V_initializer=pp. in (-55,, 2.11 entries=('I' Nonel)
14
self.delas (ww.aaDeLayyeeeL..NN.. self.E 1g..dpn.Projj1l1ggmosss11111
16
17
Comm-bp.dne don. iee. (self..uu self. num, prob-p, weight-0.6) 0)
18
syn-bp.. dyn. (xpon.dess tau=5.),
out=bp.dyn. dyn. COBA desc (B=
21
post=self.N
22
23
self.I bp. drw.Prrojjlllnnpcsssggl1 (sc.. self. num,
comm=bp.c dnn www..i..eccnnttttttoooooooooooee
prob-p, weight=6. tau-10.),
26
Eetynn.xppo....m http:e.ddn.xxxo.eeeecc. out-bp.dy dyn..ddd..cccm COBA. cosa 80.), Po
27
28
30
31
def update (self, inp):
32
spk self.delay,de I')
(elf.E(epX(:self
Ielf.I(spx exc_)x (elf.delay(self... ((in)) return self.N. spike.value
35
36
37
39 model EINet (scale=10.)
40 indi ces bm. arange (1000) 100 ms
41 spks = bm. for for_loop (1 ambda i: model model. .ste_rrun (1, 20-), indices)
bp. vi sual i ze. raster_plot (indica ces, spks, show-(True)
Listing S6: Implementation the balance network model using the connectivity operators.
O DATA AVAILABILITY
BrainPy is pip installable Python package and available at the following GitHub repos- itory: https///gftaabbbccoo////aaa/aggs http///aannpb/aiiaayy... brainpy readthedocs..oo/
with documentation at https:// /7 All the code to reproduce the result presented the pa- per can be found in the following GitHub repository: htt::::///ithubb.com//ra bra awaimnyyy.aa177--rrprroouul11i11
The MNIST dataset can be found in http::////eeee. hww..//cm/.//ec/////iiiiss.////
it is also conveniently accessible in Python via the brainpy dat ascts package. The processed KTH data can be available in (Antonik al. 2019). The multi-area connectivity data can be found in (Markov
et al. 2014) and
P SUPPLEMENTARY FIGURES
30

---

Published as a conference paper at ICLR 2024
Ell Network
E/l Network
Eil Network
200
Space
Space
Space
Space (c89)
Sparse (G8)
150
Space [Gi]
Time
Time
Time
0
10
10
10
10
Network size
Network size
Network size
Figure S16: Additional experiments of E/I balanced network. (A) Depicting the data from Figure 5C on log-scale x-axis. (B) E/I balanced network simulation with thousand synapses per neuron. Under this condition, the network still shows the linear scaling property. (C) E/I balanced network simulation where fixed connection probability (p=0.001) used and the weight is rescaled.
31


Reviews:
[{"review_id": "dG7dMpDhzk", "reviewer": "Reviewer_kvGz", "summary": "In this paper, the authors introduced a package called BrainPy. It inherits the JAX and provides support for brain simulation and SNN training. Overall, the package is interesting and useful in the stated scenarios.", "questions": "Please see the weakness.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. It demonstrates the improvement in efficiency. \n2. It provides the support for both neuroscience and DL research.", "weaknesses": "1. The package itself is more like a collection of course scripts rather than a Python package. Thus I suggest that the authors improve the engineering quality and documents for the current package.\n\n2. The comparison to existing methods is not sufficient. For example, there are existing tools like SpikingJelly for SNN. The simulation of neurons is also not sufficiently new. Thus the unique character of the current package could also be strengthened. \n\n3. I am not quite sure about the standard of package paper for ICLR. From my own understanding, the contribution and optimization to system design can be clarified in a more clear way as well.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "rm5BhPuQeT", "reviewer": "Reviewer_g6hc", "summary": "This paper introduces a brain simulator named BrainPy, which is designed to bridge the gap between brain simulation and brain-inspired computing (BIC). This paper describes the infrastructure implementation that facilitates flexible, efficient, scalable, and biological detailed brain simulations. It also describes an example project that employs this BrainPy to construct a biologically plausible spiking model to demonstrate the differentiable simulation capability of this tool.", "questions": "- What might be the biological evidence that supports the design of parameter sharing within the \"AlignPre\" and \"AlignPost\" projections?\n- How does BrainPy's speed and scalability compare to CARLsim, another brain simulator known for efficient and large-scale brain simulations?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- Clear presentation. Comprehensive comparisons with existing tools.\n- Leverages modernized tooling such as Jax and XLA, provides a user-friendly interface, and is compatible with various computing hardware.\n- Technical complexity and thoughtful designs that optimize speed and memory usage.", "weaknesses": "- Despite the exciting endeavor towards bridging the gap between brain simulators and BIC libraries, this paper appears to have limited relevance to this conference due to the lack of original theories or empirical evidence.\n- Quantitative comparisons with BIC libraries seem to be missing.\n- While the paper takes the stance of bridging brain simulators and DL frameworks, discussions about deep learning models seem to be missing.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "rXRpcZla6c", "reviewer": "Reviewer_wQtY", "summary": "The authors develop a new framework, brainpy, which allows to run (biophysically realistic) networks of neurons in a differentiable manner, thereby allowing integrations with deep learning (DL) frameworks. In addition, since it is implemented in JAX, it supports JIT compilation.", "questions": "No questions.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "I think this is an important and potentially impactful work. The paper is well written, the figures are clear, and the authors carry out many empirical experiments to demonstrate the abilities of brainpy.", "weaknesses": "The paper has the following major weaknesses:\n\n1) It does not evaluate the cost of compilation. How high is the cost of compilation compared to the runtime? Is this a clear disadvantage as compared to, e.g., NEURON? Does the compilation speed depend on whether CPU or GPU are used? How does it scale with the number of neurons?\n\n2) In section 4.2, the authors claim that there method is significantly more memory-efficient than others. Maybe I am misunderstanding this, but: do the gains that the authors claim here stem from an assumption that the connectivity matrix is low-rank? How else would they possibly be able to store connections of 86 Billion neurons?\n\nMinor: \n\n1) I believe that it would be good if the authors clarified that all JIT capabilities are due to the fact that brainpy relies on JAX, and are not implemented from scratch. Section 4.5 reads as if the authors implemented this themselves.\n\n2) The statement `It is important to note that this hierarchical composition property is not shared by other brain simulators.` is not true, see for example NetPyNE.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Np3VwalCMb", "reviewer": "Reviewer_afG7", "summary": "The authors present a programming framework that enables fast and differentiable implementation of simulations of brain circuits and similar computing systems.  The framework achieves major speed and memory benefits by taking advantage of the sparsity of these circuits in space, and in time for the case of spiking neurons.", "questions": "Presumably the JIT weight generation only works for random weights, not those learned by synaptic plasticity rules?\n\nMore detail on differentiability in spiking networks would be useful.  Equation (28) isn’t clear: neither x nor spike’ are defined, and is width the same as V_th?   As well as this, a more basic question:  do the computational benefits of sparse activity carry through to the derivatives?  For example, even if two neurons are connected with zero weight, the derivative of the objective function with respect to this weight need not be zero.", "limitations": "", "rating": 10, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "This is a very valuable contribution to the world of brain simulation, which seems likely to find a lot of users due to its speed and differentiability.", "weaknesses": "No major weaknesses identified.  \n\nI am new to ICLR reviewing so I don’t know how well this work fits within the remit of the conference.  However it is certainly a very valuable contribution to computational neuroscience.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "sInZ6pc4Hf", "reviewer": "Reviewer_zrYr", "summary": "The paper describes a package for building brain-inspired trainable network models. It is build around JAX and provides efficient just-in-time compiled code for CPUs, GPUs and TPUs. Compared to classic simulators for biological neural networks, models implemented in BrainPy benefit from auto-differentiation, and compared to standard machine learning frameworks, BrainPy provides an environment focused on building bio-inspired models with e.g. spiking interactions and detailed synaptic/neuronal dynamics.", "questions": "Due to limited time I'd like to state clearly that I did not do an in-depth review of all parts of the manuscript.\n\nProbing a subset of the results, I did not find a concrete description of the simulations underlying Fig.4C in the supplementary material. Especially, I wondered why the NEST version used by the authors is 2.20 instead of the more recent 3.x versions, and I could not find the corresponding code implementing the NEST simulation, or generally creating Fig4C, in the files supplied. \n\nOn p.8 it is mentioned concerning Fig 5C that the E/I network was scaled up to 4 million neurons with 80 synapses each. In biological networks, the (local) connectivity is typically significantly more dense, with hundreds or thousands of synapses per neuron.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- The manuscript describes a substantive library which is in an advanced stage of development.\n - It is correct that there is a need for an extensive and modern framework for training larger biological network models. The described package is therefore a significant; demonstrating its capabilities + speed comparisons to other libraries is a useful contribution.", "weaknesses": "- Except for Fig.5 A/B dealing with matrix multiplication, the figures vary the network/system size on a linear scale, and not over several orders of magnitude. This does not seem suitable to demonstrate the scaling behavior for these network models.\n- There does not seem to be a demonstration of distributed simulation/training of a large-scale model on a CPU or GPU cluster.\n- In section 4 the package is at length described as efficient, extensive, scalable, etc. without really getting into the concrete design and implementation. In particular, after reading this section I did not end up with a clear picture of the package structure and components. Maybe the description could be shortened, or made more concrete.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)