An Analysis of Linear Complexity Attention Substitutes with BEST-RQ
Ryan Whetten Titouan Parcollet Adel Moumen Marco Dinarelli Yannick Estève



Rate-Adaptive Generative Semantic Communication Using Conditional Diffusion Models
    Pujing Yang, Student Member, IEEE, Guangyi Zhang, Student Member, IEEE, and Yunlong Cai, Senior 
Member, IEEE
		P. Yang, G. Zhang, and Y. Cai are with the College of Information Science and Electronic Engineering, Zhejiang University, Hangzhou 310027, China (e-mail: yangpujing@zju.edu.cn; zhangguangyi@zju.edu.cn;  ylcai@zju.edu.cn). 
	 
    September 9, 2024
======================================================================================================================================================================================================================================================================================================================================================


§ ABSTRACT 
Self-Supervised Learning (SSL) has proven to be effective in various domains, including speech processing. However, SSL is computationally and memory expensive. This is in part due the quadratic complexity of multi-head self-attention (MHSA). 
Alternatives for MHSA have been proposed and used in the speech domain, but have yet to be investigated properly in an SSL setting.
In this work, we study the effects of replacing MHSA with recent state-of-the-art alternatives that have linear complexity, namely, HyperMixing, Fastformer, SummaryMixing, and Mamba. We evaluate these methods by looking at the speed, the amount of VRAM consumed, and the performance on the SSL MP3S benchmark. Results show that these linear alternatives maintain competitive performance compared to MHSA while, on average, decreasing VRAM consumption by around 20% to 60% and increasing speed from 7% to 65% for input sequences ranging from 20 to 80 seconds.



 



self-supervised learning, speech, efficiency, linear complexity




§ INTRODUCTION



Self-supervised learning (SSL) is an approach to training machine learning models where pseudo-targets are extracted from the data itself.  
Since SSL is unsupervised, these models can be pre-trained on immense amounts of unlabeled data, and then obtain good results on downstream tasks using minimal amounts of labeled data. 

SSL methods have proven to be useful in a variety of domains including in speech processing <cit.>, where SSL has reached state-of-the-art performance in tasks like Automatic Speech Recognition (ASR) <cit.>, Emotion Recognition (ER) <cit.>, Automatic Speaker Verification (ASV) <cit.>, Spoken Language Understanding (SLU) <cit.>, and Automatic Speech Translation (AST) <cit.>.

Despite their performance, training SSL models is still very costly in terms of the amount of data, GPUs, and time needed. For example, Google USM was trained on 12 million hours (or over 1,369 years) of audio <cit.>, and the base and large XLS-R models were trained with 128 and 200 GPUs, respectively <cit.>. Even in efforts to make state-of-the-art models like HuBERT and data2vec more efficient, these SSL models still require around 1,000 A100 GPU hours of training <cit.>. 


In a study of the architectures of SSL models for speech <cit.>, the authors identified three main culprits: (i) the Acoustic Feature Extractor (AFE), which transforms the raw waveform into a latent representation; (ii) the context encoder, which is often a large Transformer <cit.> or Conformer <cit.>; and (iii) the SSL training objective. 















When looking at these three culprits, BEST-RQ <cit.> seems to be theoretically one of the most efficient SSL models for speech that has been proposed. BEST-RQ starts with Mel Filterbanks, addressing culprit (i). Then, it creates pseudo labels using a frozen, randomly initialized linear projection and codebook coupled with cross-entropy training, addressing culptrit (iii).

In contrast, other models, such as wav2vec 2.0 <cit.>, use a learnable codebook and typically use a combination of objectives, slowing down training. For instance, previous studies showed that BEST-RQ obtain comparable downstream performance to wav2vec 2.0 while being 2.4 times faster to train <cit.>. However, for culprit (ii), the context encoder, BEST-RQ uses Conformer layers, which are computationally expensive.







Conformers, like transformers, are expensive partly due to multi-head self-attention (MHSA) time complexity being quadratic with respect to the input sequence length. Therefore, to address culprit (ii), one needs to search for an efficient alternative to MHSA. Many studies have been conducted to reduce the complexity of MHSA, but only few have been applied to speech tasks and none have have been applied to SSL for speech. Examples of such methods include HyperMixing <cit.>, Fastformer <cit.>, SummaryMixing <cit.>, and Mamba <cit.>. 


The main contribution of this work is to address culprit (ii) by, for the first time, evaluating the most promising linear time complexity alternatives to MHSA in an SSL for speech setting. 
Downstream experiments conducted following the MP3S benchmark <cit.> show that our linear-time complexity BEST-RQ maintains performance with an equivalent MHSA BEST-RQ while decreasing VRAM consumption by around 20% to 60% and increasing inference speed from 7% to 65% for input sequences from 20 to 80 seconds. 
As a second contribution, we open-source the code in the widely used SpeechBrain toolkit <cit.>, enabling the community to experiment with efficient SSL models for speech[<https://github.com/whettenr/brq-att-alt-exp>].





§ BACKGROUND
 

Multi-head self-attention (MHSA) has quadratic time complexity with respect to the input length due attention weights being calculated by a dot product between every query-key token pair. Despite its complexity, this dot product operation enables each token to access the global context which is important for reaching high performance <cit.>.





 
In research on reducing the complexity of MHSA, there are two overarching methods: (i) those aiming to approximate this pair-wise token computation with a lower cost, and (ii) those that do not seek to mimic this pair-wise token computation, but instead introduce global context in another fashion.

Some methods in the first family include sparse attention, such as BigBird and Longformer <cit.>, which use a combination of sliding window, global, and random attention to achieve linear complexity. The main issue with these sparse attention methods is that they cannot fully model global context as they operate on windows. Other methods in this family, such as Linformer and Linear Transformer <cit.> approximate attention by computing low-rank approximates of the key and value matrices or use the dot-product of kernel feature maps and make use of the associative property to achieve linear complexity. Yet, in practice, the aforementioned methods are still computationally expensive <cit.>. 


 

In contrast, Fastformer <cit.> has proven to perform well on text data, outperforming previously mentioned sparse attention methods and low-rank approximates. This is done by making use of additive attention and element-wise multiplication to summarize the query and key matrices as vectors resulting in linear complexity while fully modeling global context. Fastformer has also been applied to speech and proved to work well on ASR tasks with the branchformer architecture <cit.>. Due to these performances and the availability of its implementation, Fastformer is considered as a good representative candidate of the first family of methods.

The other family of solutions, which do not aim to approximate the self-attention mechanism, include, HyperMixing <cit.>, SummaryMixing <cit.>, and Mamba <cit.>. HyperMixing was first introduced with the HyperMixer <cit.>. The HyperMixer is an extension of the MLPMixer <cit.>. The disadvantage of MLPMixer is that it can not handle inputs of varying length, making it not suitable for NLP or speech tasks. HyperMixer extends MLPMixer by using hypernetworks <cit.> to capture global context for inputs that vary in length. HyperMixer proved to have good performance on text tasks and fully-supervised ASR with the HyperConformer <cit.>. 

Alternatively, SummaryMixing does not use hypernetworks, but instead the input is passed through a parametrized function, such as a multi layer perceptron, which is averaged across all time steps, resulting in a summary vector. To introduce global context to the input sequence, this summary vector is concatenated to the input at each time step and then fed through another parametrized function, which becomes the final output. SummaryMixing proved to perform just as well as MHSA on ASR, keyword spotting, and SLU, while decreasing the amount of required memory and training time.  

Lastly, Mamba is a method for sequence modeling based on state space models <cit.>. State space models can be thought of as a combination of recurrent neural networks and convolutional neural networks, which scale linearly with respect to the sequence length. Mamba has shown good performance on text, audio (which were in an auto regressive SSL setting), and a two supervised speech tasks <cit.>: ASR and speech enhancement. 

HyperMixing, SummaryMixing, and Mamba all have available implementations, and are therefore selected as representatives of the second family of methods.

Despite the good performance of these substitutes for MHSA on supervised tasks, their performance in an SSL setting for speech tasks has not yet been studied, and the focus of these studies has been mostly on ASR. Furthermore, these methods have not been scaled to large models, e.g. from only about 5 to 100M parameters for speech tasks <cit.>, which contrasts with the current scale of SSL models. Thus, in this work we explore these unexamined aspects of linear complexity alternatives by (i) experimenting in an SSL setting, (ii) looking at performance on a variety of speech tasks using a benchmark from the community, and (iii) scaling up model size to above 300M parameters.







§ ATTENTION ALTERNATIVES IN BEST-RQ



In this section, we give an overview of BEST-RQ and describe how Fastfastformer, SummaryMixing, HyperMixing, and Mamba achieve linear complexity. Let us define the input to the self-attention module as 𝐗∈ℝ^T × d = [𝐱_1, 𝐱_2, ...,𝐱_T] or a sequence of vectors of dimension d with T time steps. The difference between methods lies in how this input is transformed to contain information about the global context. 


BEST-RQ is, to this day, the most efficient SSL paradigm. It is notably used by large scale models such as Google USM <cit.> and Universal-1 from AssemblyAI. BEST-RQ begins with Mel Filterbanks which are passed through two paths, (i) the model, and (ii) the random-projection quantizer. For path (i), a random portion of the Mel Filterbanks are masked and then passed through two convolutional layers, a series of conformer layers, and then a linear layer. For path (ii) the unmasked Mel Filterbanks are passed through a randomly initialized frozen linear layer. Then a codebook look up is performed following:


    y = min_i norm_l2(c_i) - norm_l2(Am) ,


where m is a stacked portion of four Mel Filterbank frames, A is the linear projection, and c is the codebook. The index, y, from this codebook look up is used as the pseudo-target for pre-training for that portion of the input. Then, the cross-entropy loss is calculated between the masked sections of the output of the linear layer from the path (i) and the corresponding pseudo-targets from path (ii). BEST-RQ originally uses self-attention and, therefore, exhibits quadratic time complexity.





Fastformer. The complexity is reduced to linear by using additive attention to summarize the attention matrices as a single vector. To show how this is done, let 𝐐,  𝐊∈ℝ^T × d, be the standard linear transformations from the transformer composed of 𝐐 = [𝐪_1, 𝐪_2, ...,𝐪_T] and 𝐊 = [𝐤_1, 𝐤_2, ...,𝐤_T]. 𝐐 is summarized as a query vector 𝐪 by the following:


    α_t = softmax(𝐰_q^T 𝐪_t/√(d));
       𝐪 = ∑^T_t=1α_t 𝐪_t.


where 𝐰_q ∈ℝ^d is a learnable vector used with each vector of 𝐐 to generate the attentions score α_t. The summary query vector 𝐪 is calculated as a weighted some of the attention scores with their respective 𝐪_t. Then, 𝐪 is multiplied by each vector in 𝐊, resulting in a global context-aware key matrix that relies on the more efficient element-wise multiplication instead of dot products. 
 







HyperMixing reduces the complexity by using a token mixing multi-layered perceptron (TM-MLP) from the MLP-Mixer <cit.>, which can be described as:


    𝐓𝐌-𝐌𝐋𝐏(𝐗) = 𝐋𝐚𝐲𝐞𝐫𝐍𝐨𝐫𝐦(𝐖_1(σ(𝐖^T_2 𝐗^T))),



where W_1, W_2 are weight matrices and σ represents some non-linear activation function. This is a standard MLP except the linear operations are performed on the transposed 𝐗. This can be taught of as operating on the tokens instead of the channels which introduces global context as it allows information to pass between tokens. The key difference between MLP-Mixer and HyperMixing is that 𝐖_1and 𝐖_2 are generated by another MLP and thus can vary in length. 



SummaryMixing. The input, 𝐗, is passed into two functions f and s taking the form, in practice, of two MLPs. The output of s is averaged across all time steps, resulting in a summary vector s̅. To introduce global context to the input sequence, s̅ is concatenated to the output of f at each time step, t, and then fed through another function, or MLP in this case, called c. The output of c becomes the final output 𝐡. The overall SummaryMixing mechanism can be expressed as:


    s̅ = 1/T∑_t=1^T s(𝐱_t);
      𝐡 = c(f(𝐗), s̅).


Calculating the summary vector s̅ is an average and thus is linear with respect to the input length. As a result, SummaryMixing is able to introduce global context with linear complexity.



Mamba. The memory complexity is reduced to linear by processing the input sequence in an unidirectional manner similar to a recurrent network. Being a selective state space model, Mamba represents the past information of 𝐗 as an hidden state h_t with constant memory consumption, hence solving the problematic quadratic time complexity from MHSA. To do so, the model employs a discretized state transition matrix A̅∈ℝ^N × N, an input discretized projection matrix B̅∈ℝ^H × 1, and an output projection matrix C ∈ℝ^1 × H to compute the hidden state h_t as follow:










    h_t = A̅h_t-1 + B̅X_t, 
       
      y_t = Ch_t.

 
 To reduce the   computation overhead brought by the recurrence, an hardware-aware parallel scan algorithm introduced in <cit.> allows to unroll equation <ref> as the input sequence 𝐗 convolved with a structured kernel composed of A̅, B̅, and C fixed. With this improvement, the throughput of Mamba is five times higher than a Transformer. Finally, to introduce global context to the model, similar to MHSA, we use a bidirectional Mamba as motivated by <cit.>.




















§ EXPERIMENTS


In this section, we describe our SSL pre-training protocol and summarize the downstream tasks, giving a brief overview of the datasets and evaluation metrics for each task.



 §.§ Self-supervised Pre-training


All models are open-sourced in the SpeechBrain <cit.> library along with the hyperparameters for each model. 


Architecture details. We use the implementation of BEST-RQ from <cit.> which uses Conformer layers <cit.> with relative sinusoidal positional embedding and the PyTorch multi-head self-attention. For all models the pre-training framework stays the same. We only replace the MHSA cell with either Hypermixing, Fastformer, SummaryMixing, or Mamba. The small models have 12 conformer layers and the large models have 24 layers. 
 We adjust the hidden dimensions to make all the models have around the same number of parameters, that is, 95M and 315M for the small and large ones respectively. The detailed list of hyperparameters can be found in the open-sourced SpeechBrain recipes. 
 


Pre-training details. We pre-train all our models for the same amount of steps, set to 200k using the 960 hours of training data from the LibriSpeech <cit.> dataset. Although 200k steps is a lower number than other state-of-the-art models trained for 400k or 800k steps <cit.>, previous research has empirically shown that,  200k steps is sufficient to compare SSL model performance <cit.>. Dynamic batching is used, meaning audio files are grouped or bucketed together with those of a similar length, and the number of samples per batch varies based on the bucket size in order to keep the number of seconds of input per batch similar. In our experiments, and following the literature <cit.>, we set the batch size to 1.7 hours for all models.





    




























        








        






































    

















 §.§ Downstream Tasks and Metrics

For the downstream evaluation we use a portion of the tasks from the Multi-Probe Speech Self-Supervision benchmark or MP3S <cit.>. 
In MP3S, pre-trained models are frozen and a learned weighted sum of the outputs of the hidden layers of the pre-trained models are fed into a downstream model called a probe. Because results were shown to vary depending on the probe’s architecture, for each task two probes are provided. Furthermore, and as SSL models are commonly fine-tuned with the downstream use case, we added an evaluation with a full fine-tuning for ASR.
 


Automatic Speaker Verification (ASV) is a binary classification task where given 2 utterances, the goal is to determine whether the speakers are the same. The evaluation metric for ASV is the Equal Error Rate (EER). For the dataset, we use VoxCeleb1 <cit.>, which is made of utterances from celebrities sourced from YouTube. The dataset is divided into train and test splits, which we used accordingly. The probes for this task are the X-Vectors <cit.> and ECAPA-TDNN <cit.>.
 

Intent Classification (IC) is a classification task of predicting the main purpose or objective of an utterance. We use accuracy as the evaluation metric. For this task, we use the SLURP dataset <cit.> containing 18 intents coming from single-turn user interaction with a voice assistant. Some examples of intents are email, calendar, or play (as in play a song). The probes for this task are a linear probe, in which the output of the SSL model are average-pooled along the time dimension and then passed into a linear classifier, and an LSTM probe, which consists of a two-layered BiLSTM followed by a linear classification layer.
  

Emotion Recognition (ER) is the task of predicting the emotion of a speaker. Similar to IC, we use accuracy to measure performance. We use the IEMOCAP dataset, which consists of 10 actors performing scripts with four different emotions (neutral, happy, sad and angry). The probes for this task are a linear probe, like IC, and an ECAPA-DNN probe, like ASV.
 

Automatic Speech Recognition (ASR) is the task of transcribing what was said in an utterance. For ASR, we measure performance by the word error rate (WER). We use the LibriSpeech train-clean-100 for training, dev-clean for validation, and test-other for final testing. We report on the final test WER without a language model and with the official 4-gram language model[<openslr.org/11/>] using beam search and shallow fusion. Also, and to test performance in low-resource and out-of-domain settings, which is one use-case of SSL models, we use the two low-resource language ASR datasets proposed in MP3S, Welsh and Basque from the CommonVoice 11.0 dataset <cit.>. For LibriSpeech, we use the 2-layered BiLSTM (LSTM) and ContextNet <cit.> probes, and for CommonVoice we use the linear and LSTM probes offered by MP3S. As mentioned, it is common for pre-trained models to be fine-tuned on a given task instead of being frozen. As a representative of this, we fine tune the models using the 100 hours of labeled data in the train-clean-100 split of LibriSpeech and evaluate on the test-clean and test-other splits. The downstream architecture for this task is a feed-forward neural network with CTC loss.




§ RESULTS



We first evaluate the speed and memory gains on a controlled toy task (Figure <ref>). The findings are then extended to real SSL pre-training and the MP3S benchmark. 


Speed and Memory Evaluation. We give the models randomly generated data of lengths that range from 10 to 80 seconds with 10 seconds intervals. We set the batch size to 6, which was chosen to prevent running out of memory too quickly with MHSA. We perform 10 runs at each input length and time the forward pass as well as measure the VRAM consumption. Only the forward pass is evaluated to concur with a deployment scenario. Averages of computation time and peak VRAM over the 10 runs alongside the 95% bootstrapped confidence interval are plotted in Figure <ref>. Measurements were taken on an isolated node with a Nvidia A100 80GB GPU.  

For input sequences of 10 seconds, the amount of time and VRAM needed is relatively similar for all models. However, a difference starts to appear with 20 seconds, where the alternatives, on average, use 24% less memory, and run 7% faster relative to MHSA. On the other extreme, at an input of 80 seconds, the alternatives use on average 64% less memory and run 65% faster. SummaryMixing and Fastformer were the fastest compared to MHSA with 15% and 11% increase in speed with an input of 20 seconds and 70% and 71% increase in speed at 80 seconds, respectively. In terms of memory, Mamba and Fastformer proved to be the most memory efficient with 28% and 29% decrease in peak memory with an input of 20 seconds, and 67% and 66% decrease in memory at 80 seconds, respectively. These findings, however, need to be validated with real-scale experiments to make sure that these alternatives can reach decent downstream performance.

For pre-training time estimates, we run 5 epochs and measure the time and max VRAM on Nvidia V100 32GB GPUs keeping the batch size at 1.7 hours as in the full pre-training. We scale these numbers up to 200k steps to get an estimate of the full number GPU hours. All of the alternatives prove to be faster and use less memory except the HyperMixing-LG and Mamba-LG models. We report these in Table <ref>. 

























Speech Recognition on LibriSpeech. With the LSTM probe (Table <ref>), MHSA performed better than Mamba, HyperMixing, and SummaryMixing. For the base models with the Contextnet probe, the base version of Mamba outperformed MHSA. For larger models with the Contextnet probe, SummaryMixing performed better than MHSA except with the LM on the test-other split, in which SummaryMixing was 0.01 behind MHSA. Despite a dedicated parameter tuning, the large Fastformer falls significantly behind the others.

We then fine-tuned the SSL models on the train-100 subset of LibriSpeech (Table <ref>). Confidence intervals come from 1000 bootstraps on the test-clean set with a LM. Interestingly, the confidence intervals of HyperMixing and Mamba base models and the SummaryMixing Large model overlap with MHSA best performance. It is therefore unclear if MHSA would always outperform these alternatives. 










Speech Recognition on CommonVoice. This dataset tells a different story and we hypothesis that this is due to a major issue with most speech SSL model benchmarks. Indeed, Librispeech is always used both during the pre-training and downstream evaluation, introducing a clear bias. With CommonVoice (Table <ref>), and with both probes, MHSA is not the best model anymore. With the linear probe, HyperConformer and SummaryMixing both performed better than MHSA on Welsh and Bask. With the LSTM probe, Mamba performed the best followed by HyperMixing, then SummaryMixing or MHSA depending on the language. This demonstrates the ability of these alternative methods to generalize well to out-of-domain and low-resource datasets compared to MHSA. 




Speaker Verification. As for CommonVoice, results varied depending on the probe used, however, MHSA clearly is not the best performing solution (Table <ref>). For the ECAPA probe HyperMixing proved to be the best with an EER of 3.65 and 3.54 for the base and large models respectively. With the Xvectors probe, HyperMixing performed best out of the small models with an EER of 8.54 and for the large model SummaryMixing performed best with and EER of 8.30. 


Intent Classification. The accuracies observed on this task validate our previous findings (Table <ref>). All alternatives, except Fastformer, and with both probes performed significantly better than MHSA. 








Emotion Recognition. This task tells a similar story to intent classification, but with a slightly different leader board (Table <ref>). Again, MHSA is not the best performing solution. Indeed, Mamba performed the best with the linear probe reaching an accuracy of 64.10% and 66.15% for the base and large model respectively, compared to 60.1% and 62.1% for MHSA. For the ECAPA-TDNN probe, SummaryMixing performed the best with and accuracy of 64.34% and 65.21% for the base and large model respectively. 



























































































§ DISCUSSION AND CONCLUSION


One of the goals of SSL pre-training is to develop a model that can represent data without the need of labeled data in a way that is useful for a variety of downstream tasks. 

As this process is expensive, making SSL speech models less resource intensive is an active area of research. Part of this expense is due to multi-head self-attention. While alternatives exist, when it comes to speech tasks, they have only been applied in fully supervised tasks. In this work, we explore replacing multi-head self-attention with four state-of-the-art alternatives: HyperMixing, Fastformer, SummaryMixing and Mamba in a SSL setting for speech tasks. We show that for sequences 20 seconds or more these alternatives are substantially faster and consume less memory. 

Based on the toy test, and considering that about 75% of LibriSpeech is below 15 seconds, we believe that would see a greater difference in the pre-training time and memory when pre-training with a set such as <cit.> where all pre-training audio was cropped to be between 32 and 64 seconds.

However, for small input under approximately 20 seconds the amount of time and memory a model consumes are all relatively similar. This threshold, however, could reduce drastically if the model does not combine Mel Filterbanks with a two-dimensional CNN, as the acoustic feature extractor plays a critical role in the length of the sequence arriving to the self-attention module. For instance, one may expect that replacing this CNN with a one-dimensional one would bring this threshold to 10 seconds. Nevertheless, many files in common speech datasets are shorter than 20 seconds and many speech tasks do not require global context from anything over 20 seconds. Splitting audio files based on silence for long audio file processing or streaming is also possible which brings to question the necessity of these alternatives to MHSA. 

We believe that future work could involve processing large audio files and developing Large Audio Foundation Models, similar to research trends to include more context in Large Language Models (LLM). With this trend in mind, one could imagine doing common LLM tasks, such as summarization, that require long context directly from audio without the need of a transcription.

Nevertheless, as a result of our findings, we believe that unless specifically working with long audio files or from the raw waveform, further speed and memory gains will not be obtained by replacing MHSA with alternatives with linear complexity. We believe that more efficient SSL for speech models might be reached by other architectural changes, pruning/quantization methods, and data selection.



IEEEbib


