P-SPIKESSM: HARNESSING PROBABILISTIC SPIKING STATE SPACE MODELS FOR LONG-RANGE DEPENDENCY TASKS

Thank you for your insightful questions, suggestions and positive feedback. We have addressed your comments below and incorporated your suggestions in the revised version of our paper.

Comment 1:

IMHO, the main weakness is that this network is not fully spiking, and it's not clear if it could be implemented on existing neuromorphic chips (e.g., Intel Loihi). The temporal convolution of the SSM (eq 6) could probably be implemented by delays. But the gelu in eq 10, and the non-local normalization (in ClampFused layers) are problematic. I suggest the authors discuss these issues and possible solutions. Also, I think the authors should cast their network as "Partial SNN" in Table 1-3

Response: Thank you for your insightful comment. To further validate our model's biological plausibility, we replaced the GELU activation in Eq. 10 with a hardware friendly ReLU function [1]. The results of this modification have been included in our ablation studies. Notably, this substitution had minimal impact on accuracy; for example, on the LRA Text task, the model achieved 80.4% accuracy with ReLU compared to 81.2% with GELU. While ReLU offers a straightforward alternative to GELU, future work could focus on developing biologically plausible variants of GELU that are compatible with neuromorphic hardware. In the ablation studies we have also performed an experiment by removing normalization during training.

Comment 2:

IMHO psMNIST is not challenging enough (SOTA is nearly 100%, which may mask differences between approaches); I would recommend trying as well on sequential CIFAR10/100 (much more challenging).

Response: Thank you for your suggestion. We tested our model on the sequential CIFAR-10 dataset, achieving an accuracy of 85.6%. This result has been added to the Additional Experimental Results section in the Appendix for reference.

Comment 3:

The authors seem to ignore the sliding PSN neuron (http://arxiv.org/abs/2304.12760), which has many similarities with their model. Both use spikes. Both use temporal convolutions instead of stateful units and for this reason, both avoid BPTT and are parallelizable. A comparison with the PSN (both qualitative and quantitative) would be useful.

Response: The research on PSN neurons is indeed fascinating, and we have now included a discussion of this work in the Related Works section of our paper.

Comment 4:

I think the authors could bypass the SpikeSampler layer and send the real-valued probabilities directly to the next layer. Do you confirm? Have you tried? Of course, the computational advantages of spikes would be lost, but I expect an increase in accuracy, and it would be interesting to quantify it.

Response: Thank you for your insightful observation. As you rightly pointed out, directly passing the probabilities does lead to higher overall accuracy. For instance, on the seq-CIFAR10 task, using real-valued probabilities yields an accuracy of 87.2%. However, as you mentioned, this approach sacrifices the energy efficiency advantages provided by the spike sparsity in our original model.

Comment 5:

Would it be possible to encourage even sparser activity via some additional term in the loss function?

Response: Thank you for this excellent observation. Although we did not explore this aspect in our paper, incorporating sparsity regularization techniques [2] could be a promising direction to further reduce spiking activity. This approach would enhance the computational efficiency of the model by promoting sparsity.

Suggestions

"A is a parameter controlling the evolution" -> "A is a matrix controlling the evolution"

Response: We have modified the paper accordingly.

the y dimension is useless here. I would recommend plotting the curve y = p(t) instead of a heat map.

Response: Thank you for the insightful suggestion. We have updated the figure in our paper accordingly.

References:

[1] Timcheck, Jonathan, Sumit Bam Shrestha, Daniel Ben Dayan Rubin, Adam Kupryjanow, Garrick Orchard, Lukasz Pindor, Timothy Shea, and Mike Davies. "The intel neuromorphic DNS challenge." Neuromorphic Computing and Engineering 3, no. 3 (2023): 034005.

[2] Yan, Y., Chu, H., Jin, Y., Huan, Y., Zou, Z. and Zheng, L., 2022. Backpropagation with sparsity regularization for spiking neural network learning. Frontiers in Neuroscience, 16, p.760298.

Thank you for your insightful suggestions and valuable feedback. In this rebuttal, we have addressed your comments.

Comment 1:

P-SpikeSSM is an interesting exploration, but the elements that play a key role in it seem to be just an application of what SSM proposes, such as temporal convolution to enable parallel computation.

Response: Thank you for this insightful comment. While our approach leverages the dynamics of an underlying SSM, our computational model diverges significantly from prior SSM-based architectures. First, unlike previous SSM-based models that process real-valued data, our SSM-based neuronal model operates on sequences of spikes. This distinction allows for a different processing paradigm. Additionally, while previous typical spiking models simply applied an LIF neuron to the output of an SSM to generate spikes, we introduce a probabilistic mechanism for spike generation based on the underlying SSM-based neuronal model. This leads to several key advantages: (a) it achieves improved performance than the baselines, (b) it enables parallelizable training using a surrogate function while leveraging a probabilistic spike sampling approach, and (c) it results in sparse spiking patterns, yielding significant computational efficiency, as detailed in Section 4.2 and illustrated in Fig. 4.

Comment 2:

The network cannot be considered a fully spiking neural network, and is best called a partial SNN because of the use of gelu in SpikeMixer Block (Eq. 10). In addition, the introduction of gelu greatly increases the difficulty of deploying this network on neuromorphic hardware.

Response: Thank you for highlighting this aspect. In response, we explored replacing the GELU layers with the on-chip-compatible ReLU function [3]. The results of this modification are included in our ablation studies. Notably, the substitution had minimal impact on accuracy; for instance, on the LRA Text task, the model achieved 80.4% accuracy with ReLU, compared to 81.2% with GELU. While ReLU serves as a straightforward alternative to GELU, future work could focus on developing biologically plausible GELU variants that are optimized for compatibility with neuromorphic hardware.

Comment 3:

Considering that the network is a partial SNN, it performs only marginally better than BinaryS4D, which is also a partial SNN, on some of the LRA benchmark tasks, and even worse on others (Table 2). Given that P-SpikeSSM and BinaryS4D are both based on SSM, it's debatable whether their advantages over other transformer-based networks mostly stem from SSM.

Response: The primary distinction between our proposed P-SpikeSSM and other SSM based approaches, including the BinaryS4D model, lies in the computational paradigm: our neuronal model's SSM operates on sequences of spikes, unlike the real-valued inputs used in BinaryS4D. This design offers a significant computational advantage, as shown in Eqn. [6], by performing convolution operations on sparse spiking inputs, enabling efficient and parallelizable inference. In our approach, the underlying SSM serves as the core neuronal model, responsible for generating probabilities for spike sampling. This eliminates the need for an additional LIF neuron layer, which is a requirement in BinaryS4D. Following the review comments, replacing GELU with ReLU enables our model to be implemented using operations that are compatible with neuromorphic chips [3].

Moreover, with further hyperparameter optimization, our model achieves improved performance on the LRA Text dataset, thus outperforming BinaryS4D on 4 out of 5 datasets in the LRA benchmark. Additionally, the sparse spiking activity throughout our model, as discussed in Section 4.2, provides substantial computational efficiency. Unlike BinaryS4D, which primarily utilizes spiking inputs only in the linear layers, our approach uses spikes for communication across all layers thereby not compromising on computational benefits offered by spiking architectures.

Comment 4:

For the experiment on the Speech Command dataset (Table 3), Spiking LMUFormer achieved 96.12% accuracy on the full 35-category dataset, whereas this work is only validated on a subset of 10 categories and is not compared to some recent advanced baselines [1, 2].

Response: Thank you for highlighting this aspect. In response, we ran our model on the full 35-category dataset and achieved an accuracy of 96.23%. This surpasses the performance of [2], which reports 94.51% on the same dataset and that of SpikingLMUFormer [1]. We have included these updated results in the revised paper (Section 4.1).

Thanks for your feedback. These clarifications and additional results have addressed most of my concerns. However, I still think that the performance of the model is mainly due to SSM and I'm not sure how significant the use of spikes is for the development of the research community. Overall, I believe that the work falls short of the acceptance threshold.

We sincerely thank you for your positive review and constructive feedback. We are thrilled that you appreciated our work. Below, we have provided our responses to your comments.

Comment 1:

On line 141 and the caption for Figure 2, has it truly been proven to be the case that LIF neurons cannot be parallelized? What is the distinction between a non-linear LIF neuron and a linear LIF neuron? (To me, all LIF neurons are non-linear because they have reset terms.) Do efficient algorithms exist for LIF neurons? E.g., see “EXODUS: Stable and Efficient Training of Spiking Neural Networks” by Bauer et al., or the work SLAYER that they reference?

Response: Thank you for your comment. The stateful nature and reset mechanism of the vanilla LIF neuron pose challenges to parallelization. Some implementations address this by removing the reset mechanism to create a linear version of the LIF neuron. In line with works such as EXODUS (which leverage Implicit Function Theorem), another promising learning framework for LIF neuron based SNNs is implicit differentiation at equilibrium [1], which has shown impressive empirical performance. However, [1] has primarily been explored in vision-based tasks and has not been explored in the context of complex long-range dependency tasks, such as those in the LRA benchmark. We have added this discussion in the revised paper.

Comment 2:

On line 77, the authors claim the surrogate operation Expection[S_t] is novel. I believe such a surrogate operation has already been introduced. E.g., see “Automatic Differentiation of Programs with Discrete Randomness” by Arya et al. NeurIPS 2022.

Response: Thank you for highlighting this interesting work. [2] primarily focuses on stochastic behavior in the domain of automatic differentiation, which is used to construct new programs that compute the derivative of an original program. In contrast, our approach analyzes the stochastic nature of spiking events within the context of SNNs. We specifically leverage Expectation[S_t], where S_t represents the random variable for the spiking event, to create a surrogate function that enables efficient and parallel training.

Comment 3:

In Table 1, it is unclear to me the model size/computation budget of these various models. I could imagine model size/compute budget is a key reason why one model would outperform another, so I would like to clearly understand how this table presents a fair comparison. E.g., could you show me how the P-SpikeSSM model is iso-parameter-count?

Response: Thank you for this comment. We have added a detailed memory footprint analysis in Appendix D. On average, our models have approximately 250K parameters across all five LRA tasks, achieving a 2.4× reduction compared to the 600K parameters of the baseline transformer architectures. Our model also demonstrates substantial improvements over the LIF-based approach in both training time and inference speed, achieving superior efficiency during training and inference. Training one epoch (batch size 64) takes approximately 1.5 minutes with our model, compared to 6.1 minutes for the LIF-based variant on the ps-MNIST dataset. Similarly, inference on the test set requires only 9 seconds for our model versus 17 seconds for the LIF-based approach. This difference stems primarily from the sequential processing bottleneck inherent in LIF-based models, which also depends on Backpropagation Through Time (BPTT) for training. In contrast, our framework leverages a parallelizable architecture, enabling efficient single-pass training with standard backpropagation. All experiments were conducted on an Nvidia RTX A5000 GPU with 24GB of memory, as specified in Section 4.This comparison with conventional LIF neurons has been included in the Ablation Studies section (Section 4.1.2), specifically in lines 459–470.

Comment 4:

On line 102, the authors imply that P-SpikeSSM is a “fully” spiking model. While “fully” spiking is not defined in the literature, one might argue that P-SpikeSSM is not fully spiking because each spiking neuron contains a non-spiking SSM within it. I might suggest dropping the “fully” adjective. On line 415, “fully spiking” is implied to be defined as “not requiring floating point MAC operations,” but I am now confused regarding how P-SpikeSSM is fully spiking even though there is an SSM inside every neuron.

Response: Thank you for this suggestion, we have modified our paper accordingly and refrained from using such a terminology. Though we have an SSM inside every neuron, we can represent the internal dynamics as a convolution over sparse input spikes (Eqn. 6) in the parallel convolution based inference phase, which allows us to avoid floating point multiplicative operations.

We sincerely thank you for your positive review and constructive feedback. We have added our response to your comments below.

Comment 1:

The authors do not provide strong evidence that the sampling-based technique outperforms the conventional LIF neuronal model. The absence of an input/output step may also hinder the flow of internal dynamics information. I suggest that they compare their architecture using both approaches to clearly demonstrate the benefits of their method. While it is evident that their approach facilitates the parallelization of the algorithm, this does not necessarily imply that the solution can be obtained faster, as more steps might be required to achieve the appropriate spiking rate.

Response: We implemented a version of our model using conventional LIF neurons to conduct an ablation study. The experimental setup mirrors the one outlined in Section 4.1 for the ps-MNIST dataset and both the models have similar number of parameters. Having trained both the models till convergence of validation loss (for 100 epochs), the LIF-based approach achieves a test accuracy of 97.7% on ps-MNIST dataset, compared to 98.4% attained by our stochastic model. For an iso-accuracy convergence comparison, our stochastic model achieves the same 97.7% after only 42 epochs of training whereas, the LIF neuron based model requires 76 epochs.

Notably, our model also significantly outperforms the LIF-based approach in training and inference efficiency. Training one epoch (batch size 64) takes approximately 1.5 minutes with our model, compared to 6.1 minutes for the LIF-based variant on the ps-MNIST dataset. Similarly, inference on the test set requires only 9 seconds for our model versus 17 seconds for the LIF-based approach. This difference stems primarily from the sequential processing bottleneck inherent in LIF-based models, which also depends on Backpropagation Through Time (BPTT) for training. In contrast, our framework leverages a parallelizable architecture, enabling efficient single-pass training with standard backpropagation.

All experiments were conducted on an Nvidia RTX A5000 GPU with 24GB of memory, as specified in Section 4. This comparison with conventional LIF neurons has been included in the Ablation Studies section (Section 4.1.2), specifically in lines 459–470.