GradNorm：多任务学习中的梯度平衡方法

Zhao Chen, Vijay Badrinarayanan, Chen-Yu Lee, Andrew Rabinovich. GradNorm: Gradient Normalization for Adaptive Loss Balancing in Deep Multitask Networks. ICML 2018.

在深度学习领域，多任务学习 让单个神经网络同时学习多个任务，这不仅提高了计算效率，还能通过共享特征增强模型泛化能力。然而，多任务网络训练面临一个根本挑战：如何平衡不同任务的学习过程。传统方法通常需要繁琐的手动调整，或耗时的网格搜索，来确定各任务的损失权重。发表于 ICML 2018 的 GradNorm 算法，通过动态调整梯度大小，自动实现了多任务训练的平衡，显著提升了训练效率与最终性能。

---

---

一、多任务学习的核心挑战：训练不平衡

假设我们用一个网络同时完成深度估计（回归任务）和语义分割（分类任务）。在训练初期，深度估计任务的损失可能下降得更快，导致其产生的梯度“声音过大”，主导了网络的更新方向，从而使语义分割任务“学不进去”。这种 梯度范数的不平衡 是多任务训练不稳定的原因之一。

传统方法的局限：

• 静态权重：手动设置固定的损失权重 $w_i$，无法适应训练动态。
• 网格搜索：通过大量实验寻找最优权重，计算成本随任务数指数增长。
• 不确定性加权（Kendall et al., 2017）：基于任务的不确定性调整权重，但容易导致权重无约束增长，破坏训练稳定性。

二、GradNorm 的核心思想：智能任务权重调节器

GradNorm 将多任务训练比作一场多方会议，每个任务是一个“与会代表”，目标是让所有代表有效发言，和谐、高效地推进会议。GradNorm 就像一个智能音量调节器，通过调节任务权重，动态调节每个代表的“麦克风音量”（即梯度大小），确保没有代表被压制或忽视。

关键指标定义：

梯度范数（代表 $i$ 的当前音量）：
$$G_W^{(i)}(t)=w_i(t)\Vert {\nabla }_W{L}_i(t){\Vert }_2$$

表示任务 $i$ 对共享权重 $W$ 的梯度的范数。其中，$w_i(t)$ 是任务 $i$ 的任务权重，${\nabla }_W{L}_i(t)$ 是任务 $i$ 对共享权重 $W$ 的梯度

相对训练速率（代表 $i$ 的发言效率）：

• 首先计算损失比 ${\tilde{L}}_i(t)=L_i(t)/L_i(0)$，值越小表示任务学得越快。
• 相对逆速率 $r_i(t)={\tilde{L}}_i(t)/\text{所有任务的平均}\tilde{L}(t)$。$r_i$ 越大，代表任务 $i$ 的 loss 降低越慢、学习越缓慢。$r_i>1$：学得比平均慢； $r_i<1$：学得比平均快。

三、GradNorm 算法流程

GradNorm 通过以下步骤动态调整损失权重 $w_i(t)$：

3.1 计算一个理想中的梯度范数 $G_{\text{target}}^{(i)}(t)$

理想情况下，学得慢的任务应获得更大的梯度（提高音量），学得快的任务则应减小梯度（降低音量）。目标梯度范数设定为：
$$G_{\text{target}}^{(i)}(t)={\bar{G}}_W(t)\times [r_i(t){]}^{\alpha }$$
其中：

• ${\bar{G}}_W(t)$ 是所有任务梯度范数的平均值，作为“基准音量”。
• $\alpha$ 是一个超参数，控制平衡力度。
  • $\alpha =0$：任何非零实数的零次方 = 1，即，强制所有任务梯度范数相等。
  • $\alpha >0$：根据训练速率差异，进行差异化调节。任务学的越慢，$r_i$ 越大，梯度范数就越大。

3.2 构造梯度损失并更新权重

我们通过一个额外的损失函数来调节 $w_i(t)$，使其实际梯度范数逼近目标：
$$L_{\text{grad}}(t)=\sum _{i=1}^T\left|G_W^{(i)}(t)-G_{\text{target}}^{(i)}(t)\right|$$
这里的关键操作：固定 target 值 $G_{\text{target}}^{(i)}(t) $ 和 ${\nabla }_W{L}_i(t)$ 值，仅对 $w_i(t)$ 求导。更新规则为：
$$w_i(t)\leftarrow w_i(t)-\eta \frac{\partial L_{\text{grad}}}{\partial w_i}$$
更新后，对权重进行归一化，保持 ${\sum }_i{w}_i(t)=T$，以保持整体稳定性。

3.3 整体训练循环（Algorithm 1）

1. 前向传播：计算各任务损失 $L_i(t)$。
2. 计算梯度范数：得到 $G_W^{(i)}(t)$ 与平均值 ${\bar{G}}_W(t)$。
3. 计算训练速率：得到 $r_i(t)$。
4. 计算目标梯度范数：$G_{\text{target}}^{(i)}(t)$。
5. 更新损失权重：通过最小化 $L_{\text{grad}}$ 更新 $w_i(t)$，并归一化。
6. 更新网络权重：使用新的 $w_i(t)$ 计算总损失 $L=\sum w_i{L}_i$，进行常规反向传播。

四、实验验证与效果

4.1 实验设置

• 数据集：
  • 合成数据：模拟多任务回归，验证算法基础有效性。
  • NYUv2：室内场景数据集，包含深度估计、语义分割、表面法线预测任务。
  • MTFL：人脸关键点检测与属性分类数据集。
• 评估指标：各任务的测试误差（如 RMS 误差、分类错误率），目标是整体性能最优。

4.2 主要结果

1. 性能提升：在 NYUv2 和 MTFL 数据集上，GradNorm 在多数任务上达到或超过单任务网络性能，显著优于等权重和 2017 年提出的不确定性加权方法。
2. 高效调参：仅需调节一个超参数 $\alpha$，即可逼近网格搜索的最优效果，极大降低了调参成本。
3. 权重可视化：学到的动态权重 $w_i(t)$ 稳定收敛，且其时间平均值与网格搜索得到的最优静态权重高度吻合。

4.3 超参数 $\alpha$ 的影响

• $\alpha$ 控制任务间的不对称性。
• 在 NYUv2 上，$\alpha =1.5$ 效果最佳；在高度对称的合成任务中，$\alpha =0.12$ 已足够。
• 实验表明，只要 $\alpha$ 在合理范围内（如 0.5 ~ 2.0），GradNorm 均能带来稳定提升。

五、结论与启示

GradNorm 通过直接操作梯度空间，为多任务学习提供了一种简洁、高效的自适应平衡方法。其核心优势在于：

• 直观有效：从梯度不平衡这一根源问题入手，动态调节各任务影响力。
• 计算高效：避免网格搜索，大幅减少训练开销。
• 泛化性强：在回归、分类、不同网络架构中均表现稳定。

未来，该方法可扩展至类别不平衡学习、序列到序列模型等更多存在梯度冲突的场景。GradNorm 不仅是一个实用的多任务训练工具，也为我们理解与操纵深度学习的优化过程提供了新视角。