Variational Inference for Self-Supervised Speech Models Fine-tuning on Downstream Tasks

Dear Reviewer,

Thank you for your time and for providing a detailed review. We appreciate you recognize the strengths of our work, including its interesting insights, the novel use of variational methods for impact distribution from layer outputs, and the extensive experiments and ablation studies on three speech-related tasks.

Regarding the weaknesses you pointed out and your questions:

1. Is the proposed layer aggregation approach applicable to SSL model? Feels like they can work for models trained from scratch?

   • We aimed to propose methods for pre-trained speech self-supervised models, as questions about layer importance for different downstream tasks have been raised (see Section 6.2). However, in general, the method can be easily adapted to any transformer-based model, whether pre-trained or trained from scratch, for different tasks. Let us note, that training model from scratch is not a common as pre-trained SSL-models has superior performance.

2. Section 2.2: How does it enforce V to learn speaker-discriminative information only, while K has phonetic information?

   • This is an assumption authors do in [1], please refer to Section 3.2 and 4.2.1 in their paper for more details.

   [1] Junyi Peng, Oldrich Plchot, Themos Stafylakis, Ladislav Mosner, Luk´ as Burget, and Jan Cernocky. An attention-based backend allowing efficient fine-tuning of transformer models for speaker verification. In IEEE Spoken Language Technology Workshop, SLT 2022, Doha, Qatar, January 9-12, 2023, pp. 555–562. IEEE, 2022. doi: 10.1109/SLT54892.2023.10022775.

3. The authors mentioned that MHFA requires additional requires searching additional hyper-parameters. However, in section 5.2, the regularization weight also needs to be tuned. This doesn’t seem to be a fair comparison to prior studies.

   • While both VARAN and MHFA require additional hyper-parameters to be searched, which can be seen as a limitation, VARAN outperforms MHFA on SV and the extended version of MHFA for the SER task (see Tables 2 and 4). To ensure a fair comparison, we also searched for the optimal hyper-parameters for the MHFA layer-aggregation method. Please refer to Appendix D (Tables 11 and 13, where llrd factor, lr transformer, lr mhfa, and weight finetuning reg are all hyper-parameters are MHFA-specific).

4. “Choosing the right layer or combination of layers is challenging, as the training data does not have the optimal layer distribution.” Please explain more on this. Does it mean the empirically learned weights from training data (is it from pretraining or finetuning?) are not exactly the same as their real distribution?

   • We get our weights through approximating posterior distribution $p(z|x)$ through approximation $q_\theta(z|x)$ learned via variational inference, like in [1]. This weights learned on the fine-tuning stage. Procedure do not require special pre-training. We like to think about the method in a similar to VAE way. VAE aim to train a good representation $z$ from data $x$ ($z$ is not given in the training data), so that we can reconstruct original data from this representation. In our work we aim to “reconstruct” target variable y from learned layer distribution $q(z|x)$ and $x$ itself.

   [1] beta-VAE: Learning Basic Visual Concepts with a Constrained Variational Framework. Irina Higgins and Loic Matthey and Arka Pal and Christopher Burgess and Xavier Glorot and Matthew Botvinick and Shakir Mohamed and Alexander Lerchner. International Conference on Learning Representations 2017.

I would appreciate all the explanations from the authors! They do address some of my questions. However, there are still a couple questions remaining.

1. From the discussions of question #3, it seems that the proposed approach still needs hyper-params tuning, so this is not an actual benefit compared to prior work. I would suggest to remove this argument from the manuscript to make it a fair comparison to prior work.

2. There is no experimental results to address question #7 (although it is working in progress), so I would still challenge on the improvements. This is my biggest concern, and I'm happy to adjust my ratings if this question is addressed.

Dear Reviewer "WFfX",

We would like to remind you that we have updated our paper since your last review, and we believe the experiments we have provided address your concerns. In light of these revisions, we kindly ask you to reconsider your score. If you have any further questions, we would be happy to address them.

Dear Reviewer,

Thank you for your time and for providing a detailed review. We appreciate your mention of the paper's clarity and conciseness and your recognition that VARAN adopts variational inference to semi-automatically extract layer-wise information, providing novel insights.

We noticed that you mentioned the PR task; however, we would like to clarify that there is no such task in our experiments section. For more details, please refer to Section 4.

Regarding the weaknesses you pointed out and your questions:

1. How is discrete variational inference (both prior and posterior) modeled? Is Gumbel-Softmax used?
   • The prior distribution is modeled using the distribution's parameters (e.g., $p$ for the geometric distribution); note that this distribution is not learnable. The posterior distribution is modeled via the MHSA block (Figure 1). The input to the MHSA block consists of hidden states from the Transformer Encoder block. Formally, this is a list with a size equal to the number of layers, where each element is a tensor with dimensions [batch size, sequence length, hidden dimension]. We then apply mean pooling along the sequence length and concatenate the results into a single tensor, resulting in dimensions [batch size, number of layers, hidden dimension]. This tensor serves as input to the standard MHSA mechanism, functioning as a feature mixer between layers. Finally, we apply a linear layer to the last dimension to obtain a tensor with dimensions [batch size, number of layers, 1], and softmax is applied along the number of layers dimension to produce an individual posterior distribution for each sample in the batch. A detailed description will be provided in Section 3.2.
2. VARAN Evaluation on ASR Using More Complex Benchmarks:
   • We intentionally did not use a language model in our ASR task evaluation setup. Our goal is to explore the best utilization of SSL model layers, and comparing layer-aggregation methods after LM-decoding could introduce bias. While we agree that incorporating strong models into the evaluation pipeline can improve GER and minimize differences between layer-aggregation methods in terms of WER score, we believe this does not discredit the method overall.
3. The architecture appears to differ between speaker-dependent and speaker-independent tasks - this needs elaboration:
   • We trained all models in a speaker-independent setup. If we misunderstood your question, please let us know so we can address it more accurately.
4. State-of-the-art ASR systems typically incorporate language model integration. Why was this omitted?
   • In this paper, we do not propose a state-of-the-art ASR system but rather a layer-aggregation method applicable to various SSL architectures and downstream tasks. As mentioned previously, in the experimental setup, we omit rescoring SSL model predictions with the Language Model to avoid introducing bias in the comparison of layer-aggregation methods.
5. The reliance on human knowledge for prior distribution is problematic. Specifically, finding appropriate prior distributions for each task is impractical in real-world applications:
   • As mentioned in Section 3.3, for most tasks where we do not have any prior knowledge of layer importance, we suggest selecting the best prior distribution using a hyperparameter search of the chi-squared distribution, as it can be easily adapted to most tasks. Please refer to Section 3.2 and Appendix B for more details.
6. Even within a single task like ASR, it's unclear whether a single prior distribution can effectively handle variations across languages, dialects, and noise conditions:
   • As mentioned above, the prior distribution is parameterized by the parameters of a statistical distribution, which can be considered a hyperparameter of the model. Thus, it should not be generalized for multiple languages, dialects, and noise conditions. Like other hyperparameters, such as learning rate or batch size, it should be chosen individually based on the validation set.

Thank you for your thoughtful response! We would like to provide some clarifications to ensure that we are on the same page.

For instance, while the numerical values between layer 1 and layer 10 show large gaps, these should be treated as two distinct classes rather than continuous values.

Parametrization of our posterior distribution allows us to handle this dependency (gradients from the posterior flow through the entire model, except the feature extractor, so that all model parameters are learned to provide a better distribution). If there is no need to apply weights to the distanced layers, the model will learn it.

The ability of the model to allocate weights to distant layers is a strength, not a flaw. An example can be imagined where both low and high-level features are needed to correctly predict a label (e.g. emotion recognition depends on text and intonation).

such as Gumbel softmax

We believe Gumbel softmax is used for other cases where you need to output one value from your discrete distribution during inference, allowing us to obtain a differential approximation of the argmax. In our case, we combine our predictions with continuous weight from posterior PDF output during inference. Gumbel softmax also introduces additional variance and yields loosen lower bound during training.

Let us also note that Gumbel softmax still requires logits as input, so parameterizing the logits is completely independent of using Gumbel. One could use MHSA, RNN, or even a linear layer to obtain logits. Therefore, there is no reason that the probability distribution obtained after Gumbel softmax will have only one mode.

Evaluating the method solely on relatively simple ASR benchmarks like LibriSpeech is insufficient to demonstrate its effectiveness

We are currently working on the additional results and hope to get it in a 2 days.

Let us note that we believe the goal of our paper is to present a novel method for feature aggregation for a wide variety of tasks, rather than a SoTA ASR model.

If you have any questions or need further clarification, we are more than happy to answer.

Dear Reviewer "YerL",

We would like to remind you that we have updated our paper since your last review, and we believe the experiments we have provided address your concerns. In light of these revisions, we kindly ask you to reconsider your score. If you have any further questions, we would be happy to address them.

Dear Reviewer,

Thank you for your time and for providing a detailed review. We appreciate your recognition of VARAN as a novel layer aggregation method that adjusts the prior distribution to control layer output weight, and your acknowledgment of our experiments showing alignment between the actual posterior and prior distributions, along with strong performance on various downstream tasks.

Regarding the weaknesses you pointed out and your questions:

1. How was the prior distribution chosen, and what were the criteria for selecting the non-central Chi-squared distribution?

   • Initially, we suggest that VARAN can be used with different downstream tasks and backbone models. Therefore, we designed a process that does not require any prior knowledge of the task, aiming to simplify it by allowing a grid search to find the optimal statistical distribution parameters. The non-central Chi-squared distribution is a good choice because, by adjusting its parameters and the ability to reverse the result distribution, we can construct priors with modes corresponding to different layers.

2. In Figure 2, experiments were conducted to set various beta values for the non-central Chi-squared distribution. Is he optimal beta value found simply by grid search, or if another search method was used?

   • Yes, the optimal beta is found simply through grid search. Our results show that $\beta = 0.5$ is a strong coefficient for posterior distribution regularization, making it almost identical to the prior, which causes suboptimal performance. Conversely, setting $\beta < 0.01$ results in too little regularization, also leading to suboptimal performance. Therefore, the suggested values for grid search are between 0.01 and 0.5.

3. In the ASR experiments, the evaluation was conducted using LibriSpeech data, the domain employed in the pre-training. This method does not align with the goal of assessing generalizability, suggesting a need to test with other speech datasets.

   • Although audiobooks from the LibriSpeech domain were used during the pre-training stage, we do not believe this impacts the experiments. The advantage of having domain-specific data would equally benefit all layer-aggregation methods, and similarly, all methods might be affected by its absence, but the performance ranking would remain unchanged.

4. This paper conducted experiments on automatic speech recognition, speaker verifcation, and emotion recognition. Additional datasets are necessary to demonstrate more generalized performance compared to existing methods.

   • In our paper, we test the proposed layer aggregation method on three different tasks, each using a separate dataset. We believe these experiments are enough to support our claims about generalization to other datasets, as they show the method's effectiveness across various tasks. This approach is similar to what was used in the WavLM [1] and wav2vec 2.0 [2] papers, where they also demonstrated their models' broad applicability by testing on diverse tasks using single dataset.

   [1] Sanyuan Chen, Chengyi Wang, Zhengyang Chen, Yu Wu, Shujie Liu, Zhuo Chen, Jinyu Li, Naoyuki Kanda, Takuya Yoshioka, Xiong Xiao, et al. Wavlm: Large-scale self-supervised pre-training for full stack speech processing. IEEE Journal of Selected Topics in Signal Processing, 16(6):1505–1518, 2022.

   [2] Alexei Baevski, Yuhao Zhou, Abdelrahman Mohamed, and Michael Auli. wav2vec 2.0: A framework for self-supervised learning of speech representations. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural In- formation Processing Systems, volume 33, pp. 12449–12460. Curran Associates, Inc., 2020.

5. Why were the WavLM and Data2vec chosen among various self-supervised models, and why did you not experiment with other models?

   • WavLM and data2vec were chosen from among other speech self-supervised models because they are the newest, widely used, and have achieved the best performance. Architecturally, all speech self-supervised models are quite similar, so other models like wav2vec 2.0 and HuBERT could easily be extended to use the VARAN layer-aggregation method. However, we did not experiment with these other models due to computational constraints.

We have now answered all the questions we could without conducting additional experiments. During the rebuttal stage, we will strive to expand the experiments section as much as possible, incorporating your comments about method evaluation.

Thank you once again for your valuable feedback.

Dear Reviewer "32Z2",

We would like to remind you that we have updated our paper since your last review, and we believe the experiments we have provided address your concerns. In light of these revisions, we kindly ask you to reconsider your score. If you have any further questions, we would be happy to address them.

Dear Reviewer,

Thank you for your time and for providing a detailed review. We appreciate your recognition of VARAN's ability to adapt layer weights for each individual sample and are glad you found this approach highly effective, especially for the SER task.

Regarding the weaknesses you pointed out and your questions:

1. The authors do not compare VARAN to adapters, which are (i) a more cost-effective fine-tuning method (involving fewer parameters) and (ii) better suited for adapting models to different downstream tasks, especially when there is limited data available for fine-tuning:

   • While fine-tuning SSL models with adapters [1-3] is indeed more cost-effective, we believe that VARAN should not be directly compared to them, as the methods are orthogonal. VARAN is a layer aggregation method that can be applied to various fine-tuning techniques, as it is an essential step for making predictions. In this paper, we focused on exploring the method itself and demonstrating its application across multiple tasks. Setting up different fine-tuning techniques is beyond the scope of this research; please note that the proposed method can also be applied to fine-tuning with adapters.

     [1] Neil Houlsby and Andrei Giurgiu and Stanislaw Jastrzebski and Bruna Morrone and Quentin de Laroussilhe and Andrea Gesmundo and Mona Attariyan and Sylvain Gelly, Parameter-Efficient Transfer Learning for NLP, ICML, 2019

     [2] Edward J. Hu and Yelong Shen and Phillip Wallis and Zeyuan Allen-Zhu and Yuanzhi Li and Shean Wang and Lu Wang and Weizhu Chen, LoRA: Low-Rank Adaptation of Large Language Models, 2022

     [3] Zih-Ching Chen and Yu-Shun Sung and Hung-yi Lee, CHAPTER: Exploiting Convolutional Neural Network Adapters for Self-supervised Speech Models, ICASSPW, 2022

2. Use of Multi-Head Self-Attention (MHSA) for Predicting a Posterior Distribution in Figure 1:

   • The input to MHSA block (Figure 1) consists of hidden states from the Transformer Encoder block. Formally, this is a list with a size equal to the number of layers, where each element is a tensor with dimensions [batch size, sequence length, hidden dimension]. We then apply mean pooling along the sequence length and concatenate the results into a single tensor, resulting in dimensions [batch size, number of layers, hidden dimension]. This tensor serves as input to the standard MHSA mechanism, functioning as a feature mixer between layers. Finally, we apply a linear layer to the last dimension to obtain a tensor with dimensions [batch size, number of layers, 1], and softmax is applied along the number of layers dimension to produce an individual posterior distribution for each sample in the batch. A detailed description will be provided in Section 3.2.

3. Why Data2Vec Was Used for the SER Task and Not for All Tasks:

   • For a fair comparison of data2vec across all tasks, we needed to run a grid search with optimal hyper-parameters for each task, which is very computationally-intensive. That is why we chose the SER task specifically, as it requires less training time compared to ASR or SV, given that the IEMOCAP dataset is smaller. Despite this limitation, VARAN is designed to work with various upstream models, including wav2vec 2.0 and HuBERT, due to their similar architectures. We believe that the success of the data2vec model utilizing VARAN's layer aggregation method on the SER task demonstrates its generalization ability, but it would be interesting to compare the relative performance improvements on other tasks as well.

During the rebuttal stage, we will strive to expand the experiments section as much as possible, incorporating your comments about method evaluation.

Thank you once again for your valuable feedback.