The Need for Speed: Pruning Transformers with One Recipe

We appreciate reviewer meSa’s recognition of our comprehensive experiments and well-vetted design choices. We provide responses to the questions below in an effort to bring more clarity to our work:

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

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

$**[Pruning Time/Memory]**$ We thank the reviewer for this question. This helps elucidate one of our main contributions. As such, we have attached a table below detailing the time taken to prune various models on different datasets. We’d like to reiterate that our method appears to scale well across a variety of architectures and tasks, achieving competitive performance-to-compression ratios on the scale of minutes. There is some overhead time required in computing the activations however this is included in our pruning time listed below. In terms of memory allocation, all our pruning experiments, across both language, and vision (including semantic segmentation), can be run with 1x Nvidia RTX2080 with less than 10 GB of memory.

$**Average Time Analysis for ** OPTIN$

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.

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

$**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.

$** First row includes the approximate baseline for each dataset.*$ $** Pruning time is insured only once. Activations are stored and re-used for various compression ratios.*$ $** Relative Parameter Ratio accounts for the hidden dimension and attention heads.*$ $** 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.

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

We thank reviewer RsPn for their thoughtful and insightful comments about our paper. We appreciate their recognition of our comprehensive experiments and robust performance. We hope to provide more insight on their questions below:

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

$**[Time Requirement]**$ We agree with the reviewer that a time comparison can help validate the practicality of our method and showcase the true time benefit of an approach that is one shot. As such, we have added a table below to showcase the time costs of our method. This table shows the practicality of our method as it scales well across the tasks and model architectures, achieving competitive results on the scale of minutes. If successful, we hope to integrate this into our main draft to better depict the motivation and implications of our work.

$**Average Time Analysis for ** OPTIN$

$**[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]

$**[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]

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