[论文解读] f-Divergence Minimization for Sequence-Level Knowledge Distillation

Yuqiao Wen, Zichao Li+2 more|arXiv (Cornell University)|Jul 27, 2023

Topic Modeling被引用 4

一句话总结

该论文提出 f-distill，一种统一的序列级知识蒸馏框架，通过最小化教师模型与学生模型分布之间的 f-散度实现。通过将蒸馏表述为 f-散度最小化，并结合离线采样进行分步分解，f-distill 在四项文本生成任务中均优于现有方法，其中对称散度（JS 和 TVD）有效缓解了模式平均化与模式崩溃问题。

ABSTRACT

Knowledge distillation (KD) is the process of transferring knowledge from a large model to a small one. It has gained increasing attention in the natural language processing community, driven by the demands of compressing ever-growing language models. In this work, we propose an f-DISTILL framework, which formulates sequence-level knowledge distillation as minimizing a generalized f-divergence function. We propose four distilling variants under our framework and show that existing SeqKD and ENGINE approaches are approximations of our f-DISTILL methods. We further derive step-wise decomposition for our f-DISTILL, reducing intractable sequence-level divergence to word-level losses that can be computed in a tractable manner. Experiments across four datasets show that our methods outperform existing KD approaches, and that our symmetric distilling losses can better force the student to learn from the teacher distribution.

研究动机与目标

• 为解决如 KL 和反向 KL 散度等非对称知识蒸馏方法的局限性，前者存在模式平均化问题，后者存在模式崩溃问题。
• 将现有的序列级蒸馏方法（如 SeqKD 和 ENGINE）统一在广义 f-散度最小化框架下。
• 通过分步分解与离线教师采样，实现序列级 f-散度的高效、可计算的计算方法。
• 评估对称 f-散度（JS 和 TVD）在提升学生模型与教师分布对齐效果方面的有效性。
• 在多种文本生成任务中展示一致的性能提升，包括摘要生成、机器翻译和对话生成。

提出的方法

• 将序列级知识蒸馏表述为最小化教师与学生输出分布之间的广义 f-散度。
• 在 f-distill 框架下推导出四种变体：KL、反向 KL、Jensen–Shannon (JS) 和总变差距离 (TVD) 蒸馏。
• 提出将序列级 f-散度分步分解为词级损失，从而实现可计算的优化。
• 引入一种离线采样策略，预先计算教师模型输出，以降低对称散度的训练成本。
• 采用贪婪的、基于梯度的优化方法，通过在学生采样中应用停止梯度，避免强化学习。
• 所有 f-散度变体使用相同的训练目标，从而实现一致的比较与消融实验。

实验结果

研究问题

• RQ1 基于 f-散度最小化的统一框架是否能够超越现有基于 KL 的方法，在序列级知识蒸馏中实现性能提升？
• RQ2 对称 f-散度（JS 和 TVD）与非对称 f-散度（KL 和反向 KL）相比，在缓解模式平均化与模式崩溃方面表现如何？
• RQ3 分步分解与离线采样在多大程度上使序列级 f-散度最小化在计算上变得可行？
• RQ4 f-distill 是否在多种文本生成任务中持续优于现有蒸馏基线方法？
• RQ5 对称蒸馏损失在实践中是否能带来学生与教师分布之间更好的对齐？

主要发现

• f-distill 在四项文本生成数据集（DART、XSum、WMT16 EN-RO 和 Commonsense Dialogue）上持续优于现有分布匹配型知识蒸馏方法。
• 对称 f-散度变体（JS 和 TVD）相比非对称变体（KL 和 RKL）表现更优，证实了模式平均化与模式崩溃的负面影响。
• 当与表示匹配型知识蒸馏结合时，f-distill 实现了显著的性能提升，证明其作为插件模块的兼容性与有效性。
• f-散度的分步分解使原本难以计算的序列级散度实现了可计算性。
• 从教师模型进行离线采样有效降低了对称散度的计算成本，且未牺牲性能。
• 实证结果表明，所提方法在所有评估基准上均达到最先进性能，且在各类指标与设置下均保持一致的性能增益。