本文介绍了上海科技大学屠可伟课题组与阿里巴巴达摩院的一项合作研究，提出了在结构预测问题上一种较为通用的结构化知识蒸馏方法。该论文已被 ACL 2021 接受为长文。
论文标题：
Structural Knowledge Distillation: Tractably Distilling Information for Structured Predictor
论文地址：
http://faculty.sist.shanghaitech.edu.cn/faculty/tukw/acl21kd.pdf
简介
知识蒸馏（knowledge distillation，简称 KD）是模型压缩方法的一种，用于将知识从一个复杂的教师模型转移到一个简单的学生模型中。KD 的基本思想是希望学生模型的预测尽量接近教师模型的预测。自然语言处理和计算机视觉的很多任务要求结构化的输出，例如图片的像素级标签、句子的单词级标签序列等。这时就需要使用结构知识蒸馏（structural KD）在模型之间转移这种结构化的信息。但结构输出空间一般是输出尺寸的指数大小，直接求解和优化 structural KD 的目标函数是不可行的。之前的工作一般是在特定情形中采用近似求解或将问题转换成非结构化知识蒸馏求解。本文则根据很多常见结构预测模型会把输出结构的打分分解成多个子结构的打分之和，且子结构空间远小于结构输出空间的特性，提出了在一定条件下 structural KD 目标函数的多项式复杂度精确求解的方法。
背景知识
2.1 知识蒸馏
知识蒸馏框架通常包含一个或多个大型的已训练的教师模型和小型的学生模型，训练时知识蒸馏会鼓励学生模型的预测尽量接近教师模型的预测，一般来说就是将两者的交叉熵作为目标函数。假设我们要训练一个多分类模型，类别集合为 。在样本 x 上教师模型的预测分布为 ，学生模型的预测分布为 ，则知识蒸馏要最大化两者的交叉熵：
同时我们使用训练集上的样本 训练学生模型，目标函数为：
2.2 词级别知识蒸馏
给定输入序列 ，输出序列 。如果将所有输出 都看作相互独立的，那么就可以在序列的每个位置分别做知识蒸馏，即词级别知识蒸馏（token-wise knowledge distillation），对每个位置的 求和得到整个序列的知识蒸馏目标函数：
2.3 结构知识蒸馏
但实际上不同位置的输出往往并不是相互独立的，比如用 BIOES 标注 NER 标签的例子：“希尔顿离开北京机场了”，其标签为”B-PER, I-PER, E-PER, O, O, B-LOC, I-LOC, I-LOC, E-LOC, O“。I 标签前只能是 B 标签，不能是其他标签，类似的还有 E 标签前必须是 B 或者 I 标签等规则。假设这个例子中只有 PER 和 LOC 两种 NER 标签，那么每个词的标注有 9 种，长度为 10 的输出序列理论上有 种组合，但实际上违反了 BIOES 标注规则的结构都不可能出现。因此常见的序列结构预测模型例如线性链 CRF 会建模不同位置输出之间的相关性。同理，很多其他形式的结构预测模型（例如树或图的预测）也会建模输出结构不同部分之间的相关性。我们希望能在结构预测模型的知识蒸馏中，让学生模型学习到教师模型对完整结构的预测，亦即结构知识蒸馏（structrual knowledge distillation）。令所有可能的结构输出序列构成结构输出空间 ，则结构知识蒸馏的目标函数为：
由于结构输出空间大小 往往是输出序列长度 n 的指数函数，因此直接按上式计算是不可行的。我们去年在 ACL2020 发表的论文“Structure-Level Knowledge Distillation For Multilingual Sequence Labeling”提出了两种近似方法。一种方法是只考虑 中概率最大的若干输出结构，在这些输出序列上求和，计算目标函数的近似值。如可以使用动态规划算法预测出前 k 个概率最高的输出结构 ，并假设教师模型在这 k 个输出序列上是均匀分布，从而得到 top-K 知识蒸馏：
如果对教师模型在这 k 个输出结构上的预测概率进行一个估计，就得到 weighted top-K 知识蒸馏：
另一种方法是在一些特定场景下，通过全局结构分布求出局部的边缘分布，将问题强行变成一个非结构化知识蒸馏问题。Posterior Distillation 方法就是在线性链 CRF 模型之间进行知识蒸馏时，用前向后向算法求出模型在每个 token 上输出的边缘分布，然后再使用 token-wise KD。上述第一种方法舍弃了概率不大的结构输出，第二种方法则优化的不再是真正的结构知识蒸馏目标函数，两者都是近似解决结构知识蒸馏。
本文的工作
3.1 Structural KD
上面介绍的方法都是近似方法，本文则介绍了在一定条件下一种精确求解结构知识蒸馏目标函数的方法。很多结构预测模型会将结构输出分解成许多子结构，结构输出的分数是由子结构的分数求和得到。因此完整结构输出的概率公式如下：
其中 y 是某个结构输出， 是模型对 y 的打分， 表示 u 是 y 的一个子结构；我们定义子结构空间 ，即所有可能子结构的集合； 是 partition function。例如前例中的序列标注任务。如果是最大熵（MaxEnt）模型，模型对每个位置预测的标签打分，整个标签序列的分数等于各个标签的分数之和：，y 的子结构即每个位置的输出 。子结构空间大小为 。如果使用线性链 CRF 模型，则模型对每个位置的输出 和相邻两个位置的输出对 打分，y 的子结构即所有 和 ，因而，子结构空间大小为 。如果结构知识蒸馏的学生模型是基于这种结构分解模型，我们可以简化结构知识蒸馏目标函数如下：
如果同时满足条件：1. 学生模型的子结构空间 是多项式大小，2. 教师模型可快速计算学生模型子结构的边缘分布，那么这个目标函数可以在多项式时间内精确求解。
3.2 四类场景
下面按照教师和学生模型子结构分解的差异分成四类场景举例讨论，其中所有场景都讨论了命名实体识别（NER）任务，场景 1 和 2 还讨论了依存分析（dependency parsing）任务。
应用场景举例：在线服务需要快速响应，因此需要将准确性很高但使用了大型预训练上下文相关词向量嵌入的教师网络压缩成只使用静态词向量嵌入的轻量级学生网络。
本场景教师模型和学生模型都是线性链 CRF，两者有相同的子结构空间。我们将发射分数 合并到了转移分数 中。此时学生模型的子结构空间。教师模型的子结构分数，子结构边缘分布可由计算，其中 由经典的前向后向算法得到。Case 1b: 基于图的依存分析到基于序列标注的依存分析
教师模型和学生模型分别是基于图和基于序列标注的依存分析模型。教师模型的准确率要远远高于学生模型，但学生模型的速度更快，因此两者之间的知识蒸馏是有意义的。
这里依存分析使用 head-selection 方式，对输入句子的每个单词 预测一个 dependency head 和 label ，即结构输出 ，。学生模型是一个MaxEnt模型，对每个单词预测 ，即子结构空间。教师模型是一个 Biaffine 模型，对两个单词 之间的 dependency arc 进行建模，预测 是否是 的 head，以及预测这个 dependency arc 的 label。即对每个单词独立预测 和 并分别得到 ，所以其子结构边缘分布为。
以往的研究表明序列标注任务上线性链 CRF 的性能一般都要高于 MaxEnt 模型，但是往往实际应用中我们希望能获得 MaxEnt 模型的速度和效率，因此这个场景下我们设计了在线性链 CRF 模型和 MaxEnt 模型之间进行知识蒸馏。
此时教师模型的子结构空间为，学生模型的子结构空间为 。这时我们可以使用前向后向算法求出教师模型在学生模型的子结构上的边缘分布：。Case 2b: 二阶依存分析到基于序列标注的依存分析这个场景类似于 Case 1b，不同之处是我们将教师模型替换成了性能更高的二阶依存分析器。教师模型对 dependency arc 的存在建模使用了二阶依存分析，即对所有共享一个单词的两个 arc 进行打分；对 arc 的类型预测则与 Case 1b 一样。因此教师模型的子结构空间包含所有共享一个单词的 dependency arc 对。学生模型同 Case 1b，子结构空间。教师模型可以使用平均场变分推断方法估计 arc 的存在概率 ，arc 的类型预测 同 Case 1b，因此子结构边缘分布。
这个场景与 Case 2a 的蒸馏方向相反。一个应用场景是零样本跨语言命名实体识别任务：现有大部分基于预训练多语言 BERT 的模型使用 MaxEnt 解码，我们可以将其迁移到使用静态词向量嵌入的线性链 CRF 模型以获得更快的速度。
由 Case 2a 可知 ，。容易得到子结构边缘分布。Case 4: NER as Paring 模型到最大熵模型
这里的教师模型是当前准确度最高的命名实体识别模型之一，但是其时间复杂度比通常的序列标注模型要高，因此在这个场景中我们将其蒸馏到一个 MaxEnt 模型中。
教师模型用依存分析的方式做命名实体识别任务，即对一个候选实体的头尾两个词 之间的 dependency arc 打分，用 arc 的类型表示命名实体的类型。而学生模型是最大熵模型使用 BIOES 标注。学生模型的子结构空间是 ， 取值为 BIOES 标注的 NER 标签，如“B-PER”。教师模型使用动态规划可以得到子结构边缘分布 。
实验
本文在上述四个场景的不同 case 进行了实验。命名实体识别任务使用了 CoNLL2002/2003 和 WikiAnn 数据集，用到的无标签数据从 WikiAnn 数据集采样得到。对于依存分析使用了 Penn Treebank（PTB）3.0 数据集，用到的无标签数据从 BLLIP 数据集得到。作为对比的基准知识蒸馏方法是前面介绍的 token-wise KD，weighted top-K KD 和 Posterior KD，表中记为：Token. KD, Top-WK KD 和 Pos. KD。文中提出的 structural KD 方法记为 struct. KD。
这是在有标签数据上进行上述实验的结果，w/o KD 为学生模型进行监督学习的结果。然后在训练时加入了无标签数据。先用教师模型对 unlabel 数据预测获得伪标签，Case 3 只包含带伪标签的训练数据，其他 Case 使用的训练数据是有标签和伪标签数据的混合。Top-1 是仅使用这样的训练数据来训练学生模型。
token-wise KD 和 structural KD 的比较Case 2a 中学生模型都是 MaxEnt，比较了 MaxEnt 教师模型加 token-wise KD和 CRF 教师模型加 structural KD 的组合。Case 3 中教师都是 MaxEnt 模型，比较了 Token-wise KD 加 MaxEnt 学生模型和 Structural KD 加 CRF 学生的组合。
由上述各表可以看出各种条件下 structural KD 的效果一般都是最好的。
结语
本文提出了在一定条件下可以多项式时间精确求解和优化结构知识蒸馏目标函数的方法，从而扩展了结构知识蒸馏的应用范围，是对知识蒸馏方法的一个非常有意义的探索。上海科技大学信息学院屠可伟老师研究组主要从事自然语言处理、机器学习等人工智能领域的研究，目前侧重于研究语言结构的表示、学习与应用。研究组近几年已发表顶会论文数十篇，2021 年至今已发表 ACL 长文 7 篇、ACL Findings 长文 1 篇、NAACL 长文 1 篇。研究组现招收硕士研究生（推免生）、博士后和研究助理，欢迎有兴趣的同学联系屠老师。
