Task Weighting through Gradient Projection for Multitask Learning
This work addresses multitask learning challenges for AI practitioners by providing an incremental improvement to existing gradient projection methods.
The paper tackles the problem of task conflicts in multitask learning by adapting the Gradient Projection algorithm to include task prioritization, replacing uniform projection probabilities with a probability distribution that applies weighting only during conflicts. Experiments on nuScenes, CIFAR-100, and CelebA datasets show significant improvements in performance metrics for most tasks compared to the baseline.
In multitask learning, conflicts between task gradients are a frequent issue degrading a model's training performance. This is commonly addressed by using the Gradient Projection algorithm PCGrad that often leads to faster convergence and improved performance metrics. In this work, we present a method to adapt this algorithm to simultaneously also perform task prioritization. Our approach differs from traditional task weighting performed by scaling task losses in that our weighting scheme applies only in cases where tasks are in conflict, but lets the training proceed unhindered otherwise. We replace task weighting factors by a probability distribution that determines which task gradients get projected in conflict cases. Our experiments on the nuScenes, CIFAR-100, and CelebA datasets confirm that our approach is a practical method for task weighting. Paired with multiple different task weighting schemes, we observe a significant improvement in the performance metrics of most tasks compared to Gradient Projection with uniform projection probabilities.