多任务学习全流程细节

1. 为什么需要多任务学习

  1. 样本偏差
  2. 行为路径

2. 共享机制：模型架构和特征组合

两类共享

  1. 模型架构

    • 共享embedding
    • 共享中间层某些隐藏单元
    • 共享模型某一层或最后一层结果
    • 共享之外的部分各自独立
    • 模型设计中，层间关系自由组合搭配

  2. 特征组合

    • 多个任务可以采用不同的特征组合
    • 有的特征只属于模型架构的某个部分
    • 有些特征整个模型都可以使用

经典参数共享机制

  1. 参数的硬共享机制（基于参数的共享，Parameter Based）

    • 共享特征、特征embedding和隐层
    • 最后一层 FC + Softmax 实现不同任务
    • 最后通过一个线性融合来实现多目标排序

  2. 参数的软共享机制（基于约束的共享，Regularization Based）

    • 每个任务都有自己的参数和模型结构
    • 网络设计可以选择哪些共享哪些不共享
    • 通过正则化的方式拉近模型参数之间的距离（例如，L2正则，Dropout）

多任务学习的4种可能效果

  1. Well Done：所有任务共同提升，或者牺牲 辅助任务 实现 主任务 提升(主任务+辅助任务 结构)
  2. 还不错：所有任务不降低，至少一个任务提升（帕累托优化）
  3. 不理想：跷跷板现象，多个任务有涨有跌
  4. 无法接受：负迁移现象，所有任务都不如前

3. 多任务学习的两大优化方向

两大方向

  1. 网络结构（Architectures）

    • 目前多任务研究和应用的主要焦点
    • 设计更好的参数共享位置和方式
    • 那些参数共享、在什么位置共享、如何进行共享

  2. 优化策略（Optimization Strategy）

    • 宏观上看是达到多任务的平衡优化，微观上是缓解梯度冲突和参数撕扯
    • 设计更好的优化策略，以提升 反向传播 过程中的多任务平衡
    • 平衡Loss数量级（magnitude）
    • 调节Loss更新速度（velocity）
    • 优化梯度Grad更新方向（direction）

网络结构优化

经典结构变迁 Share Bottom -> MMoE -> PLE

  1. Share Bottom：早期常用方法（Hard or Soft）
  2. Google MMoE 2018：将 Share Bottom 分解成多个 Expert，引入门控网络自动控制梯度贡献
  3. Google SNR 2019：使用 NAS 对 Sub-Network 进行搜索组合
  4. 腾讯 PLE 2020：在 MMoE 基础上为每个任务增加自有 Expert，仅由本任务对其梯度进行更新
  • Share Bottom -> MMoE：不明显增加模型参数，且自动捕获任务之间的差异性。
  • MMoE -> PLE：更好地降低了相关性不强的任务间进行信息传递的可能隐患。

核心论文

  1. MMoE：《Modeling task relationships in multi-task learning with multi-gate mixture-of-experts》
    • MMoE 结构设计中的 Multi-gate 对于任务差异带来的冲突有一定的缓解作用，即使在多任务之间的的相关性不高的情况下，也有不错的效果。 MMoE 中不同的 expert 负责学习不同的信息内容，然后通过 gate 来组合这些信息，通过不同任务 gate 的 softmax 的热力分布差异，来表明expert对不同的目标各司其责，从而提升了效果。

MMoE核心代码参考：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
import torch
import torch.nn as nn
import torch.nn.functional as F

class MMoE_Layer(nn.Module):
    def __init__(self, expert_dim, n_expert, n_task):
        super(MMoE_Layer, self).__init__()
        self.n_task = n_task
        self.expert_dim = expert_dim
        self.n_expert = n_expert

        # 定义多个专家网络
        self.experts = nn.ModuleList([nn.Sequential(
            nn.Linear(expert_dim, expert_dim),
            nn.ReLU()
        ) for _ in range(n_expert)])
        
        # 定义多个门控网络，每个任务一个门
        self.gates = nn.ModuleList([nn.Sequential(
            nn.Linear(expert_dim, n_expert),
            nn.Softmax(dim=1)
        ) for _ in range(n_task)])
    
    def forward(self, x):
        # 通过所有专家网络
        expert_outputs = [expert(x) for expert in self.experts]  # List of [batch_size, expert_dim]
        expert_tensor = torch.stack(expert_outputs, dim=1)       # [batch_size, n_expert, expert_dim]
        
        towers = []
        for gate in self.gates:
            g = gate(x)                                       # [batch_size, n_expert]
            g = g.unsqueeze(2)                                # [batch_size, n_expert, 1]
            # 加权求和专家输出
            tower = torch.sum(expert_tensor * g, dim=1)       # [batch_size, expert_dim]
            towers.append(tower)
        
        return towers  # List of [batch_size, expert_dim] for each task

  1. SNR：《SNR: Sub-Network Routing forFlexible Parameter Sharing in Multi-Task Learning》

    • 思路与网络自动搜索（NAS）接近，通过动态学习产出多任务各自采用的 sub-network。研究思路是希望在更相似的任务下能学习到共享多一些的结构。

  2. PLE：《Progressive layered extraction (ple): A novel multi-task learning (mtl) model for personalized recommendations》

    • Share Bottom 结构的梯度冲突问题：在 Share Bottom 的结构上，整个共享参数矩阵如同质量较大的物体，在梯度更新的环节，两个 Loss 反向计算的梯度向量分别是 g1 和 g2，是这个物体受到的两个不同方向不同大小的力，这两个力同时来挪动这个物体的位置，如果在多次更新中两个力大概率方向一致，那么就能轻松达到和谐共存、相辅相成。反之，多个力可能出现彼此消耗、相互抵消，那么任务效果就会大打折扣。
    • MMoE 通过『化整为零，各有所得』解决思路：把一个共享参数矩阵化成多个结合 gate 的共享 Expert，这样不同的loss在存在相互冲突的时候，在不同的 expert 上，不同 Loss 可以有相对强弱的表达，那么出现相互抵消的情况就可能减少，呈现出部分 experts 受某 task 影响较大，部分 experts 受其他 task 主导，形成『各有所得』的状态。
    • PLE 通过增加 Spcific Experts：能进一步保障『各有所得』，保证稳定优化。

PLE核心代码参考：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
import torch
import torch.nn as nn
import torch.nn.functional as F

class PleLayer(nn.Module):
    def __init__(self, n_task, n_experts, expert_dim, n_expert_share, dnn_reg_l2=1e-5):
        super(PleLayer, self).__init__()
        self.n_task = n_task
        self.n_experts = n_experts  # List indicating number of experts per task
        self.expert_dim = expert_dim
        self.n_expert_share = n_expert_share

        # 定义每个任务的特定专家网络
        self.task_experts = nn.ModuleList([
            nn.ModuleList([
                nn.Sequential(
                    nn.Linear(expert_dim, expert_dim),
                    nn.ReLU()
                ) for _ in range(n_experts[i])
            ]) for i in range(n_task)
        ])

        # 定义共享专家网络
        self.share_experts = nn.ModuleList([
            nn.Sequential(
                nn.Linear(expert_dim, expert_dim),
                nn.ReLU()
            ) for _ in range(n_expert_share)
        ])

        # 定义门控网络，每个任务一个门
        self.gates = nn.ModuleList([
            nn.Sequential(
                nn.Linear(expert_dim, n_expert_share + n_experts),
                nn.Softmax(dim=1)
            ) for n_experts in n_experts
        ])

    def forward(self, x):
        # 通过共享专家网络
        share_outputs = [expert(x) for expert in self.share_experts]  # List of [batch_size, expert_dim]

        towers = []
        for i in range(self.n_task):
            # 通过任务特定专家网络
            task_specific_outputs = [expert(x) for expert in self.task_experts[i]]  # List of [batch_size, expert_dim]
            
            # 合并共享专家和任务特定专家
            combined_experts = share_outputs + task_specific_outputs  # List of [batch_size, expert_dim]
            combined_tensor = torch.stack(combined_experts, dim=1)    # [batch_size, n_expert_share + n_experts[i], expert_dim]
            
            # 通过门控网络获取权重
            g = self.gates[i](x)                                     # [batch_size, n_expert_share + n_experts[i]]
            g = g.unsqueeze(2)                                       # [batch_size, n_expert_share + n_experts[i], 1]
            
            # 加权求和专家输出
            tower = torch.sum(combined_tensor * g, dim=1)            # [batch_size, expert_dim]
            towers.append(tower)
        
        return towers  # List of [batch_size, expert_dim] for each task

优化策略改进

从 Loss 和 梯度 的维度去思考不同任务之间的关系，以期在优化过程中缓解梯度冲突，参数撕扯，尽量达到多任务的平衡优化。目前各式各样的多任务多目标优化方法策略，主要集中在3个问题：

  • Magnitude-Loss量级：Loss有大有小，取值较大的Loss会主导更新

    • 最可能的原因是不同任务使用的 Loss 类型不一样
    • 典型的例子是二分类任务 + 回归任务的多目标优化，L2 Loss 和 交叉熵损失 的 Loss 大小与梯度大小的量级和幅度可能差异很大，如果不处理会对优化造成很大干扰。

  • Velocity-Loss学习速度：不同任务在训练和优化过程中学习速度不一致

    • 不同任务因为样本的稀疏性、学习的难度不一致，在训练和优化过程中，存在 Loss 学习速度不一致的情况。
    • 如果不加以调整，可能会出现某个任务接近收敛甚至过拟合的时候，其他任务还是欠拟合的状态。

  • Direction-Loss梯度冲突：不同任务对参数的梯度大小和方向不一致

    • 相同参数被多个梯度同时更新的时候，由于不同任务产生的梯度存在不同的大小和方向，可能会出现冲突，导致相互消耗抵消，进而出现跷跷板、甚至负迁移现象。

经典方法

  1. Magnitude方向

    • Uncertainty weight 2018：自动学习任务的uncertainty，给uncertainty大的任务小权重、uncertainty小的任务大权重。

    • GradNorm 2018：综合任务梯度的二范数和Loss下降速度，设计关于权重的损失函数—— Gradient Loss，并通过梯度下降更新权重。

  1. Velocity方向
    • DWA 2019：用Loss下降速度来衡量任务的学习速度，直接得到任务的权重。

  1. Direction方向

    • MGDA(Parteo) 2018：提出一种基于Frank-Wolfe的优化方法，能适应大规模问题。并且为优化目标提供一个上限，通过优化上限可以通过单次反向传捆来更新梯度，而无需单独更新特定任务的梯度，减小了MGDA的计算开销。

    • PE-LTR(Parteo) 2019：在Parteo基础上为每个任务引入优先级约束，并提供一个两步解法来求解新的优化目标。

    • PCGrad 2020：两个梯度冲突时，直接把一个任务的梯度投影到另一个任务的法向量上以减轻梯度干扰。

    • Gradvac 2021：利用任务相关性设置梯度的相似性目标，并且自适应地对齐任务梯度以实现这些目标。

Uncertainty Weight

《Multi-task learning using uncertainty to weigh losses for scene geometry and semantics》

  1. 简单的多任务学习

往往是把所有的 $Loss$ 进行联合优化，通常需要手动调整它们的 $weights$：

然而这种方式通常存在如下问题：

  • 模型最终学习效果对于 $weights$ 非常敏感，否则很难同时获得对于多个任务都比较优的模型。
  • 同时手动调整这些 $weights$ 也是非常费时费力的工作。
  1. 通过 $uncertainty$ 来指导权重的调整

核心思路：$Loss$ 大的任务，包含的 $uncertainty$ 也应该多，而它的权重就应该小一点。
令 $uncertainty$ 是一个可学习的参数（$\sigma_1, \sigma_2$）：

模型优化过程中会倾向于“惩罚高 $Loss$ 而低 $\sigma$”的情况

  • 如果一个任务的 $Loss$ 高，同时 $\sigma$ 又小的话，这一项就会很大，优化算法就会倾向于优化它。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
import torch
import torch.nn as nn
import torch.nn.functional as F

# 定义多任务学习模型
class MultiTaskModel(nn.Module):
    def __init__(self, input_dim, shared_dim, task_dims, num_tasks):
        super(MultiTaskModel, self).__init__()
        self.num_tasks = num_tasks
        
        # 共享层
        self.shared_layer = nn.Sequential(
            nn.Linear(input_dim, shared_dim),
            nn.ReLU()
        )
        
        # 每个任务的特定输出层
        self.task_layers = nn.ModuleList([
            nn.Linear(shared_dim, task_dims[i]) for i in range(num_tasks)
        ])
        
        # 不确定性参数（log_var）
        self.log_vars = nn.Parameter(torch.zeros(num_tasks))
    
    def forward(self, x):
        shared_representation = self.shared_layer(x)
        task_outputs = [self.task_layers[i](shared_representation) for i in range(self.num_tasks)]
        return task_outputs

# 定义Loss
class MultiTaskLoss(nn.Module):
    def __init__(self, num_tasks):
        super(MultiTaskLoss, self).__init__()
        self.num_tasks = num_tasks
    
    def forward(self, log_vars, y_true, y_pred):
        loss = 0
        for i in range(self.num_tasks):
            precision = torch.exp(-log_vars[i])
            task_loss = F.mse_loss(y_pred[i], y_true[i], reduction='mean')
            loss += precision * task_loss + log_vars[i]
        return loss

# 示例参数
input_dim = 10
shared_dim = 64
task_dims = [1, 1]  # 两个回归任务
num_tasks = len(task_dims)

# 初始化模型、损失函数和优化器
model = MultiTaskModel(input_dim, shared_dim, task_dims, num_tasks)
criterion = MultiTaskLoss(num_tasks)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-5)

# 示例数据
batch_size = 3
y_true_task1 = torch.tensor([[1.0], [2.0], [3.0]])
y_true_task2 = torch.tensor([[1.5], [2.5], [3.5]])
inputs = torch.randn(batch_size, input_dim)

# 训练步骤
model.train()
optimizer.zero_grad()
y_pred = model(inputs)
loss = criterion(model.log_vars, [y_true_task1, y_true_task2], y_pred)
loss.backward()
optimizer.step()

print(f"Total Loss: {loss.item()}")
print(f"Log Vars: {model.log_vars.data}")
GradNorm（未完待整理）

《Gradnorm: Gradient normalization for adaptive loss balancing in deep multitask networks》

核心思想：

  • 希望不同的任务的 Loss 量级是接近的
  • 希望不同的任务以相似的速度学习

具体实现：

  • 定义两种类型的Loss：
    • Label Loss 和 Gradient Loss
    • 两种类型的Loss独立优化，不进行运算
DWA

《End-to-End Multi-Task Learning with Attention》
这篇paper中直接定义了一个指标来衡量任务学习的快慢，然后来指导调节任务的权重。

用这一轮 Loss 除以上一轮 Loss ，这样可以得到这个任务 Loss 的下降情况用来衡量任务的学习速度，然后直接进行归一化得到任务的权重。

  • 当一个任务 Loss 比其他任务下降的慢时，这个任务的权重就会增加
  • 下降的快时权重就会减小
  • 可以视作只考虑了任务下降速度的简化版的 Gradient normalization，简单直接
PCGrad（未完待整理）

《Gradient surgery for multi-task learning》

GradVac（未完待整理）

《Investigating and improving multi-task optimization in massively multilingual models》

总结

  1. 首先关注业务场景，思考业务目标优化重点，进而确定多任务的组合形式：

    • 主任务 + 主任务：解决业务场景既要又要的诉求，多个任务都希望提升
    • 主任务 + 辅任务：辅助任务为主任务提供一些知识信息的增强，帮助主任务提升，需要考虑不同任务间的重要度和相似性，考虑清楚辅助任务和主任务的关系

  2. 实际训练过程中，可以训练优化其中1个任务，观察其他任务的 Loss 变化

    • 其他任务 Loss 同步下降，则关联性较强
    • 其他任务 Loss 抖动或有上升趋势，要回到业务本身思考是否要联合多任务训练

  3. 网络结构选择 MMoE 或者 PLE

  4. 训练过程中关注 LossLoss 的量级，如果不同任务之间差异很大，注意约束和控制

  5. 训练过程的优化策略，可以尝试 PCGrad 等方法对梯度进行调整，并观察效果