梯度手术-多任务学习优化方法[NeurIPS 2020]

  • 2021 年 1 月 15 日
  • AI

 分享一篇NeurIPS 2020上关于多任务学习的论文。


在之前的一篇文章《“一心多用”的方法——多任务学习》中,对多任务学习(Multi-task Learning,MTL)的背景、应用、挑战做了一些简单的介绍。在文章末尾,还提了我个人认为的待研究的方向:

针对这些挑战,在多任务学习的领域中,可以从以下几个研究方向考虑能做一些工作:

    1. 多任务网络结构设计

    2. 多任务损失函数设计

    3. 辅助学习(任务的筛选方法)

    4. 多任务网络的训练方法

今天分享的这一篇论文,便是从第4点“多任务网络的训练方法”着手,提出名为PCGrad的方法,在训练多任务模型时,通过调整梯度,尽量降低任务之间的干扰。并且,这一方法与模型无关:它可以与其他多任务网络结构相结合,进一步提升多任务学习的性能

图源[1]

首先需要明确的是,在多任务网络的训练过程中,需要解决一些什么问题。

假设有下式中的两个优化任务,对应损失函数分别为L1和L2,θ对应模型的参数(在这个例子中做了简化,只考虑两个参数θ1和θ2)。对多任务学习而言,优化目标是学习模型的参数θ,以最小化L(L1与L2的和)。

式源[1]

定义了这一简化的问题后,下图(a)可视化了这一多任务学习问题的目标空间(即L和(θ1, θ2)的关系),图(b)和图(c)分别为任务1和任务2的目标空间,即为图(a)的拆解版。其中,颜色越深的位置表示损失函数L越小。

图源[1]

对于一个任务1+任务2的多任务问题,其目标空间为下图(d):最深颜色的一条线为两个任务共有的目标最优(损失最小)区域。

图源[1] (d)中箭头朝向表示梯度的方向,箭头长度表示梯度的大小

为了解决这一多任务的优化问题,一个naive的做法是:使用经典的优化器,分别对两个任务的损失函数进行梯度计算得到g1和g2,然后进行向量加法得到最终的梯度g (参见下面的Definition 4)。

这一做法存在显著的问题,在上图(d)中两个箭头的起点处,进行梯度计算,得到任务1梯度g1的红色向量(在该点处任务1的目标空间比较陡峭(参见上图b),故梯度大)和任务2的梯度g2的蓝色向量(在该点处任务2的目标空间比较平缓(参见上图c),故梯度小)。这两个任务的梯度向量:

  • 数值相差巨大;

  • 方向南辕北辙。

因此,简单的对梯度进行向量加法虽然看似能够整体上优化损失函数,但并不能达到多任务学习的优势,因为上述问题可能带来:

  • 减速训练过程;

  • 带来性能下降,无法达到最优点。

具体来说,可以将问题分解为:

  1. 梯度冲突:不同任务梯度之间但更新方向不同,定义上来说,当cosθ<0时,两个梯度存在冲突;

    图源[1]

  2. 存在主导梯度:两个任务的梯度,虽然方向并不冲突,但其中一个梯度数值上远大于另一个梯度,则较大梯度会形成主导从而影响较小梯度。文中定义一个梯度幅值相似度(gradient magnitude similarity),用于衡量两个梯度的大小差异;

  3. 高曲率:在高曲率的位置,主导任务的性能提升会被高估,其性能下降会被低估。


已有的一些改进方法,不能很好的解决上述的三个问题,例如:

  • 对不同任务的损失函数增加权重(通过经验、或者通过Uncertainty的计算[2]),但这只改变了梯度的数值差异。

这篇论文的作者提出的PCGrad (projecting conflicting gradients)方法,同时处理数值和方向的差异。下图(e)中,通过PCGrad与Adam的结合,能够沿着空间中正确的优化方向找到最小损失的点。

图源[1]

简单来说,就是要处理冲突的梯度(下图a),将第i个任务的梯度投影到与第j个任务正交的方向(下图b),或者相反(下图c)。下图d表示的是对不冲突的梯度,不做处理。

图源[1]

算法流程如下

  1. 对一个batch中对每个任务Ti,随机采样一个其他任务Tj,计算Ti和Tj对应梯度向量之间的余弦相似度;

  2. 如果余弦相似度为正,表示梯度无冲突,不做梯度修改。若为负,表示是相互冲突的梯度,根据下式计算投影;

  3. 重复上述步骤,直到这一batch中的任务遍历完成;

  4. 最终的梯度方向为,投影后的梯度向量之和。

图源[1] PCGrad的梯度更新流程

论文还提供了两个定理及其详细证明,论述了PCGrad方法的收敛性(定理一)以及其相较于标准梯度下降方法能达到更低误差(定理二):

这里不讲定理和证明部分,分享一些我个人认为有意思的前提条件

  1. 任务的损失函数是凸的且可微的,凸(convexity)对于神经网络方法来说虽然不满足,但是很多用于神经网络的优化方法也是基于凸的假设;

  2. 定理2的(a)表示,两个梯度的方向冲突不能太小,且两个梯度的幅值差异要足够大,否则可能影响收敛效果;

  3. 定理2的(b)表示,多任务的曲率H有下界,且曲率需要足够大;

  4. 定理2的(c)表示,学习的步长/学习率要足够大,并且这个值与曲率有关,曲率越大,步长也要越大。因此,在使用PCGrad方法时,需要设置较大的学习率。

虽然并不是说不满足上述前提条件就一定不work了,但在一定程度上对于算法的应用效果能够给予分析的参考。

最后,来看一看实际的实验效果。作者在多个计算机视觉领域监督学习数据集和机器人领域强化学习数据集上做了实验,均取得了好的效果。具体可以查阅原文,这里摘录一个在CityScapes和NYUv2数据集上的实验结果:通过和其他多任务学习的方法(例如Cross-Stitch, MTAN,Uncertainty)进行结合,能够取得更好的性能。

表源[1]

总结一下

  • 代码开源在[3]中;

  • 要说这一方法有什么明显的缺点或者代价呢,那就是会带来额外的训练时间和显存占用

  • 但这一方法总体思路清晰(提出了优化中的三个问题并针对性解决),实现简单,并且模型无关,能够与其他多任务方法共同使用;

  • 类似的思路,或许也可以用于其他领域,比如说持续学习、或者多模态学习中。

参考资料:

[1] Tianhe Yu, Saurabh Kumar, Abhishek Gupta, Sergey Levine, Karol Hausman, and Chelsea Finn. Gradient surgeryfor multi-task learning, 2020.

[2] Alex Kendall, Yarin Gal, and Roberto Cipolla. Multi-task learning using uncertainty to weigh losses for scenegeometry and semantics, 2018.Al

[3] //github.com/tianheyu927/PCGrad


– END –

新朋友们可以看看我过往的相关文章

【相关推荐阅读】


“一心多用”的方法——多任务学习

多模态:自动驾驶中的多模态学习

模式识别学科发展报告丨前言

深度学习在计算机视觉中应用的三大问题

从2018到2020的AI发展预测


长按二维码关注我