Further Compressing Distilled Language Models via Frequency-aware Partial Sparse Coding of Embeddings (2024)

Abstract: Although pre-trained language models (PLMs) are effective for natural language understanding (NLU) tasks, they demand a huge computational resource, thus preventing us from deploying them on edge devices. Researchers have therefore applied compression techniques for neural networks, such as pruning, quantization, and knowledge distillation, to the PLMs. Although these generic techniques can reduce the number of internal parameters of hidden layers in the PLMs, the embedding layers tied to the tokenizer arehard to compress, occupying a non-negligible portion of the compressed model. In this study, aiming to further compress PLMs reduced by the generic techniques, we exploit frequency-aware sparse coding to compress the embedding layers of the PLMs fine-tuned to downstream tasks. To minimize the impact of the compression on the accuracy, we retain the embeddings of common tokens as they are and use them to reconstruct embeddings of rare tokens by locally linear mapping. Experimental results on the GLUE and JGLUE benchmarks for language understanding in English and Japanese confirmed that our method can further compress the fine-tuned DistilBERT models models while maintaining accuracy.

## introduction transformer @xcite -based language models (lms) have been extensively used to solve natural language processing (nlp) tasks via pretrain and fine-tuning @xcite ; the accuracy of the fine-tuned lms can be improved by scaling up the model and pre-training data sizes @xcite . pre-trained lms (plms) thereby became larger and larger, which prevents us from deploying them on resource-constrained environments. thus, we cannot leverage powerful plms to text with privacy concerns in end-user devices or confidential documents in small businesses. * currently, he works for mirai translate, inc. to make plms faster and smaller while maintaining the accuracy, researchers have utilized common compression techniques for neural networks (surveyed in @xcite ); the techniques include pruning, quantization, and knowledge distillation, mainly focusing on compressing hidden layers which occupy the largest part in the plms with deep transformer layers @xcite . in the distilled plms, however, parameters in those other than hidden layers account for a large proportion of the entire parameters (figure 1 ), and most of them are accounted for by the embedding layers. for instance, the parameters of the embedding layer account for about 34.9% of distilbert-base-uncased foot_0 @xcite , whereas they occupy about 21.4% of the original 12-layer bert-base-uncased foot_1 @xcite . therefore, we subject the embedding layers to further compression. in this study, given a plm fine-tuned to the target downstream task, we propose to compress the embedding layer of the plm by using sparse coding of embeddings, which represents embeddings with a sparse linear combination of basis embeddings @xcite . the issue here is that the sparse coding introduces approximation errors, or noises, into the fine-tuned embeddings. to reduce the impact of these noises on the plm's outputs, we perform a frequency-aware partial sparse coding of embeddings; namely, we regard a small number of common token embeddings as basis embeddings to reconstruct the remaining rare token embeddings, as employed in @xcite for recurrent neural network lms. since the embeddings of the recent plms will be contextualized through deep transformer layers and noisy rare token embeddings will be supplemented by intact embeddings of surrounding common tokens, we adopt simple locally-linear embeddings @xcite to choose a few basis (common token) embeddings for each rare token embedding, thereby enabling sparser coding of embeddings. each rare token embedding is thereby represented as a weighted linear sum of the nearest neighbor common token embeddings. finally, we save the weight and the ids of the common tokens to dynamically reconstruct embeddings during inference. we applied our method to english and japanese distilbert models fine-tuned to glue @xcite and jglue datasets @xcite , respectively. we then compared our methods with three baselines; the two of them approximate the same rare token embeddings as our method, @xmath0, respectively. the other approximates the entire embedding layers using sparse vectors to select vectors to sum up from shared chunks of vectors (additive quantization). the contributions of this paper are as follows: ## proposed method the major difficulty in compressing the embeddings of a fine-tuned transformer-based plm is that if the compression introduces some approximation (noises) in the embeddings, they will severely affect the latter processing in the deep transformer layers. @xcite confirmed on sentiment classification that the accuracy of bert decreased greatly (>10%) when one or more words take perturbations mimicking typos. motivated by this observation, we adopt partial sparse coding, which reconstructs only a subset of plm's embeddings whose approximation errors (noises) will not severely affect the plm's behavior. to reduce the memory footprint, we divert common token embeddings to the candidate of basis embeddings that represent the rare token embeddings. our partial sparse coding consists of the following two steps: step 1: splitting vocabularies. we first split the vocabulary of the plm, v, into two portions, v c (source vocabularies) and v @xmath1, in which the embeddings of the source vocabularies (source embeddings are used as basis embeddings in sparse coding to approximate the embeddings of the target vocabularies (target embeddings). step 2: reconstructing target embeddings. we then compute compact representations for the target embeddings, to facilitate the compression, we choose a small subset of size k, n i , among the source embeddings x to approximate each target embedding y @xmath2. we then represent target embedding y i by compact (sparse vector) representations, ŷi = {(j, α ij ) | x @xmath3, namely, k pairs of embedding id @xmath4. we then reduce the target embeddings y by replacing them with their sparse vector representations and dynamically reconstruct the target embedding y i during the inference by referring to ŷi , thus obtaining a model with a smaller embedding layer. the resulting embedding layer is composed of the original parameters (embeddings) x for the source embeddings and the compact representations for the target embeddings y. the modified model has (d -2k)|v r | fewer parameters, 3 which greatly reduces the number of parameters in the embedding layer when k ≪ d and |v c | ≪ |v r |. ## experimental setup we evaluate our frequency-aware sparse coding of embeddings on distilled plms fine-tuned to nlu tasks in terms of the model size and performance. 4 because the target rare token embeddings may not be updated for the target task and are in the semantic space of plm instead of the fine-tuned plm), we may be able to obtain better embeddings by computing weights for the summation in the semantic space of the plm and by using the weights to reconstruct the target embeddings in the semantic space of the fine-tuned plm. however, our preliminary experiments revealed that the fine-tuning did not change the embeddings much, this did not contribute to the accuracy improvements. ## related work in the development of neural network-based nlp, how to embed a sequence of discrete symbols in languages into the continuous space has been an important issue, and various compact representations of embeddings have been explored. in what follows, we first review approaches to compressing word embeddings ( § 5.1). we next introduce finer-grained tokenization than words, which results in compact embedding layers ( § 5.2). we then discuss a method of predicting embeddings of out-of-vocabulary words ( § 5.3). ## conclusions we proposed a simple yet effective sparse coding method to compress the embedding layer of a given fine-tuned plm. we keep common tokens that appear frequently in the fine-tuning data and only compress the embeddings of rare tokens that do not appear in the fine-tuning data. we select only a small subset of the nearest neighbor source (common token) embeddings to approximate the target (rare token) embeddings so that we represent the target embeddings with only a small number of parameters. our experimental results confirmed that our frequency-aware partial sparse coding can greatly compress the embedding layer while preventing performance degradation. our method works effectively without carefully choosing the number of the source embedding for compression. in future work, we will apply a method to select the target tokens for compression from the vocabulary while considering the easiness of reconstruction as well as the frequency. we also plan to apply our method to decoder-only and encoderdecoder lms, although there are issues as stated in the limitations section.