Paper: The Need for Speed: Pruning Transformers with One Recipe (ICLR-2024)

Reviewer comment (Experiments):
It seems like PTF is the most relevant and competitive baseline. Per 4.1, the authors compare with an abridged version of PTF however. If you use the same tricks that PTF used (e.g., mask rearrangement -- maybe mask tuning), would your method outperform PTF?

Author response:
$\textbf{[Applying PTF Tricks]}$
In the language domain, PTF presents the most competitive baseline. We are augmenting our results using the associated tricks from PTF and have provided them in the below table. Improvement tricks, such as those applied here, often boost performance in narrow applications and generally require data or time that a practitioner may find difficult to execute. For these reasons, we didn’t initially report these augmented results; our key motivation, in this paper, is to show generality over a variety of tasks and domains. 

| Method      | MNLI | QQP | QNLI | SST | STS-B | MRPC |
| ----------- | ----------- |----------- |----------- |----------- |----------- |----------- |
| $PTF_{+tricks}$      | 82.51 | 90.35 | 90.06 | 92.49 | 88.00 | 85.27 | 
|$OPTIN_{+tricks}$   | $\textbf{82.74}$ | $\textbf{90.43}$ | $\textbf{90.35}$ |$\textbf{92.73}$ | $\textbf{88.21}$ | $\textbf{85.68}$ |

Reviewer comment (Presentation):
Figure 1 is a bit complicated to look at. One possibility is to break up this figure into two (the left vs. the right). The details are nice, but they detract from the main point (e.g., the activation difference and the final difference in logit distribution). The logits could be illustrated as two histograms, highlighting their difference.

Author response:
$\textbf{[Improving Figures]}$
We appreciate the reviewer’s comments on Figure 1. We have made another attempt to separate the details from the framework diagram to develop a more pleasing and easy-to-understand image. We also greatly appreciate your comment on our Figure 2. We had made every effort to add an additional distinguishing feature for different methods, (i.e. marker shape + color in each plot). We had missed this one in our review, and are updating it. We also split the plot into two, as we too can understand the difficulty posed by the overlay. We will reply in this thread once the main paper has been updated to reflect this change.

Reviewer comment (Presentation):
The double axis in Figure 2 right is really confusing. Maybe just add a 5th plot? Might not be so easy to see for anyone with color blindness.

Author response:
$\textbf{[Improving Figures]}$
We appreciate the reviewer’s comments on Figure 1. We have made another attempt to separate the details from the framework diagram to develop a more pleasing and easy-to-understand image. We also greatly appreciate your comment on our Figure 2. We had made every effort to add an additional distinguishing feature for different methods, (i.e. marker shape + color in each plot). We had missed this one in our review, and are updating it. We also split the plot into two, as we too can understand the difficulty posed by the overlay. We will reply in this thread once the main paper has been updated to reflect this change.

Reviewer comment (Novelty):
The paper's motivation could be stronger. The abstract and title focus heavily on "without retraining" but it seems like many previous papers already focus on the post-training regime (as the paper cites). The more important motivation is then: What is PFT lacking that is present in this method? *Why is that difference important, intuitively?

Author response:
$\textbf{[Motivation + Italics]}$
We appreciate the reviewer’s comments on motivation. The practical applications of our methods is central to the paper’s design, and it is thus crucial that a reader understands the novelty and applicability of our method. Restating our contributions, in reference to PTF: 

* [C1] Our OPTIN framework leverages the trajectory measure, which efficiently captures long-range dependencies between parameters in the network. PTF approximates the fisher-information matrix to determine the relative importance of each parameter in the network. However, this method of ranking parameters is highly subject to parameters that may exhibit high gradients and thus fails to sufficiently capture relationships between parameters across layers. In comparison, our method uses the forward pass to establish and capture relationships between successive parameters ultimately achieving superior performance.

* [C2] The OPTIN framework has exhaustively verified superior performance across language and vision domains, including an extension to CNN architectures and segmentation models. Meanwhile, PTF has only shown applicability in the language domain. This provides significant motivation for our method, as it can be easily adopted as a plug-and-play across various domains in deep learning. We would really like to highlight this generalizability to multiple domains and model architectures, as many works depend on the intricacies of a particular type of data or model in order to prune efficiently. 

We are in the process of refreshing our Abstract, Introduction, and Related Works to best elicit this motivation, primarily focusing on the multi-modality and widescale applicability of our pruning framework. 

Finally, we thank the reviewer for noticing the italicized words. This was most definitely a formatting error/typo on our side. We meant to italicize only once in the abstract and once in the introduction, and are correcting this in the updated draft.

We hope these responses have helped clarify any questions the reviewer may have. We are adjusting the main paper to better focus on said motivation as well as improve the quality and presentation of our figures. We hope you will consider our changes and improvements in order to best evaluate the potential of our work. Please feel free to reach out with any questions, we’d be more than happy to follow up with additional improvements or changes.

Reviewer comment (Experiments):
I think this paper is a ok paper. For discussion, I want to see the method performance on more complicated network architecture or datasets.

Author response:
We thank reviewer oBzK for recognizing the practical problem we are aiming to solve and appreciate the reviewer’s comments on the clarity and comprehensibility of our methodology. 

$\textbf{[More Architectures or Datasets]}$Following the suggestions for more experimentation, we have extended our method to the task of language generation using GPT-2 to show versatility (see below). We would however like to reiterate the diversity of tasks and modalities already included in the main paper. We were able to show competitive performance in both the natural language and image classification benchmarks. We further demonstrated applicability to tasks including transfer learning, and the high-resolution semantic segmentation. We highlight the complexity of semantic segmentation as the practical improvements in this domain could lead to huge potential impacts, particularly in the field of autonomous driving. Semantic segmentation on the cityscapes datasets operates at a resolution of 1024x2048, which is significantly higher than common resolutions for classification, making inference speed a huge bottleneck. By showing applicability to this style of model, we open avenues for future research to investigate the multi-modal capability of efficient model compression algorithms. Further, we demonstrate the robustness of our OPTIN framework as it is able to achieve competitive compression ratios on a variety of tasks. We would like to highlight that our main contribution is the generality and versatility of our one-shot approach; Some methods exhibit tricks or artifacts to improve performance in specific applications, however, our design approaches the problem from a more general perspective of a model compression framework that scales well across modalities. 

$\textbf{Extending OPTIN to Language Generation}$
To show versatility, we produce an initial experiment on Language Generation. Our framework has primarily focussed on model compression techniques for transformers. Given their prevalence in modern deep learning, the research into efficient deep learning attempts to solve a practical problem across the industry. In particular, we’ve shown the flexibility and extensibility of our method over a variety of tasks. In this section, we produce some initial findings by extending our framework to language generation tasks. In particular, we evaluate the GPT-2 model on the WikiText2, PTB, and LAMBDA datasets. Unlike our previous results in the main paper, these datasets are used to evaluate a model's ability in language generation. This differs from the standard GLUE benchmark in NLP that mainly measures accuracy/correlation. We hope to extend our tests in the main paper, as showing generalizability of our method to a primarily decoder-related task, such as language generation, opens multiple avenues for future extensions, including those in image-generation.

| Ratio (%)  | Model | Pruning Time (Hours) | ~ $\pm \triangle WikiText2$ (PPL)  | ~ $\pm \triangle PTB_{Text} $ (PPL)  | ~ $\pm \triangle LAMBDA_{Text} $ (PPL)  |
| ----------- | ----------- |----------- | ----------- | ----------- |----------- |
| 1.0 | GPT-2| -|(28.2) | (31.9) | (42.3) |
| 0.9  | GPT-2 | 0.3 | +3.8|+4.1|+2.5|
| 0.8  | GPT-2 | 0.3 | +4.1|+6.3|+5.1|

$\textit{* First row includes the approximate baseline for each dataset.}$
$\textit{* Pruning time is insured only once. Activations are stored and re-used for various compression ratios.}$
$\textit{* Relative Parameter Ratio accounts for the hidden dimension and attention heads.}$
$\textit{* PPL referenced Perplexity, a lower measure is better for language generation}$

We hope that our main results on Semantic Segmentation, and the additional introductory experiments here provide sufficient inclination towards the future potential impact of such a generalizable model compression framework. We believe that there is strong potential for future avenues of research that can stem from achieving a generalizable approach to compressing transformer models across modalities, and we hope that our work compellingly conveys the message. We would like to thank the reviewer again for their time and interest in our work.

Reviewer comment (Evaluation):
**Performance on Vision Models**: As highlighted in Table 3, the performance of OPTIN (without token reduction) appears to be inferior to that of VTP and PoWER for both evaluated models. It seems that the impressive performance of $OPTIN_\tau$ can be mainly attributed to token reduction. Does it mean OPTIN is more suitable for language models rather than vision models?

Author response:
$\textbf{[Addressing Weaknesses]}$
We very much appreciate the time taken by reviewer RsPn to understand our work. With regard to the competitive methods in the vision domain, it seems that each method has its own approach to the granularity of pruning. In particular, VTP performs layer-wise patch-wise pruning while also removing parameters in the QKV and linear layers within each transformer block. PoWER introduces a soft-extract layer in each self-attention module to determine which tokens are kept at each self-attention layer and CP-ViT performs head and patch pruning under a dynamic search. The focus on patch-based reductions enables these methods to achieve larger FLOP reduction rates with minimal parameter changes. Each method exhibits its own artifacts which make it hard to compare fairly. In our case, we justify the inclusion of token reduction (sometimes referred to as patch reduction), to achieve a more fair comparison across the methods. Our base $\beta$ configuration directly removes parameters, while the augmented $\tau$ configuration leverages token reduction to further achieve more competitive flops reduction ratios. Further, we’d like to reiterate that even under token reduction techniques, our method still outperforms state-of-the-art approaches, as seen in Table  4. For these reasons, we find our method to be applicable in the vision domain.

We hope that these answers have helped to clarify any questions you may have about our work. Please feel free to ask any follow-up questions you may have to help us best exemplify our work.

Reviewer comment (Experiments):
In Table 2, the results for PTF are directly cited from the original paper. Could the authors confirm that the pre-trained checkpoints and fine-tuning settings are consistent across all experiments to ensure a fair comparison?

Author response:
$\textbf{[Table 2 Results and Checkpointing]}$
In Table 2, we use the checkpoints provided by the authors of PTF to replicate the baseline performance. Applying the OPTIN method to the baseline model incurs no fine-tuning or additional training costs as our method masks the least important parameters, and thus we can ensure both methods are compared fairly as they start with the same model.

Reviewer comment (Writing):
Could the authors give further explanations on 'IM-Dense (After dense embedding layer)'. What does it mean?

Author response:
$\textbf{[Ablation Experiment in Table 1a - IM Dense]}$
The purpose of this ablation is to decide at which location to extract the features used in our distillation loss. Following recent works, including [1,2], we noted that there are multiple locations from which we can extract meaningful features for our distillation loss. At the end of each transformer block there exists a 2-layer MLP (in the case of BERT models). The first layer increases the hidden dimension, while the second is complimentary in operation. In this ablative study, we compare the location to compute distillation error in subsequent transformer blocks. FFN refers to extracting after the 1st  linear layer at the increased embedding dimension, thus leveraging a more robust feature comparison, while IM-Dense refers to extracting after the 2nd linear layer. Our study ultimately concluded that a higher degree of information regarding parameter importance could be extracted from the distillation error computed along the larger hidden embedding dimension. We will revise the in-text terminology to make this clearer.

*[1] Peng et al. 2019: Correlation Congruence for Knowledge Distillation
*[2] Hao et al. 2022: [see main paper]

Reviewer comment (Reproducibility):
How does OPTIN determines layer-wise sparsity? Does OPTIN rank weights globally?

Author response:
$\textbf{[Sparsity]}$
Previous works have shown the benefits of using a global sparsity as opposed to layerwise sparsity [3,4]. In line with these works, we too exploit a global sparsity constraint as we rank the importance of each parameter globally in the network. See Details in Appendix A.1.

*[3] Lin et al. 2018: Accelerating Convolutional Networks via Global dynamic Filter Pruning
*[4] Song et al. 2022: [see main paper]

Reviewer comment (Experiments):
In Table 1.(d), why $\lambda_c$ is chosen to be 0.1 even though 0.01 provides better results?

Author response:
$\textbf{[Lambda Value]}$
We appreciate the reviewer pointing this out. The goal of our ablations was to test the sensitivity to different parameters and validate our design choices. In the early inception of our project, we defaulted to using 0.1 for language models, to favor the trajectory measurement in computing parameter importance. We hoped this ablative study could help elucidate the importance of trajectory, however, we note that the values of our main language table could be improved with this lambda parameter, and have rerun the language experiments with the updated lambda value. As can be seen in the table below, there is an improvement across most tasks in the GLUE benchmark when compared to the previous configuration.

| Method      | MNLI | QQP | QNLI | SST | STS-B | MRPC |
| ----------- | ----------- |----------- |----------- |----------- |----------- |----------- |
| $OPTIN_{\lambda=0.1}$      | 81.90 | 90.06  | 88.49 | 92.24 | $\textbf{87.25}$ | 85.13|
| $OPTIN_{\lambda=0.01}$   | $\textbf{82.12}$ |$\textbf{90.08}$ | $\textbf{88.54}$ | $\textbf{92.36}$ |87.19 |  $\textbf{85.21}$|