[Paper Review] AdaTask: A Task-aware Adaptive Learning Rate Approach to Multi-task Learning

Enneng Yang, Junwei Pan+7 more|arXiv (Cornell University)|Nov 28, 2022

Domain Adaptation and Few-Shot Learning42 references4 citations

TL;DR

AdaTask proposes a task-aware adaptive learning rate method that separates task-specific accumulative gradients in multi-task learning (MTL) to mitigate parameter-level task dominance. By decoupling learning rates per task per parameter, AdaTask achieves state-of-the-art average performance and significantly improves underperforming tasks without degrading dominant ones.

ABSTRACT

Multi-task learning (MTL) models have demonstrated impressive results in computer vision, natural language processing, and recommender systems. Even though many approaches have been proposed, how well these approaches balance different tasks on each parameter still remains unclear. In this paper, we propose to measure the task dominance degree of a parameter by the total updates of each task on this parameter. Specifically, we compute the total updates by the exponentially decaying Average of the squared Updates (AU) on a parameter from the corresponding task.Based on this novel metric, we observe that many parameters in existing MTL methods, especially those in the higher shared layers, are still dominated by one or several tasks. The dominance of AU is mainly due to the dominance of accumulative gradients from one or several tasks. Motivated by this, we propose a Task-wise Adaptive learning rate approach, AdaTask in short, to separate the \emph{accumulative gradients} and hence the learning rate of each task for each parameter in adaptive learning rate approaches (e.g., AdaGrad, RMSProp, and Adam). Comprehensive experiments on computer vision and recommender system MTL datasets demonstrate that AdaTask significantly improves the performance of dominated tasks, resulting SOTA average task-wise performance. Analysis on both synthetic and real-world datasets shows AdaTask balance parameters in every shared layer well.

Motivation & Objective

• Address the lack of metrics to quantify task dominance in multi-task learning (MTL) parameter optimization.
• Identify that task dominance in MTL arises from imbalanced accumulative gradients in adaptive optimizers like Adam and RMSProp.
• Propose a method to decouple task-specific gradient accumulation to balance learning rates across tasks per parameter.
• Improve performance of dominated tasks while maintaining or slightly reducing performance on dominant tasks for overall SOTA performance.
• Demonstrate that AdaTask achieves better balance in parameter updates across all shared layers in MTL models.

Proposed method

• Introduce a novel metric, task-wise exponentially decaying average of squared parameter updates (AU), to quantify task dominance on each parameter.
• Define the ratio of AU from a single task to total AU as the rAU metric to measure dominance of a task on a parameter.
• Propose AdaTask to separate accumulative gradients per task per parameter in adaptive optimizers (e.g., Adam, RMSProp, AdaGrad).
• Modify existing adaptive optimizers to maintain independent first and second-order moments for each task on shared parameters.
• Apply the task-specific learning rate update rule: learning rate for each task is computed independently based on its own accumulated gradient statistics.
• Ensure compatibility with all major adaptive optimizers (Adam, RMSProp, AMSGrad, etc.) by integrating task-wise gradient tracking into the optimization loop.

Experimental results

Research questions

• RQ1 To what extent do existing MTL methods suffer from task dominance in parameter updates, especially in higher shared layers?
• RQ2 Can task dominance be quantified using a metric based on the exponentially decaying average of squared parameter updates (AU)?
• RQ3 Does separating task-specific accumulative gradients in adaptive optimizers reduce parameter-level task dominance and improve performance on dominated tasks?
• RQ4 Can AdaTask achieve Pareto-optimal performance across all tasks without significantly degrading dominant tasks?
• RQ5 How well does AdaTask balance learning rates and parameter updates across all shared layers in real-world MTL models?

Key findings

• AdaTask significantly improves the performance of dominated tasks—e.g., task A in synthetic, CityScapes, and TikTok datasets—while maintaining strong overall performance.
• On the CityScapes dataset, 99% of shared parameters in EqualWeight and 95% in GradNorm are dominated by task B, but only 7% in AdaTask, demonstrating effective dominance reduction.
• After a few epochs, over 98% of shared parameters in all four layers of the synthetic dataset have rAU values in the (40%, 60%] range under AdaTask, indicating strong balance.
• AdaTask achieves state-of-the-art average task-wise performance across four datasets (synthetic, CityScapes, TikTok, WeChat), outperforming MGDA, GradNorm, UW, PCGrad, and CAGrad.
• The method achieves a Pareto-stationary solution, as no baseline outperforms AdaTask on all tasks simultaneously.
• Learning rate dominance is eliminated in AdaTask, as each task computes its own learning rate per parameter, preventing interference from dominant tasks.