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 direct 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 time 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 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 to 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 by Lu et al.
(2016) uses task groupings 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
Our goal is to rate balance our multitask network, ensuring that tasks train at similar rates.
In our case, we measure training rate by the loss ratio of each task, L (cid:48) i :=
L i /L (0) i , where L (0)
i is the initial loss of task i .
This effectively initializes the gradients per task to all be the same magnitude at the start of training.
We shall see this loss ratio helps balance both classiﬁcation and regression tasks and can be easily adapted to other loss functions in Section 5.2.
Rate balancing is enforced at ﬁlter F when:
[Equation]
where the α hyperparameter can be thought of as 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 the tasks back towards a common training rate more forcefully.
When tasks are fully 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 ﬁlter F .
Our desired proportionality in Equation 1 leads us naturally to the loss function for task i at ﬁlter F :
[Equation]
The expected value E j ( norm ( ∇ F L j )) is a mean of the gradient norms of L j at ﬁlter F across each task j .
The above loss is for one task; the full loss is just the mean of the individual task losses,
L grad =
(1 /T )
(cid:80) i
L ( i ) grad .
In simple terms, norm ( ∇ F L i ) is pulled to the value of (relative training rate of task i ) α × (average gradient among all tasks).
If the relative training rate is the same amongst all tasks (or if α = 0 ), this loss pulls all gradient norms towards the mean gradient norm (Figure 1).
In practice, we only apply GradNorm to one downstream ﬁlter in our network, and L grad only updates the task weights w i ( t ) to simplify the problem; in our tests GradNorm only incurs a ∼ 5% slowdown at training.
Downstream normalization also should backpropagate beneﬁts upstream.
For convolutional ﬁlters, we separately consider the gradient norms of each ﬁlter channel, while for dense layers we consider the matrix norm of the full afﬁne transformation matrix.
Bias terms are left unnormalized.
After every update, we 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
f i ( x ) =
σ i tanh
( A i x ) , (3) where tanh acts element-wise and the matrices A i are generated from a common baseline matrix,
A i = B + (cid:15) i .
The common baseline is crucial as it allows our model to ﬁnd useful shared features.
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 the A i map inputs to 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 .
This metric is appropriate as it gives equal weight to all tasks.
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 tends to grow weights too large and too quickly, so although networks train quickly at the onset the training is less stable in the long run.
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 use an expanded version of the NYUv2 (Nathan Silberman & Fergus (2012)) dataset, with 40,000 images complete with pixel-wise depth, surface normals, and room keypoint labels.
Keypoint labels are obtained through professional labeling, while surface normals are generated from camera parameters and the depth maps through standard methods.
As it is an unofﬁcial dataset, expanded
NYUv2 gives noisier labels than the normal NYUv2, but compensates with a 50-fold increase in training data.
Inputs are downsampled to 320 x 320 pixels and outputs to 80 x 80 pixels.
We also test GradNorm on the standard NYUv2 dataset, which carries depth, surface normals, and semantic segmentation labels (which we cluster into 13 distinct classes).
This dataset is much smaller, with a training split of ∼ 800 examples, and we also downsample this dataset to the same 320 x 320 input with 80 x 80 output.
This downsampled NYUv2 dataset allows us to demonstrate that our methods work with a cross-entropy classiﬁcation (segmentation) task added to the mix.
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 to maximize 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, our VGG model has ∼ 29 million parameters to our thin ResNet’s ∼ 14 million.
By designing two extremely different network topologies, we will 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 our regression losses are quadratic (for normals, 1 −| cos θ | ≈ θ 2 + O ( θ 4 ) ), but we mix in cross-entropy loss to show the robustness of GradNorm across various loss functions.
In Section 5.2 we also discuss how to generalize our methods to other loss functions.
5. N ETWORK P ERFORMANCE
[Table]
In Tables 1 and 2 we display the performance of our networks on the expanded and downsampled
NYUv2 datasets, respectively.
In both cases, GradNorm networks outperform other multitask methods, and either matches (within noise) or surpasses the performance of single-task networks.
[Table]
In Tables 1 and 2, we also ﬁnd that the GradNorm Converged Weights network performs quite well.
This network is derived by calculating the GradNorm time-averaged weights E t [ w i ] for each task (e.g. by averaging the curves in Figure 5), 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 the Section 5.3 and show that these weights lie very close to the optimal weights extracted from exhaustive grid search.
It may be surprising that GradNorm can balance regression and classiﬁcation losses in the same network.
Some theoretical insight for why this might be the case can be found in Kendall et al.
(2017), where the authors derive that (under some mild assumptions) the estimated intrinsic task uncertainty σ i follows the same rule σ i ∼ √ L i for both cross entropy and squared losses.
In the case where we use a loss function that is a transform φ of either cross-entropy or squared error (for example, an L 1 loss has φ ( L ) = √ L with respect to squared error), then the loss ratio needs to be updated to L (cid:48)
=
φ − 1 ( L/L (0) )
so that the learning rates L (cid:48) are all on an appropriate shared scale.
Mixing in more exotic loss functions may require additional statistical reasoning to ﬁnd the right form of L (cid:48) , but cross entropy and squared error (and functional transforms thereof) cover most loss functions used in practice, and so GradNorm can be easily adapted to various scenarios.
Our optimal GradNorm network corresponds to α = 1 . 5 , revealing a signiﬁcant asymmetry in learning dynamics between the different tasks in NYUv2.
GradNorm handles the asymmetry by regularizing the network and reining in the training rates of overﬁtting tasks.
Figure 3 displays how
GradNorm improves test-time depth error on the expanded NYUv2 dataset by ∼ 5% , even though its depth training loss is worse.
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.
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 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. 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 5 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.
5. 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 done 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.
5. Q UALITATIVE R ESULTS
Figure 6 shows visualizations of the VGG SegNet outputs on test set images along with the ground truth, for both the expanded and downsampled NYUv2 datasets.
The improvements are incremental, but the GradNorm network consistently outputs smoother, more detailed pixel map predictions.
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 operates under an attractively simple heuristic of rate balancing, and can accommodate problems of varying com- plexities within the same uniﬁed model using a single hyperparameter representing task asymmetry.
A GradNorm network also can 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.