Fast Vocabulary Transfer for Language Model Compression (2022)

Abstract: Real-world business applications require a trade-off between language model performance and size. We propose a new method for model compression that relies on vocabulary transfer. We evaluate the method on various vertical domains and downstream tasks. Our results indicate that vocabulary transfer can be effectively used in combination with other compression techniques, yielding a significant reduction in model size and inference time while marginally compromising on performance.

## introduction in the last few years, many nlp applications have been relying more and more on large pre-trained language models (lm) @xcite @xcite . because larger lms, on average, exhibit higher accuracy, a common trend has been to increase the model's size. some lms like gpt-3 @xcite and bloom 1 have reached hundreds of billion parameters. however, these models' superior performance comes at the cost of a steep increase in computational footprint, both for development and for inference, ultimately hampering their adoption in real-world business use-cases. besides models that only a few hi-tech giants can afford, like gpt-3, even smaller lms with hundreds of million parameters could be too expensive or infeasible for certain products. for one thing, despite being tremendously cheaper than their bigger cousins, fine-tuning, deploying and maintaining large numbers of such models (one for each downstream task) soon becomes too expensive. furthermore, latency and/or hardware requirements may limit their applicability to specific 1 https://bigscience.huggingface.co/blog/bloom use-cases. for all these reasons, significant efforts -in both academic and industry-driven researchare oriented towards the designing of solutions to drastically reduce the costs of lms. recently, several attempts have been made to make these models smaller, faster and cheaper, while retaining most of their original performance @xcite . notably, knowledge distillation (kd) @xcite is a teacher-student framework, whereby the teacher consists of a pre-trained large model and the student of a smaller one. the teacher-student framework requires that both the teacher and the student estimate the same probability distribution. while the outcome is a smaller model, yet, this procedure constrains the student to operate with the same vocabulary as the teacher in the context of language modeling. in this work, we explore a method for further reducing an lm's size by compressing its vocabulary through the training of a tokenizer in the downstream task domain. the tokenizer @xcite @xcite is a crucial part of modern lms. in particular, moving from word to subwordlevel, the tokenization solves two problems: vocabulary explosion and unknown words. moreover, the capability to tokenize text effectively in any domain is key for the massive adoption of pre-trained general-purpose lms fine-tuned on downstream tasks. indeed, tokenizers are still able to process out-of-distribution texts at the cost of producing frequent word splits into multiple tokens. however, the language varies significantly in vertical domains or, more generally, in different topics. hence, ad-hoc tokenizers, trained on the domain statistics, may perform a more efficient to-figure 1 : sketch of the vt procedure. first, the vocabulary is constructed on the in-domain data, then an embedding is assigned to each token, transferring information from the pre-trained representations of the general-purpose language model. kenization, reducing on average the length of the tokenized sequences. this is important since compact and meaningful inputs could reduce computational costs, while improving performance. indeed, memory and time complexity of attention layers grows quadratically with respect to the sequence length @xcite . furthermore, a vertical tokenizer may require a smaller vocabulary, which also affects the size of the embedding matrix, hence further reducing the model's size. following this intuition, we propose a vocabulary transfer (vt) technique to adapt lms to in-domain, smaller tokenizers, in order to further compress and accelerate them. this technique is complementary to the aforementioned model compression methods and independent of the type of tokenizer. as a matter of fact, we apply it in combination with kd. our experiments show that vt achieves an inference speed-up between x1.07 and x1.40, depending on the downstream task, with a limited performance drop, and that a combination of vt with kd yields an overall reduction up to x2.76. the paper is organized as follows. after reviewing related works in section 2, we present the methodology in section 3, we then outline the experiments in section 4 and draw our conclusions in section 5. ## related works the goal of model compression is to shrink and optimize neural architectures, while retaining most of their initial performance. research on lm compression has been carried out following a variety of approaches like quantization @xcite , pruning @xcite knowledge distillation @xcite @xcite , and combinations thereof @xcite . a most popular distillation approach in nlp was proposed by @xcite . the obtained model, called distilbert, is a smaller version of bert, with the same architecture but half the layers, trained to imitate the full output distribution of the teacher (a pre-trained bert model). dis-tilbert has a 40% smaller size than bert and retains 97% of its language understanding capabilities. this enables a 60% inference-time speedup. further compression was achieved by @xcite by adding transformer-layer, predictionlayer and embedding-layer distillation. the resulting model, tinybert, is 10 times smaller than bert, with only four layers and reduced embeddings sizes. related methods were proposed @xcite , achieving similar compression rates. all these works focus on the distillation of general-purpose language models. @xcite investigated the interaction between kd and domain adaptation. little focus has been devoted thus far to the role of tokenization in the context of model compression. even in domain adaptation @xcite , the vocabulary was kept the same. both the versatility of the subword-level tokenization, and the constraints imposed by the teacherstudent framework (same output distribution), discouraged such investigations. recently, @xcite presented an approach for transferring the vocabulary of an lm into a new vocabulary learned from new domain, with the purpose of boosting the performance of the fine-tuned model. to the best of our knowledge, we are the first to study vt in the scope of model compression. ## vocabulary transfer let us consider a lm, trained on a general-purpose domain d gen and associated with a vocabulary v gen . such a vocabulary is used by the lm's tokenizer in order to produce an encoding of the input string via an embedding matrix e gen defined on v gen . more specifically, a tokenizer is a function that maps a textual string into a sequence of symbols of a given vocabulary v. let t be a tokenizer associated with a vocabulary v and a string s, we have t : hence, the vocabulary of the tokenizer determines how words in a text are split, whether as words, subwords, or even characters. these symbols, which define the lm's vocabulary, are statistically determined by training the tokenizer to learn the distribution of a dataset. now, let us consider a vertical domain d in , also referred as in-domain. for the reasons discussed earlier, a vocabulary v in specialized on d in itself better fits the language distribution than v gen . unfortunately, with a new vocabulary, embedding representations associated with the tokens of v gen would be lost. thus, vt aims to initialize v in by re-using most of the information learned from the lm pre-trained on d gen . once the new tokenizer t in has been trained on the in-domain dataset d in using a given vocabulary size, t in will be different from the lm's tokenizer t gen . however, the two tokenizers' vocabularies v gen and v in may still have a large portion of their symbols in common. our objective is to transfer most of the information from v gen into v in . to this end, we first define a mapping between each symbol in v in and a set of symbols in v gen . then, we define an assignment criterion, based on the mapping, to obtain the embeddings for the tokens of t in . the existence of a meaningful mapping is the only underlying assumption for vt, which holds for all the phonographic languages. one such criterion, called vocabulary initialization with partial inheritance (vipi), was defined by @xcite . whenever a token is in v in but not in v gen , vipi calculates all the partitions of the new token with tokens from v gen , then takes the minimal partitions and finally averages them to obtain an embedding for the new token. differently, we define a simplified implementation of vipi called fvt for fast vocabulary transfer. instead of calculating all tokenizations, fvt uses a straightforward assignment mechanism, whereby input: he was initially treated with interferon alfa. tgen: he, was, initially, treated, with, inter,##fer, ##on, al, ##fa, . t100: he, was, initially, treated, with, interferon, alfa, . figure 2 : example of different tokenizations using a pre-trained or an adapted tokenizer. in the latter case, domain-specific words are not broken down into multiple word pieces. ## experiments in the experiments we measure the impact of fvt on three main kpis: quality (f1 score), size of the models and speedup in inference. ## conclusion the viability and success of industrial nlp applications often hinges on a delicate trade-off between computational requirements, responsiveness and output quality. hence, language model compression methods are an active area of research whose practical ramifications are self-evident. one of the factors that greatly contribute to a model's inference speed and memory footprint is vocabulary size. vt has been recently proposed for improving performance, but never so far in the scope of model compression. in this work, we run an extensive experimental study on the application of a lightweight method for vt, called fvt. an analysis conducted on various downstream tasks, application domains, vocabulary sizes and on its possible combination with knowledge distillation indicates that fvt enables a strategic trade-off between compression rate, inference speed and accuracy, especially, but not only, in more specialized domains. importantly, fvt appears to be orthogonal to other model compression methods. in the future, we plan to fully integrate vocabulary transfer within knowledge distillation during the learning process in order to maximize the information transfer. we also plan to define a unified metric that combines all the kpis, to facilitate model selection.