The Pupil Has Become the Master: Teacher-Student Model-Based Word Embedding Distillation with Ensemble Learning

1 Introduction

As deep learning starts dominating the field of machine learning, there have been growing interests in deploying deep neural models for real applications. [Hinton et al., 2014] stated that academic research on model development had mostly focused on accuracy improvement, whereas the deployment of deep neural models would also require the optimization of other practical aspects such as speed, memory, storage, power, etc. To satisfy these requirements, several neural model compression methods have been proposed, which can be categorized into the following four: weight pruning [Denil et al., 2013; Han et al., 2015; Jurgovsky et al., 2016], weight quantization [Han et al., 2016; Jurgovsky et al., 2016; Ling et al., 2016], lossless compression [Van Leeuwen, 1976; Han et al., 2015], and distillation [Mou et al., 2016]. This paper focuses on distillation methods that can remarkably reduce the model size, resulting in much less memory usage and fewer computations.

Distillation aims to extract core elements from a complex network and transfer them to a simpler network so it gives comparable results to the complex network. It has been shown that the core elements can be transferred to various types of networks i.e., deep to shallow networks [Ba and Caruana, 2014], recurrent to dense networks [Chan et al., 2015], and vice versa [Romero et al., 2014; Tang et al., 2016]. Lately, embedding distillation was suggested [Mou et al., 2016], which transferred the output of the projection layer in the source network as input to the target network, although accuracy drop was expected with this approach. Considering the upper bound of a distilled network, that is the accuracy achieved by the original network [Ba and Caruana, 2014], enough room is left for the improvement of embedding distillation. Distilled embeddings can significantly enhance the efficiency of deep neural models in NLP, where the majority of model space is occupied by word embeddings.

In this paper, we first propose a new embedding distillation method based on three teacher-student frameworks, which is a more advanced way of embedding distillation, because the previous one [Mou et al., 2016] is a standalone embedding distillation (Section 2.4) with limited knowledge transfer. Our distilled embeddings not only enable the target network to outperform the previous state of the art [Mou et al., 2016], but also are eight times smaller than the original word embeddings yet allow the target network to achieve compatible (sometimes higher) accuracy to the source network. We then present a novel ensemble approach which extends this distillation framework by allowing multiple teacher models when training a student model. After learning from multiple teachers during training, the student model runs on its own during decoding such that it performs faster and lighter than any of the teacher models yet pushes the accuracy much beyond them with just 1.25% (50/4000) the size of other typical ensemble models. All models are evaluated on seven document classification datasets; our experiments show the effectiveness of the proposed frameworks, and our analysis illustrates an interesting nature of the distilled word embeddings. To the best of our knowledge, this is the first time that embedding distillation is thoroughly examined for natural language processing and used in ensemble to achieve such promising results. 1

1 https://github.com/bgshin/distill demo

2 Background

Our embedding distillation framework is based on teacherstudent models [Ba and Caruana, 2014; Sau and Balasubramanian, 2016; Hinton et al., 2014], where teacher models are trained on deep neural networks and transfer their knowledge to student models on simpler networks. The following subsections describe three popular teacher-student methods applied to our framework. The main difference between these three methods is in their cost functions (Figure 1). The last subsection discusses embedding encoding that is used to extract distilled embeddings from the projection layer.

Throughout this section, a logit refers to a vector representing the layer immediately before the softmax layer in a neural network, where zi and vi are the teacher's and student's logit values for the class i, respectively. Note that the student models are not necessarily optimized for only the gold labels but also optimized for the logit values from the teacher models in these methods.

2.1 Logit Matching (LM)

Proposed by [Ba and Caruana, 2014], the cost function of this teacher-student method is defined by the logit differences between the teacher and the student models (D: the total number of classes):

2.2 Noisy Logit Matching (NLM)

Proposed by [Sau and Balasubramanian, 2016], this method is similar to Logit Matching except that Gaussian noise is introduced during the distillation, simulating variations in the teacher models, which gives a similar effect for the student to learn frommultiple teachers. The cost function takes random noise η drawn from Gaussian distribution such that the logit of each teacher model is z′i = (1+η) · zi. Thus, the final cost function becomes LNLM = 1 2D ∑∀i |z′i - vi|2 .

2.3 Softmax TauMatching (STM)

Proposed by [Hinton et al., 2014], this method is based on softmax matching where softmax values are compared between the teacher and the student models instead of logits. Later, [Hinton et al., 2014] added two hyperparameters to further generalize this method. The first hyperparameter, λ, is for the weighted average of two sub-cost functions, where the first sub-cost function measures a cross-entropy between the student's softmax predictions and the truth values, represented as L1 = - ∑∀i yi log pi (i indexes classes, y is the gold label, pi ∈ (0, 1) is the prediction for a sample). Another cost function involves the second hyperparameter, τ, that is a temperature variable normalizing the output of the teacher's logit value:

Given si, the second cost function can be defined as L2 = - ∑∀i si(z, τ) log pi. Therefore, the final cost function becomes LSTM = λL1 + (1 - λ)L2. If λ weights more on L1, the student model values more on the gold labels than teacher's predictions. If τ is greater, the teacher's output becomes more uniformed, implying that the probability values are spread out more throughout all classes.

2.4 Embedding Encoding (ENC)

Embedding distillation was first proposed by [Mou et al., 2016] for NLP tasks. Unlike our framework, their method does not rely on teacher-student models, but rather directly trains a single model with an encoding layer inserted between the embedding layer and its upper layer in the network (Figure 2a). Each word wi is entered to an embedding layer φ that yields a large embedding vector φ(wi). This vector is projected into a smaller embedding space by Wenc with an activation function f. As a result, a smaller embedding vector φ′ (wi) is produced for wi as follows (benc: a bias for the projection):

The smaller embedding φ′ (wi) generated by this projection contains distilled knowledge from the larger embedding φ(wi). The cost function of this method simply measures the crossentropy between gold labels and the softmax output values.

3 Embedding Distillation

Our proposed embedding distillation framework begins by training a teacher model using the original embeddings. After training, the teacher model generates the corresponding logit value for each input in the training data. Then, a student model that comprises a projection layer is optimized for the logit (or softmax) values from the teacher model. After training the student model, small embeddings are distilled from the projection layer (Figure 3a).

The original large embeddings as well as weights in the projection layer are discarded for deployment such that the small embeddings can be referenced directly from the word indices in the student model during decoding (Figure 3b). Such distillation significantly reduces the model size and computations in the network, which is welcomed in production.

3.1 Distillation via Teacher-Student Models

Projecting a vector into a lower dimensional space generally entails information loss, although it does not have to be the case under two conditions. First, the source embeddings comprise both relevant and irrelevant information for the target task, therefore, there is a room to discard the irrelevant information. Second, the projection layer in the target network is capable of preserving the relevant information from the source embeddings. The first condition is met for NLP because most vector space models such as Word2Vec [Mikolov et al., 2013], Glove [Pennington et al., 2014], or FastText [Bojanowski et al., 2017] are trained on a vast amount of text, where only a small portion is germane to a specific task.

To meet the second condition, teacher-student models are adapted to our distillation framework, where the projection layer in a student model learns the relevant information for the target task from a teacher model. The output dimension of the projection layer is generally much smaller than the size of the embedding layer in a teacher model. Note that it is possible to integrate multiple projection layers in a student model; Experiment section discusses performance difference by adding different numbers of projection layers in the student model.

(a) Embedding distillation.

(b) Model deployment.

• (a) Embedding distill ensemble.

(b) Model deployment.

Figure 4: (a) Our proposed embedding distillation ensemble model. One representing logit (R. LOGIT) is calculated from a set of multiple teachers' logits by the proposed ensemble methods. (b) No need to evaluate teachers at deployment, unlike other ensemble methods.

3.2 Projection Layer Initialization

Unlike [Mou et al., 2016] who randomly initialized the projection layer, it is initialized with vectors pre-trained by an autoencoder [Hinton and Zemel, 1994] in our framework. This initialization stabilizes and improves optimization of neural networks during training, resulting more robust models.

4 Distillation Ensemble

Ensemble methods generally achieve higher accuracy than a standalone model; however, slow speed is expected due to the runs from multiple models in ensemble. This section presents a novel ensemble approach based on our distillation framework using logit matching that produces a light-weighted student model trained by multiple teachers (Figure 4).

The premise of this approach is that it is possible to have multiple teacher models train a student model by combining their logit values such that the student no longer needs the teachers during decoding because it already learned "enough" from them during training. As a result, our ensemble approach ensures higher efficiency for the student model than for any of the teacher models during decoding. The following sections describe two different ensemble methods applied to our framework.

4.1 Routing by Agreement Ensemble (RAE)

This method gives more weights to the majority, by adopting the dynamic routing algorithm presented by [Sabour et al., 2017]. It first collects the consensus of all teachers, then boosts weights of the teachers who strongly agree with that consensus whereas suppresses the influence of the teachers who do not agree as much. The procedure of calculating the representing logit is described in Algorithm 1.

The squash function in lines 4 and 9 is a non-linear activation function that ensures the norm of the output vector to be in [0, 1] [Sabour et al., 2017]. This vectorized activation is important for a routing algorithm because both magnitude and direction play an important role in the agreement calculation by the dot product, which enforces the consensus into the direction with strong confidence.

Algorithm 1: Get R. LOGIT for RAE and RDE

Input: Teachers' logits Z ∈ RT×C , and an algorithm selector b ∈ {RAE, RDE} 1 . Output: The representing logit, zrep ∈ RC . 2 k ← 1 if b is RAE else -1 3 for t ∈ {1, . . . , T} do 4 xt ← squash(zt ), wt ← 0 5 while n iterations do 6 c ← softmax(w) 7 zrep = ∑Tt=1 ct · zt 8 s ← squash(zrep ) 9 if not last iteration then 10 for t ∈ {1, . . . , T} do 11 wt ←wt +kxt · s 12 return zrep

4.2 Routing by Disagreement Ensemble (RDE)

This method focuses on the minority vote instead, because minority opinions may cast important information for the task. The algorithm is the same as RAE, except for the sign of the weight update (line 2 in Algorithm 1).

5 Experiments

5.1 Datasets

All models are evaluated on seven document classification datasets in Table 1. MR, SST-*, and CR are targeted at the task of sentiment analysis while Subj, TREC, and MPQA are targeted at the classifications of subjectivity, question types, and opinion polarity, respectively. About 10% of the training sets are split into development sets for SST-* and TREC, and about 10/20% of the provided resources are divided into development/evaluation sets for the other datasets, respectively.

5.2 Word Embeddings

For sentiment analysis, raw text from the Amazon Review dataset 2 is used to train word embeddings, resulting 2.67M word vectors. For the other tasks, combined text from Wikipedia and the New York Times Annotated corpus 3 are used, resulting 1.96Mword vectors. For each group, two sets of embeddings are trained with dimensions of 50 and 400 by Word2Vec [Mikolov et al., 2013]. While training, default hyper-parameters are used without an explicit hyper-parameter tuning.

2 http://snap.stanford.edu/data/web-Amazon.html

3 https://catalog.ldc.upenn.edu/LDC2008T19

5.3 Network Configuration

Two types of teacher models are developed using Convolutional Neural Networks (CNN) and Long Short-Term Memory Networks (LSTM); comparing different types of teacher models provides more generalized insights for our distillation framework. All teacher models use 400 dimensional word embeddings, and all student models are based on CNN. The CNN-based teacher and student models share the followings: filter sizes = [2, 3, 4, 5], # of filters = 32, dimension of the hidden layer right below the softmax layer = 50. Teacher models add a dropout of 0.8 to the hidden layer, wheres student models add dropouts of 0.1 to both the hidden layer and the the projection layer. On the other hand, the LSTM-based teacher models use two bidirectional LSTM layers. Only the last output vector is fed into the 50 dimensional hidden layer, which becomes the input to the softmax layer. A dropout of 0.2 is applied to all hidden layers, both in and out of the LSTM. For both CNN and LSTM ensembles, all 10 teachers share the same model structures with different initializations. Although this limited teacher diversity, our ensemble method produces remarkably good results (Section 5.7).

Each student model may consist of one or two projection layers. The one-layered projection adds one 50 dimensional hidden layer above the embedding layer that transfers core knowledge from the original embeddings to the distilled embeddings. The two-layered projection comprises two layers; the size of the lower layer is the same as the size of teacher's embedding layer, 400, and the size of the upper layer is 50, which is the dimension of the distilled embeddings. This twolayered projection is empirically found to be more robust than various combinations of network architectures including wider and deeper layers from our experiments. 4

5.4 Pre-trainedWeights

An autoencoder comprising a 50-dimensional encoder and a 400-dimensional decoder is used to pre-train weights for the two-layered projection in student models, where the encoder and the decoder have the same and inversed shapes as the upper and lower layers of the projection. Note that results by using pre-trained weights for the one-layered projection are not reported in Table 2 due to the limited space, but we consistently see robust improvement using pre-trained weights for the projection layers.

5.5 Embedding Distillation

Six models are evaluated on the seven datasets in Table 1 to show the effectiveness of our embedding distillation framework: logit matching (LM), noisy logit matching (NLM), softmax tau matching (STM) models with teacher-student based embedding distillation (TSED), and another three models with the autoencoder pre-trained weights (*+PT) and the two layered projection network (*+2L). Teacher models using 400-dim embeddings (*-400) are also presented along with the baseline model using 50-dim word embeddings (*-50) and the previous distillation model (ENC). The comparison to these two models highlights the strength of our distilled models, significantly outperforming them with the same dimensional word embeddings (50-dim).

4 Among convolutional, relational, and dense-networks with different configurations, the dense-network with the reported configuration produces the best results.

Table 2 shows the results achieved by all models. While the two existing models, *-50 and ENC, show marginal differences, our proposed models, *+TSED, outperform the previous embedding distillation SOTA (ENC). Our final models, *+PT+2L, outperform all the other models, reaching similar (or even higher) accuracy to the teacher models (*-400), which confirms that the proposed embedding distillation framework can successfully transfer the core knowledge from the original embeddings with respect to the target tasks, independent from the network structures of the teacher models.

The fact that the best model for each task comes from a different teacher-student strategy appears to be random. However, considering that it is the nature of neural network models whose accuracy deviates at every training, this is not surprising although it signifies the need of ensemble approaches to take advantage of multiple teachers and produce a more robust student model.

5.6 Lexical Analysis

Embedding distillation for a task such as sentiment analysis can be viewed as a vector transformation that adjusts similarities between word embeddings with respect to their sentiments. Thus, we hypothesize that distilled word embeddings from sentiment analysis should bring similar sentiment words together while disperse opposite ones in vector space.

To verify this, lexical analysis on both the original and the distilled word embeddings is conducted using the four sentiment datasets: SST-1, SST-2, MR, and CR. First, positive and negative words are collected from two publicly available lexicons, the MaxDiff Twitter Sentiment Lexicon [Kiritchenko et al., 2014] and the Bing Liu Opinion Lexicon [Hu and Liu, 2004]. Then, two groups of sentiment word sets are constructed, (Pt, Nt) and (Po, No), where P∗ and N∗ compose positive and negative words, and ∗t and ∗o are collected from the Twitter and the Opinion lexicons, respectively. Next, the intersection between each type of sentiment word sets is found, that are Pto = Pt ∩Po and Nto = Nt ∩No, where |Pto| = 72 and |Nto| = 89. Finally, the intersections between ∗to and the vocabulary set from each of the four sentiment datasets are found (e.g., Pℓ = Aℓ ∩Pto, where ℓ is one of the four datasets and Aℓ is the set of all words in ℓ):

• ℓ = SST1 ⇒ |Pℓ| = 66 and |Nℓ| = 83.
• ℓ = SST2 ⇒ |Pℓ| = 67 and |Nℓ| = 83.
• ℓ =MR⇒ |Pℓ| = 67 and |Nℓ| = 82.
• ℓ = CR⇒ |Pℓ| = 19 and |Nℓ| = 57.

For each ℓ, cosine similarities are measured for all possible word pairs (wi,wj) ∈ PNℓ × PNℓ where PNℓ = Pℓ ∪ Nℓ, using the original and distilled embeddings. Figure 5 illustrates the similarity distributions. It is clear that similar sentiment words generally give high similarity scores with the distilled embeddings (the plots drawn by red circles), whereas opposite sentiment words give low similarity scores (the plots drawn by red crosses). On the other hand, low similarity scores are shown for any case with the original embeddings (the plots drawn by blue circles and crosses). The normal distributions derived from the distilled embeddings (the red lines) are more symmetric and spread than the ones from the original embeddings (the blue lines), implying that distilled embeddings are more stable for the target task.

5.7 Distillation Ensemble

All ensemble models are based on our distillation framework using logit matching (LM) where the teacher models compose of 10 CNN-based or 10 LSTM-based models. Two ensemble methods are evaluated: Routing by Agreement (RAE) and Routing by Disagreement (RDE), and Figure 6 shows comparisons between these ensemble models against the teacher models.

The most notable finding is that RDE significantly outperforms the teacher, if the dataset is big. For example, RDE outperforms the teacher models on average, except for CR and MPQA, whose training set is relatively small (Table 1). The insight behind this trend might be that if there are many data samples, then the probability of exploring different knowledge from minor opinions could be increased, which would positively affect to the ensemble.

5.8 Model Reduction

The deployment models from either distillation or ensemble are notably smaller than the teacher models. Since word embeddings occupy a large portion of neural models in NLP, reducing the size of word embeddings through distillation decreases the size of the entire model roughly by the same ratio as its reduction; in our case, eight times (400/50). Furthermore, if the proposed distillation ensemble method is compared to other typical ensemble methods, this reduction ratio becomes even larger. Table 3 shows the number of neurons required for previous ensemble methods and the proposed one. When training, the proposed one comprises 10% more parameters due to the distillation process. However, when deploying, the reduction of neuron is about eighty times (4000/50). This is because the proposed framework doesn't require repetitive evaluation of teacher models when testing, while other ensemble methods require evaluation of all sub models (teachers).

It is worth mentioning that the deployment models produced by distillation ensemble not only outperform the teacher models in accuracy but also are significantly smaller such that they operate much faster and lighter than the teacher models upon deployment. This is very welcoming for those who want to embed these models into low-resource platforms such as mobile environments.

6 Conclusion

This paper proposes a new embedding distillation framework based on several teacher-student methods. Our experiments show that the proposed distillation models outperform the previous distillation model and give compatible accuracy to the teacher models, yet they are significantly smaller. Lexical analysis on sentiments reveals the comprehensiveness of the distilled embeddings. Moreover, a novel distillation ensemble approach is proposed, which shows huge advantage in both speed and accuracy over any teacher model. Our distillation ensemble approach consistently shows more robust results when the size of training data is sufficiently large.