Context-aware Dynamic Pruning for Speech Foundation Models

Thank you for your encouraging response.

Concern about the low-resource languages

Thank you for pointing this out. To address this, we reviewed the pretraining data used by OWSM-V3.1 to confirm the data volume for each language. Based on this analysis, we incorporated Hungarian and Bulgarian into Europarl-ST for additional training. The data for these two languages was sourced from the Fleurs dataset and combined with Europarl-ST. The data volume used for OWSM pretraining for each language is as follows:

Below, we present the evaluation results for decoders and encoders pruned individually. A sparsity of 0% represents the baseline model, which is simply fine-tuned without pruning.

Encoders

Decoders

Pruning the decoder led to declines in accuracy, similar to what was observed for higher-resource languages such as French and German. However, the impact of pruning the encoder appears to be more significant for low-resource languages. Even at lower sparsity ratios, encoder pruning caused a greater degradation in performance.

This highlights the critical role of the encoder in preserving accuracy for low-resource languages. Since these languages typically have less training data, the encoder's ability to effectively capture the input feature is essential. In contrast, the decoder appears less sensitive to sparsity changes, likely because its task relies on the representations from the encoder.

We have added this point to the Appendix B.

How critical is Gumbel-softmax for dynamic pruning? Could a simpler sigmoid-based selection function perform comparably in this setting? How are the temperatures for Gumbel-softmax adjusted?

Attempting to fine-tune OWSM-v3.1 with a sigmoid-based function is unlikely to work effectively. Through various experiments with different hyperparameter settings, we found that while training may succeed when the gate predictor outputs continuous values, issues often arise during inference. Specifically, even when the softmax temperature was lowered to make the gate predictor outputs more sharply peaked, we observed that the gate probabilities frequently converged to values around 0.4–0.6. This means that during training, some module outputs were scaled by 0.4, but during inference, those same modules were completely skipped. This mismatch caused a significant drop in inference accuracy.

The Gumbel-Softmax approach we adopted addresses this issue by leveraging the Straight-through Gumbel-Softmax Estimator (SGSE), which enables the use of strictly binary values as pruning masks. With SGSE, the pruning mask is fully binary (0 or 1) during the forward pass in training, ensuring consistency between training and inference and maintaining robust performance. We recognize that this critical detail was not adequately explained in the paper. To address this, we have added a description in Section 3.3.

Based on this study, can we reasonably infer that the E-Branchformer architecture contains additional components that could be pruned for efficiency, potentially aligning it more closely with the Conformer architecture?

While there are similarities between the pruned E-Branchformer and Conformer architectures, the parallel structure of the E-Branchformer appears to be beneficial in the later layers, where modules remain fully active even at sparsity levels between 50% and 70%.

In models with a large number of stacked E-Branchformer layers, our results suggest the possibility of reducing the reliance on self-attention in the earlier layers. While the computational cost of Conformer and E-Branchformer layers is comparable[1], replacing some of the earlier layers with convolution-based layers that place less emphasis on self-attention could be an approach to improve efficiency.

Reference [1] Peng, Y., Kim, K., Wu, F., Yan, B., Arora, S., Chen, W., Tang, J., Shon, S., Sridhar, P., Watanabe, S. (2023) A Comparative Study on E-Branchformer vs Conformer in Speech Recognition, Translation, and Understanding Tasks. Proc. INTERSPEECH 2023, 2208-2212, doi: 10.21437/Interspeech.2023-1194

Is elapsed time the correct metric for comparing model efficiency? In my experience, any torch implementation does not yet improve efficiency in terms of flops as torch is not yet designed for it. Can the authors show metrics in terms of FLOPs maybe?

We believe that elapsed time provides a practical perspective for comparing model efficiency in real-world scenarios, as it captures hardware and implementation-specific factors. Additionally, we appreciate your suggestion concerning complexity, and we included FLOPs for theoretical efficiency in table 1.

The FLOPs for the encoder, decoder, and encoder-decoder at various sparsity levels are shown below. These values, measured in GFLOPs, were calculated using the fvcore package. The table presents the target sparsity ratios in descending order: 0.1, 0.3, 0.5, 0.7, and 0.9, with the non-sparse baseline having a value of 3781. The significant reduction in FLOPs for the decoder is attributed to its autoregressive inference process, where the beam size is treated as the batch size. In this analysis, we set the beam size to 5, meaning the encoder processes a batch size of 1 per inference, while the decoder processes a batch size of 5. These results have been incorporated into Section 4.

Please add what conclusion we should draw from a figure to the caption of each figure. This makes the figures difficult to follow.

Thank you for your valuable feedback. The figures in our paper can be broadly categorized into two types: those illustrating pruning patterns and those evaluating performance metrics. For performance evaluation graphs, we agree with your observation and have updated the captions to include key conclusions drawn from the figures. As for the figures related to pruning patterns, they are primarily used for comparative analysis across multiple figures. The detailed discussions and insights derived from these comparisons are provided in the main text.

I would appreciate if the authors can point out whats novel in their "novel context-aware pruning technique" (except the gating technique) which is difficult for me to understand.

We apologize for not clearly stating the differences in our use of Gumbel-Softmax during training. We employed the Straight-Through Gumbel-Softmax Estimator, which ensures that the computational graphs are fully aligned between training and inference. We have added detailed descriptions in Section 3.3.

Thank you very much for your thoughtful and insightful comments.

A major weakness of the paper is the lack of baselines and comparison with prior-art.

We sincerely appreciate this insightful comment. To address this, we have incorporated new results using unstructured pruning methods, specifically magnitude-based pruning and PARP[1]. The table below presents a comparison between these two approaches. Fine-tuning for PARP was performed on a training dataset that included both ASR and ST, and WER was compared on the German ASR task.

Even with unstructured pruning, it was evident that pruning the decoder keeps its performance. However, compared to module-level results at the same sparsity ratio, module-level pruning, especially on the encoder side, maintained accuracy more effectively. We mentioned this additional experiments in section 2 and have been added to the Appendix D. We are currently doing our best to tune the parameters for both encoder and decoder side. If we are able to train the model with better result, we will update the table and parameter settings in the Appendix.

The paper feels like an analysis paper.

We fully acknowledge that our work builds upon the Peng et al. (2023). However, our approach introduces significant differences, particularly in the depth of analysis conducted, with a specific focus on module-wise pruning strategies. Furthermore, we have made key adjustments that allow the model to be trained to achieve a targeted sparsity level. While Peng et al. (2023) confined their analysis to sparsity levels ranging from approximately 60–90%, our work extends this analysis to include efficient module selection strategies across a broader range of sparsity levels, including lower sparsity settings. Additionally, we have modified the use of Gumbel-Softmax in our approach. Instead of employing the standard Gumbel-Softmax, we utilized the Straight-through Gumbel-Softmax Estimator (SGSE) to enable training with strictly binary pruning masks. This important distinction was not adequately detailed in the initial version of our paper; therefore, we have added a comprehensive explanation in Section 3.3 to clarify this point.

The scope of the experimental section is rather narrow.

Similar concerns regarding the dataset were also raised by other reviewers, and we have conducted additional experiments incorporating VoxForge, Fleurs, and LibriSpeech from different perspectives. Regarding the suggestion to expand the scope of model analysis, as stated in Section 3.2, we determined that OWSM is more suitable for these experiments compared to Whisper.

Reference
[1] Lai, Cheng-I. Jeff, et al. "Parp: Prune, adjust and re-prune for self-supervised speech recognition." Advances in Neural Information Processing Systems 34 (2021): 21256-21272.

Thank you for providing such a valuable review.

Limited Baseline Comparisons

Based on reviewer's feedback, we conducted additional experiments involving layer skipping, reproduction of , and unstructured pruning. However, considering the objective of our study—to analyze and address the research question—we would like to emphasize that skipping layers reduces the granularity of analysis compared to module-based pruning. This distinction is important for ensuring clarity in our findings.

Reproducibility:

We conducted our experiments with ESPnet, and uploaded our experimental codes to anonymized github. https://anonymous.4open.science/r/anonymized_dynamic_pruning-C546/

We have also added the citation for Gumbel-Softmax, and added more detailed explanation on how we utilized the Straight-through Gumbel Softmax Estimator. Thank you very much for bringing this to our attention.

Need for more In-Depth Ablation Discussion:

Thank you for your observation regarding the need for a more in-depth ablation discussion based on pre-training dataset of OWSM model. Since all the pretraining data is available, we checked the data distribution according to the paper, specifically for French, German, and Italian. The data sizes used for OWSM pre-training is as follows:

Given the smaller amount of data for Italian, we conducted additional analyses to investigate whether there were significant differences compared to other languages or tasks. However, we did not identify any noteworthy findings. In response, we are also examining languages with even smaller pretraining datasets, such as Hungarian and Bulgarian, which contain 97 hours and 18 hours of data, respectively. We incorporated these analyses and findings into the revised manuscript, specifically in Appendix B.

Interpretability of Pruning Choices:

We deeply appreciate your insightful comments. To address your feedback, we included additional discussions based on the characteristics of each tasks on the following points:

Joint Pruning in Section 4.2.1

• When only the decoder is pruned, early src-ATT layers are removed. In contrast, when both the encoder and decoder are pruned together, earlier src-ATT layers in the decoder become more active.
• This difference likely stems from whether all encoder modules are fully utilized. If encoder capacity is limited, the decoder compensates by computing self-ATT and FFN layers on the audio features. If the encoder is fully available, the decoder focuses on incorporating more contextual information from output tokens before src-ATT.

Differences Between ASR and ST in Figure 7

• Self-ATT is more prominent before src-ATT in ST than in ASR.
• This is likely because alignment via attention is more critical in ST, whereas ASR prioritizes directly integrating audio features, leading to differences in the timing of self-ATT computations. We will incorporate these discussions into the manuscript, including a detailed analysis in Section 4.2.1 and 4.2.2.

Contribution Clarification:

In this work, we utilized the Straight-through Gumbel-Softmax Estimator (SGSE) to create a trainable pruning mask that outputs strictly binary values for training. We regret that this important detail was not clearly articulated in the paper.

For example, during training, the gate probabilities would be multiplied to the output of each modules , and during inference, activation is determined using a threshold value ranges from 0 to 1. After training our model, we realized that there are several values close to the threshold. This means that use partially activate those modules during training, while we completely ignore these modules during inference. We solved this issue by using the SGSE.

We have added a paragraph to Section 3.3 to explain this approach and included a detailed problem formulation in the Appendix A.

Have you observed any patterns in performance across different subtasks or languages, particularly with varying sparsity targets?

We understand this comment as relating to the need for an in-depth analysis, as previously noted. For further details, please refer to that section.

I believe the pdf was not updated (at least for me it seems to be the same thing). Can you check that? I couldn't verify the results you added.

Also the authors did not address the question about significance of the results. Several conclusions were made on pretty narrow differences, so I wonder if authors can provide more information about it.

Does your method require specifying the source language at inference, or can the model identify the language automatically, as Whisper does?

Our method requires specifying the source language at inference time; it does not incorporate an automatic language identification process as Whisper does. This is because our objective is to analyze how a pre-trained foundation model manages task-specific and language-specific differences. We believe that incorporating language ID estimation is not necessary to our work.

Would it be possible to provide code or configuration files to facilitate reproducibility?

We have uploaded our experimental code to an anonymous github. Please refer to the following link: https://anonymous.4open.science/r/anonymized_dynamic_pruning-C546/

As a side note, I recommend adding baseline (without pruning) to the tables as this would help visualization and understanding.

Thank you for your suggestion. While we have not yet included the baseline (without pruning) in the current tables due to the large number of models being compared, we explor ways to adjust the table layout to incorporate this information more effectively. We will make our best effort to present this in a concise manner in the revised version.

We also calculated the FLOPs for these layer-skipped models. Since different modules are utilized depending on whether module-level or layer-level pruning is applied, we examined the resulting differences. The table below shows the FLOPs when the encoder is pruned. As a result, we found that module-level pruning results in slightly lower FLOPs compared to layer-level pruning. These findings further emphasize the importance of module-level pruning from a performance perspective. We will address these points in Section 3.2 and include additional details in the Appendix C.

GFLOPs

While the paper focuses on reducing inference time, it doesn't address the computational overhead introduced by the dynamic pruning process itself and searching the best sparsity number during finetuning.

Thank you. We calculated the FLOPs for each sparsity level and added them to Table 1. The computation overhead for determining sparsity during fine-tuning, including both the fine-tuning and evaluation processes, required approximately two days on a single A40 GPU.

Lack majour novelty

Although our method is indeed an extension of Peng et al. (2023), we have introduced novel contributions compared to Peng et al. (2023) in the following aspects:

• Multilingual and Multitask Training: We trained our model in a multilingual and multitask setting.
• Different Utilization of Gumbel-Softmax: We employed Gumbel-Softmax differently in our approach: that is, we used the straight-through Gumbel-Softmax estimator for training pruning masks.

We realized that we did not adequately describe the second point in our paper. To address this, we included additional explanations in the section 3.3 of the Methodology. Moreover, we noticed that our method lacks a detailed problem formulation. We have addressed this by pointing it out in section 3.3 and adding the formulation in the Appendix A.

Typo ? [line 130]

Thank you for catching that! We have corrected it.

I think [at line 285] the sentence is incomplete.

Thank you again for pointing this out. We have revised it to ensure the sentence is complete.

In Figure 1, it is not clear how mean pooling can be applied to the input of the transformer decoder given its autoregressive nature.

Thank you, the diagram could indeed be clearer. The mean pooling applies to the audio information, which is used to control the decoder. This does not interfere with the auto-regressive processing on the decoder side during inference. To clarify this further, we revised Figure 1 (a) and (b) to make it more apparent that this pertains to the audio information.

Not able to understand after reading the paper, if the time reduction is from using the joint finetuned model or individual task finetuned model. I think it is for individual task finetuned models.

We apologize for not making this clear in the paper. The model used for inference is trained separately for each task, with individual models fine-tuned for ASR and ST tasks. We included this information in the caption of table 1.

Please provide more detailed comparison of their approach to previous work and their novel contributions beyond Peng et al. (2023b).

We provided a detailed explanation from the perspective of lack of novelty in response to the previous comment. Specifically, we clarified the differences in the use of Gumbel-Softmax in Section 3.3 and added a detailed formulation in the Appendix A. We also reproduced their results for comparison and included these results in Appendix B.

Reference

1. Yifan Peng, Jaesong Lee, and Shinji Watanabe. I3d: Transformer architectures with input-dependent dynamic depth for speech recognition. In ICASSP 2023 - 2023 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 1–5, 2023

Thank you very much for reviewing our paper and for providing positive feedback.

The size of the dataset (20 hours) used is too small to make general claims.

We appreciate your observation. Indeed, the limited size of the dataset was a valid concern. To address this, we conducted additional experiments by incorporating the VoxForge dataset, which includes the same languages, into the original dataset. This resulted in a joint dataset for ASR task, whose sizes are as follows:

It is true that the absolute increase in data size is not substantial. However, we believe that this amount of additional data is sufficient to evaluate the impact of data increase on the performance of our method. Moreover, as our approach relies on fine-tuning a pre-trained model, using a dataset of this scale ensures practical and reproducible results. The results of the models trained on this expanded dataset are presented below. The table highlights the performance of the German ASR model with the decoder side pruned, showing improvements in accuracy compared to models trained on Europarl-ST alone:

The performance degradation across different sparsity levels is comparable between Europarl-ST and joint dataset. To reflect these results, we have updated Section 4 and Appendix B to include these additional experiments.

The authors do not compare their results with any other pruning methods.

Thank you for your feedback. In response to the feedback, we conducted additional experiments comparing module-based pruning, layer-level skipping, and unstructured pruning. These experiments were carried out using our joint ASR and ST dataset. We report the WER for German ASR, where the encoders and decoders were pruned separately. The sparsity patterns used in these experiments are illustrated in Figure 13 in Appendix C.

Encoder

Decoder

The degradation in accuracy on the encoder side is significant, leading us to believe that layer-level pruning for encoder side is not well-suited for speech foundation models. With layer-level pruning, even cgMLP, which is preserved in module-level pruning, is forcibly removed. This results in the loss of parameters that play a critical role on the encoder side, highlighting a key disadvantage of this approach.

For the decoder side, the results indicate that up to 70% sparsity, there is no significant difference between skipping at the layer level and skipping at the module level. However, at lower sparsity levels, the layer-level approach appears to perform better, whereas at higher sparsity levels, the module-level approach demonstrates superior performance. Based on these findings, we believe that at higher sparsity levels, the roles and processing order of individual modules are more effectively optimized compared to the more coarse-grained approach of skipping entire layers, leading to the observed results.