跳转至

MAdam: Metric-Aware Multi-Objective Adam

会议: NeurIPS 2026
arXiv: 2606.03904
论文: https://arxiv.org/abs/2606.03904
代码: https://github.com/lfb-1/MAdam
领域: 优化/理论
关键词: 多目标优化, Adam优化器, Fisher信息矩阵, 度量感知预条件, 帕累托多任务学习

一句话总结

揭示多目标优化解算器与 Adam 默认耦合导致的权重边际化与欧氏几何畸变两大失配,提出即插即用的度量感知封装器 MAdam,通过在线估计跨目标 Fisher 矩阵在 Adam 更新前对调和方向实施偏好条件化白化,使更新步遵循正确流形度量,在多任务学习、Pareto 前沿拟合、PINN 与医学图像超分/分割上全面超越 Adam。

研究背景与动机

多目标优化(Multi-Objective Optimization, MOO)在现代机器学习中扮演着基础支柱角色,广泛涵盖多任务表征学习、物理信息神经网络(PINN)多尺度损失平衡以及医学图像多保真度联合重建。为了协调相互竞争的多目标,学术界发展出三大主流求解范式:自适应调整目标权重的损失平衡方法(如 DWA、FAMO)、显式对梯度求交投影消除几何冲突的梯度平衡方法(如 PCGrad、CAGrad),以及沿超曲面追踪折衷前沿的 Pareto 引导方法(如 PaMaL、PaLoRA)。尽管各族求解器在设计原理上各有侧重,但在最终执行参数更新时,均无一例外地将重构协调后的单一下降方向交付给通用自适应优化器 Adam 执行。

然而,梯度下降的本质由构造下降方向与沿特定黎曼度量执行步长共同决定。Adam 默认的二阶动量指数移动平均(EMA)在平稳性假设下实质上逼近标量目标的经验 Fisher 信息矩阵对角线,但这与 MOO 求解器的设计初衷存在深刻的内在冲突。首先是权重失配:在时变偏好 \(\bm{\lambda}^{(t)}\) 驱动下,Adam 的二阶动量将瞬时偏好与各任务梯度统计量强行杂糅在同一个时间窗口内,导致瞬时偏好被历史平均“边际化”稀释,使原本意在探索多样化 Pareto 前沿的轨迹退化塌缩为均匀混合状态。其次是几何失配:MOO 求解器大多基于标准欧几里得内积空间进行冲突投影或线性标量化,而 Adam 施加的对角 RMS 变换扭曲了这一度量,在非凸或各向异性损失地形中将原本一致协作的任务方向人为翻转为伪冲突,甚至引发特定任务损失不降反升。

这种两难困境表明,单纯在求解器侧修正下降方向却放任 Adam 强加未经校准的统计度量,注定无法达成预期的优化目标。本文的切入点是回归标量化多目标的内在曲率几何,求解在瞬时激活偏好下真正支配更新步的度量结构。核心 idea:将求解器调和方向视作标量化目标的欧氏梯度,利用独立维护的跨任务梯度二阶矩在线合成瞬时偏好条件化的对角 Fisher 矩阵,以此对调和方向进行白化预处理;在该白化输入下 Adam 二阶矩退化为单位矩阵,使最终参数更新严格受偏好条件化黎曼度量引导。

方法详解

整体框架

在多目标优化问题中,共享模型参数 \(\bm{\theta} \in \mathbb{R}^d\) 同时优化 \(C\) 个目标损失 \(\{\ell_i(\bm{\theta})\}_{i=1}^C\)。在每一步迭代 \(t\),MOO 求解器根据当前偏好向量 \(\bm{\lambda}^{(t)}\) 或任务间几何关系输出调和方向 \(\mathbf{d}^{(t)} = \sum_{i=1}^C \lambda_i^{(t)} \mathbf{g}_i^{(t)}\)(其中 \(\mathbf{g}_i^{(t)} = \nabla_{\bm{\theta}} \ell_i(\bm{\theta}^{(t)})\) 为各任务梯度)。MAdam 作为不侵入底层求解器与优化器核心逻辑的即插即用封装器,通过三步闭环实现度量感知更新:首先通过随机成对采样以常数级反向传播开销增量维护跨目标 Fisher 矩阵 \(\widehat{\mathbf{F}}_{ij}\);接着根据瞬时偏好 \(\bm{\lambda}^{(t)}\) 组合出条件化曲率 \(\widehat{\mathbf{C}}_{\bm{\lambda}}^{(t)}\) 并加权升温构造预条件对角阵 \(\widehat{\mathbf{M}}_{\bm{\lambda}}^{(t)}\);最后对调和方向执行前置白化逆变换 \(\tilde{\mathbf{d}}^{(t)} = (\widehat{\mathbf{M}}_{\bm{\lambda}}^{(t)})^{-1} \mathbf{d}^{(t)}\),并直接交由标准 Adam 执行更新。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["多目标输入与瞬时偏好<br/>{l_i(θ)}, 偏好向量 λ^(t) ∈ Δ^(C-1)"] --> B["随机成对采样<br/>均匀抽取任务对 (i, j) 计算梯队 g_i, g_j"]
    B --> C["跨任务 Fisher 统计维护<br/>EMA 增量更新 F_ii, F_jj, F_ij 交叉协方差"]
    C --> D["偏好条件化度量合成<br/>显式加权求和 C_λ 并结合升温因子构造 M_λ"]
    D --> E["调和方向白化预条件<br/>d~^(t) = (M_λ)^(-1) d^(t)"]
    E --> F["标准 Adam 引擎执行<br/>二阶矩输入白化收缩至单位阵,输出度量感知更新步"]

关键设计

1. 求解器与 Adam 耦合失配机制解构:定位权重边际化与几何伪冲突两大根源

传统 MOO 流水线直接将调和梯度 \(\mathbf{d}^{(t)}\) 喂入 Adam。对于时变偏好 \(\bm{\lambda}^{(t)}\),Adam 的二阶矩估计 \(\hat{\mathbf{v}}^{(t)} \approx \mathbb{E}[\mathbf{d}^{(t)} \odot \mathbf{d}^{(t)}] = \sum_{i,j=1}^C \mathbb{E}[\lambda_i \lambda_j \mathbf{g}_i \odot \mathbf{g}_j]\)。由于 Adam 无法解耦权重与梯度的联合期望,历史偏好被粗暴平滑为静态期望 \(\mathbb{E}[\lambda_i \lambda_j]\),导致当前步的特定偏好在分母处失效,破坏了探索 Pareto 前沿所需的方向特异性。更严重的是几何失配:在一阶泰勒展开下,第 \(i\) 个任务的实际损失变化量由双线性型决定:

\[\Delta \ell_i^{(t)} \approx -\eta \, \mathbf{g}_i^{(t)\top} \bigl(\mathbf{M}_{\mathrm{Adam}}^{(t)}\bigr)^{-1} \hat{\mathbf{m}}^{(t)}\]

当目标数 \(C > 1\) 时,由于各任务梯度 \(\mathbf{g}_i^{(t)}\) 与综合方向 \(\mathbf{d}^{(t)}\) 不共线,对角缩放阵 \((\mathbf{M}_{\mathrm{Adam}}^{(t)})^{-1} = \operatorname{Diag}(1 / (\sqrt{\hat{\mathbf{v}}^{(t)}} + \varepsilon))\) 的坐标非均匀拉伸可能直接改变内积符号,使得原本在欧氏几何下与 \(\mathbf{d}^{(t)}\) 呈锐角(协同下降)的任务方向在更新后内积转为负数,诱发严重的负迁移和伪对抗。

2. 偏好条件化对角 Fisher 曲率推导:瞬时偏好与跨任务相关性的显式解耦

为消除时间窗口平滑引起的权重钝化,MAdam 在理论上推导了与瞬时标量化目标严格对应的对角 Fisher 信息矩阵 \(\mathbf{C}_{\bm{\lambda}^{(t)}}\)。不同于 Adam 维护单一下降向量的外积期望,MAdam 将当前的确定性偏好权值显式提至期望符号之外,将曲率正交分解为各任务自曲率与跨任务互协方差两大部分:

\[\mathbf{C}_{\bm{\lambda}^{(t)}} = \sum_{i,j=1}^C \lambda_i^{(t)} \lambda_j^{(t)} \mathbf{F}_{ij} = \sum_{i=1}^C (\lambda_i^{(t)})^2 \mathbf{F}_{ii} + \sum_{i \neq j} \lambda_i^{(t)} \lambda_j^{(t)} \mathbf{F}_{ij}\]

其中对角块 \(\mathbf{F}_{ii} = \mathbb{E}[\mathbf{g}_i \odot \mathbf{g}_i]\) 表征单一目标的坐标尺度,而非对角块 \(\mathbf{F}_{ij} = \mathbb{E}[\mathbf{g}_i \odot \mathbf{g}_j]\) 精确量化了第 \(i\) 与第 \(j\) 个目标在各参数坐标上的协作(正值)或冲突(负值)统计。当偏好 \(\bm{\lambda}^{(t)}\) 动态偏转时,该分解确保度量张量立即响应当前折衷需求,保留各方向上真实的曲率弹性。

3. 前置调和方向白化预条件:以不变应万变的非侵入式更新控制

为在不破坏现有深度学习框架和底层优化器实现的前提下落实该度量,MAdam 采用前置变换策略,直接对 MOO 求解器交付的调和方向执行坐标归一化预条件处理:\(\tilde{\mathbf{d}}^{(t)} = (\widehat{\mathbf{M}}_{\bm{\lambda}}^{(t)})^{-1} \mathbf{d}^{(t)}\),其中度量矩阵定义为 \(\widehat{\mathbf{M}}_{\bm{\lambda}}^{(t)} = \operatorname{Diag}(\alpha^{(t)} \sqrt{\widehat{\mathbf{C}}_{\bm{\lambda}}^{(t)}} + \varepsilon + (1 - \alpha^{(t)}))\)。当该白化向量送入标准 Adam 后,其二阶矩统计量天然收缩至单位阵:

\[\hat{\mathbf{v}}^{(t)} \approx \mathbb{E}[\tilde{\mathbf{d}}^{(t)} \odot \tilde{\mathbf{d}}^{(t)}] = (\mathbf{M}_{\bm{\lambda}}^{(t)})^{-2} \mathbf{C}_{\bm{\lambda}}^{(t)} \approx \mathbf{1} \implies \mathbf{M}_{\mathrm{Adam}}^{(t)} \approx \mathbf{I}\]

此时 Adam 的自适应分母在数学上退化为恒等变换,而一阶动量则平滑跟踪预条件方向 \(\hat{\mathbf{m}}^{(t)} \approx (\mathbf{M}_{\bm{\lambda}}^{(t)})^{-1} \mathbf{d}^{(t)}\)。于是实际参数更新步精确收敛为 \(\Delta \bm{\theta}^{(t)} = -\eta (\mathbf{M}_{\bm{\lambda}}^{(t)})^{-1} \mathbf{d}^{(t)}\),一举消除了 Adam 原始分母与一阶动量之间的非对齐,使全局更新完全由偏好条件化的自然梯度黎曼度量主导。

4. 随机成对采样与升温调制:从二次复杂度到常数级开销的工程落地

精确计算全部 \(C(C+1)/2\) 个对称 Fisher 矩阵块原本需要每步执行 \(C\) 次完整的反向传播,带来无法接受的算力开销。MAdam 设计了随机成对采样机制:在每步迭代中均匀随机抽取任务对 \((i, j) \sim \text{Uniform}([C] \times [C])\),仅计算这两个特定任务的梯度 \(\mathbf{g}_i^{(t)}\) 与 \(\mathbf{g}_j^{(t)}\),并单步更新 \(\widehat{\mathbf{F}}_{ii}\)、\(\widehat{\mathbf{F}}_{jj}\) 和 \(\widehat{\mathbf{F}}_{ij}\) 这 3 个 EMA 缓存项,其余矩阵块保持原值。该机制将每步反向传播次数严格限制为 2 次,复杂度降至 \(O(1)\),无视任务总数 \(C\) 的增长。同时,引入单调递增的升温系数 \(\alpha^{(t)} \in [0, 1]\),在训练初期 Fisher 统计量方差较大时平滑回退为标准 Adam,待协方差 EMA 稳定后无缝过渡至完整度量感知更新。

损失函数 / 训练策略

MAdam 保持底层各任务损失形式不变。对包含 \(C\) 个子目标的网络,每步以统一动量超参(复用 Adam 的 \(\beta_2=0.999\))维护 Fisher 矩阵缓存,初始阶段设定 warmup 步数完成 \(\alpha^{(t)}\) 从 0 到 1 的线性爬坡。调和梯度的白化操作在 PyTorch 中通过逐元素向量除法完成,整体开销主要体现为额外维护 \(C(C+1)/2\) 份共享参数尺寸的对角浮点张量。

实验关键数据

主实验

在多任务学习(MTL)经典基准 MultiMNIST 与 UTKFace 上,MAdam 广泛叠加于线性标量化(LS)、动态权重平均(DWA)、不确定度加权(UW)、公平多任务(IMTL)及冲突规避梯度(CAGrad)等各族解算器,均取得一致的精度提升。在 7 关节机器人逆动力学回归基准 SARCOS 上,MAdam 在所有解算器下均显著降低了多任务均方误差。

数据集 求解器基线 基础模型指标 + MAdam (本文) 提升幅度 / 相对变化
MultiMNIST (C=2) LS (线性标量化) 平均精度 94.08 ± 0.37% 95.11 ± 0.32% +1.03%
MultiMNIST (C=2) DWA 平均精度 94.08 ± 0.34% 95.08 ± 0.25% +1.00%
MultiMNIST (C=2) UW (不确定度加权) 平均精度 94.34 ± 0.31% 95.14 ± 0.44% +0.80%
MultiMNIST (C=2) IMTL 平均精度 93.95 ± 0.29% 95.12 ± 0.25% +1.17%
MultiMNIST (C=2) CAGrad 平均精度 94.05 ± 0.24% 95.10 ± 0.35% +1.05%
SARCOS (C=7) LS (线性标量化) 平均 MSE (×100) 10.99 ± 0.59 10.00 ± 0.07 相对降低 9.0%
SARCOS (C=7) DWA 平均 MSE (×100) 11.11 ± 0.65 10.09 ± 0.13 相对降低 9.2%
SARCOS (C=7) UW (不确定度加权) 平均 MSE (×100) 10.83 ± 0.71 9.92 ± 0.31 相对降低 8.4%
SARCOS (C=7) IMTL 平均 MSE (×100) 12.10 ± 0.29 11.05 ± 0.39 相对降低 8.7%
SARCOS (C=7) CAGrad 平均 MSE (×100) 10.98 ± 0.63 10.38 ± 0.43 相对降低 5.5%

在 Pareto 前沿搜索任务(Pareto MTL)中,针对时变偏好引导的 PaMaL 与 PaLoRA 模型,MAdam 在 MultiMNIST 上将超体积(Hypervolume, HV)分别从 0.8869 提升至 0.9086、从 0.9016 提升至 0.9177;在 7 目标 SARCOS 前沿恢复中,PaLoRA 的超体积从 431.21 提升至 442.17,全前沿平均误差由 16.62 压降至 13.91。

消融实验

在 SARCOS 数据集(7 目标,LS 标量化)上针对 MAdam 的核心设计组件展开剥离实验,验证各模块的不可或缺性:

配置 平均误差 (MSE×100) ↓ 关键设计验证结论
LS + MAdam (完整方法) 10.00 ± 0.07 具备完整的偏好条件化 Fisher 估计与平稳升温,方差极低
去掉 Fisher 升温机制 (w/o FIM rampup) 11.74 ± 0.54 性能急剧恶化并劣于基线,早期未收敛的高方差 Fisher 导致破坏性更新
去掉随机采样/全量二次遍历 (w/o stochastic sampling) 10.33 ± 0.61 每步显式计算所有 \(O(C^2)\) 项,收敛精度与采样相当,证明随机成对估计的无偏与稳健
去掉非对角跨任务块 (w/o off-diagonal \(\mathbf{F}_{ij}\)) 10.47 ± 0.70 仅使用各任务对角 Fisher \(\mathbf{F}_{ii}\),忽略任务间冲突/协同项导致性能明显回落
原始 LS 基线 (Vanilla Adam) 10.99 ± 0.59 承受 Adam 历史加权平均与对角拉伸引起的几何失真

关键发现

  • 早期升温(Rampup)是稳定二阶度量的生命线:消融数据显示,若在训练伊始直接启用未经 EMA 充分平滑的 Fisher 逆变换,误差直接飙升至 11.74(甚至劣于 Adam 原生的 10.99)。只有借助 \(\alpha^{(t)}\) 将前置度量从恒等单位阵平滑过渡到 Fisher 矩阵,才能彻底规避初期的梯度扰动。
  • 跨目标非对角协方差承载关键冲突几何:单纯使用任务内部方差 \(\mathbf{F}_{ii}\) 会损失近一半的性能增益(10.00 退化到 10.47),证明多任务优化中坐标级的正负协同项(\(\mathbf{g}_i \odot \mathbf{g}_j\))对矫正搜索流形至关重要。
  • 突破 Pareto 前沿塌缩瓶颈:在连续偏好交替(\(\bm{\lambda}_A \leftrightarrow \bm{\lambda}_B\))多阶段模拟与多任务实测中,原生 Adam 会被历史动量死锁在折衷中心,而 MAdam 凭借瞬时重构的曲率矩阵实现了近乎瞬时的响应,大幅拉宽了 Pareto 折衷边界的多样性。

亮点与洞察

  • 解耦“方向求解”与“流形度量”:绝大多数 MOO 研究致力于改进梯度投影和标量化公式(即方向 \(\mathbf{d}\)),却忽视了底层优化器 Adam 正在通过对角二阶矩悄悄施加一个与求解器假设完全背道而驰的各向异性度量。MAdam 从优化度量本源切入,指出了多目标优化领域长久忽视的基础漏洞。
  • 数学机制上使 Adam 二阶矩自然湮灭为恒等阵:巧妙利用调和梯度外积与 Fisher 矩阵之间的内在联系,通过前置乘以 \(\mathbf{M}_{\bm{\lambda}}^{-1}\) 使得送入 Adam 的输入特征被精确白化,倒逼 Adam 内部有偏的二阶矩坍缩为单位矩阵 \(\mathbf{I}\),在完全不修改 Adam 源码的前提下实现了对优化步度量的全权接管。
  • 常数级计算开销的随机采样近似:没有陷入二次复杂度 \(O(C^2)\) 梯度反向传播的泥潭,而是通过随机抽取单对任务维护历史二阶动量池,使得 MAdam 在单步耗时仅微增 12%~19% 的轻微代价下,在同等 FLOPs 预算内反而比原生 Adam 更快达到预定收敛精度。

局限与展望

  • 参数显存开销随任务数二次方增长:算法必须显式常驻保存 \(C(C+1)/2\) 个与模型共享参数同尺度的张量。在参数量达到数十亿的大语言模型或海量超多任务(如 \(C > 50\))场景下,显存空间将成为严峻瓶颈。
  • 对极度非平稳的快速振荡偏好仍存滞后:虽然显式解耦了偏好因子 \(\bm{\lambda}^{(t)}\),但底层 Fisher 块 \(\widehat{\mathbf{F}}_{ij}\) 依然依赖滑动平均参数 \(\beta_2=0.999\) 进行平滑估算,当损失函数几何形状本身随参数迭代发生极快畸变时,EMA 可能存在轻微的曲率估计延迟。
  • 未来改进方向:可进一步探索低秩近似(Low-Rank Factorization)或参数分块绑定策略压缩跨目标 Fisher 存储张量,或将度量感知思想拓展至 Muon 等新兴基于正交化/矩阵动量的优化算法中。

相关工作与启发

  • vs MTAdam / AdaTask 等专用多任务优化器:MTAdam 等方法试图针对每个任务定制专有的学习率或自适应调整动量,而 MAdam 属于更高维度的通用封装器,不仅能直接包裹任何传统 MOO 求解器,还从理论上根除了欧氏空间求解与 RMS 空间执行的几何失配。
  • vs K-FAC / Shampoo / Sophia 等二阶/自然梯度优化器:此类方法专为单目标设计,旨在通过近似逆曲率加速全局收敛,无法感知多目标间时变的偏好权重 \(\bm{\lambda}^{(t)}\);若直接应用于多目标,会面临与 Adam 相同的偏好历史边际化问题。MAdam 填补了时变偏好多目标与曲率感知优化之间的空白。
  • vs APT (Zhou et al., 2026):同期工作 APT 主要聚焦于一阶动量 \(\beta_1\) 在多任务中的历史惯性阻滞问题;而 MAdam 则靶向二阶矩 \(\beta_2\) 造成的度量畸变与权重吸收,两者分别从一阶方向保真与二阶度量对齐两个正交维度切入,具有极佳的互补潜力。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 深刻揭示了 MOO 求解器与 Adam 耦合中普遍存在的权重边际化与几何失配机理,理论推导严密且优雅。
  • 实验充分度: ⭐⭐⭐⭐⭐ 覆盖经典分类/回归多任务、Pareto 前沿多目标、20 种高维偏微分方程 PINN 以及脑部 MRI 超分/病灶分割等多个领域,证据扎实详尽。
  • 写作质量: ⭐⭐⭐⭐⭐ 逻辑推进行云流水,从问题揭示、理论推导、工程实现到实验消融层层递进,数学表达极其严谨自洽。
  • 价值: ⭐⭐⭐⭐⭐ 作为零改动嵌入即可带来显著收益的通用优化器封装器,对解决多任务深度学习与复合损失训练中的负迁移现象具有极高的实用价值。