P ROACTIVE S EQUENCE G ENERATOR VIA K NOWLEDGE A CQUISITION
Anonymous authors
Paper under double-blind review
A BSTRACT
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 K nowledge A cquisition (KA) where student models receive log q ( y t | y <
t , x ) as rewards when producing the next token y t 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.
I NTRODUCTION
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 ﬁtting.
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 D KL ( q ( y t | y <t , x ) (cid:107) p ( y t | y <t , x )) where q ( y t | y <t , x ) [teacher] or p ( y t | 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 ( y t | y <t , x ) have been retrieved among the top- k (cid:48) in p ( y t
| y <t , x ) .
However, sequence modeling actually cares about precision - the fraction of the top- k (cid:48)
tokens in p ( y t | y <t , x ) are included in the top- k ” of q ( y t | y <t , x ) , because imperfect predictions in the top- k (cid:48) cause exposure bias at inference time.
To optimize precision, we attempt to minimize KL-divergence in the reverse order, that is, D
KL ( p ( y t | y <t , x ) (cid:107) q ( y t | 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 log q ( y t
| y <t , x )) as rewards.
After receiving the feedback, “Actor” updates it’s belief accordingly, as shown in Fig.
Unlike task-speciﬁc 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 ﬁnite 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 ﬁrst 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
[Figure]
“Amazing view   of the seaside, of the boats in the sea of 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 ﬁrst to provide a depth-analysis of the effects of minimizing D KL ( p (cid:107) q ) on alleviating exposure bias.
In addition, we enable stable reinforcement training via replacing task-speciﬁc 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.
S EQUENCE GENERATION
Let ( x , y ) ∈
( X , Y ) be an example translation pair, where x is the source sentence and y is the target sentence.
y t is the token at position t and y <t is tokens before position t .
Sutskever et al.
propose sequence-to-sequence models which typically factorize the joint probability p ( y | x ) over a sequence of conditional probabilities with parameter θ :
[Equation]
where, p θ ( y t | y <t , x ) is the probability of the token y t 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
[Equation]
By substituting Eqn.(2) for Eqn.(1), we have
[Equation]
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 y t ∈ V , running on a world model M with transition function f .
Training.
We minimize the cross-entropy loss
[Equation]
where, y ref 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”.
D ISCREPANCY AMONG PROCEDURES
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 deﬁnitely 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.
K NOWLEDGE A CQUISITION
Unlike knowledge distillation, we attempt to minimize
[Equation]
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 y t given preceding tokens y <t and the source sentence x , i.e .,
[Equation]
7) “Actor”: updates the policy p θ ( y t | h t ) 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 D KL (cid:2) q (cid:107) p (cid:3) , our approach is actually a kind of reverse KD, which we call ”Knowledge
Acquisition” (KA).
3. L EARNING STRATEGY
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
[Figure]
Fig.
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 ) .
3.1. O BJECTIVE FUNCTIONS
The entropy term in Eqn.(6) goes away in KD because q ( x ) is a ﬁxed value and never be optimized during training.
It’s clear that the entropy term helps to avoid over-ﬁtting the model to human references and generalize well on unseen test data.
3.1. D ERIVATIVES
Based on Lagrangian relaxation (see Appendix A.1), the derivatives of D KL w.r.t.
p ( x ) in both orders are
[Equation]
Eqn.(8) tells us that
[Equation]
(cid:107) q ( x ) <(b) q ( x ) < p ( x ) ⇒ G q (cid:107) p ( x ) > 0 and G p (cid:107) q ( x ) >
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), (cid:12)(cid:12) G q (cid:107) p ( x ) (cid:12)(cid:12) > (cid:12)(cid:12) G p (cid:107) q
x ) (cid:12)(cid:12) while in (b), (cid:12)(cid:12) G q (cid:107) p ( x ) (cid:12)(cid:12) < (cid:12)(cid:12)
G p (cid:107) q ( x ) (cid:12)(cid:12)
In other words, minimizing D KL ( q ( x ) (cid:107) p ( x )) attempts to pull up the probability of x that underestimated by the student models while minimizing D KL ( p ( x ) (cid:107) q ( x ))
attempts to push down the the probability of x that was over-estimated by the student models.
If we deﬁne
Recall k (cid:48)
( p ( x ) , k ”) =
[Equation] [Equation]
Precision
[Equation] [Equation]
) where, p (cid:48) s Top- k (cid:48) is a set including k (cid:48) items having the highest rankings in p ( x ) and q (cid:48)
s Top- k ” is a set including k ” items having the highest rankings in q ( x ) .
Clearly, minimizing D KL ( q ( x ) (cid:107) p ( x )) is to maximize recall because its role is to increase the chance that x ’s in q (cid:48) s Top- k ” are truly picked.
By contrast, minimizing D KL ( p ( x ) (cid:107) 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 (cid:48) accumulate and make the generation brittle.
3.1. D OES THE DISTINCTION MATTER ?
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-ﬁt 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 D KL ( q (cid:107) 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.
3. T O BE P ROACTIVE WITH R ICH A DVICE
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 D KL term becomes
[Equation]
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
[Equation]
Fig.
3 shows that KD drives p 0 and p 1 to the real values very quickly, but almost ignore p 2 and psince they are not directly optimized.
In contrast, KA ﬁrst optimizes p 0 and p 2 which are the top-2.
After ∼ 100 epochs, when p 0 and p 1 become the top-2, the loss drops very fast and even lower than that of KD.
Moreover, KA drives p 2 and p 3 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
[Equation]
(11) i.e., minimizing D KL
[Equation]
p
[Equation]
y
[Equation]
y
[Equation]
| h t ) (cid:1) .
y t : t (cid:48) are tokens starting with position t and ending at position t (cid:48) .
By increasing t (cid:48) - 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 (cid:48) - t
= 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).
R ELATED W ORK
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 (cid:54) = inference.
To handle the mismatch between training and inference, previous works attempt to directly optimize the task-speciﬁc 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
[Figure]
[Table]
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
[Table]
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 signiﬁcant 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 D
KL ( q (cid:107) 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 ∗ D KL ( q (cid:107) p )+0 .
5 ∗ D KL ( p (cid:107) q ) ), but at sequence-level.
In this paper, we offer a thorough analysis of the effects of minimizing D KL ( p (cid:107) q ) on avoiding exposure bias, and we ﬁnd that minimizing D
KL ( p (cid:107) q ) performs the best among D KL ( q (cid:107) p ) and JSD.
[Figure] [Figure]
[Table]
E XPERIMENTS 5. D ATASET
We test our approach on WMT 2017 German-English with 4M sentence pairs, validate on newstest2016 and test on newstest2017.
All the sentences are ﬁrst 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.
5. T RAINING
Without a good starting point, the performance of minimizing D KL ( p (cid:107) 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 ﬁne-tune them in all the experiments by minimizing
[Equation] [Equation]
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 conﬁgurations 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 , (cid:15) = 1 e − 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.
5. R ESULTS
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 + 12 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
[Table]
,never,no,hardly,do,make, go 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.
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 ﬁrst and then drops after k = 8 .
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.
[Equation]
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 preﬁx 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 ﬁne-tuning.
This also indicates that KA motivates exploration more than KD.
C ONCLUSION
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 efﬁcient models
A A PPENDIX A. D ERIVATIVES
• D KL ( q ( x ) (cid:107) p ( x )) .
The objective is
Minimize (cid:88) q ( x
[Equation]
log( q ( x
[Equation]
( x
[Equation] [Equation]
We write the Lagrangian for Eqn.13 as
[Equation]
Method of Lagrangian multipliers involves setting the derivative of L q (cid:107) p w.r.t p ( x ) to 0,
[Equation]
Using the fact that (cid:80) p ( x ) = 1 , we can show that λ = 1 .
D KL ( p ( x ) (cid:107) q ( x ))
The objective is
[Equation]
We write the Lagrangian for Eqn.16 as
[Equation]
Method of Lagrangian multipliers involves setting the derivative of L p
[Equation]
Using the fact that (cid:80) p ( x ) = 1 , we can show that λ = − 1 .
A. P ROPERTY
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
ﬁrst consider the function
z
−
log z
[Equation]
It’s easy to prove because when z < 1 , the gradient 1 − 1 z < 0 and when z > 1 , the gradient 1 − 1 z > 0 .
Thus, the function reaches the minimum value 1 at z = 1 .
When q ( x ) > p ( x ) , we have
[Equation]
When q ( x ) < p ( x ) , we have
[Equation]