Multitask learning with single gradient step update for task balancing

Abstract

Multitask learning is a methodology to boost generalization performance and also reduce computational intensity and memory usage. However, learning multiple tasks simultaneously can be more difficult than learning a single task because it can cause imbalance among tasks. To address the imbalance problem, we propose an algorithm to balance between tasks at the gradient level by applying gradient-based meta- learning to multitask learning. The proposed method trains shared layers and task-specific layers sepa-rately so that the two layers with different roles in a multitask network can be fitted to their own pur -poses. In particular, the shared layer that contains informative knowledge shared among tasks is trained by employing single gradient step update and inner/outer loop training to mitigate the imbalance problem at the gradient level. We apply the proposed method to various multitask computer vision prob-lems and achieve state-of-the-art performance. CO 2021 Elsevier B.V. All rights reserved.