In this work, we propose a Multi-Window Masked Autoencoder (MW-MAE) fitted with a novel Multi-Window Multi-Head Attention (MW-MHA) module that facilitates the modelling of local-global interactions in every decoder transformer block through attention heads of several distinct local and global windows.
Empirical results on ten downstream audio tasks show that MW-MAEs consistently outperform standard MAEs in overall performance and learn better general-purpose audio representations, along with demonstrating considerably better scaling characteristics. Investigating attention distances and entropies reveals that MW-MAE encoders learn heads with broader local and global attention. Analyzing attention head feature representations through Projection Weighted Canonical Correlation Analysis (PWCCA) shows that attention heads with the same window sizes across the decoder layers of the MW-MAE learn correlated feature representations which enables each block to independently capture local and global information, leading to a decoupled decoder feature hierarchy.
Code for feature extraction and downstream experiments along with pre-trained models will be released publically.

§ INTRODUCTION

With rapid advances in hardware capabilities driving models of ever-increasing complexity and appetite for data, leveraging unlabelled data for learning effective deep representations has garnered significant interest. Self-supervised learning, which solves a pretext task that utilizes labels generated from the data itself, has emerged as a notable approach for training deep neural representations without labelled
data. Several methods for learning self-supervised representations from audio data have been proposed, including autoregressive methods that learn to predict the future based on the past input [1, 2, 3, 4, 5],  methods that learn contrastive representations from different views of the input [6, 7, 8, 9, 10, 11, 12], and
masked predictive modelling methods that learn to predict removed portions of the input data [13, 14, 15].

Together with the transformer architecture [16] and its successors [17, 18], masked predictive modelling has led to significant advances across natural language processing (NLP) [13, 19], computer vision [15, 20] and audio and speech processing [14, 21, 22]. Masked Autoencoders (MAEs) by [23] are a recent addition to the masked predictive modelling family.
Initially proposed for learning visual representations from randomly masked image patches, MAEs are experiencing widespread adoption across several domains [24, 25, 26, 27, 28, 29] due to their inherent scalability and simple design. In the audio domain, several recent works have adapted MAEs to learn a general-purpose audio representation from spectrogram inputs [30, 31]. These works address several challenges that are unique to the audio domain and exhaustively study the effect of masking strategies and other hyperparameters, providing vital information for training MAEs on audio data.

Recent works have shown that leveraging local information in the Multi-Head Attention (MHA) module of a transformer layer through convolutions [32], attention with local windows [18, 33] or pooling attention [34, 35, 36] can lead to improved performance. Within the framework of Masked Autoencoders, [37] evaluated the impact of local windowed attention [18] for audio representation learning, demonstrating better performance across 4 downstream audio recognition tasks. However, in these methods: (i) all the attention heads within a MHA module operate at the same local context, thus only capturing local information at the transformer layer level, and (ii) they require explicit approaches to mitigate the lack of connections across windows and to capture local-global information.

In this work, we propose Multi-Window Masked Autoencoders (MW-MAE) with both local and global attention for learning general-purpose audio representations from spectrogram inputs. Decoders in an MW-MAE are fitted with a novel Multi-Window Multi-Head Attention module (MW-MHA) (Fig <ref>). Each attention head in the proposed MW-MHA module computes self-attention over non-overlapping windows of different sizes, which facilitates modelling of local-global interactions in every decoder transformer block.
The proposed MW-MAEs outperform standard MAEs on 10 downstream audio recognition tasks.

At the same time, MW-MAEs adapt better to varying patch sizes and increasing number of patches, perform better in low-data scenarios, as well as demonstrate better performance and scaling characteristics with respect to changing encoder and decoder complexities.
Exhaustive exploratory analysis of attention distances and entropies shows that attention heads in MW-MAE encoders learn broader local-global attention as compared to standard MAEs, despite having an identical architecture. Further, analysis of feature representations from the decoder attention heads using Projection Weighted Canonical Correlation Analysis (PWCCA) [38] indicates that attention heads with the same window sizes across the decoder layers of the MW-MAE learn correlated feature representations leading to a decoupled feature hierarchy, confirming that MW-MHA modules learn local-global features in each decoder block.

§ BACKGROUND AND RELATED WORKS

MW-MHA module

An overview of the proposed Multi-Window Multi-Head Attention (MW-MHA) module, and the overall MW-MAE architecture. In MW-MHA, each attention head operates on non-overlapping windows of different sizes (coded by different colours) of the input matrices. As evident from b) MW-MAE uses the proposed MW-MHA block only in the decoder.

Recently, several works have been proposed for learning audio representations in a self-supervised manner. Most of these works can be loosely categorized into one or more of the following groups based on their underlying pretext task: (i) predictive; (ii) contrastive; and (iii) masked predictive modelling.
Several methods adopt a predictive coding approach, which aligns well with the sequential nature of audio input. Autoregressive predictive coding (APC) [2, 4, 3] is one such method which utilizes Recurrent Neural Networks (RNNs) to predict future elements of a sequential input given the past, whereas non-autoregressive approaches using Convolutional Neural Networks (CNNs) have also been proposed [5]. Contrastive predictive coding [1] optimizes a predictive coding objective in the latent space while simultaneously optimizing a contrastive objective function. This brings us to contrastive representation learning, which operates on the premise of learning a representation space that maximizes agreement between views from the same input sample while minimizing inter sample agreement. Several contrastive learning-based methods, originally proposed for computer vision [39, 40, 41], have been adapted for learning audio representations [7, 42, 43]. A widely used contrastive approach for learning speech representations is the Wav2Vec family of algorithms [9, 44, 10, 11], which learn contextualized speech representations by optimizing a contrastive objective between quantized latent representations and representations generated from masked time steps.
Finally, self-supervised learning methods based on masked predictive modelling have a simple premise: remove a portion of the input data, and learn to predict the removed content. After being (re-)popularized by the likes of BERT [13] in NLP and fueled by the recognition of the Transformer [16] as a viable cross-domain neural architecture, masked modelling has seen wide adoption in several domains, such as computer vision [20, 15], audio and speech [45] as well as a multi-domain self-supervised learning frameworks [46].
In the audio domain, several recent methods use masked predictive modelling to learn self-supervised representations [44]. These include HuBERT [14], which trains a BERT-like model to predict pre-determined cluster assignments from masked speech features, WavLM [21], which learns a joint denoising and masked prediction task, and SSaST [22], which jointly reconstructs and contrasts masked patches. More recently, [47] proposed an iterative masked modelling approach using an iterative self-distilled tokenizer that generates refined discrete labels from audio input data for the next stage of pretraining.

Masked Autoencoders (MAEs): Recently, [23] proposed Masked Autoencoders for learning self-supervised image representations. In an MAE, the input is split into non-overlapping patches, which are then linearly projected to a fixed dimension by the patch embedding layer. A large subset of these patches is masked out (, over 75% of the patches), and the unmasked patches are then encoded by a Vision Transformer (ViT) [17]. With learnable mask tokens filled in the correct positions to restore the original patch order, these encoded patches are fed to a transformer based decoder, whose objective is to learn to reconstruct the masked patches.
The high masking ratio allows large encoders to be paired with significantly smaller decoders due to the reduced encoding complexity, while simultaneously forcing the encoder to learn better contextualized representations by minimizing extrapolation from redundant neighbouring patches. Several recent works based on MAEs have been proposed for training general-purpose audio representations. [30] explored a joint discriminative and generative objective for training audio MAEs and evaluated fine-tuning performance on 5 downstream tasks.
By training shallow downstream classifiers on 15 downstream tasks in accordance with the HEAR-2021 [48] protocol, [31] investigated various hyper-parameters such as patch size and the effect of input audio clip duration on model performance. [37] investigated shifting windows based local self-attention [18] of a fixed context ($4 \times 4$ windows) in all but the last few layers of audio MAEs. In contrast, in this work we propose a Masked Autoencoder fitted with a novel Multi-Window Multi-Head Attention module that can model attention at several context levels and can capture local-global interactions in every transformer layer.

§ PROPOSED APPROACH

§.§ Multi-Window Multi-Head Attention

To better capture local-global attention, we propose a Multi-Window Multi-Head Attention (MW-MHA) module, where each attention head computes self-attention across spectrogram patches in different local windows and then combines the contribution of each attention head, as illustrated in Figure <ref>. We define MW-MHA with $h$ parallel heads as follows:

\begin{align}
\mathrm{MWMHA}(Q, K, V) &= \mathrm{Concat}(\mathrm{winHead_1}, ..., \mathrm{winHead_h})W^O\\
\mathrm{winHead}_i &= \mathrm{WinAttention}(QW^Q_i, KW^K_i, VW^V_i, win_i)
\end{align}

As opposed to MHA, each head $\mathrm{winHead}_i$ computes local self-attention over non-overlapping windows of size $win_i$ by partitioning input matrices $QW^Q_i, KW^K_i, VW^V_i \in \R^{n \times d_k}$ into $Q_{win_i}, K_{win_i}, V_{win_i} \in \R^{m \times win_i \times d_k}$, given that $n = m\times{win_i}$.
This is followed by computing standard self-attention $X_{win_i} = \mathrm{Attention}(Q_{win_i}, K_{win_i}, V_{win_i})$ [16] on these partitioned inputs. Finally, $ X_{win_i} \in \R^{m \times win_i \times d_k}$ is reshaped to $ X \in \R^{n \times d_k}$ to get the output.

In the proposed MW-MHA module, individual attention heads capture information at multiple local contexts, and the final projection matrix $W^{O}_{i} \in \R^{hd_k \times d_{m}}$ learns the contribution of each of these heads, allowing inter-window interaction and connection. This design facilitates learning both local and global time-frequency information at several granularities in every transformer block (as supported by exploratory analysis in Section <ref>). This is in contrast to shifting [18, 49], striped [33] windowed self-attention, or pooling attention [34, 35, 36], where all attention heads within the same block have the same window size and thus only perform local self-attention at the block level.
Pseudo-code for the proposed MW-MHA is provided in Appendix <ref>.

§.§ Masked Autoencoder with Multi-Window Multi-Head Attention

Patch embeddings, masking strategy and masking ratio: We use mel-spectrograms as inputs, partitioning them into non-overlapping patches,
which are then flattened and embedded into linear projections.
For encoding positional information, we use fixed sinusuidal positional embeddings, similar to [30, 31, 37]. We use a high masking ratio ($80\%$) and random unstructured masking, which have been shown to work well for audio [31, 37].

Encoder: In line with previous work [23, 31, 37], we use a Vision Transformer (ViT) [17] based encoder, which only processes non-masked patches (20% in this work).
Due to the random masking strategy, majority of the patches are not processed by the encoder at training time. This minimizes the benefit of using the proposed MW-MHA modules in the encoder transformer blocks (as evidenced by experiments in Section <ref>). Thus, transformer blocks in our encoder use standard Multi-Head Attention.

Decoder with Multi-Window Multi-Head Attention: We add fixed sinusoidal positional embeddings to the encoded visible patches concatenated with trainable masked tokens after restoring original patch order. The resulting tensor is then fed to the decoder, which is also a stack of transformer layers, followed by a linear head that reconstructs the original input spectrogram. This is consistent with previous works [23, 31, 37]. Given that the decoder processes all the patches, we replace the Multi-Head Attention module with the proposed Multi-Window Multi-Head Attention, thus modelling local-global attention in every decoder block.

Selecting window sizes: We follow a simple rule for determining the window sizes of each constituent $\mathrm{winHead}_i$: given the total number of patches $n_{p}$, we simply take all non-unary factors of $n_p$ and add two additional global self-attention heads. As an example, our default configuration yields $n_p = 250$, and thus the window sizes for each MW-MHA module in all decoder blocks will be $[2, 5, 10, 25, 50, 125, 250, 250]$ for a total of $8$ attention heads, which is a reasonable number of attention heads inline with previous research [23, 37]. Not only is this method simple to follow, but it also scales well with number of patches, effectively covering several possible local context levels.

Pre-training objective: During pre-training, we optimize a loss function that computes mean squared error (MSE) between the predicted masked patches and their corresponding input spectrogram patches. In early experiments, we observed reduced performance when using per-patch normalization, and thus we do not normalize target spectrogram patches.

§ EXPERIMENTS

§.§ Datasets and Tasks

Pre-training: We use the full AudioSet dataset [50] (AS-5k) for pre-training MAEs and MW-MAEs. With over 5000 hours of audio data distributed in  2 million 10-second weakly annotated YouTube clips spanning 527 classes, AudioSet is one of the largest publicly available audio corpora.

Downstream tasks: Recently, several standardized benchmarks have been proposed to evaluate audio representations thoroughly across a wide variety of domains, such as SUPERB [51] and HEAR [48]. While both benchmarks offer avenues for fast, reproducible and accessible comparison of audio representations, the SUPERB benchmark focuses primarily on speech-processing applications. In contrast, the HEAR benchmark consists of 19 tasks spanning diverse audio domains of speech, music and environmental sounds and redistributes standardized and preprocessed datasets.
However, some of these tasks are simply smaller subsets of one another, whereas performance on some HEAR tasks has been demonstrated to be correlated [48]. For evaluating audio representations, we utilize a subset of the HEAR benchmark which consists of ten diverse tasks spanning multiple domains, which can be found in Appendix <ref> along with the underlying selection criterion. We believe the selected tasks constitute a balanced evaluation protocol that facilitates assessment of audio representations without doing excessive evaluations. For downstream evaluation, we follow the HEAR protocol, where for each task, a shallow downstream classifier is trained on top of fixed features extracted using a pretrained model. This practice has become quite prevalent and allows the evaluation of how representations generalize to a broad range of tasks without the drawbacks of fine-tuning large, heterogeneous neural networks.

Measuring overall performance: Given the wide variety of downstream tasks and feature representations evaluated, a single metric to quantify the performance would significantly aid analysis. However, given the differing difficulty levels of the tasks as well as outliers arising from the nature of the representations evaluated, simply averaging the scores is not sufficient. To counteract this, we utilize a normalized overall score to track overall performance of a given audio representation. Mathematically, overall score $s(m) \in [0.,100.]$ of a model $m$ is given as:

\begin{equation}
s(m) = \frac{1}{|T|}\sum_{t \in T} \frac{x_t(m) - min_t}{max_t - min_t} * 100
\end{equation}

where $x_t(m)$ denotes performance of the model $m$ on task $t$, and $min_t$ and $max_t$ represent the worst and the best performance across all models on the task. By taking the relative performance of the best and the worst approach on a task into consideration, this overall score takes how hard the task is to improve on in consideration. It is worth noting that the normalized score is computed across all the evaluated methods in all upcoming sections, including ablations. This is similar to the overall score used by the public leaderboard of the SUPERB [51] benchmark, except that we do not set the normalized value of the worst performing method to $0$, and the proposed overall score has an upper range of $100.0$.

§.§ Implementation Details

Features: We use log-scaled mel spectrograms with a window size of $25$ ms, a hop size of $10$ ms and $F=80$ mel-spaced frequency bins in the $50-8000$ Hz range, extracted using the torchaudio [52] toolkit. All datasets have a sampling frequency of $16000$ Hz. Instead of normalizing by dataset statistics, we adopt a per-instance standardization scheme.

Pre-training: We use the AudioSet dataset for pre-training our Masked Autoencoders. We extract log-scaled mel spectrograms for the entire AudioSet dataset and randomly crop a segment 200 timesteps in length from each data sample.
Our default configuration consists of a ViT-B encoder. All our MAE variants accept a $200\times80$-dimensional (T $\times$ F, respectively) input corresponding to an audio duration of 2 seconds, which achieves performance on-par with longer input durations as demonstrated by [31].
For our default configuration, our patch embedding computes non-overlapping patches with a patch size of $(4\times16)$, given it's desirable performance v/s complexity tradeoff as found by [31].
A key characteristic of the Masked Autoencoder paradigm is its asymmetric design, which allows pairing small decoders with large encoders while scaling favourably for linear probe performance [23].
Thus, in contrast to [37], we adopt a smaller 4-layer deep transformer-based decoder of width 384 and 8 attention heads for our default configuration.

We train Masked Autoencoders with the proposed MW-MHA module, which are referred to as MW-MAEs, as well as their standard MAE counterparts. All MAEs are pre-trained for 100 epochs with a batch size of 1024 and a weight decay of 0.05 on a single TPU-v3 VM with 8 TPU cores, with the default configuration taking around 36 hours to train. We warm up for ten epochs to a base learning rate of 1e-5, followed by a cosine decay schedule. A masking ratio of 0.8 with unstructured random masking is used, and no other data augmentations are used during pre-training.

Training downstream models: We first extract fixed feature embeddings for all downstream tasks to train downstream models. In the MAE framework, the decoder is discarded after pretraining and feature embeddings are extracted using just the encoder. To generate scene embeddings consistent with the HEAR protocol, we use the exact patch aggregation process as [31]: we break audio clips into non-overlapping 2 second chunks, concatenating the features in time and finally taking a mean over the time axis to generate a fixed vector representation independent of the input audio duration.
The hear-eval-kit, released alongside the HEAR benchmark, was used to extract fixed feature embeddings and to train a shallow MLP classifier with a single hidden layer with 1024 neurons for each task in a reproducible manner. Experiments are repeated with at least ten random seeds for each task, resulting in 100 experiments for every evaluated representation.

Comparison with various audio representations from the literature. $95\%$ confidence intervals are reported over $10$ runs on downstream classifiers. We pre-trained all highlighted audio representations, with different gray levels indicating directly comparable MAE and MW-MAE configurations. For other pre-trained audio representations, publicly available official implementations were used. All downstream models were trained by us using the hear-eval-kit.
$s(m)$ denotes the proposed normalized overall score (Sec <ref>)
*: same configuration as MSM-200 16x4 [31], with 8 attention heads in the decoder instead of 6. For model parameter counts, refer to Appendix <ref>

1lModel     1lPT-Data     BO     CD     ESC-50     LC     Mri-S     Mri-T     NS-5h     SC-5h     F50K     VL     $s(m)$

3lNaive Baselines     5l

HEAR-Naive [48]     -     52.62.4     30.90.8     5.80.2     33.51.1     38.01.3     36.41.9     18.64.4     8.50.4     7.10.2     11.20.5     5.00.7

3lSupervised     5l

PaSST-base [53]     AS-5k     94.90.5     61.00.3     94.80.3     60.10.2     96.50.1     87.60.6     23.30.9     66.61.4     64.20.1     25.50.8     73.50.4

3lSSL     5l

W2V2-base [10]     LS-960     74.01.0     46.40.3     31.10.4     51.20.2     77.30.2     55.10.3     7.40.8     90.80.3     18.10.1     35.50.8     43.10.2

W2V2-large [10]     VP-100k     93.10.7     66.90.4     60.10.5     62.40.3     93.90.1     77.40.2     42.01.0     87.60.5     34.20.1     53.61.0     74.00.4

WavLM-base [21]     LS-960     89.40.7     56.30.2     46.60.4     63.20.3     95.10.1     83.40.2     37.30.8     57.20.8     29.90.1     22.60.6     60.50.2

WavLM-large [21]     Mix-94k     96.40.5     57.20.2     47.90.4     61.10.3     96.80.1     89.50.1     53.70.5     46.20.8     29.00.1     23.70.9     64.00.2

HuBERT-base [14]     LS-960     92.10.6     70.80.2     57.80.6     56.50.3     94.40.1     84.90.3     19.40.7     93.20.1     32.30.1     61.80.6     72.50.2

HuBERT-large [14]     LL-60k     94.10.7     70.70.1     60.30.4     59.90.2     95.30.1     83.50.3     19.30.8     83.20.7     31.50.1     66.10.9     73.40.3

SSaST-base [22]     AS+LS     93.40.9     56.50.2     68.40.4     60.70.3     96.70.1     96.30.1     66.80.7     53.51.3     38.20.1     28.50.9     71.70.2

BEATs-iter3 [47]     AS-5k     94.00.8     67.30.2     83.70.3     68.00.2     94.70.1     95.80.1     69.40.8     85.20.3     53.60.2     38.51.0     85.70.3

3lMAE based     5l

AudioMAE [37]     AS-5k     93.70.6     68.20.2     60.60.4     42.20.2     89.20.2     86.60.2     64.50.8     28.61.5     37.90.1     29.71.0     62.90.3

MAE-B-4x16-4l*     AS-5k     96.20.3     72.20.2     80.90.4     67.30.3     97.40.1     98.30.1     68.30.4     89.40.3     50.40.1     43.10.9     88.10.2

MAE-B-5x5-4l     AS-5k     96.00.4     70.90.2     80.90.4     67.60.4     97.60.1     98.40.0     69.30.4     88.40.3     49.30.2     37.70.6     86.80.2

MAE-L-4x16-8l     AS-5k     96.10.4     73.80.1     81.60.3     68.50.2     97.60.1     98.30.0     69.00.5     91.20.2     51.80.1     46.90.8     90.00.2

3lProposed     5l

MW-MAE-B-4x16-4l     AS-5k     96.00.5     73.10.3     81.20.4     68.80.2     97.40.1     97.90.1     69.30.6     90.90.2     51.20.2     44.20.9     89.20.2

MW-MAE-B-5x5-4l     AS-5k     96.60.4     73.80.4     82.00.3     70.10.4     97.50.1     98.30.1     72.90.5     91.70.2     51.30.1     44.20.6     90.60.1

MW-MAE-L-4x16-8l     AS-5k     95.90.3     76.10.2     83.60.3     69.70.3     97.40.0     98.20.1     71.20.7     93.00.1     53.50.1     51.90.7     92.60.2

§.§ Comparison with Existing Works

Table <ref> shows how MW-MAE fares against recent audio representations. The highlighted model configurations that we pre-trained from scratch on AudioSet have the following naming convention: the first substring shows the type of MAE (vanilla or proposed MW-MAE), followed by a single alphabet denoting ViT Encoder configuration. This is followed by the patch size used, and finally, the depth of the decoder.
It's worth noting that while embedding sizes of MAE and corresponding MW-MAE configurations are the same, the embedding sizes of other methods can be different. This is inline with the current consensus of evaluating self-supervised representations in the audio domain [51, 48].
MW-MAE configurations outperform all other comparable MAEs, with the largest "MW-MAE-L-16x4-8l" configuration outperforming all the methods in overall performance (92.60.2). MW-MAEs also outperform AudioMAE with standard shifting window based attention, as well as BEATs-iter3, which is the pretrained representation obtained after 3 stages of self-distilled learning as proposed by [47].
MW-MAEs perform exceptionally well on pitch perception (NS-5h), while achieving performance on-par with speech specific representations such as WavLM, HuBERT and Wav2Vec2 (denoted W2V2) for Keyword spotting (SC-5h). Perhaps more surprisingly, they outperform speech representations trained on much larger training sets on the emotion recognition (CREMA-D) as well as speaker count classification (LibriCount) tasks.
While PaSST, which is a recent state-of-the-art approach for training supervised transformers on AudioSet, outperforms every model on ESC-50 and FSD50K tasks, the overall performance of the proposed approach is significantly better. Overall, the proposed MW-MAEs learn a better general-purpose audio representation than standard MAEs, generalizing well to several audio domains and demonstrating excellent overall performance in comparison to recent audio representations.

§.§ Key Model Characteristics

We conduct several experiments to examine key differences between MAE and the proposed MW-MAE. While we have only reported overall score $s(m)$, detailed results for all these experiments can be found in Section <ref>.

2*Model     Downstream     Linear Probe     Fine-tuning

$s(m)$     (mAP)     (mAP)

MAE Base     88.10.2     18.90.0     23.80.1

MW-MAE Base (decoder only)     89.20.2     20.20.0     23.90.1

MW-MAE Base (enc+dec)     89.10.3     20.20.0     24.20.1

Performance impact of MW-MHA module placement

MW-MHA in the encoder: As previously mentioned, adding MW-MHA to the encoder block does not improve downstream performance.
Further, we also investigate the impact of linear probing as well as fine-tuning the entire encoder stack for in-domain classification on AudioSet-20k balanced subset. No data augmentations were used.
As evident from Table <ref>, when compared with including MW-MHA blocks in the decoder only, there is no performance benefit to adding MW-MHA blocks to the encoder for neither downstream performance, nor for in-domain linear probe on AudioSet-20k. However, when fine-tuning the entire encoder stack, adding MW-MHA blocks to the encoder provides a slight improvement and is worth considering.

Performance impact of various patch sizes:  In an MAE, the patch embedding layer generates non-overlapping patches from the input. Thus, the size of the patch governs the number of patches as well as the time-frequency resolution that the transformer layers work at, making it an important hyperparameter to investigate.

Patch size

Encoder complexity

Decoder depth

Pre-training data

Ablation experiments comparing standard MAE v/s proposed MW-MAE at different patch sizes (a), encoder complexity (b), decoder depth (c) as well as amount of pre-training data used (d). $s(m)$ is the proposed overall score (Sec <ref>). Detailed results can be found in Appendix <ref>.

Figure <ref> shows how different patch sizes affect downstream performance. The proposed MW-MAE model, with an overall score of 90.60.1, outperforms standard MAE for every patch size for identical decoder configurations. It's also worth noting that MAE performance degrades as we decrease the patch size beyond $4\times16$, whereas MW-MAE performance continues to improve. These observations show that the proposed MW-MAE adapts better to varying patch sizes and time-frequency resolutions, while scaling well with increasing number of patches.

Encoder size: As shown in Figure <ref>, we investigate how encoders of five different complexities affect overall performance. All the trained models have the same decoder configuration (384 neurons, depth=4, h=8). With an overall score of 89.20.2, MW-MAE with the ViT-Base encoder performs better than MAEs with encoders of any size in this experiment. The most prominent performance gap is observed for the ViT-Large setting, where MAE and MW-MAE attain overall scores of 88.20.2 and 92.30.2, respectively. The drop in performance for the ViT-Huge encoder for both MAEs and MW-MAEs suggests possible overfitting.

Decoder depth: In Figure <ref>, we show how increasing decoder complexity by increasing decoder depth affects overall performance.
As expected, increasing decoder depth improves performance for both methods. For decoder depth=$8$, MW-MAE (89.90.2) outperforms MAE (88.20.1) by a considerable margin in overall performance.
We also observed that with an overall score of 88.30.2, MW-MAE with depth=$2$ performs on par with MAEs with up to $4$ decoder blocks.
This observation complements the inherent asymmetric nature of Masked Autoencoders, and thus the proposed MW-MAE performs favourably in terms of complexity and scalability.

Pre-training data: Finally, Figure <ref> depicts how performance varies as we reduce the amount of data used for pre-training. Overall, performance for both the MAE and the proposed MW-MAE methods continues to decrease monotonically as we remove more and more data.
However, the performance loss trend for MW-MAE is much more favourable. A 90% reduction in the amount of pre-training data results in a $28.17\%$ reduction in performance for standard MAEs (from 88.10.2 to 63.30.2), whereas MW-MAE only suffers a $13.5\%$ drop in performance (from 89.20.2 to 77.20.3).
Thus, we conclude that the proposed MW-MAEs are more adept at handling low-data scenarios in comparison to standard MAEs.

§ EXPLORATORY ANALYSIS

§.§ Inspecting encoder attention heads

Attention entropies in encoder blocks

Mean attention distances for the select transformer blocks.

Investigating MAE and MW-MAE encoder attention heads. (a) depicts average entropies of encoder attention heads over the course of pretraining in a every encoder transformer block. (b) Depicts mean attention distance distributions of the first two and the last two transformer blocks at different amounts of pretraining data used.

Analyzing attention entropies: We first analyze individual attention heads in a ViT-Medium encoder (depth=12, h=8). Figure <ref> shows scatter plots of average entropies of individual encoder attention heads computed over the entire NSynth Pitch 5h validation set on a block-by-block basis at different stages during pre-training. It's worth noting that the higher the entropy, the more global the attention, with lower attention mass spent on closer tokens [54], and thus, a higher variance in entropies of individual attention heads highlights more spread out local and global attention. In the early epochs, MAE encoders actually have higher variance in entropy distribution, especially in the latter transformer layers. As pretraining goes on, interestingly, this effect is reversed, and the attention heads in the MW-MAE encoder now start converging towards high entropy variance configurations in the early layers.

Analyzing attention distances: We analyze mean attention distances for attention heads in the first two and the last two encoder blocks. Similar to [55], we compute attention-weighted patch distances between the query patch position and the locations it attends to for each attention head, averaging it for all patches positions. This is repeated for all inputs in the FSD50K validation set. Figure <ref> depicts the distribution of mean attention distance for MAE and MW-MAE encoders (base configuration) pretrained with different amounts of training data. We can observe that MW-MAE attention heads demonstrate a broader distribution of attention distances, modelling local-global attention better than the MAE encoder especially in the first two transformer blocks.

Comparing features learned by different attention heads in the encoder and the decoder of a standard MAE and the proposed MW-MAE using PWCCA. Each tick separates the attention heads of a transformer block from the next.

From these observations, we can conclude that in an MW-MAE, the decoder fitted with an MW-MHA can force the encoder to better capture local-global interactions even without explicit windowed attention modules, leading to improved performance.

§.§ Comparing attention feature representations through PWCCA

Several recent works have used Canonical Correlation Analysis (CCA) to compare feature representations and learning dynamics of deep neural networks [56, 57].
We use Projection Weighted CCA (PWCCA) [38], which computes a weighted mean of the CCA vectors to compare the representations learned by individual attention heads of the encoder and the decoder in identically configured MAE and MW-MAE (ViT-M encoder: depth=12, h=8; Default decoder: 384 neurons, depth=4, h=8).
MW-MAE decoder uses default attention head window sizes as specified in Sec <ref>.
Figure <ref> depicts correlation matrices of measured PWCCA score between attention heads. We can observe a remarkable difference in correlation between the decoders: feature representations from the MW-MAE decoder attention heads with the same window sizes are strongly correlated across decoder layers, whereas attention heads with global self-attention (7, 8, 15, 16, 23, 24, 31, 32) are the least correlated, consistent with observations made for the MAE decoder. These observations suggest a decoupling of different aspects of the feature hierarchy in the MW-MAE decoder, as attention heads of specific window sizes in each decoder block capture local information at a specific granularity, which is in line with our original hypothesis. These observations are also corroborated by decoder depth ablation experiments from Sec <ref>, where we observed that a MW-MAE with a single transformer block performs on par with MAEs fitted with up to 4 decoder blocks. Finally, the difference in correlation matrices between the encoders is much less stark, which is expected since both use standard MHA blocks.

§ CONCLUSION

This work presents Multi-Window Masked Autoencoder (MW-MAE) for learning general-purpose audio representations.
Decoders in MW-MAEs are fitted with a novel Multi-Window Multi-Head Attention (MW-MHA) module, which learns information captured at multiple granularities of local-global context by its constituent attention heads computing self-attention over different non-overlapping windows.
Empirical experiments on ten downstream tasks show that the proposed MW-MAEs consistently outperform standard MAEs in overall performance when pre-trained on the AudioSet dataset, demonstrating better scaling characteristics.

Exploratory analyses highlight key differences between the attention representations learned by standard MAEs and the proposed MW-MAEs. Based on attention entropy and mean attention distance analysis, we discover that encoder attention heads in an MW-MAE better capture local-global interactions, even without explicit local-global attention modules. We also learn that attention heads of the same window size across the transformer blocks of the MW-MAE decoder are correlated, learning a decoupled feature hierarchy allowing transformers to capture relevant information at the block level, supporting our original motivation.

§ ACKNOWLEDGEMENTS

This project is supported by the Pioneer Centre for Artificial Intelligence, Denmark.[<https://www.aicentre.dk>]. We are also thankful to the TPU Research Cloud Program [<https://sites.research.google/trc/about>], a Google Research Initiative, for providing TPU v2 and v3 devices used in this project.

[1]
A. v. d. Oord, Y. Li, and O. Vinyals, “Representation learning with contrastive predictive coding,” arXiv preprint arXiv:1807.03748, 2018.

[2]
Y.-A. Chung, W.-N. Hsu, H. Tang, and J. R. Glass, “An unsupervised autoregressive model for speech representation learning.” in Interspeech 2019, 2019, pp. 146–150.

[3]
Y.-A. Chung and J. Glass, “Generative pre-training for speech with autoregressive predictive coding,” in 2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), 2020, pp. 3497–3501.

[4]
Y.-A. Chung, H. Tang, and J. R. Glass, “Vector-quantized autoregressive predictive coding.” in Interspeech 2020, 2020, pp. 3760–3764.

[5]
A. H. Liu, Y.-A. Chung, and J. R. Glass, “Non-autoregressive predictive coding for learning speech representations from local dependencies,” in Interspeech 2021, 2021.

[6]
M. Tagliasacchi, B. Gfeller, F. de Chaumont Quitry, and D. Roblek, “Pre-training audio representations with self-supervision,” IEEE Signal Processing Letters, vol. 27, pp. 600–604, 2020.

[7]
A. Saeed, D. Grangier, and N. Zeghidour, “Contrastive learning of general-purpose audio representations,” in 2021 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), 2021, pp. 3875–3879.

[8]
E. Fonseca, D. Ortego, K. McGuinness, N. E. O'Connor, and X. Serra, “Unsupervised contrastive learning of sound event representations,” in 2021 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), 2021, pp. 371–375.

[9]
S. Schneider, A. Baevski, R. Collobert, and M. Auli, “wav2vec: Unsupervised Pre-Training for Speech Recognition,” in Proc. Interspeech 2019, 2019, pp. 3465–3469.

[10]
A. Baevski, Y. Zhou, A. Mohamed, and M. Auli, “wav2vec 2.0: A framework for self-supervised learning of speech representations,” in Advances in Neural Information Processing Systems, vol. 33, 2020, pp. 12 449–12 460.

[11]
W.-N. Hsu, A. Sriram, A. Baevski, T. Likhomanenko, Q. Xu, V. Pratap, J. Kahn, A. Lee, R. Collobert, G. Synnaeve, and M. Auli, “Robust wav2vec 2.0: Analyzing domain shift in self-supervised pre-training,” in Interspeech 2021, 2021.

[12]
A. K. Sarkar, Z.-H. Tan, H. Tang, S. Shon, and J. Glass, “Time-contrastive learning based deep bottleneck features for text-dependent speaker verification,” IEEE/ACM Transactions on Audio, Speech, and Language Processing, vol. 27, no. 8, pp. 1267–1279, 2019.

[13]
J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova, “BERT: Pre-training of deep bidirectional transformers for language understanding,” in Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers).1em plus 0.5em minus 0.4emMinneapolis, Minnesota: Association for Computational Linguistics, Jun. 2019, pp. 4171–4186.

[14]
W.-N. Hsu, B. Bolte, Y.-H. H. Tsai, K. Lakhotia, R. Salakhutdinov, and A. Mohamed, “Hubert: Self-supervised speech representation learning by masked prediction of hidden units.” IEEE Transactions on Audio, Speech, and Language Processing, pp. 1–1, 2021.

[15]
Z. Xie, Z. Zhang, Y. Cao, Y. Lin, J. Bao, Z. Yao, Q. Dai, and H. Hu, “Simmim: A simple framework for masked image modeling,” in Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2022, pp. 9653–9663.

[16]
A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin, “Attention is all you need,” Advances in neural information processing systems, vol. 30, 2017.

[17]
A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, J. Uszkoreit, and N. Houlsby, “An image is worth 16x16 words: Transformers for image recognition at scale,” in International Conference on Learning Representations, 2021.

[18]
Z. Liu, Y. Lin, Y. Cao, H. Hu, Y. Wei, Z. Zhang, S. Lin, and B. Guo, “Swin transformer: Hierarchical vision transformer using shifted windows,” in Proceedings of the IEEE/CVF international conference on computer vision, 2021, pp. 10 012–10 022.

[19]
M. Lewis, Y. Liu, N. Goyal, M. Ghazvininejad, A. Mohamed, O. Levy, V. Stoyanov, and L. Zettlemoyer, “BART: Denoising sequence-to-sequence pre-training for natural language generation, translation, and comprehension,” in Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics.1em plus 0.5em minus 0.4emOnline: Association for Computational Linguistics, Jul. 2020, pp. 7871–7880.

[20]
H. Bao, L. Dong, S. Piao, and F. Wei, “BEit: BERT pre-training of image transformers,” in International Conference on Learning Representations, 2022.

[21]
S. Chen, C. Wang, Z. Chen, Y. Wu, S. Liu, Z. Chen, J. Li, N. Kanda, T. Yoshioka, X. Xiao et al., “Wavlm: Large-scale self-supervised pre-training for full stack speech processing,” IEEE Journal of Selected Topics in Signal Processing, vol. 16, no. 6, pp. 1505–1518, 2022.

[22]
Y. Gong, C.-I. Lai, Y.-A. Chung, and J. Glass, “Ssast: Self-supervised audio spectrogram transformer,” in Proceedings of the AAAI Conference on Artificial Intelligence, vol. 36, 2022, pp. 10 699–10 709.

[23]
K. He, X. Chen, S. Xie, Y. Li, P. Dollár, and R. Girshick, “Masked autoencoders are scalable vision learners,” in Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2022, pp. 16 000–16 009.

[24]
C. Feichtenhofer, Y. Li, K. He et al., “Masked autoencoders as spatiotemporal learners,” Advances in neural information processing systems, vol. 35, pp. 35 946–35 958, 2022.

[25]
C. Wei, H. Fan, S. Xie, C.-Y. Wu, A. Yuille, and C. Feichtenhofer, “Masked feature prediction for self-supervised visual pre-training,” in Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2022, pp. 14 668–14 678.

[26]
Z. Hou, X. Liu, Y. Cen, Y. Dong, H. Yang, C. Wang, and J. Tang, “Graphmae: Self-supervised masked graph autoencoders,” in Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, 2022, pp. 594–604.

[27]
Y. Pang, W. Wang, F. E. Tay, W. Liu, Y. Tian, and L. Yuan, “Masked autoencoders for point cloud self-supervised learning,” in Computer Vision–ECCV 2022: 17th European Conference, Tel Aviv, Israel, October 23–27, 2022, Proceedings, Part II.1em plus 0.5em minus 0.4emSpringer, 2022, pp. 604–621.

[28]
R. Bachmann, D. Mizrahi, A. Atanov, and A. Zamir, “Multimae: Multi-modal multi-task masked autoencoders,” in Computer Vision–ECCV 2022: 17th European Conference, Tel Aviv, Israel, October 23–27, 2022, Proceedings, Part XXXVII.1em plus 0.5em minus 0.4emSpringer, 2022, pp. 348–367.

[29]
Y. Seo, D. Hafner, H. Liu, F. Liu, S. James, K. Lee, and P. Abbeel, “Masked world models for visual control,” in Conference on Robot Learning.1em plus 0.5em minus 0.4emPMLR, 2023, pp. 1332–1344.

[30]
A. Baade, P. Peng, and D. Harwath, “MAE-AST: Masked Autoencoding Audio Spectrogram Transformer,” in Proc. Interspeech 2022, 2022, pp. 2438–2442.

[31]
D. Niizumi, D. Takeuchi, Y. Ohishi, N. Harada, and K. Kashino, “Masked spectrogram modeling using masked autoencoders for learning general-purpose audio representation,” in HEAR: Holistic Evaluation of Audio Representations (NeurIPS 2021 Competition), ser. Proceedings of Machine Learning Research, vol. 166.1em plus 0.5em minus 0.4emPMLR, 13–14 Dec 2022, pp. 1–24.

[32]
A. Gulati, J. Qin, C.-C. Chiu, N. Parmar, Y. Zhang, J. Yu, W. Han, S. Wang, Z. Zhang, Y. Wu, and R. Pang, “enConformer: Convolution-augmented Transformer for Speech Recognition,” in enInterspeech 2020.1em plus 0.5em minus 0.4emISCA, Oct. 2020, pp. 5036–5040.

[33]
X. Dong, J. Bao, D. Chen, W. Zhang, N. Yu, L. Yuan, D. Chen, and B. Guo, “enCSWin Transformer: A General Vision Transformer Backbone With Cross-Shaped Windows,” in enProceedings of the IEEE/CVF international conference on computer vision, 2022, pp. 12 124–12 134.

[34]
H. Fan, B. Xiong, K. Mangalam, Y. Li, Z. Yan, J. Malik, and C. Feichtenhofer, “Multiscale vision transformers,” in Proceedings of the IEEE/CVF international conference on computer vision, 2021, pp. 6824–6835.

[35]
Y. Li, C.-Y. Wu, H. Fan, K. Mangalam, B. Xiong, J. Malik, and C. Feichtenhofer, “Mvitv2: Improved multiscale vision transformers for classification and detection,” in Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2022, pp. 4804–4814.

[36]
W. Zhu and M. Omar, “Multiscale audio spectrogram transformer for efficient audio classification,” in ICASSP 2023-2023 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP).1em plus 0.5em minus 0.4emIEEE, 2023, pp. 1–5.

[37]
P.-Y. Huang, H. Xu, J. Li, A. Baevski, M. Auli, W. Galuba, F. Metze, and C. Feichtenhofer, “Masked autoencoders that listen,” Advances in Neural Information Processing Systems, vol. 35, pp. 28 708–28 720, 2022.

[38]
A. Morcos, M. Raghu, and S. Bengio, “Insights on representational similarity in neural networks with canonical correlation,” in Advances in Neural Information Processing Systems, vol. 31.1em plus 0.5em minus 0.4emCurran Associates, Inc., 2018.

[39]
T. Chen, S. Kornblith, M. Norouzi, and G. Hinton, “A simple framework for contrastive learning of visual representations,” in ICML 2020: 37th International Conference on Machine Learning, vol. 1, 2020, pp. 1597–1607.

[40]
X. Chen and K. He, “Exploring simple siamese representation learning,” in Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2021, pp. 15 750–15 758.

[41]
J.-B. Grill, F. Strub, F. Altché, C. Tallec, P. H. Richemond, E. Buchatskaya, C. Doersch, B. A. Pires, Z. D. Guo, M. G. Azar, B. Piot, K. Kavukcuoglu, R. Munos, and M. Valko, “Bootstrap your own latent: A new approach to self-supervised learning,” in Advances in Neural Information Processing Systems, vol. 33, 2020, pp. 21 271–21 284.

[42]
D. Niizumi, D. Takeuchi, Y. Ohishi, N. Harada, and K. Kashino, “Byol for audio: Self-supervised learning for general-purpose audio representation,” in 2021 International Joint Conference on Neural Networks (IJCNN), 2021, pp. 1–8.

[43]
G. Elbanna, N. Scheidwasser-Clow, M. Kegler, P. Beckmann, K. El Hajal, and M. Cernak, “Byol-s: Learning self-supervised speech representations by bootstrapping,” in HEAR: Holistic Evaluation of Audio Representations.1em plus 0.5em minus 0.4emPMLR, 2022, pp. 25–47.

[44]
A. Baevski, S. Schneider, and M. Auli, “vq-wav2vec: Self-supervised learning of discrete speech representations,” in International Conference on Learning Representations, 2019.

[45]
A. T. Liu, S.-w. Yang, P.-H. Chi, P.-c. Hsu, and H.-y. Lee, “Mockingjay: Unsupervised speech representation learning with deep bidirectional transformer encoders,” in ICASSP 2020-2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP).1em plus 0.5em minus 0.4emIEEE, 2020, pp. 6419–6423.

[46]
A. Baevski, W.-N. Hsu, Q. Xu, A. Babu, J. Gu, and M. Auli, “data2vec: A general framework for self-supervised learning in speech, vision and language,” in Proceedings of the 39th International Conference on Machine Learning, ser. Proceedings of Machine Learning Research, vol. 162.1em plus 0.5em minus 0.4emPMLR, 17–23 Jul 2022, pp. 1298–1312.

[47]
S. Chen, Y. Wu, C. Wang, S. Liu, D. Tompkins, Z. Chen, W. Che, X. Yu, and F. Wei, “BEATs: Audio pre-training with acoustic tokenizers,” in Proceedings of the 40th International Conference on Machine Learning, ser. Proceedings of Machine Learning Research, vol. 202.1em plus 0.5em minus 0.4emPMLR, 23–29 Jul 2023, pp. 5178–5193.

[48]
J. Turian, J. Shier, H. R. Khan, B. Raj, B. W. Schuller, C. J. Steinmetz, C. Malloy, G. Tzanetakis, G. Velarde, K. McNally, M. Henry, N. Pinto, C. Noufi, C. Clough, D. Herremans, E. Fonseca, J. Engel, J. Salamon, P. Esling, P. Manocha, S. Watanabe, Z. Jin, and Y. Bisk, “HEAR: Holistic Evaluation of Audio Representations,” in Proceedings of the NeurIPS 2021 Competitions and Demonstrations Track.1em plus 0.5em minus 0.4emPMLR, Jul. 2022, pp. 125–145, iSSN: 2640-3498.

[49]
K. Chen, X. Du, B. Zhu, Z. Ma, T. Berg-Kirkpatrick, and S. Dubnov, “Hts-at: A hierarchical token-semantic audio transformer for sound classification and detection,” in ICASSP 2022-2022 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP).1em plus 0.5em minus 0.4emIEEE, 2022, pp. 646–650.

[50]
J. F. Gemmeke, D. P. Ellis, D. Freedman, A. Jansen, W. Lawrence, R. C. Moore, M. Plakal, and M. Ritter, “Audio set: An ontology and human-labeled dataset for audio events,” in 2017 IEEE international conference on acoustics, speech and signal processing (ICASSP).1em plus 0.5em minus 0.4emIEEE, 2017, pp. 776–780.

[51]
S.-w. Yang, P.-H. Chi, Y.-S. Chuang, C.-I. Jeff Lai, K. Lakhotia, Y. Y. Lin, A. T. Liu, J. Shi, X. Chang, G.-T. Lin, T.-H. Huang, W.-C. Tseng, K.-t. Lee, D.-R. Liu, Z. Huang, S. Dong, S.-W. Li, S. Watanabe, A. Mohamed, and H.-y. Lee, “SUPERB: Speech Processing Universal PERformance Benchmark,” in Proc. Interspeech 2021, 2021, pp. 1194–1198.

[52]
Y.-Y. Yang, M. Hira, Z. Ni, A. Chourdia, A. Astafurov, C. Chen, C.-F. Yeh, C. Puhrsch, D. Pollack, D. Genzel, D. Greenberg, E. Z. Yang, J. Lian, J. Mahadeokar, J. Hwang, J. Chen, P. Goldsborough, P. Roy, S. Narenthiran, S. Watanabe, S. Chintala, V. Quenneville-Bélair, and Y. Shi, “Torchaudio: Building blocks for audio and speech processing,” arXiv preprint arXiv:2110.15018, 2021.

[53]
K. Koutini, J. Schlüter, H. Eghbal-zadeh, and G. Widmer, “Efficient Training of Audio Transformers with Patchout,” in Proc. Interspeech 2022, 2022, pp. 2753–2757.

[54]
K. Clark, U. Khandelwal, O. Levy, and C. D. Manning, “What Does BERT Look at? An Analysis of BERT's Attention,” in Proceedings of the 2019 ACL Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP.1em plus 0.5em minus 0.4emFlorence, Italy: Association for Computational Linguistics, Aug. 2019, pp. 276–286.

[55]
M. Raghu, T. Unterthiner, S. Kornblith, C. Zhang, and A. Dosovitskiy, “Do vision transformers see like convolutional neural networks?” Advances in Neural Information Processing Systems, vol. 34, pp. 12 116–12 128, 2021.

[56]
M. Raghu, J. Gilmer, J. Yosinski, and J. Sohl-Dickstein, “SVCCA: Singular Vector Canonical Correlation Analysis for Deep Learning Dynamics and Interpretability,” in Advances in Neural Information Processing Systems, vol. 30.1em plus 0.5em minus 0.4emCurran Associates, Inc., 2017.

[57]
A. Pasad, J.-C. Chou, and K. Livescu, “Layer-Wise Analysis of a Self-Supervised Speech Representation Model,” in 2021 IEEE Automatic Speech Recognition and Understanding Workshop (ASRU), Dec. 2021, pp. 914–921.

[58]
M. Tian, A. Srinivasamurthy, M. Sandler, and X. Serra, “A study of instrument-wise onset detection in beijing opera percussion ensembles,” in 2014 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), 2014, pp. 2159–2163.

[59]
H. Cao, D. G. Cooper, M. K. Keutmann, R. C. Gur, A. Nenkova, and R. Verma, “Crema-d: Crowd-sourced emotional multimodal actors dataset,” IEEE transactions on affective computing, vol. 5, no. 4, pp. 377–390, 2014.

[60]
K. J. Piczak, “ESC: Dataset for Environmental Sound Classification,” in Proceedings of the 23rd Annual ACM Conference on Multimedia.1em plus 0.5em minus 0.4emACM Press, 2015, pp. 1015–1018.

[61]
F.-R. Stöter, S. Chakrabarty, B. Edler, and E. A. Habets, “Classification vs. regression in supervised learning for single channel speaker count estimation,” in 2018 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP).1em plus 0.5em minus 0.4emIEEE, 2018, pp. 436–440.

[62]
F.-R. Stöter, S. Chakrabarty, E. Habets, and B. Edler, “Libricount, a dataset for speaker count estimation,” April 2018.

[63]
A. Anantapadmanabhan, A. Bellur, and H. A. Murthy, “Modal analysis and transcription of strokes of the mridangam using non-negative matrix factorization,” in 2013 IEEE International Conference on Acoustics, Speech and Signal Processing, 2013, pp. 181–185.

[64]
J. Engel, C. Resnick, A. Roberts, S. Dieleman, D. Eck, K. Simonyan, and M. Norouzi, “Neural audio synthesis of musical notes with wavenet autoencoders,” 2017.

[65]
P. Warden, “Speech commands: A dataset for limited-vocabulary speech recognition,” 2018.

[66]
E. Fonseca, X. Favory, J. Pons, F. Font, and X. Serra, “Fsd50k: an open dataset of human-labeled sound events,” IEEE/ACM Transactions on Audio, Speech, and Language Processing, vol. 30, pp. 829–852, 2021.

[67]
B. Kim, M. Ghei, B. Pardo, and Z. Duan, “Vocal imitation set: a dataset of vocally imitated sound events using the audioset ontology.” in DCASE, 2018, pp. 148–152.

§ APPENDIX

§ MULTI-WINDOW MULTI-HEAD ATTENTION

def WinAttention(Q, K, V, win_i):
n, d_k = Q.shape[-2:]
# partition inputs along patch dimension
# into non-overlapping windows
Q = Q.reshape(-1, win_i, d_k)
K = K.reshape(-1, win_i, d_k)
V = V.reshape(-1, win_i, d_k)
# compute self-attention
X = softmax(Q.(K.transpose()) / sqrt(d_k)).V
# reshape results
X = X.reshape(-1, n, d_k)
return X

Pseudocode for WinAttention

§ MORE ABOUT DOWNSTREAM TASKS

Overview of tasks for downstream evaluation. All these tasks are a part of the HEAR [48] benchmark.

Short Hand     Name         Description       Size (in Hours)      Metric

BO     Beijing Opera [58, 48]     Classifying percussion instruments     0.3     Accuracy

CD     Crema-D  [59]      Emotion Recognition      $\sim10$      Accuracy

ESC-50     ESC-50  [60]       Environmental Sound Classification     2.77     Accuracy

LC     LibriCount  [61, 62]     Speaker Count Identification, Simulated Cocktail Party     $\sim8$     Accuracy

Mri-S     Mridangam Stroke [63]     Stroke classification in pitched percussion instruments     1.57     Accuracy

Mri-T     Mridangam Tonic [63]     Tonic classification in pitched percussion instruments     1.57     Accuracy

NS-5h     NSynth Pitch 5h  [48, 64]     88-way Pitch Classification, reduced training subset     $\sim5.5$     Accuracy

SC-5h     Speech Commands 5h  [48, 65]     Keyword Spotting, reduced training subset     $\sim6.5$     Accuracy

F50K     FSD50K [66]      Multilabel, large scale Audio Tagging     $\sim100$     mAP

VL     VoxLingua107 Top10 [48, 67]     Spoken language identification     5     Accuracy

The following is our reasoning behind excluding the other tasks from the HEAR benchmark suite:

* Nsynth-Pitch 50hr and Speech Commands Full because we already use the smaller subsets.

* Gunshot Triangulation: Gunshot is an event in both AudioSet and FSD50k ontology, and is thus redundant.

* GTZAN Music Speech: FSD50k already has music and speech labels, and the model performance correlation study in the HEAR paper [48] shows high correlation with FSD50k.

* GTZAN Genre: highly correlated results with FSD50K and ESC-50 (surprisingly) as per [48]

* Vocal Imitations: high correlation with LibriCount [48].

* Bee Hive state Classification: large runtime costs, niche task.

* MAESTRO 5hr and DCASE 2016 Task 2: significant complexity (storage, runtime, timestep based evaluation).

§ EXPERIMENTAL DETAILS AND HYPERPARAMETERS

In this section, we provide additional experimental details. Apart from AudioSet, all other datasets are obtained directly from the HEAR [<https://hearbenchmark.com/hear-tasks.html>], where they are pre-processed to $16000$ Hz and distributed in a standard format.

Similar to [37], our effective learning rate ($lr_{\text{eff}}$) depends on the base learning rate ($lr_{\text{base}}$) and the batch size as follows: $lr_{\text{eff}}=lr_{\text{base}}*\frac{\text{batch size}}{256}$. In early experiments, we did not find strong augmentations at pre-training time to improve downstream performance, hence no augmentations are used. For more details, refer to  Table <ref>. As previously mentioned, hear-eval-kit[<https://github.com/hearbenchmark/hear-eval-kit>] was used for downstream experiments, and along with the details provided here should allow for consistent, reproducible downstream experimentation.

Pre-training (PT) and Downstream (FT) hyperparameters. *: For ViT-L and ViT-H based models, smallest batch size that didn't give OOM was used.

Configuration     AS-5k Pre-training     Downstream

Optimizer     AdamW     1cAdam

Optimizer momentum     $\beta_1=0.9$, $\beta_2=0.999$     1c$\beta_1=0.9$, $\beta_2=0.95$

Weight decay     1c0.05     N/A

Base learning rate     0.000015     0.0001

Learning rate schedule     linear-warmup + cosine decay     1cfixed

Minimum learning rate     0.0     1c0.0001

Dropout     0.     0.25

Warm-up epochs     10     N/A

Epochs     100     500

Early Stopping     N/A     20

Batch size     1024*     1024

Accelerators     8x TPU-v3 cores     1 Nvidia-A40

§ ADDITIONAL MODALITY TESTED: IMAGENET

Given the in-depth ablations and exploratory analysis, as well as resource constraints, we couldn't dive deeper into full scale testing of an additional modality given. However, as a proof of concept, we pre-trained MAE and MW-MAEs with ViT-Medium encoder on ImageNet for 100 epochs, followed by evaluating linear probe performance (training for 50 epochs).

Proof of concept ImageNet experiments. Pre-training was done only for 100 epochs. Validation accuracy denotes linear probe performance on the ImageNet validation set.

Model     Params     Validation Accuracy

MAE Medium Encoder     38 M     29.00.0

MW-MAE Medium Encoder     38 M     29.80.0

§ PARAMETER COUNT, AVERAGES AND OVERALL SCORES

Models, number of parameters, plain average scores and overall scores of models omitted from Table <ref> due to space constraints.

Model     # Params     Average     $s(m)$

HEAR-Naive     -     24.3±0.5    	5.0±0.7

PaSST-base     86 M     67.5±0.3     73.5±0.4

Wav2Vec2-base     94.4 M     48.7±0.1     43.1±0.2

Wav2Vec2-large     315.4 M     67.1±0.2     74.0±0.4

WavLM-base     94.4 M     58.1±0.1     60.5±0.2

WavLM-large     315.4 M     60.1±0.1     64.0±0.2

HuBERT-base     94.4 M     66.3±0.1     72.5±0.2

HuBERT-large     315.4 M     66.4±0.2     73.4±0.3

SSaST-base     89 M     65.9±0.1     71.7±0.2

BEATs-Iter3     90 M     75.0±0.2     85.7±0.3

3lMAE based

AudioMAE     86.0 M     60.1±0.2     62.9±0.3

MAE-B-4x16-4l     86.0 M     76.4±0.1     88.1±0.2

MAE-B-5x5-4l     86.0 M     75.6±0.1     86.8±0.2

MAE-L-4x16-8l     302.4 M     77.5±0.1     90.0±0.2

MW-MAE-B-4x16-4l     86.0 M     77.0±0.1     89.2±0.2

MW-MAE-B-5x5-4l     86.0 M     77.8±0.1     90.6±0.1

MW-MAE-L-4x16-8l     302.4 M     79.1±0.1     92.6±0.2

§ DETAILED ABLATION RESULTS

Results from Patch size ablation experiments. ViT-B encoder was used for all experiments. n denotes total number of patches, and h denotes the number of attention heads in each decoder transformer block.

1lModel     BO     CD     ESC-50     LC     Mri-S     Mri-T     NS-5h     SC-5h     F50K     VL     $s(m)$

3lPatch Size=(8$\times$16), n=125, h=4     5l

MAE     94.9±0.8     70.2±0.3     80.4±0.5     66.0±0.3     97.4±0.1     97.7±0.1     65.9±0.7     88.9±0.5     49.4±0.1     40.6±0.5     85.9±0.3

MW-MAE     95.9±0.5     72.3±0.2     81.2±0.3     68.4±0.3     97.3±0.1     97.8±0.1     67.4±0.8     90.0±0.3     50.8±0.1     41.9±0.5     88.0±0.2

3lPatch Size=(4$\times$16), n=250, h=8     5l

MAE     96.2±0.3     72.2±0.2     80.9±0.4     67.3±0.3     97.4±0.1     98.3±0.1     68.3±0.4     89.4±0.3     50.4±0.1     43.1±0.9     88.1±0.2

MW-MAE     96.0±0.5     73.1±0.3     81.2±0.4     68.8±0.2     97.4±0.1     97.9±0.1     69.3±0.6     90.9±0.2     51.2±0.2     44.2±0.9     89.2±0.2

3lPatch Size=(8$\times$8), n=250, h=8     5l

MAE     96.1±0.6     72.5±0.2     81.3±0.2     66.0±0.3     97.5±0.1     98.1±0.0     68.5±0.7     89.5±0.4     50.2±0.1     42.3±0.5     87.7±0.2

MW-MAE     96.3±0.4     73.0±0.1     82.6±0.3     69.3±0.3     97.5±0.1     98.1±0.1     70.3±0.8     90.5±0.1     51.4±0.1     42.3±0.5     89.4±0.1

3lPatch Size=(4$\times$8), n=500, h=12     5l

MAE     96.7±0.2     71.3±0.3     79.0±0.4     67.8±0.3     97.7±0.0     98.5±0.0     68.7±0.4     89.0±0.4     49.8±0.2     39.2±0.7     87.2±0.1

MW-MAE     95.6±0.7     74.1±0.2     81.9±0.3     70.1±0.3     97.6±0.1     98.2±0.1     72.0±0.7     91.2±0.3     51.6±0.1     44.0±0.8     90.3±0.2

3lPatch Size=(5$\times$5), n=640, h=16     5l

MAE     96.0±0.4     70.9±0.2     80.9±0.4     67.6±0.4     97.6±0.1     98.4±0.0     69.3±0.4     88.4±0.3     49.3±0.2     37.7±0.6     86.8±0.2

MW-MAE     96.6±0.4     73.8±0.4     82.0±0.3     70.1±0.4     97.5±0.1     98.3±0.1     72.9±0.5     91.7±0.2     51.3±0.1     44.2±0.6     90.6±0.1

Effect of encoder size on performance. Patch size of 4$\times$16 was used for all experiments.

1lModel     BO     CD     ESC-50     LC     Mri-S     Mri-T     NS-5h     SC-5h     F50K     VL     $s(m)$

3lEncoder=ViT-T     5l

MAE     95.6±0.5     63.2±0.2     70.1±0.5     64.6±0.3     97.1±0.1     97.4±0.1     66.4±0.7     74.3±0.8     41.6±0.1     26.4±0.6     77.6±0.3

MW-MAE     93.3±1.0     64.4±0.2     71.9±0.5     65.5±0.3     97.1±0.1     97.6±0.1     68.1±0.4     77.0±0.6     43.4±0.1     28.6±1.1     79.0±0.3

3lEncoder=ViT-M     5l

MAE     95.2±0.7     69.5±0.2     77.8±0.3     67.4±0.3     97.4±0.0     98.0±0.1     66.6±0.7     88.0±0.4     48.1±0.1     38.3±0.8     85.3±0.2

MW-MAE     95.9±0.3     71.8±0.3     80.3±0.4     69.7±0.1     97.2±0.1     97.8±0.1     68.1±0.5     88.8±0.6     49.6±0.1     39.8±0.8     87.5±0.2

3lEncoder=ViT-B     5l

MAE     96.2±0.3     72.2±0.2     80.9±0.4     67.3±0.3     97.4±0.1     98.3±0.1     68.3±0.4     89.4±0.3     50.4±0.1     43.1±0.9     88.1±0.2

MW-MAE     96.0±0.5     73.1±0.3     81.2±0.4     68.8±0.2     97.4±0.1     97.9±0.1     69.3±0.6     90.9±0.2     51.2±0.2     44.2±0.9     89.2±0.2

3lEncoder=ViT-L     5l

MAE     95.8±0.6     72.4±0.1     79.7±0.3     66.8±0.4     97.5±0.1     98.2±0.1     69.5±0.6     90.9±0.2     50.7±0.1     43.6±0.4     88.3±0.2

MW-MAE     95.7±0.5     75.5±0.2     82.5±0.5     70.1±0.3     97.4±0.0     98.1±0.1     70.7±0.6     93.2±0.1     53.3±0.1     51.9±0.8     92.3±0.2

3lEncoder=ViT-H     5l

MAE     96.8±0.2     71.1±0.2     78.3±0.4     67.1±0.2     97.5±0.0     98.5±0.0     67.6±0.6     89.6±0.1     49.5±0.2     40.0±0.7     86.9±0.1

MW-MAE     96.8±0.2     74.8±0.1     81.6±0.4     69.5±0.4     97.4±0.0     98.2±0.1     70.8±0.5     92.4±0.2     52.1±0.1     47.5±0.6     91.1±0.2

Effect of decoder depth on downstream performance. ViT-B encoder, patch size of 4$\times$16 were used for each experiment.

1lModel     BO     CD     ESC-50     LC     Mri-S     Mri-T     NS-5h     SC-5h     F50K     VL     $s(m)$

3ldepth=1     5l

MAE     96.4±0.2     69.8±0.3     78.9±0.3     67.4±0.3     97.4±0.1     97.9±0.1     66.4±0.8     88.5±0.2     49.4±0.2     39.0±1.1     86.1±0.2

MW-MAE     96.6±0.5     72.4±0.2     79.0±0.4     68.7±0.3     97.5±0.1     98.0±0.1     68.8±0.5     90.2±0.3     50.6±0.1     39.1±0.8     87.8±0.2

3ldepth=2     5l

MAE     96.8±0.3     71.3±0.3     78.8±0.2     68.8±0.2     97.4±0.1     98.2±0.0     67.2±0.6     90.0±0.2     49.6±0.2     39.4±0.7     87.3±0.1

MW-MAE     96.0±0.7     73.1±0.2     79.4±0.3     69.2±0.3     97.4±0.1     98.2±0.1     69.0±0.6     90.6±0.2     50.7±0.2     40.1±0.6     88.3±0.3

3ldepth=4     5l

MAE     96.2±0.3     72.2±0.2     80.9±0.4     67.3±0.3     97.4±0.1     98.3±0.1     68.3±0.4     89.4±0.3     50.4±0.1     43.1±0.9     88.1±0.2

MW-MAE     96.0±0.5     73.1±0.3     81.2±0.4     68.8±0.2     97.4±0.1     97.9±0.1     69.3±0.6     90.9±0.2     51.2±0.2     44.2±0.9     89.2±0.2

3ldepth=8     5l

MAE     96.3±0.3     71.7±0.3     81.6±0.4     67.4±0.3     97.4±0.0     98.1±0.1     67.8±0.7     89.9±0.3     50.8±0.2     43.4±0.6     88.2±0.1

MW-MAE     96.2±0.5     73.2±0.2     82.2±0.4     69.7±0.3     97.3±0.0     98.1±0.1     69.4±0.5     91.3±0.2     52.0±0.2     44.7±0.8     89.9±0.2

Amount of pre-training dataset used v/s downstream performance.

1lModel     BO     CD     ESC-50     LC     Mri-S     Mri-T     NS-5h     SC-5h     F50K     VL     $s(m)$

3l10% of AS-5k     5l

MAE     93.6±0.7     51.3±0.2     49.5±0.3     48.4±0.4     97.1±0.1     96.4±0.1     61.1±0.7     70.4±0.9     29.7±0.2     17.3±0.5     63.3±0.2

MW-MAE     94.1±0.3     63.9±0.3     67.1±0.3     60.5±0.2     97.3±0.1     97.6±0.0     64.4±0.5     82.0±0.4     40.9±0.2     30.1±1.1     77.2±0.3

3l25% of AS-5k     5l

MAE     96.2±0.6     57.5±0.3     64.9±0.4     56.9±0.3     97.4±0.1     97.5±0.1     65.0±0.6     79.3±0.4     39.2±0.1     24.2±0.7     73.6±0.2

MW-MAE     96.1±0.5     68.0±0.2     75.5±0.4     67.2±0.3     97.3±0.1     98.0±0.1     65.9±0.4     86.5±0.2     46.4±0.1     35.7±0.6     83.8±0.2

3l50% of AS-5k     5l

MAE     97.2±0.3     65.5±0.3     74.1±0.3     64.3±0.3     97.5±0.1     98.1±0.1     67.0±0.6     85.3±0.6     45.1±0.1     32.4±0.8     81.9±0.2

MW-MAE     95.9±0.5     70.9±0.2     79.1±0.3     69.1±0.4     97.4±0.1     98.1±0.1     68.4±0.7     88.5±0.2     49.1±0.1     39.5±0.5     87.0±0.2

3l75% of AS-5k     5l

MAE     95.3±0.5     70.2±0.2     79.0±0.3     67.4±0.2     97.4±0.1     98.1±0.1     67.4±0.6     88.8±0.3     49.2±0.1     39.5±0.7     86.2±0.2

MW-MAE     96.0±0.5     72.6±0.3     80.5±0.4     69.5±0.3     97.4±0.1     97.9±0.1     68.3±0.4     89.9±0.2     50.5±0.1     41.7±0.8     88.4±0.2

3l100% of AS-5k     5l

MAE     96.2±0.3     72.2±0.2     80.9±0.4     67.3±0.3     97.4±0.1     98.3±0.1     68.3±0.4     89.4±0.3     50.4±0.1     43.1±0.9     88.1±0.2

MW-MAE     96.0±0.5     73.1±0.3     81.2±0.4     68.8±0.2     97.4±0.1     97.9±0.1     69.3±0.6     90.9±0.2     51.2±0.2     44.2±0.9     89.2±0.2

§ LIMITATIONS

The direct limitations of our work are:

* Pre-training data scale: As opposed to text corpus used in NLP [13] as well as speech representations [10, 14], AudioSet is several order of magnitudes smaller. While MW-MAEs demonstrate good performance characteristics in low-data scenarios, analysis on larger scales of data is definitely warranted.

* Computational demands: transformer based models are computationally expensive to train, and despite their favourable generalization characteristics, MW-MAEs are no different. MW-MAEs and as well as previous works [31, 37] have showed the efficacy of MAEs when pretrained with AudioSet, however, training on longer duration audio data is still  a challenge.

* Runtime Overhead: While theoretically MW-MHA should be faster at runtime, the overhead of calling multiple attention heads individually results in slight slowdown as compared to optimized CUDA kernel implementations for MHA. A custom kernel for the operation should be able to improve this.


# WHAT IS THE BEST DATA AUGMENTATION FOR 3D BRAIN TUMOR SEGMENTATION?

###### Abstract

Training segmentation networks requires large annotated datasets, which in
medical imaging can be hard to obtain. Despite this fact, data augmentation
has in our opinion not been fully explored for brain tumor segmentation. In
this project we apply different types of data augmentation (flipping,
rotation, scaling, brightness adjustment, elastic deformation) when training a
standard 3D U-Net, and demonstrate that augmentation significantly improves
the network’s performance in many cases. Our conclusion is that brightness
augmentation and elastic deformation work best, and that combinations of
different augmentation techniques do not provide further improvement compared
to only using one augmentation technique. Our code is available at
https://github.com/mdciri/3D-augmentation-techniques.

Index Terms—  Data augmentation, 3D brain tumor segmentation, MRI, 3D U-Net,
deep learning, artificial intelligence

##  1 Introduction

Deep learning is increasingly being used in medical imaging, as it often
provides better results compared to traditional analysis methods [1]. Compared
to other computer vision tasks, it is in medical imaging more difficult to
acquire a large number of training images due to ethics and data protection
regulations (e.g. GDPR), and data augmentation is therefore even more
important to increase the number of images for training and testing.
Surprisingly, there are very few papers that investigate how important
different types of augmentation (e.g. rotations, random flipping, scaling,
elastic deformations) are for training convolutional neural networks (CNNs)
for image classification or image segmentation, and this is especially true
for 3D CNNs. In [2] it is reported that the number of papers about brain tumor
segmentation has been increasing exponentially in the last decade, and most of
them utilizing artificial intelligence. Moreover, just considering the past
Multimodal Brain tumor Segmentation (BraTS) Challenges from 2017 to 2020 the
number of participants were 34, 64, 93, and 55 respectively. Many of these
papers briefly mention that data augmentation was used, but do not mention
details like the range of the random rotations or elastic deformations, or if
the random rotations were drawn from a normal distribution or from a uniform
distribution.

Nalepa et al. [3] provide an overview of data augmentation for brain tumor
segmentation, showing that flipping, rotation, scaling and pixel-wise (e.g.
adding noise or changing the brightness) augmentation were most common, while
translations and elastic deformations were less common. They also provide
their own comparison of augmentation techniques, but it lacks details such how
large the elastic deformations are. Furthermore, scaling was not tested, no
statistical evaluation is done, and the used segmentation network seems to be
2D and not 3D. Shorten et al. [4] provide a more general survey on data
augmentation, but barely mention elastic deformations and brightness
augmentation. A number of papers have proposed to learn the best augmentation
[5, 6], but a common drawback is a much longer training time. Here we
therefore provide a comparison of different data augmentation techniques for
brain tumor segmentation, so that new researchers in this field know what kind
of augmentation to apply.

##  2 Data

The MR images used for this project are from the BraTS Challenge 2020 [2, 7,
8]. The BraTS 2020 training dataset contains MR volumes of shape $240\times
240\times 155$ from 369 patients, and for each patient four types of MR images
were collected: native (T1), post-contrast T1-weighted (T1Gd), T2-weighted
(T2), and T2 Fluid Attenuated Inversion Recovery (FLAIR). The BraTS 2020
validation dataset contains the same type of MR images from 125 patients,
without the ground truth annotations. The images were acquired from 19
different institutions with different clinical protocols. The training set was
segmented manually, by one to four raters, following the same annotation
protocol, and their annotations were approved by experienced neuro-
radiologists. Moreover, all data were co-registered to the same anatomical
template, interpolated to the same resolution (1 mm3) and skull-stripped.

##  3 Methods

In this project we use a standard 3D U-Net [9] architecture, one of the most
used nowadays, which is trained with 4 MR images (T1, T1Gd, T2, FLAIR) to
perform a 4 class segmentation: background, whole tumor (WT), tumor core (TC),
and enhancing tumor (ET). Although more advanced segmentation networks have
been proposed, a standard U-Net [10] won the BraTS 2020 challenge. Our network
is trained on sub-volumes of $128\times 128\times 128$ voxels. The U-Net has 4
encoder and decoder steps: each step is made by two 3D convolutions: for the
encoder a 3D convolution layer with stride 1 followed by another one with
stride 2; whereas for the decoder a 3D Transpose convolution with stride 2,
concatenated with the respective encoder step output, followed by a 3D
convolution with stride 1. Each convolution has kernel size $4\times 4\times
4$, He weights initialization [11] and same padding. The first 3D convolution
uses 64 filters which are doubled at each encoder step, viceversa for the
decoder. Data pre- and post-processing were done exactly as in [12]. Adam [13]
is used as the network’s optimizer with learning rate $\lambda=10^{-4}$, and
the loss chosen is the generalized Dice loss [14]. The augmentation techniques
used for this projects are:

  * •

Patch extraction: from each original volume a sub-volume of shape $128\times
128\times 128$ is extracted around its center. In this way each sub-volume
mostly contains brain tissue and not the surrounding background.

  * •

Flipping: random flipping of one of the three different axes with 1/3
probability.

  * •

Rotation: rotation applied to each axis with angles randomly chosen from a
uniform distribution with range between 0°and 15°, 30°, 60°, or 90°.

  * •

Scale: scaling applied to each axis by a factor randomly chosen from a uniform
distribution with range $\pm 10$% or $\pm 20$%.

  * •

Brightness: power-law $\gamma$ intensity transformation with its gain ($g$)
and $\gamma$ parameters chosen randomly between 0.8 - 1.2 from a uniform
distribution. The intensity ($I$) is randomly changed according to the
formula: $I_{new}=g\cdot I^{\gamma}$.

  * •

Elastic deformation: elastic deformation with square deformation grid with
displacements sampled from a normal distribution with standard deviation
$\sigma=2$, 5, 8, or 10 voxels [9], where the smoothing is done by a spline
filter with order 3 in each dimension.

Moreover, in order to report a robuster evaluation, 3-fold cross validation is
applied to each model, and these 3 models are ensembled by averaging their
softmax layer outputs.

Successively, the authors ranked all the augmentation approaches as in the
BraTS 2020 challenge and handled the ties as in [15] to determine which of of
these augmentation techniques and parameters yield the best performance on the
validation set. There will be shown two rankings: one for all the different
augmentation techniques and one including also their combinations.
Furthermore, the techniques with higher rank are also combined between each
other with a probability of 0.5 for each patch.

Each model is trained, with one or more augmentation techniques, over 200
epochs with early stopping after 25 epochs in case the validation loss does
not decrease. The segmentation networks were trained with Nvidia Tesla V100
and Nvidia Quadro RTX 8000 graphics cards. In the end, all the trained models
were evaluated on the 125 subjects of the BraTS 2020 challenge validation set.
The metrics (Dice score and Hausdorff distance 95 percentile) are calculated
by the CBICA Image Processing Portal (https://ipp.cbica.upenn.edu), while the
rank scores were calculated by us.

##  4 Results

First of all, we want to investigate whether data augmentation significantly
increases the Dice score. Hence, we applied a non-parametric permutation test
to the Dice scores from the 125 validation subjects. As the Dice scores are
from the same 125 validation subjects in all cases, we used a paired t-test to
test if the mean Dice difference is significantly different from zero. In non-
parametric statistics, a paired t-test can be performed with a sign flipping
test, where the sign of each pairwise difference is randomly flipped a large
number of times, to obtain the null distribution [16]. We used 100,000 sign
flips per test and the p-values are given in Table 1. Brightness augmentation
and elastic deformations with a $\sigma=2$ result in significantly higher Dice
scores for all 3 tumor classes, scaling with $\pm$ 20% significantly improves
the Dice scores for 2 classes, while flipping and 90°rotation only
significantly improve one tumor class.

Table 1: Non-parametric p-values for comparing different types of data augmentation, obtained through a sign flipping test using the 125 validation subjects. ET = enhancing tumor, WT = whole tumor, TC = tumor core. The p-values have been multiplied with 36 (36 one sided tests) as Bonferroni correction for multiple comparisons. The bold numbers are those that are smaller than 0.05. | p-value for Dice score
---|---
Comparison |  ET |  WT |  TC
Flipping + PE $>$ PE |  0.0082 |  1.0 |  0.19116
Brightness + PE $>$ PE |  0.00036 |  0.00036 |  0.00036
Scale $\pm 10$% + PE $>$ PE |  0.00036 |  1.0 |  1.0
Scale $\pm 20$% + PE $>$ PE |  1.0 |  0.00036 |  0.00396
Rotation 0°- 15°+ PE $>$ PE |  1.0 |  1.0 |  1.0
Rotation 0°- 30°+ PE $>$ PE |  1.0 |  1.0 |  1.0
Rotation 0°- 60°+ PE $>$ PE |  0.4676 |  0.2422 |  1.0
Rotation 0°- 90°+ PE $>$ PE |  0.00036 |  1.0 |  1.0
Elastic deformation 2 + PE $>$ PE |  0.00036 |  0.00036 |  0.00036
Elastic deformation 5 + PE $>$ PE |  0.00036 |  1.0 |  0.44208
Elastic deformation 8 + PE $>$ PE |  0.00036 |  0.00071 |  0.13464
Elastic deformation 10 + PE $>$ PE |  0.00036 |  1.0 |  1.0

On top of Figure 1 are illustrated the distributions of the Dice score and
Hausdorff distance 95 percentile on the three different brain tumor classes
over the 125 subjects in the validation set, for each tested augmentation
approach. Each box represents the interquartile range (IQR) between the 25 and
75 percentile, the line inside the median, the whiskers are set at 10 and 90
percentiles of the metrics. On the other hand, on the bottom of Figure 1 are
represented the rank scores for the different augmentation techniques used,
firstly without combining any technique, secondly including combinations of
them with their best parameters. Looking at Figure 1, it is possible to say
that brightness and elastic deformations with $\sigma=2$ are the two best
augmentation techniques for this dataset, having the best data distributions
and 1 and 2 position respectively. Moreover, the other most effective
techniques are scaling with $\pm 20$%, rotation with random angle chosen
between 0°- 90°, and flipping having 4, 7, and 8 rank position respectively.
Hence, four more trainings were done combining these techniques, and a second
ranking was performed including them shown in Rank 2 in Figure 1. Brightness
and elastic deformations with $\sigma=2$ kept their rank position, whereas the
combination of them and the one including also scaling with $\pm 20$% take the
4 and 5 position respectively.

Fig. 1: On top the Dice score and Hausdorff distance 95 percentile
distributions over the different augmentation techniques applied during the 3D
U-Net’s training. Each box represents the IQR between the 25 and 75
percentile, the orange line inside the median, whereas the whiskers are set at
10 and 90 percentiles of the metrics for the ET, WT, and TC class highlighted
in red, blue and green respectively. On the bottom the ranks scores calculated
without (blue) and with (orange) considering multiple augmentation.

Note that the CBICA systems reports 0% and 373mm as Dice score and Hausdorff
95 percentile when the network detects a class (usually the ET one) that is
not present, or viceversa. In the case of the combination of elastic
deformation with $\sigma=2$ and brightness, the ET whiskers go till 0% and
373mm.

##  5 Discussion

The aim of this paper is not to propose a state-of-the-art image segmentation
method with high performance as in [12, 10], but to investigate how different
augmentation techniques affect a CNN’s learning. Our results show that data
augmentation significantly improves the performance on the validation data in
many cases compared to only using patch extraction as baseline technique (see
Table 1). A possible explanation why data augmentation has not been fully
explored for brain tumor segmentation is that the BraTS training set is rather
large (369 subjects for the 2020 version), and several papers suggest that
data augmentation would not help much [17, 18, 19]. Since all subjects in
BraTS have been registered to a common space, augmentation is important to
show the network brains from different angles, while augmentation may not be
as important for a dataset where the brains have not been registered.

In Figure 1 reports also the rank scores for our tested approaches. It is
clear that augmentation improves the training performance because patch
augmentation, which is our baseline, has the lowest rank in both cases (13 and
17). Augmentation techniques have to be chosen smartly in order to create new
training images that still represent and/or look like the original ones. For
this particular dataset:

  * •

elastic deformation generates realistic tumors and achieves its best scores
with $\sigma=2$, which is also the technique with the best rank position.
Additionally, even using a value of $\sigma=5$ or 8, the scores are still
suitable, achieving the 6 and 3 rank position respectively, whereas
$\sigma=10$ it is too high and it deforms the brains excessively;

  * •

$\gamma$ correction (brightness) augmentation is the second best form of
augmentation here, which is probably explained by the fact that the BraTS 2020
dataset contains data from 19 different sites. Moreover, according to Nalepa
et al. [3], only Isensee et al. [10] used gamma correction in their comparison
of the BraTS 2018 participants;

  * •

random scaling and rotation create bigger and/or smaller tumors that are
distributed everywhere in the training volumes’ space. This is probably the
reason why greater scaling and rotation ranges result in better scores;

  * •

flipping is very easy to implement technique, but it does not increase
significantly the network’s performance.

Moreover, Figure 1 shows that the best combination is the one with elastic
deformation with $\sigma=2$ and brightness, followed by the one with scaling
$\pm 20$, but it is interesting to notice that combining different
augmentation techniques does not improve the network’s performance. Indeed,
the models trained with only elastic deformation or brightness augmentation
are higher ranked than the ones trained on multiple augmentation techniques
(see Rank 2). This may be due to the fact that each augmentation technique is
applied with a probability 0.5, so, if more augmentation techniques are
applied, fewer original images are shown to the network during the training.
Indeed, the percentage of original images used during the training is 50%,
25%, 12.5%, 6.25%, and 3.125% combining 1, 2, 3, 4 and 5 different techniques
respectively. Combining different types of augmentation can potentially be
done differently, to always guarantee that for example 20% of the shown
patches are always the original ones.

To summarize, the authors recommend everyone who is going to use this type of
dataset in the future to apply, at least, elastic deformation and brightness
adjustment as augmentation techniques, and in case they would like to use more
augmentation techniques, to be sure that the network considers a suitable
percentage of original images during its training.

##  6 Acknowledgments

This study was supported by LiU Cancer, VINNOVA Analytic Imaging Diagnostics
Arena (AIDA), and the ITEA3 / VINNOVA funded project Intelligence based
iMprovement of Personalized treatment And Clinical workflow supporT (IMPACT).
Anders Eklund has previously received hardware from Nvidia, otherwise the
authors have no conflicts of interest to declare.

## References

  * [1] Geert Litjens, Thijs Kooi, Babak Ehteshami Bejnordi, Arnaud Arindra Adiyoso Setio, Francesco Ciompi, Mohsen Ghafoorian, Jeroen Awm Van Der Laak, Bram Van Ginneken, and Clara I Sánchez,  “A survey on deep learning in medical image analysis,”  Medical image analysis, vol. 42, pp. 60–88, 2017.
  * [2] Bjoern H Menze, Andras Jakab, Stefan Bauer, Jayashree Kalpathy-Cramer, Keyvan Farahani, Justin Kirby, Yuliya Burren, Nicole Porz, Johannes Slotboom, Roland Wiest, et al.,  “The multimodal brain tumor image segmentation benchmark (BRATS),”  IEEE transactions on medical imaging, vol. 34, no. 10, pp. 1993–2024, 2014.
  * [3] Jakub Nalepa, Michal Marcinkiewicz, and Michal Kawulok,  “Data augmentation for brain-tumor segmentation: A review,”  Frontiers in Computational Neuroscience, vol. 13, 2019.
  * [4] Connor Shorten and Taghi M Khoshgoftaar,  “A survey on image data augmentation for deep learning,”  Journal of Big Data, vol. 6, no. 1, pp. 60, 2019.
  * [5] Ekin D Cubuk, Barret Zoph, Dandelion Mane, Vijay Vasudevan, and Quoc V Le,  “Autoaugment: Learning augmentation strategies from data,”  in Proceedings of the IEEE conference on computer vision and pattern recognition, 2019, pp. 113–123.
  * [6] Ju Xu, Mengzhang Li, and Zhanxing Zhu,  “Automatic data augmentation for 3D medical image segmentation,”  in International Conference on Medical Image Computing and Computer-Assisted Intervention. Springer, 2020, pp. 378–387.
  * [7] Spyridon Bakas, Hamed Akbari, Aristeidis Sotiras, Michel Bilello, Martin Rozycki, Justin S Kirby, John B Freymann, Keyvan Farahani, and Christos Davatzikos,  “Advancing the cancer genome atlas glioma MRI collections with expert segmentation labels and radiomic features,”  Scientific data, vol. 4, pp. 170117, 2017.
  * [8] Spyridon Bakas, Hamed Akbari, Aristeidis Sotiras, Michel Bilello, Martin Rozycki, Justin Kirby, John Freymann, Keyvan Farahani, and Christos Davatzikos,  “Segmentation labels and radiomic features for the pre-operative scans of the TCGA-LGG collection,”  The Cancer Imaging Archive, 2017.
  * [9] Olaf Ronneberger, Philipp Fischer, and Thomas Brox,  “U-net: Convolutional networks for biomedical image segmentation,”  in International Conference on Medical image computing and computer-assisted intervention. Springer, 2015, pp. 234–241.
  * [10] Fabian Isensee, Philipp Kickingereder, Wolfgang Wick, Martin Bendszus, and Klaus H Maier-Hein,  “No new-net,”  in International MICCAI Brainlesion Workshop. Springer, 2018, pp. 234–244.
  * [11] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun,  “Delving deep into rectifiers: Surpassing human-level performance on imagenet classification,”  in Proceedings of the IEEE international conference on computer vision, 2015, pp. 1026–1034.
  * [12] Marco Domenico Cirillo, David Abramian, and Anders Eklund,  “Vox2Vox: 3D-GAN for Brain Tumour Segmentation,”  in International MICCAI Brainlesion Workshop. Springer, 2020.
  * [13] Diederik P Kingma and Jimmy Ba,  “Adam: A method for stochastic optimization,”  arXiv preprint arXiv:1412.6980, 2014.
  * [14] Carole H Sudre, Wenqi Li, Tom Vercauteren, Sebastien Ourselin, and M Jorge Cardoso,  “Generalised Dice overlap as a deep learning loss function for highly unbalanced segmentations,”  in Deep learning in medical image analysis and multimodal learning for clinical decision support, pp. 240–248. Springer, 2017.
  * [15] Maurice G Kendall,  “The treatment of ties in ranking problems,”  Biometrika, pp. 239–251, 1945.
  * [16] Anderson M Winkler, Gerard R Ridgway, Matthew A Webster, Stephen M Smith, and Thomas E Nichols,  “Permutation inference for the general linear model,”  Neuroimage, vol. 92, pp. 381–397, 2014.
  * [17] Gregor Urban, M Bendszus, F Hamprecht, and J Kleesiek,  “Multi-modal brain tumor segmentation using deep convolutional neural networks,”  MICCAI BraTS (brain tumor segmentation) challenge. Proceedings, winning contribution, pp. 31–35, 2014.
  * [18] Mark Lyksborg, Oula Puonti, Mikael Agn, and Rasmus Larsen,  “An ensemble of 2D convolutional neural networks for tumor segmentation,”  in Scandinavian Conference on Image Analysis. Springer, 2015, pp. 201–211.
  * [19] Mohammad Havaei, Axel Davy, David Warde-Farley, Antoine Biard, Aaron Courville, Yoshua Bengio, Chris Pal, Pierre-Marc Jodoin, and Hugo Larochelle,  “Brain tumor segmentation with deep neural networks,”  Medical image analysis, vol. 35, pp. 18–31, 2017.