Paper: A Parallel Multi-compartment Spiking Neuron For Multi-scale Sequential Modeling (ICLR-2024)

Reviewer comment (Theory):
The implementation seems to require feed-forward structure as well as positive inputs, which seems quite restrictive (in particular the latter).

Author response:
We greatly appreciate your time in reviewing our manuscript and providing constructive feedback. As you have rightly pointed out, we did assume that $I_h[t]$ to be non-negative in Equation 15. This is a compromise to enable the parallelization of the reset mechanism, and we provide a detailed proof of Equation 15 based on this assumption in Appendix A.4.2.
However, it's worth mentioning that this assumption only applies to the last compartment. It does not require the input to our PMSN model, namely $I(t)$, to be always positive. In fact, the negative input current to other compartments (i.e., 1 to n-1) in Equation 5  is allowed. As specified below the Equation 3, its value is formulated as $I^l(t)=\mathcal{W}^l S^{l-1}(t)$. Therefore, we **don’t need to impose any constraint on the input signal $I(t)$ of the PMSN model**, but only apply a clamp function on the intermediate states (output of Equation 12) before it is fed into Equation 15. We will highlight this point in our revised manuscript to avoid confusion.

In the following, we would like to discuss the **impact of removing this assumption**. Firstly, according to our analysis, if negative inputs are allowed for the output compartment, the state update of the last compartment can only be done in a serial manner, while other compartments can still be computed in parallel. Following your suggestion, we have conducted an ablation study to quantify its impact on the simulation speed. The time elapse (seconds/epoch) for PMSN and last-compartment-serial PMSN are 62 vs 158 on sequential CIFAR-10 (T=32), and 73 vs 5182 on sequential CIFAR-10 (T=1024) tasks. These exploded simulation time underscore the necessity of the proposed parallel solution. Interestingly, we also observed that the last-compartment-serial PMSN model, which allows negative input for the last compartment, performs worse than the PMSN model. It exhibits a 3.98% and 4.5% accuracy drop on sequential CIFAR-10 and sequential CIFAR-100 (T=32) tasks, respectively. We suspect that negative input can result in negative membrane potentials that adversely affect the gradient flow. This is because when the membrane potential is far from the firing threshold, the gradient g’[t] computed from the surrogate function will become zero. This issue hinders the effective gradient propagation, therefore, leads to poorer classification accuracies than our proposed one.

Reviewer comment (Evaluation):
The neuron model does not seem to improve the benchmark results dramatically over other approaches in all tasks (ie recurrent RNNs are similar for S-MNIST) and it thus remains open how useful the model will be.

Author response:
We appreciate your insightful comment. While it is true that the performance enhancement of our model over recurrent SNNs may not seem dramatic in the context of the S-MNIST task, it’s worth noting that this task is relatively easy and models can perform well with limited memory capacity. The existing methods have already reached a very high accuracy and the margin for improvement is limited. Specifically, the state-of-the-art (SOTA) model TC-LIF with an RSNN architecture achieves an error rate of just 0.8%. Nevertheless, our model managed to reduce this error rate to 0.47% further.

We understand that the S-MNIST dataset alone may not fully demonstrate the potential benefits of our proposed model. Thus, to further illustrate the enhanced sequential modeling capabilities of our proposed model, we have extended our comparisons to include more challenging sequential tasks. In these scenarios, our model consistently outperforms other SOTA models by significant margins. For instance, in the SHD speech command recognition task, our model achieves an accuracy improvement of 1.8% and 5.3% over the SOTA Adaptive Axonal Delay (feedforward) and TC-LIF (recurrent) model respectively, while maintaining a comparable number of parameters. Moreover, in the Sequential CIFAR10 (T=32, 1024) and Sequential CIFAR100 (T=32) tasks, our model surpasses the accuracy of other SOTA models by 2.52%, 11.91%, and 3.87%, respectively. It is worth noting that, the accuracy improvements increase along with the increasing complexity of the task.

We believe these results underscore the potential of our model in handling more complex sequential modeling tasks, and we look forward to exploring its capabilities further in future research.

Reviewer comment (Theory):
I don't follow why one can assume that $I(t)$ is always positive for the “parallel implementation of the output compartment with reset”. In typical SNNs, the input current is given by $I(t)=WS(t-1)$ (without a synapse model) and the weights are positive or negative real numbers, so in general, the inputs *can* be negative. Isn't that assumed here? However, if negative inputs are allowed, then the floor of the accumulated input with the threshold $\theta$ (in Eq 15) will not result in the number of spikes over the time period. Maybe I am missing something here. If only positive currents are allowed, it should be discussed as it seems to be a major restriction. If so it would be interesting to know what the impact of this step would be (on the simulation speed) if not done in parallel to allow for negative currents.

Author response:
We greatly appreciate your time in reviewing our manuscript and providing constructive feedback. As you have rightly pointed out, we did assume that $I_h[t]$ to be non-negative in Equation 15. This is a compromise to enable the parallelization of the reset mechanism, and we provide a detailed proof of Equation 15 based on this assumption in Appendix A.4.2.
However, it's worth mentioning that this assumption only applies to the last compartment. It does not require the input to our PMSN model, namely $I(t)$, to be always positive. In fact, the negative input current to other compartments (i.e., 1 to n-1) in Equation 5  is allowed. As specified below the Equation 3, its value is formulated as $I^l(t)=\mathcal{W}^l S^{l-1}(t)$. Therefore, we **don’t need to impose any constraint on the input signal $I(t)$ of the PMSN model**, but only apply a clamp function on the intermediate states (output of Equation 12) before it is fed into Equation 15. We will highlight this point in our revised manuscript to avoid confusion.

In the following, we would like to discuss the **impact of removing this assumption**. Firstly, according to our analysis, if negative inputs are allowed for the output compartment, the state update of the last compartment can only be done in a serial manner, while other compartments can still be computed in parallel. Following your suggestion, we have conducted an ablation study to quantify its impact on the simulation speed. The time elapse (seconds/epoch) for PMSN and last-compartment-serial PMSN are 62 vs 158 on sequential CIFAR-10 (T=32), and 73 vs 5182 on sequential CIFAR-10 (T=1024) tasks. These exploded simulation time underscore the necessity of the proposed parallel solution. Interestingly, we also observed that the last-compartment-serial PMSN model, which allows negative input for the last compartment, performs worse than the PMSN model. It exhibits a 3.98% and 4.5% accuracy drop on sequential CIFAR-10 and sequential CIFAR-100 (T=32) tasks, respectively. We suspect that negative input can result in negative membrane potentials that adversely affect the gradient flow. This is because when the membrane potential is far from the firing threshold, the gradient g’[t] computed from the surrogate function will become zero. This issue hinders the effective gradient propagation, therefore, leads to poorer classification accuracies than our proposed one.

Reviewer comment (Reproducibility):
It is not clearly stated whether the flux parameters and decay rates of and between the compartments are learned with SGD.

Author response:
Thank you for pointing this out. In this work, the flux parameters and decay rates are learned jointly with weight parameters using AdamW Optimizer. We have mentioned this point below Equation 9, and we will further emphasize their biophysical meanings in the revised version of our manuscript. Furthermore, we will provide a detailed summary of our training configurations in our Appendix and make our code publicly available after the revision process to ensure reproducibility and facilitate understanding of our work.

Reviewer comment (Writing):
In the equations (eg Eq. 3) the neural indices are written as super-script ($v^{i}$) which is easily confused with power. Better to use subscripts or write $v^{(i)}$ to avoid confusion.

Author response:
We sincerely appreciate your suggestion. We will revise these indices in our manuscripts to improve clarity.

Reviewer comment (Experiments):
Regarding Figure 4: while the parallel model was compared to its serial implementation in terms of speed up, what is the ratio/training time difference between PMSN and PSN? If PSN is faster, since it is only a single compartment model, how significant is this difference? I am not referring to the computational cost but instead training acceleration only.

Author response:
Thanks for your insightful comment. Following your suggestions, we have conducted a comparative analysis of the training speed between the PMSN and PSN models. This involves measuring each model's training time across various tasks with different time window lengths, with identical training configurations and network architectures for a fair and direct comparison. The results, including time elapsed (seconds/epoch) and classification accuracy (%), are as follows:

| Dataset | S-MNIST (T=784) | Sequential CIFAR10 (T=32) | Sequential CIFAR10 (T=1024) | 
| ------- | --------------- | ------------------ |-------------------- | 
| Speed (PMSN) | 35.3235 | 61.7562 | 73.2060 | 
| Accuracy (PMSN) | 99.40% | 90.97% | 82.14% | 
| Speed (PSN) | 21.1776 |  34.0634 | 49.0240 |
| Accuracy (PSN) | 97.90% | 88.45% | 55.24% | 

The results clearly show that the PMSN model, with more complex and sophisticated multi-compartment dynamics (five compartments in PMSN vs. one in PSN), trains slightly slower than the PSN model. However, the difference in training speed is not significant (with a ratio of less than 2 across all the experiments), demonstrating the efficiency of our parallelization approach. Furthermore, this minor trade-off in speed is outweighed by the superior sequential modeling capabilities of PMSN, particularly evident in its performance on the Sequential CIFAR-10 dataset. This underlines the value of the PMSN model in applications where enhanced temporal signal processing capacity is critical.

Reviewer comment (Reproducibility):
Is the parallel implementation/methodology that incorporates the reset mechanism general enough to be applied to single compartment models?

Author response:
We appreciate your insightful question. This reset mechanism can indeed be seamlessly transferred to the single-compartment Integrate-and-Fire (IF) model, thereby enabling parallelizable training for this model. The generalization is straightforward and involves a simple modification. Specifically, it requires replacing the input of Equation 13, denoted as $I_h[t]$, with the input of the single-compartment model. For instance, the replaced input can be denoted as $I^l[t]=\mathcal{W}^l S^{l-1}[t]$, which adheres to a common practice of single-compartment models regardless of the dynamics of synaptic decay. Consequently, Equations 13 and 14 effectively formulate a parallel IF model that incorporates our reset mechanism.