GradDrop：让多任务学习不再“梯度拔河”

Chen, Zhao, Jiquan Ngiam, Yanping Huang, Thang Luong, Henrik Kretzschmar, Yuning Chai, and Dragomir Anguelov. Just pick a sign: Optimizing deep multitask models with gradient sign dropout. NeurIPS 2020

在多任务学习中，一个模型需要同时学习完成多个相关任务。这通常通过一个共享的主干网络和多个任务特定的头部来实现，优化目标是最小化所有任务损失的总和。这种设计虽然能促进任务间的知识共享，但也引入了一个根本性挑战：梯度冲突，如下文所介绍。发表于 NeurIPS 2020 的 GradDrop 方法，旨在通过令梯度更新符号统一（均为正 / 均为负）来解决这个问题，加速多任务学习。

---

---

一、核心挑战：梯度冲突与“伪平衡点”

想象一个网络要同时学习两个任务。在更新某个共享参数时，任务 A 的梯度是 +0.8，建议增加该参数，而任务 B 的梯度是 -0.5，建议减少该参数。标准的优化器会简单地将梯度相加，即 +0.3，然后参数被往“增加”的方向推了一点点。但这样两个任务可能都不满意：A 觉得推得不够，B 觉得推反了。这就是梯度冲突，导致优化效率低下，最终可能停在一个对哪个任务都不够好的折中点上。

更糟糕的是，优化可能停滞在一个 “伪平衡点” —— 并非所有任务的损失都达到极小值，只是它们的梯度正负相抵消了。（先前方法 PCGrad 无法解决这个问题。）这是我们希望解决的核心问题。

二、现有方法及其局限

为解决梯度冲突，研究者们提出了多种方法：

• GradNorm：动态调整不同任务的损失权重，平衡各任务的学习速度。
• MGDA：寻找一个能同时减小所有损失的梯度下降方向。
• PCGrad：当两个任务的梯度冲突时，将它们投影到彼此的正交方向上以消除冲突。

这些方法大多属于**“减法”思路**：通过削弱、移除或投影掉冲突的梯度成分，来避免内耗。然而，这可能会抑制任务间有益的竞争，且计算复杂度往往较高，例如需要迭代求解或成对投影。

三、GradDrop：梯度符号丢弃

GradDrop 提出 梯度符号丢弃，其核心思想非常直观：与其让正负梯度在“拔河”中相互抵消，不如在每次更新时，统一所有梯度的方向（符号）。

3.1 关键直觉：符号一致性即共识

GradDrop 的出发点是，当多个梯度试图更新同一个参数时，冲突直接体现在它们的符号上。因此，我们追求每次更新都是“符号纯净”的。

3.2 算法三步走

GradDrop 作为一个可插入任何共享层的模块，在反向传播中工作。

步骤一：计算梯度符号纯度
对于某个参数位置，我们汇集所有任务在该处的梯度值。我们计算一个 梯度符号纯度 分数：
$$\mathcal{P}=\frac{1}{2}\left(1+\frac{{\sum }_i\nabla L_i}{{\sum }_i|\nabla L_i|}\right)$$
这个值在 0 到 1 之间。它量化了“增加”该参数的共识程度。

例子，如何计算这个纯度：两个梯度为 +0.8 和 -0.5。

• 分子和：$+0.8+(-0.5)=+0.3$
• 分母和：$|+0.8|+|-0.5|=1.3$
• $\mathcal{P}=0.5\ast (1+0.3/1.3)\approx 0.615$
• 这意味着约有 61.5% 的“梯度力量”支持增加该参数。

步骤二：基于纯度的随机掩码
我们并不简单地服从多数，即总是选择符号占优的方向，而是引入一个随机抽签机制来增加探索性。具体的，我们生成一个 0 到 1 之间的随机数 $U$，并与 $f(\mathcal{P})$ 比较（$f$ 通常为恒等函数）。

• 如果 $f(\mathcal{P})>U$：则保留所有正梯度，丢弃所有负梯度。
• 如果 $f(\mathcal{P})<U$：则保留所有负梯度，丢弃所有正梯度。
  在我们的例子中，$f(0.615)=0.615$。假设 $U=0.42$，由于 $0.615>0.42$，我们保留正梯度。

步骤三：执行更新
使用掩码过滤后的梯度（在我们的例子中为 +0.8）进行参数更新。下一次更新时，会重新进行抽签。

以上例子中，梯度只有一维。当梯度包含多个维度时，我们对每个维度都进行上述操作，为每个维度计算一个符号纯粹的梯度值。也就是说，如果 task 1 的梯度是 (1, -1)，task 2 的梯度是 (-1, 1)，最后我们更新的梯度，可能是 (1, 1) ，也可能是 (-1, -1)。

3.3 GradDrop 的理论保证

GradDrop 并非简单的启发式方法，它有坚实的理论保证。

命题一：稳定点即联合最优点

GradDrop 产生的梯度更新为零，当且仅当在该参数处，所有任务的梯度都为零。

这意味着优化过程不会停留在梯度相互抵消的“伪平衡点”上，只有真正对所有任务都接近最优的点才会成为稳定点。

命题二：对每个任务敏感

GradDrop 更新的大小，会随着参数偏离任一任务的局部最优点而单调增加*。*

即使当前总损失很小，但只要某个任务表现很差（其梯度很大），GradDrop 仍有很大概率产生一个大的更新信号来优化这个落后任务，防止其被其他任务“淹没”。

命题三：统计效率对齐 SGD

在期望上，GradDrop 带来的损失下降量与标准随机梯度下降（Stochastic Gradient Descent）相同。

其引入的随机性是可控且有益的，帮助模型逃离狭窄的局部最优，寻找更宽广的“联合最小值”盆地。

四、实验结果

论文在多个标准设定下验证了 GradDrop 的有效性。

任务类型	数据集	关键指标	GradDrop 表现
多任务学习	CelebA (40 个属性)	分类错误率 / F1 分数	全面优于基线及 GradNorm、PCGrad等 方法，且训练速度最快。
迁移学习	ImageNet → CIFAR-100	Top-1 错误率 / 测试损失	显著降低错误率和损失，与 GradNorm 结合效果更佳。
复杂单任务	Waymo 3D 检测	3D 平均精度 / 带朝向的精度	在 AP 与 APH 指标上全面领先，尤其在方向敏感性指标上提升显著。

关键观察：

1. 普遍有效：无论在经典多任务、迁移学习还是内部包含多损失的检测任务上，GradDrop 均能带来提升。
2. 计算高效：仅需简单的逐元素计算，速度快于需要迭代或投影的复杂方法。
3. 模块化兼容：可与 GradNorm 等方法协同使用，获得叠加效果。

五、总结

GradDrop 通过一个简单而巧妙的机制：基于符号一致性的概率性梯度选择，优雅地解决了多任务学习中的梯度冲突问题。它强迫优化方向在每次更新时保持一致性，从而高效地引导模型寻找对所有任务都友好的“联合最小值”。其坚实的理论保证、极低的计算开销和即插即用的特性，使其成为多任务优化工具箱中一个非常具有吸引力的新工具。