Masked Autoencoders with Multi-Window Local-Global Attention Are Better Audio Learners

Continued..

3. It could have been better to provide a quantitative run-time comparison...

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.

4. [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.

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.

Conclusion

We thank the reviewer for the valuable feedback and for engaging in a discussion with us! We hope we have addressed your concerns, and we would love to hear back!

References

[a] https://hearbenchmark.com/hear-leaderboard.html
[b] Schmid et al, 2023, “Dynamic Convolutional Neural Networks as Efficient Pre-trained Audio Models”

We thank the reviewer for their comments. We have addressed your comments one by one.

Multiscale approaches are not particularly new but this might be the first application of it to audio representation learning with masked AE.

Various approaches to local and global attention have been applied to audio representation learning before (as several of our references indicate), but we have proposed a novel method with local-global windows within the same multihead module. Our approach is not merely novelty in application, it’s a new approach to local-global attention modelling.

Weaknesses And Questions

1. Marginal gains with the proposed approach .... MW-MAE is better by a small margin.

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].

2. The local vs. global feature aspect of the model is not well-supported...global features at utterance level.

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.

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.

Continued..

We thank the reviewer for their comments!

The paper is well-structured and easy to follow.

We are glad that the reviewer found the paper easy to read! We gave it our best!

The proposed approach to enhancing multi-head attention is straightforward, yet it demonstrates a performance boost.

Thanks for the words of encouragement! It elates us that the reviewer also sees the strength in our simple yet effective approach! Now, allow us to address your concerns!

Concerns regarding applicability and contribution

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.

[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.

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.

Questions

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

Addressed above.

Typo: (6p, Sec. 4.3.) ...

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

Conclusion

Again, we thank the reviewer for their input and for engaging in a discourse! We hope we were able to address your concerns and look forward to answering any more questions!

References

[a] Liu et al. 2021, “Swin transformer: Hierarchical vision transformer using shifted windows.”
[b] Niizumi et al. 2022, "Masked Spectrogram Modeling using Masked Autoencoders for Learning General-purpose Audio Representation”.

Dear Reviewer fDLB,

We have responded to your review, incorporated your suggestions, as well as added additional experiments w.r.t. your concerns. Given that the discussion period will end soon, we would appreciate it if you could go through our response, and maybe we can discuss any further questions you might have!

Best,
Authors

The results of the auxiliary experiments they provided demonstrate that the proposed method is effective not only in self-supervised learning for audio data but also in ASR

Thank you so much! It feels great that you are also enthusiastic about the applicability of our approach!

Additionally, as mentioned in the "Author's Response to Reviewer 4eUM," the increase in training time due to the introduced architecture somewhat negates the efficacy of the proposed method.

The performance reduction in training time is solely due to implementation: Multi Head Attention has gone through several rounds of optimization of the kernel implementation in low-level code (CUDA, ROCm etc) and calls all attention heads in a single kernel call, which significantly optimizes issues such as memory transfer and kernel overheads. Our implementation as of now, as we pointed out to reviewer 4eUM, calls underlying attention heads individually which leads to overhead. We have complete confidence that a better implementation of native windowed attention can iron out these issues!

Conclusion

Overall, we would like to thank the reviewer for taking the time to review our paper, and providing feedback which has proved of immense benefit, and finally, for investing the time to go through our rebuttal and additional experiments.

We thank the reviewer for their comments. Allow us to address your concerns one by one

Weakness 1

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.

Weakness 2

In terms of the experimental aspect ..... further ablation experiments pertaining to this aspect would be valuable.

Since this weakness overlaps with the question posed by the reviewer, we have addressed both these issues here. We have already conducted extensive experiments w.r.t. network hierarchy (as per Sec 4.4). Thus, we focus on further ablations exploring the downstream performance impact of different window sizes and the number of attention heads, as shown in Table below. All evaluated methods have an identical encoder (ViT-Base, 4x16 patch size), number of patches $n_p = 250$, the same number of decoder layers ($l=4$) with the same hidden dimension ($d=384$). We will add detailed results to the appendix (we were having trouble squeezing in the full markdown table in a textbox).

As evident, irrespective of the window size used, MW-MAE outperforms MAE in overall performance by a considerable margin.

Weakness 3

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.

Continued..

Continued ...

Weakness 4

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.

Questions

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.

Conclusion

We thank the reviewer for engaging in a dialogue with us and for valuation suggestions. We hope we were able to address your concerns, and hope to hear back!

References

[a] Liu et al. 2021, “Swin transformer: Hierarchical vision transformer using shifted windows.”
[b] Dong et al. 2022, “CSWin Transformer: A General Vision Transformer Backbone With Cross-Shaped Windows”.
[c] Beltagy et al., “Longformer: The Long-Document Transformer.”
[d] Fan et al. 2021, “Multiscale vision transformers".
[e] Li et al. 2022, “Mvitv2: Improved multiscale vision transformers for classification and detection”.
[f] Zhu and Omar, 2023, “Multiscale audio spectrogram transformer for efficient audio classification”.
[g] Chen at al, 2022, “A hierarchical token-semantic audio transformer for sound classification and detection”.
[h] Saeed et al., "Contrastive Learning of General-Purpose Audio Representations", ICASSP 2021.
[i] Niizumi et al., "BYOL for Audio: Self-Supervised Learning for General-Purpose Audio Representation", IJCNN 2021.
[j] https://hearbenchmark.com/hear-leaderboard.html

Dear reviewer M3PT,

We have incorporated your comments and suggestions, as well as conducted additional experiments for window/attention head ablations. Given that the discussion period will end soon, we were hoping you could go through our response and tell us if it addresses your concerns, as well as answer any further questions you might have. We would highly appreciate it!

Best,
Authors.

Dear Reviewer M3PT,

We have incorporated your suggestions, responded to your questions, and as per your recommendation, conducted additional full pre-training experiments to elucidate the effect of window size and attention heads on downstream performance.

We thank you for your invaluable feedback! As the discussion period will end in a few hours, please take a look at our rebuttal and discourse with the other reviewers, and consider taking part in the discussion before the discussion period ends! We would greatly appreciate it.

Best,
Authors

We thank the reviewer for their comments. We are happy that the reviewer finds our analysis in Sec 5 appropriate! Let us address your concerns.

Weaknesses

1. Typos such as "we use fixed sinusuidal positional embeddings"

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

2. lack of parameter explanation in 3.1

$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

3. 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.

[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.

4. Introduction section and Related Works section overlapped by a lot.

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.

5. There should be another line "enc only" in table 2.

Thank you for the suggestion. We have now also included this extra experiment on MW-MAE Base (encoder only) setting. We didn’t believe originally that it was necessary, given that enc+dec doesn’t provide substantial improvements for pre-training. Following will be the updated Table 2:

While there is no difference in performance when fine-tuning, there is a considerable decline in performance in the encoder only setting for linear probe performance. The overall conclusion from the paper, regarding the addition of MW-MHA modules, doesn’t change.

Questions

1. I'm wondering whether the results of audioMAE and MW-MAE are comparable, what is the patch size and spectrogram size of audioMAE you are using? Have you tried re-training audioMAE with your current settings?

we have answered this in the section above.

2. 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?

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].

Conclusion

We thank the reviewer for engaging in the discussion and hope we were able to address your concerns. We look forward to answering any further questions the reviewer might have!

References

[a] Niizumi et al., 2022, “Masked spectrogram modeling using masked autoencoders for learning general-purpose audio representation”.
[b] Turian et al., 2022, “HEAR: Holistic Evaluation of Audio Representations”
[c] Yang et al., 2021, “SUPERB: Speech Processing Universal PERformance Benchmark”

| Is there any reference that proves [a] is a better baseline than AudioMAE under the same settings?

In paper [a], Table 5, authors conducted downstream experiments on several HEAR datasets with the same patch size as AudioMAE (16x16) and 2-second audio input: MSM-MAE-208.

Early on in our study before we committed to a reference implementation to build upon, we basically compared these numbers with the results we obtained using official AudioMAE pre-trained weights and found that AudioMAE downstream results were worse than MSM-MAE-208. For instance, AudioMAE gave us 68.2±0.2 and 37.9±0.1 for CD and FSD50K. For the same tasks, MSM-MAE numbers from [a] were 70.4 and 50.8, respectively. Considering the efforts that go into reproducing a paper, the decision to build upon the stronger baseline seemed justified to us. And thus, the performance gap cannot be chalked merely to difference in patch and spectrogram size.

| I'm wondering how much that MW-MAE outperforms AudioMAE under the same settings

We understand your concern and point of view. We have been running experiments comparing AudioMAE (essentially SWIN Transformer in the decoder) with our proposed approach under identical settings: smaller patch size than original AudioMAE, smaller decoder than original AudioMAE (4 layers instead of 16), 2-second audio crops vs 10. We'll get back to you once we have the results.

Dear reviewer a4K8,

As we mentioned in our previous comment, we have conducted additional experiment on AudioMAE (essentially SWIN blocks in the decoder) in identical setting as the proposed MW-MAE for a direct comparison. We conduct experiments for the base configuration: ViT-Base encoder, 4x16 patch size, decoder with 4 layers and 384 hidden dimensions.

Following are the results on the different tasks. We have added results of comparable MW-MAE method from the main text for easy viewing.

Compared to the official AudioMAE pretrained model [1], retraining on 2-sec inputs and 4x16 patch size improves AudioMAE performance drastically. However, MW-MAEs ($s(m)=89.2±0.2)$ still demonstrate a considerable absolute improvement of 3.1% in overall score over the AudioMAE model ($s(m)=86.1±0.3$).

Given that we have now compared the two methods in identical settings, we can conclude that this performance gap between AudioMAE and MW-MAE is due to better local-global attention modelling capabilities of the proposed approach.

We thank the reviewer for the suggestion, as this further demonstrates the strength of our approach. We hope this addresses your concerns and we eagerly wait to hear back from you.

*EDIT: Added MAE results as well.

References

[1] Huang et al., 2022, "Masked Autoencoders that Listen"

Thank you for these additional experiments and figures, now I understand that MWMHA has some potential on other tasks. However, the main topic of this paper is the application of MWMHA in MAE on the audio representation task, and the improvements from this specific application of MWMHA is not as much as those from other tasks. I'd like to clarify my statement, I think the effectiveness of this specific application of MWMHA in MAE is restricted.

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.

On the other hand, I think this specific application of MWMHA on MAE could be much restricted

We would like to bring to the reviewer's attention, in case they haven't seen our response to Reviewer fDLB, our preliminary results investigating MW-MHA for Supervised ASR. For more details, you can view our response to Reviewer fDLB, but the gist of it is that while true that it is just preliminary experiments, we observed considerable improvement in WER performance on LibriSpeech, which highlights potential applicability of MW-MHA beyond MAEs.

Conclusion

We thank the reviewer for their suggestions, and most of all, for continuously engaging in discussions with us! We had a lot of work cut out for us (justifiably) for the rebuttal based on reviewer feedback, so hearing back from reviewers really bolsters our confidence! Thanks for taking the time to review our paper.

We have added the MAE results in the table in our previous comments. You can see that MAE with Swin blocks performs worse than our MAE baseline on some tasks (4/10) and at par on others. MAE with Swin blocks does have considerably worse overall performance than our MAE baseline, although not as much of a disparity compared to official pre-trained AudioMAE model.