Feed m Birds with One Scone: Accelerating Multi-task Gradient Balancing via Bi-level Optimization

Feed m Birds with One Scone: Accelerating Multi-task Gradient Balancing via Bi-level Optimization

论文信息: arXiv:2603.07389 | 2026年3月8日 | Meta AI
作者: Xuxing Chen, Yun He, Jiayi Xu, Minhui Huang, Xiaoyi Liu, Boyang Liu, Fei Tian, Xiaohan Wei, Rong Jin, Bo Long, Xue Feng 等12人

1. 问题与动机

多任务学习（MTL）在工业推荐和广告排序系统中极为普遍——一个模型需同时优化点击率、转化率、质量评分等多个目标。核心难题是梯度冲突：不同任务的梯度方向可能相互抵消。现有梯度平衡方法（MGDA、CAGrad、SDMGrad等）存在O(md)的计算复杂度瓶颈，对拥有数十亿参数的大型模型不可行。FAMO虽将复杂度降至O(d)，但缺乏充分理论支撑且与Adam等自适应优化器存在兼容性问题。

2. 方法概述

提出 MARIGOLD 框架，将梯度平衡重新定义为双层优化问题：下层以当前任务权重λ训练模型参数θ，上层优化λ使最坏情况下的损失减量最大化（min-max问题）。关键创新在于用零阶方法估计超梯度——对λ施加随机扰动，用函数值之差近似梯度方向，将复杂度压缩至O(d)。

3. 核心技术贡献

1. 无线性化的梯度平衡重构：直接在真实损失上定义上层目标，绕开线性化近似，支持任意优化器（Adam、AdaGrad等）
2. 高效零阶超梯度估计：单点零阶估计每次迭代仅需一次额外前向传播，复杂度O(d)
3. 与任意优化器兼容：下层训练可对接任何现成优化器，与工业系统无缝集成
4. 通用双层框架：可扩展到辅助学习等场景，仅需替换上层目标函数

4. 实验结果分析

公开数据集（NYU-v2）：分割mIoU 41.01（FAMO 38.60，+2.41），平均排名3.56（FAMO 5.44）。计算时间152秒/epoch（FAMO 182秒，MGDA 375秒），比FAMO快17%，比MGDA快60%。

Cityscapes：分割mIoU 75.30（FAMO 74.76），平均排名4.50（FAMO 7.75）。

Meta工业实验：在广告排序系统上，点击预测NE +0.08%，转化|点击NE +0.03%，绝对转化NE +0.07%，蒸馏NE +0.14%。所有四个任务均正向收益，无负迁移。

5. 局限与未来方向

• 缺乏收敛率理论分析
• 温度系数β需调参（建议{0.03, 0.1, 0.3}）
• 单点零阶估计在高维空间方差较大
• 任务数量极多时的可扩展性未系统研究

6. 工业实践意义

• 计算成本可接受：O(d)复杂度仅需一次额外前向传播，对日训万亿样本的系统可行
• 与Adam无缝兼容：少数明确支持任意优化器的梯度平衡方法
• Pareto改善有实证：Meta内部所有任务均正向，最坏情况设计天然对抗负迁移
• 落地建议：可作为固定权重多任务系统的即插即用升级，替换现有梯度平衡方案