[论文解读] Knowledge Distillation for Multi-task Learning

Weihong Li, Hakan Bilen|arXiv (Cornell University)|Jul 14, 2020

Domain Adaptation and Few-Shot Learning参考文献 31被引用 5

一句话总结

本文提出了一种基于知识蒸馏的多任务学习方法，通过可学习适配器将多任务网络特征与特定任务模型对齐，从而改善训练平衡性和性能。通过联合最小化特定任务损失和蒸馏损失，该方法在NYU-V2和SVHN & Omniglot等基准数据集上，于语义分割、深度估计和表面法线预测等多样化任务中实现了更优且均衡的性能表现。

ABSTRACT

Multi-task learning (MTL) is to learn one single model that performs multiple tasks for achieving good performance on all tasks and lower cost on computation. Learning such a model requires to jointly optimize losses of a set of tasks with different difficulty levels, magnitudes, and characteristics (e.g. cross-entropy, Euclidean loss), leading to the imbalance problem in multi-task learning. To address the imbalance problem, we propose a knowledge distillation based method in this work. We first learn a task-specific model for each task. We then learn the multi-task model for minimizing task-specific loss and for producing the same feature with task-specific models. As the task-specific network encodes different features, we introduce small task-specific adaptors to project multi-task features to the task-specific features. In this way, the adaptors align the task-specific feature and the multi-task feature, which enables a balanced parameter sharing across tasks. Extensive experimental results demonstrate that our method can optimize a multi-task learning model in a more balanced way and achieve better overall performance.

研究动机与目标

• 解决多任务学习中因损失量级和特性不同而导致的训练主导问题。
• 通过防止高性能任务掩盖低性能任务，提升所有任务的整体性能。
• 通过特定任务适配器对齐特征，实现在语义层次和复杂度各异的任务之间的均衡参数共享。
• 利用预训练的单任务模型作为多任务网络的知识源，实现更好的泛化性和效率。
• 相比损失加权或梯度调制方法，提供更稳定、可控的优化过程，后者常无法有效防止任务主导现象。

提出的方法

• 在离线阶段为每个任务单独训练并冻结其特定任务模型的参数。
• 优化多任务网络，使其同时最小化特定任务损失和蒸馏损失，后者用于对齐其特征与冻结的特定任务模型的特征。
• 引入小型可学习特定任务适配器，将多任务网络的特征映射到每个特定任务模型的特征空间。
• 使用线性或非线性适配器（例如带ReLU的1×1卷积）实现跨不同任务的灵活且高效的特征对齐。
• 在多个网络层（例如中间层和最后一层共享层）应用蒸馏损失，以增强特征对齐并提高优化稳定性。
• 端到端训练多任务网络，采用组合损失：选定层上特定任务损失与蒸馏损失的加权和。

实验结果

研究问题

• RQ1 从特定任务模型进行知识蒸馏是否能改善多任务学习中的平衡性与整体性能？
• RQ2 将多任务网络特征与单任务模型特征对齐，是否能带来更好的泛化能力并减少任务主导现象？
• RQ3 不同适配器架构（线性与非线性）对特征对齐和最终性能有何影响？
• RQ4 在多个网络层而非仅最后一层应用蒸馏损失，其影响是什么？
• RQ5 该方法是否能在平衡多任务优化方面超越现有的损失加权与梯度调制策略？

主要发现

• 所提方法在NYU-V2上达到SOTA性能，语义分割的mIoU达到20.75，优于MTAN及其他基线方法。
• 在NYU-V2上，同时在中间层和最后一层应用蒸馏损失，使mIoU从18.75提升至20.75，证明了多层蒸馏的优势。
• 该方法显著平衡了训练动态：蒸馏损失收敛更均衡，SVHN和Omniglot的蒸馏损失分别为0.089和0.041，而对应的任务特定损失分别为0.137和1.492。
• 在NYU-V2上，分割、深度估计和法线估计的蒸馏损失值（0.534、0.381、0.364）比任务特定损失（0.027、0.127、0.039）更均衡，表明优化过程更加平衡。
• 线性适配器已足够实现性能提升，非线性适配器未带来显著改进，表明简单投影在特征对齐中已足够有效。
• 移除适配器后，mIoU下降至17.11，证实通过适配器实现的特征对齐对性能和平衡至关重要。