[Paper Review] Multitask Learning with Single Gradient Step Update for Task Balancing

Sungjae Lee, Youngdoo Son|arXiv (Cornell University)|May 20, 2020

Domain Adaptation and Few-Shot Learning33 references4 citations

TL;DR

This paper proposes a multitask learning method that balances task gradients using single gradient step updates within a meta-learning framework. By decoupling shared and task-specific layers and optimizing the shared layer via inner/outer loop training with one-step updates, the method mitigates task imbalance and achieves state-of-the-art performance on multiple computer vision benchmarks.

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 separately so that the two layers with different roles in a multitask network can be fitted to their own purposes. 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 problems and achieve state-of-the-art performance.

Motivation & Objective

• To address task imbalance in multitask learning, where some tasks dominate training due to gradient dominance.
• To improve generalization and computational efficiency in multitask deep learning.
• To decouple shared and task-specific network components for more effective learning.
• To develop a gradient-level balancing mechanism using meta-learning principles.
• To achieve state-of-the-art performance across diverse multitask computer vision tasks.

Proposed method

• The method separates the network into shared and task-specific layers, each trained with distinct objectives.
• The shared layer is optimized using a single gradient step update within an inner/outer loop meta-learning framework.
• Inner loop updates are performed on task-specific heads, while outer loop updates refine the shared weights using a meta-objective.
• The meta-optimization process minimizes the weighted sum of task losses, implicitly balancing gradients across tasks.
• The approach enables end-to-end training with improved gradient flow and reduced interference between tasks.
• The method is applied to multiple computer vision tasks, including object detection and semantic segmentation.

Experimental results

Research questions

• RQ1 How can task imbalance in multitask learning be effectively mitigated at the gradient level?
• RQ2 Can a single gradient step update strategy improve performance and stability in multitask training?
• RQ3 Does decoupling shared and task-specific learning improve model generalization and efficiency?
• RQ4 Can meta-learning principles be effectively adapted to balance gradients across diverse tasks?
• RQ5 Does the proposed method achieve state-of-the-art results on standard multitask vision benchmarks?

Key findings

• The proposed method achieves state-of-the-art performance on multiple multitask computer vision benchmarks.
• Task imbalance is significantly reduced, as evidenced by more uniform gradient updates across tasks.
• The single gradient step update strategy improves training stability and convergence speed.
• Decoupling shared and task-specific layers leads to better feature representation and task-specific adaptation.
• The method reduces computational overhead while maintaining high performance, making it efficient for real-world applications.
• Empirical results show consistent improvements over baseline multitask learning methods across diverse vision tasks.