Source-linked AI summary
Just Pick a Sign: Optimizing Deep Multitask Models with Gradient Sign Dropout
Zhao Chen, Jiquan Ngiam, Yanping Huang, Thang Luong, Henrik Kretzschmar, Yuning Chai, Dragomir Anguelov
TL;DR
Multiloss training can suffer when constituent gradients conflict, especially because summed-loss critical points may be poor for individual tasks. GradDrop probabilistically masks gradient signs according to their consistency, and the paper reports more stable convergence and efficacy across multitask, transfer-learning, and complex single-task settings.
Problem
Summing multiple gradient signals can create conflicting updates and poor minima for constituent losses in multitask and other multiloss models.
Method
GradDrop is a modular layer that probabilistically masks opposing gradient signs at an activation layer according to sign purity.
Results
GradDrop yields more stable convergence points than naïve gradient descent in theory and simulation and demonstrates efficacy across multitask, transfer-learning, and complex single-task models.
Takeaways & Limitations
GradDrop encourages joint minima while preserving stochasticity and can be used across network architectures without inference-time compute.
Takeaways & Limitations
The method's broader-impact discussion notes that multitask training can reinforce biases when improperly trained, although GradDrop may help combat related overfitting.
Abstract
from arXiv · showhide
The vast majority of deep models use multiple gradient signals, typically corresponding to a sum of multiple loss terms, to update a shared set of trainable weights. However, these multiple updates can impede optimal training by pulling the model in conflicting directions. We present Gradient Sign Dropout (GradDrop), a probabilistic masking procedure which samples gradients at an activation layer based on their level of consistency. GradDrop is implemented as a simple deep layer that can be used in any deep net and synergizes with other gradient balancing approaches. We show that GradDrop outperforms the state-of-the-art multiloss methods within traditional multitask and transfer learning settings, and we discuss how GradDrop reveals links between optimal multiloss training and gradient stochasticity.
1 Introduction
GradDrop addresses conflicting gradient signals in multiloss networks by masking gradients according to sign consistency, aiming for joint minima and stable convergence. The paper introduces a modular layer applicable across network architectures and learning settings.
- Motivation: Multiple losses can create conflicting gradients, making naïvely summing their signals unlikely to produce the best solution.The issue also affects models not traditionally labeled multitask, including multiclass classifiers and object detectors.
- Motivation: Following conflicting gradients blindly can produce tug-of-wars and critical points where some constituent gradients remain large.Such points can correspond to poor performance on individual tasks.
- Method: GradDrop selects one gradient sign according to the distribution of gradient values and masks gradients with the opposite sign.The procedure requires updates to be pure in sign at each update position.
- Method: GradDrop adds stochasticity where gradients conflict more, continuing until the model finds a joint minimum for all losses.The paper connects this behavior to the role of gradient noise in escaping narrow, low-quality minima.
- Contributions: GradDrop is a modular layer usable in any network with multiple gradient signals and adds no inference-time computation.The contribution is presented for multitask learning, transfer learning, and complex single-task models such as 3D object detectors.
- Contributions: GradDrop is reported to yield more stable convergence points than naïve gradient descent in theory and simulation.The paper also demonstrates its efficacy across multitask, transfer-learning, and 3D object-detection settings.
2 Related Work
Related work frames gradient optimization, multitask learning, and applications with multiple losses as connected challenges in deep learning. Existing approaches span optimization methods, architectures, and gradient balancing.
- Gradient optimization: Gradient optimization is central to deep learning, with research targeting faster, more stable training and links between stochasticity and robust minima.The cited relationship motivates using noisy gradients to repel poor local minima.
- Multitask learning: Multitask learning combines multiple loss surfaces and supports generalization and information transfer between correlated tasks.Prior research also develops architectures and loss-related methods for this setting.
- Applications: Object detection, instance segmentation, sequence-to-sequence language models, domain adaptation, and transfer learning can involve multiple interacting gradient signals.These applications extend multiloss optimization beyond traditionally defined multitask models.
3 Gradient Dropout
GradDrop computes sign purity across gradient signals and stochastically masks opposing signs, with extensions for batch-separated gradients. Its theory links the resulting updates to joint minima, stable optimization, and controllable stochasticity.
- Basic concepts: GradDrop is inserted as a layer, usually before the prediction head, to modify backpropagation while avoiding inference-time computation.The layer is typically placed at the final shared layer to reduce compute overhead and benefit backpropagation.
- Basic concepts: The sign-purity measure P lies in [0, 1] and reflects the distribution of positive gradients at each scalar.P equals 0 when all gradients are negative and 1 when all are positive.
- Basic concepts: Each gradient receives a stochastic mask that keeps either positive or negative components according to f(P) and uniform random variables U.The masked gradients are combined into the final gradient, with f commonly set to the identity.
- Batch-separated gradients: For batch-separated gradients, GradDrop uses input-weighted gradient contributions so signals can be meaningfully summed across the batch dimension.This avoids blindly combining gradients whose information depends on different inputs.
- Full algorithm: The full algorithm computes P at every gradient location and constructs a mask for each of T task-specific gradient signals.Optional leak parameters allow some original gradient through when losses have different priorities.
- Theoretical properties: GradDrop updates are zero if and only if every component loss has zero gradient at the position, identifying joint minima as stable points.This property distinguishes its stable points from those arising solely from a summed loss.
- Theoretical properties: Expected GradDrop gradient magnitude increases with distance from each component-loss minimum, maintaining sensitivity to every task.The result holds to second order around local minima.
- Statistical properties: For k = 1, GradDrop has the same expected total-loss change as SGD, while k controls the tradeoff between following the overall gradient and adding noise.Smaller k increases the penalty/noise scale for inconsistent gradients.
4 Experiments with GradDrop
Experiments across toy optimization, CelebA multitask learning, CIFAR-100 transfer learning, and Waymo 3D detection show that GradDrop generally outperforms comparison methods and can complement GradNorm.
- A Simple One-Dimensional Example: GradDrop achieves the best performance in the one-dimensional toy optimization, outperforming Random GradDrop and Iterative PCGrad through sign-consistency weighting.PCGrad does not train in the low-dimensional setting, so the comparison uses Iterative PCGrad.
- Multitask Learning on Celeb-A: GradDrop outperforms all other methods on CelebA and improves gradient sign consistency during training.The sign-consistency metric improves appreciably with GradDrop but not with the baseline.
- Multitask Learning on Celeb-A: GradDrop overfits less on CelebA and is the fastest multitask method tested, requiring O(T) computation per tensor position.The authors attribute the speed advantage to avoiding iterative or orthogonal-projection procedures used by competing methods.
- Transfer Learning on CIFAR-100: In CIFAR-100 transfer learning, the best model combines GradDrop with GradNorm, while GradDrop alone also performs well.The performance gap is larger when baseline models overfit later in training.
- 3D Point Cloud Detection on Waymo Open Dataset: GradDrop improves all reported Waymo detection metrics, with especially marked improvements in APH and without the substantial 3D regressions of gradient-norm methods.The experiments use AP and APH, where APH penalizes boxes with 180° orientation errors.
- Synergy with Gradient Normalization and Other Methods: GradDrop can operate alongside GradNorm, often improving results where GradNorm performs poorly, although experiments with GradDrop plus MGDA have limited success.GradDrop+GradNorm recovers much of GradNorm’s 3D AP regression while retaining strong 2D AP performance.
5 Conclusions
The paper presents GradDrop as a sign-purity method for multiloss optimization that encourages joint minima across architectures and multitask settings. It also argues that multitask optimization should account for gradient stochasticity and complex loss surfaces.
- Conclusions: GradDrop converts additive gradient signals into a sign-pure sum signal and encourages networks to seek joint minima.The conclusion describes this behavior as applying across a wide array of architectures and multitask settings.
- Conclusions: The authors argue that standard gradient-descent intuitions may be incomplete for multitask loss surfaces and should incorporate gradient stochasticity.They frame this as a discussion for understanding optimization robustness in increasingly complex models.
6 Broader Impacts
GradDrop is presented as a general modular algorithm for multitask models, with broader effects depending on how it is applied. The authors discuss possible benefits for reducing overfitting and supporting joint solutions, while warning that multitask learning can also reinforce bias when applied incorrectly.
- GradDrop is a general algorithm that can be added modularly to multitask models without application-specific assumptions.
- The authors suggest GradDrop may help multitask models find joint solutions that consistently incorporate available information.
- Multitask learning can reduce bias when it leverages diverse information, but incorrect training can reinforce biases through subpopulation-specific overfitting.
A Appendix
The appendix primarily documents experimental settings and adds intuition and experiments for batch-separated gradients, CIFAR-100 transfer learning, and the Waymo Open Dataset.
- The appendix mainly lists hyperparameters, datasets, and training settings for all experiments.
- Additional appendix material develops intuition for batch-separated gradients and reports experiments on CIFAR-100 transfer learning and the Waymo Open Dataset.
A.1 Addendum on Proposition 3 and Choice of Activation Function
The addendum characterizes how GradDrop’s activation function affects expected loss change and variance, then extends these properties beyond the specific linear activation family.
- For GradDrop, expected loss change is nonpositive and its magnitude increases monotonically with activation steepness k.
- For k = 1, GradDrop’s expected loss change equals the SGD loss-change estimate.
- GradDrop’s loss-change variance decreases monotonically as k increases over k ∈ [0, 1].
- The results extend to activation functions that are odd around (0.5, 0.5), monotonically increasing, and bounded between 0 and 1.
- Among valid activation functions, a steeper function yields no greater expected loss change and no greater variance than a less-steep function under the stated conditions.
A.2 More Intuition Regarding Batch-Separated Gradients
The batch-separated-gradient treatment addresses gradients whose values depend on separate input batches. The appendix explains why input-aware sign correction is useful and notes that the current derivation assumes an element-wise multiplicative virtual layer.
- Batch-separated gradients cannot simply be summed because gradient values depend on their batches’ input values.
- For multiplicative network operations, the relevant comparison is whether gradients push inputs in the same or opposite directions, not just their raw values.
- Premultiplying gradients by input signs corrects potential sign errors without adding the variance associated with input magnitudes.
- The proposed batch-marginalization rule is derived using a virtual element-wise multiplication layer, while other trainable layers remain potential future alternatives.
A.3 A Simple One-Dimensional Example: Addendum
The one-dimensional experiments use simple sinusoidal losses and compare GradDrop with an in-place variant of PCGrad. Training uses standard SGD with a fixed schedule.
- Experimental setup: All runs use an initial learning rate of 0.2, decay it by 0.5 every 1k steps, and train for 10k steps with standard SGD.
- Experimental setup: The experiments use losses of the form sin(ax + b) + 1.0, with the affine offset ensuring nonnegative loss values.The offset is cosmetic, and the sine periods are selected randomly rather than to emphasize a particular behavior.
- Baselines: Iterative PCGrad updates input gradients in-place instead of projecting static copies, adding noise that enables nontrivial behavior in low dimensions.The authors report that it generally performs similarly to ordinary PCGrad in other experimental settings.
A.4 Multitask Learning on Celeb-A: Addendum
The Celeb-A experiments treat each attribute as a separate binary classification task using a shallow convolutional network. GradDrop and baselines are applied after the final convolutional layer under a shared Adam training setup.
- Dataset and tasks: Celeb-A uses the standard 160k/20k split, with each attribute treated as a separate task trained using binary sigmoid classification loss.
- Architecture: The nine-layer convolutional network ends with two 512-unit dense layers and a 40-output predictive head.The listed stack includes convolutional, max-pooling, and dense layers.
- Architecture: GradDrop and competing methods are applied after the final convolutional layer, while batch normalization is used throughout except in the predictive head.
- Optimization: Training uses Adam with β1 = 0.9 and β2 = 0.999, batch size 8, an initial learning rate of 1e-3, and 0.96 annealing every 2400 steps.All baselines share these hyperparameters except MGDA, whose learning rate is reduced by 100x.
A.5 Transfer Learning on CIFAR-100: Addendum
The addendum evaluates transfer learning from ImageNet2012 to CIFAR-100 and GradDrop-based optimization for Waymo 3D detection. It examines leak-parameter choices, gradient-processing baselines, and complementary multitask methods across reported metrics.
- Transfer learning setup: CIFAR-100 uses a standard 40k/10k split, with CIFAR-100 and ImageNet2012 images resized to 32x32 and normalized to values between 0 and 1.The normalization is reported to improve training stability, especially early in training.
- Transfer learning setup: The transfer model is a 100-layer DenseNet-100-BC with k = 12, no data augmentation, and no Dropout because it appeared to degrade performance.Hyperparameters are optimized on the transfer-learning baseline before applying other baselines.
- Transfer learning setup: Separate BatchNorm trainable parameters for ImageNet2012 and CIFAR-100 alleviate domain-shift effects; without this separation, training is very unstable.
- Leak parameters: With ℓsource + ℓtransfer fixed at 1.0, CIFAR-100 performance depends clearly on ℓsource − ℓtransfer: errors remain similar near ℓsource = 1.0, then rise sharply.The same trend appears as a strong linear dependency in loss values, supporting the leak-parameter experiments shown in Table 5 and Figure 5.
- Leak parameters: The authors explain the leak result by allowing the less-important source set to overfit while maximally filtering transfer-set gradients for consistency.
- 3D detection: Waymo detection uses nine losses, adding a directional classifier to eight bounding-box losses; the authors report dramatically improved APH metrics from this addition.The accuracy curves show GradDrop outperforming PCGrad through most of training, while MGDA performs relatively poorly, especially on 3D metrics.
- 3D detection: Combining GradDrop with other multitask methods produces mixed results: the effect atop MGDA is described as murky, with regression in 2D metrics.Table 6 reports broader 2D and 3D comparisons, including MGDA+GradDrop.