SlowGuess/NAACL_2022_OCR huggingface dataset:
Yuanxin Liu $^{1,2}$ , Fandong Meng $^{3}$ , Zheng Lin $^{1,2*}$ , Peng Fu $^{1*}$ , Yanan Cao $^{1,2}$ , Weiping Wang $^{1}$ , Jie Zhou $^{3}$
$^{1}$ Institute of Information Engineering, Chinese Academy of Sciences, Beijing, China  
 $^{2}$ School of Cyber Security, University of Chinese Academy of Sciences, Beijing, China  
 $^{3}$ Pattern Recognition Center, WeChat AI, Tencent Inc, China
{liuyuanxin,linzheng,fupeng,caoyanan,wangweiping}@iie.ac.cn, {fandongmeng,withtomzhou}@tencent.com
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/1lyx97/TAMT.
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.
Joint work with Pattern Recognition Center, WeChat AI, Tencent Inc, China.\*Zheng Lin and Peng Fu are the corresponding authors.
Figure 1: Average downstream performance (left) and pre-training performance (right) of OMP and random subnetworks of  $\mathrm{BERT}_{\mathrm{BASE}}$ . See Appendix A.1 for the downstream results of each task.
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.
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
Human Language Technologies, pages 5840 - 5857
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 Task-Agnostic Mask Training (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 (Merit et al., 2017) for language modeling and then fine-tune 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 indeed 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.
- 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.
2.1 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. (2020). 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.
2.2 BERT Compression
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.3 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).
3.1 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:
$$
\operatorname {S e l f - A t t} _ {h} (\mathbf {H}) = \operatorname {s o f t m a x} \left(\frac {\left(\mathbf {H W} _ {Q _ {h}}\right) \left(\mathbf {H W} _ {K _ {h}}\right) ^ {\top}}{\sqrt {d _ {h}}}\right) \mathbf {H W} _ {V _ {h}} \tag {1}
$$
where  $\mathbf{H} \in \mathbb{R}^{|x| \times d_H}$  is the input;  $d_H$  and  $|\mathbf{x}|$  are the hidden size and the length of input  $\mathbf{x}$ , respectively.  $\mathbf{W}_{Q_h, K_h, V_h} \in \mathbb{R}^{d_H \times d_h}$  are the query, key and value matrices, and  $d_h = \frac{d_H}{N_h}$ . In practice, the matrices for different heads will be combined into three large matrices  $\mathbf{W}_{Q, K, V} \in \mathbb{R}^{d_H \times d_H}$ . The outputs of the  $N_h$  heads are then concatenated and linearly projected by  $\mathbf{W}_{AO} \in \mathbb{R}^{d_H \times d_H}$  to obtain the final output of the self-attention layer.
The FFN consists of two weight matrices  $\mathbf{W}_{FI} \in \mathbb{R}^{d_H \times d_I}$ ,  $\mathbf{W}_{FO} \in \mathbb{R}^{d_I \times 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.2 Subnetwork and Magnitude Pruning
Consider a model  $f(\cdot ;\pmb {\theta})$  with weights  $\pmb{\theta}$  we can obtain its subnetwork  $f(\cdot ;\mathbf{M}\odot \pmb {\theta})$  by applying a binary mask  $\mathbf{M}\in \{0,1\}^{|\pmb{\theta}|}$  to  $\pmb{\theta}$  , where  $\odot$  denotes element-wise multiplication. In terms of BERT, we extract the subnetwork from the pretrained weights  $\pmb{\theta}_0$  . Specifically, we consider the matrices of the Transformer sub-layers and the word embedding matrix, i.e.,  $\pmb{\theta}_{0} = \{\mathbf{W}_{Emb}\} \cup$ $\left\{\mathbf{W}_Q^l,\mathbf{W}_K^l,\mathbf{W}_V^l,\mathbf{W}_{AO}^l,\mathbf{W}_{FI}^l,\mathbf{W}_{FO}^l\right\}_{l = 1}^L.$
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
Figure 2: Illustration of the BERT subnetwork transfer problem and the proposed TAMT. We search the subnetworks by training binary masks on the pre-training dataset, using either the MLM loss or the KD loss (left). The identified subnetwork is then fine-tuned on a range of downstream tasks (right). The colored weights/masks are trainable and the black ones are frozen. The residual connection and layer normalization are omitted for simplicity.
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.3 Problem Formulation: Transfer BERT Subnetwork
As depicted in Fig. 2, given  $N_{\mathcal{T}}$  downstream tasks  $\mathcal{T} = \{\mathcal{T}_i\}_{i=1}^{N_{\mathcal{T}}}$ , the subnetwork  $f\left(\cdot; \mathbf{M} \odot \boldsymbol{\theta}_0, \mathcal{C}_0^{\mathcal{T}_i}\right)$  is fine-tuned on each task, together with the randomly initialized task-specific linear classifier  $\mathcal{C}_0^{\mathcal{T}_i}$ . We formulate the training algorithm for task  $\mathcal{T}_i$  as a function  $\mathcal{A}_t^{\mathcal{T}_i}\left(f\left(\cdot; \mathbf{M} \odot \boldsymbol{\theta}_0, \mathcal{C}_0^{\mathcal{T}_i}\right)\right)$  (e.g., Adam or SGD), which trains the model for  $t$  steps and produces  $f\left(\cdot; \mathbf{M} \odot \boldsymbol{\theta}_t, \mathcal{C}_t^{\mathcal{T}_i}\right)$ . After fine-tuning, the model is evaluated against the metric  $\mathcal{E}^{\mathcal{T}_i}(f(\cdot; \mathbf{M} \odot \boldsymbol{\theta}_t, \mathcal{C}_t^{\mathcal{T}_i}))$  (e.g., Accuracy or F1) for task  $\mathcal{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:
$$
\begin{array}{r l} & {\underset {\mathbf {M}} {\max } \left(\frac {1}{N _ {\mathcal {T}}} \sum_ {i = 1} ^ {N _ {\mathcal {T}}} \mathcal {E} ^ {\mathcal {T} _ {i}} \left(\mathcal {A} _ {t} ^ {\mathcal {T} _ {i}} \left(f \left(\cdot , \mathbf {M} \cdot \boldsymbol {\theta} _ {0}, \mathcal {C} _ {0} ^ {\mathcal {T} _ {i}}\right)\right)\right)\right)} \\ & {\quad \text {s . t .} \frac {\| \mathbf {M} \| _ {0}}{\left| \boldsymbol {\theta} _ {0} \right|} = (1 - \mathcal {S})} \end{array} \tag {2}
$$
where  $\| \mathbf{M} \|_0$  and  $|\pmb{\theta}_0|$  are the  $L_0$  norm of the mask and the total number of model weights respectively.
3.4 Task-agnostic Mask Training
3.4.1 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. (2020); Mallya et al. (2018). This technique involves mask binarization in the forward pass and gradient estimation in the backward pass.
As shown in Fig. 2, each weight matrix  $\mathbf{W} \in \mathbb{R}^{d_{in} \times d_{out}}$  is associated with a binary mask  $\mathbf{M} \in \{0,1\}^{d_{in} \times d_{out}}$ , which is derived from a real-valued matrix  $\overline{\mathbf{M}} \in \mathbb{R}^{d_{in} \times d_{out}}$  via binarization:
$$
\mathbf {M} _ {i, j} = \left\{ \begin{array}{l l} 1 & \text {i f} \overline {{\mathbf {M}}} _ {i, j} \geq \phi \\ 0 & \text {o t h e r w i s e} \end{array} \right. \tag {3}
$$
where  $\phi$  is the threshold that controls the sparsity. In the forward pass of a subnetwork,  $\mathbf{W} \odot \mathbf{M}$  is used in replacement of the original weights  $\mathbf{W}$ .
Since  $\mathbf{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:
$$
\overline {{\mathbf {M}}} \leftarrow \overline {{\mathbf {M}}} - \eta \frac {\partial \mathcal {L}}{\partial \mathbf {M}} \tag {4}
$$
where  $\mathcal{L}$  is the loss function and  $\eta$  is the learning rate. In other words, the gradients of  $\overline{\mathbf{M}}$  is estimated using the gradients of  $\mathbf{M}$ . In the process of mask training, all the original weights are frozen.
3.4.2 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  $\overline{\mathbf{M}}$  using OMP:
$$
\overline {{\mathbf {M}}} _ {i, j} = \left\{ \begin{array}{l l} \alpha \times \phi & \text {i f} \mathbf {M} _ {i, j} ^ {O M P} = 1 \\ 0 & \text {o t h e r w i s e} \end{array} \right. \tag {5}
$$
where  $\mathbf{M}^{OMP}$  is the binary mask derived from OMP and  $\alpha \geq 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  $\overline{\mathbf{M}}_{i,j} = \alpha \times \phi \geq \phi$ . 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  $\overline{\mathbf{M}}$  being updated, some of its entries with zero initialization will gradually surpass the threshold, and vice versa. If the threshold  $\phi$  is fixed throughout training, there is no guarantee that the binary mask will always satisfy the given sparsity. Therefore, we rank  $\overline{\mathbf{M}}_{ij}$  according to their absolute values during mask training, and dynamically adjust the threshold to satisfy the sparsity constraint.
3.4.3 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:
$$
\mathcal {L} _ {d i s t i l l} = \frac {1}{L | \mathbf {x} |} \sum_ {l = 1} ^ {L} \sum_ {i = 1} ^ {| \mathbf {x} |} (1 - \cos (\mathbf {H} _ {l, i} ^ {T}, \mathbf {H} _ {l, i} ^ {S})) (6)
$$
where  $\mathbf{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 (\cdot ,\cdot)$  is the cosine similarity.
4.1.1 Models
We examine two PLMs from the BERT family, i.e.,  $\mathrm{BERT}_{\mathrm{BASE}}$  (Devlin et al., 2019) and  $\mathrm{RoBERTa}_{\mathrm{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  $\mathrm{BERT}_{\mathrm{BASE}}$ .
4.1.2 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 WikiText-103 dataset (Merit 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.
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  $\phi$  and  $\alpha$  are set to 0.01 and 2 respectively,
Figure 3: Downstream performance of BERT subnetworks. Shadowed areas denote standard deviations.
Figure 4: Downstream performance of RoBERTa subnetworks. Shadowed areas denote standard deviations.
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. (2020). Every pruning operation in IMP is followed by resetting the remaining weights to  $\theta_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.2.1 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.