SlowGuess/NAACL_2021_OCR huggingface dataset:
Xuan Zhou<sup>1</sup>, Xiao Zhang<sup>1</sup>, Chenyang Tao<sup>2</sup>, Junya Chen<sup>2</sup>, Bing Xu<sup>1</sup>, Wei Wang<sup>1</sup>, and Jing Xiao<sup>1</sup>
$^{1}$ Ping An Technology (Shenzhen) Co., Ltd  
 $^{2}$ Duke University
{zhouxuan553,zhangxiao585,xubing962,wangweil92,xiaojing661}@pingan.com.cn
{chenyang.tao,junya.chen}@duke.edu
Although pre-trained big models (e.g., BERT, ERNIE, XLNet, GPT3 etc.) have delivered top performance in Seq2seq modeling, their deployments in real-world applications are often hindered by the excessive computations and memory demand involved. For many applications, including named entity recognition (NER), matching the state-of-the-art result under budget has attracted considerable attention. Drawing power from the recent advance in knowledge distillation (KD), this work presents a novel distillation scheme to efficiently transfer the knowledge learned from big models to their more affordable counterpart. Our solution highlights the construction of surrogate labels through the  $k$ -best Viterbi algorithm to distill knowledge from the teacher model. To maximally assimilate knowledge into the student model, we propose a multi-grained distillation scheme, which integrates cross entropy involved in conditional random field (CRF) and fuzzy learning. To validate the effectiveness of our proposal, we conducted a comprehensive evaluation on five NER benchmarks, reporting cross-the-board performance gains relative to competing prior-arts. We further discuss ablation results to dissect our gains.
The task of named entity recognition (NER) aims to put named entity mentioned in a sentence into some pre-defined categories such as the person names, organizations, locations, etc. NER is a fundamental task in nature language processing (NLP), and often serves as an information extraction tool embedded in solutions for downstream tasks such as event recognition and aspect-level sentiment analysis. And therefore, the computational and memory efficiency of an NER model is often considered crucial in many empirical settings.
Given its practical significance, considerable research effort has been devoted to NER in recent
years. One fruitful direction is to boost NER performance through the use of more sophisticated model architectures, such as Transformer and its variants (e.g., BERT (Devlin et al., 2019), ERNIE (Sun et al., 2019c), XLNet (Yang et al., 2019), RoBERTa (Liu et al., 2019b), GPT3 (Brown et al., 2020), etc.). Many of these models are based on the attention mechanism, which allows the models to adaptively focus on different parts of the sentence based on its current understanding. This enables more accurate parsing of the context, which is critical for the NER task. While such advanced models have delivered substantial improvements, a major drawback is that they are typically computationally expensive and memory intensive, preventing their applications in many cost-sensitive settings.
As such, it is often desirable to reduce those big models into more affordable counterparts, preferably without any significant performance drop. One strategy is to compress or truncate the original model, examples in this category include parameter pruning (Srinivas and Babu, 2015; McCarley, 2019), low-rank approximation (Yu et al., 2017; Ma et al., 2019) and parameter quantization (Gong et al., 2014; Wu et al., 2016). The resulting model has a similar architecture to the original model, and consequently may suffer from similar limitations of the original model. Alternatively, knowledge distillation (KD) considers transferring knowledge into models with heterogeneous architectures (Hinton et al., 2015), which allows more flexibility in the control of resource usage. In KD, a new learner, often dubbed the student model, assimilates knowledge from a pre-trained model, commonly known as the teacher. For classification tasks, this is typically achieved via minimizing their discrepancy in the output, regardless of the internal model architectures. Teacher outputs, sometimes referred to as the soft-labels, typically encode more information than what a student might receive from the
Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 5704-5716 June 6-11, 2021. ©2021 Association for Computational Linguistics
raw training feature-label pairs, thus resulting the improved learning efficiency.
However, for sequence tasks such as NER, standard techniques from KD do not readily apply. While common KD seeks to minimize the KL divergence between the output label distributions, for NER, the number of label combinations grows exponentially wrt the sequence length. Extracting teacher knowledge as if each combination is a different label category would be largely inefficient, if possible at all. On the other hand, many SOTA NER models are built on conditional random fields (CRF) to incorporate the label dependencies, and it only offers an un-normalized likelihood that is not directly amenable to the computation of KL-divergence. Another challenge for training an NER model is the lack of labeled sequences for training, while the unlabeled sequences may be plentiful. The promise to leverage unlabeled data to improve NER accuracy is appealing.
To address the above challenges, in this paper, we present an efficient knowledge distillation scheme that trains a light model (e.g., BiLSTM) which is able to retain the accuracy of its heavier counterparts, such as BERT, while significantly reduce cost. Our solution exploits the "soft surrogates", i.e., the most probable label sequences under the teacher model, to inform the student learner. To efficiently identify the most likely label sequences and determine their relative likelihood, we explore the use of Viterbi algorithm to expedite computation. We also explore the use of unlabeled data to improve performance.
In summary, we highlight the following contributions in this study: (i) We present a novel multi-grained knowledge distillation strategy for sequence labeling via efficiently selecting  $k$ -best label sequence using Viterbi algorithm; (ii) We advocate the use of a complete cross entropy loss and fuzzy distillation loss to respectively account for probability mass of un-selected sequences and the uncertainties in teacher confidence; (iii) We present a comprehensive empirical analysis to dissect the gains from each individual components. We also show our model delivers substantial improvement relative to competing solutions on a wide-range of real-world benchmarks to demonstrate its utility.
Our work is inspired by three lines of research: sequence-level knowledge distillation,  $k$ -
best Viterbi algorithm and fuzzy conditional random field (CRF). In the following we review the technical backgrounds that are needed for the construction of our model.
Sequence-level Knowledge Distillation. Knowledge distillation (KD) is originally developed for classification tasks (Hinton et al., 2015; Tang et al., 2019). When dealing with sequence outputs (e.g., machine translation, sequence labeling), where each unique combination of the output sequence is treated as different category, then the standard distillation objective is no longer appropriate, if feasible at all. This is because the number of unique combinations for a length  $L$  sequence with  $T$  possible tags scale at  $T^{L}$ . To combat such exponential scaling, Kim and Rush (2016) investigated using beam search for teacher output to select  $k$ -best candidates for KD in neural machine translation (NMT), and Mun'im et al. (2019) utilized an similar technique for KD in Large Vocabulary Continuous Speech Recognition (LVSCR) tasks.
$k$ -best Viterbi Decoding. Viterbi decoding is a dynamic programming technique to find a sequence with the highest score in an exponential-growth domain, with only linear complexity (Viterbi, 1967). Generalization has been proposed to extend its original scope to find the top- $K$  sequences that are most probable, see (Huang and Chiang, 2005; Nielsen, 2011) for details. A summary of the algorithm can be found in the supplementary material. In the context of KD in sequence tasks, a trained teacher model assigns varying probability to all sequence combinations. Our motivation is that the  $k$ -best Viterbi can be repurposed to pick out the  $K$ -most probable label sequences predicted by the teacher model, which plays an analogous role to the soft-labels (Figure 1a), without incurring unmanageable computational overhead.
Conditional Random Field. CRF (Lafferty et al., 2001) is a classic and powerful energy-based model that is capable of capturing complex spatial or temporal dependency structures. In the context of NLP, it has been generally used as a refinement layer that accounts for correlations missed by the base NLP model, which typically brings in additional performance gains for various NLP tasks (Lafferty et al., 2001; Collobert et al., 2011a; Huang et al., 2015). Implementation-wise, CRF computes an energy given a candidate output  $\mathbf{y}$  and a context  $\mathbf{x}$  (i.e., input sequence), followed by a softmax operator to obtain the conditional likelihood, i.e.,
Figure 1: Left. The model prediction accuracy improves as one includes more candidate  $K$ . Right. The non-top- $K$  probability remains non-neglectable even for moderately large  $K$ . Here, OntoNotes 4.0 is used for illustration, and Y-axis represents the proportion of samples, whose non-top- $K$  cumulative probability is greater than 0.1.
$$
\mathbb {P} (\mathbf {y} | \mathbf {x}) = \frac {e ^ {s (\mathbf {x} , \mathbf {y})}}{\sum_ {\tilde {\mathbf {y}} \in Y _ {\mathrm {a l l}}} e ^ {s (\mathbf {x} , \tilde {\mathbf {y}})}}, \tag {1}
$$
where  $Y_{\mathrm{all}}$  is the set of all possible tag sequences,  $s(\mathbf{x},\mathbf{y})$  represents the "compliance" energy score between two sequences  $\mathbf{x}$  and  $\mathbf{y}$ . In classical NLP models,  $s(\mathbf{x},\mathbf{y})$  is typically specified via handcrafted features and dependencies. In modern deep learning models, a typical decomposition of  $s(\mathbf{x},\mathbf{y})$  is given by two parts
$$
s (\mathbf {x}, \mathbf {y}) = \sum_ {j = 1} ^ {L - 1} \operatorname {l o g i t} _ {j, \mathbf {y} _ {j}} (\mathbf {x}) + A _ {\mathbf {y} _ {j}, \mathbf {y} _ {j + 1}}, \tag {2}
$$
namely the emission and transition score. The transition matrix  $A \in \mathbb{R}^{T \times T}$  characterizes the smoothness of the label sequence (probability of switching between consequent labels), and  $\operatorname{logit}_{j,k}(\mathbf{x})$  denotes the likelihood of seeing label  $k$  at position  $j$  as predicted by the base model. Note that we do not have to enumerate the exponentially many  $Y_{\mathrm{all}}$  to compute the likelihood, which can be efficiently handled by the Viterbi algorithm introduced above in linear time (Collobert et al., 2011b).
Fuzzy CRF. It has been argued that properly adjusting the uncertainty of teacher prediction usually lends better knowledge transfer to the student, which can either be sharpening or relaxing the teacher predicted distributions. In standard KD, this is achieved by the incorporation of an annealing factor. In our setup, we consider relaxation via generalizing CRF to a candidate set of label sequences (Shang et al., 2018) rather than individual ones. Formally, we define the fuzzy loss as
$$
\mathbb {P} (Y _ {\mathrm {c a n d i d a t e}} | \mathbf {x}) = \frac {\sum_ {\mathbf {y} ^ {\prime} \in Y _ {\mathrm {c a n d i d a t e}}} e ^ {s (\mathbf {x} , \mathbf {y} ^ {\prime})}}{\sum_ {\tilde {\mathbf {y}} \in Y _ {\mathrm {a l l}}} e ^ {s (\mathbf {x} , \tilde {\mathbf {y}})}}, \quad (3)
$$
where  $Y_{\text{candidate}}$  contains candidate label sequences. Here, we will use the  $k$ -best teacher predicted label sequences as the candidate set.
Figure 2:  $k$ -best cross entropy loss uses predicted weights for each candidate label sequence, while  $k$ -best fuzzy loss group the candidates together and use their aggregated weight. The two losses represent fine-grained and coarse-grained learning from teacher respectively. Both schemes lump weights for the nontop- $K$  labels.
3 Multi-Grained Distillation
In this section we detail the construction of our distillation scheme, with the overall architecture outlined in Figure 3.
3.1 Viterbi Distillation for Sequence Outputs
Now we are ready to present details of our main contribution, Viterbi distillation for sequence outputs. Our basic idea is to extract information from the teacher model via drawing a set of most probable sequences, together with the respective confidence to those sequences. Then these sequences are presented to the student model during its training, to pass on the knowledge from teacher through various loss functions.
More specifically, we apply the  $k$ -best Viterbi algorithm (see Algorithm 1 in the supplementary material) to get the pairs  $\{(\mathbf{y}_1^{(i)}, p_1^{(i,\mathbf{t})}), \dots, (\mathbf{y}_K^{(i)}, p_K^{(i,\mathbf{t})})\}$  for sample  $\mathbf{x}^{(i)}$ , where  $\mathbf{y}_j^{(i)}$  is the  $j$ -th most-likely label sequence, and  $p_j^{(i,\mathbf{t})}$  is the corresponding probability under the teacher model (indicated by the superscript  $\mathbf{t}$ ). Similarly, in our subsequent discussions we denote  $p_j^{(i,\mathbf{s})}$  as the probability produced by the student model on the label sequence  $\mathbf{y}_j^{(i)}$ .
3.1.1 Fine-grained  $k$ -best Cross Entropy Distillation
We proposed to approximate the complete cross entropy with the following  $k$ -best Vertibi approximation to avoid the exponential level of computational
Figure 3: Model schematic for multi-grained knowledge distillation.
complexity:
$$
\begin{array}{l} \mathcal {L} _ {\mathrm {S e q - c e}} (\mathbf {x} ^ {(i)}) \\ := -\sum_ {\tilde {\mathbf {y}} \in Y _ {\text {a l l}}} \mathbb {P} ^ {\mathbf {t}} (\tilde {\mathbf {y}} | \mathbf {x} ^ {(i)}) \log \left(\mathbb {P} ^ {\mathbf {s}} (\tilde {\mathbf {y}} | \mathbf {x} ^ {(i)})\right) \\ \approx - \sum_ {j = 1} ^ {K} p _ {j} ^ {(i, \mathbf {t})} \log \left(p _ {j} ^ {(i, \mathbf {s})} \left(\mathbf {y} _ {j} ^ {(i)} \mid \mathbf {x} ^ {(i)}\right)\right) \tag {4} \\ - \left(1 - p _ {\text {s u m}} ^ {\mathbf {t}}\right) \log \left(1 - p _ {\text {s u m}} ^ {\mathbf {s}}\right), \\ \end{array}
$$
where  $\mathbb{P}^{\mathbf{t}}(\cdot |\cdot)$  and  $\mathbb{P}^{\mathrm{s}}(\cdot |\cdot)$  are the CRF likelihoods respectively for teacher and student, and we collect the probability mass of the top-  $K$  sequences into terms  $p_{\mathrm{sum}}^{\mathbf{t}}$  and  $p_{\mathrm{sum}}^{\mathrm{s}}$ , i.e.,
$$
p _ {\text {s u m}} ^ {\mathbf {t}} = \sum_ {j = 1} ^ {K} p _ {j} ^ {(i, \mathbf {t})}, \quad p _ {\text {s u m}} ^ {\mathbf {s}} = \sum_ {j = 1} ^ {K} p _ {j} ^ {(i, \mathbf {s})}. \tag {5}
$$
Note that we have suppressed the dependency on  $\mathbf{x}$  for notational clarity. We emphasize the residual probability mass  $(1 - p_{\mathrm{sum}}^{\mathbf{t}})\gg 0$  (see Figure 1b), and therefore should not be excluded from the loss as in Kim and Rush (2016); Wang et al. (2020).
3.1.2 Coarse-grained Fuzzy Distillation
To more appropriately account for the uncertainty in the teacher guidance, we advocate the use of a fuzzy objective that does not discriminate between the likelihood among the top-  $K$  picks. Concretely, this is given by the binary cross entropy in terms of  $p_{\mathrm{sum}}^{\mathbf{t}}$  and  $p_{\mathrm{sum}}^{\mathbf{s}}$ ,
$$
\begin{array}{l} \mathcal {L} _ {\text {S e q - f u z z y}} \left(\mathbf {x} ^ {(i)}\right) := - p _ {\text {s u m}} ^ {\mathbf {t}} \log \left(p _ {\text {s u m}} ^ {\mathbf {s}}\right) \tag {6} \\ - \left(1 - p _ {\text {s u m}} ^ {\mathbf {t}}\right) \log \left(1 - p _ {\text {s u m}} ^ {\mathbf {s}}\right). \\ \end{array}
$$
Multi-grained Distillation. See Figure 2 for a graphical illustration of cross entropy and fuzzy objectives, which shows different granularity on learning from teacher.
3.1.3 Integrating Distillation Objectives
In addition to the "pure" distillation operations, we also allow direct learning from ground-truth labels and surrogate ones labeled by the teacher model (details found in Section 3.2). These two kinds of targets are named hard labels, and the related loss is as follows,
$$
\mathcal {L} _ {\text {h a r d}} \left(\mathbf {x} ^ {(i)}\right) := - \log \mathbb {P} \left(\mathbf {y} _ {\text {h a r d}} \mid \mathbf {x} ^ {(i)}\right). \tag {7}
$$
Thus, the final loss is the weighted sum of three terms defined above,
$$
\begin{array}{l} \mathcal {L} _ {\text {S e q - a l l}} = \frac {1}{N} \sum_ {i = 1} ^ {N} \lambda_ {1} \mathcal {L} _ {\text {h a r d}} \left(\mathbf {x} ^ {(i)}\right) \tag {8} \\ + \lambda_ {2} \mathcal {L} _ {\text {S e q - f u z z y}} (\mathbf {x} ^ {(i)}) + \lambda_ {3} \mathcal {L} _ {\text {S e q - c e}} (\mathbf {x} ^ {(i)}), \\ \end{array}
$$
where  $\lambda_{i}\geq 0,i = 1,2,3$
Automated Tuning of Loss Weights. Tuning the loss weights  $\{\lambda_i\}_{i=1}^3$  in (8) via standard techniques such as grid search can be laborious and costly. In this work, we leverage the uncertainty weighting strategy proposed in (Cipolla et al., 2018) to automate the weight selection procedure to balance multi-task objectives, both for its simplicity and robust performance. More specifically, we add uncertainty regularization terms for the weights, resulting
$$
\mathcal {L} _ {\text {S e q - a l l}} ^ {\mathrm {u}} = \mathcal {L} _ {\text {S e q - a l l}} - \frac {1}{2} \left(\log \lambda_ {1} + \log \lambda_ {2} + \log \lambda_ {3}\right), \tag {9}
$$
as our new learning objective.
3.2 Data Augmentation and Misc
Augmenting with Unlabeled Data. NER applications are often challenged with the lack of labeled instances for training. Motivated by the observations from prior studies that data augmentation using unlabeled data may improve distillation performance (Mukherjee and Awadallah, 2019; Tang
et al., 2019), we also feed unlabeled data to the teacher model to construct additional surrogate label sequences using  $k$ -best Viterbi for distillation. We pool the surrogate sequences from both labeled and unlabeled data for cross entropy and fuzzy loss. Besides, we use Top-1 sequence as hard label to educate the student.
Token-level Distillation. To verify the necessity of using a sequence-level distillation, we additionally consider token-level distillation with standard KD techniques. Specifically, we consider the following loss
$$
\mathcal {L} _ {\text {T o k e n}} := - \frac {1}{N} \sum_ {i = 1} ^ {N} \sum_ {j = 1} ^ {T} \mathcal {L} _ {i, j}, \tag {10}
$$
where  $\mathcal{L}_{i,j} = \sum_{k=1}^{L} \mathbb{P}^{\mathbf{t}}(\mathbf{y}_j = k|\mathbf{x}^{(i)})\log (\mathbb{P}^{\mathbf{s}}(\mathbf{y}_j = k|\mathbf{x}^{(i)}))$  is the token-level KL-divergence between the teacher's and student's predicted distributions for each token, rather than the token sequence distribution. This breaks the dependencies between tokens to avoid the combinatory complexity, thus enabling standard KD. The probability  $\mathbb{P}(\mathbf{y}_j = k|\mathbf{x}^{(i)})$  can be simply calculated by emission score, in this case, the loss is denoted as  $\mathcal{L}_{\mathrm{Token - em}}$ . To recover the spatial dependencies, we can add transition matrix to derive the probability  $\mathbb{P}(\mathbf{y}_j = k|\mathbf{x}^{(i)})$  by calculating the marginal distribution of token from the probability distribution  $\mathbb{P}(\mathbf{y}|\mathbf{x}^{(i)})$  given by CRF, similar to Wang et al. (2020). We denote this corresponding loss as  $\mathcal{L}_{\mathrm{Token - pos}}$ .
NER without BERT. Extensive investigations in NER have been conducted without appealing to the BERT architecture, and instead using simpler architectures such as BiLSTM or convolutional neural network (CNN) (Huang et al., 2015; Strubell et al., 2017a). Special effort has been made to adapt the NER network architectures for better handling of lexicons. Liu et al. (2019a) leveraged hybrid semi-Markov CRFs to improve NER recognition with Gazetteers, where segments are used as the basic units instead of words. In Zhang and Yang (2018) a lattice-structured LSTM is proposed to encode a sequence of input characters as well as all potential words that match a lexicon. Gui et al. (2019) incorporates lexicons using a rethinking mechanism under a CNN setup, which renders faster inference compared with lattice-LSTM. Peng et al. (2019)
explores weighted word embedding to match all probable words given the lexicons. In the method given by Ghaddar and Langlais (2018), a lexical representation is computed for each word with a 120-dimensional vector, where element encodes the similarity of the word with an entity type. Recently, Ding et al. (2019) leveraged the graph neural network (GNN) to exploit the additional rich information captured by the gazetteers, setting new SOTA performance. Our distillation work is orthogonal to these developments, and the techniques can be combined for further improvements.
Compressing BERT with Distillation. Various efforts have been made to reduce the size and cost of BERT or other big models while maximally maintaining their outstanding performance, and KD offers an appealing alternative to the direct amputation of the original models. Along this line, Tang et al. (2019) studied on compressing BERT to BiLSTM for classification, resulting a model with comparable performance to ELMo but  $100 \times$  fewer parameters and  $15 \times$  faster inference. Tsai et al. (2019) successfully applied KD for multilingual sequence labeling model, enabling SOTA results on an MiniBERT model afforded by a single CPU. Other developments include distillation on intermediate representations (Sun et al., 2019b) and student pre-training (Turc et al., 2019). The value of unlabeled data in KD has also been explored (Mukherjee and Awadallah, 2019). Our work complements these studies via presenting a framework of using top- $K$  soft-surrogate labels for KD.
Distilling with top- $K$  Picks. Parallel to our work, Tang et al. (2020) also considered a top- $K$  scheme for KD. This work distincts from our proposal via assigning uniform weights to the un-selected labels sequences, a practice that can be largely inefficient. Another work close to our setup is Wang et al. (2020), where the authors tried to distill the structural knowledge from multiple monolingual teachers to a single student. In their loss, the probability mass of the non-top- $K$  picks is discarded. In Figure 1 (right), we show that the probability for non-top- $K$  picks is non-neglectable, which is properly accounted for in our proposed complete cross entropy loss.
To validate the proposed solution and elaborate the gains, we benchmark it against state-of-the-art methods, through a wide range of experiments on
real-world datasets. All experiments are implemented with Tensorflow and executed on a single NVIDIA P100 GPU. Details of the experimental setup are provided in supplementary material, due to space limits, and our code will be available at https://github.com/11zhouxuan/multi_grained_kd_ner.
The following real-world datasets are considered in our study. Detailed summary statistics of the datasets can be found in supplementary material.
CoNLL-2003 NER (Sang and De Meulder, 2003) consists of newswire from Reuters RCV corpus. Unlabeled data from the Reuters RCV corpus is used for our data augmentation experiments.
OntoNotes (Consortium et al., 2011) is an annotated multilingual corpus consists of texts from a wide variety of sources, such as telephone conversation, broadcast and newswire. For our NER experiment, we consider
- English NER is derived from OntoNotes Release 5.0 and processed according to Pradhan et al. (2013).
- Chinese NER is derived from OntoNotes Release 4.0 and processed according to Che et al. (2013).
Text data from other OntoNotes tasks are used as unlabeled data.
MSRA (Levow, 2006) is a Chinese NER dataset with its corpus derived from news domain. We use the Chinese word segmentation dataset MSR (Levow, 2006) for unlabeled data.
Weibo (Peng and Dredze, 2015) is a Chinese NER dataset derived from social media contents. Only a small fraction of data is labeled in this dataset.
5.2 Model specification