SlowGuess/NAACL_2022_OCR huggingface dataset:
Qinyuan Ye $^{1\dagger}$  Madian Khabsa $^{2}$  Mike Lewis $^{2}$  Sinong Wang $^{2}$  Xiang Ren $^{1}$  Aaron Jaech $^{2}$ $^{1}$ University of Southern California  $^{2}$ Meta AI {qinyuany, xiangren}@usc.edu  
<mkhabsa, mikelewis, sinongyang, ajaech>@fb.com
Distilling state-of-the-art transformer models into lightweight student models is an effective way to reduce computation cost at inference time. The student models are typically compact transformers with fewer parameters, while expensive operations such as self-attention persist. Therefore, the improved inference speed may still be unsatisfactory for real-time or high-volume use cases. In this paper, we aim to further push the limit of inference speed by distilling teacher models into bigger, sparser student models – bigger in that they scale up to billions of parameters; sparser in that most of the model parameters are n-gram embeddings. Our experiments on six single-sentence text classification tasks show that these student models retain  $97\%$  of the RoBERTa-Large teacher performance on average, and meanwhile achieve up to 600x speed-up on both GPUs and CPUs at inference time. Further investigation reveals that our pipeline is also helpful for sentence-pair classification tasks, and in domain generalization settings.
Large pre-trained Transformers (Devlin et al., 2019; Liu et al., 2019) are highly successful, but their large inference costs mean that people who host low-latency applications, or who are simply concerned with their cloud computing costs have looked for ways to reduce the costs. Prior work mainly achieves this by leveraging knowledge distillation (Hinton et al., 2015), which allows for the capabilities of a large well-performing model known as the teacher to be transferred to a smaller student model. For example, DistillBERT (Sanh et al., 2019) is a smaller transformer model distilled from BERT (Devlin et al., 2019), which reduces BERT's size by  $40\%$  and becomes  $60\%$  faster during inference. However, such speed-up may be still
Figure 1: Performance vs. Inference Speed. With Deep Averaging Network (DAN; Iyyer et al. 2015) and knowledge distillation, we obtain a student model with competitive performance on IMDB dataset, while being  $607\mathrm{x}$  faster than RoBERTa-Large, and  $20\mathrm{x}$  faster than bi-directional LSTMs at inference time.
insufficient for high-volume or low-latency inference tasks. In this paper, we aim to further push the limit of inference speed, by introducing Sparse Distillation, a framework that distills the power of state-of-the-art transformer models into a shallow, sparsely-activated, and richly-parameterized student model.
Counter to the convention of using "smaller, faster, [and] cheaper" (Sanh et al., 2019) student models, our work explores a new area of the design space, where our fast and cheap student model is actually several times larger than the teacher. The student model we use is modified from Deep Averaging Network (DAN) in Iyyer et al. (2015). DANs take a simple architecture by mapping the n-grams in the input sentence into embeddings, aggregating the embeddings with average pooling, and then using multiple linear layers to perform classification (see Fig. 2). This architecture is reminiscent of the high expressive power of billion-parameter n-gram models (Buck et al., 2014; Brants et al., 2007) from before the existence of pre-trained language
$^{\dagger}$ Work partially done while interning at Meta AI.  
 $^{1}$ Code available at https://github.com/INK-USC/spa rse-distillation.
Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 2361 - 2375 July 10-15, 2022 ©2022 Association for Computational Linguistics
models. By selecting the n-gram vocabulary and the embedding dimension, DANs also scale up to billions of parameters. Meanwhile, the inference costs are kept low as DANs are sparsely-activated.
One weakness of DANs is that they are restricted in modeling high-level meanings in long-range contexts, as compared to the self-attention operator in Transformers. However, recent studies have shown that large pre-trained Transformers are rather insensitive to word order (Sinha et al., 2021) and that they still work well when the learned self-attention is replaced with hard-coded localized attention (You et al., 2020) or convolution blocks (Tay et al., 2021). Taken together, these studies suggest that on some tasks it may be possible to get competitive results without computationally expensive operations such as self-attention.
To verify our hypothesis, we use six single-sentence text classification tasks $^2$  and apply knowledge distillation to DANs. We observe that the resulting student models retain  $97\%$  of the RoBERTa-Large teacher performance on average. We also show that our method falls outside of the Pareto frontier of existing methods; compared to a baseline of distilling to a LSTM student, our method gives comparable accuracy at less than 1/20 the inference cost (see Fig. 1). Based on our empirical results, we conclude that faster and larger student models provide a valuable benefit over existing methods. We further examine our method (1) with QQP, a sentence-pair task, (2) in privacy-preserving settings (i.e., no access to task-specific data during distillation), and (3) in domain generalization and adaptation settings (i.e., student models are applied and adapted to new data domains), where we find our method continues to bring improvements over non-distillation baselines.
2 Sparse Distillation with DANs
Our goal is to train an efficient text classification model  $M$  for a given task  $T$ . In a  $n$ -way classification problem, the model  $M$  takes input text  $\mathbf{x}$ , and produces  $\hat{\mathbf{y}} \in \mathbb{R}^n$ , where  $\hat{y}_i$  indicates the likelihood that the input  $\mathbf{x}$  belongs to category  $i$ . The task  $T$  has a train set  $D_{train}$  and a validation/development set  $D_{dev}$ . Additionally, we assume access to a large unlabeled corpus  $C$  which is supposedly in
Figure 2: We primarily use a modified Deep Averaging Network (DAN; Iyyer et al. 2015) as the student model in this paper. DAN contains a sparse n-gram embedding table and two linear layers. Embedding dimension  $d_{e}$  is set to 3 in this figure for illustration purpose.
a domain relevant to task  $T$ . We comprehensively evaluate the efficiency of the model  $M$  by reporting: (1) accuracy on  $D_{dev}$ , (2) inference speed, and (3) the number of parameters in the model.
2.2 Method Overview
To train a text classifier that is both efficient and powerful, we employ knowledge distillation (Hinton et al., 2015), by having a powerful teacher model provide the supervision signal to an efficient student model. In particular, we are interested in using sparse n-gram based models as our student model. We explain the teacher and student model we use in §2.3, the training pipeline in §2.4, and implementation details in §2.5
Teacher Model. Fine-tuning a pre-trained transformer model is the predominant recipe for obtaining state-of-the-art results on various text classification tasks. Our teacher model is a RoBERTa-Large model (Liu et al., 2019) fine-tuned on the training set  $D_{train}$  of task  $T$ .
Student Model. Our student model is based on the Deep Averaging Network (DAN, Iyyer et al. 2015) with the modification that we operate on  $n$ -grams instead of just words. See Fig. 2 for an illustration of the model architecture. Specifically, for an input sentence  $\mathbf{x}$ , a list of  $n$ -grams  $g_{1}, g_{2}, \ldots, g_{n}$  are extracted from the sentence. These  $n$ -gram indices are converted into their embeddings (with dimension  $d_{e}$ ) using an embedding layer  $\operatorname{Emb}(.)$ . The sentence representation  $\mathbf{h}$  will be computed as the average of all  $n$ -gram embeddings, i.e.,  $\mathbf{h} = \operatorname{Mean}(\operatorname{Emb}(g_{1}), \operatorname{Emb}(g_{2}), \ldots, \operatorname{Emb}(g_{n})) \in \mathbb{R}^{d_{e}}$ . The sentence representation then goes through two
${}^{2}$  Transformers are effective at many tasks beyond text classification. We extend our method to sentence-pair tasks in later sections and leave other use cases as future work.
Stage 1: Train the Teacher Model
Stage 2: Knowledge Distillation
Figure 3: We adopt a three-stage pipeline for Sparse Distillation: (1) We fine-tune a RoBERTa-Large model on  $D_{train}$  to get the teacher model. (2) We apply teacher model to the unlabeled corpus  $C$  and  $D_{train}$ , and train the student model (DAN) to mimic the predictions of the teacher. This model is denoted as "DAN (KD)" (3) We further fine-tune the student model with  $D_{train}$ . This model is denoted as "DAN (KD+FT)".
Stage 3: Further Fine-tune
fully connected layers,  $(\mathbf{W}_1,\mathbf{b}_1)$  and  $(\mathbf{W}_2,\mathbf{b}_2)$ , to produces the final logits  $\hat{\mathbf{z}}$ , i.e.,  $\hat{\mathbf{z}} = M_s(\mathbf{x}) = \mathbf{W}_2(\mathrm{ReLU}(\mathbf{W}_1\mathbf{h} + \mathbf{b}_1)) + \mathbf{b}_2 \in \mathbb{R}^n$ . The logits are transformed into probabilities with the Softmax function, i.e.,  $\hat{\mathbf{y}} = \mathrm{Softmax}(\hat{\mathbf{z}}) \in \mathbb{R}^n$ .
Remarks on Computation Complexity. Multi-headed self-attention is considered the most expensive operation in the teacher transformers, where the computation complexity is  $O(m^2)$  for a sequence with  $m$  sub-word tokens. The student model, Deep Averaging Network (DAN), can be considered as pre-computing and storing phrase representations in a large embedding table. By doing so, the computation complexity is reduced to  $O(m)$ . However, unlike the teacher, the context is limited to a small range, and no long-range information (beyond n-gram) is taken into account by the student model.
2.4 Training Pipeline
Our training pipeline is illustrated in Fig. 3. It has three stages: (1) We first fine-tune a RoBERTa-Large model on the train set  $D_{train}$  of task  $T$ , and use the resulting model as the teacher model. (2) We train the student model by aligning the predictions of the teacher  $(\hat{y})$  and the predictions of the student  $(\hat{y})$  on the union of unlabeled corpus  $C$  and the train set  $D_{train}$ . We align the predictions by minimizing the KL divergence between the two distributions, i.e.,  $L = \sum_{j=1}^{n} \tilde{y}_j \log \frac{\tilde{y}_j}{\hat{y}_j}$ . The resulting student model is denoted as "DAN (KD)". (3) We further fine-tune the student model from step (2) with the task train set  $D_{train}$ , and get a new student model. This model is denoted as "DAN (KD+FT)". This third stage is optional.
2.5 Implementation Details
Determine N-gram Vocabulary. Our student model takes in n-grams as input. We determine
the n-gram vocabulary by selecting the top  $|V|$  frequent n-grams in  $D_{train}$  and  $C$ . For each downstream dataset, we compute the vocabulary separately. We use CountVectorizer with default whitespace tokenization in sklearn (Pedregosa et al., 2011) to perform this task. We set n-gram range to be  $(1,4)$  and set  $|V| = 1,000,000$ ,  $d_e = 1,000$ , unless specified otherwise.
Optimization. The architecture of DAN is sparsely-activated, and thus can be sparsely-optimized to reduce memory footprint. To facilitate this, we design a hybrid Adam optimizer, where we use SparseAdam<sup>3</sup> for the sparse parameters (i.e., the embedding layer), and regular Adam for dense parameters. This implementation helps to improve speed and reduce memory usage greatly – we can train a 1-billion parameter DAN with the batch size of 2048 at the speed of 8 batches/second, on one single GPU with 32 GB memory.
Additional Details. Due to space limit, we defer details such as hyper-parameters settings and hardware configurations in Appendix A.
3 Experiment Settings
Downstream Datasets. Following Tay et al. (2021), we mainly use six single-sentence classification datasets as the testbed for our experiments and analysis. These datasets cover a wide range of NLP applications. We use IMDB (Maas et al., 2011) and SST-2 (Socher et al., 2013) for sentiment analysis, TREC (Li and Roth, 2002) for question classification, AGNews (Zhang et al., 2015) for news classification. We use Civil Comments
<sup>3</sup>Source code: https://torch.org/docs/master/generated/torch_optim.SparseAdam.html. Please refer to Appendix B.2 for a brief introduction on SparseAdam.
Dataset D	|Dtrain|	|Ddev|	Avg. l	Distillation Corpus C	|C|
IMDB	25,000	25,000	300	Amazon Reviews and ★	75m
SST-2	67,349	872	11	Amazon Reviews	75m
TREC	5,452	500	11	PAQ	65m
AGNews	120,000	7,600	55	CC-News	418m
CCom	1,804,874	97,320	67	Reddit News and ★	60m
WToxic	159,571	63,978	92	★	37m
Table 1: Datasets and Distillation Corpus Used in Our Study.  $|.|$  represents the size of a dataset. "Avg.  $l$  " represents the average number of tokens in the input sentence.  $\star$  represents the unlabeled data released with the original dataset.
(Borkan et al., 2019) and Wiki Toxic (Wulczyn et al., 2017) dataset for toxicity detection.
Knowledge Distillation Corpora. We manually select a relevant unlabeled corpus  $C$  based on the task characteristics and text domain. For example, the IMDB and SST-2 models, which are tasked with classifying the sentiment of movie reviews, are paired with a corpus of unlabeled Amazon product reviews (Ni et al., 2019). TREC, a question classification task, is paired with PAQ (Lewis et al., 2021), a collection of 65 million questions. AGNews, a news classification task, is paired with CCNews corpus (Nagel, 2016). For Civil Comments, a dataset for detecting toxic news comments, we select the News subreddit corpus from ConvoKit (Chang et al., 2020), which is built from a previously existing dataset extracted and obtained by a third party and hosted by pushshift.io. Details of all datasets and corpora are listed in Table 1.
3.2 Compared Methods
To comprehensively evaluate and analyze the n-gram student models, we additionally experiment with (1) training a randomly-initialized DAN model with  $D_{train}$ , without knowledge distillation ("from scratch"); (2) directly fine-tuning general-purpose compact transformers, e.g., DistilBERT (Sanh et al., 2019), MobileBERT (Sun et al., 2020); (3) using other lightweight architectures for the student model, such as DistilRoBERTa (Sanh et al., 2019), Bi-LSTM (Tang et al., 2019) and Convolution Neural Networks (Chia et al., 2019), in task-specific distillation setting. We also quote performance from (Tay et al., 2021) when applicable.
How well can DANs emulate the performance of the teacher? In Table 2, we present the results on 6 single-sentence classification datasets. Firstly, we find that in 5 out of the 6 datasets, the gap between the teacher and the student model is within  $3\%$ . This suggests the power of simple n-gram models may be underestimated previously, as they are typically trained from scratch, without modern techniques such as pre-training and knowledge distillation. This also echoes with a series of recent work that questions the necessity of word order information (Sinha et al., 2021) and self-attention (You et al., 2020), in prevalent transformer architectures. Secondly, we observe that knowledge distillation help close more than half the gap between the teacher model and the student model trained from scratch. The effect is more significant with TREC dataset ( $13\%$  improvement), a 46-way classification problem, whose train set has a small size of 5,452. It is hard to estimate parameters of a large sparse model with merely 5,452 examples; however, supervising it with large-scale corpus and distillation target effectively densified the supervision signals and help address the sparsity issues during model training.
How fast are DANs? We have previously hypothesized that DANs will have superior inference speed due to its simple and sparse architecture. In this section we quantify this advantage by comparing the student model with the RoBERTa-Large teacher model. We also include the baselines listed in §3.2 for a comprehensive comparison. For simplicity, we use BPE tokenizer and re-use the embedding table from RoBERTa-Large for our student Bi-LSTM and CNN model. We use 2-layer BiLSTM with hidden dimension of 4, 64, 256 and 512. For the CNN model, we use one 1D convolution layer with hidden dimension of 128 and context window of 7.
We provide speed comparison across all datasets in Table 3. We provide more fine-grained comparison on IMDB dataset in Table 4 and Fig. 1. DAN achieves competitive performance and the fastest inference efficiency among all different student model architectures. The speed-up differs across datasets, ranges from  $4\mathrm{x}$  to  $1091\mathrm{x}$ . It is most significant on Civil Comments  $(1091\mathrm{x})$ , Wiki Toxic  $(668\mathrm{x})$  and IMDB dataset  $(607\mathrm{x})$ , as they have
<sup>4</sup>It is possible that a careful comparison of different distillation corpora can result in better performance. For the purpose of this study, we leave this as future work.
Model	IMDB	SST-2	TREC	AGNews	CCom	WToxic	QQP
DAN (from scratch)	88.3	79.5	78.4	91.1	95.7	92.2	82.0
DAN (KD)†	92.0	87.0	91.8	90.0	96.2	93.9	63.2
DAN (KD)	93.2	86.4	91.8	90.6	96.3	94.0	84.1
DAN (KD+FT)	93.5	88.5	92.6	93.0	96.3	92.5	84.2
DistilBERT (Sanh et al., 2019)	92.2	90.8	92.8	94.5	96.9	93.1	89.4
MobileBERT (Sun et al., 2020)	93.6	90.9	91.0	94.6	97.0	93.5	90.5
Transformer-Base (Tay et al., 2021)	94.2	92.1	93.6	93.5	-‡	91.5	-
ConvNet (Tay et al., 2021)	93.9	92.2	94.2	93.9	-‡	93.8	-
RoBERTa-Large (Liu et al., 2019)	96.3	96.2	94.8	95.4	96.3	94.1	92.1
Table 2: Performance Comparison on 6 Single-sentence Tasks and 1 Sentence-pair Task. We report accuracy for all datasets. For single-sentence tasks, the gap between the teacher model (RoBERTa-Large) and the n-gram based student model (DAN(KD)/DAN(KD+FT)) is within  $3\%$  in most cases. Also, we observe that knowledge distillation helps close more than half the gap between the teacher model and the n-gram model trained from scratch. Knowledge distillation is performed without task data  $(D_{train})$ , assuming that the task data is private (see §4.3). The dataset we obtain from public sources differs from the one in Tay et al. (2021).
Model	IMDB	SST-2	TREC	AGNews	CCom	WToxic	QQP
RoBERTa-Large	29 (1x)	298 (1x)	549 (1x)	147 (1x)	35 (1x)	72 (1x)	240 (1x)
DistilBERT	176 (6x)	1055 (4x)	930 (2x)	740 (5x)	188 (5x)	426 (6x)	1201 (5x)
MobileBERT	158 (5x)	736 (3x)	402 (1x)	751 (5x)	187 (5x)	400 (6x)	943 (4x)
DANs	17557 (607x)	3020 (10x)	2236 (4x)	24084 (164x)	38024 (1091x)	48133 (668x)	35708 (149x)
Table 3: Inference Speed Comparison (Unit: samples per second). DANs greatly improves inference speed, with the speed-up ranging from  $4\mathrm{x}$  to  $1091\mathrm{x}$ . Speed-up is most significant with classification tasks with long sequences as input, e.g., Civil Comment, Wiki Toxic, and IMDB.
Parameter Count
Total/Sparse/Dense	IMDB
Acc.	GPU Speed	CPU Speed
RoBERTa-Large	355M/51M/304M	96.3	29 (1x)	1 (1x)
DistilBERT	66M/23M/43M	92.2	176 (6x)	11 (8x)
MobileBERT	25M/4M/21M	93.6	158 (5x)	8 (6x)
*DistilRoBERTa	83M/39M/44M	95.9	176 (6x)	8 (6x)
*LSTM (21-512d)	62M/51M/11M	95.9	362 (12x)	31 (22x)
*LSTM (21-256d)	56M/51M/5M	95.8	665 (23x)	52 (37x)
*LSTM (21-64d)	53M/51M/2M	95.3	818 (28x)	101 (73x)
*LSTM (21-4d)	52M/51M/<1M	93.1	813 (28x)	146 (105x)
*CNN (11-256d)	53M/51M/2M	89.2	3411 (109x)	251 (181x)
*DAN (this work)	1001M/1000M/1M	93.5	17558 (607x)	923 (663x)
longer input sequences, and the complexity grows quadratically with sequence length in transformer models. Moreover, as shown in Table 4, DAN has an acceptable CPU inference speed, which greatly reduce the hardware cost for inference. We believe all these characteristics makes student DAN model as an ideal option for production or real-time use on single-sentence classification tasks.
Simplest is the best: Exploring different design choices for DAN. We try several modifications
Table 4: Detailed Inference Speed Comparison on IMDB. DANs achieves better accuracy and inference speed compared to other lightweight architectures such as LSTMs and CNNs. Moreover, DANs achieves acceptable inference speed on CPUs.  $\star$  indicates the model is trained with task-specific distillation; no  $\star$  indicates the model is trained with direct fine-tuning.
Variations	Acc.	Variations	Acc.
1. Pooling Methods	2. Dense Layers
Mean Pooling (*)	93.2	1000 → 1000 → 2 (*)	93.2
Max Pooling	91.8	1000 → 1000 → 256 → 2	93.1
Attentive Pooling	93.0	1000 → 1000 → 256 → 64 → 2	93.0
Sum	92.9
3. Embedding Initialization	4. Parallel Training
Without initialization (*)	93.2	1 GPU, param. 1b (*)	93.2
With initialization	93.2	2 GPUs, param. 2b	93.1
Table 5: Variations made to the student model and the performance on IMDB.  $\star$  represents the design we adopt in our main experiments.
to our current experiment pipeline, including (1) replace average pooling with max pooling, attentive pooling, or taking sum in the DAN model; (2) pre-compute a n-gram representation by feeding the raw n-gram text to a RoBERTa-Large model, and using the representations to initialize the embedding table of the student model; (3) attach more dense layers in the DAN; (4) use even larger student models by leveraging parallel training across multiple GPUs. More details about these variations are in Appendix B.1. We experiment with IMDB dataset and list the performance in Table 5. In general, we do not observe significant performance improvements brought by these variations. Thus, we keep the simplest design of DAN for all other experiments.
Figure 4: Trade-off between the vocabulary size and the embedding dimension. Given a fixed parameter budget, empirical results suggest that a larger embedding dimension and a smaller vocabulary size should be selected.
4.2 Controlling the Parameter Budget
Given a fixed parameter budget, how to allocate it wisely to achieve optimal performance? We discuss this question in two scenarios: the users wish to control the parameter budget (1) during knowledge distillation (KD), or (2) during inference.
During KD: Trade-off between vocabulary size and embedding dimension. We explore how the configuration of vocabulary size and embedding dimension influence the student model performance. We train student models on the IMDB dataset with 19 configurations, and show the results graphically in Figure 4. Detailed results are deferred in Table 8 in Appendix C. All else being equal, having more parameters in the student model is beneficial to the performance. For a fixed parameter budget, higher accuracy was achieved by increasing the embedding dimension and making a corresponding reduction in the vocabulary size. Our best performing model has  $|V| = 1,000,000$  and  $d_{e} = 1,000$ . We keep this configuration for the main experiments in previous sections.
During inference: Reduce the model size with n-gram pruning. The model size of DANs is flexible even after training, by excluding the least frequent n-grams in the vocabulary. We test this idea on IMDB and AGNews dataset and plot the performance in Fig. 5. We try two ways to estimate n-gram frequency: (1) using distillation corpus  $C$  and the training set  $D_{train}$ ; (2) using  $D_{train}$  only. We observe that: (1) n-gram frequencies estimated on  $D_{train}$  are more reliable, as  $D_{dev}$  has a n-gram distribution more similar to  $D_{train}$  compared to  $C + D_{train}$ ; (2) DANs maintain decent accuracy ( $>90\%$ ) even when the model size is cut to  $3\%$  of
Figure 5: Post-hoc pruning according to n-gram frequency. We disable the least frequent n-grams during inference to further reduce model size. When the n-gram frequencies are estimated appropriately, DANs maintain decent performance (acc. > 90%) even when model is 3% of its original size.  $C + D_{train} / D_{train}$  represent different ways to estimate n-gram frequencies.
its original size. In this case, users of DANs can customize the model flexibly based on their needs and available computational resources.
4.3 Privacy-preserving Settings