如何融合多任务学习 (Multi-Task Learning ) 损失函数loss

多任务学习 (Multi-task learing) 关注的一个问题是 如何优化一个包含多个目标损失函数的模型，通常最直接的方法是通过一个线性函数组合这些损失函数：
$$L_{total}=\sum _i{w}_i{L}_i$$
每个损失函数的权重 $w_i$是一个先验的超参数，而调这些参数代价是很大的，而且模型对这些权重参数的选择特别敏感，权重参数没有选择好，很有可能就会导致整个模型学的不好。所以需要一个更加方便的能够让模型去学习一个最优的权重参数，以下是目前在多任务学习中，权重参数的学习方案总结。

1 Uncertainty Weighting

uncertainty Weighting
Aleatoric uncertainty and epistemic uncertainty

1.1 基础概念

在贝叶斯模型里，通常有两种不确定性：

• 认知不确定性 (Epistemic uncertainty)：认知不确定性主要来自模型本身，若输入的数据是模型训练中并未见过的数据，则就会存在认知不确定性，这种不确定性可以通过增加训练数据来解释。
• 偶然不确定性 (Aleatoric uncertainty) ：偶然不确定性是不能通过数据来解释的，而这种不确定性主要来自其它的因素的干扰，比如数据本身就存在噪声等，而偶然不确定性 (Aleatoric uncertainty)又可以进一步分成两个子类：

1）Data-dependent 或者 异方差不确定性 （Heteroscedastic uncertainty)：主要依赖于具体的输入数据，不同的输入具有不同的不确定性。
2）Task-dependent 或者 同方差不确定性 (Homoscedastic uncertainty)：不依赖输入，所有的输入具有相同的不确定性。

1.2 方法

通过用task uncertainty来衡量不同task任务之间的相对置信度。论文推导了一个基于不确定性的高斯似然最大化的multi-task loss函数。假设$f^W(x)$是模型的输出，其中$W$是权重参数，$x$是输入值，我们可以定义如下高斯概率函数：
$$p(y|f^W(x))=N(f^W(x),{\sigma }^2)$$
其中$\sigma$是一个噪声值，在最大化似然函数中，通过最大化模型的log似然函数，所以可以重新定义如下：
$$logp(y|f^W(x))\propto -\frac{1}{2{\sigma }^2}||y-f^W(x)|{|}^2-log\sigma$$
其中$\sigma$是模型观测的噪声参数，用来衡量输出的噪声大小程度。通过最大化log似然函数优化模型的参数权重$W$和观测的噪声参数$\sigma$。
假设我们的模型输出$y_1$,$y_2$，每一个都属于高斯分布，则：
$$p(y_1,y2|f^W(x))=p(y_1|f^W(x)).p(y_2|f^W(x))$$
$$=N(y_1;f^W(x),{\sigma }_1^2).N(y_2;f^W(x),{\sigma }_2^2)$$
最小化损失函数如下：
$$L(W,{\sigma }_1,{\sigma }_2)=-logp(y_1,y_2|f^W(x))$$
$$\propto \frac{1}{2{\sigma }_1^2}||y_1-f^W(x)|{|}^2+\frac{1}{2{\sigma }_2^2}||y_2-f^W(x)|{|}^2+log{\sigma }_1{\sigma }_2$$
$$=\frac{1}{2{\sigma }_1^2}{L}_1(W)+\frac{1}{2{\sigma }_1^2}{L}_1(W)+log{\sigma }_1+log{\sigma }_2$$

其中${\sigma }_1$和${\sigma }_2$是噪声参数，分别控制着$L_1(W)$和$L_2(W)$损失loss的相对权重，从直观上理解上述公式，若噪声参数${\sigma }_1$越大，则对应的损失函数$L_1(W)$的权重就越小，但由于模型的会尽可能的让损失函数为0，则会使得$\sigma$变得很大，完全忽视了数据的影响，因此对噪声项增加了正则化项$\log \sigma$。

2 GradNorm

2.1 原理

GradNorm，通过梯度归一化方法，动态的调整每个任务的梯度值，防止模型偏向某个梯度较大的任务，导致模型学习不充分。通过使用梯度量化自动均衡算法，提升了multi-task多任务学习模型的收敛速度以及整体性能，而且算法只有一个超参数需要调整，相比其他的方法，减少了所需调整的超参数，并且取得了较好的试验效果。

2.2 方法

在multi-task的loss损失函数中，通常是对多个单任务的loss $L_i$进行线性相加：
$$L=\sum _i{w}_i{L}_i$$
在论文中，权重$w_i$是自适应学习出来的，对于每一个训练step，$w_i=w_i(t)$。首先了解下论文中的几个基本符号代表的含义：

• $W$ 表示的模型网络的参数权重
• $G_W^{(i)}(t)=||{\nabla }_W{w}_i(t)L_i(t)|{|}_2$表示的是单个任务loss $w_i(t)L_i(t)$对权重$W$的梯度的$L_2$范数
• ${\overline{G}}_W(t)=E_{task}[G_W^{(i)}(t)]$表示的所有任务在训练时间t的平均梯度归一化值
• ${\tilde{L}}_i(t)=L_i(t)/L_i(0)$表示的是任务$i$在时间$t$的loss比例。${\tilde{L}}_i(t)$用来测量的是任务$i$的训练速率的倒比例，若${\tilde{L}}_i(t)$的值越小对应的是较快的学习速度
• $r_i(t)={\tilde{L}}_i(t)/E_{task}[{\tilde{L}}_i(t)]$表示的是任务$i$的相对倒学习速率

因此希望每个任务的梯度接近如下形式：
$$G_W^{(i)}(t)\to {\overline{G}}_W(t)\times [r_i(t){]}^{\alpha }$$
则每个权重计算形式如下：
$$w_i(t)=\frac{{\overline{G}}_W(t)\times [r_i(t){]}^{\alpha }}{G_W^{(i)}(t)}$$
其中$\alpha$是一个超参数，在tasks差距很大时，$\alpha$可以设置更多点，有一个更强的训练速率起到平衡的作用。上述公式给出了每个任务$i$的梯度归一化值，通过更新loss权重$w_i(t)$，使得每个任务的梯度值朝着梯度归一化值移动。最终定义的损失函数为$L_1$损失函数，衡量的是每个任务的实际梯度和目标归一化梯度值之间的差距，所以对于每个任务，损失函数定义如下：
$$L_{grad}(t;w_i(t))=\sum _i|G_W^{(i)}(t)-{\overline{G}}_W(t)\times [r_i(t){]}^{\alpha }{|}_1$$
在对$w_i(t)$求梯度过程${\nabla }_{w_i}{L}_{grad}$，让目标梯度归一化值${\overline{G}}_W(t)\times [r_i(t){]}^{\alpha }$作为一个固定的常量值，防止$w_i(t)$不符合逻辑的直接变为0，$L_{grad}$对于$w_i$是可导的，计算的梯度${\nabla }_{w_i}{L}_{grad}$可以用标准的梯度下降法更新权重$w_i$。为了将梯度归一化与全局的学习率解耦，每次我们会重新归一化权重$w_i(t)$，使得${\sum }_i{w}_i(t)=T$。
最终我们的参数$w_i(t)$和$W$更新公式如下：
$$w_i(t)\to w_i(t+1)using{\nabla }_{w_i}{L}_{grad}$$
$$W(t)\to W(t+1)using{\nabla }_{W}L(t)$$
$$renormalize{w}_i(t+1)\to \sum _i{w}_i(t+1)=T$$

3 Multi-Objective Optimisation

3.1 原理

Multi-Objective Optimisation将多任务学习multi-task learning当做multi-objective optimization多个目标优化问题。在MTL问题中，通常是对输入空间$X$以及对应的一系列任务空间${Y^t}_{t\in [T]}$，数据形式为: $\{x_i,y_i^1,y_i^2,...,y_i^T{\}}_{i\in [N]}$，其中$T$是任务的个数，$N$表示的是数据量，$y_i^t$是第$i$个数据在第$t$个任务的label。考虑每个任务的预测函数为: $f^t(x;{\theta }^{sh},{\theta }^t):X\to Y^t$，其中参数${\theta }^{sh}$是所有任务的共享参数，参数${\theta }^t$是每个任务单独的参数。通常一般的MTL任务优化函数如下：
$$\underset{{\theta }^{sh},{\theta }^1,...,{\theta }^T}{\min }\sum _{t=1}^T{c}^t{\hat{L}}^t({\theta }^{sh},{\theta }^t)$$
其中$c^t$表示的是每个任务的权重，${\hat{L}}^t({\theta }^{sh},{\theta }^t)$是任务$t$的损失函数，定义为${\hat{L}}^t({\theta }^{sh},{\theta }^t)=\frac{1}{N}{\sum }_{i}L(f^t(x_i;{\theta }^{sh},{\theta }^t),y_i^t)$

3.2 方法

论文中，作者将MTL优化问题看成multi-objective优化问题，优化的是一组存在竞争的目标：
$$\underset{{\theta }^{sh},{\theta }^1,...,{\theta }^T}{\min }L({\theta }^{sh},{\theta }^1,...,{\theta }^T)$$
$$=\min ({\hat{L}}^1({\theta }^{sh},{\theta }^1),...,{\hat{L}}^T({\theta }^{sh},{\theta }^T){)}^T$$
注意的是这里和MTL不同的是，最后的形式是一个vector，而不是一个scalar。而求解的目标是帕累托最优( Pareto optimality)。MLT的帕累托最优定义如下：
对于不同的$a^t$取值，我们可以得到不同的参数$\theta$，若对于$\theta$在每个任务中${\hat{L}}^t({\theta }^{sh},{\theta }^t)\le {\hat{L}}^t({\overline{\theta }}^{sh},{\overline{\theta }}^t)$，则称$\theta$要比$\overline{\theta }$好 （dominiate)，若参数${\theta }^{\ast }$比任何一个参数在每个任务中的loss都要小，则称${\theta }^{\ast }$为帕累托最优。
而multi-objective 优化问题可以用梯度下降法求解，其中一个方法是Multiple Gradient Descent Algorithm (MGDA)，MGDA利用KKT条件，对于共享参数${\theta }^{sh}$和独立参数${\theta }^t$，KKT条件是：

• 存在$a^1,...,a^t\ge 0$，使得${\sum }_{t=1}^T{a}^t=1$和${\sum }_{t=1}^T{a}^t{\nabla }_{{\theta }^{sh}}{\hat{L}}^t({\theta }^{sh},{\theta }^t)=0$
• 对于所有任务$t$，${\nabla }_{{\theta }^t}{\hat{L}}^t({\theta }^{sh},{\theta }^t)=0$
  根据 MGDA for multiobjective optimization，证明了下面式子的解要么是帕累托最优的必要条件，要么是一个能优化所有任务的好的优化方向：
  $$\underset{a^1,...,a^t}{\min }\{{\left|\begin{array}{c}\left|\begin{array}{c}{\sum }_{t=1}^T{a}^t{\nabla }_{{\theta }^{sh}}{\hat{L}}^t({\theta }^{sh},{\theta }^t)\end{array}\right|\end{array}\right|}_2^2|\sum _{t=1}^T{a}^t=1,a^t\ge 0\forall a\}$$

所以MTL的问题可以变为在任务相关的参数${\theta }^T$上做梯度下降，再用上述的解在共享的参数上做梯度下降。

4 Geometric Loss

4.1 原理

geometric loss 将MTL任务中的多个loss用几何方式组合，相比直接对多个loss求平均的组合方式，能够更好的应对不同的task之间的收敛速度不一样的问题。

4.2 方法

基于几何方式组合形式的最终loss方式组合如下：
$$L_{Total}=\prod _{i=1}^n\sqrt[n]{L_i}$$
其中$n$表示的是总任务个数，$L_i$表示的是第$i$个任务产生的loss。这样的组合方式，比起对所有任务的loss相加求平均，可以缓解任务中loss差距较大而导致某些任务学习受到影响，如果我们提前已知某些任务更重要，那么可以对重要的task的loss加重，形式如下：
$$L_{Total}=\prod _{i=1}^n\sqrt[n]{L_i}\times \prod _{j=1}^m\sqrt[m]{L_j}$$
其中$m<n$，且这$m$个任务的重要程度是按照从高到低依次排序的。在实际应用中，用log函数，将连乘转换成加法形式：
$$L_{Total}=\log (\prod _{i=1}^n\sqrt[n]{L_i}\times \prod _{j=1}^m\sqrt[m]{L_j})$$
$$=\frac{1}{n}\sum _i^n\log L_i+\frac{1}{m}\sum _i^m\log L_i$$

5 HydaLearn

5.1 原理

HydaLearn （Highly Dynamic Learning）主要针对MTL过程中的两个存在的问题：

• 在学习过程中，辅助task有可能逐渐漂移，降低main task的效果
• 对于基于mini-batch的这样梯度迭代，最优的task weights应该取决于每次的mini-batch的样本组成

所以针对上述两个问题，论文提出了HydaLearn，将main task的收益与每个任务的梯度关联起来，在每次的mini-batch梯度更新上，能够自适应的调整每个task的weight。

5.2 方法

令$T_m$表示main task，$T_a$表示辅助auxiliary task，$L_m$和$L_a$分别表示主任务和辅助任务的loss，最后这两部分的loss组合形式如下：
$$L({\theta }_{s,t},{\theta }_{m,t},{\theta }_{a,t})=w_{m,t}{L}_m({\theta }_{s,t},{\theta }_{m,t})+w_{a,t}{L}_a({\theta }_{s,t},{\theta }_{a,t})$$
其中$w_{m,t}$,$w_{a,t}$分别表示主任务和辅助任务的权重分值，${\theta }_{s,t},{\theta }_{m,t},{\theta }_{a,t}$分别表示模型的共享参数和task-specific层的参数。那么其中$w_{m,t},w_{a,t}$怎么设置呢？计算形式如下：
$$\frac{w_{m,t}}{w_{a,t}}\approx \frac{{\delta }_{m,m,t}}{{\delta }_{m,a,t}}$$
其中${\delta }_{m,m,t},{\delta }_{m,a,t}$分别表示的是main task和auxiliary task在训练step $t$步中获得的增益gain值，而这个增益值，我们可以通过梯度更新过程中，前后的loss差异值计算求得，得到这个比值后，我们就可以得到权重$w_{m,t},w_{a,t}$的比值，最后我们可以设定一个限制，使得$w_{m,t}+w_{a,t}=w$，这样，我们每次根据任务的loss gain的比值来得到权重的比值，然后分别得到主任务和辅助任务的权重，再得到组合后的loss。

6 Coefficient of variations Weighting (CoV-Weighting)

6.1 原理

coefficient of variations 动态计算每个task的权重，在训练过程中，通过loss的均值和标准差之间的变化情况来计算每个task的权重。认为当一个task的损失loss的方差趋于零时，该task的优化目标就可以了。我们来看下标准的MTL的loss组合方式如下：
$$L_{total}=\sum _i{a}_i{L}_i$$

6.2 方法

首先定义coefficient变异系数，也叫相对标准偏差，公式如下：
$$c_L=\frac{{\sigma }_L}{{\mu }_L}$$
其中${\sigma }_L,{\mu }_L$分别表示loss $L$的标准差和均值。而这个值是与loss的大小无关的，我们知道不同的loss在大小上是有差异的。变异系数将loss的大小和loss的权重解耦，因此当一个损失loss复杂多变的时候，小幅度的值也有可能有较大的影响。而一个更大的loss值，若在训练过程中loss已经平稳了，会分配一个更小的权重。
论文中提出了一种基于loss_ratio变化的测量方法，而不是loss本身，变化的ratio计算形式如下：
$$r_t=\frac{L_t}{\mu L_{t-1}}$$
其中$L_t$表示的是在step t步中观测到的loss，$\mu L_{t-1}$表示的是累积到$t-1$步中的loss平均值。最后第$i$个任务在时间$t$步中的权重$a_{i,t}$计算表达式如下：
$$a_{i,t}=\frac{1}{z_t}{c}_{r_{it}}=\frac{1}{z_t}\frac{{\sigma }_{r_{it}}}{{\mu }_{r_{it}}}$$
其中$z_t$是一个归一化因子：$z_t={\sum }_i{c}_{r_{it}}$，这样我们就可以保证所有的任务权重之后${\sum }_i{a}_{it}=1$。

7 Scaled Loss Approximate Weighting

7.1 原理

SLAW: Scaled Loss Approximate Weighting for Efficient Multi-Task Learning 主要通过选择每个任务的权重，使得各任务的梯度范数相等，来平衡训练过程中的不同task产生的影响。

7.2 方法

每个任务loss的加权梯度范数相等定义如下：
$$w_i||g_i||=w_j||g_j||\forall i,j$$
如果我们将$w_i$作为未知变量，$||g_i||$作为已知变量，则对于$n$个task，我们将有n-1个这样的等式，对于$w_i$我们将会有无穷个解，其中一种解决方案如下：
$$w_i=\frac{1}{||g_i||}$$
为了与全局的学习率进行解耦，我们强行将权重$w_i$之和限制到一个固定值：
$$\sum _{i=1}^n{w}_i=n$$
所以我们可以得到每个权重的$w_i$的权重计算如下：
$$w_i=\frac{n}{||g_i||}/\sum _{j=1}^n\frac{1}{||g_j||}$$
从上面的式子可以看到，需要额外的每次计算梯度$g_i$，为了避免这种计算代价，论文给出了一种近似解，其中$s_i$替代$g_i$，而$s_i$具体计算公式如下：
$$a_i\leftarrow \beta a_i+(1-\beta )L_i^2$$
$$b_i\leftarrow \beta b_i+(1-\beta )L_i$$
$$s_i\leftarrow \sqrt{a_i-b_i^2}$$
其中$\beta$是移动平均线的一个参数，其中$s_i$评估的是$L_i$的标准差，根据论文提供的理论：

我们可以得到%=$s_i$是$||g_i||$的近似估计值，所以，最终权重$w_i$的计算公式如下：
$$w_i=\frac{n}{s_i}/\sum _{j=1}^n\frac{1}{s_j}$$

8 总结

如下是每种方案，任务$i$对应的权重$w_i$计算方式：

Method	权重$w_i$计算
Uncertainty Weighting	$1/{\sigma }_i^2+\frac{lo{g}_{{\sigma }_i}}{L_i}$
GradNorm	$\frac{L_i(t)/L_i(0)}{g_i(t)}$ (简写公式）
Multi-objective	${\sum }_i{w}_i{\nabla }_{{\theta }_s}{L}_i=0$
Geometric Loss	$n{L}_i/\log L_i$
HydaLearn	$\frac{w_i}{w_j}=\frac{{\delta }_i}{{\delta }_j}$
CoV-Weighting	${\sigma }_{r_i}/{\mu }_{r_i}$
Scaled Loss Approximate Weighting	$\frac{L_i(t)/L_i(0)}{s_i(t)}$ (简写公式）