458 Chapter 13. 神经机器翻译模型训练 肖桐 朱靖波
13.5 知识蒸馏
理想的机器翻译系统应该是品质好、速度快、存储占用少。不过,为了追求更好
的翻译品质,往往需要更大的模型,但是相应的翻译速度会降低,模型的体积会变
大。在很多场景下,这样的模型无法直接使用。比如,Transformer-Big 等“大”模型
通常在专用服务器上运行,在手机等受限环境下仍很难应用。
但是,直接训练“小”模型的效果往往并不理想,其翻译品质与“大”模型相比
仍有比较明显的差距。既然直接训练小模型无法达到很好的效果,一种有趣的想法
是把“大”模型的知识传递给“小”模型。这类似于,教小孩子学习数学,不是请一
个权威数学家(即数据中的标准答案)进行教学,而是会请一个小学数学教师(即
“大”模型)来教小孩子。这就是知识蒸馏的基本思想。
13.5.1 什么是知识蒸馏
通常,知识蒸馏可以被看作是一种知识迁移的手段
[549]
。如果把“大”模型的知
识迁移到“小”模型,这种方法的直接结果就是模型压缩(Model Compression)。当
然,理论上也可以把“小”模型的知识迁移到“大”模型,比如,将迁移后得到的
“大”模型作为初始状态,之后继续训练该模型,以期望取得加速收敛的效果。不过,
在实践中更多是使用“大”模型到“小”模型的迁移,这也是本节讨论的重点。
知识蒸馏基于两个假设:
•“知识”在模型间是可迁移的。也就是说,一个模型中蕴含的规律可以被另一
个模型使用。最典型的例子就是预训练语言模型(见第九章)。使用单语数据学
习到的表示模型,在双语的翻译任务中仍然可以发挥很好的作用。也就是,把
在单语语言模型学习到的知识迁移到双语翻译过程中对句子的表示任务上。
• 模型所蕴含的“知识”比原始数据中的“知识”更容易被学习到。比如,机器
翻译中大量使用的回译(伪数据)方法,就把模型的输出作为数据让系统进行
学习。
这里所说的第二个假设对应了机器学习中的一大类问题
——
学习难度(Learning
Difficulty)。所谓难度是指:在给定一个模型的情况下,需要花费多少代价对目标任
务进行学习。如果目标任务很简单,同时模型与任务很匹配,那学习难度就会降低。
如果目标任务很复杂,同时模型与其匹配程度很低,那学习难度就会很大。在自然语
言处理任务中,这个问题的一种表现是:在质量很高的数据中学习的模型的翻译质
量可能仍然很差。即使训练数据是完美的,但是模型仍然无法做到完美的学习。这可
能是因为建模的不合理,导致模型无法描述目标任务中复杂的规律。在机器翻译中
这个问题体现的尤为明显。比如,在机器翻译系统输出的 n-best 结果中挑选最好的
译文(称为 Oracle)作为训练样本让系统重新学习,系统仍然达不到 Oracle 的水平。
知识蒸馏本身也体现了一种“自学习”的思想。即利用模型(自己)的预测来教
模型(自己)。这样既保证了知识可以向更轻量的模型迁移,同时也避免了模型从原
13.5 知识蒸馏 459
始数据中学习难度大的问题。虽然“大”模型的预测中也会有错误,但是这种预测是
更符合建模的假设的,因此“小”模型反倒更容易从不完美的信息中学习到更多的
知识
15
。类似于,刚开始学习围棋的人从职业九段身上可能什么也学不到,但是向一
个业余初段的选手学习可能更容易入门。另外,也有研究表明:在机器翻译中,相比
于“小”模型,“大”模型更容易进行优化,也更容易找到更好的模型收敛状态
[612]
。
因此在需要一个性能优越,存储较小的模型时,也会考虑将大模型压缩得到更轻量
模型
[613]
。
通常把“大”模型看作是传授知识的“教师”,被称作教师模型(Teacher Model)
;把“小”模型看作是接收知识的“学生”,被称作学生模型(Student Model)。比
如,可以把 Transformer-Big 看作是教师模型,把 Transformer-Base 看作是学生模型。
13.5.2 知识蒸馏的基本方法
知识蒸馏的基本思路是让学生模型尽可能去拟合教师模型
[549]
,通常有两种实现
方式
[550]
:
• 单词级的知识蒸馏(Word-level Knowledge Distillation)。该方法的目标是使得学
生模型的预测(分布)尽可能逼近教师模型的预测(分布)。令 x = {x
1
,...,x
m
}
和 y = {y
1
,...,y
n
} 分别表示输入和输出(数据中的答案)序列,V 表示目标语
言词表,则单词级的知识蒸馏的损失函数被定义为:
L
word
= −
n
X
j=1
X
y
j
∈V
P
t
(y
j
|x)logP
s
(y
j
|x) (13.22)
这里,P
s
(y
j
|x) 和 P
t
(y
j
|x) 分别表示学生模型和教师模型在 j 位置输出的概率。
公式(13.22)实际上在最小化教师模型和学生模型输出分布之间的交叉熵。
• 序列级的知识蒸馏(Sequence-level Knowledge Distillation)。除了单词一级输出
的拟合,序列级的知识蒸馏希望在序列整体上进行拟合。其损失函数被定义为:
L
seq
= −
X
y
P
t
(y|x)logP
s
(y|x) (13.23)
公式(13.23)要求遍历所有可能的译文序列,并进行求和。当词表大小为 V ,序
列长度为 n 时,则序列的数量有 V
n
个。因此,会考虑用教师模型的真实输出
序列 ˆy 来代替整个空间,即假设 P
t
(ˆy|x) = 1。于是,目标函数变为:
L
seq
= −logP
s
(ˆy|x) (13.24)
15
很多时候,“大”模型和“小”模型都是基于同一种架构,因此二者对问题的假设和模型结构都是
相似的。
460 Chapter 13. 神经机器翻译模型训练 肖桐 朱靖波
这样的损失函数最直接的好处是,知识蒸馏的流程会非常简单。因为只需要利
用教师模型将训练数据(源语言)翻译一遍,之后把它的输出替换为训练数据
的目标语言部分。之后,利用新得到的双语数据训练学生模型即可。图13.12对
比了词级和序列级知识蒸馏方法。
0.7
0.2
0.4
0.3
0.6
0.8
教师
模型
输出
分布
⟨eos⟩
I
am
fine
good
词
表
Student
⟨eos⟩
I
am
fine
真实数据:
(a) Word-level
1
I
1
am
1
good
1
⟨eos⟩
教师
模型
输出
译文
⟨eos⟩
I
am
fine
good
词
表
Student
⟨eos⟩
I
am
good
教师译文:
(b) Sequence-level
图 13.12 词级和序列级知识蒸馏的差异
本质上,单词级的知识蒸馏与语言建模等问题的建模方式是一致的。在传统方
法中,训练数据中的答案会被看作是一个 One-hot 分布,之后让模型去尽可能拟合这
种分布。而这里,答案不再是一个
One-hot
分布,而是由教师模型生成的真实分布,
但是损失函数的形式是一模一样的。在具体实现时,一个容易出现的问题是在词级
别的知识蒸馏中,教师模型的 Softmax 可能会生成非常尖锐的分布。这时需要考虑
对分布进行平滑,提高模型的泛化能力,比如,可以在 Softmax 函数中加入一个参数
α,如 Softmax(s
i
) =
exp(s
i
/α)
P
i
′
exp(s
i
′
/α)
。这样可以通过 α 控制分布的平滑程度。
除了在模型最后输出的分布上进行知识蒸馏,同样可以使用教师模型对学生模
型的中间层输出和注意力分布进行约束。这种方法在第十四章中会有具体应用。
13.5.3 机器翻译中的知识蒸馏
在神经机器翻译中,通常使用公式(13.24)的方法进行知识蒸馏,即通过教师模
型构造伪数据,之后让学生模型从伪数据中学习。这样做的好处在于,系统研发人
员不需要对系统进行任何修改,整个过程只需要调用教师模型和学生模型标准的训
练和推断模块即可。
那么如何构造教师模型和学生模型呢?以 Transformer 为例,通常有两种思路:
• 固定教师模型,通过减少模型容量的方式设计学生模型。比如,可以使用容量较
大的模型作为教师模型(如:Transformer-Big 或 Transformer-Deep),然后通过将
神经网络变“窄”、变“浅”的方式得到学生模型。例如,可以用 Transformer-Big
做教师模型,然后把 Transformer-Big 的解码器变为一层网络,作为学生模型。
13.5 知识蒸馏 461
• 固定学生模型,通过模型集成的方式设计教师模型。可以组合多个模型生成更
高质量的译文。比如,融合多个 Transformer-Big 模型(由不同参数初始化方式
训练的),之后学习一个 Transformer-Base 模型。
此外还可以采用迭代式知识蒸馏的方式。首先,通过模型集成得到较强的教师
模型,再将知识迁移到不同的学生模型上,随后继续使用这些学生模型集成新的教
师模型。不断的重复上述过程可以逐步提升集成模型的性能,如图13.13所示。值得
注意的是,随着迭代次数的增加,集成所带来的收益也会随着子模型之间差异性的
减小而减少。
学
生
模
型
1
学
生
模
型
2
学
生
模
型
3
学
生
模
型
4
学
生
模
型
5
教
师
模
型
1
教
师
模
型
2
集成
学
生
模
型
1
学
生
模
型
2
学
生
模
型
3
学
生
模
型
4
学
生
模
型
5
教
师
模
型
1
教
师
模
型
2
集成
学
生
模
型
1
学
生
模
型
2
学
生
模
型
3
学
生
模
型
4
学
生
模
型
5
教
师
模
型
1
教
师
模
型
2
集成
轮
数
1
轮
数
2
知识蒸馏
知识蒸馏
图 13.13 迭代式知识蒸馏
如果倾向于使用更少的存储,更快的推理速度,则可以使用更小的学生模型。值
得注意的是,对于 Transformer 模型来说,减少解码端的层数会给推理速度带来巨大
的提升。特别是对于基于深层编码器的 Transformer-Deep,适当减少解码端层数往往
不会带来翻译品质的下降。可以根据不同任务的需求,选择适当大小的学生模型,来
平衡存储空间、推断速度和模型品质之间的关系。