多任务学习优化总结 Multi-task learning（附代码）

目录

一、多重梯度下降multiple gradient descent algorithm (MGDA)

二、Gradient Normalization (GradNorm)

三、Uncertainty

---

多任务学习的优势不用说了，主要是可以合并模型，减小模型体积，只用一次推理也可以加快速度。对于任务表现的提升，可能不是那么明显，理论上来说相似的任务确实可以提高彼此的表现，但是在实际应用中，在减小模型的同时仅仅想要保持每一个任务的表现都不是那么容易。

总结来说，多任务学习的发展主要有两个方向：

1. 调整网络的结构以适应于多任务的学习。这种方法需要根据具体问题去更改网络结构，繁琐又难以再现。一般来说调整网络结构只能根据经验，调优起来也很不方便。
2. 调整网络的loss来进行不同的任务融合。在这种思路中，更好的加权方式应该是动态的，根据不同任务学习的阶段，学习的难易程度，甚至是学习的效果来进行调整。

如果想要仔细了解多任务学习的发展历程，可以看看这篇Survey：Multi-Task Learning for Dense Prediction Tasks: A Survey。

以下主要介绍几种常用的多任务学习方法：MGDA，GradNorm，Uncertainty。

一、多重梯度下降multiple gradient descent algorithm (MGDA)

这个方法来自Intel Labs 2019年的一篇文章Multi-Task Learning as Multi-Objective Optimization，这篇文章的优化部分有点难，需要很多时间去理解。 如果要节省时间，建议直接拉GitHub上的代码下来对照着学习：https://github.com/intel-isl/MultiObjectiveOptimization。

作者的motivation是，由于任务之间不是完全竞争或者不竞争的关系，而是一种相互博弈的关系，这时候单纯的线性解就没那么有用了，所以需要去找到一个帕累托最优解来优化多个任务的表现，也就是，把多任务学习变成多目标优化问题。而在多目标优化算法中，就有一种叫MGDA的方法，它可以在共享参数的过程中优化多个任务的梯度。

为了利用MGDA方法，定义一个优化问题：

 这个问题的解要么是满足KKT条件的点（鞍点），要么是梯度下降的方向。对不同的参数有不同的梯度下降：

• 在task-specific参数上做一般的梯度下降（下图第2行）
• 在shared参数上做的梯度下降（下图第5行）

 怎么去解这个问题呢？作者用的是2013年提出的Frank-Wolfe算法去解：

我们可以对比代码来看（以下MGDA及MGDA-UB的代码中，shared参数在低层，task- specific参数在高层，一般情况都是如此）：

# This is MGDA 只有这一段与MGDA-UB有区别
for t in tasks:
     # Comptue gradients of each loss function wrt parameters
     optimizer.zero_grad()
     rep, mask = model['rep'](images, mask)  # 先基于shared参数进行推理
     out_t, masks[t] = model[t](rep, None)  # 再基于task-specific参数进行推理
     loss = loss_fn[t](out_t, labels[t])
     loss_data[t] = loss.data[0]
     loss.backward()  
     grads[t] = []
     # 一般的梯度下降（loss不缩放）
     for param in model['rep'].parameters():
         if param.grad is n