G RAD N ORM : G RADIENT N ORMALIZATION FOR A DAPTIVE L OSS B ALANCING IN D EEP M ULTITASK N ETWORKS
Anonymous authors
Paper under double-blind review
A BSTRACT
Deep multitask networks, in which one neural network produces multiple predictive outputs, are more scalable and often better regularized than their single-task counterparts.
Such advantages can potentially lead to gains in both speed and performance, but multitask networks are also difﬁcult to train without ﬁnding the right balance between tasks.
We present a novel gradient normalization (GradNorm) technique which automatically balances the multitask loss function by directly tuning the gradients to equalize task training rates.
We show that for various network architectures, for both regression and classiﬁcation tasks, and on both synthetic and real datasets, GradNorm improves accuracy and reduces overﬁtting over single 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 which 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 hope to 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.
I NTRODUCTION
Single-task learning in computer vision has enjoyed much success in deep learning, with many models now performing at or beyond human accuracies for a wide array of tasks.
However, a system that strives for full scene understanding cannot focus on one problem, but needs to perform many diverse perceptual tasks simultaneously.
Such systems must also be efﬁcient, especially within the restrictions of limited compute environments in embedded systems such as smartphones, wearable devices, and robots/drones.
Multitask learning most naturally lends itself to this problem by sharing weights amongst different tasks within the same model and producing multiple predictions in one forward pass.
Such networks are not only scalable, but the shared features within these networks tend to be better regularized and boost performance as a result.
In the ideal limit, we can thus have the best of both worlds with multitask networks: both more efﬁciency and higher performance.
The key difﬁculty in multitask learning lies in the balancing of tasks, and perhaps the simplest way to control this balance is to choose the correct joint loss function.
In practice, the multitask loss function is often assumed to be linear in the single task losses, L = (cid:80)
i w i L i , where the sum runs over T tasks.
The challenge is then to ﬁnd the best value for each w i that balances the contribution of each task for optimal model training.
Our proposed method is furthermore an adaptive method, allowing w i to vary with the training step t , and so w i = w i ( t ) .
Our key insight lies in the observation that these w i ( t ) inﬂuence training only because they control the magnitude of the gradients generated from task i .
As such, manipulating the gradient norms themselves would be a more direct way to control the training dynamics.
More speciﬁcally, we propose a simple heuristic that penalizes the network when backpropagated gradients from any task are too large or too small.
The correct balance is struck when tasks are training at similar rates; if task i is training relatively quickly, then its weight w i ( t ) should decrease relative to other task weights w j ( t ) | j (cid:54) = i to allow other tasks more inﬂuence on the network.
Our method can be said to be a form of batch normalization (Ioffe & Szegedy (2015)) for backpropagation, ensuring that gradients from each task per batch lie on a common statistical scale.
We will show that, when implemented, gradient normalization leads across-the-board improvements in accuracy and suppresses overﬁtting.
Our main contributions to the ﬁeld of multitask learning are as follows:
An attractively simple heuristic for multitask loss balancing involving training rate equalization, which is implemented through a novel gradient loss function.
A simpliﬁcation to exhaustive grid search (which has compute complexity O ( N T ) for N grid points in one dimension) that only involves tuning one robust hyperparameter.
Demonstration that direct interaction with gradients provides a powerful way of reasoning about multitask learning.
R ELATED W ORK
Multitask learning has existed well before the advent of deep learning (Caruana (1998);
Bakker &
Heskes (2003)), but the robust learned features within deep networks have spurned renewed interest.
Although our primary application area is computer vision, multitask learning has applications in multiple other ﬁelds, from natural language processing (Hashimoto et al.
Collobert & Weston (2008); Søgaard & Goldberg (2016)) to speech synthesis (Wu et al.
; Seltzer & Droppo (2013)), from very domain-speciﬁc applications like trafﬁc prediction (Huang et al.
to very general cross-domain work (Bilen & Vedaldi (2017)).
Multitask learning is very well-suited to the ﬁeld of computer vision, where making multiple robust predictions is crucial for complete scene understanding.
Deep networks have been used to solve various subsets of multiple vision tasks, from 3-task networks (Eigen & Fergus (2015); Teichmann et al.
to much larger subsets as in UberNet (Kokkinos (2016)).
Often, single computer vision problems can even be framed as multitask problems, such as in Mask R-CNN for instance segmentation (He et al.
(2017)) or YOLO-9000 for object detection (Redmon & Farhadi (2016)).
Researchers often assume a ﬁxed loss function or network architecture, but there has also been signiﬁcant work on ﬁnding optimal ways to relate tasks to each other in a multitask model.
Clustering methods have shown success beyond deep models (Kang et al.
; Jacob et al.
(2009)), while constructs such as deep relationship networks (Long & Wang (2015)) and cross-stich networks (Misra et al.
(2016)) search for meaningful relationships between tasks and learn which features to share between them.
Work in Warde-Farley et al. (2014) and Lu et al.
(2016) use groupings amongst labels to search through possible architectures for learning.
Perhaps the most relevant to the current work, Kendall et al.
(2017) uses a joint likelihood formulation to derive task weights based on the intrinsic uncertainty in each task.
M ETHODOLOGY 3. A G RADIENT L OSS F UNCTION B ASED ON R ATE B ALANCING
We begin with the standard multitask loss function with time dependency, L ( t ) = (cid:80) w i ( t ) L i ( t ) , and our goal is to learn the functions w i ( t ) . We argued in Section 1 that w i ( t ) is intimately related to the norm of gradients from each task backpropagated into the network. We thus must motivate a set of desirable gradient magnitudes, and use those desired magnitudes to set the task weights w i ( t ) .
Consider the norms of gradients from task i on some set of weights W within the network, norm ( ∇ W L i ( t )) (speciﬁc choices for W to be discussed later).
Our method of gradient normalization (hereafter referred to as GradNorm) works in two steps: (1) We ﬁrst scale all gradient norms to an equal value as a neutral starting point.
This value is most naturally chosen to be the average gradient norm amongst tasks, E task [ norm ( ∇ W L i ( t ))] , where we use E task [ X ] to denote the average value of a task-dependent quantity X across tasks.
(2) We then modify gradient norms with a rate balancing term that ensures no task trains relatively too slowly.
The gradient norms of task i should grow when task i trains relatively slowly, thereby boosting more sluggish tasks.
Gradient norms thus should be an increasing function of the relative inverse training rate for each task.
To quantify training rates, we choose the loss ratio of task i at training step t , L (cid:48)
i ( t ) : = L i ( t ) /L i (0) , as a measure of task
i ’s inverse training rate; smaller values of L (cid:48) i ( t ) would mean that task i has trained more. If L (cid:48) i ( t ) denotes the inverse training rate of task i , then the relative inverse training rate is just L (cid:48) i ( t ) /E task
[ L (cid:48) i ( t )] . Using this simple loss ratio metric is valid for both regression squared loss and classiﬁcation cross-entropy loss, as we will see in Section 5.2 1 .
Our desired gradient norms are therefore:
[Equation]
where α is an additional hyperparameter. α sets the strength of rate balancing in the multitask problem, and also is a measure of the asymmetry between tasks.
In cases where tasks are very different in their complexity, leading to different learning dynamics, a higher value of α should be used to pull tasks back towards a common training rate more forcefully.
When tasks are more symmetric (e.g. the synthetic examples in Section 4), a lower value of α is appropriate.
Note that α = 0 will always try to pin the norms of backpropped gradients from each task to be equal at W .
Equation 1 sets a desired target for our gradient norms, and we want to update our loss weights w i ( t ) to move gradient norms towards this target. To accomplish this, GradNorm is implemented as a loss function L grad which is just the L1 distance between actual gradient norms and the targets in Equation 1:
[Equation]
The above loss is for one task; the full loss is just the mean of the individual task losses,
L grad ( t ; W ) = (1 /T )
(cid:80) i
L ( i ) grad ( t ; W ) .
L grad is then differentiated with respect to each w i ( t ) , and its gradients are applied via standard update rules to update these weights (see Figure 1 for a schematic view).
In principle, it is also possible to update all network weights (not just w i ( t ) ) based on gradient of L grad , but in practice this adds undue complexity to the problem and often degrades performance.
We can choose W , the weights upon which we rate balance gradient norms, to be any subset of weights within layers of our network. In practice, in order to save on compute overhead, we choose W to be the weights in the last layer which is shared amongst all three tasks.
This simpliﬁcation greatly shortens the number of layers L grad must be backpropagated through, and with this choice of W in our experiments GradNorm only adds ∼ 5% of additional compute time.
After every update step, we also renormalize the weights w i ( t ) so that (cid:80) i
w i ( t ) =
T in order to decouple gradient normalization from the global learning rate.
A S IMPLE T OY E XAMPLE
To illustrate GradNorm on a simple system, we consider T regression tasks onto the functions
[Equation]
where tanh acts element-wise. We use squared loss to train each task.
The matrices B and (cid:15) i have elements generated IID from N (0 , 10) and N (0 , 3 . 5) , respectively. Our task is thus to perform regression on multiple tasks with shared information B along with information speciﬁc to each task, (cid:15) i .
The σ i are ﬁxed scalars which set the variance of the outputs f i .
Higher values of σ i induce higher values of squared loss for that task.
These tasks are harder to learn due to the higher variances in their response values, but they also backpropagate larger gradients.
Classically, such a scenario can lead to suboptimal training dynamics as the higher σ i tasks tend to dominate the training.
All toy problem runs use a 4-layer fully-connected ReLU-activated network with 100 neurons per layer as a common trunk.
A ﬁnal afﬁne transformation per task gives T ﬁnal predictions.
Inputs are in R 250 , and outputs lie in R 100 .
To ensure consistency, we only compare models initialized to the same random values and fed data generated from a ﬁxed random seed.
The asymmetry α is set low to 0.12 for these experiments, as the output functions f i are all of the same form.
In these toy problems, we measure the task-normalized test-time loss, which is the sum of the test loss ratios for each task, (cid:80) i L (cid:48)
i ( t ) .
A simple sum of losses is wholly inadequate to judge the overall performance of a multitask network, as it biases itself towards tasks with higher loss scales, and there exists no general metric by which to judge multitask performance in any setting. Luckily, our toy problem was designed with tasks which are statistically identical except for their loss scales σ i . For this simple example, there is therefore a clear measure of overall network performance, which is the sum of losses with each loss normalized to its σ i - precisely the sum of loss ratios.
In the case of T = 2 , we choose the values ( σ 0 , σ 1 ) =
Classically, task 1 can suppress task 0’s inﬂuence during training due to its higher loss scale.
As shown in the top panels of Figure 2, gradient normalization remedies the issue by increasing w 0 ( t ) to counteract the larger gradients coming from T 1 , and the improved task balance results in better test-time performance.
The possible beneﬁts of gradient normalization become even clearer when the number of tasks increases.
For T = 10 , we sample the σ i from a normal distribution and plot the results in the bottom row of Figure 2.
GradNorm signiﬁcantly improves test time performance over naively weighting each task the same.
Like T = 2 , for T = 10 the w i ( t ) grow larger for smaller σ i tasks; GradNorm is giving tasks with smaller loss scales more breathing room.
For both T = 2 and T = 10 , GradNorm is more stable and outperforms the uncertainty weighting proposed by Kendall et al.
Uncertainty weighting, which enforces that w i ( t ) ∼ 1 /L i ( t ) , tends to grow weights too large and too quickly as the loss for each task drops. Although such networks train quickly at the onset, the training soon crashes as the global learning rate grows too large. This issue is exacerbated as uncertainty weighting allows w i ( t ) to change unconstrained (compared to GradNorm which ensures (cid:80) w i ( t ) =
T always), which pushes global learning rate up even further.
Overall, the traces for each w i ( t ) during a single GradNorm run seem fairly stable and convergent.
In fact, in Section 5.3 we will see how the time-averaged weights E t
[ w i ( t )] lie close to the optimal static weights, suggesting GradNorm can greatly simplify the tedious grid search procedure.
A PPLICATION TO A L ARGE R EAL -W ORLD D ATASET
We primarily use NYUv2 as our dataset of choice. The standard NYUv2 dataset carries depth, surface normals, and semantic segmentation labels (which we cluster into 13 distinct classes). NYUv is quite small as a dataset, with a training split of ∼ 800 examples, but contains both regression and classiﬁcation labels, making it a good choice to test the robustness of GradNorm. To show GradNorm in action on a more large-scale multitask dataset, we also expand NYUv2 to 40,000 images complete with pixel-wise depth, surface normals, and room keypoint labels.
Keypoint labels are obtained through professional human labeling services, while surface normals are generated from camera parameters and the depth maps through standard methods.
Following Lee et al.
(2017), the state-of-the-art in room layout prediction, all inputs are downsampled to 320 x 320 pixels and outputs to 80 x 80 pixels.
These resolutions also speed up training without compromising complexity in the inputs or labels.
5. M ODEL AND I NDIVIDUAL T ASK L OSSES
We try two different models: (1) a SegNet (Badrinarayanan et al.
; Lee et al.
network with a symmetric VGG16 (Simonyan & Zisserman (2014))
encoder/decoder, and (2) an FCN (Long et al.
network with a modiﬁed ResNet-50 (He et al.
encoder and shallow ResNet decoder.
The VGG SegNet reuses maxpool indices to perform upsampling, while the ResNet FCN learns all upsampling ﬁlters.
The ResNet architecture is further thinned (both in its ﬁlters and activations) to contrast with the heavier, more complex VGG SegNet: stride-2 layers are moved earlier and all 2048-ﬁlter layers are replaced by 1024-ﬁlter layers.
Ultimately, the VGG SegNet has 29M parameters versus 15M for the thin ResNet.
Although we will focus on the VGG SegNet in our more in-depth analysis, by designing and testing on two extremely different network topologies we will further demonstrate that GradNorm is very robust to the choice of base model.
We use standard pixel-wise loss functions for each task: cross entropy for segmentation, squared loss for depth, and cosine similarity for normals.
As in Lee et al.
(2017), for room layout we generate
Gaussian heatmaps for each of 48 room keypoint types and predict these heatmaps with a pixel-wise squared loss.
Note that all regression tasks are quadratic losses (our surface normal prediction uses [Table] a cosine loss which is quadratic to leading order), allowing us to use the loss ratio L (cid:48)
i ( t ) of each task as a direct proxy for each task’s inverse training rate.
5. N ETWORK P ERFORMANCE
In Table 1 we display the performance of GradNorm on the NYUv2 dataset (with input/output resolutions as described in Section 5).
Speciﬁc training schemes for all NYUv2 models are detailed in
Appendix A. We see that GradNorm improves the performance of all three tasks with respect to the equal-weights baseline (where w i ( t ) = 1 for all t , i ), and that GradNorm either surpasses or matches (within statistical noise) the best performance of single networks for each task.
The GradNorm Converged Weights network is derived by calculating the GradNorm time-averaged weights E t [ w i ( t )] for each task (e.g. by averaging curves like those found in Appendix B), and retraining a network with weights ﬁxed to those values.
GradNorm thus can also be used to extract good values for static weights.
We pursue this idea further in Section 5.3 and show that these weights lie very close to the optimal weights extracted from exhaustive grid search.
[Table]
To show how GradNorm can perform in the presence of a much larger dataset, we also perform extensive experiments on the expanded NYUv2 dataset, which carries a factor of 50x more data.
The results are shown in Table 2.
As with the standard NYUv2 runs, GradNorm networks outperform other multitask methods, and either matches (within noise) or surpasses the performance of singletask networks
[Table] Figure 3 shows test and training loss curves for GradNorm ( α = 1 . 5 ) and baselines on the expanded
NYUv2 dataset for our VGG SegNet models. GradNorm improves test-time depth error by ∼ 5% , despite ending with much higher training loss.
GradNorm achieves this by aggressively rate balancing the network (enforced by a high asymmetry α = 1 . 5 ), and ultimately suppresses the depth weight w depth ( t ) to lower than 0.10 (see Appendix B for more details).
The same trend exists for keypoint regression, and is a clear signal of network regularization.
In contrast, the uncertainty weighting technique (Kendall et al.
(2017)) causes both test and training error to move in lockstep, and thus is not a good regularizer.
Only results for the VGG SegNet are shown here, but the Thin ResNet FCN produces consistent results.
GradNorm versus an equal weights baseline and uncertainty weighting (Kendall et al.
. 5. G RADIENT N ORMALIZATION F INDS O PTIMAL G RID -S EARCH W EIGHTS IN O NE P ASS
For our VGG SegNet, we train 100 networks from scratch with random task weights on expanded
NYUv2.
Weights are sampled from a uniform distribution and renormalized to sum to T = 3 .
For computational efﬁciency, we only train for 15000 iterations out of the normal 80000, and then compare the performance of that network to our GradNorm α = 1 .
5 VGG SegNet network at the same 15000 steps.
The results are shown in Figure 4.
Even after 100 networks trained, grid search still falls short of our GradNorm network.
But even more remarkably, there is a strong, negative correlation between network performance and task weight distance to our time-averaged GradNorm weights.
At an L 2 distance of ∼ 3 , grid search networks on average have almost double the errors per task compared to our GradNorm network.
GradNorm has effectively allowed us to “cheat” and immediately ﬁnd the optimal grid search weights without actually performing grid search, simplifying a process that is usually notoriously laborious.
5. Q UALITATIVE R ESULTS
Figure 5 shows visualizations of the VGG SegNet outputs on test set images along with the ground truth, for both the expanded and downsampled NYUv2 datasets. Ground truth labels are juxtaposed with outputs from the equal weights network, 3 single networks, and our best GradNorm network.
The qualitative improvements are incremental, but we ﬁnd the GradNorm network tends to output smoother, more detailed pixel map predictions when compared to the other two baselines.  C ONCLUSIONS
Gradient normalization acts as a good model regularizer and leads to superb performance in mul- titask networks by operating directly on the gradients in the network.
GradNorm is driven by the attractively simple heuristic of rate balancing, and can accommodate problems of varying complex- ities within the same uniﬁed model using a single hyperparameter representing task asymmetry. A GradNorm network can also be used to quickly extract optimal ﬁxed task weights, removing the need for exhaustive grid search methods that become exponentially more expensive with the number of tasks. We hope that our work has not only introduced a new methodology for quickly balancing multitask networks, but also has shown how direct gradient manipulation can be a powerful way to reason about task relationships within a multitask framework.
Appendices A G ENERAL T RAINING C HARACTERISTICS
All runs are trained at a batch size of 24 across 4 Titan X GTX 12GB GPUs and run at 30fps on a single GPU at inference.
NYUv2 runs begin with a learning rate of 2e-5.
Expanded NYUv2 runs last 80000 steps with a learning rate decay of 0.2 every 25000 steps.
Downsampled NYUv2 runs last 20000 steps with a learning rate decay of 0.2 every 6000 steps.
Updating w i ( t ) is performed at a learning rate of 0.025 for both GradNorm and the uncertainty weighting (Kendall et al.
(2017)) baseline.
All optimizers are Adam, although we ﬁnd that GradNorm is insensitive to the optimizer chosen.
We implement GradNorm using TensorFlow v1.2.1.
B E FFECTS OF TUNING THE ASYMMETRY α
The only hyperparameter in our technique is the asymmetry α . The optimal value of α for NYUv lies near α = 1 . 5 , while in the highly symmetric toy example in Section 4 we used α = 0 .
This observation reinforces why we call α an asymmetry parameter.
Tuning α leads to performance gains, but we found that for NYUv2, almost any value of 0 < α <will improve network performance over an equal weights baseline.
Figure 6 shows that higher values of α tend to push the weights w i ( t ) further apart, which more aggressively reduces the inﬂuence of tasks which overﬁt or learn too quickly (in our case, depth). Remarkably, at α = 1 . 75 (not shown) w depth ( t ) is suppressed to below 0.02 at no detriment to network performance on the depth task.