Domain-specific knowledge distillation yields smaller and better models for conversational commerce (2022)

Abstract: We demonstrate that knowledge distillation can be used not only to reduce model size, but to simultaneously adapt a contextual language model to a specific domain. We use Multilingual BERT (mBERT; Devlin et al., 2019) as a starting point and follow the knowledge distillation approach of (Sahn et al., 2019) to train a smaller multilingual BERT model that is adapted to the domain at hand. We show that for in-domain tasks, the domain-specific model shows on average 2.3% improvement in F1 score, relative to a model distilled on domain-general data. Whereas much previous work with BERT has fine-tuned the encoder weights during task training, we show that the model improvements from distillation on in-domain data persist even when the encoder weights are frozen during task training, allowing a single encoder to support classifiers for multiple tasks and languages.

## introduction encoders and language models such as bert (devlin et al., 2019), roberta (liu et al., 2019) and elmo (peters et al., 2017) are the backbone of many nlp technologies. they are typically trained on data from wikipedia, commoncrawl, or large homogeneous collections of text; however, language varies widely in real-world settings and the type of language used in some contexts is not well represented in the data used to train these models. in particular, the language used in e-commerce, and more specifically, conversational commerce, such as conversations pertaining to customer service in the context of online shopping or banking, exhibits both syntactic structures and vocabulary that are under-represented in the wikipedia data used to train multilingual bert. at the same time, these models are too large to deploy in many industry settings, where computational resources and inference-speed are concerns. model size is often reduced using methods such as quantization (whittaker and raj, 2001; shen et al., 2020), pruning (han et al., 2015, 2016) and knowledge distillation (hinton et al., 2015; sanh et al., 2019). however, even leveraging these techniques, the memory footprint of the typical encoder can easily be three orders of magnitude greater than that of the typical classifier, and it follows that encoding is much more time-intensive than classification. @xcite in conversational commerce, a variety of classifiers are required to model different aspects of the conversation. in this case, it is beneficial for efficiency, to use a single encoder for all of the classifiers as illustrated in figure 1 (left), rather than using a separate encoder for each classification task (figure 1 , right). @xcite thus text can be encoded only once and passed to any number of downstream classifiers. typically, domain adaptation with language models is accomplished using back-propagation during task training (see inter alia devlin et al., 2019; liu et al., 2019; sanh et al., 2019). however, this approach requires a separate encoder for each classifier. instead, we adapt the encoder to a particular domain before classifier training. we show that knowledge distillation, a common approach for reducing model size, is very adept for domain adaptation. this allows us to accomplish two goals, classifiers encoder encodings (vectors) messages (texts) classifiers encoder encoder encoder encoder encoder encoder encoder encoder figure 1 : the difference in architecture between using one encoder for multiple tasks vs. one encoder per task size reduction and domain adaptation, with a single training process. evaluating on five languages and two domains, we show that distilling on unlabeled data from the domain of interest results in a smaller model that is domain-specific and outperforms the f1-score of a model distilled on domain-general data by 2.3% on average and the larger teacher model by an f1 of 1.2%. the improvement in performance persists even when relatively little training data is used. we show that the domain-adapted encoder performs better than the domain-general model both when encoder weights are fine-tuned, as in previous work, and when they are frozen, leaving them task agnostic. furthermore, the boost in performance from distillation on in-domain data is greater than the improvement from fine-tuning the encoder during task training. we begin with an overview of previous work in domain adaptation and knowledge distillation, highlighting the benefit of doing both at once ( §2). this is followed by a description of the domains, data and tasks with which we evaluate domain adaptation through knowledge distillation ( §3). we detail our approach in section 4, investigating how much data is necessary ( §5) and examining the impact of domain adaptation on sentence embeddings ( §6). we evaluate on two domains and five languages in section 8, considering both training scenarios where encoder weights are frozen for task training and where they are fine-tuned. ## knowledge distillation method for knowledge distillation, we use the established and open-source approach of sanh et al. (2019), 910 which follows liu et al. (2019)'s proposed best practices for bert training. these include dynamic masking, large batches to leverage gradient accumulation and training on the masked language modeling task but not next sentence prediction. sanh et al.'s implementation of knowledge distillation trains the student model using the distillation loss of the soft target probabilities of the teacher. because of this, the student model benefits from the the teacher model's full distribution during training. due to this rich input, we expect that high performance can be achieved with less training data. we use mbert foot_10 as the teacher model and our student models have the same general architecture, hidden-size dimension and number of word embeddings. we reduce the model size by removing the token-embeddings and pooling and reducing the number of layers from 12 to 6. this reduces the total number of parameters by 43 million or 24% and increases the inference speed by 38% ( ## data requirements because knowledge distillation takes advantage of an existing model, which was already trained on a large amount of data, we expect that distillation training will be relatively economical in its use of data. furthermore, many research objectives focus on a single domain and do not require the breadth of nlu capability of a domain-general model, but instead benefit from a depth of capability in one domain. here we attempt to establish how much data is enough for knowledge distillation for a single domain and where we reach diminishing returns. as a case-study, we use increasing quantities of stackoverflow english data for knowledge distillation and compare the performance of these models to both the teacher model (mbert) and huggingface's multilingual distilled bert model (distilmbert), which was distilled using the same approach. to measure the impact of domain adaptation from knowledge distillation alone, we freeze the encoder weights during task training and present the results in we distilled models with english stackoverflow data, using increments of 0.3gb. we found that a minimum of 1.5gb was needed for convergence, but 2.1gb was enough to outperform distilmbert and perform on par with the teacher mbert. improvements stop after 3gb. we conclude that 2.1 in sanh et al. 2019, however, that paper considers only the english bert model, while we use multilingual bert. because the multilingual model has a significantly larger vocabulary (or number of word embeddings), which is not reduced by this distillation process, the proportionate difference in model size is for distilled mbert models is less than for distilled bert. gb is sufficient and 3gb is optimal for adaptation to the td domain, while more data increases training time without improving performance. we contextualize this finding by considering the data quantity used to train mbert and dis-tilmbert. while it is hard to ascertain the exact amount of data used to train these models, we estimated by following the data sampling procedure used by the creators of those models. @xcite by our best estimate, roughly 222gb of uncompressed text was used. @xcite in contrast, only 2.1gb of uncompressed text was needed to outperform distilmbert on our in-domain task. thus, for distillation for a single domain and language, the required amount of training data is reduced by two orders of magnitude. @xcite in this experiment, the new data matches on both domain and language. to test whether the language match is responsible for the improvement, we distilled a model on 3gb of english wikipedia data. we sampled this data by randomly selecting 3000 1mb chunks of text from the english wikipedia dump. this model (wiki3.0) under-performs the one distilled on the same amount of stackoverflow data, showing that the language match alone is insufficient to explain the improved performance and suggests that the domain match is more important. ## vector changes under domain adaptation to better understand the differences between an encoder trained on domain-general data versus indomain data, we compare sentence embeddings produced by the encoder that we adapted to the technical domain (td3.0) and the encoder distilled on the same amount of domain-general data (wiki3.0). the td domain has lots of homonyms like 'python' and 'float' that have both a technical word-sense and a non-technical one. we expect models trained on the td domain to pay attention to the dominant technical word-senses, and models trained on wikipedia to pay greater attention to the non-technical word-senses. by extension, a distance function derived from a td model is expected to be more sensitive to technical word-senses than a distance function derived from a wikipedia model. thus we expect the distance function for 'python' and a non-technical synonym (i.e., 'snake') to be closer when derived from a domain-general model and the distance function between 'python' and another programming language (i.e., 'php') to be closer when derived from a model trained on technical data. because bert embeddings are contextual, we provide a context for each word pair by creating sentences for each, such that either word may appear in the sentence. sentences designed for technical word pairs are biased towards a technical context, and sentences for non-technical word pairs are biased towards a non-technical sense. @xcite as an example, the technical sentences used to compare @xcite training with 3 gb @xmath0. @xcite the full collection of sentences is in the appendix. java and c# are given below in list items 1 and 2 and the non-technical sentences used for java and coffee are given in list items 3 and 4. substituting each word from a pair into the sentence, we have a pair of sentences like items 1 and 2 and embed each with both the td model and the domain-general model. we take the mean of the token embeddings as a representation of the sentence foot_15 and then take the cosine similarities between the sentence embedding produced by the td model and the embedding from the domain-general model. these cosine similarities are given for both the technical and non-technical pairs in we find that cosine distance is smaller for sentences that capture the general word-sense when they are encoded by the general model. similarly, it is smaller for sentences that capture the technical word-sense when they are encoded with the td model. this suggests that the td model has adapted its representations for these homonyms to their technical meanings. ## experiments we compare the performance of two domainadapted encoders that were trained using the method described in section 4. cc-distilmbert was trained with data from four languages from the cc domain ( §3.1.1) and td-distilmbert with data from five languages from the td domain ( §3.2.1). the amount of data used to distil each model is summarized in for classifier training, we train a structured selfattention classifier head on each encoder used in evaluation. @xcite we conduct two experiments: in the first, we fine-tune the encoder weights, as has been done in previous work such as devlin et al. 2019; sanh et al. 2019; in the second, we freeze the encoder weights to demonstrate that this approach can be used in contexts where the same underlying encoder is to be used by multiple classifiers, removing the need to encode a text every time it is classified by a different classifier. we compare our domain-adapted, distilled models using the tasks described in section 3 with two baselines: the teacher model used for distillation (mbert) and sanh et al.'s domain-general distilled model (distilmbert). in each case, we train and evaluate separate classifiers for each language in the dataset. we evaluate model performance two ways, first fine-tuning the encoder weights during task training and second freezing the encoder weights to test the generalizability of a single encoder to multiple classifiers. ## results the results for each language and encoder are broken down for each experiment in performance is better for all models when the encoder weights are fine-tuned during task training. still, domain-adapted models perform better @xcite code and data for reproduction are included in supplementary materials and will be made publicly available upon publication. each result in for japanese, in the case of the td model, very little data was available for domain adaptation (0.06 gb; shown in ## conclusion we addressed the problem of encoder scalability in the context of conversational commerce by showing that knowledge distillation with domainspecific data reduces model size, while simultaneously improving model performance. this approach allows for the training of an encoder that can be used across a variety of languages, is smaller than a state-of-the-art model like mbert, and performs better on domain-specific tasks. because our approach uses only training data for the domain and languages of interest, less data is necessary for training, reducing the time, cost and environmental impact of training, while accommodating limited data availability. a key advantage of domain adaptation during encoder rather than classifier training is that it allows for the deployment of a single encoder, which can serve multiple classifiers at runtime. this reduces storage and maintenance cost, and due to the much larger size of the encoders compared with the classifiers, provides dramatically better scalability in real-world, e-commerce applications that must support multiple languages and tasks.