A Diagonal Structured State Space Model on Loihi 2 for Efficient Streaming Sequence Processing

1. While this work presents engineering value, its algorithmic contributions lack substantial novelty. The authors introduce a neuromorphic variant of the S4 model without normalization layers and residual connections. Although omitting these layers may not lead to significant accuracy drops on simpler datasets, such as s-MNIST and s-CIFAR, this approach may not scale effectively to more complex tasks requiring long-range dependency modeling

2. While its exciting that authors have directly reported results after running their model on Intel Loihi 2 chip, the empirical evaluation can be more rigorous. Instead of just MNIST, CIFAR the authors can evaluate running their model on the long range arena benchmark or Speech Command dataset. This is to better capture the long-range dependency capturing capability of the proposed architecture.

Despite the thorough and clear documentation of the work, the reviewer has to list the following concerns that underline the assigned score for this submission.

1. ICLR might not be the appropriate target for this work. The major contribution of the paper is the successful deployment of S4D on the Loihi 2 chip. With this implementation centric work, the paper might be better placed at a systems conference where much more space can be attributed to a detailed description of this core contribution (MLSys, AICAS etc). The algorithmic contributions are very incremental in terms of model architecture (replacing activation functions) and quantization. Why is the deployment the core contribution? Quantization methods for SSMs were investigated in [1.], and 8-bit quantization is an established standard in the ML community with most deployment frameworks supporting 8-bit precision at least for matmuls. The A, B, C matrices of SSMs might be a different matter, but does not provide data or comparison against established methods and references [1.] for this sake. At the same time, the task evaluation on sequential MNIST and CIFAR-10 is quite weak to claim contributions towards quantization methods. E.g. line 361f highlights that the accuracy drop introduced by 8-bit quantization is very small compared to the baseline. This is not a surprise on small scale datasets. We know for example from language models that naive quantization aware training works well up to a certain scale around 7B parameters and only breaks afterwards.
2. The hardware aware implementation is a strength of the paper. Unfortunately, a comparison to the obvious naive implementation that the authors discuss as well is missing. This would give readers a clearer picture of the significance of the hardware aware implementation. Without this data added, it is not clear that the proposed implementation is actually a significant contribution. A comparison to existing spiking neural network implementations on Loihi 2 would further add value to the paper.
3. The hardware-aware implementation on Loihi 2 is evaluated against a JIT compiled full precision pytorch model deployed on a Jetson Orin Nano. There are a couple of issues with this comparison.
   • Integer precision for matrix operations poses a significant reduction in energy spend on computational operations and memory movement. Hence, comparing a quantized model against a full precision model is not fair.
   • The implementation on Loihi 2 is tuned towards the system architecture of the Loihi system, while the authors mention that the Jetson implementation was just-in-time compiled from a torch model. It is not clear from the presentation of the paper, which optimizations were leveraged by the JIT compiler. The authors are transparent about this issue and even point out that they were not able to use NVIDIA's TensorRT framework for efficient deployment. It is hence likely that a highly optimized Loihi 2 implementation is compared to a poorly optimized CUDA implementation on the Jetson Orin Nano. To strengthen the paper, it is recommended to provide an implementation of similar sophistication or at least provide insights into which optimizations were part of the compilation procedure.
   • The Jetson Orin Nano is oversized for the networks implemented. An Orin Nano has 512 to 1024 CUDA cores plus 3rd generation tensor cores as well as 4 to 8 GB DRAM. Running networks with up to 275.000 parameters favors the smaller Loihi 2 chip with only 128 cores and no DRAM. It is not clear from the paper if the Jetson was fully utilized. For example the power of the large DRAM might significantly contribute to the energy consumption of the Jetson despite not requiring DRAM at all for such small models. The results in table 2 suggest that only the large batch size of 64 could fully utilize the Jetson system - and in this setting the Jetson outperforms the Loihi 2 implementation. To allow for a fair comparison, it would be valuable to add results of an optimized implementation on a small deep learning accelerator of similar size as the Loihi 2 system. For example the Hailo-8 M.2 chip might be a better system for comparison than the oversized Jetson system. Another alternative to strengthening the results would be to run larger networks that saturate the compute and memory capacity of the Jetson. Perhaps the comparison with a low-power CPU with sufficiently large cache to host the 275.000 kb for the parameters might be a fairer comparison than the Jetson.
   • To the best of the knowledge of the reviewer, there is no other implementation of SSMs on inference hardware that reports energy or latency numbers. In this environment, it would be valuable to add related works that optimize implementations of related recurrent neural networks for example on FPGAs. This would contribute to setting the present paper in the context of the RNN inference landscape.

1. Q-S5: Towards Quantized State Space Models

• The primary modification to the model structure involves using ReLU instead of GLUs/GeLUs as the activation function to increase activation sparsity. It would be better if the authors could visualize the activation maps to substantiate this claim.

• The paper notes that different bit-widths were chosen for activations, synaptic weights, and the A, B, C matrices of the S4D layers. Including an ablation study on the bit-width selection would strengthen the paper by illustrating the trade-off between accuracy and efficiency.

• The generalizability of this approach is limited. Showing results for the spiking version of other SSM models, such as Mamba or Mamba2, would enhance the paper’s contribution.

1. I do not see a lot of technical novelty in this work. The SSM models and the quantization techniques are all well-explored in the community. Though the authors evaluate the SSMs on Loihi, I do not see any significant challenges and novel solutions to address the Loihi-specific constraints.

2. The evaluation benchmarks adopted are too simple and outdated. The authors should report results on long range modeling benchmarks, such as Long range Arena and other modern NLP benchmarks.

3. SSMs can be mapped to programmable spiking neurons as shown in [1]. The comparison with this work should be included.

[1] https://arxiv.org/abs/2406.02923v1