GradNorm

背景

在MultiTask任务中，不同task的loss尺度是不一样的，可能相差很大，这样会导致共享的权重被大尺度的task主导，导致小尺度loss的task学不好，为了解决这个问题，常用的一个方法是对权重加权：\[ loss(t) = w_A \times loss_A(t) + w_B \times loss_B(t) \]

这样做的缺点就是，不同task的loss大小在训练中是会发生变化的，而上面的权重参数是固定的，不能很好的适应训练过程中loss的变化

因此GradNorm提出了一种新的方法，即把loss的权重w也当作一个参数来学习：\[ \text{loss}(t) = w_A(t) \times \text{loss}_A(t) + w_B(t) \times \text{loss}_B(t) \]

 

方法

相关符号定义：

• \( W \)：实际应用 GradNorm那部分权重，为节省计算成本，\( W \) 通常选为最后一个共享层的权重参数
• \( G_W^{(i)}(t) = \|\nabla_W w_i(t) L_i(t)\|_2 \)：所选权重 \( W \)对应的加权单任务损失 \( w_i(t) L_i(t) \) 的梯度的 \( L_2 \) 范数
• \( \overline{G}_W(t) = \mathbb{E}_{\text{task}}[G_W^{(i)}(t)] \)：训练时刻 \( t \) 所有任务的平均梯度范数
• $\tilde{L}_i(t) = L_i(t)/L_i(0)$：任务 $i$ 在时间 $t$ 时和初始时刻的损失比率，$\tilde{L}_i(t)$ 反比于训练速率
• $r_i(t) = \tilde{L}_i(t)/E_{\text{task}}[\tilde{L}_i(t)]$：任务 $i$ 相对反比训练速率

 

GradNorm的目标是根据task的学习速率动态的调整不同task的loss权重，使不同task的收敛速度保持在相同水平，为了达到这个目标，在正常的loss之外定义了一个 Gradient loss 来优化不同task的权重 $w_i(t)$：\[ L_{\text{grad}}(t; w_i(t)) = \sum_{i} \left\lvert G_{W}^{(i)}(t) - \overline{G}_{W}(t) \times \left[ r_i(t) \right]^{\alpha} \right\rvert_1 \] 

在微分 \( L_{\text{grad}} \) 时：

• \( \bar{G}_w(t) \times [r_i(t)]^\alpha \) 被看作为常数
• 在每次更新 \( w_i(t) \) 前，都会重新标准化 \( w_i(t) \) 使得 \( \sum_i w_i(t) = T \)，这使得 GradNorm 方法对 \( w_i \) 的调整与全局学习率的设置解耦

 

算法流程：

 

 

参考资料

Multi-Task Learning：GradNorm

多目标 | 样本权重：GradNorm和DWA利用学习速度调权