Super Tickets in Pre-Trained Language Models: From Model Compression to Improving Generalization (2021)

Abstract: The Lottery Ticket Hypothesis suggests that an over-parametrized network consists of ”lottery tickets”, and training a certain collection of them (i.e., a subnetwork) can match the performance of the full model. In this paper, we study such a collection of tickets, which is referred to as ”winning tickets”, in extremely over-parametrized models, e.g., pre-trained language models. We observe that at certain compression ratios, the generalization performance of the winning tickets can not only match but also exceed that of the full model. In particular, we observe a phase transition phenomenon: As the compression ratio increases, generalization performance of the winning tickets first improves then deteriorates after a certain threshold. We refer to the tickets on the threshold as ”super tickets”. We further show that the phase transition is task and model dependent — as the model size becomes larger and the training data set becomes smaller, the transition becomes more pronounced. Our experiments on the GLUE benchmark show that the super tickets improve single task fine-tuning by 0.9 points on BERT-base and 1.0 points on BERT-large, in terms of task-average score. We also demonstrate that adaptively sharing the super tickets across tasks benefits multi-task learning.

## introduction the lottery ticket hypothesis (lth, @xcite ) suggests that an over-parameterized network consists of "lottery tickets", and training a certain collection of them (i.e., a subnetwork) can 1) match the performance of the full model; and 2) outperform randomly sampled subnetworks of the same size (i.e., "random tickets"). the existence of such a collection of tickets, which is usually referred to as "winning tickets", indicates the potential of training a smaller network to achieve the full model's performance. lth has been widely explored in across various fields of deep learning @xcite @xcite @xcite . aside from training from scratch, such winning tickets have demonstrated their abilities to transfer across tasks and datasets @xcite @xcite . in natural language processing, @xcite ; @xcite have shown existence of the winning tickets in pre-trained language models. these tickets can be identified when fine-tuning the pre-trained models on downstream tasks. as the pre-trained models are usually extremely overparameterized (e.g., @xcite , @xcite , @xcite ), previous works mainly focus on searching for a highly compressed subnetwork that matches the performance of the full model. however, behavior of the winning tickets in lightly compressed subnetworks is largely overlooked. in this paper, we study the behavior of the winning tickets in pre-trained language models, with a particular focus on lightly compressed subnetworks. we observe that generalization performance of the winning tickets selected at appropriate compression ratios can not only match, but also exceed that of the full model. in particular, we observe a phase transition phenomenon (figure 1 ): the test accuracy improves as the compression ratio grows until a certain threshold (phase i); passing the threshold, the accuracy deteriorates, yet is still better than that of the random tickets (phase ii). in phase iii, where the model is highly compressed, training collapses. we refer to the set of winning tickets selected on that threshold as "super tickets". we interpret the phase transition in the context of trade-offs between model bias and variance @xcite . it is well understood that an expressive model induces a small bias, and a large model induces a large variance. we classify the tickets into three categories: non-expressive tickets, lightly expressive tickets, and highly expressive tickets. the full model has a strong expressive power due to over-parameterization, so that its bias is small. yet its variance is relatively large. in phase i, by removing non-expressive tickets, variance of the selected subnetwork reduces, while model bias remains unchanged and the expressive power sustains. accordingly, generalization performance improves. we enter phase ii by further increasing the compression ratio. here lightly expressive tickets are pruned. consequently, model variance continues to decrease. however, model bias increases and overturns the benefit of the reduced variance. lastly for phase iii, in the highly compressed region, model bias becomes notoriously large and reduction of the variance pales. as a result, training breaks down and generalization performance drops significantly. we conduct systematic experiments and analyses to understand the phase transition. our experiments on multiple natural language understanding (nlu) tasks in the glue @xcite benchmark show that the super tickets can be used to improve single task fine-tuning by 0.9 points over bertbase @xcite and 1.0 points over bertlarge, in terms of task-average score. moreover, our experiments show that the phase transition phenomenon is task and model dependent. it becomes more pronounced as a larger model is used to fit a task with less training data. in such a case, the set of super tickets forms a compressed network that exhibits a large performance gain. the existence of super tickets suggests potential benefits to applications, such as multi-task learning (mtl). in mtl, different tasks require different capacities to achieve a balance between model bias and variance. however, existing methods do not specifically balance the bias and variance to accommodate each task. in fact, the finetuning performance on tasks with a small dataset is very sensitive to randomness. this suggests that model variance in these tasks are high due to overparameterization. to reduce such variance, we propose a tickets sharing strategy. specifically, for each task, we select a set of super tickets during single task fine-tuning. then, we adaptively share these super tickets across tasks. our experiments show that tickets sharing improves mtl by 0.9 points over mt-dnn base @xcite and 1.0 points over mt-dnn large , in terms of task-average score. tickets sharing further benefits downstream fine-tuning of the multi-task model, and achieves a gain of 1.0 task-average score. in addition, the multi-task model obtained by such a sharing strategy exhibits lower sensitivity to randomness in downstream fine-tuning tasks, suggesting a reduction in variance. we summarize our contributions as follows: ## background we briefly introduce the transformer architecture and the lottery ticket hypothesis. ## finding super tickets we identify winning tickets in bert through structured pruning of attention heads and feed-forward layers. specifically, in each transformer layer, we associate mask variables ξ h to each attention head and ν to the ffn @xcite : here, we set ξ h , ν ∈ {0, 1}, and a 0 value indicates that the corresponding structure is pruned. we adopt importance score @xcite ) as a gauge for pruning. in particular, the importance score is defined as the expected sensitivity of the model outputs with respect to the mask variables. specifically, in each transformer layer, the importance score is closely tied to expressive power. a low importance score indicates that the corresponding structure only has a small contribution towards the output. such a structure has low expressive power. on the contrary, a large importance score implies high expressive power. we compute the importance scores for all the mask variables in a single backward pass at the end of fine-tuning. we perform one-shot pruning of the same percent of heads and feed-forward layers with the lowest importance scores. we conduct pruning multiple times to obtain subnetworks, or winning tickets, at different compression ratios. we adopt the weight rewinding technique in @xcite : we reset the parameters of the winning tickets to their values in the pre-trained weights, and subsequently fine-tune the subnetwork with the original learning rate schedule. the super tickets are selected as the winning tickets with the best rewinding validation performance. ## multi-task learning with tickets sharing in multi-task learning, the shared model is highly over-parameterized to ensure a sufficient capacity for fitting individual tasks. thus, the multi-task model inevitably exhibits task-dependent redundancy when being adapted to individual tasks. such redundancy induces a large model variance. we propose to mitigate the aforementioned model redundancy by identifying task-specific super tickets to accommodate each task's need. specifically, when viewing an individual task in isolation, the super tickets can tailor the multi-task model to strike an appealing balance between the model bias and variance (recall from section 3 that super tickets retain sufficient expressive power, yet keep the model variance low). therefore, we expect that deploying super tickets can effectively tame the model redundancy for individual tasks. given the super tickets identified by each task, we exploit the multi-task information to reinforce fine-tuning. specifically, we propose a tickets sharing algorithm to update the parameters of the multitask model: for a certain network structure (e.g., an attention head), if it is identified as super tickets by multiple tasks, then its weights are jointly updated by these tasks; if it is only selected by one specific task, then its weights are updated by that task only; otherwise, its weights are completely pruned. see figure 2 for an illustration. in more detail, we denote the weight parameters in the multi-task model as θ. suppose there are n tasks. for each task @xmath0, . . . , n }, we denote as the collection of the mask variables, where ℓ is the layer index and h is the head index. then the parameters to be updated in task i are denoted as θ @xmath1, ω i ), where m (•, ω i ) masks the pruned parameters according to ω i . we use stochastic gradient descenttype algorithms to update θ i . note that the taskshared and task-specific parameters are encoded by the mask variable ω i . the detailed algorithm is given in algorithm 1. tickets sharing has two major difference compared to sparse sharing @xcite initialize the super tickets for task i: θ @xmath2, ω i ). ## domain adaptation to demonstrate that super tickets can quickly generalize to new tasks/domains, we conduct few-shot domain adaptation on out-of-domain nli datasets. ## 8 analysis sensitivity to random seed. to better demonstrate that training with super tickets effectively reduces model variance, we evaluate models' sensitivity to changes in random seeds during single task fine-tuning and multi-task downstream finetuning. in particular, we investigate fitting small tasks with highly over-parametrized models (variance is often large in these models, see section 5 and 6). as shown in ## discussion structured lottery tickets. lth hypothesizes that a subset of unstructured parameters can be trained to match the full model's performance. instead, we question whether a subset of structured weight matrices, e.g., ffn layers and attention heads, can also be trained to match the full model's performance. this question is more practically important than the unstructured one: training and inference on structured matrices are better optimized for hardware acceleration. our results give a positive answer to this question, while previous works show that the structured tickets do not exist in highly compressed models @xcite . searching better generalized super tickets. we select winning tickets according to the sensitivity of the model outputs with respect to the mask variables of each structure @xcite , as this measure is closely tied to the structure's expressive power (section 3). in addition, we conduct an one-shot pruning for computational simplicity. we leave other importance measures and pruning schedules, which may help identifying better generalized super tickets, for future works @xcite @xcite @xcite . searching super tickets efficiently. determining the compression ratio of the super tickets requires rewinding models at multiple sparsity levels. to leverage super tickets in practice, a potential direction of research is to find heuristics to determine this ratio prior or early-on in training. we leave this for future works. ## conclusion we study the behaviors of the structured lottery tickets in pre-trained bert. we observe that the generalization performance of the winning tickets exhibits a phase transition phenomenon, suggesting pruning can improve generalization when models are lightly compressed. based on the observation, we further propose a tickets sharing strategy to improve multi-task fine-tuning. our analysis paves the way for understanding the connection between model compression and generalization.