Just Pick a Sign: Optimizing Deep Multitask Models with Gradient Sign Dropout (2010.06808v1)

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.

• The paper introduces Gradient Sign Dropout (GradDrop), a novel method to optimize deep multitask models by addressing conflicting gradients from different tasks using probabilistic sign consistency.
• GradDrop enforces consistent gradient directions by probabilistically masking gradients based on a sign purity score, theoretically driving the model towards robust joint minima for all tasks.
• Experimental results demonstrate GradDrop's superior performance and efficiency over existing methods on various multitask and transfer learning tasks, highlighting its potential for wide adoption.

Overview of "Just Pick a Sign: Optimizing Deep Multitask Models with Gradient Sign Dropout"

In the paper "Just Pick a Sign: Optimizing Deep Multitask Models with Gradient Sign Dropout," the authors present a method aimed at addressing the challenges of optimizing deep multitask models. The method, termed Gradient Sign Dropout (GradDrop), focuses on mitigating the issue of conflicting gradient directions that arise when multiple gradients, corresponding to various loss components, attempt to update the shared weights in a multitask learning framework.

Core Concepts and Approach

GradDrop is introduced as a solution to the common problem in multitask models where naive summation of gradients corresponding to different tasks can lead to suboptimal updates. This is due to potential conflicts in the gradient directions, leading to a gradient tug-of-war that hinders convergence towards a shared optimal solution. This approach is particularly notable in scenarios characterized by complex multitask loss surfaces, where individual tasks have separate local minima. The fundamental idea of GradDrop is to enforce a consistent gradient direction at each update step by probabilistically masking gradients based on their sign consistency.

The GradDrop method constructs a Gradient Positive Sign Purity score, $\mathcal{P}$, which evaluates the sign consistency across gradients. It then utilizes a probabilistic masking procedure to selectively keep or discard gradients, ensuring that only gradients of a consistent direction (either positive or negative) are considered. This approach inherently introduces stochasticity that nudges the model towards robust minima, thereby improving convergence stability.

Theoretical Contributions

The authors provide several theoretical results to justify the effectiveness of GradDrop. They establish that GradDrop updates drive the system towards joint minima, i.e., convergence points where the gradients for all tasks simultaneously vanish. They also demonstrate that GradDrop is effective in encouraging broader and more robust solutions by introducing controlled stochasticity, which helps escape sharp, poor-quality minima.

Additionally, the paper compares GradDrop's performance to baseline methods like vanilla gradient descent, as well as other multitask methods such as MGDA and PCGrad. The GradDrop algorithm achieves notable improvements in convergence, outperforming traditional approaches in various multitask and transfer learning settings.

Experimental Evaluation

The experimental results presented in the paper include applications in diverse settings, such as Celeb-A attribute prediction, transfer learning from ImageNet to CIFAR-100, and 3D object detection from point clouds on the Waymo Open Dataset. Across these experiments, GradDrop consistently outperforms existing multitask learning techniques, demonstrating its efficacy in both achieving lower error rates and higher predictive accuracy.

Crucially, the experiments illustrate GradDrop's flexibility and efficiency, as it incurs minimal computational overhead compared to other methods while still providing significant performance gains. The approach also proves to be synergistic with existing optimization techniques like GradNorm, highlighting its potential as a modular enhancement for deep learning pipelines.

Implications and Future Directions

The introduction of GradDrop contributes to a more nuanced understanding of gradient dynamics in multitask learning. It highlights the importance of addressing gradient conflicts systematically to enhance model robustness and generalization. This approach could be further developed to incorporate other forms of gradient modification techniques, potentially leading to even more sophisticated methods for optimizing complex networks.

In conclusion, GradDrop provides a significant advancement in the field of deep multitask learning by tackling gradient inconsistencies through a probabilistic dropout mechanism. Its demonstrated success across varying domains suggests that GradDrop could be widely adopted in future applications, especially in environments where tasks are inherently interdependent and feature complex interactions. Future research could explore extending the GradDrop framework to other paradigms, such as meta-learning or reinforcement learning, where gradient interactions are similarly complex.