相关定义与分类

多任务学习(Multi-task learning, MTL)最早可以追溯到1993年的一篇文章,它描述的是一种学习范式——多个任务的数据一起来学习,学习的效果有可能要比每个任务单独学习的结果要好。本质上是利用多个任务的共享信息来提高在所有任务上的泛化性。

多任务学习可以认为是迁移学习的一种表现形式

Wasi Ahmad的《Multi-Task Machine Learning》

参数共享形式

  • 参数的硬共享机制:硬共享是神经网络的多任务学习中最常见的一种方式。一般来讲,它可以应用到所有任务的所有隐层上,而保留任务相关的输出层。硬共享机制降低了过拟合的风险。有文献证明了该机制过拟合风险为O(#Tasks)O(\#Tasks).
  • 参数的软共享机制:软共享机制下的每个任务都有自己的模型和自己的参数,但是对模型参数的距离进行了正则化来保证参数相似,例如 L2正则化,迹正则化(trace norm)等.

目标函数设计

minθsh,θ1,,θTt=1TctL^t(θsh,θt)\min_{\substack{\boldsymbol{\theta}^{s h}, \boldsymbol{\theta}^{1}, \ldots, \boldsymbol{\theta}^{T}}} \sum_{t=1}^{T} c^{t} \hat{\mathcal{L}}^{t}\left(\boldsymbol{\theta}^{s h}, \boldsymbol{\theta}^{t}\right)

待更

不确定性加权

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

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

p(yfW(x))=N(fW(x),σ2)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 为观测噪声,将其作为标准差时,其大小代表了模型拟合效果的难易程度。

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

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

p(y1,y2fW(x))=p(y1fW(x))p(y2fW(x))=N(y1;fW(x),σ12)N(y2;fW(x),σ22)\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}

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

L(W,σ1,σ2)=logp(y1,y2fW(x))12σ12y1fW(x)2+12σ22y2fW(x)2+logσ1σ2=12σ12L1(W)+12σ22L2(W)+logσ1σ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(yfW(x))=Softmax(1σ2fW(x))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)

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

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

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

1
2
3
4
5
6
7
8
9
10
11
12
13
14
# 伪代码
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/σ21/\sigma^2 也是模型学习得到的,所以面对那些很难拟合的子任务L(W)\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θshθ1,,θTL(θsh,θ1,,θT)=minθshθ1,,θT(L^1(θsh,θ1),,L^T(θsh,θT))\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 条件来寻找满足帕累托最优的驻点

  1. 对于所有的任务t,θtL^t(θsh,θt)=0\forall t, \nabla_{\boldsymbol{\theta}^{t}} \hat{\mathcal{L}}^{t}\left(\boldsymbol{\theta}^{s h}, \boldsymbol{\theta}^{t}\right)=0;
  2. α1,,αT0 s.t. t=1Tαt=1∃ \alpha^{1}, \ldots, \alpha^{T} \geq 0 \ s.t. \ \sum_{t=1}^{T}\alpha^t=1 使得t=1TαtθshL^t(θsh,θt)=0\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.

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

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

minα1,,αT{t=1TαtθshL^t(θsh,θt)22t=1Tαt=1,αt0t}\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算法就是求解这类带约束的凸优化问题的一种算法。

参数更新算法

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

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

链式改进

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

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

ft(x;θsh,θt)=(ft(;θt)g(;θsh))(x)=ft(g(x;θsh);θt)zi=g(x;θsh)\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}

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

t=1TαtθshL^t(θsh,θt)22Zθsh22t=1TαtZL^t(θsh,θt)22\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}

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

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

minα1,,αT{t=1TαtZL^t(θsh,θt)22t=1Tαt=1,αt0t}\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\}

文章证明了这种近似依然能达到帕累托最优。

参考

  1. 当我们在谈论Multi-Task Learning(多任务/多目标学习) - 知乎
  2. 机器学习技术:多任务学习综述! - 知乎
  3. 多任务学习概述 - Harlin’s blog
  4. 论文讲解~CVPR2021跨领域多任务学习 Cross-Domain Multi-task Learning for Object Detection and Saliency Estimation - 知乎