Source-linked AI summary
GradNorm: Gradient Normalization for Adaptive Loss Balancing in Deep Multitask Networks
Zhao Chen, Vijay Badrinarayanan, Chen-Yu Lee, Andrew Rabinovich
TL;DR
Deep multitask networks must balance competing tasks while retaining their efficiency and shared representations. This paper introduces GradNorm, which dynamically tunes task loss weights through gradient normalization and training-rate balancing. Across synthetic and real settings, it improves multitask performance, reduces overfitting, and matches or surpasses exhaustive grid search with one asymmetry hyperparameter.
Problem
Multitask networks are efficient and can improve performance, but unequal task contributions produce imbalanced gradients that make shared-feature training difficult.
Method
GradNorm adaptively tunes task loss weights by directly controlling gradient magnitudes and balancing the training rates of different tasks.
Results
GradNorm improves multitask test-time performance across synthetic and real datasets, while matching or surpassing exhaustive grid search and other adaptive weighting methods.
Takeaways & Limitations
Direct gradient interaction provides a powerful way to control multitask learning with substantially less search effort than exhaustive loss-weight tuning.
Takeaways & Limitations
Pathological tasks that remain stuck at majority-class predictions can cause GradNorm to keep increasing their loss weights.
Abstract
from arXiv · showhide
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 $α$. 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.
1. Introduction
Multitask networks promise efficient, robust multi-output prediction but are difficult to train because task imbalance appears as unequal backpropagated gradients. GradNorm addresses this by efficiently tuning task loss weights through direct gradient-magnitude control.
- Motivation: Multitask learning shares weights across tasks, enabling multiple inferences in one forward pass with potential efficiency and performance benefits.The motivation includes embedded systems such as smartphones, wearables, and robots or drones.
- Challenge: Task imbalance impedes robust shared-feature learning because dominant tasks induce relatively large backpropagated gradients.Earlier approaches primarily manipulated forward-pass relationships or network architectures rather than directly addressing gradient imbalance.
- Approach: GradNorm is an efficient multitask loss-balancing algorithm that directly tunes gradient magnitudes through adaptive task weights.Its gradient normalization operates across tasks and uses training-rate balancing as the objective.
- Contribution: GradNorm matches or surpasses exhaustive grid-search performance while requiring tuning of only a single hyperparameter.The paper contrasts this with the substantially more expensive search procedures used for loss balancing.
- Contribution: The method provides a way to control multitask learning through direct gradient interaction.This is presented as a central contribution of the paper.
2. Related Work
Related work develops multitask systems by sharing features, exploiting task relationships, and learning architectures or task weights. GradNorm differs by targeting the gradient imbalance underlying difficult multitask optimization.
- Scope: The broader literature spans deep multitask applications in computer vision and other fields including natural language processing and speech synthesis.The paper’s primary application area is computer vision.
- Multitask vision: Multitask learning supports complete scene understanding by combining multiple robust visual predictions in shared deep networks.Prior systems range from three-task networks to larger collections of vision tasks.
- Task relationships: Prior research exploits task relationships through clustering, deep relationship networks, cross-stitch networks, and architecture-search methods.These approaches learn meaningful relationships or which features to share between tasks.
- Adaptive weighting: Related work also derives task weights from intrinsic task uncertainty using a joint likelihood formulation.This provides a comparison point for adaptive weighting methods.
3. The GradNorm Algorithm
GradNorm adaptively balances multitask training by tuning loss weights so gradient magnitudes reflect relative training rates. It uses a gradient loss at selected shared weights and renormalizes task weights after updates.
- 3.1. Definitions and Preliminaries: The algorithm learns task-weight functions to place gradient norms on a common scale and make different tasks train at similar rates.The multitask loss is linear in the weighted single-task losses, allowing weights to couple directly to backpropagated gradient magnitudes.
- Implementation: GradNorm is generally applied at the last shared layer to reduce computation and increases NYUv2 training time by only approximately 5%.After each update, weights are renormalized so their sum equals T, decoupling normalization from the global learning rate.
- 3.1. Definitions and Preliminaries: GradNorm measures each weighted task gradient with an L2 norm, averages these norms across tasks, and defines relative inverse training rates from normalized loss ratios.The loss ratio is Li(t)/Li(0); lower values indicate faster training.
- 3.2. Balancing Gradients with GradNorm: The asymmetry hyperparameter α controls the restoring force toward a common training rate.Higher α is intended for tasks with more different learning dynamics, while α = 0 targets equal gradient norms.
- 3.2. Balancing Gradients with GradNorm: GradNorm computes an L1 gradient loss between actual and target task gradient norms, differentiates it with respect to task weights, and updates those weights.The target is treated as fixed during differentiation to prevent weights from drifting spuriously toward zero.
4. A Toy Example
Toy regression experiments isolate loss-scale asymmetry across tasks and show that GradNorm improves balance and test performance over equal weighting and uncertainty weighting. Its learned weights are stable and close to effective static weights.
- 4. A Toy Example: The toy system uses regression tasks with similar loss functions but different output scales, creating task asymmetry through different expected squared-loss magnitudes.Higher-output-scale tasks are harder to learn and generate larger gradients.
- 4. A Toy Example: The experiments use a shared four-layer ReLU trunk with task-specific final predictions and fixed initialization and data seeds for valid comparisons.The synthetic tasks share information while retaining task-specific components.
- Evaluation: Task-normalized test-time loss sums per-task loss ratios because unnormalized loss sums overemphasize tasks with larger loss scales.This metric is especially appropriate here because the toy tasks are statistically identical except for their loss scales.
- Results: For T = 2 with scales (1.0, 100.0), equal weighting lets the high-scale task suppress the other task, whereas GradNorm increases the smaller-scale task’s weight and improves test performance.The weight adjustment counters the larger gradients from the high-scale task.
- Results: For T = 10, GradNorm significantly improves test-time performance over equal weighting and grows weights for tasks with smaller output scales.For both T = 2 and T = 10, it is more stable and outperforms uncertainty weighting.
- Results: GradNorm weight traces are stable and convergent, with time-averaged weights lying close to optimal static weights.This suggests that adaptive training can simplify exhaustive static-weight search.
5. Application to a Large Real-World Dataset
Experiments on NYUv2 evaluate GradNorm across augmented datasets, architectures, multitask baselines, and static-weight searches. GradNorm improves multitask performance, reduces overfitting, and approaches or exceeds strong single-task and grid-search comparisons.
- Datasets: NYUv2 provides indoor-scene depth, surface-normal, segmentation, and keypoint tasks for testing GradNorm across regression and classification settings.The experiments use NYUv2+seg and an augmented NYUv2+kpts dataset.
- Models and training: The evaluation uses VGG16 SegNet and a thinner ResNet-50 FCN, with task-specific losses for segmentation, depth, normals, and keypoints.The losses are cross entropy, squared loss, cosine similarity, and pixel-wise squared loss, respectively.
- Main results: GradNorm α = 1.5 improves all three NYUv2+seg tasks over equal weights and matches or surpasses single-network performance.Its time-averaged weights also lie close to weights obtained through exhaustive grid search.
- Main results: GradNorm networks on the 50x-augmented NYUv2+kpts dataset outperform other multitask methods and match or surpass single-task networks.This extends the observed results to a substantially larger dataset.
- Training dynamics: ∼5% lower test-time depth error occurs despite higher training loss, while GradNorm suppresses the depth weight below 0.10.The authors interpret the divergence between test and training behavior as a regularization signal; keypoint regression shows the same trend.
- Grid-search comparison: Even after 100 random-weight networks, grid search falls short of GradNorm; at L2 distance ∼3, average errors are almost doubled per task.The result indicates that GradNorm reaches the optimal grid-search weights in one training run.
- Asymmetry parameter: Higher α values push task weights farther apart, and nearly every 0 < α < 3 improves NYUv2 performance over equal weights.At α = 1.75, the depth weight is suppressed below 0.02 without harming depth performance.
6. Conclusions
GradNorm balances multitask loss weights by training rates and improves test-time performance across synthetic and real datasets. It can outperform adaptive weighting methods and match or surpass exhaustive grid search while requiring less time.
- GradNorm tunes multitask loss weights by balancing the training rates of different tasks.
- GradNorm improves multitask test-time performance across a variety of synthetic and real-data scenarios.
- GradNorm offers superior performance over state-of-the-art adaptive weighting methods while matching or surpassing exhaustive grid search with substantially less time.
- GradNorm may extend beyond multitask learning to class-balancing and sequence-to-sequence models involving conflicting gradient signals.
7. GradNorm: Gradient Normalization for Adaptive Loss Balancing in Deep Multitask Networks: Supplementary Materials
The supplementary experiments examine how α affects performance and how GradNorm handles mixed regression and classification tasks. Results show broadly positive gains and strong performance, with instability at large α and limitations for pathological tasks.
- Performance versus α: Low α keeps task gradient norms similar, whereas high α penalizes tasks whose losses decrease too quickly and favors slower-learning tasks.
- Performance versus α: Performance gains occur at almost all positive, numerically stable α values and peak around α = 1.5 for NYUv2+kpts across both backbones.
- Performance versus α: At α ≥3 in NYUv2, some task weights approach zero and GradNorm updates become unstable.
- Performance versus α: The ResNet architecture is more robust to α than VGG, although both achieve similar gains with an appropriate setting.
- MTFL experiments: The MTFL dataset combines five facial-landmark regression targets with four classification labels and is evaluated at 40x40 and 160x160 resolutions.
- MTFL experiments: For both resolutions, GradNorm outperforms other methods on nearly all tasks and especially improves gender and smile classification.
- MTFL experiments: GradNorm identifies undertrained classification tasks and boosts them, while keypoint and failure-rate improvements may fall within error bars.
- MTFL limitations: Glasses and pose classification remain stuck at the majority classifier, causing GradNorm to keep increasing their loss weights.