Learning to Win Lottery Tickets in BERT Transfer via Task-agnostic Mask Training
Abstract
Recent studies on the lottery ticket hypothesis (LTH) show that pre-trained language models (PLMs) like BERT contain matching subnetworks that have similar transfer learning performance as the original PLM.
These subnetworks are found using magnitude-based pruning.
In this paper, we find that the BERT subnetworks have even more potential than these studies have shown.
Firstly, we discover that the success of magnitude pruning can be attributed to the preserved pre-training performance, which correlates with the downstream transferability.
Inspired by this, we propose to directly optimize the subnetwork structure towards the pre-training objectives, which can better preserve the pre-training performance.
Specifically, we train binary masks over model weights on the pre-training tasks, with the aim of preserving the universal transferability of the subnetwork, which is agnostic to any specific downstream tasks.
We then fine-tune the subnetworks on the GLUE benchmark and the SQuAD dataset.
The results show that, compared with magnitude pruning, mask training can effectively find BERT subnetworks with improved overall performance on downstream tasks.
Moreover, our method is also more efficient in searching subnetworks and more advantageous when fine-tuning within a certain range of data scarcity.
Our code is available at https://github.com/llyx97/TAMT .
Introduction
The NLP community has witnessed a remarkable success of pre-trained language models (PLMs).
After being pre-trained on unlabelled corpus in a self-supervised manner, PLMs like BERT (Devlin et al., 2019) can be fine-tuned as a universal text encoder on a wide range of downstream tasks. [Figure]
However, the growing performance of BERT is driven, to a large extent, by scaling up the model size, which hinders the fine-tuning and deployment of BERT in resource-constrained scenarios.
[-o r e G a p −3 − −
-]
At the same time, the lottery ticket hypothesis (LTH) (Frankle and Carbin, 2019) emerges as an active sub-field of model compression.
The LTH states that randomly initialized dense networks contain sparse matching subnetworks , i.e., winning tickets (WTs), that can be trained in isolation to similar test accuracy as the full model.
The original work of LTH and subsequent studies have demonstrated that such WTs do exist at random initialization or an early point of training (Frankle et al., 2019, 2020).
This implicates that it is possible to reduce training and inference cost via LTH.
Recently, Chen et al.
(2020) extend the original
LTH to the pre-training and fine-tuning paradigm, exploring the existence of matching subnetworks in pre-trained BERT.
Such subnetworks are smaller in size, while they can preserve the universal transferability of the full model.
Encouragingly, Chen et al.
(2020) demonstrate that BERT indeed contains matching subnetworks that are transferable to multiple downstream tasks without compromising accuracy.
These subnetworks are found using iterative magnitude pruning (IMP) (Han et al., 2015) on the pre-training task of masked language modeling (MLM), or by directly compressing BERT with oneshot magnitude pruning (OMP), both of which are agnostic to any specific task.
In this paper, we follow Chen et al.
(2020) to study the question of LTH in BERT transfer learning.
We find that there is a correlation, to certain extent, between the performance of a BERT subnetwork on the pre-training task (right after pruning), and its downstream performance (after fine-tuning).
As shown by Fig.
1, the OMP subnetworks significantly outperform random subnetworks at 50% sparsity in terms of both MLM loss and downstream score.
However, with the increase of model sparsity, the downstream performance and pre-training performance degrade simultaneously.
This phenomenon suggests that we might be able to further improve the transferability of BERT subnetworks by discovering the structures that better preserve the pre-training performance.
To this end, we propose to search transferable BERT subnetworks via T ask-
A gnostic M ask
T raining (TAMT), which learns selective binary masks over the model weights on pre-training tasks.
In this way, the structure of a subnetwork is directly optimized towards the pre-training objectives, which can preserve the pre-training performance better than heuristically retaining the weights with large magnitudes.
The training objective of the masks is a free choice, which can be designed as any loss functions that are agnostic to the downstream tasks.
In particular, we investigate the use of MLM loss and a loss based on knowledge distillation (KD) (Hinton et al., 2015).
To examine the effectiveness of the proposal, we train the masks on the WikiText dataset (Merity et al., 2017) for language modeling and then finetune the searched subnetworks on a wide variety of downstream tasks, including the GLUE benchmark (Wang et al., 2019) for natural language understanding (NLU) and the SQuAD dataset (Rajpurkar et al., 2016) for question answering (QA).
The empirical results show that, through mask training, we can in- deed find subnetworks with lower pre-training loss and better downstream transferability than OMP and IMP Compared with
IMP, which also involves training (the weights) on the pre-training task, mask training requires much fewer training iterations to reach the same performance. Moreover, the subnetworks found by
mask training are generally more robust
when
being fine-tuned
with
reduced
data, as
long as the training data is not extremely scarce.
In summary, our contributions are:
• We find that the pre-training performance of a BERT subnetwork correlates with its downstream transferability, which provides a useful insight for the design of methods to search transferable BERT subnetworks.
• Based on the above finding, we propose to search subnetworks by learning binary masks over the weights of BERT, which can directly optimize the subnetwork structure towards the given pre-training objective.
• Experiments on a variety of NLP tasks show that subnetworks found by mask training have better downstream performance than magnitude pruning.
This suggests that BERT subnetworks have more potential, in terms of universal downstream transferability, than existing work has shown, which can facilitate our understanding and application of LTH on BERT.
Related Work 2. The Lottery Ticket Hypothesis
The lottery ticket hypothesis (Frankle and Carbin, 2019) suggests the existence of matching subnetworks, at random initialization, that can be trained in isolation to reach the performance of the original network.
However, the matching subnetworks are found using IMP, which typically requires more training cost than the full network.
There are two remedies to overcome this problem: Morcos et al.
(2019) proposed to transfer the WT structure from source tasks to related target tasks, so that no further searching is required for the target tasks.
You et al.
(2020) draw early-bird tickets (prune the original network) at an early stage of training, and only train the subnetwork from then on.
Some recent works extend the LTH from random initialization to pre-trained initialization (Prasanna et al., 2020; Chen et al., 2020; Liang et al., 2021; Chen et al., 2021b).
Particularly, Chen et al.
(2020) find that WTs, i.e., subnetworks of the pre-trained
BERT, derived from the pre-training task of MLM using IMP are universally transferable to the downstream tasks.
The same question of transferring WTs found in pre-training tasks is also explored in the CV field by Chen et al. (2021a)
; Caron et al.
EarlyBERT (Chen et al., 2021b)
investigates drawing early-bird tickets of BERT.
In this work, we follow the question of transferring WTs and seek to further improve the transferability of BERT subnetworks.
In the literature of BERT compression, pruning (LeCun et al., 1989; Han et al., 2015) and KD (Hinton et al., 2015) are two widely-studied techniques.
BERT can be pruned in either unstructured (Gordon et al., 2020; Sanh et al., 2020; Mao et al., 2020) or structured (Michel et al., 2019;
Hou et al., 2020) ways.
Although unstructured pruning is not hardware-friendly for speedup purpose, it is a common setup in LTH, and some recent efforts have been made in sparse tensor acceleration (Elsen et al., 2020;
Tambe et al., 2020). In BERT KD, various knowledge are explored, which includes the soft-labels (Sanh et al., 2019), the hidden states (Sun et al., 2019; Hou et al., 2020; Liu et al., 2021) and the attention relations (Jiao et al., 2020), among others.
Usually, pruning and KD are combined to compress the fine-tuned BERT.
By contrast, the
LTH compresses BERT before fine-tuning.
Another way to obtain more efficient BERT with the same transferability as the original one is to pretrain a compact model from scratch.
This model can be trained either with the MLM objective (Turc et al., 2019) or using pre-trained BERT as the teacher to perform KD (Wang et al., 2020a;
Sun et al., 2020; Jiao et al., 2020).
By contrast, the LTH extracts subnetworks from BERT, which is about exposing the knowledge already learned by BERT, rather than learning new knowledge from scratch.
Compared with training a new PLM, the LTH in
BERT is still underexplored in the literature.
2. Learning Subnetwork Structure via Binary Mask Training
To make the subnetwork structure trainable, we need to back-propagate gradients through the binary masks.
This can be achieved through the straight-through estimator (Bengio et al., 2013) or drawing the mask variables from a hard-concrete distribution (Louizos et al., 2018) and then using the re-parameterization trick. Mask training has been utilized in model compression (Wang et al., 2020b; Sanh et al., 2020), and parameter-efficient training (Mallya et al., 2018; Zhao et al., 2020; Radiya-Dixit and Wang, 2020).
However, unlike these works that learn the mask for each task separately ( task-specific ), we learn the subnetwork structure on pre-training task and transfer it to multiple downstream tasks ( task-agnostic ).
Methodology 3. BERT Architecture
BERT consists of an embedding layer and L Transformer layers (Vaswani et al., 2017).
Each Transformer layer has two sub-layers: the self-attention layer and the feed-forward network (FFN).
The self-attention layer contains N h parallel attention heads and each head can be formulated as:
[Equation]
where H ∈ R | x |× d H is the input; d H and | x | are the hidden size and the length of input x , respectively.
W Q h ,K h ,V h ∈ R d H ×
d h are the query, key and value matrices, and d h = d H N h .
In practice, the matrices for different heads will be combined into three large matrices W Q,K,
V ∈ R d H × d H .
The outputs of the N h heads are then concatenated and linearly projected by W AO ∈ R d H × d H to obtain the final output of the self-attention layer.
The FFN consists of two weight matrices
W FI ∈ R d H × d
I , W FO ∈ R d I × d H with a GELU activation (Hendrycks and Gimpel, 2016) in between, where d I is the hidden dimension of FFN.
Dropout (Srivastava et al., 2014), residual connection (He et al., 2016) and layer normalization (Ba et al., 2016) are also applied following each sublayer.
Eventually, for each downstream task, a classifier is used to give the final prediction based on the output of the Transformer module.
3. Subnetwork and Magnitude Pruning
Consider a model f ( · ; θ ) with weights θ , we can obtain its subnetwork f ( · ; M ⊙ θ ) by applying a binary mask M ∈ { 0 , 1 } | θ | to θ , where ⊙ denotes element-wise multiplication.
In terms of
BERT, we extract the subnetwork from the pretrained weights θ 0 .
Specifically, we consider the matrices of the Transformer sub-layers and the word embedding matrix, i.e., θ {+
+}
= {
[Equation]
Magnitude pruning (Han et al., 2015) is initially used to compress a trained neural network by setting the low-magnitude weights to zero.
It can be conducted in two different ways: 1) Oneshot magnitude pruning (OMP) directly prunes the trained weights to target sparsity while 2) iterative magnitude pruning (IMP) performs pruning and retraining iteratively until reaching the target spar-
𝐖 𝑬𝒎𝒃 ⊙𝐌 𝑬𝒎𝒃 𝐿× 𝐖 𝑸 ⊙𝐌 𝑸
Self-Att
FFN 𝐖 𝑬𝒎𝒃 𝐿× 𝐖 𝑸 𝐖 𝑲 𝐖 𝑽 𝐖 𝑨𝑶 Self-Att 𝐖 𝑭𝑰
FFN 𝐖 𝑭𝑶 𝐖 𝑸
⊙𝐌 𝑸 𝐖
𝑸
⊙𝐌 𝑸 𝐖 𝑨𝑶
⊙𝐌 𝑨𝑶 𝐖 𝑭𝑰 ⊙𝐌 𝑭𝑰 𝐖 𝑭𝑶 ⊙𝐌
𝑭𝑶 𝓒 𝒎𝒍𝒎 𝓛 𝑚𝑙𝑚
𝓛 𝑑𝑖𝑠𝑡𝑖𝑙𝑙 𝓣 𝒊 … 𝐖 𝑬𝒎𝒃 ⊙𝐌 𝑬𝒎𝒃 𝐖 𝑸 ⊙𝐌 𝑸
Self-Att
FFN 𝐖 𝑸
⊙𝐌 𝑸 𝐖
𝑸
⊙𝐌 𝑸 𝐖 𝑨𝑶
⊙𝐌 𝑨𝑶 𝐖 𝑭𝑰 ⊙𝐌 𝑭𝑰 𝐖 𝑭𝑶 ⊙𝐌 𝑭𝑶 𝓒 𝓣 𝒊 𝓛 𝓣 𝒊
Pre-training Data 𝓣 [Figure] 𝒯 …
[Figure]
𝐿×
[Figure]
sity.
OMP and IMP are also widely studied in the literature of LTH as the method to find the matching subnetworks, with an additional operation of resetting the weights to initialization.
3. Problem Formulation: Transfer BERT Subnetwork
As depicted in Fig. 2, given N T downstream tasks
T =
{T i } N T i =1 , the subnetwork f (cid:16) · ; M ⊙ θ 0 , C T i 0 (cid:17) is fine-tuned on each task, together with the randomly initialized task-specific linear classifier C T i 0 . We formulate the training algorithm for task T i as a function A T i t (cid:16) f (cid:16) · ; M ⊙ θ 0 , C T i 0 (cid:17)(cid:17) (e.g., Adam or
SGD), which trains the model for t steps and produces f (cid:16) · ; M ⊙ θ t , C T i t (cid:17) .
After fine-tuning, the model is evaluated against the metric E T i ( f ( · ; M ⊙ θ t , C T i t ))
(e.g., Accuracy or F1) for task T i .
In this work, we focus on finding a BERT subnetwork, that maximally preserves the overall downstream performance given a particular sparsity S , especially at the sparsity that magnitude pruning performs poorly.
This can be formalized as:
[Equation]
where ∥ M ∥ 0 and | θ 0 | are the L 0 norm of the mask and the total number of model weights respectively.
3. Task-agnostic Mask Training 3.4. Mask Training with Binarization and Gradient Estimation
In order to learn the binary masks, we adopt the technique for training binarized neural networks (Hubara et al., 2016), following Zhao et al.
Mallya et al.
This technique involves mask binarization in the forward pass and gradient estimation in the backward pass.
As shown in Fig.
2, each weight matrix W ∈
R d in × d out is associated with a binary mask M ∈ { 0 , 1 } d in × d out , which is derived from a real-valued matrix M ∈ R d in ×
d out via binarization:
[Equation] [Equation] where ϕ is the threshold that controls the sparsity.
In the forward pass of a subnetwork, W ⊙ M is used in replacement of the original weights W .
Since M i,j are discrete variables, the gradient signals cannot be back-propagated through the binary mask.
We therefore use the straight-through estimator (Bengio et al., 2013) to approximate the gradients and update the real-valued mask:
[Equation]
where L is the loss function and η is the learning rate.
In other words, the gradients of M is estimated using the gradients of M .
In the process of mask training, all the original weights are frozen.
3.4. Mask Initialization and Sparsity Control
The real-valued masks can be initialized in various forms, e.g., random initialization.
Considering that magnitude pruning can preserve the pre-training knowledge to some extent, and OMP is easy to implement with almost zero computation cost, we directly initialize M using OMP:
M i
[Equation]
j
[Equation]
where M OMP is the binary mask derived from
OMP and α ≥ 1 is a hyper-parameter.
In this way, the weights with large magnitudes will be retained at initialization according to Eq. 3, because the corresponding M i,j = α × ϕ ≥ ϕ .
In practice, we perform OMP over the weights locally based on the given sparsity, which means the magnitudes are ranked inside each weight matrix.
As M being updated, some of its entries with zero initialization will gradually surpass the threshold, and vice versa.
If the threshold ϕ is fixed throughout training, there is no guarantee that the binary mask will always satisfy the given sparsity.
Therefore, we rank M ij according to their absolute values during mask training, and dynamically adjust the threshold to satisfy the sparsity constraint.
3.4. Mask Training Objectives
We explore the use of two objectives for mask training, namely the MLM loss and the KD loss.
The MLM is the original task used in BERT pretraining.
It randomly replaces a portion of the input tokens with the [MASK] token, and requires the model to reconstruct the original tokens based on the entire masked sequence.
Concretely, the MLM objective is computed as cross-entropy loss on the predicted masked tokens.
During MLM learning, we allow the token classifier (i.e., the C mlm in Fig. 2) to be trainable, in addition to the masks.
In KD, the compressed model (student) is trained with supervision from the original model (teacher).
Under our framework of mask training, the training signal can also be derived from the unpruned
BERT.
To this end, we design the KD objective by encouraging the subnetwork to mimic the representations of the original BERT, which is shown to be a useful source of knowledge in BERT KD (Sun et al., 2019;
Hou et al., 2020).
Specifically, the distillation loss is formulated as the cosine similarity between the teacher’s and student’s representations:
L distill =
L | x | L (cid:88) l = | x | (cid:88) i = (1 − cos (cid:0) H Tl,i , H Sl,i (cid:1) ) (6) where H l,i is the hidden state of the i th token at the l th layer; T and S denote the teacher and student respectively; cos( · , · ) is the cosine similarity.
Experiments 4. Experimental Setups 4.1. Models
We examine two PLMs from the BERT family, i.e., BERT BASE (Devlin et al., 2019) and
RoBERTa BASE (Liu et al., 2019).
They have basically the same structure, while differ in the vocabulary size, which results in approximately 110M and 125M parameters respectively.
The main results of Section 4.2.1 study both two models.
For the analytical studies, we only use BERT BASE .
4.1. Baselines, Datasets and Evaluation
We compare our mask training method with IMP,
OMP as well as subnetworks with random structures.
Following Chen et al.
(2020), we use the
MLM loss during IMP training.
For TAMT, we consider three variants, namely TAMT-MLM that uses
MLM as training objective, TAMT-KD that uses the KD objective (Eq. 6), and TAMT-MLM+KD that equally combines MLM and KD.
We build our pre-training set using the WikiText103 dataset (Merity et al., 2017) for language modeling.
For downstream fine-tuning, we use six datasets, i.e., CoLA, SST-2, RTE, MNLI, MRPC and STS-B from the GLUE benchmark for NLU and the SQuAD v1.1 dataset for QA.
Evaluations are conducted on the dev sets.
For the downstream tasks, we follow the standard evaluation metrics (Wang et al., 2019).
For the pretraining tasks, we calculate the MLM and KD loss on the dev set of WikiText-103.
More information about the datasets and evaluation metrics can be found in Appendix B.1.
4.1. Implementation Details
Both TAMT and IMP are conducted on the pretraining dataset.
For mask training, we initialize the mask using OMP as described in Section 3.4.2.
We also provide a comparison between OMP and random initialization in Section 4.2.4.
The initial threshold ϕ and α are set to 0 . 01 and 2 respectively, [Figure]
[Figure]
which work well in our experiments. For IMP, we increase the sparsity by 10% every 1/10 of total training iterations, until reaching the target sparsity, following Chen et al.
Every pruning operation in IMP is followed by resetting the remaining weights to θ 0 .
In the fine-tuning stage, all the subnetworks and the full PLMs are trained using the same set of hyper-parameters unless otherwise specified.
For TAMT, IMP and random pruning, we generate three subnetworks with different seeds, and the result of each subnetwork is also averaged across three runs, i.e., the result of every method is the average of nine runs in total.
For OMP, we can only generate one subnetwork, which is fine-tuned across three runs.
More implementation details and computing budgets can be found in Appendix B.2.
4. Results and Analysis 4.2. Main Results
Fig.
3 and Fig. 4 present the downstream performance of BERT and RoBERTa subnetworks, respectively.
We can derive the following observations:
There is a clear gap between random subnetworks and the other ones found with certain inductive bias.
At 50% sparsity for BERT and 30% for RoBERTa, all the methods, except for “Rand”, maintain 90% of the full model’s overall performance.
As sparsity grows, the OMP subnetworks degrade significantly.
IMP, which is also based on magnitude, exhibits relatively mild declines.
TAMT further outperforms IMP with perceivable margin.
For BERT subnetworks, the performance of TAMT variants are close to each other, which have advantage over IMP across 60% ∼ 90% sparsity. When it comes to RoBERTa, the performance of TAMT-KD is undesirable at 70% ∼ 90% sparsity, which only slightly outperforms IMP.
In comparison, TAMT-MLM consistently surpasses IMP and TAMT-KD on RoBERTa. Combining MLM and KD leads to comparable average performance as TAMT-MLM for BERT, while slightly improves over TAMT-MLM for [Figure] RoBERTa.
This suggests that the two training objectives could potentially benefit, or at least will not negatively impact each other.
In Section 4.2.2, we will show that the MLM and KD objectives indeed exhibit certain consistency.
At 90% sparsity, all the methods perform poorly, with average scores approximately half of the full model.
On certain tasks like CoLA, RTE and MRPC, drastic performance drop of all methods can even be observed at lower sparsity (e.g., 60% ∼ 80% ).
This is probably because the number of training data is too scarce in these tasks for sparse PLMs to perform well.
However, we find that the advantage of TAMT is more significant within a range of data scarsity, which will be discussed in Section 4.2.6.
We also note that RoBERTa, although outperforms BERT as a full model, is more sensitive to task-agnostic pruning.
A direct comparison between the two PLMs is provided in Appendix A.2.
4.2. The Effect of Pre-training Performance
As we discussed in Section 1, our motivation of mask training is to improve downstream transferability by preserving the pre-training performance.
To examine whether the effectiveness of TAMT is indeed derived from the improvement on pretraining tasks, we calculate the MLM/KD dev loss for the subnetworks obtained from the mask training process, and associate it with the downstream performance. The results are shown in Fig.
5, where the "Avg Score" includes CoLA, SST-2, MNLI, STS-B and SQuAD. In the following sections, we also mainly focus on these five tasks.
We can see from Fig. 5 that:
There is a positive correlation between the pretraining and downstream performance, and this trend can be observed for subnetworks across different sparsities.
Compared with random pruning, the magnitude pruning subnetworks and TAMT subnetworks reside in an area with lower MLM/KD loss and higher downstream score at 50% sparsity.
As sparsity increases, OMP subnetworks gradually move from the upper-left to the lower-right area of the plots.
In comparison, IMP is better at preserving the pre-training performance, even though it is not deliberately designed for this purpose.
For this reason, hypothetically, the downstream performance of IMP is also better than OMP.
TAMT-MLM and TAMT-KD have the lowest
MLM and KD loss respectively, which demonstrates that the masks are successfully optimized towards the given objectives.
As a result, the downstream performance is also elevated from the OMP initialization, which justifies our motivation.
Moreover, training the mask with KD loss can also optimize the performance on MLM, and vice versa, suggesting that there exists some consistency between the objectives of MLM and KD.
It is also worth noting that the correlation between pre-training and fine-tuning performance is not ubiquitous.
For example, among the subnetworks of OMP, IMP and TAMT at 50% sparsity, the decrease in KD/MLM loss produces little or no downstream improvement; at 60% ∼ 80% sparsity, OMP underperforms random pruning in
MLM, while its downstream performance is better.
These phenomenons suggest that some properties about the BERT winning tickets are still not wellunderstood by us. [Figure]
[Figure]
4.2. The Effect of Pre-training Cost
We have shown that mask training is more effective than magnitude pruning.
Now let us take a closer look at the results of TAMT and IMP with different iterations of pre-training, to evaluate their efficiency in subnetwork searching.
For TAMT, we directly obtain the subnetworks from varied pretraining iterations.
For IMP, we change the pruning frequency to control the number of training iterations before reaching the target sparsity.
Fig.
6 presents the downstream results with increased pre-training iterations and time.
We can see that for all the methods, the fine-tuning performance steadily improves as pre-training proceeds.
Along this process, TAMT advances at a faster pace, reaching the best score achieved by IMP with 8 .
× fewer iterations and 8 .
7 × fewer time.
This indicates that directly optimizing the pre-training objectives is more efficient than the iterative process of weight pruning and re-training.
[Figure]
O M P T A M T - K D - 1 T A M T - K D - 2 T A M T - K D - 3 T A M T - M L M - 1 T A M T - M L M - 2 T A M T - M L M - 3 I M P - 1 I M P - 2 I M POMP TAMT-KD-1 TAMT-KDTAMT-KD-3 TAMT-MLMTAMT-MLM-2 TAMT-MLMIMP-1 IMPIMP- 0.
M a s k S
i m il a r i t
y
4.2. The Effect of Mask Initialization
In the main results, we use OMP as the default initialization, in order to provide a better start point for TAMT.
To validate the efficacy of this setting, we compare OMP initialization with random initialization.
Concretely, we randomly sample some entries of the real-valued masks to be zero, according to the given sparsity, and use the same α and ϕ for the non-zero entries as in Eq. 5.
The results are shown in Fig.
We can see that, for random initialization, TAMT can still steadily improve the downstream performance as pre-training proceeds.
However, the final results of TAMT-MLM/KD (Rand_init) are significantly worse than TAMT-MLM/KD, which demonstrates the necessity of using OMP as initialization.
4.2. Similarity between Subnetworks
The above results show that the subnetworks found by different methods perform differently.
We are therefore interested to see how they differ in the mask structure.
To this end, we compute the similarity between OMP mask and the masks derived during the training of TAMT and IMP.
Following Chen et al.
(2020), we measure the Jaccard similarity between two binary masks M i and M j as | M i ∩ M j | | M i ∪ M j | , and the mask distance is defined as [Equation]
From the results of Fig. 8, we can find that:
) With different objectives TAMT produces dif- ferent mask structures.
The KD loss results in masks in the close proximity of OMP initialization, while the MLM masks deviate away from OMP.
Among the four methods, IMP and TAMT-MLM have the highest degree of dissimilarity, despite the fact that they both involve MLM training.
[Figure] Although IMP, TAMT-KD and TAMT-MLM are different from each other in terms of subnetwork structure, all of them clearly improves over the
OMP baseline.
Therefore, we hyphothesize that the high-dimensional binary space { 0 , 1 } | θ | might contain multiple regions of winning tickets that are disjoint with each other.
Searching methods with different inductive biases (e.g., mask training versus pruning and KD loss versus MLM loss) are inclined to find different regions of interest.
4.2. Results of Reducing Fine-tuning Data
To test the fine-tuning results with reduced data, we select four tasks (CoLA, SST-2, MNLI and
SQuAD) with the largest data sizes and shrink them from the entire training set to 1,000 samples.
Fig.
9 summarizes the results of subnetworks found using different methods, as well as results of full BERT as a reference.
We can see that the four datasets present different patterns.
For MNLI and
SQuAD, the advantage of TAMT first increases and then decreases with the reduction of data size.
The turning point appears at around 10,000 samples, after which the performance of all methods, including the full BERT, degrade drastically (note that the horizontal axis is in log scale).
For SST-2, the performance gap is enlarged continuously until we have only 1,000 data.
With regard to CoLA, the gap between TAMT and IMP shrinks as we reduce the data size, which is not desirable.
However, a decrease in the gap between full BERT and IMP is also witnessed when the data size is reduced under 5,000 samples.
This is in part because the Mcc of IMP is already quite low even with the entire training set, and thus the performance decrease of IMP is limited compared with TAMT.
However, the results on CoLA, as well as the results on MNLI and SQuAD with less than 10,000 samples, also suggest an inherent difficulty of learning with limited data for subnetworks at high sparsity, which is also discussed in the main results.
Conclusions
In this paper, we address the problem of searching transferable BERT subnetworks.
We first show that there exist correlations between the pre-training performance and downstream transferablility of a subnetwork.
Motivated by this, we devise a subnetwork searching method based on task-agnostic mask training (TAMT).
We empirically show that
TAMT with MLM loss or KD loss achieve better pre-training and downstream performance than the magnitude pruning, which is recently shown to be successful in finding universal BERT subnetworks.
TAMT is also more efficient in mask searching and produces more robust subnetworks when being fine-tuned within a certain range of data scarsity.
Limitations and Future Work
Under the framework of TAMT, there are still some unsolved challenges and interesting questions worth studying in the future work: First, we focus on unstructured sparsity in this work, which is hardware-unfriendly for speedup purpose. In future work, we are interested in investigating TAMT with structured pruning or applying unstructured BERT subnetworks on hardware platforms that support sparse tensor acceleration (Elsen et al., 2020;
Tambe et al., 2020).
Second, despite the overall improvement achieved by TAMT, it fails at extreme sparsity or when the labeled data for a task is too scarce. Therefore, another future direction is to further promote the performance of universal PLM subnetworks on these challenging circumstances. To achieve this goal, thirdly, a feasible way is to explore other task-agnostic training objectives for TAMT beyond MLM and hidden state KD, e.g., self-attention KD (Jiao et al., 2020) and contrastive learning (Gao et al., 2021).
An in-depth study on the selection of TAMT training objective might further advance our understanding of TAMT and the LTH of BERT.
Acknowledgments
This work was supported by the National Natural Science Foundation of China under Grants 61976207 and 61906187.
A More Results and Analysis A. Single Task Downstream Performance of OMP and Random Pruning
In Fig.
1 of the main body of paper, we show that the pre-training and overall downstream performance of OMP, as well as the gap between “OMP" and “Rand", degrade simultaneously as sparsity increases.
The detailed results of each downstream task are presented in Fig.
As we can see, the general pattern for every task is similar, with the exception that the gap between “OMP" and “Rand" slightly increases before high sparsity on tasks RTE, MNLI and SQuAD.
A. Comparison Between BERT and RoBERTa Subnetworks
In the main results of Fig. 3 and Fig.
4, we compare the fine-tuning performance of subnetworks of the same PLM but found using different methods.
In this section, we give a comparison between subnetwords of BERT BASE and RoBERTa BASE .
As shown in Fig.
11, RoBERTa consistently outperforms BERT as a full model.
However, as we prune the pre-trained weights accroding to the magnitudes, the performance of RoBERTa declines more sharply than BERT, leading to worse results of
RoBERTa subnetworks when crossing a certain sparsity threshold.
This phenomenon suggests that, compared with BERT, RoBERTa is less robust to task-agnostic magnitude pruning.
More empirical and theoretical analysis are required to understand the underlying reasons.
A. Pre-training Performance and Single Task Downstream Performance
The relation between pre-training performance and overall downstream performance is illustrated in
Fig.
Here in this appendix, we provide the detailed results about each single downstream task, as shown in Fig.
12 and Fig. 13.
As we can see, the pattern in each single task is general the same as we discussed in Section 4.2.2.
When the model sparsity is higher than 50% , TAMT promotes the performance of OMP in terms of both pre-training tasks and downstream tasks, and improves over
IMP with perceivable margin.
As shown in Fig.
3 and Fig. 4 of the main paper, both IMP and TAMT display no obvious improvement over OMP on MRPC and RTE (but no degradation as well).
Therefore, we do not report the comparison on these two datasets.
A. Pre-training Iteration and Single Task Downstream Performance
In Fig.
6, we show the overall downstream performance at 70% sparsity with the increase of mask training iterations.
Here, we report the results of each single downstream task from 60% ∼ 80% sparsities, which are shown in Fig.
14, Fig.and Fig.
We can see that: 1) The single task performance of both TAMT-MLM and TAMT-KD grows faster than IMP at 60% and 70% sparsity, with the only exception of STS-B, where TAMTMLM and IMP are comparable in the early stage of pre-training.
The MLM and KD objectives are good at different sparsity levels and different tasks.
TAMT-KD performs the best at 60% sparsity, sur-
[Table] [Table] passing TAMT-MLM on all the five tasks.
In contrast, TAMT-MLM is better at higher sparsities. 3)
At 80% sparsity, the searching efficiency of the KD objective is not desirable, which requires more pretraining steps to outperform IMP on CoLA, STS-B, SQuAD and the overall performance.
However, the advantage of TAMT-MLM is still obvious at 80% sparsity.
A. Subnetwork Similarity at Different Sparsities
In Section 4.2.5, we analyse the similarity between subnetworks at 70% sparsity.
In Fig.
17, we present additional results of subnetworks at different sparsities.
We can see that the general pattern, as discussed in Section 4.2.5, is the same across 60% , 70% and 80% sparsities.
However, as sparsity grows, different searching methods becomes more distinct from each other.
For instance, the similarity between TAMT-MLM and IMP subnetworks decreases from 0.75 at 60% sparsity to less than 0.6 at 80% sparsity.
This is understandable because the higher the sparsity, the lower the probability that two subnetworks will share the same weight.
B More Information about Experimental Setups B. Datasets and Evaluation
For pre-training, we adopt the WikiTextdataset 1 for language modeling.
WikiText- is a collection of articles on Wikipedia and has over 100M tokens.
Such data scale is relatively small for PLM pre-training. However, we find that it is sufficient for mask training and IMP to discover subnetworks with perceivable downstream improvement. For the downstream tasks, we use six datasets from the GLUE benchmark and the SQuAD v1. dataset 2 .
The GLUE benchmark is intended to train, evaluate, and analyze NLU systems.
Our experiments include the tasks of CoLA for linguistic acceptability, SST-2 for sentiment analysis, RTE and MNLI for natural language inference, MRPC and STS-B for semantic matching/similarity.
The
SQuAD dataset is for the task of question answering.
It consists of questions posed by crowdworkers on a set of Wikipedia articles.
Tab. 1 summarizes the dataset statistics and evaluation metrics.
All the datasets are in English language.
B. Implementation Details
The hyper-parameters for pre-training and finetuning are shown in Tab. 1.
The pre-training setups of IMP basically follow (Chen et al., 2020), except for the number of training epochs, because we use different pre-training datasets.
Since we aim at finding universal PLM subnetworks that are agnostic to the downstream tasks, we do not perform hyper-parameter search for TAMT based on the downstream performance.
The pre-training hyper-
[Table]
parameters in Tab. 1 are determined as they can guarantee stable convergence on the pre-training tasks.
For fair comparison between TAMT and IMP, we control the number of pre-training iterations (i.e., the number of gradient descent steps) to be the same.
Considering that the IMP subnetworks of different sparsities are obtained from different pre-training iterations, we adjust the pre-training iterations of TAMT accordingly.
Specifically, we set the maximum number of pre-training epochs to 2 for IMP, which equals to 27.92K training iterations.
Thus, the sparsity is increased by 10% every 2.792K iterations.
Tab. 2 shows the number of pretraining iterations for IMP and TAMT subnetworks at 20% ∼ 90% sparsity.
Note that the final training iteration does not equal to 27.92K at 100% sparsity according to Tab. 2.
This is because we prune to 10% sparsity at the 0 th iteration, which follows the implementation of Chen et al.
The hyper-parameters for downstream finetuning follow the standard setups of (Wolf et al., 2020; Chen et al., 2020).
We use the same set of hyper-parameters for all the subnetworks, as well as the full models.
We perform evaluations during the fine-tuning process, and the best result is reported as the downstream performance.
Training and evaluation are implemented on
Nvidia V100 GPU.
The codes are based on the
Pytorch framework 3 and the huggingface Transformers library 4 (Wolf et al., 2020).
Tab.
3 shows the pre-training time of IMP and TAMT.
[Figure]
[Figure]
C o L A
M cc
[Figure]
M cc
[Figure]
cc
[Figure] [Figure]
cc
[Figure] [Figure] [Figure]
[Figure]
single task, with increased pre-training iterations.
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
OMP
TAMT-KDTAMT-KD-
[Table]
TAMT-KD-
[Table]
TAMT
[Table]
MLM-
[Table]
TAMT
[Table]
MLM
[Table]
TAMT
[Table]
MLM
[Table]
IMPIMP-
IMP-
[Table]
[Figure]
[Table]
[Figure]
Mask Similarity
[Figure]
TAMT
[Figure]
TAMT
[Figure]
TAMT
[Figure]
TAMT-MLM
[Figure]
TAMT
[Figure]
MLM
[Figure]
TAMT
[Figure]