[Paper Review] Just Pick a Sign: Optimizing Deep Multitask Models with Gradient Sign Dropout

Chen Zhao, Jiquan Ngiam+5 more|arXiv (Cornell University)|Oct 14, 2020

Domain Adaptation and Few-Shot Learning51 references64 citations

TL;DR

GradDrop introduces a gradient masking layer that selects gradient signs based on a consistency score, promoting joint minima across multiple losses and improving multitask and transfer learning performance.

ABSTRACT

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.

Motivation & Objective

• Motivates that naïve summation of multiple gradient signals can hinder multitask training due to conflicting directions.
• Proposes GradDrop to promote joint minima by selectively masking gradient signs.
• Demonstrates GradDrop effectiveness across multitask learning, transfer learning, and complex single-task models.
• Explores theoretical properties and synergy with existing gradient-balancing methods.

Proposed method

• Define Gradient Positive Sign Purity P = 1/2(1 + sum_i ∇L_i / sum_i |∇L_i|).
• Compute mask M_i for each gradient using a monotonic function f and a random uniform U to decide which sign to keep, producing new gradient as sum_i M_i ∇L_i.
• Apply GradDrop as a modular layer before the prediction heads, with optional leak parameters ℓ_i to bias where necessary.
• Extend GradDrop to batch-separated gradients by summing gradients across the batch with a virtual layer to compute P and M_i.
• Provide a full algorithm for BACKWARD pass of the GradDrop layer, including normalization and optional gradient leakage.
• Prove that GradDrop ensures stable points only at joint minima and that gradient magnitude remains sensitive to each loss.

Experimental results

Research questions

• RQ1 Can GradDrop reliably steer optimization toward joint minima across multiple losses in multitask settings?
• RQ2 How does GradDrop compare to existing multitask gradient methods (MGDA, PCGrad, GradNorm) in various tasks and architectures?
• RQ3 Does GradDrop interact beneficially with transfer learning and other gradient-based regularizers?
• RQ4 What are the theoretical guarantees and statistical properties of GradDrop updates?

Key findings

• GradDrop outperforms state-of-the-art multitask methods on CelebA, CIFAR-100 transfer, and Waymo 3D detection in key metrics.
• In CelebA, GradDrop achieves the lowest error rate (8.52%) and highest max F1 (29.57) with comparable speed to baselines.
• GradDrop provides notable gains in transfer learning (CIFAR-100) and 3D detection metrics, and synergizes with GradNorm.
• GradDrop maintains expected total loss movement and increases gradient sensitivity to individual tasks, promoting joint minima.
• GradDrop remains computationally efficient, adding minimal overhead at inference and generally lower training time than some alternatives.