Year: 2019

Abstract:
Sequence-to-sequence models such as transformers, which are now being used in a wide variety of NLP tasks, typically need to have very high capacity in order to perform well. Unfortunately, in production, memory size and inference speed are all strictly constrained. To address this problem, Knowledge Distillation (KD), a technique to train small models to mimic larger pre-trained models, has drawn lots of attention. The KD approach basically attempts to maximize recall, i.e., ranking Top-k” tokens in teacher models as higher as possible, however, whereas precision is more important for sequence generation because of exposure bias. Motivated by this, we develop Knowledge Acquisition (KA) where student models receive logq(yt|y<t,x) as rewards when producing the next token yt given previous tokens y<t and the source sentence x. We demonstrate the effectiveness of our approach on WMT’17 De-En and IWSLT’15 Th-En translation tasks, with experimental results showing that our approach gains +0.7-1.1 BLEU score compared to token-level knowledge distillation.

heading: 1 INTRODUCTION
text: We have recently witnessed rapid progress sequence models for natural language processing (NLP). In particular, transformers (Vaswani et al., 2017; Dai et al., 2019) and BERT (Devlin et al., 2018; Liu et al., 2019) have been established as state-of-the-art approaches for sequence generation and language modeling respectively. However, the resulting sequence models are often exceedingly large. For example, a big transformer with 6 layers, 1024 dim and 16 heads has∼ 100M parameters, which can make production deployment impractical, especially on mobile devices. Unfortunately, learning small sequence models is very challenging because of the nature of deep learning - curve fitting. Training deep neural networks is data hungry and needs bigger “brains”. To deal with the challenges, Kim & Rush (2016) consider applying Knowledge Distillation approaches (Hinton et al., 2015) to match the distribution of the teacher model. At token-level, KD approaches minimize DKL(q(yt|y<t,x)‖p(yt|y<t,x)) where q(yt|y<t,x) [teacher] or p(yt|y<t,x) [student] is a probability distribution over tokens in the vocabulary conditioned on previous tokens y<t and the source sentence x. In other words, the KD approaches maximize recall - the fraction of the top-k” tokens in q(yt|y<t,x) have been retrieved among the top-k′ in p(yt|y<t,x). However, sequence modeling actually cares about precision - the fraction of the top-k′ tokens in p(yt|y<t,x) are included in the top-k” of q(yt|y<t,x), because imperfect predictions in the top-k′ cause exposure bias at inference time. To optimize precision, we attempt to minimize KL-divergence in the reverse order, that is, DKL(p(yt|y<t,x)‖q(yt|y<t,x)). Our approach can be thought of as a kind of Actor-Critic approach where the “Actor” is a student model asking for advice when learning to generate every single token and the “Critic” is a teacher model that gives logq(yt|y<t,x)) as rewards. After receiving the feedback, “Actor” updates it’s belief accordingly, as shown in Fig. 1. Unlike task-specific metrics such as BLEU score, teacher models summarize training data points and structure knowledge in an abstract, hierarchical manner. Therefore, they are able to offer more useful advice than matching n-grams with a finite number of human references. For example, assume there is a training example {SRC: “Die Landschaft am Meer ist wunderschn”, TRG: “Amazing view along the sea” }. During learning, given the first token “Amazing”, student models might ask “How about saying amazing view?”, teacher models most likely give high rewards since it’s expected to go with
“Amazing view of the seaside, of the boats in the seaof the beach, from the boats· · · , · · · along the sea, of the birds flying over the water  ”, where most of them express similar meanings with the human references even though they have mismatched n-grams. Contributions. To the best of our knowledge, we are the first to provide a depth-analysis of the effects of minimizing DKL(p‖q) on alleviating exposure bias. In addition, we enable stable reinforcement training via replacing task-specific metrics with high-capacity models. Finally, we improved +1.1 BLEU score on WMT’17 de-en task and +0.7 BLEU score on IWSLT’15 th-en task over token-level KD with up to relative 200% improvement.

heading: 2 SEQUENCE GENERATION
text: Let (x,y) ∈ (X ,Y) be an example translation pair, where x is the source sentence and y is the target sentence. yt is the token at position t and y<t is tokens before position t. Sutskever et al. (2014) propose sequence-to-sequence models which typically factorize the joint probability p(y | x) over a sequence of conditional probabilities with parameter θ:
p ( y|x ) = T∏ t=0 pθ ( yt | y<t,x ) (1)
where, pθ(yt | y<t,x) is the probability of the token yt given previous tokens y<t and the source sentence x. Basically, the preceding tokens y<t are encoded into the hidden states via a state transition function
ht = f(ht−1,yt−1;x) (2) By substituting Eqn.(2) for Eqn.(1), we have
p ( y|x ) = T∏ t=0 pθ(yt | ht)︸ ︷︷ ︸ Policy π
(3)
This tells us that the sequence models, e.g., RNNs, acts like a stochastic policy which picks a discrete action, i.e., producing a token yt ∈ V , running on a world modelM with transition function f . Training. We minimize the cross-entropy loss
LCE = − ∑ t logpθ ( yreft | y ref <t ,x ) (4)
where, yref denotes human references. At each position, the models are conditioned on the groundtruth tokens annotated by humans no matter what tokens are predicted by themselves. Inference. In sequence generation tasks, exact inference is intractable due to exponentially large search space. Instead, we apply an approximation inference algorithm - Beam Search (BS). BS is a greedy heuristic search that maintains the top-B most likely partial sequences through the search tree, where B is referred to as the beam size. At each position, BS expands these B partial sequences to all possible beam extensions and then selects the B highest scoring among the expansions. Unlike training, ground-truth tokens are not available. The models have to use their own predictions in decoding. Evaluation. To evaluate the quality of generated sequences, we typically use metrics such as BLEU score to measure their n-grams overlap with human references. However, the standard metrics are problematic and none of them correlate strongly with human evaluation at the level of individual sentences. For example, given a human reference “Amazing view along the sea”, the sequence ”The scenery of the seaside is beautiful” gets low BLEU score because there is no matching n-grams of order 2, 3, or 4. In addition, human references are a few sentences lacking in diversity. For example, when asking more people, they might say “A nice beach” or “ what amazing view of the seashore”.

heading: 2.1 DISCREPANCY AMONG PROCEDURES
text: Exposure bias. During training, the models only explore the training data distribution, but never get exposed to their own predictions. Searching in under-explored space causes errors. More importantly, such errors accumulate over time because of greedy search. For the example in Sec. 1, assume there are no sentences in training data starting with “Amazing”.
At position 1: pick a token from p(y|“Amazing”,x)⇒ “cup” At position 2: pick a token from p(y|“Amazing cup”,x)⇒ “on” At position 3: pick a token from p(y|“Amazing cup on”,x)⇒ “table”
We can see that the poor token “cup” caused by the noisy distribution p(y|“Amazing”,x) makes the situation even worse. The distributions become more and more noisy and the generation goes far away. Sub-optimal models. The training objective is different with the metrics used in evaluation. To address this issue, some works attempt to directly optimize the metrics. It definitely helps improving the score, but we suspect that this might hurt the models because poor metrics discourage learning the semantic meanings. For example, the low score of “The scenery of the seaside is beautiful” makes the meanings of “Amazing” and “beautiful” or “seaside” and “sea” far from the other.

heading: 3 KNOWLEDGE ACQUISITION
text: Unlike knowledge distillation, we attempt to minimize DKL [ pθ(yt | ht)‖qφ(yt | ht) ] (5)
= − [ Eyt∼pθ(yt|bt) [ logqφ(yt | ht) ] +H [ pθ(yt | ht) ]] (6)
Eqn.6 can be explained as an Actor-Critic algorithm (Konda & Tsitsiklis, 2000) that maximizes the expected rewards with an entropy term, where
“Critic”: (aka teacher models) estimate the rewards of producing token yt given preceding tokens y<t and the source sentence x, i.e., Q(a = yt, s = ht) = logqφ(yt | ht) (7) “Actor”: updates the policy pθ(yt | ht) guided by the ”Critic”
As opposed to behavior cloning, student models ask teacher models to evaluate their behaviors and then update accordingly when producing the next token. By comparing with KD that attempts to minimize DKL [ q‖p ] , our approach is actually a kind of reverse KD, which we call ”Knowledge Acquisition” (KA).

heading: 3.1 LEARNING STRATEGY
text: Without loss of generality, we simplify the problem to a 1-D problem where the student model p(x) is a single-modal Gaussian and the teacher model q(x) is a mixture of two Gaussians, as shown in
Fig. 2. Let’s take a depth-analysis of their objective functions and their derivatives w.r.t. p(x) to see what changes after swapping p(x) and q(x).

heading: 3.1.1 OBJECTIVE FUNCTIONS
text: The entropy term in Eqn.(6) goes away in KD because q(x) is a fixed value and never be optimized during training. It’s clear that the entropy term helps to avoid over-fitting the model to human references and generalize well on unseen test data.

heading: 3.1.2 DERIVATIVES
text: Based on Lagrangian relaxation (see Appendix A.1), the derivatives of DKL w.r.t. p(x) in both orders are
∂DKL(q(x)||p(x)) ∂p(x) = 1− q(x) p(x)︸ ︷︷ ︸
Gq‖p(x)
, ∂DKL(p(x)||q(x))
∂p(x) = log
p(x) q(x)︸ ︷︷ ︸ Gp‖q(x)
(8)
Eqn.(8) tells us that
(a) q(x) > p(x)⇒ Gq‖p(x) < 0 and Gp‖q(x) < 0 (b) q(x) < p(x)⇒ Gq‖p(x) > 0 and Gp‖q(x) > 0
That is, the probability of x is pulled up when q(x) dominates p(x), and pushed down otherwise. However, they take efforts at different levels (see Appendix A.2).
In (a), ∣∣Gq‖p(x)∣∣ > ∣∣Gp‖q(x)∣∣ while in (b), ∣∣Gq‖p(x)∣∣ < ∣∣Gp‖q(x)∣∣
In other words, minimizing DKL(q(x)‖p(x)) attempts to pull up the probability of x that underestimated by the student models while minimizing DKL(p(x)‖q(x)) attempts to push down the the probability of x that was over-estimated by the student models. If we define
Recallk′(p(x), k”) =
|p′s Top-k′ ∩ q′s Top-k”| k”
(9)
Precisionk”(p(x), k′) =
|p′s Top-k′ ∩ q′s Top-k”| k′ (10)
where, p′s Top-k′ is a set including k′ items having the highest rankings in p(x) and q′s Top-k” is a set including k” items having the highest rankings in q(x). Clearly, minimizing DKL(q(x)‖p(x)) is to maximize recall because its role is to increase the chance that x’s in q′s Top-k” are truly picked. By contrast, minimizing DKL(p(x)‖q(x)) is to maximize precision since its role is to avoid x’s that are over-estimated being picked by mistake. Overall, sequence modelling cares about precision because at test time, the generation is sensitive to the model’s predictions. The errors caused by inaccurate predictions in the top-k′ accumulate and make the generation brittle.

heading: 3.1.3 DOES THE DISTINCTION MATTER?
text: One might reasonably ask whether the distinction matters, given that both learning strategies aim to eventually reach the same equilibrium point, i.e., p(x) = q(x)? In practice, the answer is yes due to the following reasons: Low-capacity p(x). In practice, student models are of limited capacity due to memory and inference speed restrictions, especially in production. For example, running a big transformer on mobile device with real-time response is impractical. In this case, the equilibrium state p(x) = q(x) is unreachable. Thus, different learning strategies lead to different local optima, and result in different performance. Imperfect q(x). The teacher models are generally better approximations of the real underlying distributions from which the training data is, but they are still not 100% accurate. In particular, the exponentially-large search space makes sequence models easily over-fit to more frequently occurring data points, while being uncertain about the data points that are less likely observed. Kim & Rush (2016) also point out a mixture of DKL(q‖p) and cross-entropy loss reaches the best performance. In summary, with the limited capacity and noisy q(x), picking what to learn smartly turns out to be important.

heading: 3.2 TO BE PROACTIVE WITH RICH ADVICE
text: To further study the capability in exploration, we explicitly restrict the knowledge to be partially observed by allowing only the top-k to be available. The DKL term becomes∑
x∈Top-k p(x)
p ( x) ( logp ( x ) − logq ( x ))
where, Top-k p(x) is a set of k items with the highest ranking. Only p(x) and q(x) according to the x in the top-k are calculated while the rest is discarded. Assume x is a discrete variable and takes value 0, 1, 2, and 3. We build a quite simple model with a single soft-max layer to produce a distribution over x, and set top-k=2. The real distribution is
q0 = 0.4, q1 = 0.3, q2 = 0.2, q3 = 0.1
Fig. 3 shows that KD drives p0 and p1 to the real values very quickly, but almost ignore p2 and p3 since they are not directly optimized. In contrast, KA first optimizes p0 and p2 which are the top-2. After ∼ 100 epochs, when p0 and p1 become the top-2, the loss drops very fast and even lower than that of KD. Moreover, KA drives p2 and p3 much closer to the real values since the model has explored 3 states, i.e., x = 0, 1, 2. In sequence generation tasks, if student models receive rewards at each position, the entire distribution is fully visible. KA implicitly explores in the above way. In theory, we can generalize our approach to multi-step
Eyt:t′∼pθ(yt:t′ |ht) [ logqφ(yt:t′ | ht) ] +H [ pθ(yt:t′ | ht) ] (11)
i.e., minimizing DKL ( pθ(yt:t′ |ht)‖qφ(yt:t′ |ht) ) . yt:t′ are tokens starting with position t and ending at position t′. By increasing t′-t, the student models can capture long-term dependency and are much more capable of handling exposure bias. However, the training is expensive even with t′-t=2 since the search space explodes exponentially. In this case, picking top-k candidates is a solution. However, in practice we only model the reward for a single step (but use the signal from all steps in each batch update).

heading: 4 RELATED WORK
text: Our work closely relates to two lines of work: knowledge transfer and dealing with the mismatch between training and inference procedures. Both of them have already been explored in the literature. Learning 6= inference. To handle the mismatch between training and inference, previous works attempt to directly optimize the task-specific metric at test time. Ranzato et al. (2015) propose sequence-level training algorithm in reinforcement learning. The models receive rewards until the completion of the entire sequence. Considering that the search space in sequence generation is exponentially large, i.e., O(|V|T ), where V is a set of tokens in the vocabulary (∼ 10K or more) and
T is the length of the sequence (∼ 20 or more), the rewards are extremely sparse. This makes the training unstable. To alleviate the sparse rewards problem, Liu et al. (2017) use Monte Carlo rollouts to estimate immediate rewards at each position. Unfortunately, the estimation is very noisy and of high variance due to the exponentially large search space. Moreover, the training is computationally expensive. In this paper, we attempt to use Actor-Critic algorithms, where the “Critic” is a teacher model used to estimate action-value function, i.e., the rewards of producing the current token given previous tokens, and the “Actor” is a student model that learns to produce a token at each position. To deal with exposure bias, Bengio et al. (2015) propose to gradually replace tokens from human references to generated tokens during training. The training starts with tokens from human references and ends up with using generated tokens. However, the rewards which are the matching n-grams with a few human references limits the capability of the models in exploration. In this paper, we equip the models with knowledgeable teachers which offer smart advice based on semantic meanings. Transfer learning. In the context of sequence generation, pre-training with monolingual data (Radford et al., a;b) has had significant success in language understanding. Another line of work is knowledge distillation (KD) (Hinton et al., 2015), in particular when models are of different architectures and sizes. Kim & Rush (2016) propose two levels of knowledge distillation. At sequencelevel, student models are trained on the augmented dataset with outputs of running beam search with teacher models. At token-level, they get the conditional probability of each token given preceding tokens closer to that of teacher models. In other words, this is equivalent to minimizing DKL(q‖p), where q is the distribution of teacher models and p is the distribution of student models. Yu et al. (2017) use adversarial training to encourage the models producing human-like sequences by learning a sequence-level discriminator to distinguish generated sequence and human references. In fact, this is equivalent to minimizing Jensen-Shannon Divergence (JSD, i.e., 0.5∗DKL(q‖p)+0.5∗DKL(p‖q)), but at sequence-level. In this paper, we offer a thorough analysis of the effects of minimizing DKL(p‖q) on avoiding exposure bias, and we find that minimizing DKL(p‖q) performs the best among DKL(q‖p) and JSD.

heading: 5 EXPERIMENTS

heading: 5.1 DATASET
text: We test our approach on WMT 2017 German-English with 4M sentence pairs, validate on newstest2016 and test on newstest2017. All the sentences are first tokenized with moses tokenizer and then segmented into 40K joint source and target byte pair encoding (Sennrich et al., 2015). Another Thai-English dataset comes from IWSLT 2015. There are 90k sentence pairs in training and we take 2010/2011 as the dev set and 2012/2013 as the test set. Byte pair encoding is of size 10K.

heading: 5.2 TRAINING
text: Without a good starting point, the performance of minimizing DKL(p‖q) degrades because the models are more than likely stuck on the current best tokens and unlikely to explore. We therefore pre-train student models and fine-tune them in all the experiments by minimizing
LALL(y;x, θ, φ) = (1− λ)LNLL(y;x, θ) + λ ∑ t DKL ( pθ(yt | ht)‖qφ(yt | ht) ) ︸ ︷︷ ︸
DKL(qφ‖pθ)for KD
(12)
where, λ is trade-off parameters. Basically, “Critic” and “Actor” are optimized simultaneously. However, in this paper, we simply freeze teacher models and leave joint training to future work. Model. Our teacher models and students models all use the transformer architecture, which has achieved state-of-the-art accuracy and is widely used in recent NMT research. Model configurations are listed in Table.1. We train all transformer models using the implementation in Fairseq (Ott et al., 2019). We use Adam optimizer (Kingma & Ba, 2014) with β0 = 0.9, β1 = 0.98, = 1e−8. Learning rate is 0.0001 and dropout rate is 0.3. At inference time, we use beam search with a beam size of 6 and length penalty 1.0.

heading: 5.3 RESULTS
text: Our results on NMT tasks are reported in Table.2. We tune hyper-parameters λ and top-k on validation set where λ = 0.5 for KD which is consistent with that in Kim & Rush (2016). We observed that KA consistently outperforms KD on both De-En (high-resource) and Th-En (low-resource) tasks by 0.7 - 1.1 BLEU score. In addition, we also test JSD (i.e., 12 KA + 1 2 KD), which is equivalent to GAN. The accuracy lies between KA and KD. The results say that KA does help avoiding exploration bias and further close the gap between training and inference. Exploration. Similar to Sec.3.2, we evaluate the performance by allowing only top-k tokens to be available. We vary k from 2 to 16 and conduct the experiments on validation set. In Fig. 4a, we
see that the distribution is noisy because the accuracy of either KD or KA eventually goes down and when using the full information, i.e., k = |V|, KD achieves 35.41 BLEU while KA achieves 35.79 BLEU, which are far away from the best. Moreover, KA is able to learn more from noise because KD directly goes down as k increases, while the accuracy of KA goes up first and then drops after k = 8.
Figure 5: Entropy over epochs.
To further analyze the capability in exploration, we attempt to count the tokens in the top-k which have never been included in previous epochs.∣∣∣{yt ∈ Top-k piθ(yt|ht)}
&{yt /∈ Top-k pjθ(yt|ht),∀j < i} ∣∣∣
where the superscript i denotes epoch. We randomly sample 2000 sentence pairs from WMT’17 De-En training data. In Fig. 4b, we see that there are much more novel tokens in KA than that in KD. Table.3 demonstrates two examples. a) Given the prefix tokens “I will” and the source sentence, KA probably imagines
try to attend? give attention? see to?
which provide different ways to express the meaning of attend, while KD just try to explore concern. b) when predicting the token no, KA proposes a set of negative tokens
fail,resist,refrain
, and even go which may relate to the next token move. Fig. 5 shows that the entropy of KA goes up while the entropy of KD dropping during fine-tuning. This also indicates that KA motivates exploration more than KD.

heading: 6 CONCLUSION
text: We proposed a learning principle where student models actively ask for advice from teacher models. The theoretical analysis proves that it helps to alleviate exposure bias in sequence generation and improve the capability of generalization on unseen data. In general, our approach is extendable to phase-level or sequence-level, but we simply focus on token-level due to practical issues. However, our token-level KA approach is complementary to sequence-level KD by using the outputs of teacher models as extra data. Experimental results on benchmark translation data sets show that our approach improves over previous methods of leveraging large models (or ensembles) to improve the quality of smaller, more efficient models.

heading: A APPENDIX
text: A.1 DERIVATIVES
• DKL(q(x)‖p(x)). The objective is Minimize ∑ q(x) log( q(x)
p(x) ), s.t.
∑ p(x) = 1 (13)
We write the Lagrangian for Eqn.13 as Lq‖p(x, λ) = ∑ q(x) log( q(x)
p(x) ) + λ
(∑ p(x)− 1 ) (14)
Method of Lagrangian multipliers involves setting the derivative of L q‖p w.r.t p(x) to 0,
Gq‖p = ∂Lq‖p ∂p(x) = λ− q(x) p(x) = 0 (15)
Using the fact that ∑
p(x) = 1, we can show that λ = 1. • DKL(p(x)‖q(x)). The objective is
Minimize ∑ p(x) log( p(x)
q(x) ), s.t.
∑ p(x) = 1 (16)
We write the Lagrangian for Eqn.16 as Lp‖q(x, λ) = ∑ p(x) log( p(x)
q(x) ) + λ
(∑ p(x)− 1 ) (17)
Method of Lagrangian multipliers involves setting the derivative of L p‖q w.r.t p(x) to 0,
Gp‖q = ∂Lp‖q
∂p(x) = 1 + λ+ log
p(x) q(x) = 0 (18)
Using the fact that ∑ p(x) = 1, we can show that λ = −1.
A.2 PROPERTY
We’ll prove the key properties in Sec.3.1.2.
(a) When q(x) > p(x), we have 1 − q(x)p(x) < 0 and log p(x) q(x) < 0. And, when q(x) < p(x),
we have 1− q(x)p(x) > 0 and log p(x) q(x) > 0.
(b) Let’s first consider the function z − log z z − log z { > 1 if z 6= 1 = 1 if z = 1
(19)
It’s easy to prove because when z < 1, the gradient 1− 1z < 0 and when z > 1, the gradient 1− 1z > 0. Thus, the function reaches the minimum value 1 at z = 1.
When q(x) > p(x), we have∣∣Gq‖p(x)∣∣− ∣∣Gp‖q(x)∣∣ = ∣∣1− q(x) p(x) ∣∣− ∣∣ log p(x) q(x) ∣∣ (20) = q(x)
p(x) − log q(x) p(x) − 1 > 0 (21)
When q(x) < p(x), we have∣∣Gq‖p(x)∣∣− ∣∣Gp‖q(x)∣∣ = ∣∣1− q(x) p(x) ∣∣− ∣∣ log p(x) q(x) ∣∣ (22) = 1− (q(x)
p(x) − log q(x) p(x) ) < 0 (23)

Review summaries:
In this paper, the authors proposed a new method called knowledge acquisition (KA) for distilling the learned knowledge from the teacher model to the student model. Unlike the conventional KL-divergence based knowledge distillation method, the authors take the reverse version which learns the student to increase the precision. The paper gave a thorgough analysis on the proposed KA strategy and compared with other strategies like KL and JS divergence. On the sequence generation task (translation), the authors showed that the proposed KA strategy achieved better performance compared with KD based methods when distilling the knowledge from a teacher model to a student model.