[论文解读] BERT-EMD: Many-to-Many Layer Mapping for BERT Compression with Earth Mover's Distance

Jianquan Li, Xiaokang Liu+4 more|arXiv (Cornell University)|Oct 13, 2020

Topic Modeling参考文献 24被引用 7

一句话总结

该论文提出 BERT-EMD，一种新颖的 BERT蒸馏方法，通过使用地球移动距离（EMD）实现教师模型与学生模型之间的多对多层映射，以优化知识迁移。通过采用成本注意力机制学习自适应的层权重，BERT-EMD 实现了最先进的压缩性能，在仅使用 50% 参数量和推理时间的情况下，于 9 个 GLUE 任务中的 7 个任务上超越了 BERT-Base 模型。

ABSTRACT

Pre-trained language models (e.g., BERT) have achieved significant success in various natural language processing (NLP) tasks. However, high storage and computational costs obstruct pre-trained language models to be effectively deployed on resource-constrained devices. In this paper, we propose a novel BERT distillation method based on many-to-many layer mapping, which allows each intermediate student layer to learn from any intermediate teacher layers. In this way, our model can learn from different teacher layers adaptively for various NLP tasks. %motivated by the intuition that different NLP tasks require different levels of linguistic knowledge contained in the intermediate layers of BERT. In addition, we leverage Earth Mover's Distance (EMD) to compute the minimum cumulative cost that must be paid to transform knowledge from teacher network to student network. EMD enables the effective matching for many-to-many layer mapping. %EMD can be applied to network layers with different sizes and effectively measures semantic distance between the teacher network and student network. Furthermore, we propose a cost attention mechanism to learn the layer weights used in EMD automatically, which is supposed to further improve the model's performance and accelerate convergence time. Extensive experiments on GLUE benchmark demonstrate that our model achieves competitive performance compared to strong competitors in terms of both accuracy and model compression.

研究动机与目标

• 为解决 BERT 蒸馏中一对一层映射的局限性，该方法缺乏理论指导，且无法适应多样化的自然语言处理任务。
• 使每个学生层能够动态地从任意中间教师层中学习，以适应特定任务的语言知识需求。
• 通过将教师-学生层间距离建模为最优传输问题，利用地球移动距离（EMD）优化知识迁移。
• 通过成本注意力机制自动学习每一教师层的重要性，从而提升性能并加速收敛。
• 在不牺牲性能的前提下实现高效率压缩，使模型可部署于资源受限设备上。

提出的方法

• 提出一种多对多层映射策略，使每个中间学生层可关注所有中间教师层，实现在不同任务间的自适应知识迁移。
• 使用地球移动距离（EMD）计算从教师到学生知识转换的最小累积成本，建模隐藏表示的最优传输。
• 引入成本注意力机制，学习基于任务的层权重，动态分配不同教师层的重要性。
• 采用运输问题的公式化方法求解基于 EMD 的知识蒸馏，确保整体且高效的知识迁移。
• 使用基于 EMD 和注意力加权层匹配的知识蒸馏损失进行学生模型训练，实现端到端优化。
• 在 https://github.com/lxk00/BERT-EMD 上发布代码与数据以支持可复现性。

实验结果

研究问题

• RQ1 与一对一映射相比，多对多层映射策略是否能提升 BERT 压缩中的知识蒸馏效果？
• RQ2 在 BERT 蒸馏中，使用地球移动距离（EMD）作为距离度量是否能带来更好的知识迁移与模型性能？
• RQ3 可学习的成本注意力机制是否能提升蒸馏过程的适应性与收敛速度？
• RQ4 BERT-EMD 是否在多种自然语言处理任务中实现了显著减小的模型尺寸与推理时间，同时保持具有竞争力的性能？
• RQ5 在不同语言复杂度的任务中（如 RTE 与 MNLI），模型性能如何变化？

主要发现

• BERT-EMD 仅使用 6 层，在 9 个 GLUE 任务中的 7 个上超越了原始的 12 层 BERT-Base 模型，展现出强大的压缩效率。
• 在 RTE 任务中，BERT-EMD 相较于 BERT-Base 实现了 5.3% 的准确率提升，凸显其在复杂推理任务中的有效性。
• 在 STS-B 任务中，BERT-EMD 的斯皮尔曼等级相关系数较 BERT-Base 提升 1%，表明其在回归类任务中表现更优。
• 消融实验表明，若移除 EMD 模块，性能显著下降，尤其在 STS-B 任务上，验证了 EMD 在知识最优传输中的关键作用。
• 成本注意力机制进一步提升了性能并加速了收敛，证实其在学习任务特定层重要性方面的价值。
• 流量矩阵的可视化显示，不同任务强调不同的教师层——例如，RTE 使用非对角线映射，而 MNLI 则不同，证明了自适应、任务感知的知识迁移。