多任务学习优化Gradient Normalization

多任务学习优化

1、多任务学习模型

对于多任务的loss，最简单的方式是直接将这两个任务的loss直接相加，得到整体的loss，那么loss函数为：

$$L=\sum L_i\text{\hspace{1em}(1.1)}$$

在这里，模型的整体 loss 是所有子任务的 loss 之和。

这种 loss 计算方式的不合理之处是显而易见的，不同任务 loss 的量级很有可能不一样， loss 直接相加的方式有可能会导致多任务的学习被某个任务所主导或学偏。当模型倾向于去拟合某个任务时，其他任务的效果往往可能受到负面影响，效果会相对变差。

如果对loss函数进行调整，对每一个子任务配置不同的权重。那么loss函数变成如下所示：

$$L=\sum w_i\times L_i\text{\hspace{1em}(1.2)}$$

可以看到相比较于各个子任务的loss直接想加，该方式会对每个任务的loss进行加权。调整了每一个任务在整个优化目标中占比。但是固定的权重在训练的过程中不会发生改变。

这种对子任务loss加权方式是对loss直接想加的进化。但是不同任务的训练收敛过程是不一样的。比如 A 任务已经收敛，但是 B 任务还没开始收敛。这个时候继续训练下去直至B收敛的话A任务就会过拟合，如果此时不训练了A就不能收敛。

因此更好的训练方式是不同的时间、不同的阶段、不同的学习难易程度、不同的学习效果的情况下不同的子任务的权重应该是不同的。

$$L=\sum w_i(t)\times L_i\text{\hspace{1em}(1.3)}$$

2、Gradient Normalization—梯度标准化

《Gradnorm: Gradient normalization for adaptive loss balancing in deep multitask networks》，ICML 2018，Cites：17

论文地址

2.1、主要思想

• 希望不同的任务的loss量级是接近的。
• 希望不同的任务以相似的速度学习。

2.2、实现方式

• 文章定义了两种类型的loss。Label Loss 和Gradient Loss。这两种loss独立优化，不进行运算。
• Label Loss即多任务学习中，每个任务的真实的数据标签与网络预测标签计算的loss。Label Loss由学习任务的性质决定，比如分类的loss或者回归的loss。Label Loss通过对不同任务的Loss进行加权求和来实现：$L=\sum w_i(t)\times L_i$​​。Label Loss 是关于网络参数$W$的函数。
• Gradient Loss，用来衡量每个任务的loss的权重 $w_i(t)$ 的好坏，Gradient Loss是关于权重$w_i(t)$的函数。
• 每个任务的权重 $w_i(t)$​ 是一个变量（注意这里 $w$ 与网络参数 $W$ 是不同的），$w$ 也通过梯度下降进行更新，$t$ 表示当前处于网络训练的第 $t$ 步。

2.3、实现公式

下面的几个公式是Gradient Normalization的核心，目的是将Gradient Loss表示为一个关于loss权重的函数

1. 首先定义了一些变量来衡量任务的loss的量级，其中公式2.1用来表示某个任务loss的量级：
   $$G_W^{(i)}(t)=\Vert {\nabla }_W{w}_i(t)L_i(t){\Vert }_2\text{\hspace{1em}(2.1)}$$
   $${\bar{G}}_W(t)=E_{task}[G_W^{(i)}(t)]\text{\hspace{1em}(2.2)}$$

• $W$是真个多任务学习网络参数的子集。文章中作者选择了网络share bottom部分最后一层的参数。
• $G_W^{(i)}(t)$ 是任务 $i$ 梯度标准化的值，是任务 $i$ 的权重$w_i(t)$与第$i$个任务的loss损失 $L_i(t)$ 的乘积对参数$W$求梯度的$L2$范数，$G_W^{(i)}(t)$ 可以衡量某个任务loss的量级，在这里我们认为$G_W^{(i)}(t)$越大该任务的loss的量级就越大。
• ${\bar{G}}_W(t)$ 是全局梯度标准化的值（即所有任务梯度标准化值的期望值），通过所有$G_W^{(i)}(t)$求均值实现。

2. 其次定义了一些变量来衡量任务的学习速度，其中：

$${\tilde{L}}_i(t)=L_i(t)/L_i(0)\text{\hspace{1em}(2.3)}$$

$$r_i(t)={\tilde{L}}_i(t)/E_{task}[{\tilde{L}}_i(t)]\text{\hspace{1em}(2.4)}$$

• $L_i(0)$与$L_i(t)$分别代表子任务 $i$ 的第 $0$ 步和第$t$步的loss； ${\tilde{L}}_i(t)$ 在一定程度上衡量了任务$i$的反向的训练速度，${\tilde{L}}_i(t)$ 越大，表明loss减小的很慢，几乎没有减小。便表示网络训练越慢；
• $E_{task}[{\tilde{L}}_i(t)]$表示了各个任务反向训练速度的期望。
• $r_i(t)$ 是任务的相对反向训练速度，$r_i(t)$越大，表示任务 $i$ 在所有任务中训练越慢。

3. 最终，Gradient Loss（GL）表示为：
   $$L_{grad}(t;w_i(t))=\sum _i\mid G_W^{(i)}(t)-{\bar{G}}_W(t)\times [r_i(t){]}^{\alpha }{\mid }_1\text{\hspace{1em}(2.5)}$$

• 公式2.5中 $\alpha$ 是超参。
• $G_W^{(i)}(t)$是公式2.1的结果。代表的是任务 $i$ 在 $t$ 时刻或者是说第 $t$ 步loss的量级。该值越大该任务该时刻的loss量级就越大。
• ${\bar{G}}_W(t)$是公式2.2的结果。代表所有任务量级的期望（平均值）。
• $r_i(t)$是公式2.4的结果。代表的是任务 $i$ 在所有任务中训练的速度。其中子任务训练越快的话，$r_i(t)$越小。$r_i(t)$越小的话式子2.5会越来越大。这样的话速度越快loss越大。
• 从公式中可以看出，当某个人物的loss过大或者过小的时候$G_W^{(i)}(t)-{\bar{G}}_W(t)$会变大。导致Gradient Loss（GL）会增加。优化GL的过程就是促进模型选出合适的子任务的$w_i$。使得各个子任务中间的量级和速度保持大概差不多的地步。保证子任务更新梯度的同步性。

4. 最后根据Gradient Loss 对子任务中$w_i(t)$ 进行更新

$$w_i(t+1)=w_i(t)-\lambda \times Grandient(GT,w_i(t))\text{\hspace{1em}(2.6)}$$

2.4、训练流程

• 算法流程如下

1. 初始化各个子任务loss的权重为1，初始化网络参数$W$，配置 $\alpha$ 的值，选择参数$W$，通常选择某个子集用来评估该任务的loss量级。
2. 通过加权求和的方式计算网络的loss。
3. 计算梯度标准化的值$G_W^{(i)}(t)$，相对反向训练速度$r_i(t)$；
4. 计算全局梯度标准化的值${\bar{G}}_W(t)$和 ;
5. 计算Gradient Loss；
6. 计算Gradient Loss对$w_i(t)$求导的值；
7. 通过第6步的结果更新$w_i(t)$；
8. 通过第2步的结果更新整个网络的参数；
9. 把第7步的结果renormalize一下，使得$w_i(t)$的和等于任务数量

2.4、总结

我们都知道神经网络中损失函数的优化过程就是对神经网络中的网络参数$W$更新的过程。label loss就是用来更新网络参数的损失，目的是让该模型推理出来的结果和真实的标签差距越来越小。

而Gradient Loss是用来更新多任务学习中各个子任务的 $w_i$ 的值而存在的。目的是让该模型在训练的过程中各个子任务训练的量级和速度差不多。从而让各个子任务同步的学习，都收敛，然后通过各个子任务的平均间接的提升模型的预测能力。