Multi-Task Learning：一招鲜能否吃遍天？

不确定性加权

Kendall A, Gal Y, Cipolla R. Multi-task Learning using Uncertainty to Weigh Losses for Scene Geometry and Semantics. CVPR 2018.

不确定性加权（Uncertainty Weighting）是 Alex Kendall 等人提出的一种动态调整不同任务的权重实现多任务学习的方法。

作者引入了贝叶斯学派的思想，假设每一个子任务的预测结果对真值的后验概率服从于正态分布：

$p\left(\mathbf{y} \mid \mathbf{f}^{\mathbf{W}}(\mathbf{x})\right)=\mathcal{N}\left(\mathbf{f}^{\mathbf{W}}(\mathbf{x}), \sigma^{2}\right)$

其中 $\sigma$ 为观测噪声，将其作为标准差时，其大小代表了模型拟合效果的难易程度。

直观理解就是预测值总是在真值附近波动.

接下来考虑多任务场景，特别地以两个回归任务为例，其联合分布就可以表示为：

$\begin{aligned} p\left(\mathbf{y}_{1}, \mathbf{y}_{2} \mid \mathbf{f}^{\mathbf{W}}(\mathbf{x})\right) &=p\left(\mathbf{y}_{1} \mid \mathbf{f}^{\mathbf{W}}(\mathbf{x})\right) \cdot p\left(\mathbf{y}_{2} \mid \mathbf{f}^{\mathbf{W}}(\mathbf{x})\right) \\&=\mathcal{N}\left(\mathbf{y}_{1} ; \mathbf{f}^{\mathbf{W}}(\mathbf{x}), \sigma_{1}^{2}\right) \cdot \mathcal{N}\left(\mathbf{y}_{2} ; \mathbf{f}^{\mathbf{W}}(\mathbf{x}), \sigma_{2}^{2}\right)\end{aligned}$

我们的目标就是最大化该联合分布的似然估计，从而得出以模型参数 $\bf W$ 和各个子任务的方差 $\sigma^2$ 为分布参数的估计值。这等价于最小化下面的目标函数：

$\begin{aligned} \mathcal{L}(\mathbf W,\sigma_1,\sigma_2)&=-\log p\left(\mathbf{y}_{1}, \mathbf{y}_{2} \mid \mathbf{f}^{\mathbf{W}}(\mathbf{x})\right) \\ &\propto \frac{1}{2 \sigma_{1}^{2}}\left\|\mathbf{y}_{1}-\mathbf{f}^{\mathbf{W}}(\mathbf{x})\right\|^{2}+\frac{1}{2 \sigma_{2}^{2}}\left\|\mathbf{y}_{2}-\mathbf{f}^{\mathbf{W}}(\mathbf{x})\right\|^{2}+\log \sigma_{1} \sigma_{2} \\ &=\frac{1}{2 \sigma_{1}^{2}} \mathcal{L}_{1}(\mathbf{W})+\frac{1}{2 \sigma_{2}^{2}} \mathcal{L}_{2}(\mathbf{W})+\log \sigma_{1} \sigma_{2} \end{aligned}$

对于分类任务，可作如下假设：

$p\left(\mathbf{y} \mid \mathbf{f}^{\mathbf{W}}(\mathbf{x})\right)=\operatorname{Softmax}\left(\frac{1}{\sigma^{2}}\mathbf{f}^{\mathbf{W}}(\mathbf{x})\right)$

从而通过近似可以得到相同的目标函数。回归和分类混合的多任务同样如此。

在实践中，其实就是将噪声方差 $\sigma^2$ 也作为模型的可学习对象，从而使得模型可以自动地调整目标函数中各个损失项的权重，然后在目标函数末尾加上正则项 $\log(\sigma)$ 防止参数 $\sigma$ 过大，权重为 0。

同时具体到代码上，其实是参数化 $s:=\log\sigma^2$ ，然后参与计算时在指数计算回去，即 $\sigma^2=\exp(s)$ . 因为这能够保证参数化时满足 $\sigma^2\gt0$ 的约束，同时也保证了权重 $1/\sigma^2$ 不会为 Nan 。

# 伪代码
log_var_a = torch.zeros((1,), requires_grad=True)
log_var_b = torch.zeros((1,), requires_grad=True)

params = ([p for p in model.parameters()] + [log_var_a] + [log_var_b])
optimizer = optim.Adam(params)

def criterion(y_pred, y_true, log_vars):
  loss = 0
  for i in range(len(y_pred)):
    precision = torch.exp(-log_vars[i])
    diff = (y_pred[i]-y_true[i])**2.
    loss += torch.sum(precision * diff + log_vars[i], -1)
return torch.mean(loss)

值得注意的是，因为损失项的权重 $1/\sigma^2$ 也是模型学习得到的，所以面对那些很难拟合的子任务 $\mathcal{L}(\mathbf W)$ ，模型会本能地/贪心地抑制这一项所占的比例。所以该方法反而会使模型关注那些简单的任务。如果作业时目标是希望学习到应对难任务的模型，那么该方法反而可能会降点 23333。

帕累托最优

Sener O, Koltun V. Multi-task learning as multi-objective optimization[J]. Advances in neural information processing systems, 2018, 31.

处理多任务损失的另一个方法是最优化领域的视角。一个很直观的理解就是给多个损失项乘以或静态或动态的权重，其实就是将多目标优化问题转变为了优化加权后的单目标优化问题。但这种调权重的方法的问题是，如果多个任务间存在竞争关系——一个任务变好，其它任务就会变差，那么就很难找到一组有效的权重了。

帕累托最优（Pareto optimal）是经济学中的一种资源分配的理想状态，而在多目标优化领域，就代表在刚开始优化时，每个任务都能变好，直到优化到某一步时，发现一个任务变好但至少有一个其他任务会变差，那么我们就认为达到了帕累托最优（的这样一个临界状态），优化结束。

问题转换

应用帕累托最优策略，我们需要先将MTL原来的单目标优化恢复成多目标优化：

$\min _{\substack{\boldsymbol{\theta}^{sh} \\ \boldsymbol{\theta}^{1}, \ldots, \boldsymbol{\theta}^{T}}} \mathbf{L}\left(\boldsymbol{\theta}^{s h}, \boldsymbol{\theta}^{1}, \ldots, \boldsymbol{\theta}^{T}\right)=\min _{\substack{\boldsymbol{\theta}^{sh} \\ \boldsymbol{\theta}^{1}, \ldots, \boldsymbol{\theta}^{T}}}\left(\hat{\mathcal{L}}^{1}\left(\boldsymbol{\theta}^{s h}, \boldsymbol{\theta}^{1}\right), \ldots, \hat{\mathcal{L}}^{T}\left(\boldsymbol{\theta}^{s h}, \boldsymbol{\theta}^{T}\right)\right)^{\top}$

根据优化方向的一篇数学论文里提出的 Multiple Gradient Descent Algorithm (MGDA)，我们可以通过求其 KKT 条件来寻找满足帕累托最优的驻点。

对于所有的任务 $\forall t, \nabla_{\boldsymbol{\theta}^{t}} \hat{\mathcal{L}}^{t}\left(\boldsymbol{\theta}^{s h}, \boldsymbol{\theta}^{t}\right)=0$ ;
$∃ \alpha^{1}, \ldots, \alpha^{T} \geq 0 \ s.t. \ \sum_{t=1}^{T}\alpha^t=1$ 使得 $\sum_{t=1}^{T} \alpha^{t} \nabla_{\boldsymbol{\theta}^{s h}} \hat{\mathcal{L}}^{t}\left(\boldsymbol{\theta}^{s h}, \boldsymbol{\theta}^{t}\right)=0$ .

对于第一个条件，因为是对子任务自己的网络参数 $\boldsymbol{\theta}^{t}$ 求梯度，要优化满足这个条件等同于平时正常地使用梯度下降来求解。

对于第二个条件，MGDA给出了证明，如果能最小化下面的目标函数为0，那么就得出了对应的驻点；如果没法为0，该目标函数的优化同样能得出可以改善性能的解。

$\min _{\alpha^{1}, \ldots, \alpha^{T}}\left\{\left\|\sum_{t=1}^{T} \alpha^{t} \nabla_{\boldsymbol{\theta}^{s h}} \hat{\mathcal{L}}^{t}\left(\boldsymbol{\theta}^{s h}, \boldsymbol{\theta}^{t}\right)\right\|_{2}^{2} \bigg|\sum_{t=1}^{T} \alpha^{t}=1, \alpha^{t} \geq 0 \quad \forall t\right\}$

该优化问题，本质上就是在一个凸包里找一个点，使其距离所有的已知点都尽可能最近。这是在优化方向被广泛研究的问题。例如 Frank-Wolfe算法就是求解这类带约束的凸优化问题的一种算法。

参数更新算法

综上所述，我们可以得到一个完整的更新网络参数的算法：

先梯度下降更新任务自己的网络参数 $\boldsymbol{\theta}^{t}$ ，然后使用Frank-Wolfe算法求解得到合适的 $\alpha$ ，通过加权更新共享网络参数 $\boldsymbol{\theta}^{sh}$ .

链式改进

上述算法虽然可以求解问题，但效率太低。直观表现就是每个任务都要计算针对共享参数的梯度 $\nabla_{\boldsymbol{\theta}^{t}} \hat{\mathcal{L}}^{t}\left(\boldsymbol{\theta}^{s h}, \boldsymbol{\theta}^{t}\right)$ ，从而反向传播的计算量就和任务的数量成正比；当共享参数量很大时，非常低效。

作者专门针对 encoder-decoder的网络，提出了以encoder的输出 $\bf Z$ 为界的优化方案，将反向传播的链条分成两部分：

$\begin{aligned} f^{t}\left(\mathbf{x} ; \boldsymbol{\theta}^{s h}, \boldsymbol{\theta}^{t}\right)&=\left(f^{t}\left(\cdot ; \boldsymbol{\theta}^{t}\right) \circ g\left(\cdot ; \boldsymbol{\theta}^{s h}\right)\right)(\mathbf{x})=f^{t}\left(g\left(\mathbf{x} ; \boldsymbol{\theta}^{s h}\right) ; \boldsymbol{\theta}^{t}\right)\\ \mathbf{z}_i&=g\left(\mathbf{x} ; \boldsymbol{\theta}^{s h}\right) \end{aligned}$

将压缩层的表征向量记为 $\bf Z$ ，从而可以得出原优化目标的理论上界，即：

$\left\|\sum_{t=1}^{T} \alpha^{t} \nabla_{\boldsymbol{\theta}^{s h}} \hat{\mathcal{L}}^{t}\left(\boldsymbol{\theta}^{s h}, \boldsymbol{\theta}^{t}\right)\right\|_{2}^{2} \leq\left\|\frac{\partial \mathbf{Z}}{\partial \boldsymbol{\theta}^{s h}}\right\|_{2}^{2}\left\|\sum_{t=1}^{T} \alpha^{t} \nabla_{\mathbf{Z}} \hat{\mathcal{L}}^{t}\left(\boldsymbol{\theta}^{s h}, \boldsymbol{\theta}^{t}\right)\right\|_{2}^{2}$

这样做的好处是，反向传播时每个任务的只需要计算对 $\bf Z$ 的梯度即可， $\bf Z$ 的量级相比 $\boldsymbol{\theta}^{sh}$ 小很多，因此计算量会大大减小。此外，上界存在与 $\alpha$ 无关系数 $\left\|\frac{\partial \mathbf{Z}}{\partial \boldsymbol{\theta}^{s h}}\right\|_{2}^{2}$ ，使用可以去掉。

用该上界替换原优化目标，优化问题就近似成：

$\min _{\alpha^{1}, \ldots, \alpha^{T}}\left\{\left\|\sum_{t=1}^{T} \alpha^{t} \nabla_{\mathbf{Z}} \hat{\mathcal{L}}^{t}\left(\boldsymbol{\theta}^{s h}, \boldsymbol{\theta}^{t}\right)\right\|_{2}^{2} \mid \sum_{t=1}^{T} \alpha^{t}=1, \alpha^{t} \geq 0 \quad \forall t\right\}$