# A differentiable brain simulator bridging brain simulation and brain-
inspired computing

Chaoming Wang1,∗, Tianqiu Zhang1,∗ , Sichao He2, Yifeng Gong3, Hongyaoxing
Gu4,
Shangyang Li1, Si Wu1
1School of Psychological and Cognitive Sciences, IDG/McGovern Institute for
Brain Research,
Peking-Tsinghua Center for Life Sciences, Center of Quantitative Biology,
Academy for
Advanced Interdisciplinary Studies, Peking University
2School of Software Engineering, Beijing Jiaotong University
3School of Computer Science & Technology, Beijing Institute of Technology
4Institute of Software, Chinese Academy of Sciences
∗The authors contributed equally
<EMAIL_ADDRESS>

###### Abstract

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

##  1 Introduction

Brain simulation aims to elucidate brain functions by building dynamical
models that mimic the structure and dynamics of the brain (Gerstner et al.,
2014), while brain-inspired computing aims to develop intelligent systems by
learning from the structure and computational principles of the brain (Mehonic
& Kenyon, 2021). 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 a 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), Brian/Brian2 (Goodman &
Brette, 2008; Stimberg et al., 2019) and PyNN (Davison et al., 2008), are
designed for simulating brain dynamics models with high fidelity and accuracy.
They rely on customized numerical solvers and data structures that are not
compatible with automatic differentiation, and hence cannot support training
models with standard gradient-based methods. On the other hand, by leveraging
the automatic differentiation functionality of deep learning (DL) frameworks
like PyTorch (Paszke et al., 2019) and TensorFlow (Abadi et al., 2016),
existing BIC libraries, such as snnTorch (Eshraghian et al., 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 are not suitable to simulate large-scale brain
networks with realistic biophysical properties.

In this paper, we propose BrainPy as an innovative solution to bridge this
gap. Unlike traditional 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 a 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 a novel abstraction for executing diverse synaptic
projections. Additionally, to tackle the memory-intensive demands of brain
dynamics, BrainPy employs an object-oriented just-in-time (JIT) compilation
approach. Leveraging 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.

##  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 et al., 2019; Vitay et al., 2015), JSON (Dai et
al., 2020; Dura-Bernal et 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), snnTorch (Eshraghian et al.,
2021), Norse (Pehle & Pedersen, 2021), SpikingJelly (Fang et al., 2020), and
BrainCog (Zeng et al., 2022). These libraries utilize DL frameworks, such as
PyTorch (Paszke 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.

##  3 The Design of BrainPy

BrainPy is designed to take the combined advantages of brain simulators and DL
frameworks. This innovative tool is specifically engineered to leverage the
strengths of JAX (Frostig et al., 2018), a powerful AI framework, while
simultaneously expanding upon it to enable flexible, efficient, and scalable
simulation of various brain dynamics models.

Figure 1: The overview of BrainPy architecture.

By building upon JAX, BrainPy inherits the robust computational capabilities
and extensive library support that JAX provides (see Figure 1A). 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.dnn.FromFlax 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 feature, for instance, enables the interpretation of a
dynamical system in BrainPy as a Flax recurrent cell, allowing brain models
developed in BrainPy to be utilized within a machine-learning context.

However, BrainPy goes beyond being a mere extension of JAX. It introduces
novel and fundamental features to empower users to simulate and analyze the
intricate dynamics of the brain (refer to Figure 1B and Section 4). (1)
BrainPy provides flexibility in modeling brain dynamics. The brain exhibits a
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 can 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 dynamics 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
large-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 seamless
integration with DL techniques, BrainPy offers an integrated platform for
comprehensive simulation, training, and analysis of brain dynamics models (see
Figure 1C). Firstly, it enables efficient simulation of models at 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 online and offline learning
algorithms, and detailed spiking neural networks and rate-based recurrent
neural networks can be trained using its back-propagation-based 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
a novel object-oriented JIT transformation, BrainPy enables direct
optimization of class-based models into executable processes using XLA on CPU,
GPU, or TPU devices during simulation, training, and analysis tasks (see
Figure 1D).

##  4 Methodology

In this section, we present the specific optimizations implemented in BrainPy
for brain modeling, and highlight the advancements achieved at 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-driven computation

Compared to DNN models, brain dynamics models perform computation in a
different way. They usually have sparse connections (neurons in a 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., 2006; Roy
et al., 2019)). These unique features greatly hinder the running efficiency of
brain dynamics models if conventional dense array operators are used (see
Section 5.1). Traditional brain simulators utilize a specific data structure
and accompanying event-driven computation code to address this problem (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 functions in brain simulation as reusable primitive
operators, so 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 a set of sparse and event-driven operators in its
brainpy.math.sparse and brainpy.math.event modules, respectively. Sparse
operators can not only store synaptic connectivity 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. In contrast, BrainPy’s sparse operators prioritize the optimization
of sparse computations specifically for brain dynamics modeling. As a 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 gradient-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 a near quadratic scaling
requirement111Empirically, the best fit between the number of synapses $N_{s}$
and the number of neurons $N_{n}$ across different animal species is described
by $N_{s}=N_{n}^{1.4}$ (Lansner & Diesmann, 2012). as 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 a broad range of researchers. This
poses significant obstacles for researchers seeking to engage in large-scale
brain modeling endeavors.

To address this challenge, BrainPy introduces a JIT connectivity algorithm
with extremely low sampling complexity (see Appendix C). JIT connectivity is a
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 a 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 neuron fires. BrainPy
provides a comprehensive range of JIT connectivity operators crafted for
performing matrix-matrix multiplication and matrix-vector multiplication
within the brainpy.math.jitconn module. Compared to standard dense or 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 an easy avenue for large-scale brain simulation
for researchers.

###  4.3 Integrating models in brain simulation and AI 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 a 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-based, 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 AlignPre and AlignPost projections offer superior advantages. Both
projections assume homogeneous parameters governing synaptic dynamics within a
projection222We utilize the exponential synapse model to exemplify the
homogeneity. The dynamics of the model is described by
$\dot{g}=-g/\tau+\sum_{k}\delta(t-t^{k})$, where $g$ the conductance, $\tau$
the time constant, $t_{k}^{\mathrm{sp}}$ the pre-synaptic spike. We consider
the projection homogeneous if the value of $\tau$ remains consistent across
all synapses.. It is important to note that these projections are not
approximations. They accurately compute the same dynamics as the original
projections while providing new benefits.

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

#### Minimal device memory consumption.

In general, a projection requires storing and updating $pmn$ synaptic
variables, where $m$ and $n$ are the numbers of pre- and post-synaptic
neurons, and $p$ is the connection probability. However, AlignPre and
AlignPost 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 S7A and
Figure S7B for the computing workflow). Another aspect that showcases the
memory efficiency of AlignPre and AlignPost projections is their ability to
automatically merge duplicate synapse variable creation and updating across
multiple projections. AlingPre is particularly suitable for handling one-to-
many connections since it keeps a trace of synaptic dynamics induced by pre-
synaptic neurons (Figure S7C). This trace can be shared and reused across
multiple post-synaptic groups if the synaptic parameters (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 a single trace of variables for modeling the
postsynaptic currents. We showcase the advantages of this automatic synaptic
merging by constructing a 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 (for detailed
information, please refer to Appendix F).

#### Decoupling brain dynamics from its communication.

Furthermore, the AlignPre and AlignPost projections offer a novel perspective
on the incorporation of conventional DL components within brain simulation
models. As depicted in Figure 2A, there is a 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,
dropout, and normalization can also serve as alternative communication
mechanisms for propagating brain dynamics (Figure 2B).

###  4.4 Flexible modeling with a multi-scale model building interface

Figure 3: Multi-scale model building interface of BrainPy. Here bp is 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 a modular, composable, and flexible
programming interface centered around the brainpy.DynamicalSystem class. The
key idea underlying this multi-scale model-building paradigm is that models at
any level of resolution can be defined as brainpy.DynamicalSystem 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 3 presents an illustrative example of hierarchical model composition.
It demonstrates the recursive construction of a model, progressing from
channels (Figure 3A) to neurons (Figure 3B), networks (Figure 3C), and finally
systems (Figure 3D). It is important to note that this hierarchical
composition property is not shared by other brain simulators. 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 as building blocks to 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 integrate-and-fire (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-intensive operators, the memory-intensive nature of
brain dynamics models poses significant challenges for efficient simulation
using native Python. To overcome this, BrainPy heavily leverages the XLA
compiler (Artemev et al., 2022) for JIT compilation. JIT compilation executes
models outside of the Python interpreter and optimizes memory-intensive
operators by automatic kernel fusion. This fusion technique effectively
reduces off-chip memory access, kernel launching overhead, and CPU-GPU
switching delays, making XLA an ideal choice for compiling brain dynamics
models and mitigating their computational overhead within the Python
interface. Particularly, BrainPy incorporates a set of object-oriented JIT
transformations into its multi-scale model-building interface, as outlined in
Appendix G. These transformations enable any BrainPy model defined using the
modular and composable programming interface to be seamlessly converted into
an optimized XLA process on CPU, GPU, and TPU platforms.

##  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 a biologically plausible spiking network model on
working memory tasks.

###  5.1 Efficiency of BrainPy

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-vector 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.math.event.csrmv operator, which is used to
compute $\mathbf{y}=\mathbf{M}\mathbf{v}$, where
$\mathbf{M}\in\mathbb{R}^{m*n}$ is the synaptic connectivity,
$\mathbf{v}\in\mathbb{R}^{n}$ the presynaptic spikes, and
$\mathbf{y}\in\mathbb{R}^{m}$ the postsynaptic currents. Unlike the
alternative dense matrix-vector multiplication, it takes advantage of the CSR
representation to store $\mathbf{M}$. Different from the corresponding sparse
operator, it makes full use of the event property of the $\mathbf{v}$ and
computes only at positions corresponding to the spiking event (see Appendix
B). Therefore, brainpy.math.event.csrmv is capable of significantly reducing
temporal and spatial costs associated with synaptic computations. In practice,
we compare brainpy.math.event.csrmv with the corresponding sparse and dense
operators in JAX and PyTorch. Each operator performs the synaptic computation
in 1 s (10,000 times with the time step 0.1 ms), where we randomly sample
spiking events in $\mathbf{v}$ 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 on 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.

To verify the efficiency of BrainPy on realistic brain simulation models, we
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., 2018). We use E/I balanced networks with the LIF neuron (i.e.,
the COBA-LIF model (Vogels & Abbott, 2005)) and the Hodgkin-Huxley neuron
(i.e., the COBA-HH model (Brette et al., 2007)) as benchmarks. The E/I
balanced network typically exhibits sparse and event-driven properties, and is
highly suitable to apply brainpy.math.event.csrmv for its computation. As
shown in Figure 4C 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 Section 4.2). We investigate the memory usage and
execution speed of dense, sparse, and JIT connectivity operator
brainpy.math.jitconn.mv_prob_uniform when performing the matrix-vector
multiplication $\mathbf{y}=\mathbf{J}\mathbf{v}$. Here, $\mathbf{J}$
represents the connection matrix with a connectivity probability of $p$, and
the weights at nonzero positions are sampled from a normal distribution
$N(\mu,\sigma^{2})$ using the seed $s$. $\mathbf{J}$ is stored as a dense
matrix in the dense operator, a CSR sparse matrix in the sparse operator, and
four scalars $(p,\mu,\sigma,s)$ in the JIT connectivity operator. Our
experiments reveal that as the matrix size increases, the JIT connectivity
operator maintains nearly constant memory consumption (Figure 5A) 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.

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. We have successfully
scaled up this excitatory/inhibitory balanced network model to over 4 million
neurons and 80 incoming synapses per neuron on a single GPU device. The memory
and computing time scale linearly with network size (Figure 5C).

Furthermore, these large-scale operators can be applied in brain-inspired
reservoir models (Lukoševič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 (Antonik et 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 a testing accuracy of 91.3%. In contrast, our
implementation with JIT connectivity operators allows easy scaling up to a
reservoir with 30,000 nodes, resulting in a 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 5E), on par with the state-of-art machine learning
algorithms. Additionally, the reservoir model using the JIT connectivity
operator is twice as 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 5F).

###  5.3 Training functional brain dynamics models

Figure 6: BrainPy helps to train a data-driven, biologically realistic SNN
model for performing a working memory task. (A) The delayed match-to-sample
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 back-propagation
algorithm to train a data-driven SNN model of the prefrontal cortex (PFC)
(refer to Appendix I) with a working memory task (Figure 6A and Figure 6B). We
use the generalized leaky integrate-and-fire (GIF) model to fit the spiking
patterns of PFC neurons as observed in 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 6D). Furthermore, we consider the exponential synapse model
to model AMPA and GABA synapses (Figure 6E). We train the network to solve the
delayed match-to-sample (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 6A). 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 to make a 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 6F). 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 et al. (2016)).

##  6 Conclusion

In conclusion, BrainPy provides a 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 AI 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.

## Reproducibility Statement

We put all the codes in our supplementary materials for reproducing our
experiment results which are mainly demonstrated in Section 5. 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
5.2 and Section 5.3 , a complete description of the model and training
parameters are given in Appendix H and Appendix I.

## References

  * Abadi et al. (2016) Martín Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, et al.  Tensorflow: A system for large-scale machine learning.  In _12th USENIX symposium on operating systems design and implementation (OSDI 16)_ , pp. 265–283, 2016.
  * Antonik et al. (2019) Piotr Antonik, Nicolas Marsal, Daniel Brunner, and Damien Rontani.  Human action recognition with a large-scale brain-inspired photonic computer.  _Nature Machine Intelligence_ , pp. 1–8, 2019.
  * Artemev et al. (2022) Artem Artemev, Tilman Roeder, and Mark van der Wilk.  Memory safe computations with xla compiler.  _ArXiv_ , abs/2206.14148, 2022.
  * Babuschkin et al. (2020) 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, Iurii 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, Wojciech Stokowiec, and Fabio Viola.  The DeepMind JAX Ecosystem, 2020.  URL http://github.com/deepmind.
  * Blundell et al. (2018) Inga Blundell, Romain Brette, Thomas A Cleland, Thomas G Close, Daniel Coca, Andrew P Davison, Sandra Diaz-Pier, Carlos Fernandez Musoles, Padraig Gleeson, Dan FM Goodman, et al.  Code generation in computational neuroscience: a review of tools and techniques.  _Frontiers in neuroinformatics_ , pp. 68, 2018.
  * Brette et al. (2007) Romain Brette, Michelle Rudolph, Ted Carnevale, Michael Hines, David Beeman, James M Bower, Markus Diesmann, Abigail Morrison, Philip H Goodman, Frederick C Harris, et al.  Simulation of networks of spiking neurons: a review of tools and strategies.  _Journal of computational neuroscience_ , 23(3):349–398, 2007.
  * Carvalho et al. (2020) 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.
  * Compte et al. (2000) Albert Compte, Nicolas J.-B. Brunel, Patricia S. Goldman-Rakic, and X. J. Wang.  Synaptic mechanisms and network dynamics underlying spatial working memory in a cortical network model.  _Cerebral cortex_ , 10 9:910–23, 2000.
  * Constantinidis et al. (2016) Christos Constantinidis, Xue-Lian 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.
  * Dai et al. (2020) Kael Dai, Sergey L Gratiy, Yazan N Billeh, Richard Xu, Binghuang Cai, Nicholas Cain, Atle E Rimehaug, Alexander J Stasik, Gaute T Einevoll, Stefan Mihalas, et al.  Brain modeling toolkit: An open source software suite for multiscale modeling of brain circuits.  _PLOS Computational Biology_ , 16(11):e1008386, 2020.
  * Davison et al. (2008) Andrew P. Davison, Daniel Brüderle, Jochen M. Eppler, Jens Kremkow, Eilif B. Müller, Dejan Pecevski, Laurent Udo Perrinet, and Pierre Yger.  Pynn: A common interface for neuronal network simulators.  _Frontiers in Neuroinformatics_ , 2, 2008.
  * Dura-Bernal et al. (2019) Salvador Dura-Bernal, Benjamin A Suter, Padraig Gleeson, Matteo Cantarelli, Adrian Quintana, Facundo Rodriguez, David J Kedziora, George L Chadderdon, Cliff C Kerr, Samuel A Neymotin, et al.  Netpyne, a tool for data-driven multiscale modeling of brain circuits.  _Elife_ , 8:e44494, 2019.
  * Eshraghian et al. (2021) Jason K Eshraghian, Max Ward, Emre Neftci, Xinxin Wang, Gregor Lenz, Girish Dwivedi, Mohammed Bennamoun, Doo Seok Jeong, and Wei D Lu.  Training spiking neural networks using lessons from deep learning.  _arXiv preprint arXiv:2109.12894_ , 2021.
  * Fang et al. (2020) Wei Fang, Yanqi Chen, Jianhao Ding, Ding Chen, Zhaofei Yu, Huihui Zhou, Yonghong Tian, and other contributors.  Spikingjelly.  https://github.com/fangwei123456/spikingjelly, 2020.  Accessed: YYYY-MM-DD.
  * Frostig et al. (2018) Roy Frostig, Matthew Johnson, and Chris Leary.  Compiling machine learning programs via high-level tracing.  2018\.
  * Gerstner et al. (2014) 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.
  * Gewaltig & Diesmann (2007) Marc-Oliver Gewaltig and Markus Diesmann.  Nest (neural simulation tool).  _Scholarpedia_ , 2(4):1430, 2007.
  * Gleeson et al. (2010) Padraig Gleeson, Sharon Crook, Robert C Cannon, Michael L Hines, Guy O Billings, Matteo Farinella, Thomas M Morse, Andrew P Davison, Subhasis Ray, Upinder S Bhalla, et al.  Neuroml: a language for describing data driven models of neurons and networks with a high degree of biological detail.  _PLoS computational biology_ , 6(6):e1000815, 2010\.
  * Goodman & Brette (2008) Dan FM Goodman and Romain Brette.  Brian: a simulator for spiking neural networks in python.  _Frontiers in neuroinformatics_ , 2:5, 2008.
  * Hass et al. (2016) Joachim Hass, Loreen Hertäg, and Daniel Durstewitz.  A detailed data-driven network model of prefrontal cortex reproduces key features of in vivo activity.  _PLoS Computational Biology_ , 12, 2016.
  * Hazan et al. (2018) Hananel Hazan, Daniel J. Saunders, Hassaan Khan, Devdhar Patel, Darpan T. Sanghavi, Hava T. Siegelmann, and Robert Kozma.  Bindsnet: A machine learning-oriented spiking neural networks library in python.  _Frontiers in Neuroinformatics_ , 12:89, 2018.  ISSN 1662-5196.  doi: 10.3389/fninf.2018.00089.  URL https://www.frontiersin.org/article/10.3389/fninf.2018.00089.
  * Heek et al. (2020) Jonathan Heek, Anselm Levskaya, Avital Oliver, Marvin Ritter, Bertrand Rondepierre, Andreas Steiner, and Marc van Zee.  Flax: A neural network library and ecosystem for JAX, 2020.  URL http://github.com/google/flax.
  * Hines & Carnevale (1997) Michael L Hines and Nicholas T Carnevale.  The neuron simulation environment.  _Neural computation_ , 9(6):1179–1209, 1997.
  * Joglekar et al. (2017) 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 of the primate cortex.  _Neuron_ , 98:222–234.e8, 2017.  URL https://api.semanticscholar.org/CorpusID:4572302.
  * Jordan et al. (2018) 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.  URL https://api.semanticscholar.org/CorpusID:3646839.
  * Kingma & Ba (2014) Diederik P. Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  _CoRR_ , abs/1412.6980, 2014.
  * Knight & Nowotny (2020) James C. Knight and Thomas Nowotny.  Larger gpu-accelerated brain simulations with procedural connectivity.  _Nature Computational Science_ , 1:136 – 142, 2020.
  * Kunkel et al. (2012) 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 Neuroinformatics_ , 5, 2012.  URL https://api.semanticscholar.org/CorpusID:6261003.
  * Lansner & Diesmann (2012) Anders Lansner and Markus Diesmann.  Virtues, pitfalls, and methodology of neuronal network modeling and simulations on supercomputers.  2012\.  URL https://api.semanticscholar.org/CorpusID:59848490.
  * Lukoševičius (2012) Mantas Lukoševičius.  A practical guide to applying echo state networks.  In _Neural Networks_ , 2012.
  * Lytton et al. (2008) William W. Lytton, Ahmet Omurtag, Samuel A. Neymotin, and Michael L. Hines.  Just-in-time connectivity for large spiking networks.  _Neural Computation_ , 20:2745–2756, 2008.
  * Markov et al. (2014) Nikola T Markov, Maria M Ercsey-Ravasz, AR Ribeiro Gomes, Camille Lamy, Loic Magrou, Julien Vezoli, Pierre Misery, Arnaud Falchier, Rene Quilodran, Marie-Alice Gariel, et al.  A weighted and directed interareal connectivity matrix for macaque cerebral cortex.  _Cerebral cortex_ , 24(1):17–36, 2014.
  * Mehonic & Kenyon (2021) Adnan Mehonic and A. J. Kenyon.  Brain-inspired computing needs a master plan.  _Nature_ , 604:255 – 260, 2021.
  * Meunier et al. (2010) David Meunier, Renaud Lambiotte, and Edward T Bullmore.  Modular and hierarchically modular organization of brain networks.  _Frontiers in neuroscience_ , 4:200, 2010.
  * Mihalas & Niebur (2009) Stefan Mihalas and Ernst Niebur.  A generalized linear integrate-and-fire neural model produces diverse spiking behaviors.  _Neural Computation_ , 21:704–718, 2009.
  * Paszke et al. (2019) Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al.  Pytorch: An imperative style, high-performance deep learning library.  _Advances in neural information processing systems_ , 32, 2019.
  * Pehle & Pedersen (2021) Christian Pehle and Jens Egholm Pedersen.  Norse - A deep learning library for spiking neural networks, January 2021.  URL https://doi.org/10.5281/zenodo.4422025.  Documentation: https://norse.ai/docs/.
  * Plank et al. (2021) Philipp Plank, A. Rao, Andreas Wild, and Wolfgang Maass.  A long short-term memory for ai applications in spike-based neuromorphic hardware.  _Nature Machine Intelligence_ , 4:467 – 479, 2021.
  * Potjans & Diesmann (2012) Tobias C. Potjans and Markus Diesmann.  The cell-type specific cortical microcircuit: Relating structure and activity in a full-scale spiking network model.  _Cerebral Cortex (New York, NY)_ , 24:785 – 806, 2012.
  * Rasmussen (2018) Daniel Rasmussen.  Nengodl: Combining deep learning and neuromorphic modelling methods.  _Neuroinformatics_ , 17:611–628, 2018.
  * Ros et al. (2006) 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 to characterize neuronal dynamics.  _Neural Computation_ , 18:2959–2993, 2006.
  * Roy et al. (2019) Kaushik Roy, Akhilesh R. Jaiswal, and Priyadarshini Panda.  Towards spike-based machine intelligence with neuromorphic computing.  _Nature_ , 575:607 – 617, 2019.
  * Song et al. (2016) 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.
  * Stimberg et al. (2014) Marcel Stimberg, Dan FM Goodman, Victor Benichoux, and Romain Brette.  Equation-oriented specification of neural models for simulations.  _Frontiers in neuroinformatics_ , 8:6, 2014.
  * Stimberg et al. (2019) Marcel Stimberg, Romain Brette, and Dan FM Goodman.  Brian 2, an intuitive and efficient neural simulator.  _Elife_ , 8:e47314, 2019.
  * Sussillo & Abbott (2009) David Sussillo and L. F. Abbott.  Generating coherent patterns of activity from chaotic neural networks.  _Neuron_ , 63:544–557, 2009.
  * Swadlow (1990) 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.
  * Vitay et al. (2015) Julien Vitay, Helge Ülo Dinkelbach, and Fred H Hamker.  Annarchy: a code generation approach to neural simulations on parallel hardware.  _Frontiers in neuroinformatics_ , 9:19, 2015.
  * Vogels & Abbott (2005) Tim P. Vogels and L. F. Abbott.  Signal propagation and logic gating in networks of integrate-and-fire neurons.  _The Journal of Neuroscience_ , 25:10786 – 10795, 2005\.
  * Zeng et al. (2022) Yi Zeng, Dongcheng 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 Du, Qingqun Kong, Zizhe Ruan, and Weida Bi.  Braincog: A spiking neural network based brain-inspired cognitive intelligence engine for brain-inspired ai and brain simulation.  _ArXiv_ , abs/2207.08533, 2022.

##  Appendix A Environment settings

In this paper, all evaluations and benchmarks were conducted in a Python 3.10
environment, which was installed on a system running Ubuntu 22.04.2 LTS. For
CPU experiments, we used Intel(R) Xeon(R) W-2255 CPU @ 3.70GHz, 64GB RAM @
3200MHz. For GPU experiments, we used the NVIDIA RTX${}^{\text{TM}}$ 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 2.20.0, Brian2
(Stimberg et al., 2019) at version 2.5.1, ANNArchy (Vitay et al., 2015) in
version 4.7.2, and BindsNet (Hazan et al., 2018) in version 0.3.2,

For brain-inspired computing, we compared BrainPy with several spiking neural
networks (SNN) training packages, including snnTorch (Eshraghian et al., 2021)
in version 0.6.1, SpikingJelly (Fang et al., 2020) in version 0.0.0.0.14, and
Norse (Pehle & Pedersen, 2021) 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.

##  Appendix B Event-driven matrix-vector multiplication

Event-driven matrix-vector multiplication $\mathbf{y}=\mathbf{M}\mathbf{v}$ in
BrainPy is used for synaptic computation, where $\mathbf{v}$ is the
presynaptic spikes, $\mathbf{M}$ the synaptic connection matrix, and
$\mathbf{y}$ the postsynaptic current. Specifically, it performs matrix-vector
multiplication in a sparse and efficient way by exploiting the event property
of the input vector $\mathbf{v}$. Instead of multiplying the entire matrix
$\mathbf{M}$ by the vector $\mathbf{v}$, which can be wasteful if $\mathbf{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.

In particular, BrainPy implements brainpy.math.event.csrmv, in which the
connection matrix $\mathbf{M}$ is stored as the compressed sparse row (CSR)
sparse matrix. The computation of this operator requires several parameters:

  1. 1.

data: The weights of $\mathbf{M}$ contain the non-zero elements in the row-
major order.

  2. 2.

indices: The array contains the column indices of the non-zero elements in the
matrix $\mathbf{M}$.

  3. 3.

indptr: The array contains the starting index of each row.

  4. 4.

events: The presynaptic spiking vector $\mathbf{v}$.

  5. 5.

shape: A tuple with two integers representing the shape of the matrix
$\mathbf{M}$.

brainpy.math.event.csrmv makes full use of the event property of $\mathbf{v}$,
and computes only at positions where the spike in $\mathbf{v}$ is True. The
pseudo-code of this operator can be written in Listing S1.

⬇

1def csrmv(data, indices, idnptr, events, outs):

2 for i, event in enumerate(events):

3 if event:

4 for j in range(idnptr[i], idnptr[i + 1]):

5 outs[indices[i]] += data[j]

Listing S1: The Python pseudo-code of brainpy.math.event.csrmv, where data,
indices, and idnptr corresponds to the matrix $\mathbf{M}$, events indicates
the vector $\mathbf{v}$, and outs represents the postsynaptic current
$\mathbf{y}$.

##  Appendix C Matrix-vector multiplication with the just-in-time connectivity

Synaptic connectivity storage imposes a memory bottleneck for large-scale
neuronal network simulations, as it scales quadratically with the number of
neurons. Previous studies have demonstrated that static synaptic parameters in
brain modeling can 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-vector product operation $\mathbf{y}=\mathbf{M}\mathbf{v}$, where
$\mathbf{M}$ is the synaptic connection matrix to be regenerated, $\mathbf{v}$
is the presynaptic spike vector, and $\mathbf{y}$ is 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 a Bernoulli process with
success probability $p$.

A naive way of drawing from the Bernoulli process is to sample repeatedly from
the uniform distribution $U[0,1]$ and create a synapse if the sample is
smaller than $p$. However, this is highly inefficient for sparse connectivity
($p\ll 1$).

A more efficient way of drawing the sampling is using the geometric
distribution $\text{Geo}[p]$ Knight & Nowotny (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
$\frac{{\mathrm{log}}(\text{U}[0,1])}{{\mathrm{log}}(1-{P})}$. However, this
sampling method is still costly on CPU and GPU devices due to the existence of
log operation.

We propose to sample the distance between two consecutive connection positions
using the uniform random integers from 1 to $\lfloor\frac{2}{p}-1\rfloor$. The
expectation of $U[1,\lfloor\frac{2}{p}-1\rfloor]$ is $\frac{1}{p}$, ensuring
that the sampling has the desired connection probability of $p$. In practice,
sampling from $U[1,\lfloor\frac{2}{p}-1\rfloor]$ is one order of magnitude
faster than sampling from $\text{Geo}[p]$. The corresponding Python pseudo-
code of our proposed solution is shown in Listing S2.

⬇

1import math, random

2

3def jitconn_prob_homo(events, prob, weight, seed, outs):

4 random.seed(seed)

5 max_cdist = math.ceil(2/prob - 1)

6 for event in events:

7 if event:

8 post_i = random.randint(1, max_cdist)

9 while post_i < len(outs):

10 outs[post_i] += weight

11 post_i += random.randint(1, max_cdist)

Listing S2: The Python pseudo-code of our just-in-time connectivity operator,
where events indicates the presynaptic spikes $\mathbf{v}$, prob, weight, and
seed corresponds to the matrix $\mathbf{M}$, and outs represents the
postsynaptic current $\mathbf{y}$. The seed parameter guarantees the
reproducibility and consistency of matrix regeneration across multiple
invocations of this function. Note here, all nonzero elements in the matrix
$\mathbf{M}$ have the save value weight.

##  Appendix D Synaptic projections with AlignPre and AlignPost

AlignPre holds true because with homogeneous synaptic parameters, the spike
train coming from the same presynaptic neuron will lead to the same synaptic
dynamics. As a result, all synapses originating from the same presynaptic
neuron can share a single dynamical variable. The AlignPre projection is
suitable for all types of synapse models. In contrast, the applicability of
AlignPost is limited to synapse models that exhibit exponential dynamics or
are composed of exponential components. For the exponential synapse, the
conductance $g$ of a post-synaptic neuron is updated according to $g\leftarrow
g+1$ whenever a spike arrives and regardless of which presynaptic neuron
emitted this spike (refer to Appendix E for details).

Figure S7: Overview of Synaptic Projections. (A) The workflow of the AlignPre
model. (B) The workflow of the AlignPost model. (C) The AlignPre model is
suitable for one-to-many connection. (D) The AlignPost model is suitable for
many-to-one connection.

AlignPre is capable of modeling scenarios where the pre-synaptic neuron group
projects to multiple downstream post-synaptic neuron groups. This mechanism
imposes a constraint on the size of synapse variables, ensuring alignment with
the number of pre-synaptic neuron groups. Consequently, if these projections
share the same delay, only a single synapse variable is stored for all these
projections. The operational sequence of AlignPre 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 is adjusted to
match the number of post-synaptic neuron groups. The operational sequence of
AlignPost closely resembles that of AlignPre, with the notable difference
being that the calculation of synaptic variables occurs after the synaptic
communication layer. Specifically, AlignPost combines delayed pre-synaptic
spikes with synaptic weights to 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.

##  Appendix E Exponential synapse model enables the AlignPost projection

For the exponential synapse, the conductance $g$ of a post-synaptic neuron is
governed by

$g(t)=\sum_{j}^{n}\mathrm{exp}(-\frac{t-t^{sp}_{j}}{\tau}),$ (1)

where $n$ is the total number of spikes the post-synaptic neuron received at
the current time $t$, $t^{sp}_{j}$ the spiking time of the received $j$-th
spike, and $\tau$ the synaptic time constant.

Eq (1) can be rewritten as a differential equation:

$\dot{g}=-\frac{g(t)}{\tau}+\sum_{i,k}\delta(t-t^{\mathrm{sp}}_{i,k}),$ (2)

where $i$ is the index of the connected pre-synaptic neuron, and
$t^{\mathrm{sp}}_{i,k}$ is the time of $k$-th spike of the pre-synaptic neuron
$i$.

In the discrete version, Eq (2) is equivalent to the following equations:

$g(t)=\mathrm{exp}(-\Delta t/\tau)g(t-\Delta t),$ (3)

and whenever a pre-synaptic spike arrives, $g(t)$ undergoes an update
according to:

$g(t)\leftarrow g(t)+1.$ (4)

Therefore, the exponential dynamics property enables us to track and record
information using a single variable for each post-synaptic neuron.

##  Appendix F Multi-area spiking neural network

We build a spiking network model which is inspired by with 30 brain areas.
Simulations are performed using a network of leaky integrate-and-fire neurons,
with the local circuit and long-range connectivity structure derived from
(Joglekar et al., 2017). Each of the five areas consists of 4000 neurons, with
3200 excitatory and 800 inhibitory neurons.

For each neuron, the membrane potential dynamics are described by the
following equations:

$\tau\frac{dV}{dt}=-(V(t)-V_{rest})+RG(t)$ (5)

when $V(t)>V_{th}$, the neuron generates a spike and $V(t)$ is set to
$V_{reset}$. $V$ is the membrane potential, $V_{rest}$ is the resting membrane
potential, $V_{reset}$ is the reset membrane potential, $V_{th}$ is the spike
threshold, $\tau$ is the time constant, $R=1\omega$ is the resistance, and
$G(t)$ is the time-variant synaptic inputs.

For the synapse, we use conductance-based synaptic interactions. Particularly,
$G(t)$ is given by:

$\displaystyle G(t)=-\sum_{j}g_{ji}(t)\left(V_{i}-E_{j}\right),$ (6)

where $V_{i}$ is the membrane potential of neuron $i$. The synaptic
conductance from neuron $j$ to neuron $i$ is represented by $g_{ji}(t)$, while
$E_{j}$ signifies the reversal potential of that synapse. For excitatory
synapses, $E_{j}$ was set to 0 mV, whereas for inhibitory synapses, it was -80
mV.

The dynamics of the synaptic conductance is given by:

$\displaystyle\frac{dg_{ji}}{dt}=-\frac{g_{ji}}{\tau_{\mathrm{decay}}}+g_{\mathrm{max}}\sum_{k}\delta(t-t_{j}^{k}),$
(7)

where $t_{j}^{k}$ is the spiking time of the presynaptic spike. Whenever a
spike occurred in neuron $j$, the synaptic conductance $g_{ji}$ experienced an
immediate increase by a fixed amount $g_{\mathrm{max}}$. Subsequently, the
conductance $g_{ji}$ decayed exponentially with a time constant of
$\tau_{\mathrm{decay}}=5$ ms for excitation and $\tau_{\mathrm{decay}}=10$ ms
for inhibition.

The connection density is set according to the experimental connectivity data
(Markov et al., 2014). The inter-areal connectivity is measured as a 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 E/I
balanced 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 a
distance matrix based on 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 we measure the memory usage after constructing the model, the
compilation time during JIT compilation, and the execution time when
simulation the model. The data is listed in Figure S8.

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.

##  Appendix G OO transformation comparison

The JIT compilation interface of JAX/XLA is designed to work with pure Python
functions, which is not compatible with the modular and composable programming
interface in BrainPy for building brain dynamics models (see Section 4.4). To
bridge this gap, we provide brainpy.math.jit for automatic object-oriented JIT
transformation.

In this object-oriented transformation, variables that will change during
execution should be declared as brainpy.math.Variable (otherwise, they will be
treated as constants and compiled into the transformed models). Then, a single
line of calling model_transforemd = brainpy.math.jit(model) will transform the
model onto the optimized XLA process model_transforemd. In this way, the
object-oriented JIT compilation in BrainPy provides a robust infrastructure
for optimizing the performance of brain dynamics models.

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 to separate variables from computing logic. In Haiku, state
variables have a complex syntax for declaration and calculation, which
increases the difficulties for users to use these time-dependent variables in
complex brain dynamical systems. In BrainPy, variables that will change during
execution should be declared as brainpy.math.Variable (otherwise, they will be
treated as constants and compiled into the transformed models). Then, a single
line of calling model_transforemd = brainpy.math.jit(model) will transform the
model onto the optimized XLA process model_transforemd. In this way, the
object-oriented JIT compilation in BrainPy provides a robust infrastructure
for optimizing the performance of brain dynamics models. Although we could
have built on existing ML libraries, we felt a ground-up approach optimized
for brain modeling was needed. We believe this Pythonic approach makes the
object-oriented JIT much more accessible for neuroscientists without deep ML
expertise and lowers the barriers compared to Flax or Haiku.

##  Appendix H The reservoir computing model

#### Reservoir model.

The dynamics of the reservoir model used here is given by (Lukoševičius,
2012):

$\displaystyle\mathbf{x}(t)$
$\displaystyle=(1-\alpha)\cdot\mathbf{x}(t-1)+\alpha\cdot
f(\mathbf{W}_{in}\,\mathbf{u}(t)+\mathbf{W}_{rec}\,\mathbf{x}(t-1)),$ (8)
$\displaystyle\mathbf{y}(t)$ $\displaystyle=\mathbf{W}_{out}\,\mathbf{x}(t)$
(9)

where $\mathbf{x}(t)$ is the reservoir state, $\mathbf{y}(t)$ the readout
value, $\mathbf{W}_{in}$ and $\mathbf{W}_{rec}$ are input and recurrent
connection matrices, respectively, $\mathbf{W}_{out}$ the readout weight
matrix which can be trained by either offline learning or online learning
algorithms, $\alpha\in(0,1]$ the leaky rate, $\mathbf{u}(t)$ the input at time
step $t$, and $f$ the nonlinear activation function.

#### Model implementation.

$\mathbf{W}_{in}$ and $\mathbf{W}_{rec}$ are fixed and randomly initialized,
and they are highly suitable to be implemented with the just-in-time
connectivity operators.

Since $\mathbf{W}_{in}$ is usually initialized with the uniform distribution
$U[-s,s]$, we implement the input operation $\mathbf{W}_{in}\mathbf{u}(t)$
with brainpy.math.jitconn.mv_prob_uniform(vector, w_low, w_high, conn_prob,
seed), where vector corresponds to the input $\mathbf{u}(t)$, w_low
corresponds to $-s$, w_high corresponds to $s$, conn_prob corresponds to the
connection probability of the input matrix $\mathbf{W}_{in}$, and seed the
random seed controlling the reproducibility of the input matrix.

$\mathbf{W}_{rec}$ is usually initialized with the normal distribution
$N(0,\sigma)$ with a desirable spectral radius $\rho$. Usually, we have the
relationship of $\sigma=\rho/\sqrt{m*p}$, where $p$ is the connection
probability of the matrix, $m$ the matrix size, and $\rho$ the spectral radius
of the recurrent weight $\mathbf{W}_{rec}$. Therefore, we implement the
recurrent matrix operation $\mathbf{W}_{rec}\,\mathbf{x}(t-1)$ with
brainpy.math.jitconn.mv_prob_normal(vector, w_mu, w_sigma, conn_prob, seed),
where vector corresponds to the reservoir state $\mathbf{x}(t-1)$, w_mu is 0,
w_sigma corresponds to $\sigma$, conn_prob corresponds to the connection
probability of the recurrent matrix $\mathbf{W}_{rec}$, 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
$\mathbf{W}_{out}$ that minimizes the square error between $\mathbf{y}(t)$ and
$\mathbf{y}^{target}(t)$. The common way to learn the linear output weight
$\mathbf{W}_{out}$ is using the ridge regression algorithm (Lukoševičius,
2012). However, ridge regression is 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 FORCE learning method
(Sussillo & Abbott, 2009) to train the model. Contrary to 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 a model one sample of data at a time, which can
significantly reduce computational costs and enables training a large-scale
reservoir to be possible.

#### Experiment details.

We conducted several experiments to investigate the performance of large-scale
reservoir models on 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 which can be captured by reservoir models.
Then we evaluate the reservoir model on the MNIST dataset, which is a static
image dataset without temporal information. All parameter selections for
classifying two datasets are shown in Table S1.

Table S1: The parameter table of the reservoir model on different datasets. Parameter | KTH dataset | MNIST dataset
---|---|---
$\mathbf{W}_{in}$ connection probability | [0.01, 0.005] | 0.1
$\mathbf{W}_{rec}$ connection probability | [0.001, 0.0002, 0.0001] | [0.1, 0.01]
Distribution of $\mathbf{W}_{in}$ | Uniform distribution | Uniform distribution
Distribution of $\mathbf{W}_{rec}$ | Normal distribution | Normal distribution
Spectral radius $\rho$ | 1.0 | 1.3
Input scaling $s$ | 0.1 | 0.3
Leaky rate $\alpha$ | 0.9 | 0.6
Number of training epoch | 10 | 5
FORCE learning rate | 0.1 | 0.1

We set the connection probabilities of $\mathbf{W}_{in}$ and
$\mathbf{W}_{rec}$ as follows: For the KTH dataset, $\mathbf{W}_{in}$ has a
connection probability of 0.01 for size $\in$ [2000, 4000, 8000, 10000, 20000]
and 0.005 for size $\in$ [30000]. $\mathbf{W}_{rec}$ has a connection
probability of0.001 for size $\in$ [2000, 4000, 8000], 0.0002 for size $\in$
[10000], and 0.0001 for size $\in$ [20000, 30000]. For the MNIST dataset,
$\mathbf{W}_{in}$ has a connection probability of 0.1 for all reservoir sizes,
and $\mathbf{W}_{rec}$ has a connection probability of 0.1 for size $\in$
[2000, 4000, 8000, 10000], and 0.01 for size $\in$ [20000, 30000, 40000,
50000].

##  Appendix I Recurrent neural networks for performing the DMS task

#### Network structure.

The architecture of recurrent networks used here (both the rate- and spiking-
based models) is 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-varying outputs. The input layer
encodes the sensory information relevant to the task, while the readout layer
produces a decision in terms of an abstract decision variable.

Figure S9: The schematic diagram of the recurrent neural network model for
implementing a working memory task.

#### The input layer.

The input layer consists of $N$ motion direction-tuned input neurons (for the
spiking network, $N=100$; for the rate model, $N=24$). The tuning of the
motion direction-selective neurons followed a von Mises distribution, such
that the firing rate activity of the input neuron $i$ during the stimulus
period was

$\displaystyle
u_{t}^{i}=A\exp\left(\kappa\cos\left(\theta-\theta_{\mathrm{pref}}^{i}\right)\right),$
(10)

where $\theta$ is the direction of the stimulus, $\theta_{\mathrm{pref}}^{i}$
is the preferred direction of input neuron $i$, $\kappa$ was set to 2.0, and
$A$ is the maximum firing rate and was set to 40 Hz. Furthermore, for all
periods, input neurons keep a constant background firing rate
$r_{\mathrm{bg}}$. Here, $r_{\mathrm{bg}}$ was set to 1 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 S10 demonstrates the cases of match (when the stimuli
during the sample and test periods are identical) and non-match (when they are
different).

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.

#### The recurrent layer in the rate model.

In the rate-based model, the recurrent layer consists of 80 excitatory neurons
and 20 inhibitory neurons. The firing activity $\mathbf{r}$ of the rate-based
recurrent neurons was modeled with the following dynamical equation (Song et
al., 2016):

$\displaystyle\tau\frac{\mathrm{d}\bm{r}}{\mathrm{d}t}=-r+f\left(W^{\mathrm{rec}}\bm{r}+W^{\mathrm{in}}\bm{u}+\bm{b}^{\mathrm{rec}}+\sqrt{2\tau}\sigma_{\mathrm{rec}}\zeta\right),$
(11)

where $\tau$ is the neuron time constant, $f(*)$ represents the activation
function, $W^{\mathrm{rec}}$ and $W^{\mathrm{in}}$ are the recurrent synaptic
weights and input synaptic weights, respectively, $b^{\mathrm{rec}}$ is a bias
term, $\zeta$ is independent Gaussian white noise with zero mean and unit
variance applied to all recurrent neurons and $\sigma_{\mathrm{rec}}$ is the
noise intensity. The activation function used in this study is
$f(x)=\max(0,x)$.

In practice, the simulation of the Eq. 11 was performed with the Euler
integration method:

$\displaystyle\bm{r}_{t}=(1-\alpha)\bm{r}_{t-dt}+\alpha
f\left(W^{\mathrm{rec}}\bm{r}_{t-dt}+W^{\mathrm{in}}\bm{u}_{t}+\bm{b}^{\mathrm{rec}}+\sqrt{\frac{2}{\alpha}}\sigma_{\mathrm{rec}}N(0,1)\right),$
(12)

where $\alpha=dt/\tau$, $dt$ the numerical integration step, $N(0,1)$ denotes
the standard normal distribution.

To maintain the separation of excitatory and inhibitory neurons, we decomposed
the recurrent weight matrix $W^{\mathrm{rec}}$ as the product between a
trainable non-negative matrix $W^{\mathrm{rec},+}$ and a fixed diagonal matrix
$D$ (Song et al., 2016):

$\displaystyle W^{\mathrm{rec}}=W^{\mathrm{rec},+}D$ (13) $\displaystyle
D=\left[\begin{array}[]{lll}1&&\\\ &\ddots&\\\ &&-1\end{array}\right]$ (17)

#### The recurrent layer in the spiking model.

For the spiking model, the recurrent cell was implemented with the spiking
neuron model. In this study, the spiking neuron is modified from the
generalized integrate-and-fire neuron model (Mihalas & Niebur, 2009). In
particular, this model has two internal currents, one fast and one slow. Its
dynamic behavior is given by

$\displaystyle\tau_{I1}\frac{d\mathbf{I_{1}}}{dt}$
$\displaystyle=-\mathbf{I_{1}},$ fast internal current (18)
$\displaystyle\tau_{I2}\frac{d\mathbf{I_{2}}}{dt}$
$\displaystyle=-\mathbf{I_{2}},$ slow internal current (19)
$\displaystyle\tau_{V}\frac{d\mathbf{V}}{dt}$
$\displaystyle=-\mathbf{V}+V_{rest}+R(\mathbf{I_{1}}+\mathbf{I_{2}}+\mathbf{I_{ext}}),$
membrane potential (20)

When $V^{i}$ of $i$-th neuron meets $V_{th}$, the modified GIF model fires:

$\displaystyle I_{1}^{i}\leftarrow A_{1},$ (21) $\displaystyle
I_{2}^{i}\leftarrow I_{2}^{i}+A_{2},$ (22) $\displaystyle V^{i}\leftarrow
V_{rest},$ (23)

where $\tau_{I1}$ denotes the time constant of the fast internal current,
$\tau_{I2}$ the time constant of the slow internal current, $\tau_{V}$ the
time constant of membrane potential, $R$ the resistance, $\mathbf{I_{ext}}$
the external input, $V_{rest}$ the resting potential, and $A_{1}$ and $A_{2}$
the spike-triggered currents.

For matching the statistical data of electrophysiological experiments in Hass
et al. (2016), we assigned 80 excitatory neurons and 20 inhibitory neurons.
Inhibitory neurons exhibit a bursting firing pattern, while excitatory neurons
show adaptive spikings. We set $A_{1}=8.0$ for inhibitory neurons, $A_{1}=0$
for excitatory neurons, $A_{2}=-0.6$, $\tau_{I1}=10$ ms, $\tau_{I2}$ was
sampled from a uniform distribution $U[100,3000]$ ms.

The external current $\mathbf{I_{ext}}$ in the network was modeled through the
exponential synapse model, whose dynamics is given by:

$\displaystyle\frac{d{I_{\mathrm{ext}}^{i}}}{dt}=-\frac{{I_{\mathrm{ext}}^{i}}}{\tau_{\mathrm{decay}}}+\sum_{j}W^{\mathrm{rec}}_{ij}z^{j}[t-d^{ij}]+\sum_{j}W^{\mathrm{in}}_{ij}u^{j}[t],$
(24)

where $\tau_{\mathrm{decay}}$ is the time constant of the synaptic state,
$t^{k}$ the time of the presynaptic spike, $W^{\mathrm{rec}}$ the recurrent
weight, $W^{\mathrm{in}}$ the input to recurrent weight, and $d^{ij}$ the
synaptic delay. Here, $d^{ij}=0$ ms. $\tau_{\mathrm{decay}}=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:

$\displaystyle\mathbf{I_{1}}[t+\Delta
t]=\mathrm{where}(\mathbf{z}[t]==1,A_{1},\alpha_{I_{1}}\mathbf{I_{1}}[t])$
(25) $\displaystyle\mathbf{I_{2}}[t+\Delta
t]=\alpha_{I_{2}}\mathbf{I_{2}}[t]+A_{2}\mathbf{z}[t]$ (26)
$\displaystyle\mathbf{I_{\mathrm{ext}}}[t+\Delta
t]=\alpha_{I_{\mathrm{ext}}}\mathbf{I_{\mathrm{ext}}}[t]+W_{\mathrm{rec}}\mathbf{z}[t]+W_{\mathrm{in}}\mathbf{u}[t]$
(27) $\displaystyle\mathbf{V}[t+\Delta
t]=\alpha_{V}\mathbf{V}[t]+(V_{rest}+R(\mathbf{I_{1}}[t+\Delta
t]+\mathbf{I_{2}}[t+\Delta t]+\mathbf{I_{ext}}[t]))\Delta t$ (28)
$\displaystyle\mathbf{z}[t+\Delta t]=\mathrm{where}(\mathbf{V}[t+\Delta
t]>V_{th},1,0)$ (29)

where $z$ is the spiking state,
$\alpha_{I_{1}}=e^{-\frac{1}{\tau_{I_{1}}}\Delta t}$,
$\alpha_{I_{2}}=e^{-\frac{1}{\tau_{I_{2}}}\Delta t}$,
$\alpha_{V_{th}}=e^{-\frac{1}{\tau_{V_{th}}}\Delta t}$,
$\alpha_{V}=e^{-\frac{1}{\tau_{V}}\Delta t}$, and
$\alpha_{I_{\mathrm{ext}}}=e^{-\frac{1}{\tau_{\mathrm{decay}}}\Delta t}$.

Similar to the rate-based model, we decomposed the recurrent weight matrix
$W^{\mathrm{rec}}$ as the product between a trainable non-negative matrix
$W^{\mathrm{rec},+}$ and a fixed diagonal matrix $D$ (Song et al., 2016):

$\displaystyle W^{\mathrm{rec}}=W^{\mathrm{rec},+}D$ (30)

To make the non-differentiable spiking activation compatible with the back-
propagation algorithm, we considered a surrogate gradient:

$\displaystyle\mathrm{spike}^{\prime}(x)=\text{ReLU}(\alpha*(\mathrm{width}-|x|))$
(31)

where $\mathrm{width}=1.0$, and $\alpha=0.3$.

#### The readout layer.

The recurrent neurons project linearly to the output neurons. For the rate-
based model, the readout layer is given by:

$\displaystyle\mathbf{y}[t]=W^{\mathrm{out}}\mathbf{r}[t]+b^{\mathrm{out}},$
(32)

where $W^{\mathrm{out}}$ are the synaptic weights between the recurrent and
output neurons, and $b^{\mathrm{out}}$ the bias.

For the spiking network, the output neuron is the leaky neuron, whose dynamics
is given by:

$\displaystyle\tau_{\mathrm{out}}\frac{d\mathbf{y}}{dt}=-\mathbf{y}+W^{\mathrm{out}}\mathbf{z}+b^{\mathrm{out}},$
(33)

where $\tau_{\mathrm{out}}$ is the time constant of the output neuron,
$W^{\mathrm{out}}$ the synaptic weights between the recurrent and output
neurons, and $b^{\mathrm{out}}$ the bias. In the discrete description, the
output dynamics is written as:

$\displaystyle\mathbf{y}[t+\Delta
t]=\alpha_{\mathrm{out}}\mathbf{y}[t]+(W^{\mathrm{out}}\mathbf{z}[t]+b^{\mathrm{out}})\Delta
t,$ (34)

where $\alpha_{\mathrm{out}}=e^{-\frac{1}{\tau_{\mathrm{out}}}\Delta t}$.

#### Weight initialization.

Initial input, recurrent, and readout weights were drawn from a Gaussian
distribution $W_{ji}\sim\sqrt{\frac{s}{n_{\mathrm{in}}}}\mathscr{N}(0,1)$,
where $n_{\mathrm{in}}$ is the number of afferent neurons, $\mathscr{N}(0,1)$
is the zero-mean unit-variance Gaussian distribution, and $s$ is the weight
scale. For inhibitory neurons, $s=0.2$; for excitatory neurons, $s=0.05$.

#### Training methods.

The training was performed using the BPTT algorithm. The integration time step
$dt$ is 1 ms for the spiking neural network, while $dt$ 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, we add an additional voltage
regularization loss that penalizes voltages that fall significantly outside
the support of the surrogate gradient function (Plank et al., 2021):

$\displaystyle
L_{\mathrm{voltage}}[t]=\lambda_{v}\sum_{i}\sum_{t}(\text{ReLU}(V^{i}[t]-0.4))^{2}+(\text{ReLU}(-V^{i}[t]-2.0))^{2},$
(35)

where $\lambda_{v}$ is the strength of the voltage regularization.

#### The performance of the rate-based model.

In contrast to the spiking-based model, the rate model employed here exhibits
a slow convergence of training (see Figure S11A). It requires dozens of epochs
of training to achieve a high accuracy. Figure S11 B 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 a high
firing rate.

Figure S11: The rate-based recurrent model for implementing a DMS working
memory task. (A) The training accuracy of the rate-based recurrent neural
network for performing a 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.