Robust Lottery Tickets for Pre-trained Language Models (2022)

Abstract: Recent works on Lottery Ticket Hypothesis have shown that pre-trained language models (PLMs) contain smaller matching subnetworks(winning tickets) which are capable of reaching accuracy comparable to the original models. However, these tickets are proved to be notrobust to adversarial examples, and even worse than their PLM counterparts. To address this problem, we propose a novel method based on learning binary weight masks to identify robust tickets hidden in the original PLMs. Since the loss is not differentiable for the binary mask, we assign the hard concrete distribution to the masks and encourage their sparsity using a smoothing approximation of L0 regularization. Furthermore, we design an adversarial loss objective to guide the search for robust tickets and ensure that the tickets perform well bothin accuracy and robustness. Experimental results show the significant improvement of the proposed method over previous work on adversarial robustness evaluation.

## introduction large-scale pre-trained language models (plms), such as bert @xcite , roberta @xcite and t5 @xcite have achieved great success in the field of natural language processing. as more transformer layers are stacked with larger self-attention blocks, the complexity of plms increases rapidly. due to the over-parametrization of plms, some transformer heads and even layers can be pruned without significant losses in performance @xcite @xcite . the lottery ticket hypothesis suggests an overparameterized network contains certain subnetworks (i.e., winning tickets) that can match the performance of the original model when trained in isolation @xcite . chen et al. (2020); prasanna et al. (2020) also find these winning tickets exist in plms. @xcite prune bert in an unstructured fashion and obtain winning tickets at sparsity from 40% to 90%. @xcite aim at finding structurally sparse tickets for bert by pruning entire attention heads and mlp. previous works mainly focused on using winning tickets to reduce model size and speed up training time @xcite , while little work has been done to explore more benefits, such as better adversarial robustness than the original model. as we all know, plms are vulnerable to adversarial examples that are legitimately crafted by imposing imperceptible perturbations on normal examples @xcite @xcite . recent studies have shown that pruned subnetworks of plms are even less robust than their plm counterparts @xcite . @xcite observe that when fine-tuning the pruned model again, the model yields a lower robustness. @xcite clarify the above phenomenon further: the compressed models overfit on shortcut samples and thus perform consistently less robust than the uncompressed large model on adversarial test sets. in this work, our goal is to find robust plm tickets that, when fine-tuned on downstream tasks, achieve matching test performance but are more robust than the original plms. in order to make the topology structure of tickets learnable, we assign binary masks to pre-trained weights to determine which connections need to be removed. to solve discrete optimization problem of binary masks, we assume the masks follow a hard concrete distribution (a soft version of the bernoulli distribution), which can be solved using gumbel-softmax trick @xcite . we then use an adversarial loss objective to guide the search for robust tickets and an approximate l o regularization is used to encourage the sparsity of robust tickets. robust tickets can be used as a robust substitute of original plms to fine-tune downstream tasks. experimental results show that robust tickets achieve a significant improvement in adversarial robustness on various tasks and maintain a matching accuracy. our codes are publicly available at github 1 . the main contributions of our work are summarized as follows: ## experiments we conduct several experiments to demonstrate the effectiveness of our method. we first compare the proposed method with baseline methods in terms of clean accuracy and robust evaluation. then, we perform an ablation study to illustrate the role of sparse mask learning and adversarial loss objective in our method. in addition, we try to further flesh out our method with several additional analysis experiments. following the official bert implementation @xcite , we use bert base as our backbone model for all experiments. ## discussion in this section, we study how the implementation of robust tickets affects the model's robustness. ## conclusion in this paper, we articulate and demonstrate the robust lottery ticket hypothesis for plms: the full plm contains subnetworks (robust tickets) that can achieve a better robustness performance. we propose an effective method to solve the ticket selection problem by encouraging weights that are not responsible for robustness to become exactly zero. experiments on various tasks corroborate the effectiveness of our method. we also find that pretrained weights may be a key factor affecting the robustness on downstream tasks.