Paper: Masked Autoencoders with Multi-Window Local-Global Attention Are Better Audio Learners (ICLR-2024)

Reviewer comment (Novelty):
Marginal gains with the proposed approach: Overall MW-MAE is slightly better than the original MAE. At individual task levels, sometimes MAE is better and sometimes MW-MAE is better by a small margin.

Author response:
Within the reported 95% confidence intervals, MW-MAEs consistently outperform MAEs in a host of conditions and scenarios. 

1. The overall performance of MW-MAEs is considerably better than MAEs in a wide variety of conditions such as:  
    a. low data scenarios.  
    b. For larger encoder complexities: 92.6±0.2 vs 90.0±0.2 for ViT-L is a considerably large improvement, especially considering it’s an aggregation over 10 tasks.
2. At the task level,   
    a. MW-MAEs consistently outperform their MAE counterparts on majority of the tasks. For example, ViT-L based MW-MAE from table 2 achieves an absolute performance improvement of 2.3%, 2%, 1.2%, 2.2%, 1.8%, 1.7%, 5% on CD, ESC-50, LC, NS-5h, SC-5h, FSD50k and VL tasks, respectively. That is a considerable improvement across the board on 7/10 tasks.   
    b. The only tasks it does not improve on, are BO, Mri-S and Mri-T, where MW-MAEs are within the reported 95% confidence intervals, not worse. This observation is mostly consistent across all our experiments. We suspect this is due to performance saturation on these tasks, as also suggested by the HEAR leaderboard [a].

Reviewer comment (Evaluation):
The local vs. global feature aspect of the model is not well-supported. From the multi-window approach, we can see that there are multiple scales but since there is no frame-level task in the evaluation setup, we cannot judge whether the local features/embeddings are also useful in downstream tasks. All the tasks evaluated here are based on global features at utterance level.

Author response:
Before continuing, we would like to add that just because a task is evaluated at utterance level, it doesn’t mean that local information was not leveraged at all. For example, FSD50K is a multilabel audio tagging task, with multiple audio events occurring throughout the audio clip. While prediction is done at the utterance level, rich local feature information is crucial for detecting all the audio events present in the audio clip. This local information is consolidated over the layers and assimilated into a feature representation. Also, we have shown extensive, quantitative exploratory analysis in Section 5.1 that verifies that the proposed MW-MAE captures a better mix of local and global attention.

Having said that, we appreciate the reviewer's suggestion. We have conducted additional experiments comparing MAE and MW-MAE base configurations on MAESTRO-5h, a frame-level music transcription task, which is a part of the HEAR benchmark. We chose MAESTRO-5h over DCASE 2016 Task 2 because our test suite includes ESC-50 and FSD50K, which already span office sound environments. It is evaluated in a 5-fold cross-validation setting, with the Onset FMS evaluation metric. We conducted the experiments 5 times and reported 95% confidence intervals over the cross-validation scores of individual runs.

| Model | Onset FMS |
| --- | --- |
|  |  |
| MAE-Base | 46.47±0.05 |
| MW-MAE Base | 47.31±0.07 |

MW-MAE does improve performance on MAESTRO-5h. While this might not look like much, at the time of writing this comment, **MW-MAE outperforms every method on the HEAR leaderboard [a] for MAESTRO-5h, including large multi and cross-modal ensemble methods**, so this is excellent performance for a single method. For reference, a pre-trained Wav2Vec2-Base only achieves an Onset FMS of 3.29 [a].

We didn’t compare all methods on MAESTRO to begin with, because it poses serious computation (and space) complexity: it takes longer to evaluate downstream performance on just MAESTRO than it takes to evaluate all the other downstream tasks combined. For the number of experiments we did in this paper, it was highly unfeasible, more so given that our benchmark suite includes tasks like FSD50K, where local information is known to be of importance [b]. 

Also, in response to reviewer fDLB, we have also posted proof of concept results for Supervised ASR, which might be of relevance in regard to this question.

Reviewer comment (Experiments):
It could have been better to provide a quantitative run-time comparison of MAE and MW-MAE, rather than mentioning this shortcoming in the Appendix.

Author response:
For our base configuration, the _training time_ throughput of MW-MAE (4109 samples/ sec) is 9.14 % slower than an MAE (4536 samples/sec) on the same hardware. This is primarily due to kernel call overhead of the MW-MHA module since we have to make individual calls for each attention head in an MW-MAE, whereas optimized MHA implementations do so in one kernel call.  A better native implementation can alleviate this difference, and we are working on that.

However, it is worth noting that since the decoder is discarded after pretraining, _there is no difference in throughput when running downstream experiments/inference/finetuning_ since the encoders are identical.

Reviewer comment (Presentation):
[Minor comment] It could have been better if the 10 tasks were mentioned in the main text before Table 1 rather than leaving the list to the Appendix.

Author response:
That is an oversight on our behalf. We will mention the 10 tasks in the main text, however, the table will remain in the Appendix due to space limitations.

Reviewer comment (Evaluation):
Evaluations are on utterance level tasks, so how can we make sure that the local information is also well-preserved at the frame level?

Author response:
Before continuing, we would like to add that just because a task is evaluated at utterance level, it doesn’t mean that local information was not leveraged at all. For example, FSD50K is a multilabel audio tagging task, with multiple audio events occurring throughout the audio clip. While prediction is done at the utterance level, rich local feature information is crucial for detecting all the audio events present in the audio clip. This local information is consolidated over the layers and assimilated into a feature representation. Also, we have shown extensive, quantitative exploratory analysis in Section 5.1 that verifies that the proposed MW-MAE captures a better mix of local and global attention.

Having said that, we appreciate the reviewer's suggestion. We have conducted additional experiments comparing MAE and MW-MAE base configurations on MAESTRO-5h, a frame-level music transcription task, which is a part of the HEAR benchmark. We chose MAESTRO-5h over DCASE 2016 Task 2 because our test suite includes ESC-50 and FSD50K, which already span office sound environments. It is evaluated in a 5-fold cross-validation setting, with the Onset FMS evaluation metric. We conducted the experiments 5 times and reported 95% confidence intervals over the cross-validation scores of individual runs.

| Model | Onset FMS |
| --- | --- |
|  |  |
| MAE-Base | 46.47±0.05 |
| MW-MAE Base | 47.31±0.07 |

MW-MAE does improve performance on MAESTRO-5h. While this might not look like much, at the time of writing this comment, **MW-MAE outperforms every method on the HEAR leaderboard [a] for MAESTRO-5h, including large multi and cross-modal ensemble methods**, so this is excellent performance for a single method. For reference, a pre-trained Wav2Vec2-Base only achieves an Onset FMS of 3.29 [a].

We didn’t compare all methods on MAESTRO to begin with, because it poses serious computation (and space) complexity: it takes longer to evaluate downstream performance on just MAESTRO than it takes to evaluate all the other downstream tasks combined. For the number of experiments we did in this paper, it was highly unfeasible, more so given that our benchmark suite includes tasks like FSD50K, where local information is known to be of importance [b]. 

Also, in response to reviewer fDLB, we have also posted proof of concept results for Supervised ASR, which might be of relevance in regard to this question.

Reviewer comment (Novelty):
Limited applicability and contribution. The proposed method involves applying self-attention over non-overlapping windows of varying sizes for each head in the multi-head attention. However, experimentally, the method is only applied within the confines of a self-supervised approach in the audio domain, using the masked autoencoder structure, and specifically for the decoder. Compared to the previous work, Audio-MAE, which addressed challenges in initially applying MAE to the audio domain, this study merely demonstrates the effects of improving a specific part of the structure.

Author response:
> The proposed method involves applying self-attention over .... using the masked autoencoder structure, and specifically for the decoder.

### [Proof of concept 1: Auxiliary Modality] MW-MAEs on ImageNet

We would like to bring to your attention our initial experimental results using Masked Autoencoders for computer vision on the ImageNet dataset (as presented in Appendix D). We pretrained MAE and MW-MAEs with ViT-Medium encoder and a 4-layer decoder with 8 attention heads and 512 feature dimensions on ImageNet for 100 epochs, followed by evaluating linear probe performance (training for 50 epochs). MW-MAE outperforms MAE by almost 1% in validation accuracy. 

| Model | Params | Validation Acc |
| --- | --- | --- |
| MAE- Medium Encoder | 38 M | 29.0±0.0 |
| MW-MAE Medium Encoder | 38 M | 29.8±0.0 |

### [Proof of concept 2: Auxiliary Task] MW-MHA for Supervised ASR

We have conducted some preliminary experiments to investigate the applicability of the proposed MW-MHA module beyond MAEs. We evaluated Supervised Transformer based ASR models, trained using CTC objective function. We trained on LibriSpeech 100h subset and reported Word Error Rates (WER) on the test sets, using validated recipes from the SpeechBrain toolkit. We will add these results to the appendix.

- The network architecture used consists of a 2-layer Convolutional frontend, followed by a 12-layer encoder and 4-layer decoder, each with 4 attention heads. We trained on mel-filterbanks with 80 mel bins.
- For the MW-MHA Transformer, we only replaced the MHA blocks in the encoder with MW-MHA blocks, and used attention window sizes of 1/8, 1/4, 1/2 and global attention corresponding to the number of timeframes obtained from the largest audio segment in the corpus.
- MW-MHA transformer outperforms baseline transformer by a considerable margin on both test sets.

| Model | Params | Test Clean WER | Test Other WER |
| --- | --- | --- | --- |
| Transformer | 8.7 M | 7.56 | 17.90 |
| MW-MHA Transformer | 8.7 M | 6.73 | 17.20 |

While the scope of our investigation in this paper is confined to self-supervised learning, audio and masked autoencoders, these two auxiliary experiments clearly indicate that the proposed approach has the potential for wide applicability!

---

> Compared to the previous work, Audio-MAE, which addressed challenges in initially applying MAE to the audio domain, this study merely demonstrates the effects of improving a specific part of the structure.

Audio-MAE addressed how global attention is suboptimal for spectrograms if the time-frequency information of interest is predominantly local, and to this end, they proposed to incorporate SWIN [a] based attention in some transformer blocks of _only the decoder_ of their MAE (similar to us!). We are addressing the same challenges as AudioMAE did!

Furthermore, compared to AudioMAE:

1. We perform more exhaustive evaluation of patch sizes, which is a key hyper-parameter that differs between MAEs for Vision v/s audio, as demonstrated by [b]. [b] precedes the AudioMAE paper and, chronologically, were the first to address the nuances of applying masked autoencoders to audio.
2. We evaluate on a larger, more varied selection of tasks, addressing a wider variety of challenging applications in the audio domain.
3.  We also conducted direct in-depth exploratory analysis of the local-global effects in the audio domain, measuring quantifiable variances in attention distances and entropies, which was missing from AudioMAE.

Thus, we believe our contributions are just as significant, as we go above and beyond in several facets of investigating the challenges of applying MAE to the audio domain.

Reviewer comment (Novelty):
I have concerns regarding the limited applicability and contribution of this research, as mentioned in the above weaknesses.

Author response:
> I have concerns regarding the limited applicability and contribution of this research, as mentioned in the above weaknesses.

Addressed above.

Reviewer comment (Writing):
Typo: (6p, Sec. 4.3.) ... with the largest "MW-MAE-L-16x4-8l" -> ...... with the largest "MW-MAE-L-4x16-8l"

Author response:
> Typo: (6p, Sec. 4.3.) ...

Thank you for bringing this to our attention! We will fix this.

Reviewer comment (Writing):
Typos such as "we use fixed sinusuidal positional embeddings"

Author response:
Thanks for bringing this to our attention! We’ll fix this and comb through the content for typos again.

Reviewer comment (Writing):
lack of parameter explanations in 3.1

Author response:
$Q, K, V \in \mathbb{R}^{n \times {d_m}}$ represent query, key and value input matrices, and $W_i^Q, W_i^K, W_i^V \in \mathbb{R}^{d_m \times d_k}$ are their corresponding learnable projection matrices. $d_m$ is the model's feature dimension, and $d_k = \frac{d_m}{h}$.

We’ll add this to the main text in Sec 3.1

Reviewer comment (Novelty):
Another work [1] proposed very similar ideas that introducing local attentions in the decoder can help improve the MAE learner. They tried local attentions only and hybrid attentions in which they applied global attentions on top layers. Although This paper compared MW-MAE with audioMAE and MW-MAE outperformed audioMAE by a lot, it is still unclear to me why MW-MAE is much better with such a similar design. I attached my questions in below.

Author response:
[a], which precedes AudioMAE, presented extensive experiments for training masked autoencoders on Audioset. Early on, we observed that their approach performed better than AudioMAE, and we thus build upon [a] to build a stronger baseline than AudioMAE. Key differences include:

1. For the base configuration, AudioMAE has a patch size of 16x16. Our MAE baseline has a patch size of 16x4, which gives better temporal resolution.
2. AudioMAE is trained on full 10-sec audio clips, whereas our models are trained on random 2-second crops, based on the published results of [a] where 2-sec crops offer a great performance/training time tradeoff.

We have already discussed these design decisions in section 4.2.

> Have you tried re-training audioMAE with your current settings?
> 

This stronger MAE baseline is used for fair, direct comparison with the proposed MW-MAE approach. Since we already had a stronger baseline, for consistency and being truthful to the source, we decided to conduct downstream experiments for AudioMAE using the official pre-trained model and code.

Reviewer comment (Presentation):
Introduction section and Related Works section overlaped by a lot.

Author response:
Introduction lays the ground work of the relevant literature & the Related works section expands on it with commentary on specific approaches. We will clear up repetitive references & explanations.

Reviewer comment (Evaluation):
According to paper [1], their number of audioMAE(local) on ESC-50 was 94.1 and your number of audioMAE on ESC-50 is 60.6, what is the difference?

Author response:
ESC-50 number reported in AudioMAE is a result of supervised fine-tuning of the entire encoder. This is different from our setting, where we follow the HEAR [b] protocol and conduct downstream experiments by training an MLP on top of fixed feature vectors, and do not fine-tune the entire encoder. This is inline with recent approaches for evaluating self-supervised audio representations [b, c].

Reviewer comment (Novelty):
The conceptualization of this article is relatively straightforward. Building upon the previous global-local attention mechanism, previous studies did not incorporate a hierarchical design for the local window perspective. Hence, the author implemented a mechanism with varying window sizes for each attention head within the multi-head attention mechanism. This initial approach is commendable. I am intrigued to know if the author is the originator of this multi-window local-global attention concept for the first time. No pertinent references were found at the end of the introduction, maybe leaving open the possibility of existing similar ideas in the fields of computer vision or natural language processing.

Author response:
> Hence, the author implemented a mechanism with varying window sizes for each attention head within the multi-head attention mechanism. This initial approach is commendable.
> 

We are excited that the reviewer finds our approach commendable! Thanks for the words of encouragement.

> I am intrigued to know if the author is the originator of this multi-window local-global attention concept for the first time
> 

We conducted an extensive literature review, and to the best of our knowledge, we are the first to propose a multi-windowed multihead attention module as presented in the paper.

> No pertinent references were found .... in the fields of computer vision or natural language processing.
> 

We have included key pertinent references for various related approaches in the introduction (such as sliding window attention-based methods [a-c] as well as multiscale attention-based methods [d-g]). We also elucidate how the proposed approach differs from existing approaches in Section 3.1. These references span domains beyond audio representation learning.

Reviewer comment (Experiments):
Corresponding to the points highlighted by the author in Appendix G, the experiments conducted solely focused on the (AS-5k) dataset. However, as a self-supervised model, it is expected that the author would train on larger datasets, given the significance of datasets such as librilight, which have already amassed 60k hours of data. Furthermore, exploring the model's performance on larger datasets would provide insights into its upper limits.

Author response:
> Corresponding to the points highlighted .... exploring the model's performance on larger datasets would provide insights into its upper limits
> 

While we empathize with the reviewer's point of view, there are several reasons to stick to AudioSet. Most larger audio datasets, such as those suggested by the reviewer (LibriLight) and others (e.g. VoxPopuli), comprise only speech data, which does not align with the scope of this paper (general-purpose audio representation learning). For learning audio representations that go beyond speech, AudioSet has been a defacto standard in the community, as it consists of in-the-wild audio spanning an ontology of 527 sounds such as human speech, vehicles, everyday objects, and environmental and animal sounds, to name a few.  AudioSet has been used by several key approaches for unsupervised representation learning, such as COLA [h], BYOL-A [i], SSaST and AudioMAE. For consistent comparison in line with existing literature as well as alignment with the scope of our work, picking AudioSet was the ideal choice.
Our motivation for evaluating MAEs v/s MW-MAEs when using different amounts of pre-training data (Section 4.4) was to simulate how performance treads with dataset scale, where MW-MAEs outshine MAEs consistently.

Reviewer comment (Experiments):
A meticulous examination of the author's ablation experiments on the parameters of MAE in the appendix reveals that, under different parameter configurations and tasks, MW-MAE does not exhibit significant improvements compared to MAE. For instance, [BO, Patch Size=(4×8), n=500, h=12, MAE 96.7±0.2, MW-MAE 95.6±0.7], [Mris-s all], [Mri-T all]. Similar observations can be made in the main experiments presented in Table 1. In certain tasks, a comprehensive enhancement is not evident. Consequently, I eagerly anticipate the author's elucidation regarding this aspect.

Author response:
> A meticulous examination of the author's ablation ...... author's elucidation regarding this aspect.
> 

Even in the worst cases, MW-MAEs perform on par (based on the presented 95% confidence intervals) with MAEs, including the cases highlighted by the reviewer. This phenomenon is observed only for BO, Mri-S and Mri-T tasks, as also discovered by the reviewer. Based on our observations for larger model configurations, we believe performance on these datasets to be saturated around this upper range (for e.g. no model breaks 98% mark on Mri-S, not even the very large ViT-L and ViT-H configurations). This is also supported by observations from the HEAR leaderboard [j], where even multi-modal ensembles of several large models do not go beyond this performance. 

Throughout the paper, we have shown exhaustive evidence that shows that MW-MAEs consistently outperform MAEs by a considerable margin on other, more difficult tasks (such as FSD50k: 53.5±0.1 vs 51.8±0.1 for ViT-L encoder), as well as in overall performance (92.6±0.2 vs 90.0±0.2 for ViT-L). Within 95% confidence intervals, MW-MAEs outperform MAEs consistently on 7/10 tasks while being on par on the other 3 tasks. For example, ViT-L-based MW-MAE from table 2 achieves an absolute performance improvement of 2.3%, 2%, 1.2%, 2.2%, 1.8%, 1.7%, 5% on CD, ESC-50, LC, NS-5h, SC-5h, FSD50k and VL tasks, respectively. That is a considerable improvement across the board on 7/10 tasks.

At the same time, it's worth keeping in mind that MW-MAEs demonstrate better performance characteristics for different encoder/decoder complexities as well as different amounts of pre-training data used as shown in Sec 4.4.

Reviewer comment (Experiments):
The overall engineering effort invested in this paper appears to be relatively limited, and there are certain limitations in the analysis of the multi-window multi-head attention. Could more diverse designs and comparisons be explored, specifically regarding local attention windows? For instance, is it possible to make these window sizes trainable? Alternatively, more engaging designs and comparisons could be conducted by examining the relationship between the number of attention heads, network hierarchies, and window sizes.

Author response:
> For instance, is it possible to make these window sizes trainable?
> 

While plausible, learnable window sizes for MHA are a non-trivial challenge and at this stage are beyond the scope of this paper.

> Alternatively, more engaging designs and comparisons could be conducted by examining the relationship between the number of attention heads, network hierarchies, and window sizes.
> 

We have already presented ablations for different encoder and decoder complexities in Sec 4.4 in the main text of the paper. While more configurations can certainly be explored, encoder complexities, for the sake of comparison with previous works, have been limited to known ViT encoders. 

As for more comparisons for relationship between number of attention heads and window sizes, we have just presented them in the previous comment.