跳转至

Delve into the Applicability of Advanced Optimizers for Multi-Task Learning

会议: NeurIPS 2026
arXiv: 2604.08939
代码: https://github.com/xxx
领域: 优化/理论
关键词: 多任务学习, 优化器, 动量机制, Muon, 自适应动量

一句话总结

本文揭示了现代高级优化器(如 Adam、Muon)中的指数移动平均动量会削弱多任务学习的即时梯度去冲突效果,并提出基于局部曲率的自适应动量调节与轻量方向保持策略(APT),显著提升多任务联合优化性能。

研究背景与动机

多任务学习(Multi-Task Learning, MTL)旨在利用共享网络骨干与任务专用分支同时处理多项任务,在自动驾驶、多模态感知和场景理解等领域发挥着基础性作用。近年来,基于优化的多任务学习方法通过显式操纵任务梯度(如投影正交化、最小范数逼近或帕累托妥协)或自适应加权任务损失,在平衡任务收敛速度与消除梯度冲突方面取得了重要进展,甚至在多个基准测试中逼近或超越了单任务学习(STL)上限。

然而,现有优化型多任务学习方法在设计时大多隐式预设了无动量的传统随机梯度下降(SGD)设定,完全忽视了现代高级优化器(如 Adam、AdamW 及新兴的 Muon)中普遍存在的指数移动平均(EMA)动量机制。在标准的 EMA 更新中,历史动量方向往往占据 90% 以上的更新权重,导致即使当前时刻的 MTL 算法精心计算出了完美的非冲突梯度方向,叠加进累积动量后最终的参数更新方向依然会严重偏离预定轨迹,使得即时去冲突(instant de-conflicting)蜕化为在较长时间跨度上平摊的被动平衡(amortized de-conflicting)。

另一方面,完全去除动量虽然能恢复瞬时梯度的控制权,却会彻底丧失 EMA 带来的梯度降噪与方差抑制能力,导致训练轨迹剧烈震荡且极易卡在次优解。此外,近期兴起的高效矩阵正交化优化器 Muon 通过牛顿-舒尔茨迭代归一化奇异值,在理论上天然扮演着压缩强势任务、提升弱势任务的隐式多任务学习角色,但其正交化过程对输入梯度的方向多样性与有效秩提出了极为苛刻的要求。核心 idea:本文构建了统一的 APT 框架,利用相邻步梯度余弦相似度感知局部曲率并动态调节动量系数以兼顾即时去冲突与轨迹平滑,同时针对 Muon 引入轻量方向保持机制,从根本上释放高级优化器在多任务优化中的潜能。

方法详解

整体框架

APT(Applicability of advanced oPTimizers)框架旨在系统性解决高级优化器与多任务学习机制之间的脱节问题。整体流程接收多任务网络在共享参数处反向传播得到的各个独立任务梯度,首先通过多任务梯度聚合或轻量方向保持模块合成几何信息丰富的瞬时更新方向,随后经由曲率感知模块自适应计算当前的动量衰减系数,最后输入一阶与二阶矩更新以及优化器更新步中。

对于传统 Adam 系列优化器,APT 监控当前聚合梯度与上一时刻梯度的方向夹角作为局部曲率代理,在曲率剧烈变化的高冲突区域自动降低动量系数以放大 MTL 即时去冲突梯度的权重,在平坦区域则恢复高动量以维持降噪稳定性。对于采用谱正交化的高级优化器 Muon,APT 通过中心化锚点相似度门控机制合成具有高有效秩的多任务输入梯度,协同牛顿-舒尔茨正交化实现跨任务的特征子空间重平衡。

%%{init: {'theme': 'base', 'themeVariables': {'primaryColor': '#EEF2FF', 'primaryBorderColor': '#6366F1', 'lineColor': '#4F46E5', 'fontSize': '13px'}, 'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["多任务独立梯度集合<br/>{g_1, g_2, ..., g_K}"] --> B["轻量方向保持 LDP<br/>基于均值锚点门控合成高秩梯度"]
    B --> C["瞬时多任务梯度 G_t"]
    C --> D["曲率感知自适应动量<br/>根据相邻梯度余弦动态调整 β_t"]
    D --> E["参数动量更新与正交化<br/>Adam 动量截断 / Muon 谱投影"]
    E --> F["模型共享参数更新 θ_{t+1}"]

关键设计

1. 局部曲率感知的自适应动量调节:平衡即时去冲突与时序平摊平滑 针对传统固定动量参数(如 \(\beta_1 = 0.9\))掩盖瞬时多任务去冲突梯度的缺陷,本文利用相邻优化步聚合梯度之间的余弦相似度作为损失曲面的局部曲率代理: $$ \rho_t = \cos(G_t, G_{t-1}) $$ 当优化轨迹进入任务梯度剧烈冲突的复杂鞍点或峡谷区域时,相邻梯度的方向差异显著增大,导致 \(\rho_t\) 迅速减小,反映出较高的局部曲率。此时优化器应当优先响应当前的瞬时去冲突梯度,因此动量系数 \(\beta_t\) 被动态衰减至较低水平;反之,在各任务梯度方向协调一致、曲面相对平坦的低曲率区域,\(\rho_t\) 趋近于 1,动量系数回弹至高位以享受 EMA 带来的方差削减与平稳加速收益。动量系数的闭式自适应规则为: $$ \beta_t = \beta_{\min} + (\beta_{\max} - \beta_{\min}) \cdot \operatorname{Clip}\left(\frac{\rho_t + 1}{2}, 0, 1\right) $$ 在实践中,动量上下界通常设置为 \(\beta_{\max} = 0.9\) 与 \(\beta_{\min} = 0.1\)。理论分析表明,该动态规则有效压缩了历史动量与当前 MTL 梯度的跟踪误差上界 \(\limsup_{T\to\infty} \frac{1}{T}\sum_{t=1}^T \|m_t - G_t\|^2 \le C \limsup_{T\to\infty} \frac{1}{T}\sum_{t=1}^T \beta_t^2 \|G_t - G_{t-1}\|^2\),使得模型序列在非凸多任务目标下能够收敛到半径更紧致的近似帕累托平稳邻域。

2. 隐式多任务学习机制下的 Muon 谱特性与轻量方向保持 Muon 优化器通过牛顿-舒尔茨(Newton-Schulz)迭代计算梯度矩阵 \(G_t\) 的极分解近似 \(O_t = U_t V_t^\top\)。根据奇异值分解 \(G_t = \sum_{k=1}^r \sigma_k u_k v_k^\top\),Muon 将所有非零奇异值强制归一化为 1,即 \(O_t = \sum_{k=1}^r u_k v_k^\top\)。借助切比雪夫总和不等式(Chebyshev's Sum Inequality),对于在主正交基上投影较大且与奇异值正相关的强势任务 \(i\),以及投影较小的弱势任务 \(j\),有: $$ \frac{P_i(G_t)}{P_j(G_t)} \ge \frac{P_i(O_t)}{P_j(O_t)} $$ 这证明了 Muon 在谱空间上天然具有压缩强势任务能量、提升弱势任务贡献的隐式均衡能力。然而,由于 Muon 将谱空间完全展平,若输入的聚合梯度有效秩过低或包含大量退化奇异分量,高频噪声会被同等放大。为给 Muon 提供高质且富含各任务几何特征的输入,作者提出轻量方向保持策略(Light Direction Preservation, LDP)。该策略以全任务均值梯度 \(\bar{g} = \frac{1}{K}\sum_{i=1}^K g_i\) 作为中心妥协锚点,计算每个任务梯度与锚点的方向一致性 \(c_i = \cos(g_i, \bar{g})\),并经由缩放 Sigmoid 门控完成重赋权: $$ \omega_i = \operatorname{Sigmoid}(\xi \cdot c_i), \quad \tilde{\omega}i = K \cdot \frac{\omega_i}{\sum_i g_i $$ 该设计无需运行昂贵的二次规划(QP)或 SVD 迭代求解器,即可以极小的计算开销保留高有效秩的梯度子空间,完美匹配 Muon 的正交化特性。}^K \omega_j}, \quad G_t = \sum_{i=1}^K \tilde{\omega

损失函数 / 训练策略

在训练流程中,对于结合 Adam/AdamW 的设定,针对多任务模型共享主干参数维护一阶矩 \(m_t = \beta_t m_{t-1} + (1 - \beta_t) G_t\) 与二阶矩 \(v_t = \beta_2 v_{t-1} + (1 - \beta_2) G_t^2\),并加入动态偏置校正因子 \(\rho_{1,t} = 1 - \prod_{j=1}^t \beta_j\)。对于结合 Muon 的设定,主干卷积与线性层的权重矩阵直接接受 LDP 构造的高秩梯度 \(G_t\) 进行牛顿-舒尔茨正交化迭代,偏置与一维向量参数保留标量优化;实验中搭配自适应学习率调度器(如 Reduce On Plateau),在验证指标停滞时衰减步长,以确保多任务帕累托前沿的精细收敛。

实验关键数据

主实验

在室内多任务场景理解基准 NYUv2(涵盖语义分割、深度估计、表面法向量预测 3 项任务)和城市场景理解基准 Cityscapes(分割与深度 2 项任务)上,基于 MTAN 共享骨干网络展开对比评估。下表汇总了 NYUv2 上的综合表现,其中 \(\Delta m\%\) 衡量相较于独立训练(Independent)基准的综合性能退化率(越低越好,负值表示超越独立单任务表现),MR 为多任务平均排名。

方法 (Method) 分割 mIoU ↑ 分割 Pix. Acc ↑ 深度 Abs. Err ↓ 表面法向 Mean Angle ↓ 平均排名 (MR) ↓ \(\Delta m\% \downarrow\)
Independent 38.30 63.76 0.68 25.01 - -
Linear Scalarization (LS) 39.29 65.33 0.55 28.15 12.56 5.46
MGDA 30.47 59.90 0.61 24.88 9.56 1.47
PCGrad 38.06 64.64 0.56 27.41 12.67 3.83
CAGrad 39.79 65.49 0.55 26.31 9.44 0.29
Nash-MTL 40.13 65.93 0.53 25.26 5.11 -4.04
FAMO 38.88 64.90 0.55 25.06 8.56 -4.10
FairGrad 39.74 66.01 0.54 24.84 5.00 -4.66
COST 38.06 64.71 0.54 24.47 4.67 -5.39
APT-CG 40.51 66.32 0.53 25.77 6.78 -2.67 ± 0.84
APT-Fair 38.97 65.07 0.54 24.44 5.67 -6.04 ± 0.22
APT-Muon 42.47 67.92 0.49 25.22 5.78 -5.54 ± 0.96

在 Cityscapes 双任务基准测试中,APT 同样展现出显著的增强效果:

方法 (Method) 分割 mIoU ↑ 分割像素准确率 ↑ 深度绝对误差 (Abs) ↓ 深度相对误差 (Rel) ↓ 平均排名 (MR) ↓ \(\Delta m\% \downarrow\)
Independent 74.01 93.16 0.0125 27.77 - -
LS 75.18 93.49 0.0155 46.77 11.25 22.60
CAGrad 75.16 93.48 0.0141 37.60 10.25 11.58
Nash-MTL 75.41 93.66 0.0129 35.02 4.75 6.82
FairGrad 75.72 93.68 0.0134 32.25 4.25 5.18
COST 75.73 93.53 0.0133 31.53 4.25 4.30
APT-CG 75.79 93.65 0.0129 38.13 5.25 9.31 ± 1.21
APT-Fair 74.29 93.25 0.0125 30.36 8.75 2.22 ± 1.46
APT-Muon 78.23 94.36 0.0121 27.09 1.00 -3.09 ± 1.16

消融实验

为了严格验证动量衰减系数 \(\beta_1\) 的取值对不同优化器和 MTL 算法交互特性的影响,作者在 Cityscapes 上固定不同静态 \(\beta_1\) 进行了系统测试(表内数值为综合性能表现):

动量设置 \(\beta_1\) CAGrad (Adam) CAGrad (Muon) FairGrad (Adam) FairGrad (Muon)
默认 (default=0.9) 12.38 ± 1.49 -2.37 ± 0.54 4.23 ± 3.20 6.04 ± 1.16
\(\beta_1 = 0.6\) 9.78 ± 1.14 1.95 ± 0.90 4.81 ± 2.77 1.63 ± 3.05
\(\beta_1 = 0.2\) 10.78 ± 3.03 81.12 ± 9.13 2.86 ± 1.13 44.29 ± 25.49
\(\beta_1 = 0.0\) (无动量) 9.11 ± 0.97 120.78 ± 36.70 2.38 ± 0.90 55.87 ± 28.48
APT (自适应调节) 9.31 ± 1.21 -3.09 ± 1.16 2.22 ± 1.46 -3.09 ± 1.16 (搭配LDP)

关键发现

  • 在传统的 Adam/AdamW 优化器中,较小的 \(\beta_1\) 能够释放 MTL 算子的即时去冲突威力(例如 FairGrad 的性能从默认动量下的 4.23 逐步优化到 \(\beta_1=0.0\) 下的 2.38),但完全去除动量会使优化轨迹产生严重抖动并增加方差;APT 自适应动量成功在保持稳定性的同时达到超越静态低动量的最佳权衡点(2.22)。
  • 与 Adam 相反,Muon 对动量缺失极为敏感:将 \(\beta_1\) 降至 0 会导致性能剧烈恶化崩溃(误差指标暴增至 120+),这印证了 Muon 正交化需要依赖连续稳定的历史动量谱空间,必须配合 APT 的平滑约束与 LDP 方向保持。
  • 在多达 40 个任务的 CelebA 人脸属性分类与 11 个任务的 QM9 分子图回归中,APT-Muon 在 CelebA 上取得了 0.03% 的极低退化率,APT-FAMO 在 QM9 上以 \(\Delta m\% = 55.5\%\) 刷新了该数据集的多任务 SOTA。

亮点与洞察

  • 首次揭示了“EMA 动量稀释多任务瞬时梯度”这一长期被忽视的底层矛盾:现有基于梯度的多任务算法名义上计算出正交非冲突梯度,实际上在与 90% 的历史动量混合后依然处于严重冲突状态。
  • 创新性地从矩阵谱分解与切比雪夫不等式角度给出了 Muon 充当“隐式多任务学习器”的严谨数学证明,论证了奇异值截断归一化对强弱任务的天然再平衡作用。
  • 提出的曲率感知自适应动量仅需通过相邻梯度内积即可完成轻量计算,即插即用且不引入二阶矩阵求逆或二次规划开销,能直接推广至大语言模型多任务指令微调或视觉多任务基干优化。

局限与展望

  • 作者指出,在任务数量极多(如 CelebA 的 40 任务)且任务间相关性极其发散时,以全局梯度均值作为中心锚点的有效性会轻微下降,可能需要引入动态任务分组机制。
  • 本文主要探讨了共享骨干结构的参数硬共享多任务设定,未深入探索软共享架构(如跨任务注意力、MoE 门控路由)与自适应优化器的交互影响。
  • 未来值得进一步探索将 APT 自适应动量与大规模多模态大模型预训练过程中的跨模态损失平衡相结合,探索 Muon 优化器在十亿级参数多任务模型中的可扩展性。

相关工作与启发

  • vs CAGrad / FairGrad: 现有工作仅专注于单步梯度空间内的冲突判定与凸优化求解,假定求解得到的方向会直接指导参数移动;本文指出了其在 EMA 优化器下的动量稀释现象,使现有冲突消解算法能够真正落地到现代优化器更新流中。
  • vs FAMO: FAMO 尝试通过一阶近似在损失层级时序平摊 MGDA 目标;本文从优化器第一动量更新机制切入,提供了更为普适的优化器级别自适应框架。
  • vs Muon 原生实现: 原生 Muon 针对单任务 LLM 隐藏层设计;本文首次将其迁移至多任务学习,挖掘出其内在的隐式多任务均衡特性并补充了高有效秩梯度的 LDP 构造机制。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次从优化器动量机制与谱分解视角审视多任务梯度失效问题,视角独到且理论扎实。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 4 大经典视觉、分类与图回归基准,跨越 2 至 40 任务规模,消融细致完备。
  • 写作质量: ⭐⭐⭐⭐⭐ 逻辑推导环环相扣,从经验现象观察到数学不等式证明再到算法设计,表达清晰流畅。
  • 价值: ⭐⭐⭐⭐⭐ 直击深度学习多任务优化落地中高级优化器与梯度算法脱节的普遍痛点,具有极高的工程与学术参考价值。