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

Weakness 1:

We thank the reviewer for raising concerns about the scaling experiments presented in Figure 5. To address these concerns, we have conducted additional analyses and experiments to better demonstrate scaling behavior across orders of magnitude.

First, in addition to Figure 5 A/B, Figure 5C already shows extensive characterization of scaling across several orders of magnitude from $10^4$ to $10^6$. To provide another perspective on this scaling behavior, we have added a new Figure S16 A depicting the data from Figure 5C on a log-scale x-axis.

Second, while Figures 5D/E may appear to show linear scaling without spanning orders of magnitude, this was not our intention. We observed that reservoir performance converges and stabilizes around $5\times 10^4$ units, especially for MNIST dataset. Ignoring larger sizes seemed appropriate for this reason. Additionally, the scaling of reservoirs is not only constrained by weight matrix but also training algorithms. For fast convergence, we used FORCE learning in our current manuscript, which requires large matrices. However, we can also use memory-efficient algorithms like gradient descents. We have now implemented gradient descent-based training (see the attached code), enabling scaling up to $10^6$ units.

Weakness 2:

We thank the reviewer for their valuable feedback and comments.

Inherent from the JAX capability, BrainPy naturally supports the distributed simulation or training of biological neural networks with classical parallelization methods, like jax.vmap, jax.pmap, and jax.pjit. BrainPy models have been designed for supporting such methods. For example, most models in the brainpy.dyn module support to receive one argument sharding, which is used to indicate the sharding strategy across multiple devices (for example, hh = brainpy.dyn.HH(…, sharding=[‘x_axis’, ‘y_axis’])). We also supplement an experiment to demonstrate the distributed computation support of BrainPy in Appendix L. However, this support is nothing new since they have been widely adopted in training deep neural networks, so we did not include these descriptions in limited space.

Moreover, we are working on the special support for distributed training of spiking neural network by fully considering the sparse and event-driven property of spikes. These results have not been released and are under intensive development and tests.

Weakness 3:

We appreciate the reviewer highlighting the need for more concrete software architecture and implementation details in Section 4.

In the main text of section 4, our primary focus was to highlight the new methodology contributions that BrainPy brings, rather than diving into specific details about the concrete design and implementation. For the details of these methodology, we encourage readers should refer to the supplementary files.

However, we agree including more specifics on the codebase itself would strengthen the paper. To address this, we have included more concrete information about the package structure in the supplementary materials. Specifically, Appendix M has been included to illustrate the package's structure and components, such as the core module, models module, analysis module, and utils module. This section will help readers visualize the organization of the package and how the different modules interact with each other.

Question 1:

We have upgraded NEST version to 3.6 version and post the result in Figure 4. The code for NEST simulation also append in the supplementary code. For the implementation of the EI model using JIT operators, please refer to Appendix N.

Question 2:

The configuration we adopted in our study was based on the EI balance network scaling experiment in the NEST. Usually, the scaling experiment for an EI balance network is carried out by rescaling the connection probability instead of the weight, e.g., p = 80. / N, where N is the number of neurons while the weight values remain the same in all configurations. Such an approach was used by Stimberg et al. (2020), since it retains the firing frequencies at a comparable level when the network size varies. Moreover, increasing the number of synapses per neuron will attain a similar evaluation performance, see Figure S16 B.

However, our methods can also be adopted to the scenario that a fixed connection probability is used and the weight is rescaled. We carried out this experiment in Figure S16 C. BrainPy also illustrates a very good scaling property.

Weakness:

Thank you very much for recognizing our work! ICLR is a top-tier conference attracts researchers from diverse backgrounds in biological and artificial neural network fields. Given the usefulness of BrainPy, we believe it deserves exposure to this broader range of users. Previous conferences have also featured articles pertaining to software libraries, like Betty (ICLR 2023).

Question 1:

It is true that the JIT weight generation only works for random weights (but the JIT connection (indicating connected or not connected) can be generated by a broad fixed rule). Therefore, the connections weights have to be static so synaptic plasticity rules are not suitable for the JIT connection operators. For weights with plasticity, we encourage to use sparse and event-driven operators.

Question 2:

We thank the reviewer for raising this important problem. We have revised the manuscript and fixed the missing explanations in Appendix I.

Particularly, for the forward spike function $\text{spike} = H(V(t) - V_{th}) = H(x)$, we use $x$ to represent $V(t) - V_{th}$. The gradient of $\mathtt{spike}$ is $\text{spike}'(x) = \mathbf{ReLU} (\alpha * (\text{width} - |x|))$, where $x = V(t) - V_{th}$, $\alpha$ is 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 S11.

Question 3:

We thank the reviewer for this constructive question. The event-driven operators can take advantage of the sparse activity during the forward process by using the event-driven computations. However, since the derivative of the spike is a real value, the backward of such an event-driven operator is implemented using normal sparse operators. Therefore, the training phase of a spiking neural network using back-propagation-based algorithms is not very efficient. A forward-only algorithm will significantly accelerate the training of large-scale spiking neural networks.

Weakness 1:

The compilation overhead of BrainPy is significantly lower compared to traditional brain simulators, as demonstrated through two additional experiments (see Appendix K).

Firstly, when simulating a classical EI balance network model to mimic local circuits, BrainPy exhibits much faster compilation compared to Brain2. This is because BrainPy has precompiled most of the operators used in this model, allowing the JIT compiler to simply merge them into a single computational graph. In contrast, Brain2 needs to compile the entire generated C++ model from scratch, resulting in a significantly slower compilation phase. The comparison results are presented in Figure S13.

Secondly, when simulating a multi-area connected network, BrainPy outperforms Brian2 by several orders of magnitude in terms of compilation speed. While the model's runtime in BrainPy is in minutes, Brian2 takes hours to compile the same model. BrainPy achieves this by efficiently merging scattered operations among multiple brain areas into a single operator using jax.vmap or jax.pmap.

Overall, BrainPy leverages modern JIT compilation technology to offer a significant advantage in compiling brain dynamics models compared to traditional brain simulators.

Furthermore, the compilation time in BrainPy primarily depends on the complexity of the computation graph rather than the number of neurons. We have also observed that the compilation speed exhibits minimal variation across different supported computing devices, including CPU, and GPU. Compared to the runtime, the compilation in BrainPy is usually very low. These results are detailed in the Supplementary Materials, which can be found in Figure S13.

Weakness 2:

Thank you for your question! The JIT connection operators in BrainPy do not rely on the assumption of a low-rank connectivity matrix. Instead, they assume that the connectivity is static and can be generated using a fixed rule, such as random connections with a given probability or regular grid patterns. With this assumption, the JIT connectivity algorithm regenerates the connectivity weights when a presynaptic neuron fires during computation, rather than storing the synaptic connectivity in memory. This approach has the advantage of not consuming additional device memory. However, it is important to note that JIT connection operators are only suitable for static synaptic connections.

We would like to emphasize that JIT connection operators contribute to the advancement of large-scale brain simulation and provide a means to address some of the challenges in this field. However, they are not the ultimate solution for simulating the 86 billion neurons of the human brain, as neurons and synapses are inherently plastic and subject to dynamic changes.

Minor Weakness 1:

We thank the reviewer for raising this important point! We are sorry for the unclear statements and we would like to clarify our contribution on JIT compilation. To address this, we propose updating Section 4.5 to explicitly state that BrainPy leverages JAX for just-in-time compilation. Particularly, we would like to edit:

Particularly, BrainPy integrates a collection of object-oriented JIT transformations into its multi-scale model-building interface. These transformations are built upon JAX's implementation of a function-oriented JIT interface, as detailed in Appendix G. By leveraging these transformations, any BrainPy model constructed using the modular and composable programming interface can be effortlessly converted into an optimized XLA process, compatible with CPU, GPU, and TPU platforms.

Additionally, we will comb through the corresponding Appendix to ensure there is no language that wrongly implies we built JIT functionality from scratch.

We appreciate you catching this ambiguity - thank you again for the thoughtful feedback.

Minor Weakness 2:

We thank the reviewer catching up this point! It is indeed accurate that NetPyNE provides a hierarchical interface on multiple levels: ranging from the molecular to the network level. We agree that our claim is too strong and the statement as written may be misleading. However, we would like to highlight again that BrainPy's hierarchical composition capabilities provide much greater flexibility compared to the predefined interface of NetPyNE. In NetPyNE, users can only construct models based on the hierarchical levels that are already defined. Conversely, BrainPy empowers users to create their own customizable hierarchical structures and combine multiple levels of brain models and functions, offering a more versatile approach to hierarchical composition.

Weakness 1

We disagree the reviewer's claim. BrainPy is an elaborate package that provides a collection of interconnected utilities designed to provide foundational services that enable users to easily, flexibly and efficiently perform various types of modeling for brain dynamics. Specifically, BrainPy implements

1. dedicated operators for event-driven computation based on sparse connections and JIT connectivity;
2. numerical integrators for various differential equations, the backbone of dynamical neural models;
3. a universal model-building interface for constructing multi-scale brain dynamics models and the associated JIT compilation for the efficient running of these models;
4. a toolbox specialized for brain dynamics modeling;
5. systematic components and abstractions for diverse brain dynamics models.
6. and so much more.

We also want to clarify that we have made a great effort on engineering quality and documentations including:

1. For maintainability, we organized the codebase into a modular structure with separate subpackages for the main components. This improves and enables future extensibility.
2. For understandability, we added type hints and docstrings throughout the code for clarity. And provided user-facing API documentation with usage examples.
3. For stability, we set up continuous integration pipelines to run tests and checks on every commit.
4. Moreover, we provided comprehensive and well-organized documentation that includes clear installation instructions, usage tutorials, and detailed API reference. We also have documentations for many examples written in BrainPy including multiple modeling levels ranging from detailed neuron models to large-scale whole brain networks.

We appreciate the reviewer's suggestions, and we will continue to improve the engineering quality and documentations to make BrainPy more flexible, extensible, and scalable, and to enrich the resources for a better user experience.

Weakness 2

We thank the reviewer highlighting BrainPy's distinct capabilities compared to existing BIC libraries. Unlike existing BIC libraries, BrainPy goes far beyond the simulation of neurons. Spiking neurons (like LIF) + deep learning components are the main paradigms of current BIC libraries, just like in SpikingJelly. However, BrainPy aims to provide a general-purpose brain dynamics modeling toolkit for diverse aspects of brain dynamics. As the reviewer notes, support in existing BIC libraries is lacking for key facilities needed to efficiently simulate detailed biological neural networks, like event-driven computing, sparse representations, and scalable infrastructure.

BrainPy's goal is to provide flexible, general-purpose toolkits tailored for modeling the rich dynamics found in biological neural systems, spanning complex neuron types, synaptic interactions, ion channels, and specialized computing techniques. Specifically, BrainPy offers critical modeling infrastructure oriented for biological plausibility, including: sparse, event-driven operators; scalable just-in-time connectivity; multiscale interfaces bridging levels of detail; whole-graph just-in-time optimization leveraging object-oriented design; easy Python-based operator customization; and more. With these expansive capabilities covering both biological accuracy and computational performance, BrainPy provides a foundation for efficient and scalable modeling of intricate brain dynamics, complementing existing tools focused primarily on simulating deep learning-oriented spiking neural networks. We welcome any further discussion to clarify BrainPy's approach in context of the current BIC landscape.

Weakness 3

Thanks for the question. The unique contribution and dedicated optimization of BrainPy has been elaborated in Section 4 Methodology. For the details of such optimization, we encourage to refer to the supplementary matrials.

Weakness 1:

We argue that ICLR is embracing the integration of software platforms, neuroscience, and deep learning, which is evident from its call for papers explicitly including these topics. Past conferences have also showcased articles focusing on software libraries, such as the introduction of Betty at ICLR 2023. Notably, BrainPy offers a comprehensive framework for brain dynamics modeling, encompassing systematic abstraction of brain dynamics compilation, operators, and models. Therefore, it aligns well with the specific requirements of ICLR.

Weakness 2:

We have carried out experiments for demonstrating the performance of BrainPy on BIC models.

Particularly, we conduct a comparative analysis with several popular SNN training frameworks, including snnTorch, Norse, and SpikingJelly. We first use a simple three-layer SNN model, where the input layer delivers synaptic currents to the hidden LIF layer with exponential dynamics, then the output layer readouts the hidden dynamics with an exponential synapse. The training was performed using the BPTT algorithm on the Fashion-MNIST dataset. The comparison result of the training speed per epoch is presented in Table S2. We find that BrainPy exhibits superior performances on both CPU and GPU devices. Furthermore, to test the performance of BrainPy on large-scale SNN networks, we re-implement a VGG SNN network which was originally coded in a PyTorch+SpikingJelly environment. The evaluation result is shown in Table S2. Under the same training setting, BrainPy achieves much better running performance. It reduces almost half of the time for training on GPUs.

Weakness 3:

Due to space constraints, our initial submission focused solely on optimizations enabling efficient biological spiking network simulations. However, seamless integration with deep learning models is also a key goal within BrainPy's unifying framework. To clarify BrainPy's flexible deep learning support:

First, BrainPy implements a range of fundamental deep learning layers (e.g. convolutions, pooling, normalization) within the brainpy.dnn module. Unlike JAX, these models employ a PyTorch-style class-based interface for intuitive usage - identical to other BrainPy dynamics models.

Second, BrainPy enables effortless interoperation with existing JAX machine learning libraries like Flax. The brainpy.layers.FromFlax functionality, for instance, allows any Flax model to be utilized as a BrainPy module. Moreover, BrainPy models can also be converted to formats compatible with external libraries.

Question 1:

Parameter sharing in AlignPre and AlignPost is actually a technical design for reducing the memory and computing complexity that do not affect the dynamical behavioral of the system. There is currently no direct biological evidence to support this specific correspondence. However, some general neuroscience facts/principles that might motivate these operations include:

• Synapses of the same type projecting from the same axon tend to have similar temporal dynamics. For example, AMPA synapses originating from one pyramidal neuron typically have similar decay time constants in the range of 2-10 ms. This supports using a shared variable for pre-synaptic traces in AlignPre.
• Synapses of the same type converging onto a dendrite often exhibit similar temporal dynamics. For example, the time course of NMDA synaptic currents measured at a dendrite is relatively consistent, supporting a shared variable for post-synaptic traces in AlignPost.
• Biological evidence shows synaptic strength heterogeneity is often higher than dynamics heterogeneity within a projection. AlignPre and AlignPost separate strength (weights) from dynamics, capturing this aspect.

In summary, the AlignPre and AlignPost design exploits the biological evidence of synapses having consistent dynamics but heterogeneous weights within common projections. This allows capturing the brain's connectivity constraints elegantly while optimizing computation and memory.

Question 2:

After thoroughly examining the CARLsim documentation, we unfortunately could not identify an implementation of an E/I balanced LIF network amenable to equal benchmarking.

While CARLsim does offer a Python interface through pyNN, its core simulations utilize customized C++ and CUDA kernels for efficiency. After extensive searching, we could not locate Python-accessible LIF neuron models within CARLsim for viable testing. As the reviewer notes, developing an appropriate comparison would require a non-trivial integration effort.

Given the time-constrained rebuttal period, implementing and optimizing an E/I LIF network in CARLsim for fair benchmarking is unfortunately not feasible. However, we absolutely agree that a CARLsim comparison would provide useful additional context and validity to BrainPy's performance claims.