Efficient Training of Audio Transformers with Patchout

Khaled Koutini, Jan Schlüter, Hamid Eghbal-zadeh, Gerhard Widmer

Introduction

The transformer architecture has proven very successful in sequence modeling. It allows learning dependencies between different items in the sequence regardless of their positions or their separation in the sequence. Transformers are the state-of-the-art models in different natural language processing tasks . More recently, they have been adapted to computer vision by extracting small patches from the input image and adding a learnable positional encoding to each patch. The resulting patches form a sequence that can be fed to the transformer. These vision transformer models achieve state-of-the-art performance on image classification tasks, but require large amounts of training data (e.g, in Vision Transformer (ViT) ), or heavily depend on extensive data augmentation and knowledge distillation from a CNN model (e.g, in Data-efficient Image Transformers (DeiT) ). Gong et al. further adapted vision transformers to audio spectrograms, achieving state-of-the-art performance on Audioset by using pre-trained models from computer vision and using overlapping patches from audio spectrograms for fine-tuning.

The transformer architecture consists of a series of self-attention layers . Each layer relies on calculating a distance between each pair of items from the input sequence. Although this allows each input item to attend to any other item in the sequence, this results in a complexity of $\mathcal{O}(n^{2})$ with respect to the input sequence length $n$ , in terms of both memory and computation effort. Reducing the quadratic complexity has been the target of several approaches in natural language processing, the idea being to restrict each input item (token) to attend only to a pre-selected subset of input items (tokens). One example is to allow attending only to neighbours inside a sliding window . Additionally, Kitaev et al. use locality-sensitive hashing to approximate attention, reducing the attention complexity to $\mathcal{O}(n\log{}n)$ . BigBird combines sliding windows, global attention, and random interaction between the sequence items. Masking portions of the input sequence has been shown to be an extremely effective method for training encoder/decoder transformers in NLP . In computer vision, the idea of removing patches during inference was investigated to assess the vision transformer’s robustness against input perturbations .

In this paper, we focus on applying transformers to audio processing. We address the shortcomings of current audio transformers from the aspect of computational complexity and memory requirements by introducing a new simple yet effective method for training transformers on spectrograms. In summary, the main contributions of our work are as follows:

We propose Patchout, a method that significantly reduces the computation and memory complexity of training transformers for the audio domain. Patchout also functions as a regularizer, improving the generalization of the trained transformers.

We disentangle the transformer’s positional encoding into time and frequency positional encoding, allowing for straightforward inference on audio snippets of variable length without the need for fine-tuning or interpolating positional encodings.

We investigate additional methods for reducing training complexity and demonstrate how they affect performance on the larger general-purpose Audioset as well as domain-specific downstream tasks.

Our proposed models can achieve state-of-the-art performance on several audio tagging and classification tasks using a single consumer GPU, in a relatively short time (see Figure 1). When different complexity reduction methods are combined, the models outperform CNNs in terms of training speed and memory requirements, in addition to generalization.

The Patchout faSt Spectrogram Transformer (PaSST)

The Vision Transformer (ViT) works by extracting small patches from an input image and projecting these patches linearly onto a sequence of embeddings. The sequence is then augmented by adding trainable positional encodings as biases to the input sequence. A special classification embedding (classification token) is then appended to the sequence, which is connected to a classifier after the self-attention layers. In Data-efficient Image Transformers (Deit) another special embedding for distillation (distillation token) is added. Gong et al. showed that overlapping the extracted patches improves the performance when training ViT on spectrograms. On the other hand, patch overlapping increases the total number of patches, i.e., the input sequence length. Therefore, overlapping greatly increase the memory and compute requirements for training the transformers. We propose a new method called Patchout (Section 2.2) to overcome these issues.

Figure 2 summarizes the proposed transformer architecture: The pipeline starts at the upper left with an audio spectrogram being fed into the model as input. (1) is the patch extraction and linear projection steps as explained in . In (2), frequency and time positional encodings are added, as discussed below. In (3), we apply Patchout as explained in Section 2.2, and add the classification token. In (4), we flatten the sequence and pass through $d$ layers blocks of self-attention ( $d$ is the depth of the transformer). Finally, a classifier operates on the mean of the transformed C and D tokens.

In Addition to Patchout (step 3 in Figure 2), the second main difference between our work and previous work is that we disentangle the positional encoding for the time and frequency dimensions, and as a result, we have two positional encodings: one representing the frequency, and one for time (step 2 in Figure 2). This makes inference and the tuning of the pre-trained models on downstream tasks with shorter audio length simpler. When fine-tuning or inference on shorter audio clips, we simply crop the time positional encoding parameters, without changing the frequency encoding parameter.

Multi-head attention layers rely on computing a distance between each pair of positions in the input sequence (in the form of an attention matrix), therefore having a complexity of $\mathcal{O}(n^{2})$ where $n$ is the input sequence length. As the sequence length grows – for example, when overlapping input patches or for longer audio clips – the compute and memory requirements quickly become problematic. More specifically, given an input of $b$ samples of the dimension $(n\times e)$ , where $b$ is the batch size, $n$ is the input sequence length, $e$ is the embeddings size, each multi-head attention layer projects each input sample to $h$ query $Q$ , key $K$ , and value $V$ matrices, where $h$ is the number of attention heads . Each of $Q$ , $K$ and $V$ has a shape of $(n\times\frac{e}{h})$ . The attention matrix $A$ is then computed by the matrix multiplication $Q\times K^{T}$ , scaling, and applying the soft-max activation function $A=Softmax(\frac{Q\times K^{T}}{\sqrt{d}})$ . $A$ has a shape of $(n\times n)$ and is multiplied with $V$ giving the attention output: $A\times V$ resulting in a new sequence with the same shape as the input $(n\times e)$ . As a result, the computation complexity (and memory requirements) for all the operations on the attention matrix $A$ grow quadratically $\mathcal{O}(n^{2})$ with sequence length $n$ , while the operations in the rest of the network have a linear complexity relationship with $n$ . In short, reducing the sequence length would have a large impact on the computational complexity of these models.

2 Patchout

Motivated by (a) the impact of reducing the sequence length on the computation complexity of training transformer models; (b) the fact that audio events are expected to be spread out in time and frequency in an audio clip; (c) the insight that CNNs can benefit from having a small receptive field during training for different audio tasks, as shown in , we propose Patchout, a method to efficiently train transformer models on audio spectrograms. The idea is to drop parts of the transformer’s input sequence when training, encouraging the transformer to perform the classification using an incomplete sequence. We first extract small overlapping patches from the input spectrograms and linearly project them to vectors, forming the transformer input sequence. We augment the patches with both frequency and time encoding. When training, we randomly drop parts of the sequence, reducing the sequence length, and effectively regularizing the training process. Similar to DropOut , during inference, the whole input sequence is presented to the transformer. We distinguish between different types of Patchout as follows:

Unstructured Patchout is the basic form of Patchout, where we select the patches randomly regardless of their position. We refer to models trained with this method as PaSST-U.

Structured Patchout: We randomly pick some frequency bins/time frames and remove a whole column/row of extracted patches. This structure is inspired by SpecAugment . We refer to models trained with this method as PaSST-S.

3 Further Complexity Reduction Methods

Reducing the overlap between patches results in a lower number of extracted patches, and therefore a smaller transformer input sequence length. However, Gong et al. showed that reducing the overlap (or training without overlapping) degrades the performance of the transformer on Audioset. Patchout can also be used even when there is no overlap between patches. We refer to the system without patch overlapping as PaSST-N.

3.2 Reducing the depth of the transformer

The depth of the transformer is the number of successive self-attention blocks ( d in Figure 2). The depth has a linear relationship with the overall training and inference complexity and influences the total number of parameters of the model. Since we are starting the training from models pre-trained on Imagenet (as explained in Section 3.2), we remove every other self-attention block. This allows us to benefit from the pre-training, compared to removing consecutive blocks, since the residual activations will have a less sudden change. We will refer to the model with the removed blocks as PaSST-L. It has $d=7$ self-attention blocks and 50M parameters compared to 87M in the full model.

Experiment Setup

We train our models on Audioset , the largest publicly available audio dataset, consisting of around 2 million audio clips from Youtube. The task is to tag the audio clips with labels from 527 possible classes. Furthermore, we fine-tune the models trained on Audioset on various audio classification and tagging tasks, namely, instrument recognition, environmental audio classification, and acoustic scene classification.

We use mono audio with a sampling rate of $32$ kHz. We extract Mel features from a window of $25$ ms with a hop length of $10$ ms, resulting in $128$ mel bands, similar to . Kong et al. showed the importance of balancing Audioset; therefore, we balance our training data using importance sampling. We assign a sampling weight to each sample proportional to the inverse frequency of its label $\frac{1}{\mathit{freq(label)}+100}$ . We train on $1,893,693$ (approx. 2M) training segments, and evaluate on $18,951$ audio clips. For each epoch, we sample 200k samples from the full 2M Audioset without replacement. We use the AdamW optimizer with weight decay of $10^{-4}$ , with a maximum learning rate of $10^{-5}$ . We use a linear learning rate decay from epoch 50 to 100, dropping the learning rate to $10^{-7}$ and fine-tune the model for a further 20 epochs.

2 ImageNet Pretraining

Gong et al. showed that using pre-trained models on Imagenet significantly improves their performance on Audioset. Therefore, we will use pre-tranined models in all our experiments.Our base model is DeiT B↑384 . We also achieve a comparable performance using computationally more complex ViT models such as stripped-down ViT-hug224 ; by removing half of the self-attention blocks, its depth was reduced to only 16 blocks (with the methods explained in Section 2.3.2); this will not be further explored in this paper.

3 Data Augmentation

The transformer models are very prone to overfitting, therefore data augmentation plays an essential role in the training process . In our experiments, the following augmentation strategies are used:

Two-level Mix-Up: We use Mix-up since it has been shown to improve performance . We mix both the raw waveforms randomly from the dataset as well as the final spectrograms.

Specaugment: We use SpecAugment by masking up to $48$ frequency bins and $192$ time frames similar to .

Rolling: We roll the waveforms randomly over time.

Random Gain: We multiply the audio waveforms to change the gain by $\pm$ 7 dB.

Results

Table 1 shows the mean average precision mAP (also referred to as precision-recall area under-curve) results on Audioset . As can be seen, the proposed model PaSST achieves a new state-of-the-art performance on the largest available audio tagging dataset. The proposed model outperforms AST and significantly outperforms CNNs. Using Patchout not only improves the performance of the transformer architecture, but also increases the training speed approximately 4 times, and reduces the required GPU memory to less than $25\%$ . As a result, it is possible to train PaSST on a single Nvidia RTX 2080ti (consumer GPU), achieving state-of-the-art performance in 50 hours. Furthermore, PaSST-L-S (with a scaled down depth of $d=7$ ) and PaSST-S-N (without patch overlap) significantly outperform CNNs while maintaining a higher training throughput, and with similar GPU memory requirements. PaSST-S-L and PaSST-S-N can be trained on a single GPU to reach $.459$ and $.466$ mAP in approximately 25 hours. Applying Patchout on the transformer without overlap (PaSST-S-N) outperforms the baseline PaSST-B (without Patchout) and AST while being up to 8 times faster, and requiring less than 10% of the GPU memory for training. The results are also illustrated in Figure 1. The only difference between the baseline PaSST-B and the AST model is the positional encoding. AST , like vision transformers , employs grid positional encoding. PaSST-B, on the other hand, utilises disentangled time and frequency positional encoding (see Section 2).

Table 2 shows the result of ensemble models. We ensemble models with different overlap values between input patches (Figure 2). S indicates the patches stride, S16 means no overlap between the patches. The first ensemble (2 models) averages the logits of a model with no patches overlap and a model with an overlap of 2 (stride=14). S10-S16 indicates that the models used have strides of 10,12,14 and 16.

2 Fine-tuning and Transfer to Downstream Tasks

We fine-tune the pre-trained (on Audioset) models on several downstream audio tagging and classification tasks with different dataset sizes, Table 3 summarizes the results. PaSST-(B,U,S) models use the pre-trained PaSST-S on Audioset, but for fine-tuning, we use no Patchout, unstructured Patchout, and structured Patchout respectively. It is worth noting that the transformer models can be fine-tuned using a small number of epochs. The results suggest that researchers and practitioners can use pre-trained PaSST and fine-tune them on downstream tasks without the need for large computational resources.

In summary, fine-tuning the transformer model outperforms state-of-the-art CNNs on all tasks. Patchout results in significant speedups and, in many cases, improved generalization. When combined with Structured Patchout (-S), reducing complexity by removing patch overlap (-N) performs better than reducing transformer depth (-L) and enables faster fine-tuning.

We only replace the MLP classifier in the pre-trained models for fine-tuning. When we use Patchout, we randomly remove roughly half of the input sequence. Each experiment was repeated three times, and the average results are reported. The speedup in Table 3 is relative to PaSST-B and is rounded up to the nearest integer. Details on the setup of each task can be found in our github repository. Polyphonic Musical Instrument Recognition: The task here is to detect all the instruments present in an audio clip. The OpenMIC dataset consists 20,000 audio clips. Each clip is 10 seconds long and can be assigned multiple tags out of 20 classes. The metric for the task is the mean average precision. The state-of-the-art methods for this task are CNNs with restricted receptive fields . PaSST-S-N reaches the state-of-the-art performance in less than 30 minutes on a single consumer GPU. Environmental Sound Classification: The ESC50 dataset consists of 2,000 environmental 5-second audio clips. The task is to classify each clip into one out of 50 possible classes. We report the accuracy averaged over the 5 official folds . All PaSST variants (with Patchout) can be fine-tuned on this dataset in less than 5 minutes on a single GPU. The state-of-the-art performance was achieved using the AST transformer model . The difference between AST and PaSST-B is in the positional encoding, as explained in Section 2. Acoustic Scene Classification: The task is to recognize the acoustic scene of 10-second audio clips. We use the TAU Urban Acoustic Scenes 2020 Mobile dataset as used in the DCASE 2020 challenge (DCASE20). The audio clips are recorded with different devices and further simulated devices are introduced. The performance is measured using accuracy on a dataset including unseen devices. The first place in the challenge used CNNs . Patchout accelerates training on this task, reaching state-of-the-art in less than an hour. Patchout also allows for fine-tuning on a single consumer GPU. It does, however, lead to a decrease in accuracy. Sound Event Recognition (Tagging) on FSD50K: The FSD50K dataset consists of 51K audio clips annotated with 200 sound event classes taken from the Audioset ontology . The dataset contains 100 hours of audio and is the second largest publicly available general purpose sound event recognition dataset after Audioset. Furthermore, the FSD50K evaluation set is of high quality, with each evaluation label being double-checked and assessed by two to five independent annotators . The reported results are on the official evaluation subset of FSD50K using the best model on the validation subset. The state-of-the-art in PSLA is achieved through CNN architecture and a collection of performance-improving methods such as ImageNet pre-training, label enhancement, balancing, data augmentation, and weight averaging. On this dataset, our approach significantly outperforms the current state-of-the-art. Fine-tuning PaSST-S and PaSST-S-N takes less than 2 hours and 1 hour, respectively.

Conclusion

We propose a new method for efficiently training transformers on audio spectrograms, achieving state-of-the-art performance on Audioset as well as several downstream tasks. Furthermore, Patchout significantly reduces compute complexity and memory requirements for training transformers. We investigate additional methods for reducing training complexity and propose two models, PaSST-S-L and PaSST-S-N, that outperform CNNs while having a faster training speed and comparable memory requirements. Our pre-trained models can be fine-tuned on several audio downstream tasks with little resources and little additional training time.

ACKNOWLEDGMENT

This work has been supported by the COMET-K2 Center of the Linz Center of Mechatronics (LCM) funded by the Austrian Federal Government and the Federal State of Upper Austria. The LIT AI Lab is financed by the Federal State of Upper Austria. The computational results presented have been achieved in part using the Vienna Scientific Cluster (VSC).