Paper: A differentiable brain simulator bridging brain simulation and brain-inspired computing (ICLR-2024)

Reviewer comment (Experiments):
Except for Fig.5 A/B dealing with matrix multiplication, the figures vary the network/system size on a linear scale, and not over several orders of magnitude. This does not seem suitable to demonstrate the scaling behavior for these network models.

Author response:
**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.

Reviewer comment (Experiments):
There does not seem to be a demonstration of distributed simulation/training of a large-scale model on a CPU or GPU cluster.

Author response:
**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.

Reviewer comment (Presentation):
In section 4 the package is at length described as efficient, extensive, scalable, etc. without really getting into the concrete design and implementation. In particular, after reading this section I did not end up with a clear picture of the package structure and components. Maybe the description could be shortened, or made more concrete.

Author response:
**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.

Reviewer comment (Experiments):
On p.8 it is mentioned concerning Fig 5C that the E/I network was scaled up to 4 million neurons with 80 synapses each. In biological networks, the (local) connectivity is typically significantly more dense, with hundreds or thousands of synapses per neuron.

Author response:
**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.

Reviewer comment (Experiments):
It does not evaluate the cost of compilation. How high is the cost of compilation compared to the runtime? Is this a clear disadvantage as compared to, e.g., NEURON? Does the compilation speed depend on whether CPU or GPU are used? How does it scale with the number of neurons?

Author response:
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.

Reviewer comment (Theory):
In section 4.2, the authors claim that there method is significantly more memory-efficient than others. Maybe I am misunderstanding this, but: do the gains that the authors claim here stem from an assumption that the connectivity matrix is low-rank? How else would they possibly be able to store connections of 86 Billion neurons?

Author response:
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.

Reviewer comment (Novelty):
I believe that it would be good if the authors clarified that all JIT capabilities are due to the fact that brainpy relies on JAX, and are not implemented from scratch. Section 4.5 reads as if the authors implemented this themselves.

Author response:
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.

Reviewer comment (Novelty):
The statement `It is important to note that this hierarchical composition property is not shared by other brain simulators.` is not true, see for example NetPyNE.

Author response:
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.

Reviewer comment (Novelty):
Despite the exciting endeavor towards bridging the gap between brain simulators and BIC libraries, this paper appears to have limited relevance to this conference due to the lack of original theories or empirical evidence.

Author response:
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.

Reviewer comment (Experiments):
Quantitative comparisons with BIC libraries seem to be missing.

Author response:
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.

Reviewer comment (Experiments):
While the paper takes the stance of bridging brain simulators and DL frameworks, discussions about deep learning models seem to be missing.

Author response:
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.

Reviewer comment (Experiments):
How does BrainPy's speed and scalability compare to CARLsim, another brain simulator known for efficient and large-scale brain simulations?

Author response:
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.

Reviewer comment (Reproducibility):
The package itself is more like a collection of course scripts rather than a Python package. Thus I suggest that the authors improve the engineering quality and documents for the current package.

Author response:
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.

Reviewer comment (Novelty):
The comparison to existing methods is not sufficient. For example, there are existing tools like SpikingJelly for SNN. The simulation of neurons is also not sufficiently new. Thus the unique character of the current package could also be strengthened.

Author response:
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.