Paper: LMUFormer: Low Complexity Yet Powerful Spiking Model With Legendre Memory Units (ICLR-2024)

Reviewer comment (Evaluation):
The paper could be improved by providing a more balanced view, including a discussion of potential limitations or scenarios where the LMUFormer might not perform as well. This would help readers develop a more nuanced understanding of the model’s applicability.

Author response:
Thank you for bringing up this important point. There may be a few potential limitations or scenarios where LMUFormer may not perform very well. For example, as mentioned by Reviewer **KEVY**, LMUFormer may lack the pre-training capabilities of a Transformer, that can enable high performance when fine-tuned on a range of downstream tasks. Rather, it needs to be trained on the streaming use-case directly by the user. Developing architectural optimizations for LMUFormer to scale up and benefit from pre-training is an important research direction.

Moreover, as shown in Table 4, LMUFormer, though yields higher performance compared to all transformer variants, can not surpass the S4 models on LRA dataset, that evaluates model quality under long-context scenarios. Improving the LMUFormer performance for such datasets is also an important and interesting research direction. 

These limitations have been added in Appendix A.2 of the revision.

Reviewer comment (Reproducibility):
Providing more implementation details, including hyperparameters and training procedures, would enhance the reproducibility of the results, contributing to the paper’s overall quality.

Author response:
Thank you for the important reminder. We have added the training configuration for the three tasks considered in this work in appendix A.1.

Reviewer comment (Evaluation):
Given the focus on streaming applications and real-time data processing, a more detailed analysis of the model’s performance in real-time scenarios, including potential latency issues and how they are addressed, would be valuable.

Author response:
So far, we have only explored the real-time processing theoretically; we show that our model can process each sample as soon as it arrives (after accumulating a tiny fraction of the total number of samples), which is not feasible with the traditional Transformer models. A deeper analysis may necessitate considering a variety of practical settings that may impact the real-time processing capability. For example, the hardware processing the model need to be energy efficient yet provide high utilization to ensure each input sample is accepted promptly by the model once it arrives. Moreover, the latency of the model would depend on the communication overhead of the sensor generating the input sample. To minimize the latency and enable real-time processing, this overhead should not exceed the time required by the model to process of one sample. The accurate real-time processing analysis will thus depend on the both the underlying sensor and deep learning hardware as well as the dataflow.

Reviewer comment (Experiments):
It would be nice to see a hardware simulation for SNN to learn all aspects of the new model, including energy, latency, throughput, and any overhead, etc.

Author response:
Thanks for your suggestion. We have now developed a hardware simulation framework for SNNs to analyze the energy, latency, and throughput of our non-spiking and spiking LMUFormer models. We have also incorporated the overhead due to the spike sparsity in our framework, which is minimal with high sparsity as obtained in this work. Our SNNs yield $27.2\times$ lower compute energy compared to LMUFormer on the Google speech commands with no additional overhead on latency and throughput, since our SNNs do not incur any additional overhead, unlike traditional SNNs.

Reviewer comment (Novelty):
Citation for S4 models?

Author response:
We really appreciate you catching the two errors in the paper; we have now added the citation for S4 models and the details of how we arrived at 32.03% reduction in sequence length. You are right that it is indeed (128-87)/128. 

We also apologize for a mistake we had made in the paper regarding the model size and FLOPs reduction factors. We actually showed the correct absolute model size and FLOPs numbers for the two different models, but made a mistake while computing the reduction factors. The correct reduction factor of our LMUFormer compared to the SOTA AST model is $86.933/1.622\approx53.66$ for the model size and $124/1.89\approx65.61$ for FLOPs. We have corrected these numbers in the revision.

Reviewer comment (Presentation):
Could you add "how to get 32.03%" in section 5.5? I assume it is (128-87)/128? Similarly to 70x and 140x to help reading.

Author response:
We really appreciate you catching the two errors in the paper; we have now added the citation for S4 models and the details of how we arrived at 32.03% reduction in sequence length. You are right that it is indeed (128-87)/128. 

We also apologize for a mistake we had made in the paper regarding the model size and FLOPs reduction factors. We actually showed the correct absolute model size and FLOPs numbers for the two different models, but made a mistake while computing the reduction factors. The correct reduction factor of our LMUFormer compared to the SOTA AST model is $86.933/1.622\approx53.66$ for the model size and $124/1.89\approx65.61$ for FLOPs. We have corrected these numbers in the revision.

Reviewer comment (Experiments):
In task 2, there is still some accuracy gap between the proposed method and some transformer based methods. Is there any way to compare pLMU on the same task as well as resource usage? It feels the narrative is that accuracy wise on relative complex tasks LMUformer is on par with transformer based method while out-performing pLMU.

Author response:
**Comparison between LMUFormer and pLMU**

We would like to clarify that we compare the performance of pLMU and our LMUFormer in Table 1 for psMNIST and Table 5 for Speech Commands. Note that in Table 5, the model in the first row actually represent the pLMU model, and we have clarified them in the revision. In the revision, we have added their comparison for the LRA dataset, which is shown in the table below.

| Model              | ListOps (2K) | Text (4K) | Retrieval (4K) | Image (1K) | Pathfinder (1K) | Avg   |
|--------------------|--------------|-----------|----------------|------------|-----------------|-------|
| LMUFormer          | 34.43        | **68.27** | 78.65          | 54.16      | 69.9            | 61.08 |
| Spiking LMUFormer  | **37.30**    | 65.80     | **79.76**      | **55.65**  | **72.68**       | **62.24** |
| pLMU               | 22.83        | 63.72     | 74.1           | 53.50      | 68.32           | 56.49 |
| Spiking pLMU       | 17.94        | 60.84     | 65.19          | 51.89      | 49.88           | 49.15 |
 
 For the LRA dataset, we use the same word and position embedding for all the models, replacing the patch embedding block we use for LMUFormer. This narrows the gap between LMUFormer and pLMU models on LRA compared to Google Speech Commands. Overall, the LMUFormer significantly outperforms pLMU, with comparable results only in Task 'Image' and Task 'Pathfinder'. Note the results on 'Image' task are different from the initial version of the paper, since we changed the initial learning rate and weight decay to improve the performance.

We also show a comparison of performance, model size, peak memory, training time, and inference time for the pLMU model, LMUFormer, Spiking LMUFormer and other Transformer variants, using the "Text" subtask as an example. Our results are shown in the Table below:

 | Model             | Text(4K) | Params. (M) | Peak memory (MB) | Training time (s) | Validation time (s) |
|-------------------|----------|-------------|------------------|--------------------|---------------------|
| Linear Trans.     | 65.90    | 0.174       | 644              | 5.91               | 2.37                |
| Linformer         | 53.94    | 0.698       | 1180             | 8.69               | 2.46                |
| Nystromformer     | 65.52    | 0.174       | 836              | 11.24              | 4.09                |
| LMUFormer         | 68.27    | 0.174       | 738              | 6.29               | 2.46                |
| Spiking LMUFormer | 65.80    | 0.175       | 928              | 12.02              | 5.66                |
| pLMU              | 63.72    | 0.132       | 582              | 4.60               | 1.83                |
| Spiking pLMU      | 60.84    | 0.132       | 684              | 8.68               | 3.75                |

As we can see, our LMUFormer achieves the best performance while maintaining relatively decent training and inference speeds, and requires less memory during training.

Reviewer comment (Experiments):
Besides rMNIST task, it is difficult to tell how much performance gain that the proposed method would offer while (potentially) introducing extra hardware cost.

Author response:
**Comparison between LMUFormer and pLMU**

We would like to clarify that we compare the performance of pLMU and our LMUFormer in Table 1 for psMNIST and Table 5 for Speech Commands. Note that in Table 5, the model in the first row actually represent the pLMU model, and we have clarified them in the revision. In the revision, we have added their comparison for the LRA dataset, which is shown in the table below.

| Model              | ListOps (2K) | Text (4K) | Retrieval (4K) | Image (1K) | Pathfinder (1K) | Avg   |
|--------------------|--------------|-----------|----------------|------------|-----------------|-------|
| LMUFormer          | 34.43        | **68.27** | 78.65          | 54.16      | 69.9            | 61.08 |
| Spiking LMUFormer  | **37.30**    | 65.80     | **79.76**      | **55.65**  | **72.68**       | **62.24** |
| pLMU               | 22.83        | 63.72     | 74.1           | 53.50      | 68.32           | 56.49 |
| Spiking pLMU       | 17.94        | 60.84     | 65.19          | 51.89      | 49.88           | 49.15 |
 
 For the LRA dataset, we use the same word and position embedding for all the models, replacing the patch embedding block we use for LMUFormer. This narrows the gap between LMUFormer and pLMU models on LRA compared to Google Speech Commands. Overall, the LMUFormer significantly outperforms pLMU, with comparable results only in Task 'Image' and Task 'Pathfinder'. Note the results on 'Image' task are different from the initial version of the paper, since we changed the initial learning rate and weight decay to improve the performance.

We also show a comparison of performance, model size, peak memory, training time, and inference time for the pLMU model, LMUFormer, Spiking LMUFormer and other Transformer variants, using the "Text" subtask as an example. Our results are shown in the Table below:

 | Model             | Text(4K) | Params. (M) | Peak memory (MB) | Training time (s) | Validation time (s) |
|-------------------|----------|-------------|------------------|--------------------|---------------------|
| Linear Trans.     | 65.90    | 0.174       | 644              | 5.91               | 2.37                |
| Linformer         | 53.94    | 0.698       | 1180             | 8.69               | 2.46                |
| Nystromformer     | 65.52    | 0.174       | 836              | 11.24              | 4.09                |
| LMUFormer         | 68.27    | 0.174       | 738              | 6.29               | 2.46                |
| Spiking LMUFormer | 65.80    | 0.175       | 928              | 12.02              | 5.66                |
| pLMU              | 63.72    | 0.132       | 582              | 4.60               | 1.83                |
| Spiking pLMU      | 60.84    | 0.132       | 684              | 8.68               | 3.75                |

As we can see, our LMUFormer achieves the best performance while maintaining relatively decent training and inference speeds, and requires less memory during training.

Reviewer comment (Evaluation):
Does the proposed method suit SNN better and more nature? What about the training cost and time?

Author response:
**Suitability of LMUFormer with SNNs**

LMUFormer has recurrent processing capabilities similar to SNNs. However, naively applying the SNN to LMUFormer would incur an additional temporal dimension where the SNN would integrate the membrane potential, and would consume additional energy and latency [1]. In contrast, we merge the membrane potential and spike updates of the SNN for each time step with the analogous memory and hidden state updates of the LMUFormer for each input sample to develop our spiking LMUFormer model. Thus, our spiking LMUFormer requires the same number of time steps as the number of samples in the input sequence. However, spiking LMUFormer can not be easily accelerated on traditional GPU hardware, and thus incurs higher training and inference time compared to LMUFormer there. This phenomenon is also observed in traditional SNNs. Fortunately, spiking LMUFormer can be readily deployed on neuromorphic hardware, and thus can yield high energy efficiency there.

[1] G. Datta, H. Deng, R. Aviles, Z. Liu and P. A. Beerel, "Bridging the Gap Between Spiking Neural Networks & LSTMs for Latency & Energy Efficiency," 2023 IEEE/ACM International Symposium on Low Power Electronics and Design (ISLPED), Vienna, Austria, 2023.

Reviewer comment (Reproducibility):
From eqn (4), it seems the LMU module is composed of mainly matrix-vector operations, which is similar to RNN. Is it correct that we can also parallelize it across the time dimension? Can you elaborate on how it can be solved in a non-iterative way to enable parallelized training?

Author response:
**Is it correct that we can also parallelize it across the time dimension? can you elaborate on how it can be solved in non iterative way to enable parallelized training?**

Yes, you are right. As you can see, the only recurrent connection in the LMU cell is $\boldsymbol{m}[t] = \bar{\boldsymbol{A}} \boldsymbol{m}[t-1] + \bar{\boldsymbol{B}} \boldsymbol{u}[t]$, which calculates the memory of LMU cell at time step t, $\boldsymbol{m}[t]$, based on the memory at time step t-1, $\boldsymbol{m}[t-1]$, and the input of the LMU cell at time step t, $\boldsymbol{u}[t]$. $\bar{\boldsymbol{A}}$ and $\bar{\boldsymbol{B}}$ are the state-space metrics that are fixed and can be calculated before training. 

Since it is an LTI system, we can obtain $\boldsymbol{m}[t]$ in a non-iterative way: $\boldsymbol{m}[t] = \sum_{i=1}^t \bar{\boldsymbol{A}}^{t-i} \bar{\boldsymbol{B}} \boldsymbol{u}[i]$. Moreover, as shown in the following equation, we can get the memory at all time steps simultaneously by the matrix multiplication between a pre-calculated matrix $\boldsymbol{H} $ and the inputs matrix $\boldsymbol{U}$. 

$\begin{bmatrix}
    \boldsymbol{m}[1] \\\\
    \boldsymbol{m}[2] \\\\
    \vdots \\\\
    \boldsymbol{m}[t] \\\\
\end{bmatrix} = [\bar{\boldsymbol{A}}^0\bar{\boldsymbol{B}},  \bar{\boldsymbol{A}}^1\bar{\boldsymbol{B}}, \cdots, \bar{\boldsymbol{A}}^{t-1}\bar{\boldsymbol{B}}]
\begin{bmatrix}
    \boldsymbol{u}[1] & \boldsymbol{u}[2] & \cdots & \boldsymbol{u}[t] \\\\
     & \boldsymbol{u}[1] & \cdots & \boldsymbol{u}[t-1] \\\\
     & & \ddots & \vdots \\\\
     & & & \boldsymbol{u}[1] \\\\
\end{bmatrix}= \boldsymbol{H} \boldsymbol{U}
$.

For RNN models, the parameter matrices $\bar{\boldsymbol{A}}$ and $\bar{\boldsymbol{B}} $ are not fixed as they are trainable, so the $\boldsymbol{H}$ needs to be recalculated at each forward. But they are still trained in parallel.
For more details, you can check [1].


[1] Narsimha Reddy Chilkuri and Chris Eliasmith. Parallelizing legendre memory unit training. In International Conference on Machine Learning, pp. 1898–1907. PMLR, 2021.

Reviewer comment (Evaluation):
I'm curious to know how its measured training and inference time compares to that of other linear-time transformers like Linformer on off-the-shelf hardware.

Author response:
**Additional Comparison of Model Size, Compute Complexity, Training and Inference Time**

We fully agree with you and we are happy to compare the model size and compute complexity comprehensively. However, since some studies do not open-source their model descriptions and there is some ambiguity regarding the exact model configuration, it is difficult to obtain accurate model sizes and FLOPs. Alternatively, we consider the accuracy, model size, peak memory, training time, and inference time, which allows us to compare the performance and efficiency of our LMUFormer with other models. As shown in the table below, our LMUFormer, similar to variants of the efficient Transformer Model, has a fast training speed.

 | Model             | Text(4K) | Params. (M) | Peak memory (MB) | Training time (s) | Validation time (s) |
|-------------------|----------|-------------|------------------|--------------------|---------------------|
| Linear Trans.     | 65.90    | 0.174       | 644              | 5.91               | 2.37                |
| Linformer         | 53.94    | 0.698       | 1180             | 8.69               | 2.46                |
| Nystromformer     | 65.52    | 0.174       | 836              | 11.24              | 4.09                |
| LMUFormer         | 68.27    | 0.174       | 738              | 6.29               | 2.46                |
| Spiking LMUFormer | 65.80    | 0.175       | 928              | 12.02              | 5.66                |
| pLMU              | 63.72    | 0.132       | 582              | 4.60               | 1.83                |
| Spiking pLMU      | 60.84    | 0.132       | 684              | 8.68               | 3.75                |

Reviewer comment (Novelty):
The contribution of this work is rather limited as it relies on a previously proposed module (i.e., LMU). It is not also clear why the proposed network yields a better accuracy.

Author response:
**Why does LMUFormer yield superior performance?**

1. The original LMU model directly inputs the raw input data to the memory cell. However, previous research on vision transformers have shown that a patch embedding block consisting of several convolutional layers can improve the performance of vision tasks. Inspired by this idea, we propose 1D and 2D convolutional patch embedding blocks to extract more abstract features that are inputted to the memory cell of the LMU.
2. The original LMU model is a typical RNN which focuses on the information communication between different tokens. To augment this with the information communication between channels, we explicitly add some convolutional 1D layers, which we call the channel mixer. Therefore, both per-token features and per-channel features are sufficiently integrated in our model.

Both these architectural optimizations explain the superior test accuracy of our LMUFormer, which is validated by the ablation studies presented in Table 5.

Reviewer comment (Experiments):
My main concern is the performance of such a network when pre-trained. The main advantage of transformers come from its pre-training stage which allows the model to perform well on downstream tasks during fine-tuning. To learn and store those pre-training data, transformers contains numerous parameters. As such, I am not supersized by the results reported in this paper since the transformer used for the comparison was not pre-trained and consequently it is expected to see a better performance from LMUFormer. What would be the performance of LMUFormer on well-known benchmarks such as GLUE? Can LMUFormer be pre-trained?

Author response:
**Pre-training of LMUFormer**

We agree with you that pre-training significantly enhances the capabilities of Transformer models. However, the aim of this work is not to replace Transformers or develop models for a range of downstream tasks. Rather, our aim is to consider important streaming use-cases where Transformers exhibit limitations and to develop efficient models for these specific use-cases. Transformers need to wait for the entire sequence to start processing, which incur higher latency compared to our models. It is not very clear whether pre-training on a large text corpus that empower Transformers, can improve the performance of our models on streaming use-cases. Moreover, we are not aware of any large-scale streaming dataset, such as the BookCorpus for non-streaming NLP applications, that we can use to pre-train our models. 

Current pre-trained transformers, that can obtain high performance when fine-tuned on a range of downstream tasks, have high compute complexity, and can not be easily deployed on resource constrained edge devices, as our lightweight LMUFormer. Our novel spiking version of LMUFormer further reduces the compute complexity. However, we agree with you that, given a large-scale pre-trained dataset (and enough compute power), it may be possible for LMUFormer to scale up, and achieve high performance when fine-tuned on downstream tasks. We hypothesize this may not be possible by simply stacking LMUFormer blocks, and may require network architectural optimizations, which is an important research direction.

These points have been added to the revision in Appendix A.2.

Reviewer comment (Theory):
Lack of theoretical analysis and reasoning on the superior performance of LMUFormer.

Author response:
**Why does LMUFormer yield superior performance?**

1. The original LMU model directly inputs the raw input data to the memory cell. However, previous research on vision transformers have shown that a patch embedding block consisting of several convolutional layers can improve the performance of vision tasks. Inspired by this idea, we propose 1D and 2D convolutional patch embedding blocks to extract more abstract features that are inputted to the memory cell of the LMU.
2. The original LMU model is a typical RNN which focuses on the information communication between different tokens. To augment this with the information communication between channels, we explicitly add some convolutional 1D layers, which we call the channel mixer. Therefore, both per-token features and per-channel features are sufficiently integrated in our model.

Both these architectural optimizations explain the superior test accuracy of our LMUFormer, which is validated by the ablation studies presented in Table 5.