Knowledge Distillation for Multi-task Learning

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.

We propose a knowledge distillation method for addressing the imbalance problem in multi-task learning

References in corpus (12)

• Adam: A Method for Stochastic Optimization
• Distilling the Knowledge in a Neural Network
• An Overview of Multi-Task Learning in Deep Neural Networks
• GradNorm: Gradient Normalization for Adaptive Loss Balancing in Deep Multitask Networks
• Multi-Task Learning as Multi-Objective Optimization
• Actor-Mimic: Deep Multitask and Transfer Reinforcement Learning
• Towards Understanding Knowledge Distillation
• BERT and PALs: Projected Attention Layers for Efficient Adaptation in Multi-Task Learning
• Gradient Surgery for Multi-Task Learning
• Adapting Auxiliary Losses Using Gradient Similarity
• Contrastive Representation Distillation
• Graph Representation Learning via Multi-task Knowledge Distillation

Cited by in corpus (4)

• Federated Mutual Learning
• Scaling End-to-End Models for Large-Scale Multilingual ASR
• Cross-Task Consistency Learning Framework for Multi-Task Learning
• Learning Multiple Dense Prediction Tasks from Partially Annotated Data