GradNorm: Gradient Normalization for Adaptive Loss Balancing in Deep Multitask Networks

TL;DR

GradNorm is a gradient normalization technique that adaptively balances multitask neural network training, improving accuracy and reducing overfitting across various tasks with minimal hyperparameter tuning.

Contribution

Introduces GradNorm, a simple yet effective gradient normalization method for automatic loss balancing in deep multitask networks, outperforming static and grid search approaches.

Findings

GradNorm improves accuracy across multiple tasks.

It reduces overfitting compared to baseline methods.

It matches or surpasses exhaustive grid search performance.

Abstract

Deep multitask networks, in which one neural network produces multiple predictive outputs, can offer better speed and performance than their single-task counterparts but are challenging to train properly. We present a gradient normalization (GradNorm) algorithm that automatically balances training in deep multitask models by dynamically tuning gradient magnitudes. We show that for various network architectures, for both regression and classification tasks, and on both synthetic and real datasets, GradNorm improves accuracy and reduces overfitting across multiple tasks when compared to single-task networks, static baselines, and other adaptive multitask loss balancing techniques. GradNorm also matches or surpasses the performance of exhaustive grid search methods, despite only involving a single asymmetry hyperparameter $\alpha$. Thus, what was once a tedious search process that incurred exponentially more compute for each task added can now be accomplished within a few training runs, irrespective of the number of tasks. Ultimately, we will demonstrate that gradient manipulation affords us great control over the training dynamics of multitask networks and may be one of the keys to unlocking the potential of multitask learning.