[论文解读] Less is More: Task-aware Layer-wise Distillation for Language Model Compression

Liang Chen, Simiao Zuo+4 more|arXiv (Cornell University)|Oct 4, 2022

Topic Modeling被引用 14

一句话总结

本文提出任务感知的逐层蒸馏（TED），一种利用任务特定过滤器从教师模型隐藏表征中蒸馏出相关知识的方法，仅保留与任务相关的知识，从而减少冗余并提升蒸馏效率。TED在持续预训练和微调两种场景下均达到最先进性能，在GLUE和SQuAD基准上优于现有蒸馏方法。

ABSTRACT

Layer-wise distillation is a powerful tool to compress large models (i.e. teacher models) into small ones (i.e., student models). The student distills knowledge from the teacher by mimicking the hidden representations of the teacher at every intermediate layer. However, layer-wise distillation is difficult. Since the student has a smaller model capacity than the teacher, it is often under-fitted. Furthermore, the hidden representations of the teacher contain redundant information that the student does not necessarily need for the target task's learning. To address these challenges, we propose a novel Task-aware layEr-wise Distillation (TED). TED designs task-aware filters to align the hidden representations of the student and the teacher at each layer. The filters select the knowledge that is useful for the target task from the hidden representations. As such, TED reduces the knowledge gap between the two models and helps the student to fit better on the target task. We evaluate TED in two scenarios: continual pre-training and fine-tuning. TED demonstrates significant and consistent improvements over existing distillation methods in both scenarios. Code is available at https://github.com/cliang1453/task-aware-distillation.

研究动机与目标

• 为解决逐层蒸馏中的知识冗余问题，即学生模型因教师隐藏状态中存在无关信息而难以学习有用表征。
• 通过聚焦于任务相关特征，缩小大教师模型与小学生模型之间的知识差距。
• 通过仅对齐中间层中最具预测性的知识，提升学生模型在低资源设置下的泛化能力和性能。
• 开发一种两阶段训练框架，首先识别任务相关特征，然后高效地进行蒸馏。
• 在持续预训练和微调设置下，持续优于现有蒸馏方法。

提出的方法

• TED为学生和教师模型的每一层引入任务感知过滤器，其形式为带有任务特定头（如分类任务的softmax）的神经网络。
• 在第一阶段，冻结教师和学生模型参数，训练过滤器以利用教师和学生隐藏表征预测目标任务损失，从而学习从表征中提取任务相关特征。
• 在第二阶段，冻结过滤器（不包含任务头），并将学生模型与其过滤器联合微调，以最小化各层教师和学生过滤后表征之间的均方误差（MSE）。
• 蒸馏损失定义为对应教师和学生层过滤输出之间的均方误差，以鼓励在任务相关知识上的对齐。
• 整体训练目标结合了任务特定损失、预测蒸馏损失以及可学习权重的任务感知逐层蒸馏损失。
• 过滤器可实现为线性投影、多层感知机（MLP）或Transformer层，消融实验表明不同实现存在性能权衡。

实验结果

研究问题

• RQ1 任务感知的中间隐藏表征过滤是否能提升小型学生模型的知识蒸馏效率？
• RQ2 从教师表征中过滤掉冗余信息是否能带来优于标准逐层蒸馏的学生模型性能提升？
• RQ3 TED在不同训练范式（如持续预训练和微调）下的表现如何？
• RQ4 两阶段训练框架（先过滤器预训练，再蒸馏）是否能带来优于端到端蒸馏的对齐与泛化效果？
• RQ5 不同过滤器架构（如线性、MLP、Transformer）对蒸馏性能有何影响？

主要发现

• TED在GLUE基准上显著优于标准知识蒸馏和现有逐层蒸馏方法，尤其在低数据场景下表现突出。
• 在SQuAD v1.1和v2.0问答数据集上，使用DeBERTaV3-xsmall学生模型从DeBERTaV3-base教师模型蒸馏，TED取得了最先进结果。
• 在持续预训练中，TED提升了Lambada和WikiText-103数据集上的零样本和少样本迁移性能，展现出强大的泛化能力。
• 消融研究显示，任务感知过滤器有效降低了冗余知识的负面影响，提升了蒸馏效率和学生模型拟合度。
• 两阶段训练过程相比从一开始就联合训练过滤器与学生模型，收敛性与性能更优。
• 使用更复杂的过滤器（如MLP或Transformer）相比简单线性投影能提升性能，尽管计算成本更高。