Towards Robust Pruning: An Adaptive Knowledge-Retention Pruning Strategy for Language Models (2023)

Abstract: The pruning objective has recently extended beyond accuracy and sparsity to robustness in language models. Despite this, existing methods struggle to enhance robustness against adversarial attacks when continually increasing model sparsity and require a retraining process. As humans step into the era of large language models, these issues become increasingly prominent. This paper proposes that the robustness of language models is proportional to the extent of pre-trained knowledge they encompass. Accordingly, we introduce a post-training pruning strategy designed to faithfully replicate the embedding space and feature space of dense language models, aiming to conserve more pre-trained knowledge during the pruning process. In this setup, each layer’s reconstruction error not only originates from itself but also includes cumulative error from preceding layers, followed by an adaptive rectification. Compared to other state-of-art baselines, our approach demonstrates a superior balance between accuracy, sparsity, robustness, and pruning cost with BERT on datasets SST2, IMDB, and AGNews, marking a significant stride towards robust pruning in language models.

## introduction pruning is a widely recognized compression method employed to decrease the model size and accelerate model inference @xcite @xcite . in the age of large language models @xcite @xcite @xcite @xcite , the necessity of pruning has increased because it greatly reduces deployment costs @xcite . in addition to the significant computation cost, the robustness of language models has emerged as a crucial factor that demands attention. this is primarily because models need to remain resilient against adversarial attacks, even in challenging real-world circumstances @xcite . therefore, exploring robust pruning strategies against adversarial attacks in language models could potentially yield a substantial impact @xcite . recent research has extended the pruning of language models beyond accuracy and sparsity, with an emphasis on the trade-off between accuracy, sparsity, robustness and cost @xcite @xcite . @xcite propose a joint optimization objective to guide the pruning and adversarial training simultaneously. their approach views the identified subnetworks as robust tickets, which can be trained as normal and offer enhanced robustness. despite achieving state-of-the-art results on target datasets, these methods still display vulnerabilities, as evidenced by a significant gap between metrics of clean accuracy foot_0 and accuracy under attack. moreover, the performance also rapidly declines when sparsity exceeds a moderate level. expanding on their work, @xcite propose using robust early-bird tickets to reduce the computational cost from adversarial training. however, they face similar challenges regarding the trade-off between robustness and sparsity. in summary, existing robust pruning works often demonstrate limited sparsity, insufficient robustness, and expensive cost, indicating the ongoing challenge of the balance between accuracy and the other three aspects. to address this challenge, this paper investigates why language models are susceptible to adversarial attacks. @xcite @xcite . previous studies have indicated that language models frequently capitalize on biases and artifacts inherent in datasets as predictive shortcuts, which impedes reasoning ability and skills to develop advanced semantic comprehension. @xcite @xcite . this reliance leads to a more severe loss of pre-trained knowledge during the pruning process. furthermore, the adversarial samples in natural language processing (nlp) are crafted by replacing components of sentences with semantically similar counterparts, thereby retaining high semantic similarity in the entire sentence @xcite @xcite . in this way, language models that depend on spurious features from particular words can not defend against adversarial attacks constructed by replacing those words with semantically similar alternatives. to put it more plainly, this primarily stems from the fact that, without pre-trained knowledge, the sparse language model treats the substitute word simply as an integer identifier. based on the above observation, we explore the following questions in this paper: question 1. what is the core to defend against adversarial attacks for sparse language models? this paper proposes that the robustness of sparse language models is directly proportional to the amount of pre-trained knowledge retained after pruning. intuitively, the robustness of a sparse language model is fundamentally tied to its capability to distill advanced semantic features from input sentences. this capability is largely established during the pre-training phase of dense language models, emphasizing the pivotal role of acquired semantic knowledge. the extensive experiments well support our statement. question 2. how can we efficiently prevent the loss of pre-trained knowledge in pruning to preserve or even enhance robustness? previous research has demonstrated that pruning exacerbates the model's dependency on spurious features @xcite . we further confirm that traditional pruning methods lead to a considerable loss of pre-trained knowledge and poor robustness. to prevent the above things, we propose a pruning approach that minimizes damage to the embedding space and feature space of dense language models, striving to replicate the features in each layer completely. specifically, for each layer, we iteratively eliminate a single weight at a time and counterbalance the loss by updating the remaining weights based on the hessian matrix. in this setup, the reconstruction error at each layer arises not only from its own layer but also incorporates the accumulated error from preceding layers. this is achieved by adaptively updating the pruning-dependent information in accordance with the sparse output generated by previous layers. concurrently, there's an ongoing effort to correct these errors collectively. moreover, our method, being a post-training approach, is cost-effective for current language models, as it circumvents rigorous retraining processes. extensive experiments show that our approach achieves a better trade-off between accuracy, sparsity, robustness, and pruning cost in sst2, agnews, and imdb compared with other state-of-art methods. ## methodology this section proposes a pruning method for language models that can better balance accuracy, sparsity, robustness, and pruning cost. figure 1 depicts the architecture of this method. ## experiments we first compare our method against several baseline methods, assessing accuracy, robustness, sparsity, and cost. then, an ablation study is performed to elucidate the contributions of each part in our method. finally, we augment our core findings with additional experiments and analyses to further illuminate our method. ## conclusion in this paper, we investigate the application of robust pruning methods for language models. we propose an adaptive pruning method and place a special emphasis on replicating the embedding and feature space of dense models to preserve as much pre-trained knowledge as possible. the effectiveness of this approach is confirmed through a series of experiments conducted across various tasks.