LMUFormer: Low Complexity Yet Powerful Spiking Model With Legendre Memory Units

Thanks for your positive comments and valuable suggestions to improve the quality of our work. Please see our response below.

Addressing Potential Biases

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.

Implementation Details

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

Real-Time Processing Analysis

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.

Hardware-Software Co-Design

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.

Citation for S4 models, and "how to get 32.03%" in section 5.5? Is it (128-87)/128? How 70x and 140x

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.

Thanks for your positive comments and valuable suggestions to improve the quality of our work. Please see our response below.

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.

Is there a reason why Am(t) is changed to Am[t-1] from eqn(1) to eqn(2) for discretization?

We are sorry that our omission has caused inconvenience to your understanding. The complete derivation steps using Euler method are as follows:

$

\dot{\boldsymbol{m}}(t) &= \boldsymbol{A} \boldsymbol{m}(t) + \boldsymbol{B} \boldsymbol{u}(t) \\ \frac{\boldsymbol{m}[t] - \boldsymbol{m}[t-1]}{\Delta t} &= \boldsymbol{A} \boldsymbol{m}[t] + \boldsymbol{B} \boldsymbol{u}[t] \\ \boldsymbol{m}[t] &= \boldsymbol{m}[t-1] + {\Delta t}(\boldsymbol{A} \boldsymbol{m}[t] + \boldsymbol{B} \boldsymbol{u}[t])

$

Since we assume $\Delta t$ is small enough, we can use use $\boldsymbol{m}[t-1]$ to approximate $\boldsymbol{m}[t]$. Therefore, the equation will be:

$

\boldsymbol{m}[t] &= (\boldsymbol{A} \cdot \Delta t + I)\boldsymbol{m}[t-1] + (\boldsymbol{B} \cdot \Delta t) \boldsymbol{u}[t] \\ \boldsymbol{m}[t] &= \bar{\boldsymbol{A}} \boldsymbol{m}[t-1] + \bar{\boldsymbol{B}} \boldsymbol{u}[t]

$

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.

Thanks for your positive comments and valuable suggestions to improve the quality of our work. Please see our response below.

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.

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:

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

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.

Thanks for your valuable comments and suggestions to improve the quality of our work. Please see our response below.

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.

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.