MGDA-UB：寻找多任务学习中的帕累托最优，缓解任务冲突问题

Multi-Task Learning as Multi-Objective Optimization. arxiv: 1810.04650, NeurIPS 2018.

---

---

一、问题：多任务学习中的"跷跷板"困境

多任务学习（MTL）试图用一个模型同时解决多个相关任务，比如同时识别图像中的物体、分割语义、预测深度。理想情况下，任务间互相帮助，共享知识。但现实中，任务会争夺模型容量：提升任务 A 的性能，往往导致任务 B 下滑，像跷跷板一样此消彼长。

传统方法把多个损失函数简单加权求和：$L=0.5L_1+0.5L_2$。这背后有个致命假设：任务不冲突。然而，在真实训练的时候，这个假设几乎总不成立。

例子：训练一个模型同时识别 MultiMNIST 图片的左上角数字（任务 L）和右下角数字（任务 R）。共享的卷积层容量有限，模型可能会"偏心" —— 学会了认左上角，右下角的准确率就掉下去。无论怎么调权重，总有一个任务不如单独训练的模型。

二、重要性：为什么必须解决任务冲突？

多任务学习的价值不只是省参数。真实世界的任务本就共享底层规律：自动驾驶的语义分割、深度估计、障碍物检测都依赖相同的物理规律（光学、几何）。如果模型能真正兼顾多个目标，就能学到更通用的表征。

但只要任务冲突存在，调权重就是一种"和稀泥"，找不到真正的最优解。我们需要的不是妥协，而是找到一组参数，在不伤害任何任务的前提下，无法再继续提升任一任务。这被称为帕累托最优。

各个 benchmark 下的评估指标：对每个任务独立评估。在 MultiMNIST 中看左右数字分类准确率；在 CelebA 中看40个人脸属性分类的平均错误率；在 Cityscapes 中看语义分割(mIoU)、实例分割误差、深度估计误差。目标是所有任务同时变好，而非平均性能最优。

三、难点：理论漂亮，实践卡壳

多目标优化领域早有解法。其中的 MGDA（Multiple Gradient Descent Algorithm）算法直击要害：

1. 算出每个任务对共享参数的梯度 ${\nabla }_{\theta }{L}_1,{\nabla }_{\theta }{L}_2$
2. 求解凸组合：找权重 $\alpha$，让 $\Vert {\alpha }_1{\nabla }_{\theta }{L}_1+{\alpha }_2{\nabla }_{\theta }{L}_2{\Vert }^2$ 最小
3. 用这个组合梯度更新参数

为什么这个范数是关键？

• 如果两梯度方向一致，组合向量长，说明任务协同，可以大步前进
• 如果方向相反，组合向量短，说明任务打架，必须小步调和
• 某种程度上，最小范数点对应的 α，就是最不伤害任何任务的"黄金比例"

致命问题：深度学习里，$\theta$ 有百万维，每次迭代要：

• T 次反向传播（每个任务一次）→ 训练时间 × T
• 在百万维空间求解 α → 计算量爆炸

现有方法要么放弃理论保证（GradNorm），要么只适用于小模型。

四、核心思想：从百万维降到 256 维

这篇文章提出 MGDA-UB，洞察来自深度学习模型的 编码器-解码器结构：

输入 → [共享编码器] → 表示Z (256维) → [任务特定头] → 输出

链式法则告诉我们：
$${\nabla }_{\theta }L=(\partial Z/\partial \theta )\cdot {\nabla }_{Z}L$$

• ${\nabla }_{\theta }L$：参数空间梯度，百万维
• ${\nabla }_{Z}L$：表示空间梯度，仅256维（批量大小×表示维度）
• $\partial Z/\partial \theta$：编码器的雅可比矩阵，所有任务共享

关键转换：
$$\begin{array}{rl}& \Vert {\alpha }_1{\nabla }_{\theta }{L}_1+{\alpha }_2{\nabla }_{\theta }{L}_2{\Vert }^2\\ =& \Vert (\partial Z/\partial \theta )\cdot ({\alpha }_1{\nabla }_Z{L}_1+{\alpha }_2{\nabla }_Z{L}_2){\Vert }^2\\ \le & \Vert \partial Z/\partial \theta {\Vert }^2\cdot \Vert {\alpha }_1{\nabla }_Z{L}_1+{\alpha }_2{\nabla }_Z{L}_2{\Vert }^2\end{array}$$
解释：

• $\partial Z/\partial \theta$ 与 α 无关，是个常数
• 在某种程度上，优化最后一行的第二项，等价于优化左边
• 而我们可以在一次反向传播中，从表示层 Z 分叉算出所有 ${\nabla }_Z{L}_t$

计算成本：从 O(T×百万维) 降到 O(百万维 + T × 256维)。40 个任务时，训练加速 25 倍。

理论保证：只要雅可比矩阵 $\partial Z/\partial \theta$ 满秩（实践中几乎总成立，否则任务线性相关无需权衡），优化这个上界等价于优化原问题，仍能找到帕累托最优解。

---

五、Frank-Wolfe：用最简单的投票找 α

降维后，我们要在表示空间求解：

$$\underset{\alpha }{\min }\Vert \sum {\alpha }_t{\nabla }_Z{L}_t{\Vert }^2,\text{s.t.}{\alpha }_t\ge 0,\sum {\alpha }_t=1$$

这是一个凸二次规划，但用通用求解器仍然可能很慢。Frank-Wolfe 算法用 “贪心投票” 策略，5-10 次迭代就收敛：

预计算：M_ij = g_i^T g_j  （所有梯度两两点积，T×T矩阵）

初始化：α = [1/T, ..., 1/T]  # 先平均分

重复以下步骤：
  1. 找出"最不和谐"的任务：
     t̂ = argmin_t Σ_j α_j · M_tj
     （计算当前组合方向与每个梯度的对齐程度，选最差的）
  
  2. 往这个任务靠一点：
     α_new = (1-γ)·α + γ·e_t̂
     （e_t̂是只在t̂位置为1的向量，γ∈[0,1]）
  
  3. γ用闭式解直接算出最优步长（论文Algorithm 1）
  
直到 γ≈0（再调也没用）

六、实验：理论落地，全面碾压

论文在 2 到 40 个任务的三个场景测试：

1 MultiMNIST（2 任务）

• 设定：28×28 图片，左上角和右下角各一个数字，同时分类
• 模型：LeNet 变体，编码器+两个任务头
• 结果：本文方法 97.26% / 95.90%，完全达到单任务水平。所有基线（均匀加权、GradNorm、不确定性加权）均落后 1% 以上，说明它们无法解决容量争夺。

2 CelebA（40属性分类）

• 设定：20 万张人脸，40 个二元分类任务（是否戴眼镜、是否微笑等）
• 模型：ResNet-18 编码器 + 40 个独立全连接层
• 结果：平均错误率 8.25%，显著优于 GradNorm（8.44%）和均匀加权（9.62%）。绝大多数属性上最优，证明方法在任务数爆炸时依然有效。

3 Cityscapes（3任务场景理解）

• 设定：街景图像，同时做语义分割、实例分割、深度估计

• 模型：ResNet-50 编码器 + 任务特定解码器（金字塔池化模块）

• 结果：

  方法	语义分割 (mIoU↑)	实例误差 (px↓)	深度误差 (px↓)
  单任务	60.68	11.34	2.78
  均匀加权	54.59	10.38	2.96
  本文方法	66.63	10.25	2.54

  三任务全面超越单任务，说明任务真正协同，而非争夺。

计算效率：

• 3 任务：训练时间减少 40%
• 40 任务：加速 25 倍
• 意外之喜：近似后准确率不降反升，因在更低维空间优化，梯度噪声更小。

代码：PyTorch 实现仅需约 50 行。核心是在反向传播时拦截表示层 Z 的梯度，用论文 Algorithm 2 的 Frank-Wolfe 求解器更新 α，再组合原始梯度更新参数。

七、局限与边界

1. 架构依赖：要求模型有明确的共享表示层（编码器-解码器结构）。对全共享或全独立的模型不适用。
2. 满秩假设：需雅可比矩阵 $\partial Z/\partial \theta$ 满秩。实际中几乎总成立，但若任务线性相关（如任务 1 = 2 × 任务 2），方法退化为加权平均。
3. 任务数量：Frank-Wolfe 复杂度 O(T^2)。任务数上百时可能需要更高效的近似。
4. 非凸性：理论保证收敛到帕累托平稳点（局部最优），非全局最优。