基于递进式半知识蒸馏的神经机器翻译
2021-03-17周孝青段湘煜俞鸿飞
周孝青,段湘煜,俞鸿飞,张 民
(苏州大学 计算机科学与技术学院,江苏 苏州 215006)
0 引言
机器翻译包括统计机器翻译和神经机器翻译。近年来,神经机器翻译(neural machine translation,NMT)系统[1-5]取得了巨大成功,成为机器翻译的主流方法。神经机器翻译模型的性能在很大程度上受限于神经网络参数规模。神经机器翻译模型参数过多也会存在一些问题,例如,模型的存储空间要求变高、优化时间变长等。这种情况限制神经机器翻译系统在手机等小型设备上的应用。
模型压缩是解决此类问题的一种有效方案。现有的压缩方法有剪枝、量化、低精度和知识蒸馏等。剪枝是通过移除对目标函数影响小的权重[6],或者对隐藏层单元进行阈值处理[7]。知识蒸馏是通过引入教师网络(teacher network)的相关知识或信息优化学生网络(student network)的训练[8],其中教师网络模型复杂但性能优越,学生网络精简且复杂度低。目前神经网络压缩工作包括模型参数压缩和模型存储压缩两类,前者的目标是减少模型参数,后者是减少存储。模型参数压缩有剪枝、量化、降低精度和知识蒸馏等方法。在深度神经网络压缩方法中,剪枝是基于一定的准则对权重或神经元进行剪枝;剪枝准则主要包括近似海森权重剪枝[6]和基于值的大小剪枝[9];以及移除神经元的方法[10-11]。深度神经网络压缩的其它方法包括稀疏正则化[12]、权重矩阵的低秩分解[13-16]、权重共享[17-18]和权重的二值化[19]。在机器翻译领域中,See等人提出基于权重的幅值进行修剪或按权重的类别进行剪枝[12];Kim等人提出知识蒸馏[20];林野等人对比了剪枝、量化、低精度三种模型压缩方法在Transformer和递归神经网络(recurrent neural network, RNN)两种模型上的效果[21]。
本文提出半知识蒸馏方法来压缩模型。半知识蒸馏方法将已训练的教师模型的部分知识,通过权重赋值直接传递给学生网络,学生网络以此为起点开始训练,对于压缩权重采用高斯分布初始化。为了使得整个模型完全被压缩,可以递进式运用半知识蒸馏方法来压缩另外一部分的权重,将当前的学生网络充当新的教师网络,将新的知识传递给更小参数规模的学生网络。本方法可以应用到常见的端到端模型,本文以Transformer[5]模型结构为例使用此方法。首先,通过划分Transformer参数为两个部分,压缩和未压缩部分,然后压缩和未压缩的部分通过中间映射的权重作为压缩兼容性矩阵,以保证压缩的权重可与其他未压缩权重兼容,从而不需要添加新的参数,实现压缩整个模型参数的目标。
本文进行了充分的实验,实验结果显示,递进式知识蒸馏得到的全压缩模型和半压缩模型均比相应的基准模型有显著的性能改善。半知识蒸馏可以将Transformer模型压缩成更小的学生模型,其性能甚至超过教师模型的性能。通过现有工作(词级别和句子级别的知识蒸馏[21])和本文提出的方法进行对比,在中英任务上,本文与词级别和句子级别的知识蒸馏性能可比,并且在日英任务上,本文方法获得了比词级别和句子级别知识蒸馏更好的性能。在日英翻译和中英翻译实验中,最佳性能比相应的基准模型分别提高了2.16个BLEU值和0.72个BLEU值。
1 背景
Transformer是机器翻译中最成功的NMT系统之一,同时也是一个典型的NMT系统,具有多层结构且有大规模参数的特点,在许多基于端到端序列建模任务中表现出色。因此,选择它作为基准模型进行模型压缩具有代表性和一般性。与其他基于递归神经网络(recurrent neural network,RNN)[1-3]和卷积神经网络(convolutional neural network,CNN)[4]的翻译系统相比,Transformer仅基于注意力机制进行序列建模。其结构如图1所示。
图1 Transformer 模型架构
编码器由L个相同的层堆叠组成,每个层包含两个子层: 多头自注意层和前馈神经网络层。每个子层通过残差连接(residual connection)[19]进行连接,然后是层规范化[20](layer normalization)。可以由layernorm(h +sublayer(h))表示,其中layernorm(·)表示层规范化函数的输出和sublayer(h)表示子层的输出;h表示当前子层输入的隐藏层状态。Transformer的层与层之间采用残差连接,因此模型子层输入的尺寸和子层输出的维度相同。
解码器同样由L个相同的层堆叠组成。与编码器不同的是每个层包含三个子层。每一层比编码器多一个子层,该子层是对编码器最后一层的输出进行注意力权重计算。其余与编码器一样,子层与子层之间应用残差连接和层规范化。
模型在给定原文句子序列X=(x1,…,xn)以及目标语言序列Y=(y1,…,ym)下,第i个目标词的优化目标如式(1)所示。
其中,H是所有可能候选翻译的集合,P(yi|Y
注意力机制是式(2)的一个统称,描述的是一个查询项(Query,Q)以及一组键值对(Key-Value,K-V)之间进行映射得到注意力结果。
其中,Q、K、V分别表示查询项、键和值。Q、K、V在进行注意力机制计算前,均会经过一个线性映射,得到的注意力结果最终也会经过一个线性映射得到注意力层的输出。线性映射可以表示为式(3):
(3)
其中,he∈Rn×d,hd∈Rm×d分别表示编码器的输出及解码器当前层的隐层输入;WQ,WK,WV∈Rd×q表示线性映射的权重,q表示变换后的维度信息,d表示模型维度信息。
2 知识蒸馏
神经机器翻译模型可以看作是在样本为(X,Y)的数据集上的一个具有词表大小V类的多分类,训练目标是在训练集上最小化负极大似然损失(negative likelihood loss,NLL)函数,如式(4)所示。
(4)
其中,I(·)表示指示函数,I(y=v)当y等于v时返回1,否则返回0;p(y=v|X;θ)是模型输出的概率分布。
知识蒸馏是模型压缩方法中的一类方法,表示精简的神经网络(学生),在具有更大规模的、性能更好的神经网络(教师)的指导下进行训练更新[8]。
知识蒸馏的目标函数是让学生网络拟合教师网络输出的概率分布,使得学生模型的输出与教师模型的输出概率q(y|X;θt)更加接近[9],知识蒸馏的目标函数,如式(5)所示。
(5)
其中,LKD(θs;θt)表示知识蒸馏的目标函数,θs表示学生网络的参数,θt表示教师网络的参数。
词级别的知识蒸馏是一种标准的知识蒸馏方法,在原有的最小化LNLL(θ)的基础上,新增一个目标函数LKD(θs;θt),并通过一个超参数α来平衡LNLL(θ)和LKD(θs;θt)两者对训练的重要程度,表示如式(6)所示。与词级别知识蒸馏方法不同的是,本文提出的方法在训练期间不需要计算教师网络的输出概率分布,让学生网络的概率与之接近,而是通过教师网络的权重信息指导学生网络,让学生网络的学习起点更高,缩短优化时间。
L(θs;θt)=(1-α)LNLL(θs)+αLKD(θs;θt)
(6)
其中,θs表示学生模型的参数,θt表示教师模型的参数。
句子级别的知识蒸馏是用教师网络生成的翻译结果,与训练数据的原文构成新的平行句对进行训练。与句子级别的知识蒸馏不同,本文方法进行知识蒸馏是把教师模型的权重信息传递给学生网络,而不改变训练数据集。
3 递进式半知识蒸馏
3.1 递进式半知识蒸馏方法
递进式半知识蒸馏方法总体概览图如图2所示,图中的三个精简架构是Transformer的略缩图,图中虚线是编码器和解码器的分割线;箭头代表压缩前和压缩后的模型尺寸变化的对应关系;图中采用阴影填充的模块表示被压缩维度后的权重,未填充的结构表示未被压缩维度的权重。递进式半知识蒸馏方法是通过应用半知识蒸馏两次,实现递进式压缩模型。第一次压缩: 由图2左侧模型架构应用一次半知识蒸馏方法得到中间模型架构。第二次压缩是把上一次压缩的模型(即图2中间的模型架构)作为新的教师模型,进行第二次半知识蒸馏,得到最终的全压缩网络(即图2中右侧的网络架构)。
对于学生网络,学生模型中未被压缩的参数,其维度和教师模型相同,模型被压缩的权重矩阵变成缩小的维度。递进式压缩两次,交替压缩模型的不同部分的参数(编码器和解码器),从而实现模型完整压缩,即如图2右侧的带阴影填充的全压缩模型。
图2 递进式知识蒸馏概览
3.2 结构配置
在不引入新参数的同时保证半知识蒸馏方法完美兼容Transformer,本节对Transformer的结构划分进行说明。从模型压缩的层面来划分,模型可以分为两个部分: 压缩和不压缩部分。从Transformer结构上来划分,前面提到由于Transformer采用残差连接,编码器的所有层的输入输出的维度信息要保持一致,解码器同样如此,故将Tansformer划分为编码器(左边)和解码器(右边),其中在解码器中对编码器的输出的注意力机制层中的映射权重作为划分的共享部分或称为中间部分(即每次压缩,均需要改变维度,以配合兼容编码器和解码器的对应关系)。如图1所示,左侧框内为编码器,右侧框内为解码器,同时将源端词嵌入矩阵和编码器作为模型左侧参数,目标端词嵌入(word embedding)和解码器则作为模型右侧参数部分。
中间部分是编码器的输出和解码器隐藏单元数不同而能兼容的关键。中间部分由解码器中每一个多头注意力层中输入处的线性投影(linear projection)的权重WQ、WK、WV∈Rd×q组成,如式(3)中表示。
3.3 半知识蒸馏方法
半知识蒸馏方法每次压缩模型的一部分权重。学生网络的非压缩部分直接利用教师模型的参数信息进行初始化,作为训练的起点;而学生网络的被压缩权重采用高斯分布初始化。网络参数θ划分成左边半部分参数θl和右边半部分参数θr,教师模型参数和学生模型参数分别用θt和θs表示。在递进式压缩方法中,对于先压缩左边半部分参数θl和先压缩右边半部分参数θr,在前期试验中发现影响不大。本文统一默认使用先压缩右半部分参数θr再压缩左半部分参数θl。 压缩参数和非压缩参数统一构成了比前者小的新模型(即学生模型,参数用θs表示),划分后学生网络的目标函数如式(7)所示。
左半部分则是一个知识蒸馏的过程,作为当前的学生模型,从未压缩的教师模型中获得训练所需要的知识,加快模型收敛,如式(9)所示。
对源语言和目标语言的词嵌入矩阵的处理,将源语言和目标语言的词嵌入矩阵同样划分到模型左右部分中,即每次模型左(右)半部分压缩的同时,将源语言(目标语言)词嵌入矩阵进行压缩。
4 实验
4.1 实验数据
在两个广泛使用的中英和日英数据集上进行实验。这两个翻译数据集是不同规模和不同领域的语料库,分别包含125万句对和32万句对,这有助于验证模型在不同规模训练数据上的性能。
对于日英翻译任务,数据集来自公开语料Kyoto Free Translation Task(1)http://www.phontron.com/kftt。我们使用预处理的训练集作为训练数据,其中预处理后的训练数据过滤了句长小于1和大于40的句子,有关数据处理的更多细节可以在官方网站上查阅。处理后包括33万个平行句对用于训练,对语料采用字符编码进行联合字符编码处理[21],字符编码切分规则大小设为30k。最后,随机打乱训练集并生成词汇表。
中英翻译任务的语料是由语言数据联盟(linguistic data consortium, LDC)提供,选用美国国家标准与技术研究院2002年发布的数据,其中训练数据由1.25M句子对组成;NIST02、NIST03、NIST04、NIST05和NIST08分别作为一个测试集,共5个测试集,每个测试集包含4个参考译文;NIST06作为开发集。对英文语料进行切分(tokenized)和对中文数据进行分词处理,将源端句子和目标译文的最大句子长度均限制为50个词,并将中英文词汇表的大小限制为3万,统计前30万个词作为高频词,其余词作为低频词统一替换成特殊字符
4.2 模型设置
4.2.1 基础设置
将基准(Base)模型的注意力机制的维度、源端和目标端统一设置为512,将前馈网络层的维度设置为2 048。为充分验证比较应用方法,本文在表1中准备了几个不同大小的基准模型,其中,不同维度的基准模型分别用不同的标签(Base、R、SR、S)标注,“ATT”表示注意力机制的维度信息,“FFN”表示前馈神经网络的维度信息。
表1 不同基准模型的维度信息
在编码器和解码器的各个层中的多头注意力机制均使用了8个头。在训练期间,标签平滑(label smoothing)参数=0.1,神经元的随机失活率(dropout)设为0.1。
关于解码参数设置,使用集束搜索(beam search)算法并设置束大小(beam size)B=4。 使用Adam优化器[22],其中优化器参数为β1=0.9,β2=0.98,ε=10-9,学习率使用热启动(warm-up)更新策略[5],其中warmup_steps=4 000的策略。本文实验均在开源的代码THUMT(2)https://github.com/THUNLP-MT/THUMT上进行开发实现。
4.2.2 训练设置
因为中英和日英语料规模不同,因此日英和中英的具体参数设置有所不同。
日英模型将句长接近的句子放到一起,以至于一个批次(mini-batch)的数据的句子长度接近,每一个批次6 250个词,训练2万步,模型间隔1 500步保存一次模型,模型中的解码器的词嵌入权重矩阵和解码器中输出映射到词表的权重矩阵之间共享参数。
在中英翻译实验中,训练设置将句子对按长度组合在一起,每批具有6 025个词。最大训练步数设置为20万,并且以间隔1 000个训练步保存模型。中英翻译模型在解码器的词嵌入权重矩阵和目标端映射到词表的权重矩阵之间均不共享参数。
关于模型评估,在日语到英语的翻译评估中,采用不区分大小写的机器双语互译评估[23](BLEU)得分来评估模型的翻译质量,使用multi-bleu.perl(3)https://github.com/moses-smt/mosesdecoder/blob/master/scripts/generic/multi-bleu.perl计算BLEU值。对保存的最后五个模型进行平均。
中文到英文的模型评估,采用不区分大小写的BLEU评分机制,并使用单个模型进行评估。
4.2.3 参数设置细节
词级别知识蒸馏超参数设置α=0.5,采用基准(Base)模型作为教师模型。
句子级别的知识蒸馏,以基准(Base)模型作为教师模型获取翻译结果,获取翻译结果的参数设置参照Kim等人文章[17]中的参数设置,日英翻译任务设置束大小(beam size)为B=5;中英翻译任务运行集束搜索并设置束大小B=5和B=1两种参数配置,分别返回B个候选翻译译文,最终选择具有更高的BLEU分数序列作为新的最终的翻译结果。此外,如果返回的B个候选翻译译文中有空行,则跳过该句子,以此保证句子级别知识蒸馏提供给学生模型的训练语料质量。
4.3 实验结果
中英翻译结果如表2所示,其中Seq-KD(B=1)[20]和Seq-KD(B=5)[20]分别表示束大小为1和5的情况下的句子级别的知识蒸馏;+Seq-KD(B=1)[20]和+Seq-KD(B=5)[20]分别表示半知识蒸馏和句子级别的知识蒸馏的结合;MT02,…,MT08 分别对应测试集NIST02,…,NIST08;AVG表示测试集的平均得分;BLEU差值表示当前模型平均得分与基准模型得分的差值。表2的结果包含了运用了两次递进式压缩,从512维的基准(Base)模型进行递进式压缩得到256维的全压缩模型,然后进一步进行递进式压缩,得到128维的全压缩模型。根据表2的结果显示,本文方法性能优于基准模型,在与句子级别知识蒸馏方法结合后,能进一步提升性能。
表2 各模型在中英语料上的翻译性能
表3是日英翻译的实验结果,表中,Word-KD代表词级别的知识蒸馏方法[20];Seq-KD代表句子级别的知识蒸馏方法[20];Our work 代表半知识蒸馏方法;BLEU差值代表模型翻译结果与模型相应基准模型的差值,其中加号(+)表示提升,空(-)表示基准模型自身;*表示模型与相应基准模型之间显著性检验[24],p<0.05,检验工具为开源工具。
三种不同的压缩方法都取得了显著性优于基准系统的结果,其中如表3第7~10行,本文提出的方法压缩成最终模型后,超过基准模型2.16个BLEU值,并且超过其他两种知识蒸馏方法,其中超过词级别的知识蒸馏1.15个BLEU值,超过句子级别的知识蒸馏0.28个BLEU值。
表3 日英翻译压缩率与翻译性能
5 分析
5.1 收敛性
图3显示各方法在日英数据集上的BLEU值从第0个训练步均稳步上升,在0~5 000步,本文方法高于其他三个基准模型,从与词级别知识蒸馏和句子级别知识蒸馏的曲线对比可以看出,本文方法在前期学习能力明显高于其他模型。本文方法的曲线在10 000步时开始趋于稳定,且处于其他三条曲线的上方。对比表明,本文方法能加速模型收敛,并有效改善模型泛化性能。
图3 验证集上BLEU值随训练步数变化曲线
5.2 时间和速度
首先,为了公平地衡量词级别和句子级别的时间消耗,表4统计了各方法在日英数据集上的训练时间和新增数据处理的时间,其中,数据代表新增的数据处理时间,单位: 分钟;训练代表训练时长,单位: 分钟;GPU表示GPU数目;kw/s表示千词每秒。默认的数据预处理是相同的,因此未在表中体现,主要以增加的时间作为对比依据。
表4 模型时间消耗和速度统计表
其中句子级别的知识蒸馏,需要教师模型进行翻译训练语料,从得到的翻译文本中挑选BLEU最高的句子作为学生模型的训练语料,此处统计的时间仅为教师模型翻译花费的时间,未包括挑选处理的时间。此外,关于GPU数量,在实现词级别的时候,需要计算教师模型的输出概率分布,由于实验室设备有限,无法在一个GPU上训练,故用了2个GPU,尽管可能对本文提出的方法并不公平,本文的方法总时间消耗明显小于词级别和句子级别的知识蒸馏,在训练阶段的时间消耗,本文方法与句子级别几乎持平,且小于词级别的时间消耗。在训练速度上,本文方法与句子级别的知识蒸馏持平,且高于词级别的训练速度。
6 结论
本文提出了半知识蒸馏方法,该方法从教师模型中提取权重分布,然后赋值给学生网络,从而优化学生模型的训练,加速学生模型的收敛并改善其泛化性能,达到知识蒸馏的效果。与词级别的知识蒸馏相比,本文方法不需要对每次获取教师模型的输出概率进行拟合,减少了计算量;与句子级别的知识蒸馏相比,由从教师模型产生的句子中学习变为直接从教师模型的权重中获取知识。在进行充分实验后,中英翻译实验结果表明,本文方法最佳性能高于基准模型0.72个BLEU值,优于词级别的知识蒸馏0.6个 BLEU值,并达到和句子级别知识蒸馏可比的翻译性能。日英翻译实验结果表明,本文方法比基准模型高2.16个BLEU值。此外本文方法高于词级别的知识蒸馏1.15个BLEU值和高于句子级别的知识蒸馏0.28个BLEU值。