D ISTILLED EMBEDDING : NON - LINEAR EMBEDDING FACTORIZATION USING KNOWLEDGE DISTILLATION
Anonymous authors
Paper under double-blind review
A BSTRACT
Word-embeddings are a vital component of Natural Language Processing (NLP) systems and have been extensively researched.
Better representations of words have come at the cost of huge memory footprints, which has made deploying NLP models on edge-devices challenging due to memory limitations.
Compressing embedding matrices without sacriﬁcing model performance is essential for successful commercial edge deployment.
In this paper, we propose Distilled Embedding, an (input/output) embedding compression method based on low-rank matrix decomposition with an added non-linearity.
First, we initialize the weights of our decomposition by learning to reconstruct the full word-embedding and then ﬁne-tune on the downstream task employing knowledge distillation on the factorized embedding.
We conduct extensive experimentation with various compression rates on machine translation, using different data-sets with a shared word-embedding matrix for both embedding and vocabulary projection matrices.
We show that the proposed technique outperforms conventional low-rank matrix factorization, and other recently proposed word-embedding matrix compression methods.
I NTRODUCTION
Deep Learning models are the state-of-the-art in NLP, Computer Vision, Speech Recognition and many other ﬁelds in Computer Science and Engineering.
The remarkable deep learning revolution has been built on top of massive amounts of data (both labeled and unlabeled), and faster computation.
In NLP, large pre-trained language models like BERT (Devlin et al., 2018) are state-of-the-art on a large number of downstream NLP problems.
The largest publicly available language model to date is trained with 1.6 billion parameters (Keskar et al., 2019).
On machine translation the stateof-the-art models have parameters in the order of millions.
Data privacy and server cost are some major issues, driving research towards deploying these models on edge-devices.
However, running these models on edge-devices, face memory and latency issues due to limitations of the hardware.
Thus, there has been considerable interest towards research in reducing the memory footprint and faster inference speed for these models (Sainath et al., 2013; Acharya et al., 2019; Shu & Nakayama, 2017; Shi & Yu, 2018;
Jegou et al., 2010; Chen et al., 2018; Winata et al., 2018).
The architecture of deep-learning-based NLP models can be broken down into three components.
The ﬁrst component, represents the embedding section, which maps words in the vocabulary to continuous dense vector representations of the words.
For all future references, the ﬁrst component includes both the source and target vocabulary mapping using a common embedding matrix with a shared vocabulary.
The second component, consists of a function f , typically a deep neural-network (Schmidhuber, 2015; Krizhevsky et al., 2012; Mikolov et al., 2010) which maps the embedding representation for different NLP problems (machine-translation, summarization, question-answering and others), to the output-space of function f .
The third component, is the output layer which maps the output of function f to the vocabulary-space, followed by a softmax function.
Most of these NLP problems use linear transformations, in the ﬁrst and the third components of the models.
Since, these components depend upon a large vocabulary-size, they require large number of parameters which results in higher latency and larger memory requirements.
For instance, the Transformer Base model (Vaswani et al., 2017) uses 37% of the parameters in the ﬁrst and third components using a vocabulary size of 50k, and with parameter-tying between the components.
The percentage of parameters increases to 54%, when parameters are not shared between the ﬁrst and third components.
The reader is encouraged to look at Table A.1 of the Appendix for more information.
Thus, an obvious step is to compress parametric functions used by the ﬁrst and third components.
Recently, many researchers have worked on compressing word-embedding matrices (Sainath et al., 2013; Acharya et al., 2019; Shu & Nakayama, 2017; Shi & Yu, 2018;
Jegou et al., 2010; Chen et al., 2018; Winata et al., 2018).
These techniques have proven to perform at-par with the uncompressed models, but still suffer from a number of issues.
First , embedding compression models (Shi & Yu, 2018;
Chen et al., 2018; Khrulkov et al., 2019;
Shu & Nakayama, 2017), require two hyper-parameters to be ﬁne-tuned.
These hyper-parameters inﬂuence the number of parameters in the model, and thus the compression rate.
This leads to an additional layer of complexity for optimizing the model for different NLP problems.
Additionally,
Chen et al.
(2018) requires an additional optimization step for grouping words, and lacks end-to-end training through back-propagation.
Shi & Yu (2018) also requires an additional step for performing k-means clustering for generating the quantization matrix.
Thus, most of the current state-of-the-art systems are much more complicated to ﬁne-tune for different NLP problems and data-sets.
Second , all the state-of-the-art embedding compression models compress the ﬁrst and third components separately.
In practice, state-of-the-art NLP models like Vaswani et al.
(2017) have shown to perform better with parameter sharing between the ﬁrst and third components (Press & Wolf, 2016).
Thus, there is a need for an exhaustive analysis of various embedding compression techniques, with parameter sharing.
Lastly , embedding compression models not based on linear SVD (Khrulkov et al.
; Shu &
Nakayama (2017); Shi & Yu (2018)) require the reconstruction of the entire embedding matrix or additional computations, when used at the output-layer.
Thus, the model either uses the same amount of memory as the uncompressed model or requires additional computation cost.
This makes linear
SVD based techniques more desirable for running models on edge-devices.
In this paper, we introduce Distilled Embedding, a non-linear matrix factorization method, which outperforms the current state-of-the-art methods.
Our method, ﬁrst compresses the vocabulary-space to a much smaller dimension compared to the original hidden-dimension, then applies a non-linear activation function, before recovering the original embedding-dimension.
Additionally, we also introduce an embedding distillation method, which is similar to Knowledge Distillation (Hinton et al., 2015) but we apply it to distill knowledge from a pre-trained embedding matrix and use an Lloss instead of cross-entropy loss.
To summarize our contributions are:
• We demonstrate that at the same compression rate our method outperforms existing stateof-the-art methods.
• Our proposed method is much simpler than the current state-of-the-methods, with only a single hyper-parameter controlling the compression rate.
• We introduce an embedding distillation method, which out-performs standard machine translation cross-entropy loss.
• Unlike the current state-of-the-art systems, we compress the embedding matrix with parameter sharing between the ﬁrst and third components.
We perform an exhaustive comparison of current state-of-the-art models in this setting.
• Our model uses a similar factorization as SVD, with the addition of a non-linearity, thus it is able to leverage beneﬁts of linear SVD based methods discussed above.
• We also compare all the models against a simple low-rank matrix factorization using Singular Value Decomposition (SVD) and demonstrate that it can compete with more complex algorithms.
R ELATED W ORK
We can model the problem of compressing the embedding matrix as a matrix factorization problem.
There is a considerable amount of work done in this ﬁeld, some of the popular works broadly belong to the domain of low-rank factorization (Singular Value Decomposition (SVD);Srebro &
Jaakkola (2003); Mnih & Salakhutdinov (2008), product quantization (Jegou et al., 2010) and tensor decomposition (De Lathauwer et al., 2000).
A number of prior works in embedding compression are inﬂuenced by these ﬁelds and have been applied to various NLP problems.
In this Section, we will discuss some of the signiﬁcant works across different NLP problems.
Low-rank Factorization Low-rank approximation of weight matrices, using SVD, is a natural way to compress deep learning based NLP models.
Sainath et al.
(2013) apply this to a convolutional neural network for language modeling and acoustic modeling.
Winata et al.
(2018) use SVD on all the weight matrices of an LSTM and demonstrate competitive results on question-answering, language modeling and text-entailment.
Acharya et al.
(2019) use low-rank matrix factorization for word-embedding layer during training to compress a classiﬁcation model.
However, they do not study the effects of applying a non-linear function before reconstructing the original dimension.
GroupReduce Chen et al.
(2018) apply weighted low-rank approximation to the embedding matrix of an LSTM.
They ﬁrst create a many-to-one mapping of all the words in the vocabulary into g groups, this initial mapping is done based upon frequency.
For each group g they apply weighted
SVD to obtain a lower rank estimation, the rank is determined by setting the minimum rank and linearly increasing it based upon average frequency.
Finally, they update the groups by minimizing the reconstruction error from the weighted SVD approximation.
They demonstrate strong results on language modeling and machine translation compared to simple SVD.
In their models they use different embedding matrices for input and softmax layers and apply different compression to each.
Product Quantization Jegou et al.
(2010) introduced product quantization for compressing high dimensional vectors, by uniformly partitioning them into subvectors and quantizing each subvector using K-means clustering technique.
Basically, product quantization approach assumes that the subvectors share some underlying properties which can be used to group similar subvectors together and unify their representation.
That being said, this approach breaks the original matrix into a set of codebooks coming from the center of the clusters in different partitions together with a separate index matrix which refers to the index of the clusters for each subvector.
Shi & Yu (2018) applied product quantization to a language model and were able to show better perplexity scores.
Shu &
Nakayama (2017) extended this technique by ﬁrst representing the product quantization as a matrix factorization problem, and then learning the quantization matrix in an end-to-end trainable neural network.
Li et al.
(2018) implement product quantization through randomly sharing parameters in the embedding matrix, and show good results on perplexity for an LSTM based language model.
Tensor Decomposition De Lathauwer et al.
(2000) introduced multilinear SVD, which is a generalization of SVD for higher order tensors.
Oseledets (2011) introduced an efﬁcient algorithm Tensor
Train (TT) for multilinear SVD Tensor.
Novikov et al.
(2015) applied the Tensor Train decomposition on fully connected layers of deep neural networks.
Khrulkov et al.
(2019) applied Tensor
Train algorithm to the input embedding layer on different NLP problems like language modeling, machine translation and sentiment analysis.
They demonstrate high compression rate with little loss of performance.
However, they compress only the input embedding and not the softmax layer for language modeling and machine translation.
Knowledge Distillation Knowledge distillation has been studied in model compression where knowledge of a large cumbersome model is transferred to a small model for easy deployment.
Several studies have been studied on the knowledge transfer technique (Hinton et al.
; Romero et al. (2015)).
In this paper, we propose a embedding factorization of word-embedding matrix using knowledge distillation to mimic the pre-trained word-embedding representation.
M ETHODOLOGY : D ISTILLED E MBEDDING 3. F UNNELING D ECOMPOSITION AND E MBEDDING D ISTILLATION
We present an overview of our proposed method in Figure 1.
Given an embedding matrix E ∈ R |V|× d , we can decompose it into three matrices (Equation 1), using the SVD algorithm
[Equation]
(1) where |V| is the vocabulary size and d is the embedding dimension.
Σ is a diagonal matrix containing the singular values, and matrices U and V represent the left and right singular vectors of the embedding matrix respectively.
We can obtain the reduced form of the embedding matrix by only keeping r ( < d ) largest singular values out of d .
[Equation]
[Figure]
where the matrix U = U Σ .
The reduced form of the embedding matrix will need r × ( |V| + d ) parameters compared to |V| × d .
Our proposed approach in this work, is to apply a non-linear transformation on the matrix U , before reconstructing the original embedding dimension using V (see Figure 1a), as shown in Equation 3,
[Equation]
We use the ReLU as our non-linear function f ( . ) throughout this paper.
We train a sequence to sequence model (Sutskever et al.
(2014); Vaswani et al.
(2017)) with tied input and output embedding (i.e. the output embedding is the transpose of the input embedding matrix ˜ E out =
˜ E T =
V d × r [ f ( U |V|× r ] T .
We train our model end-to-end by replacing the embedding function with Equation 3.
The matrix U and V are trainable parameters, and for the output layer we use ˜ E T , with the parameter sharing.
We train on two losses.
The standard cross entropy loss deﬁned as:
[Equation]
where M is the sequence length,
y i is the one-hot representation for the i th label and p i is the softmax probability of the i th term generated by the decoder.
In addition to the cross-entropy loss, we introduce a novel embedding reconstruction loss (Equation 5), which we refer to as embedding distillation as we distill information from the pre-trained embedding into our model,
[Equation]
where e i and ˆ e i are the embedding vectors corresponding to the i th word in the original embedding matrix E and the reconstructed embedding matrix ˆ E respectively and u i refers to the i th row of the matrix U .
We use Equation 6 as our ﬁnal loss function.
[Equation]
where α ∈ [0 , 1] is a hyper-parameter, which controls the trade-off between reconstruction and cross-entropy loss.
L recon acts as the knowledge distillation loss by which we try to distill information from the original pre-trained embedding layer as a teacher to the funneling decomposed embedding layer as a student.
The training process of our Distilled Embedding method is summarized in
Algorithm 1 .
Step 1)
Pre-training the Embedding Matrix Pre-train the sequence to sequence model with the full embedding matrix for better initialization.
Step 2)
Initializing the Weights of Funneling Decomposition Layer We extract the trained embedding matrix E from Step 1 and train our decomposed matrices U and V on reconstruction loss deﬁned in Equation 5, as shown in Figure 1a.
Step 3)
Embedding Distillation The pre-trained funneling decomposition layer is plugged into the model (replacing the original embedding matrix E ) and the entire model is trained based on Equation 6.
E XPERIMENTAL S ETUP 4. D ATASETS AND E VALUATION
We test our proposed method on machine translation which is a fundamental problem in NLP and challenging for embedding compression since we typically have at least two dictionaires and an input and output embedding.
Comparitively language modeling uses a single dictionary and classiﬁcation tasks such as sentiment analysis don’t have an output embedding.
We present results on translating three language pairs: WMT English to French (En-Fr), WMT
English to German (En-De) and IWSLT Portuguese to English (Pt-En).
We decided that these pairs are good representatives of high-resource, medium-resource and low-resource language pairs.
WMT En-Fr is based on WMT14 training data which contain 36M sentence pairs.
We used SentencePiece (Kudo & Richardson (2018)) to extract a shared vocabulary of 32k subwords.
We validate on newstest2013 and test on newstest2014.
For WMT English to German (En-De), we use the same setup as Vaswani et al.
The dataset is based on WMT16 training data and contains about 4.5M pairs.
We use a shared vocabulary of 37k subwords extracted using SentencePiece.
For the IWSLT Portuguese to English (Pt-En) dataset, we replicate the setup of Tan et al.
for training individual models.
Speciﬁcally, the dataset contains about 167k training pairs.
We used a shared vocabulary of 32k subwords extracted with SentencePiece.
For all language pairs, we measure case-sensitive BLEU score (Papineni et al.
using SacreBLEU 1 (Post (2018)).
In addition, we save a checkpoint every hour for the WMT En-Fr and WMT
En-De language pairs and every 5 minutes for the IWSLT Pt-En due to the smaller size of the dataset.
We use the last checkpoint which resulted in the highest validation BLEU and average the last ﬁve checkpoints based on this.
We use beam search with a beam width of 4 for all language pairs.
4. E XPERIMENT D ETAILS
Hyper-Parameters For WMT En-Fr and WMT En-De
we use the same conﬁguration as Transformer Base which was proposed by Vaswani et al.
Speciﬁcally, the model hidden size d model is set to 512, the feed-forward hidden size d ff is set to 2048 and the number of layers for the encoder and the decoder was set to 6.
For the IWSLT Pt-En, we use Transformer Small conﬁguration.
Speciﬁcally, the model hidden-size d model is set to 256, the feed-forward hidden size d ff is set to 1024 and the number of layers for the encoder and the decoder was set to 2.
For Transformer
Small, the dropout conﬁguration was set the same as Transformer Base.
All models are optimized using Adam (Kingma & Ba (2015)) and the same learning rate schedule as proposed by Vaswani et al.
We use label smoothing with 0.1 weight for the uniform prior distribution over the vocabulary (Szegedy et al.
; Pereyra et al.
Additionally, we set the value α of Equation 6 to 0.01.
Hardware Details We train the WMT models on 8 NVIDIA V100 GPUs and the IWSLT models on a single NVIDIA V100 GPU.
Each training batch contained a set of sentence pairs containing
[Table]
approximately 6000 source tokens and 6000 target tokens for each GPU worker.
All experiments were run using the TensorFlow framework 2 .
R ESULTS 5. M ACHINE T RANSLATION
We present BLEU score for our method and compare it with SVD,
GroupReduce (Chen et al.
(2018)), Structured Emedding (Shi & Yu (2018)), Tensor Train (Khrulkov et al.
and a smaller transformer network with the same number of parameters.
We learn a decomposition for all the methods except Tensor Train since it was pointed out in Khrulkov et al.
(2019) that there is no difference in performance between random initialization and tensor train learnt initialization.
Once initialized we plug the decomposed embedding and ﬁne-tune till convergence.
None of the weights are frozen during ﬁne-tuning.
Table 1 presents the results on English-French translation.
The hyper-parameters for tuning the competing methods are presented in Section A.1 of the Appendix.
We see that on this task our method along with Structured Embedding performs the best.
Group Reduce is the next, and SVD performs better than Tensor Train, showing that SVD is a strong baseline, when ﬁne-tuned till convergence.
We also compare against a 2 layer neural network (NN) with the same parameterization as distilled embedding which has not been initialized ofﬂine.
The results show that initializing the NN with the embedding weights is important.
On English-German translation, as seen in Table 2, our method outperforms all other methods.
The smaller transformer network does well and is only surpassed by GroupReduce amongst the competing methods.
SVD again performs better than Tensor Train.
We present the Portuguese-English translation results on Table 3.
This task presents a problem where the embedding matrix constitutes the majority of the parameters of the neural network.
The embedding dimension is smaller (256) compared to the other two tasks but embedding compression yields a BLEU score increase in all methods except Structured Embedding.
This is due to a regularization effect from the compression.
Our model again achieves the highest BLEU score.
On these three experiments we demonstrate that our funneling decomposition method with embedding distillation consistently yields higher BLEU scores compared to the competing methods.
5. A BLATION S TUDY
We present different experiments to demonstrate the effect of 1) Model Initialization, 2) Embedding
Distillation, 3) Fine-tuning strategies, 4) Compression capability and 5) Extension and generality of our method. [Table]
[Table] Initialization We do an ablation study on all the three language pairs deﬁned in Section 4.1, to conclude, if random initialization is better than model based initialization.
We conclude that model based initialization, consistently performs better (Table 4).
Embedding Distillation Table 5 presents different compression rates on the Pt-En task, and embedding distillation performs better across all of them.
In Table 4, we see that across all language pairs when we initialize our model using weights from the funneling decomposition, we improve when using Embedding Distillation during ﬁnetuning.
We performed embedding distillation with random initialization only on the smaller Pt-En dataset and observed that Embedding Distillation improves BLEU score.
Compression Rate We demonstrate in Table 5 that it is possible to compress the embedding up to 15.86x with only a 2% drop in BLEU score for Pt-En.
Re-training Fine-tuning is an important component in our method and we demonstrate through our experiments that at convergence most of the techniques are close in performance.
Table 6 shows that freezing embedding weights and re-training the network weights or vice versa leads to a sharp [Table] [Table] [Table]
[Table] [Table] drop in BLEU score, thus, we need to re-train all the weights. The use of a non-linearity and adding embedding distillation also improves BLEU score
Extension We experimented with applying two key lessons from our method, namely, using a non-linear function and embedding distillation, to a model initialized with group partitions of the
GroupReduce method (Chen et al., 2018), we refer to this method as GroupFunneling .
Tableshows that, GroupFunneling achieves a higher BLEU score on Pt-En compared to GroupReduce.
D ISCUSSION
Importance of Non-linearity We postulate that only a subset of word vector dimensions, explains most of the variance, for most word vectors in the embedding matrix.
Thus, using ReLU activation might help in regularizing the less important dimensions for a given word vector.
Importance Reconstruction Loss We propose that the embedding reconstruction might suffer from adding the ReLU activation function.
Thus, adding a loss for embedding reconstruction helps in grounding the embedding and not loose a lot of information.
Thus, the amount of regularization is controlled by the hyper-parameter α .
Our intuition is partly justiﬁed by results shown in Table
A.2, as reconstruction loss performs worse without the ReLU activation function.
C ONCLUSION AND FUTURE WORK
In this paper we proposed Distilled Embedding, a low-rank matrix decomposition with non-linearity in the bottleneck layer for a shared word-embedding and vocabulary projection matrix.
We also introduce knowledge distillation of the embedding during ﬁne-tuning using the full embedding matrix as the teacher and the decomposed embedding as the student.
We compared our proposed approach with state-of-the-art methods for compressing word-embedding matrix.
We did extensive experiments using three different sizes of datasets and showed that our approach outperforms the state-ofthe art methods on the challenging task of machine translation.
For future work, we will apply our approach to compress feed-forward and multi-head attention layers of the transformer network.
A A PPENDIX A. A DDITIONAL H YPER - PARAMETERS
WMT En-Fr Smaller Transformer Network denotes a network with the same conﬁguration as
Transformer Base but with hidden size d model of 416.
For GroupReduce, to match the same compression rate we used number of clusters c being equal to 10 and minimum rank r min to be 22.
For SVD, we decided to set the rank to 64.
For Tensor Train, we set the embedding shape to be [25 , 32 , 40] × [8 , 8 , 8] and the Tensor Train Rank to be 90.
For structured embedding we use group size as 32 and number of clusters as 2048, we then use the quantization matrix and learn the clusters from scratch.
WMT En-De Smaller Transformer Network denotes a network with the same conﬁguration as
Transformer Base but with hidden size d model of 400.
For GroupReduce, to match the same compression rate we used number of clusters c being equal to 10 and minimum rank r min to be 23.
For SVD, we decided to set the rank to 64.
For Tensor Train, we set the embedding shape to be [25 , 37 , 40] × [8 , 8 , 8] and the Tensor Train Rank to be 90.
For structured embedding we use group size as 32 and number of clusters as 2376, we then use the quantization matrix and learn the clusters from scratch.
IWSLT Pt-En Smaller Transformer Network denotes a network with the same conﬁguration as
Transformer Small but with hidden size d model of 136.
For GroupReduce, to match the same compression rate we used number of clusters c being equal to 15 and minimum rank r min to be 30.
For SVD, we decided to set the rank to 64.
For Tensor Train, we set the embedding shape to be [25 , 32 , 40] × [8 , 4 , 8] and the Tensor Train Rank to be 125.
For structured embedding we use group size as 32 and number of clusters as 4048, we then use the quantization matrix and learn the clusters from scratch.
A.
P ARAMETER COUNT
Table A.1 presents the the number of parameters in the different transfomer layers for the transformer base architecture.
[Table]
A. SVD WITH RECONSTRUCTION LOSS
Table A.2 presents BLEU scores for SVD on English French translation when combining reconstruction loss with SVD.
This result suggests that SVD does not perform well with reconstruction loss.
Model
[Table]
A. C O - RELATION OF R ECONSTRUCTION L OSS AND T RANSLATION
We present the L2 reconstruction losses for factorization of embedding matrices in table A.3.
We see that a lower reconstruction loss does not necessarily lead to better translation performance.
e.g.
[Table]
A. H YPER - PARAMETER S ENSITIVITY
We present the sensitivity of our algorithm to the hyper-parameter α for Pt-En translation in table A.4.
We observe that the algorithm works well over a wide range of α values.
In this work we did not tune α for different datasets but doing so might yield higher performance as is the case for
Pt-En translation with α = 0 .
[Table]