Weight Squeezing: Reparameterization for Model Compression Anonymous submission
Abstract
In this work, we present a novel approach for simultaneous knowledge transfer and model compression called Weight Squeezing .
In this method, we perform knowledge transfer from a pre-trained teacher model by learning the mapping from its weights to smaller student model weights, without signiﬁcant loss of model accuracy.
We apply Weight Squeezing combined with
Knowledge Distillation to a pre-trained text classiﬁcation model and compare it to various knowledge transfer and model compression methods on several downstream text classiﬁcation tasks.
We observe that our approach produces competitive results.
We also report that for some experiments, Weight Squeezing without Knowledge Distillation outperforms other baselines.
Introduction
Today, deep learning has become a key technology in natural language processing (NLP), advancing state-of-the-art results for most NLP tasks.
One of the signiﬁcant achievements in applying deep learning in NLP is transfer learning.
These methods include word vectors pre-trained on a large volume of text (Mikolov et al., 2013; Pennington et al., 2014), which are now commonly used to initialize the ﬁrst layer of neural networks for transfer learning.
In recent times, methods such as ULMFiT (Howard and
Ruder, 2018) have advanced the ﬁeld of transfer learning in NLP by using language modeling during pre-training.
Pre-trained language models now achieve state-of-the-art results on a diverse range of tasks in NLP, including text classiﬁcation, question answering, natural language inference, coreference resolution, sequence labeling, and more (Qiu et al., 2020).
Another breakthrough in NLP occurred after the introduction of Transformer (Vaswani et al., 2017) neural networks that consist mostly of linear layers that compute the attention between model inputs.
Unlike recurrent networks, Transformers have no recurrence over the spatial dimensions, making it possible to signiﬁcantly increase their size and reach state-of-the-art results for several tasks.
(Devlin et al., 2019) introduced BERT (Bidirectional
Encoder Representations from Transformers), a language representation model trained to predict masked tokens in texts from unlabeled data.
Pretrained BERT can be ﬁne-tuned to create state-ofthe-art models for a wide range of NLP tasks.
While BERT is capable of learning rich representations of text, using it for solving simple downstream tasks could be excessive.
This is especially important when running downstream models on edge devices such as mobile phones.
A common approach in such cases is model compression.
In this work, we present a novel approach to performing simultaneous transfer learning and model compression that we called Weight Squeezing In this method, we propose learning the mapping from the weights of a text classiﬁcation teacher model based on pre-trained BERT to the weights of a much smaller student model 1 .
We compared the proposed method with common approaches for model compression, including variations of Knowledge Distillation and low-rank matrix factorization.
We evaluated both Weight Squeezing combined with Knowledge Distillation and training Weight
Squeezing on its own.
Our experiments show that in most cases, Weight Squeezing achieves better performance than other baseline methods.
It also outperforms models trained from scratch.
Motivation & Problem Setting
In this section we describe our motivation for building lightweight models.
Firstly, we focus on model size.
For instance, a typical mobile app has the size of 100Mb, which signiﬁcantly restricts the model size that can be used on the device 2 .
Smaller models are also better for serving server-side.
Secondly, we keep an eye on model inference time.
A model may be efﬁcient in terms of training (i.e., by having fewer parameters), but have a signiﬁcant computational overhead regardless.
Therefore, our focus is on making the student model faster than the teacher one.
Finally, we take into account that access to training resources, such as data and computing power, can be limited.
Therefore, by focusing on taskspeciﬁc compression, we limit the data we use to what we have for speciﬁc tasks only (see 3.4.3 for details).
Related Work 3. Transfer Learning & Unsupervised Pre-training
Self-supervised training as unsupervised pretraining became one of the key techniques for solving NLP tasks.
While the amount of labeled data can be limited, we often have unlabeled data at our disposal.
This data can be effectively utilized for pre-training some parts of the models.
One way to conduct self-supervised pre-training of NLP models is language modeling.
It can be performed using autoregressive models when learning to predict the next words in a text based on previous words.
3. Transformer
The Transformer was introduced as a type of neural network architecture for machine translation (Vaswani et al., 2017).
Transformers consist of two parts: Encoder and Decoder.
The former is used to read an original sentence to gather information for the latter, which outputs the translation as an autoregressive model conditioned on encoder outputs.
Both Encoder and Decoder blocks have the same structure and consist of the same layers performing self-attention over the inputs.
For a detailed description of the architecture, see A.2.
3.2. BERT
While the Encoder part of the Transformer architecture could be removed to train a language model as unsupervised pre-training, this model’s performance can suffer because of its autoregressive nature.
Masked Language Modeling (MLM) was proposed as an alternative to Language Modeling.
Instead of training the model to predict the next word in a sentence based on previous words, we train a model to predict masked words in a text based on unmasked ones.
Since this model does not perform autoregressive computations, its output depends on all words in the input text, which makes it capable of learning more complex data patterns.
The BERT model (Devlin et al., 2019) implements the MLM approach for unsupervised pretraining using the Transformer architecture.
After the model is trained, it can be used for extracting features from text.
3. Fine-tuning
It was shown (Devlin et al., 2019) that ﬁne-tuning a pre-trained BERT model outperforms training a separate model from scratch.
Fine-tuning allows us to use learned representations to solve downstream tasks.
However, in such cases, the network architecture could be considerably over-parameterized, and therefore inefﬁcient in terms of memory and computing resources.
Trading accuracy for performance and model size is often reasonable, which brings us to model compression techniques.
3. Model Compression
There are many approaches (Ganesh et al., 2020;
Qiu et al., 2020) to compressing BERT, including pruning, quantization, parameter sharing, module replacing, and knowledge distillation.
Some of these methods can be combined to achieve better results (Mao et al., 2020).
3.4. Pruning
Pruning approaches remove less important parameters from the model.
(Sajjad et al., 2020; Fan et al., 2019) explore several strategies to dropping layers.
(Guo et al., 2019; Gordon et al., 2020; Voita et al., 2019; McCarley et al., 2019) prune some parts of a model after training.
3.4. Low-rank Matrix Factorization
Low-rank matrix factorization approaches focus on reducing the size of model parameters.
These ap- proaches include Singular-Value Decomposition (SVD), Tensor Train (TT) Decomposition (Oseledets, 2011), and others.
One of the most notable examples of using lowrank matrix factorization methods for NLP is reducing the size of the embedding matrix, which contains a signiﬁcant part of model parameters.
Instead of projecting one-hot vectors directly into the hidden space, (Lan et al., 2019)
ﬁrst project them into a lower-dimensional embedding space, and only then project it into the hidden space.
(Khrulkov et al., 2019) parameterize embedding layers based on the TT decomposition.
(Shu and Nakayama, 2019) learn discrete codes to represent embeddings.
(Acharya et al., 2019) also use low-rank matrix factorization on the word embedding layer during training.
These methods can be considered a separate case of the pruning approach where we try to reduce the size of parameters by removing parts that can be considered unnecessary.
For SVD, we drop some part of the weights which are related to small singular values, while TT can be seen as a generalization of SVD.
3.4. Knowledge Distillation
In the Knowledge Distillation (KD) approach (Ba and Caruana, 2014; Hinton et al., 2015) a smaller student model is trained to mimic a teacher BERT model.
Current KD approaches for BERT can be categorized by what exactly they try to match as follows: distillation on encoder outputs/hidden states: (Zhao et al., 2019; Jiao et al., 2019; Sun et al., 2020, 2019; Sanh et al., 2019), distillation on model output (Zhao et al., 2019; Sun et al., 2019; Sanh et al., 2019; Jiao et al., 2019;
Chen et al., 2020), distillation on attention maps (Sun et al., 2020; Jiao et al., 2019).
KD can also be split into two categories: taskagnostic and task-speciﬁc.
Task-agnostic KD involves reducing the size of BERT itself.
Most methods fall under this type of compression.
The KD methods of (Sanh et al., 2019; Sun et al., 2020) are trained on the same corpus as the one used when pre-training a BERT model from scratch.
A typical scenario when using models such as BERT is to take a model already pre-trained on a very large corpus of data, since training or ﬁne-tuning BERT is computationally expensive and could require a considerable amount of time.
In some cases, even storing a large corpus, not to mention using it for training, could be a problem in this task.
Instead of compressing BERT itself, a different approach of task-speciﬁc KD can be taken by ﬁne-tuning BERT on a downstream task ﬁrst and then applying compression techniques to train a smaller model.
(Tang et al., 2019) use KD to train a single layer BiLSTM student model from BERT.
(Mukherjee and Awadallah, 2019) show that when given a large amount of unlabeled data, student BiLSTM networks can even match the performance of the teacher.
(Turc et al., 2019) pre-train compact student models on unlabeled data and then apply
KD.
3.4. Other Compression Techniques
Quantization methods (Zafrir et al., 2019; Shen et al., 2019) compress higher-precision parameters to lower-precision ones.
BERT-of-Theseus (Xu et al., 2020) progressively substitutes modules from the source model with modules containing fewer parameters.
ALBERT (Lan et al., 2019) includes cross-layer parameter sharing.
3. Meta-learning & Weights Generation
Another related approach is using a separate model to generate weights for another model, which can also lead to model compression.
(Denil et al., 2013) show that certain parameters of a convolutional neural network (CNN) can be predicted by another neural network.
(Ha et al., 2016) with HyperNetworks are able to generate weights for LSTM and
CNN architectures.
These methods focus on reducing trainable parameters, while for inference, there is a trade-off between the number of parameters and computation time of the model.
They also do not utilize knowledge transfer from a pre-trained model.
Weight Squeezing
We now introduce a method to perform knowledge transfer and model compression by learning the mapping between teacher and student weights .
We start with a pre-trained teacher Transformer model with a large hidden state.
In our experiments, we used teacher models obtained by ﬁne-tuning pre-trained BERT models with a classiﬁer on top.
It implies that for some linear layer l , we have a weight matrix Θ tl with the shape n × m .
We explore a case where the weights of a pretrained teacher model are too big for running and storing the model on an edge device.
For this reason, we may want to train a student model with a smaller number of parameters.
Let’s say that we want the student model to make the weight matrix Θ sl at the same layer l to have the shape equal to a × b , where a < n and b < m .
In this work, we propose reparameterizing student weights Θ sl as follows:
[Equation]
where L and R are randomly initialized trainable parameters of the mapping (here and below, we omitted the l subscript for simplicity.
However, we reparameterize each parameter layer-wise, and each layer has its own mapping).
Note that if Θ s is a × b matrix and Θ t is n × m matrix, then L and
R are a × n and m × b matrices.
In this approach, instead of training student model weights from scratch, we reparameterize them as a trainable linear mapping from teacher model weights.
Doing so allows us to transfer knowledge stored in the teacher weights to the student weights.
At the same time, mapping of teacher biases and word embeddings is performed as a single linear mapping as follows:
[Equation]
where biases are matrices of size 1 × b and word embeddings have size V × b , and V is the total number of words in the vocabulary.
This reparameterization for word embeddings can be seen as a linear alignment from pre-trained embeddings.
After reparameterization of the student model weights using 1 or 2, we train mapping weights L and R using plain negative log-likelihood (Weight
Squeezing) or KD loss (Weight Squeezing combined with KD).
When the mapping is trained, we compute student weights and use them to make predictions.
4. Comparison to Similar Methods 4.1. Knowledge Distillation
Weight Squeezing (WS) is orthogonal to the Knowledge Distillation method, and the two methods can be combined to achieve better results.
We also observed that, in some cases, the stand-alone WS approach performs better than Knowledge Distillation.
We believe that the reason for that is that the nature of the WS approach allows the model to drop some of the information stored in the teacher parameters if it is excessive, similar to what is done by pruning methods.
At the same time, KD methods based on encoder outputs and attention maps (Jiao et al., 2019; Sun et al., 2020, 2019; Sanh et al., 2019) involve training student models to mimic the behavior of the teacher, even if the capacity of the small student model is not enough to store the knowledge of the larger teacher model.
To address this issue, we can use additional hyperparameters to evaluate the loss function for the KD method (e.g., extra weights to add balance between plain negative log-likelihood, KL-Divergence between the student and teacher predictions, and other regularization terms).
However, it makes the training process more complex and less stable, while
Weight Squeezing does not have this problem.
4.1. Low-rank Matrix Factorization
Compared to low-rank matrix factorization methods, we do not lose any information stored in the teacher weights when using Weight Squeezing.
In essence, even if we map teacher model weights to the weights of a signiﬁcantly smaller student model, we still have the teacher parameters represented as-is, thus allowing the student model to use all the information stored in them.
In comparison, in low-rank matrix factorization methods, we lose part of the knowledge stored in the teacher model every time we reduce the rank of the factorization.
It leads to a trade-off between the number of parameters in the student model and the amount of teacher knowledge that we can transfer.
Consider two teacher models, with 1 , 000 ,and 10 , 000 , 000 parameters, respectively.
Let’s assume that both teachers have the same structure, number of layers and other parameters, with the only difference being the size of their hidden layers.
We want to compress both of them into a student model with 100 , 000 parameters.
The compression of a smaller teacher model could be signiﬁcantly more accurate than the compression of a bigger one, since the rank of the factorization will be different for these cases.
This suggests that factorization methods can be an inappropriate choice for some compression tasks.
Furthermore, in our experiments, we observed that factorization methods could be hard to train due to their large memory footprint.
For example, we want to apply factorization using SVD (see 5.2.5 for details), for which we deﬁne the weights of linear layers as follows:
[Equation]
We now need to evaluate the result of applying this layer to some input x ˆΘ (where x and ˆΘ are b × m and m × n matrices, respectively, and b is the batch size of the input).
The ﬁrst option is to evaluate a full weight matrix ﬁrst and then multiply the resulting matrix by
[Equation]
The second option is to evaluate this layer in a sequential manner :
[Equation]
However, for Equation 3, we will create a matrix in the computational graph of size m × n (weight size of the teacher model), which we have to store for backward propagation and for which we have to evaluate gradient.
While for Equation 4 we do not have to store a big matrix in the computational graph, we do have to store a b × r matrix, the size of which, for now, depends on the batch size.
Therefore, if the batch size is big enough, methodcould require even more memory for training than in Equation 3.
For TT Decomposition, this leads to a more drastic difference in memory used for model training.
Since the evaluation of matrix multiplication for
TT cores can be deﬁned sequentially, like for Equation 4 (Novikov et al., 2020), it will create d −matrices (where d is the number of cores), the size of which depends on the batch size.
We would also like to note that the factorization rate r cannot be smaller than 1 , which is why for every speciﬁc model there is a practical limit to the compression that can be performed.
Furthermore, regardless of the factorization rate r , the result after applying the factorized layer will always retain its original size.
Thus, some operations in Transformer are not going to beneﬁt from the factorization of the model weights.
For example, if we have a Transformer with the size of the hidden layer equal to 1024 , then even with r equal to 1 , self-attention will be computed between vectors with shape 1024 .
Because of this, low rates of r did not lead to a performance boost.
In our experiments, we were able to evaluate layers sequentially for the SVD method when the teacher models were small.
However, in some cases, we had to evaluate the full layer weights to train student models.
We always had to evaluate the full weights for TT models.
Weight Squeezing is also a method that requires evaluating the full weight matrix ﬁrst.
However,
WS does not lead to a large memory footprint since the resulting layer’s size is smaller than that of the original teacher network.
For inference, we could pre-evaluate student weights once and then use them without any computational overhead.
Therefore, in some cases, low-rank methods could be considered an inappropriate choice to train small models for running them in an environment with limited resources.
Nevertheless, we still consider these methods as baselines and provide their experimental results for reference.
Experiments 5. Training Details
We trained all models on AG’s News , DBPedia , and Yelp Reviews datasets.
The datasets are described in A.1.
We used manual tuning to select the model hyperparameters, with up to 20 hyperparameter search trials.
Hyperparameter search ranges are shown in Table 9.
The hyperparameters used in the bestperforming models are shown in 5.2 and 5.3.
We selected hyperparameter conﬁgurations with the best performance, measured by model accuracy on the validation set during the training phase.
We then evaluated the models on the test set and showed the results.
We repeated each experimenttimes and showed the mean and standard deviation values of the measurements.
Since we focused on making models smaller in terms of the overall number of parameters, we trained classiﬁers with 6 Transformer layers, 2 selfattention heads, and hidden size selected from the range [16 , 32 , 64 , 128] (full list of the number of the model parameters can be found in Table 8).
For knowledge transfer, we use teachers ﬁnetuned to downstream tasks.
The teacher models were initialized from BERT models pre-trained on the Masked Language Modelling task.
For each task, we experimented with two teacher models with 6 layers, 768 and 128 hidden sizes, and 12 and 2 self-attention heads, respectively (we will refer to these sizes as 768 H 12 A and 128 H 2 A later).
5. Baselines
We experimented with various approaches to applying Weight Squeezing.
These approaches include:
Applying Weight Squeezing paired with Knowledge Distillation (WS + KD)
[Figure]
AG News
[Figure] [Figure]
DBPedia
[Figure]
Standalone Weight Squeezing (WS)
We compared Weight Squeezing with the following:
Training a student model of an appropriate size from scratch (Scratch) 2.
Knowledge Distillation (KD) 3.
Knowledge Distillation on Encoder Output (KD-EO) 4.
Low-rank Matrix Factorization combined with Knowledge Distillation (SVD + KD, TT + KD)
5.2. Teacher Models
We took BERT models pre-trained on the MLM task with conﬁgurations 768 H 12 A and 128 H 12 A and ﬁne-tuned them on downstream classiﬁcation tasks.
We used these pre-trained teacher models as the source for model compression techniques in our experiments.
5.2. Training from Scratch
For the ﬁrst experiment, we trained models of appropriate hidden sizes from scratch without utilizing any pre-training.
5.2. Knowledge Distillation
We applied the Knowledge Distillation approach to teacher models for training the student model with a smaller hidden size to match the teacher’s predictions (KD in Tables 2, 3, 4, 5).
The resulting loss function is deﬁned as:
[Figure]
[Equation]
where p t and p s are the output probabilities of student and teacher models, p ti , and p si are the ith components of teacher and student predictions, respectively, and c is the index of ground truth label.
Thus, the ﬁrst term is a negative log-likelihood of the student predictions, while the second stands for part of KL-Divergence between teacher and student predictions, which we could optimize with respect to student parameters.
In all experiments, α and β values were equal to 0 .
2 and 1 . 0 , respectively.
We also experimented with Knowledge Distillation on encoder outputs (KD-EO in Tables 2, 3, 4, 5), where we trained the linear mapping of the student’s hidden states to those of the teacher after every Transformer layer.
Training this mapping can be considered a generalization of matching values of student and teacher hidden states when they have different sizes and cannot be compared directly.
The resulting loss function is deﬁned as:
[Equation]
where the ﬁrst two terms in the equation above correspond to terms in Equation 5, while the last term is the mean squared error between teacher hidden state h tj from layer j and mapped student hidden state f j ( h sj ) from the same layer.
We optimized this loss with respect to student parameters and parameters of each mapping f j .
We used γ value equal to 1000 in our experiments, while α and β values differed from the plain KD setting
[Figure] [Figure]
Hidden size
[Figure]
Yelp-polarity
[Figure]
and were equal to 1 and 10 , respectively.
Note that there are several hidden states in each layer corresponding to different words in the input sequence.
For this reason, we used the ﬁrst hidden state of this sequence to evaluate the loss function.
5.2. Weight Squeezing
For plain Weight Squeezing (WS in Tables 2, 3, 4, 5) setting, we used ﬁne-tuned teacher models as the source of mapping for weights reparameterization.
This way, we reparameterized the parameters of all linear layers as in Equation 1 and embedding vectors as in Equation 2.
Weights of the mapping for linear layers were initialized as Xavier Normal, while the mapping for the embedding matrix was initialized as Xavier Uniform.
We trained the resulting model using plain negative log-likelihood.
We optimized the likelihood with respect to the mapping parameters used to reparameterize the student model parameters.
We also optimized the rest of the parameters of the student model which were not reparameterized (e.g, the weights of the layer normalization).
We also trained a model using Weight Squeezing with Knowledge Distillation (WS + KD in
Tables 2, 3, 4, 5).
Instead of using negative loglikelihood as the loss function, we used the KD loss function 5.
In our experiments, we used the same α and β , equal to 0 .
2 and 1 . 0 , respectively.
5.2. Low-rank Matrix Factorization
We also experimented with the matrix factorization approach (LRMF) and applied Singular Value Decomposition (SVD + KD in the results tables) to its weights.
[Equation]
where Σ is a diagonal matrix of the singular values, and U and V are the left and right singular vectors of the weight.
Note that if Θ is m × n matrix, then U is m × m , Σ is m × n and V (cid:62) is n × n matrices.
However, we could obtain the reduced form of this weight by keeping only r largest singular values as follows:
[Equation]
If we were to perform this operation on the weights of the teacher model, then instead of storing nm parameters for each weight, we would have to keep only r ( n + m ) parameters.
If r value is small enough, then the total number of parameters will be reduced compared to the original model, without introducing computation overhead.
In our experiments, we applied SVD to the parameters of the teacher model, obtained its reduced form, and then ﬁne-tuned it using the loss function from the Equation 5 with α and β , equal to 0 .
2 and 1 . 0 , respectively.
We also applied TT Decomposition (TT + KD in the results tables) to compress teacher models.
We decomposed teacher weights using TT withcores.
For training, we also used the loss function from the Equation 5 with the same α and β , equal to 0 .
2 and 1 . 0 , respectively.
Note that in these approaches, we do not directly train the student model with the speciﬁed hidden state size as in all the methods above.
Instead, we injected a bottleneck in the middle of each layer, which allowed us to reduce the total number of parameters in the model.
For this reason, we compared Weight Squeezing trained from the 128H2A teacher BERT model with hidden states with sizes in range [16, 32, 64] with the SVD model with r equal to [10, 22, 50] and with maximum r equal to [25, 50, 128] for the TT approach to reach a similar number of parameters in both of these approaches.
For the 768H12A teacher model with hidden states in the range [16, 32, 64, 128] , we used max r equal to [5, 16, 28, 49] for TT.
For
SVD trained from the 768H12A teacher model, we compared [16, 32, 64, 128] models with r equal to [1, 8, 16, 40] .
5. Optimization
For all models, we used the same optimization strategy using Rectiﬁed Adam (Liu et al., 2020) and Lookahead (Zhang et al., 2019).
For
Lookahead, we used default α equal to 0 .
5 and k equal to 6 .
For all experiments, we selected the best performing learning rate from the range [1e − 5 , 5e − 5 , 1e − 4 , 5e − 4 , 1e − 3 , 5e − 3 , 1e − 2] .
We used a linear learning rate warm-up for the ﬁrst 4000 optimization steps and linearly decayed its value to 0 for the rest of the training.
Betas were equal to (0 . 9 , 0 . 999) .
We also used dropout (Hinton et al., 2012) with a rate equal to 0 .
Gradients with norms larger than 1 were clipped.
We trained the teacher model with the batch size equal to 32 , while other models were trained with the batch size equal to 128 .
More details can be found in A.3.
Results
Plots showing the accuracy of trained models can be found in Figures 1 and 2.
See Tables 2, 3, 4,for the full list of the results.
While using WS combined with KD increases the performance of the model, plain WS outperformed plain KD method in every experiment.
We showed that plain WS method reached results comparable with KD-EO method, achieving better results in some experiments.
We observed that SVD and TT methods were competitive to WS, outperforming most baselines on the Yelp-f and the Yelp-p datasets.
However,
WS outperformed both SVD and TT on the AgNews and the DBPedia datasets in most of the experiments.
In addition, Low-Rank Factorization
Methods (SVD and TT) perform slower than other baselines in all cases (up to 560 times slower, see
A.4 for a detailed report), which makes LRMF methods harder to use when inference speed is important.
Our ﬁndings show that LRMF methods suffer from low factorization rates r .
For example, r equal to 50 for SVD trained from 128 H 2 A teacher leads to signiﬁcantly better results than r equal to 16 for
SVD trained from 768 H 12 A teacher.
Both of these models have an equal number of parameters.
All models outperformed the model trained from scratch, except for LRMF methods with small (and 1 for SVD, and 16 and 5 for TT) factorization ranks trained from 768 H 12 A teacher.
We conducted a study on the expected validation performance on the AgNews for all methods and provided the results in A.5.
We observed that WS achieved better accuracy with a lower number of hyperparameter trials than any other method.
Conclusion & Future Work
We introduced Weight Squeezing, an approach to knowledge transfer and model compression.
We showed that it can be used to compress pre-trained text classiﬁcation models and create compelling lightweight and fast models.
We showed this approach could be a competitive alternative to Knowledge Distillation methods.
While the current work was focused on transferring knowledge to task-speciﬁc models, we would like to apply Weight Squeezing for task-agnostic compression for creating more applicable BERT models trained for tasks such as Masked Language
Modelling.
It may be possible to improve WS performance through a close study of the mappings’ initialization or normalization.
We are also interested in applying this method in domains beyond NLP to compress other types of layers (convolutional, etc.).
A Appendices A. Datasets
AG news dataset.
The AG news corpus consists of news articles from AG’s corpus of news articles on the web 3 pertaining to the 4 largest classes.
DBPedia ontology dataset .
DBpedia is a crowdsourced community effort to extract structured information from Wikipedia (Lehmann et al., 2015).
The DBpedia ontology dataset was constructed by picking 14 nonoverlapping classes from DBpedia 2014.
Yelp reviews.
The Yelp reviews dataset was obtained from the Yelp Dataset Challenge.
This dataset contains review texts with two classiﬁcation tasks – Yelp-f predicting a rating the user has given, and Yelp-p predicting a polarity label.
Splitting & Preprocessing.
We randomly split the data into training and validation sets.
Input sequences were lower-cased and encoded using pretrained BPE with 30522 tokens.
We also truncated input sequences to maximum lengths.
The sizes of these splits and the maximum input length for every dataset are shown in Table 1.
[Table]
A. Transformer Architecture
The key part of the Transformer layer is the SelfAttention mechanism (Lin et al., 2017).
Below, we provide a description of this mechanism and how a
Transformer layer is constructed.
Firstly, to compute the attention score between key and query vectors, the layer uses the dot product operation.
After applying the softmax function to map the result to the [0 , 1] interval, it is multiplied by a value vector:
[Equation]
where d k is the dimension of keys and acts as a scaling factor, and Q, K, V are the query, key, and value vectors.
The above is how we have deﬁned Single Head attention.
To get Multi-Head Attention, the Transformer layer splits the vector spaces into equal parts.
The results of every attention head are concatenated into one vector:
[Equation]
n )
[Equation]
where MHA stands for Multi-Head Attention.
The reason why the mechanism is called SelfAttention is that Transformer Encoders, such as
BERT, take sequence tokens as the query, key, and value vectors.
The vectors computed after the SelfAttention mechanism pass through a residual connection.
The result goes to the Feed-Forward Network (FFN), which is deﬁned as follows:
[Equation]
where W i , b i are trainable parameters of FFN,
X is a n-dimension vector, and σ is a non-linear function.
The output of the Transformer Layer is computed by residual connection with FFN.
A. Optimization Details
For the AG News dataset, we trained the model for 50 epochs, while teacher models were trained forepochs.
For DBPedia , we trained models for 20 epochs.
The teacher model with conﬁguration 768 H 12 A was trained for 2 epochs, while the 128 H 2 A teacher was trained for 10 epochs.
Yelp-f models were trained for 16 epochs.
The 768 H 12 A teacher model was trained for 3 epochs and
the 128 H
A teacher was trained for 7 epochs.
Yelp-p models were trained for 15 epochs with teachers trained for 3 and 5 epochs for 768 H 12 A and 128 H
A conﬁgurations, respectively.
Models with inputs restricted to 128 tokens were trained for approximately 2 - 4 hours, depending on the hidden size, while training models on inputs with lengths equal to 512 took around 10 - 16 hours on a single NVIDIA Tesla T4 GPU.
The exception is TT models, for which we used 8 NVIDIA Tesla
T4 GPUs for training, with input lengths equal to 512 .
A. Inference Speed Measurements
For inference speed measurement, we compared
Non-Low Rank Matrix Factorization models (including KD, KD-EO, WS and models trained from
[Figure]
E xp ec t e d v a li d a ti on acc u r ac y
[Figure]
scratch) with models factorized using the TT and
SVD methods.
In our experiments, we used a classiﬁer with the hidden layer size equal to 16 .
For the TT method, we performed decomposition with 4 cores and a maximum decomposition rank equal to 5 andfor 768 and 128 original model hidden layer size, respectively.
For the SVD method, we decomposed matrices with rank equal to 10.
All models have a comparable number of parameters (see Table 8).
We evaluated models on sequences with input lengths of 128 and 512 .
For GPU measurements, we ran the models on 1000 samples, and onsamples for CPU.
In all experiments, we tested the speed with batch sizes equal to 1 and 16 .
We repeated the measurements 5 times and reported the mean and standard deviation values of the total computation time of the models.
For CPU measurements, we used 1.8 GHz Intel DualCore i5, while for GPU measurements, we used NVIDIA Tesla
T4.
All models were prepared with PyTorch JIT compilation, since we found that it slightly increases the speed of all models in this experiment.
The detailed results are reported in Tables(CPU) and 7 (GPU).
We observed that non-LRMF methods, including WS, outperformed SVD and TT methods in every experiment.
On CPU, SVD is 1.54-20.7 times slower than non-LRMF methods, and TT is 50.9560.0 times slower.
On GPU, SVD is 1.2-8.
times slower, and TT is 5.07-417.0 times slower than non-LRMF methods in the experiments conducted.
A. Expected Validation
Performance
Using the model evaluation method proposed in (Dodge et al., 2019), we calculated the expected validation performance given a number of hyperparameter trials for all methods on the AgNews dataset.
For the experiments, we compressed 128 H 2 A teacher model into a student classiﬁer with a hidden state size of 16 .
For LRMF methods, we used the appropriate factorization ranks to make SVD and TT approaches have approximately the same number of parameters as a plain classiﬁer.
Hyperparameter search ranges are shown in Table 9.
The results are provided in Table 3.
We observed that WS achieved better accuracy with a lower number of hyperparameter trials than any other method.
[Table]
[Table] [Table]
BERT
[Table]
[Table] [Table]
[Table] [Table]
times
[Table]
list here
[Table] [Table]
[Table] [Table]
tuned the