Paper: SummaryMixing: A Linear-Complexity Alternative to Self-Attention for Speech Recognition and Understanding (ICLR-2024)

Reviewer comment (Presentation):
Presentations need improvement. See question sections for details.

Author response:
Thanks to the reviewer for helping clarify a few aspects of the manuscript. The detailed answers are given in the question section and the paper has been updated accordingly.

Reviewer comment (Novelty):
Literature reviews on ASR are very limited. Please have an individual paragraph in Section 1.1 to describe those.

Author response:
As suggested, we added a paragraph in Section 1.1 to summarize the ASR domain and how it relates to this work. We invite the reviewer to have another look at this Section.

Reviewer comment (Novelty):
Novelty is a concern. The proposed work is an incremental work on the top of BranchFormer and HyperMixer. Although there is a paragraph describing the relationship to HyperMixer, the difference seems to be more on the choice of linear functions. Please explain more on this.

Author response:
First of all, we modified Section 2.2 to clarify the use of SummaryMixing both in Conformer and Branchformer as well as Section 2 and the Appendix to clarify the link with HyperMixer. Indeed, we also apply SummaryMixing with Conformer and achieve SOTA, but this was not clearly expressed in the first version of the manuscript. Now, if we see things from the Branchformer point of view, we agree with the reviewer, in proportion, the Sumformer implies only a few changes (and this is a strength from our point of view). However, we would like to first highlight that SummaryMixing is also applied to the Conformer architecture i.e. (not only Branchformer) as shown by Table 1 and the CTC-only experiments. Conformer with SummaryMixing beats the MHSA Conformer as well. As a more general answer, SummaryMixing can replace MHSA in any speech encoder.  Indeed, and as mentioned by the reviewer in the strength of this manuscript, SummaryMixing is about “replacing self-attention with linear layers [...]” not only for Branchformer. In this article we demonstrate that this is feasible with the two most used architectures for ASR: Conformer and Branchformer. Finally, we wish to highlight the fact that Branchformer was accepted at ICML while its contribution was to make two blocks of layers parallel instead of sequential – from our very own point of view, this does not represent much more novelty than SummaryMixing (quantitatively speaking), yet it was accepted and is now used by the community thanks to the easy open-access to the code, which is also true for SummaryMixing.

Reviewer comment (Experiments):
The improvements from the proposed method seems to depend on the use of a transformer decoder (with self-attention). Need more experiments to verify the parity of the proposed method to self-attention in other conditions. See question sections for details.

Author response:
We thank the reviewer for raising this excellent remark. We invite the reviewer to have a look at the new results added in the top comment of the reviews as well as the last paragrap of Section 3.3.1 We decided to train ASR systems with CTC only, i.e. without any attentional decoder, and reached basically the same performance between MHSA and SummaryMixing on dev-clean and tes-clean. Therefore, we can say that, in the specific context of this experimental protocol, MHSA is not needed to achieve SOTA ASR performance and that SummaryMixing can do just as well WER-wise while drastically speeding up the training and decreasing the VRAM consumption.

Reviewer comment (Writing):
Abstract: “However, attention layers in trained speech recognizers tend to not capture fine-grained pair-wise information.” Please explain what “fine-grained pair-wise information” is and why it is important (although they are explained in the following sections, abstract should be self-explanatory).

Author response:
We agree with the reviewer that this sentence is actually misleading and not adding sufficient context or value to justify its presence in the Abstract. Therefore, we decided to remove it from the abstract and instead rely on the proper explanation of this concept that appears later in the paper, as stated by the review. We kindly ask the reviewer to tell us if any further clarifications should be added.

Reviewer comment (Presentation):
Figure 1: The diagrams are a bit confusing. I understand that “T” refers to the total number of frames. However, audiences may likely think of them as different attention heads. Please consider a better way for the diagrams.

Author response:
As a part of the rewriting of Section 2, we clarified what T and the plates from the diagram mean in the text. We hope that this whole section is now clearer.

Reviewer comment (Experiments):
SLURP uses read speech - the impact of temporal averaging here may be minimal, but may be very harmful in spontaneous conversational speech. Such explorations in my view are important in this work to claim that SummaryMixing can reach top-of-the-line in speech understanding.

Author response:
We would like to address this point in two parts. First, and following the reviewer’s request, we conducted an experiment on a very noisy dataset with conversational speech: AMI. We wish to mention that there is no AMI ASR recipe available in SpeechBrain yet, and we tried our best based on the literature and the results are not far away from the numbers in other tools, like K2 or ESPnet.  From the results, the reviewer can see that SummaryMixing still works competitively compared to MHSA. It also performs much better than the previous best baseline: the Fastformer linear operation baseline fails to reach the same level of performance.

 **WERs on AMI dataset (conversational and noisy speech)**
| Model (Conformer Encoder-Decoder CTC+Attention) | Dev | Test |
| -------- | -------- | -------- |
| MHSA (41.4M)    | 21.9 | **18.2**|
| SummaryMixing (39.9M) | **21.8**| 19.4|
| Fastformer (39.9M) | 25.1 | 21.3|

Second, ASR is different from SLU, and we therefore can not affirm that such findings could be extended to SLU. As a matter of fact, we decided to follow the same experimental protocol as the [Branchformer paper](https://arxiv.org/pdf/2207.02971.pdf) accepted at ICML, hence the choice of the SLURP and Google Speech Command datasets. AMI was the only available (commercially exploitable and in SpeechBrain) dataset that we could use to try quickly SummaryMixing on very noisy and conversational speech.

Reviewer comment (Theory):
The authors claim that a score of 0.5 diagonality implies a uniform distribution and therefore, a simple average could achieve a similar token mixing. However, this is not obvious or clear to me. Could the authors explain?

Author response:
Given this comment, we revised section 1 to add more details about this issue. We would like to answer this more in detail here as well. Based on the definition of diagonality of Zhang. et. al. and Peng. et. al. , diagonality 0 indicates all the weights are pushed as far away as possible from the diagonal, diagonality 0.5 means all the weights are uniformly distributed along the sequence and diagonality 1 implies all the weights are distributed along the diagonal. For the detailed formula of “diagonality”, we invite the author to look at Zhang. et. al. (arxiv.org/pdf/2011.04906.pdf) or Peng. et. al. (https://arxiv.org/pdf/2207.02971.pdf)

Reviewer comment (Evaluation):
I was interested to know how the performance of summary mixing is impacted during inference by temporal averaging. As the sequence length increases, the sum vector becomes a mixture of more frames, degrading the ability to discriminate individual frames/tokens. Do the authors have any numbers showing ASR performance (WER) as a function of sequence length across models with self-attention and the proposed SummaryMixing?

Author response:
Following the advice of the reviewer, we added new results (in the Appendix and in the top message of these reviews) showing the decoding WER of SummaryMixing, Fastformer and MHSA as a function of the sequence length. It appears that SummaryMixing is not more negatively impacted by the sequence length than MHSA. It is actually the opposite as MHSA has a higher increase in WER following the sentence duration.

Reviewer comment (Experiments):
As SLURP uses read speech, I was wondering if the authors had performed experiments on other corpora with spontaneous speech, for example, SLUE-VoxPopuli [1] for Named Entity Recognition (NER)?

Author response:
We again agree with the reviewer. In a previous answer (see weaknesses section), we mentioned the new results obtained on the noisy and conversational speech AMI dataset. SummaryMixing still works just as well as self-attention, even in this challenging environment. We also had a look at the available recipes in SpeechBrain that we could run on more complex SLU datasets, unfortunately, we could not find any.

Reviewer comment (Novelty):
There are some papers [2,3,4] that show linear transformers can obtain comparable performance to MHSA for ASR contrary to what the authors claim. These works are also relevant and must be acknowledged in the paper.

Author response:
First of all, we agree that other linear alternatives to MHSA exist in the literature. We also understand that this might not be clear enough from the paper and this is why we modified the paper accordingly. We invite the reviewer to have a look at Section 1.1 to make sure that this now appears clearly. [2, 3] indeed passed below our radars, and we must acknowledge it (as we did in Section 1.1). We wish to mention however that no code is given to reproduce the result, while SummaryMixing is already out and usable in SpeechBrain. Also, in [3], the authors wish to keep the overall pair-wise interactions of self-attention, they simply restrict the window with a clever solution. SummaryMixing is fundamentally different – it completely removes self-attention. Finally [4] does not have linear attention at training time, the authors propose to apply a window to self-attention post-training to limit the decoding time increase with respect to the sequence length, they also use aggressive downsampling to reduce the inference time. SummaryMixing is linear at training time and decoding time, without any trick. Finally, we wish to highlight the fact that one of our baseline, Fastformer, was in fact compared to many linear SOTA alternatives to self-attention in speech. Fastformer always beats the other methods. This is why we decided to pick Fastformer as a baseline, because beating it would imply the superiority of SummaryMixing over the other techniques as well, at least in the context of the environment protocols defined in the original Fastformer paper in ours.

Reviewer comment (Writing):
The authors should clarify the ASR architecture in the final version since the current description of the ASR architecture is vague and prone to confusions -- see questions below --. The authors might be aware of this but prefer to refer readers to the SpeechBrain recipes and experimentation. This might leave a shallow reader without immediate interest in reproducing or delving into the code and recipes with a confusing architecture.

Author response:
We totally agree with the reviewer. The reason for this design choice is that, for instance, the yaml (hyperparameters definition) of the model trained on Librispeech that we have released on GitHub (please do not look for it, it’s not anonymized, we can provide it here if necessary) contains 260 lines of model definition. Hence, we initially thought that this would be better for the reader to just have access to our parameter files (**available on Github for all experiments**) as well as the SpeechBrain recipes would be sufficient. However, we added a simplified version of the model definition that hopefully will help the reader to get a better understanding of the model architecture in the Appendix. We invite the reviewer to have a read at the Appendix and let us know if any more information should be added.

Reviewer comment (Experiments):
The approach has some limitations, so it would be intriguing to examine the performance of the proposed encoder in long-context ASR.

Author response:
We agree with the reviewer and invite them to have a look at the top comment illustrating the new results on the evolution of the WER as a function of the sequence length of SummaryMixing vs Fastformer vs MHSA. We also added these results to the Appendix of the manuscript. From the results, it is clear that SummaryMixing does not suffer more than MHSA from longer sentences, it actually looks like the opposite. We also share the reviewer’s view on the application domain. However, our primary domain of expertise is speech processing, and it is unclear if 10 days would be sufficient for us to build a trustworthy experimental protocol on NLP tasks.

Reviewer comment (Experiments):
Since this is a leaning representation conference, some experiments on NLP tasks, as in the original HyperMixer paper, would have attracted more readers.

Author response:
We agree with the reviewer and invite them to have a look at the top comment illustrating the new results on the evolution of the WER as a function of the sequence length of SummaryMixing vs Fastformer vs MHSA. We also added these results to the Appendix of the manuscript. From the results, it is clear that SummaryMixing does not suffer more than MHSA from longer sentences, it actually looks like the opposite. We also share the reviewer’s view on the application domain. However, our primary domain of expertise is speech processing, and it is unclear if 10 days would be sufficient for us to build a trustworthy experimental protocol on NLP tasks.

Reviewer comment (Reproducibility):
Given the nature of speed optimization and comparison with MHSA, it would be interesting for the readers to know the specifics of the MHSA implementation as there are many very efficient implementations available nowadays.

Author response:
We agree that this is not clear from the current form of the manuscript. We invite the reviewer to have a look at section 3.1 to make sure that the revised article solves this issue. In practice, the used MHSA implementation is the one from SpeechBrain (with relative positional encoding for better ASR performance) implemented in PyTorch [here](https://github.com/speechbrain/speechbrain/blob/bd27e99119edfe6b4f530be16d46cce28af260b3/speechbrain/nnet/attention.py#L362). For the speed perspective, we compared it against the regular MHSA from PyTorch (visible [here](https://github.com/speechbrain/speechbrain/blob/bd27e99119edfe6b4f530be16d46cce28af260b3/speechbrain/nnet/attention.py#L642)  in SpeechBrain) and found no significant degradation in speed. We however did not compare this to other forms of very efficient MHSA implementation such as FlashAttention as they benefit from optimisations linked to the compilation and it would be unfair to compare this with SummaryMixing that is a raw PyTorch implementation.

Reviewer comment (Experiments):
It would have been nice to compare the proposed approach in both batch and online ASR.

Author response:
For now, we agree that only offline ASR is addressed. The primary reason for this choice is that, at the time of writing, the only available online ASR pipeline in SpeechBrain is still under Pull Request. We also fully understand that this is not a fully valid answer for not trying it on K2 for instance, but the latter would also require a complete change of environment, with different baselines (e.g. fastformer, hypermixer, the branchformer, and context net does not exist in K2). This is not realistic for a 9 page articles and within 10 days.

Reviewer comment (Experiments):
Could the MHSA improve as the context becomes larger?

Author response:
We also think that this question is interesting. Therefore, we added an experiment to the Appendix as well as the top message of this thread showing the evolution of the WER following the increase in audio duration for MHSA, SummaryMixing and Fastformer. SummaryMixing does not suffer more than MHSA from increased sentence length. Therefore, It is most likely that MHSA will not improve much over SummaryMixing with longer sentences – at least in typical speech-related tasks.

Reviewer comment (Reproducibility):
Which is the MHSA implementation you have used?

Author response:
We detailed this answer in the “Weaknesses” sub-section of this rebuttal. In short: the SpeechBrain i.e. PyTorch implementation of relative positional encoding self-attention.

Reviewer comment (Novelty):
The survey of the efficient transformer (conformed) is not sufficient. There are a lot of efficient transformers in the NLP and ML field (e.g., https://arxiv.org/pdf/2009.06732.pdf). Even if we limit the discussions to speech applications, there are a lot of methods (e.g., squeezeformer, efficient conformer, emformer, etc.).

Author response:
We agree with the reviewer that comparing with more “linear attention” is important and we modified the related work (Section 1.1) of the manuscript accordingly to incorporate many relevant methods. We invite the reviewer to take a look at this revised section. Experiment-wise, however, It is worth mentioning that our SummaryMixing was compared with two “linear-attention” alternatives - Fastformer and Hypermixer. In the original paper, Fastformer has been compared to longformer, Linformer, linear transformers, BigBird and poolingformer. All these methods are state-of-the-art alternatives to self-attention. In this original published work, Fastformer was better than all of them and has been replicated in ESPNet. This is why it appeared as a strong baseline in our case. Indeed, if Fastformer beats most existing linear alternatives, our goal should be to beat Fastformer.  Then, we wish to highlight that there exists no available implementation of any linear-attention claiming SOTA performance on ASR benchmarks -- SummaryMixing is the first one. Fastformer is implemented in ESPnet, but does not reach SOTA – SummaryMixing does on most if not all ASR datasets.

Reviewer comment (Evaluation):
The performance improvement from the conventional method (e.g., fastformer) is marginal.

Author response:
The Fastformer also can not catch the performance of MHSA, hence compute and memory costs were coming at the cost of accuracy. This is not true with SummaryMixing as it reaches or even beats MHSA. We agree with the reviewer that absolute improvements do not seem vastly game-changing. However, we wish to highlight that the BranchFormer, which simply makes the conformer parallel instead of sequential and was published at ICML, improved the performance over the Conformer by .1% on Librispeech dev-clean and .2% on test-clean. Compared to Fastformer; SummaryMixing improves the performance by .2% .1% .3% on Librispeech dev-clean, test-clean and test-others respectively. if we consider other ASR tasks, SummaryMixing always beats Fastformer, with an average absolute gain of 0.75%. The latter improvement is more significant than the one originating from Branchformer vs Conformer from the reference ICML paper. The performance might appear marginal in the reviewer’s perspective, which might be true, yet it is consistent, and more importantly, totally aligned with what we see from the improvements in ASR from the literature.

Reviewer comment (Novelty):
The method is on top of branchformer, which already uses Fastformer, and its technical novelty is weak.

Author response:
We modified Section 2.2 to clarify the use of SummaryMixing both in Conformer and Branchformer. Also, we wish to bring our perspective to the potential lack of novelty. If we see things from the Branchformer point of view, we agree with the reviewer, in proportion, the Sumformer implies only a few changes (and this is a strength from our point of view). However, we would like to first highlight that SummaryMixing is also applied to the Conformer in this paper (Table 1 and CTC only results) and it also beats or matches MHSA. SummaryMixing is expected to replace MHSA in any speech encoder equipped with it, and this article demonstrates that it works just better with the two most used architectures: Conformer and Branchformer. Second, we wish to mention that Branchformer was accepted at ICML while its contribution was to make two blocks of layers parallel instead of sequential – from our very own point of view, this does not represent much more novelty than SummaryMixing (quantitatively speaking), yet it was accepted and is now used by the community.

Reviewer comment (Writing):
Several descriptions (e.g., the surveys, distinction from HyperMixer/MLP Mixer, and detailed architectures) are unclear.

Author response:
We thank the reviewer for mentioning these issues. We revised the paper according to the three problems mentioned by the reviewer, see Section 1.1 for the survey, Section 2 and the Appendix for the Hypermixer, and Section 3.1 and the Appendix for the detailed architecture. Please let us know if the reviewer still finds the definitions confusing or lacking precision, as this only concerns text, we can easily fix it and align it with the reviewer's perspective.

Reviewer comment (Writing):
Frankly, I could not understand the distinction of the SummaryMixing block from HyperMixer/MLP Mixer only from Section 2.1, mainly due to the lack of an explanation of what HyperMixer/MLP Mixer is.

Author response:
We thank the reviewer for mentioning these issues. We revised the paper according to the three problems mentioned by the reviewer, see Section 1.1 for the survey, Section 2 and the Appendix for the Hypermixer, and Section 3.1 and the Appendix for the detailed architecture. Please let us know if the reviewer still finds the definitions confusing or lacking precision, as this only concerns text, we can easily fix it and align it with the reviewer's perspective.

Reviewer comment (Reproducibility):
It is difficult to understand the detailed architectures only from this description: "In particular, the transformation (f), summary (s), and combiner (c) functions are all implemented as a dense linear layer followed by a GeLU activation function." Similarly, how the model size of the proposed and other methods is adjusted was unclear.

Author response:
We thank the reviewer for mentioning these issues. We revised the paper according to the three problems mentioned by the reviewer, see Section 1.1 for the survey, Section 2 and the Appendix for the Hypermixer, and Section 3.1 and the Appendix for the detailed architecture. Please let us know if the reviewer still finds the definitions confusing or lacking precision, as this only concerns text, we can easily fix it and align it with the reviewer's perspective.

Reviewer comment (Novelty):
Can you expand the discussion of how this novel architecture and findings would attract the general AI and ML researchers in ICLR? This paper specializes in speech recognition (and spoken language understanding, which is a very similar task to speech recognition), and it is a narrow scope for me.

Author response:
We wish to answer the question in two parts. First, we wish to highlight that the ICLR call for papers clearly mentions that works focusing on applications to audio/language are welcome (“Representation learning for computer vision, audio, language, and other modalities”). Furthermore, plenty of articles have been published in conferences like NeurIPS / ICLR or ICML while only focusing on speech or specific application domains. For instance “Squeezeformer: An Efficient Transformer for Automatic Speech Recognition” was accepted at NeurIPS 2022 and only used CTC models on a single dataset (Librispeech). This is also true for “Understanding The Role Of Self Attention For Efficient Speech Recognition” accepted at ICLR 2022, only using the Conformer-CTC only model and the Librispeech dataset. Another example is the Branchformer, which only was evaluated on speech recognition and SLU, just like SummaryMixing, was published at ICML 2022. On other domains, MLP Mixer was strictly designed for and evaluated on clean and simple vision tasks yet published at NeurIPS 2021. MLP Mixer was, by design, not usable for variable length sequences including most speech and NLP tasks. Second, we still wish to answer the question properly. We believe that SummaryMixing would attract interest, or at least raise discussions, within the general ML community as SummaryMixing is a replacement to MHSA, evaluated on speech, but that could be extended to other modalities, very much like MLP mixer which was only designed and evaluated for vision. SummaryMixing was conceived with speech in mind because the latter domains fall within our expertise and tools, but it could be applied to vision or NLP by researchers who possess the level of knowledge in such domains.

Reviewer comment (Experiments):
Besides the above expansion, can you explain why you selected Fastformermer and ContextNet?

Author response:
We agree that the reasons for this choice are not clear. Hence, we modified Section 1.1 as well as Section 3.1 to better motivate this decision. **In short, we selected both models as baselines as they, from the literature, represent the SOTA in the tasks that we selected to evaluate SummaryMixing.** More precisely, we selected Fastformer because the original paper compares it to many SOTA linear MHSA alternatives across different tasks. Fastformer always beat them both accuracy and speed-wise. Hence, if SummaryMixing is able to beat Fastformer (and it is), then we can simply deduce that it would beat the others as well, at the very least in the context of the typical ASR / SLU environmental protocols defined in our article and in the original Fastformer/Branchformer articles. ContextNet was added as it was (at the date of writing) the best CNN-only model for ASR. In practice, that is only true on paper as no one ever succeeded in reproducing the claimed results from the original paper, but the final performance that we obtained was already pretty close to SOTA with this architecture. ContextNet was also available on SpeechBrain, making it a good choice from the reproducibility perspective of our potential readers.

Reviewer comment (Experiments):
Can you apply this to RNN-T?

Author response:
From the theoretical perspective, nothing prevents us from using SummaryMixing with RNN-T. The observed results from the article with CTC/Attention and CTC-only training clearly indicate that SummaryMixing works really well with these different decoders and the Transducer would not add any drastic change causing SummaryMixing to fail. The only difference, however, would be that a SummaryMixing RNN-T should be able to handle streaming ASR as well, as it represents the most common use case of RNN-T ASR systems. While turning SummaryMixing into a streaming-capable cell is simple, and we are more than happy to discuss it here. Building an entire experimental protocol for streaming ASR and RNN-T and adding it in a 9-page article appears to be unrealistic given the time frame of the discussion period. However, yes, SummaryMixing can be applied to RNN-T ASR.