NAACL 2022 | 机器翻译SOTA模型的蒸馏

©PaperWeekly 原创 · 作者 | BNDSBilly

研究方向 | 自然语言处理

Background

ICLR 2021 的一篇文章提出了基于 KNN 方法的机器翻译（kNN-MT），可以将 kNN 方法添加到现有的神经机器翻译模型（NMT）上，从而进一步提升推理表现。该方法帮助当时的 SOTA 德语-英语翻译模型提升了 1.5 BLEU 分数，并且还可以适应跨领域及零样本传输。

本次要分享的论文则是针对 kNN-MT 推理速度过慢的不足，提出了蒸馏方法（kNN-KD）。从而在保持 kNN-MT 表现的情况下，将推理速度提升到了与一般 NMT 模型推理速度相当的水平。

论文标题：

Nearest Neighbor Knowledge Distillation for Neural Machine Translation

收录会议：

NAACL 2022

论文链接：

https://arxiv.org/abs/2205.00479

Methods

2.1 kNN-MT

KNN-MT 方法有两个步骤：

1. Datastore creation：

根据训练集每一条样本离线构建的键值对组合，如下公式所示。其中表示样本的源语言句和目标语言句，为翻译过程中第步时已经推理出来的文本，表示第步需要推理的目标语言 token。表示经过模型 decoder 编码得到的高维向量。

2. Generation：

推理阶段的每一步时，首先根据 NMT 模型给出下一个 token 的输出概率，然后根据 kNN 方法给出下一个 token 的输出概率，最终的输出概率为。

kNN 输出概率如下：按照构造 Datastore 的方式，根据当前的测试样本先构建当前步骤的 key，然后遍历 Datastore 找到距离最近的个结果，将其距离进行一系列操作后，转化为对应 value 的输出概率，如下图所示：

在一般训练 NMT 模型时，通常使用模型预测结果和 grount-truth 的交叉熵（CE）进行训练。但在自然语言中，一个句子通常有多种表达，如果模型预测出一个合理但偏离 grount-truth 的词，CE损失也会将其视为错误并惩罚模型，导致模型泛化性变差，这就是所谓的 overcorrection。而在 KNN-MT 中，在解码阶段综合考虑了其他可能的合理解释，在一定程度上缓解了该问题，所以表现有了明显提升。

2.2 kNN-KD

针对 kNN-MT 推理速度很慢的劣势，本文作者提出了 kNN-KD 方法，步骤如下：

1. Datastore creation：与 kNN-MT 相同

2. Distillation：

对于教师模型，在训练前针对每一条训练样本的每一步骤，都按照类似 kNN-MT 中的方法输出下一 token 的生成概率。

对于学生模型，针对每一条训练样本的每一步骤，都正常输出下一 token 的生成概率。

训练过程中，蒸馏损失为教师模型和学生模型表现的交叉熵：

最终的训练损失即为：

3. Generation：在最终的推理阶段，就不需要再进行 kNN 搜索了，只要按照正常的 NMT 模型进行翻译即可。

KNN-KD 的整体工作流如下图所示：

Experiments

3.1 Setup

本文使用 IWSLT'14 德语-英语（De-En，160k 训练样本）、IWSLT'15 英语-越南语（En-Vi，113k 训练样本）和多域翻译数据集（De-En，733k 训练样本）进行了实验。使用 tst2012 作为验证集，使用 tst2013 作为测试集，分别包含和个句子。

本文所提出的 kNN-KD 是一种无架构方法，可应用于任意 Seq2Seq 模型，可以与其他提升性能的工作同时应用。因此，作者主要将 kNN-KD 与 kNN-MT 以及一些典型的 KD 方法进行比较，包括但不限于 Word-KD、Seq-KD、BERT-KD 和 Selective-KD 等。

实验中所有算法都利用 pytorch 中的 fairseq 工具包实现，在个 NVIDIA GTX 1080Ti GPU 上进行。实验模型选取层 Transformer。对于 IWSLT'14 和 IWSLT'15 模型，配置 embedding size 为，feed-forward size 为，attention heads 为。针对跨领域数据集，配置 embedding size 为，feed-forward size 为，attention heads 为。

作者提前对和（归一化温度）进行了网格搜索，并选取了验证集上的最佳 BLEU 分数对应的超参数，如下表所示，其中表示 Datastore 中数据个数：