Jacobian Descent for Multi-Objective Optimization¶
会议: NeurIPS 2026
arXiv: 2406.16232
代码: https://github.com/SimplexLab/TorchJD
领域: 优化/理论
关键词: 多目标优化, 雅可比下降, 梯度投影, Pareto前沿收敛, 样本级风险最小化
一句话总结¶
论文提出了将梯度下降推广至向量目标函数的雅可比下降（JD）框架与非冲突投影聚合器 \(\mathcal{A}_{\text{UPGrad}}\)，证明了平滑凸情形下收敛至 Pareto 前沿的强保证，并利用格拉姆矩阵将显存复杂度降至 \(\mathcal{O}(n+m^2)\)，开拓了样本级风险最小化（IWRM）新范式。
研究背景与动机¶
在深度学习中，多目标优化通常采用标量化方法（如加权求和），将多个可能冲突的损失结合为一个单一标量后使用随机梯度下降（SGD）进行训练。然而，这种简单的线性加权在缺乏各目标先验重要性信息时，极易导致某些关键目标严重劣化，甚至陷入梯度方向相互抵消、优化进展停滞的窘境。虽然帕累托（Pareto）多目标优化旨在同时兼顾所有目标且不牺牲任意单一任务，但现有多任务学习与梯度组合方法（如 MGDA、PCGrad、CAGrad 等）在处理剧烈梯度冲突与模长高度不均衡时往往表现欠佳。
这种困境的核心矛盾在于：现有梯度组合算法要么仅能收敛到极其微弱的“弱 Pareto 驻点”（Weakly Pareto Stationary Point），该点在目标空间中甚至可能距离真实的 Pareto 集合任意遥远（如在简单凸二次函数上即可能陷入坐标轴边界）；要么缺乏关于梯度模长的尺度线性不变性（例如 MGDA 当任一分量梯度趋近于零时整个更新步长直接崩溃消失）；更有甚者，如 PCGrad 在目标数大于2时无法从理论上保证更新方向无冲突，且更新严重依赖梯度的遍历顺序。
本文的切入角度是：统一所有通过聚合器将雅可比矩阵映射为单一更新向量的算法，建立形式化的雅可比下降（Jacobian Descent, JD）理论框架。核心 idea：设计一种将各个目标梯度分别投影到雅可比矩阵行向量的对偶锥上并求平均的非冲突投影聚合器 \(\mathcal{A}_{\text{UPGrad}}\)，首次在平滑凸优化下证明了严格收敛至 Pareto 前沿，并结合格拉姆矩阵链式反向累加将显存开销从 \(\mathcal{O}(m\times n)\) 骤降至 \(\mathcal{O}(n+m^2)\)。
方法详解¶
整体框架¶
针对可微向量值目标函数 \(\bm{f}: \mathbb{R}^n \to \mathbb{R}^m\)，雅可比下降（JD）在每一步迭代中计算雅可比矩阵 \(\mathcal{J}\bm{f}(\bm{x}) \in \mathbb{R}^{m \times n}\)（每一行为对应目标 \(f_i\) 的转置梯度 \(\nabla f_i(\bm{x})^\top\)），并通过聚合器 \(\mathcal{A}: \mathbb{R}^{m \times n} \to \mathbb{R}^n\) 计算更新步长： $\(\bm{x}_{t+1} = \bm{x}_t - \eta \mathcal{A}\big(\mathcal{J}\bm{f}(\bm{x}_t)\big)\)$ 当目标数 \(m=1\) 且聚合器为恒等映射时，JD 精确退化为经典梯度下降（GD）。
对于大规模样本或高维网络，算法支持两种随机变体：随机估计雅可比下降（SEJD，使用批次估计梯度）与随机子雅可比下降（SSJD，随机抽取目标子集）。在多层神经网络实现中，论文提出了基于格拉姆矩阵的纯标量反向传播管线，避免了物化和存储完整高维雅可比矩阵。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["输入：网络参数 x 与多目标向量 f(x)"] --> B["格拉姆反向累加<br/>反向传播特征雅可比并累积 G = JJ^T"]
B --> C["对偶二次规划求解<br/>求解 m 个 QP 问题获得非冲突权重 w"]
C --> D["无冲突投影聚合<br/>标量损失加权反向传播得到更新方向"]
D --> E["样本级风险最小化更新<br/>执行参数步长更新 x_{t+1} = x_t - η v"]关键设计¶
1. 理想聚合器的三大良态性质：界定多目标优化的公理化准则
为了保证在强冲突环境下稳定下降，论文形式化定义了聚合器 \(\mathcal{A}\) 必须满足的三项核心数学性质： - 非冲突性（Non-conflicting）：要求对任意 \(J \in \mathbb{R}^{m \times n}\)，均有 \(\bm{0} \mathrel{≼} J \cdot \mathcal{A}(J)\)。几何上意味着聚合更新方向 \(\mathcal{A}(J)\) 必须严格落在雅可比矩阵行向量集合的对偶锥（Dual Cone）内部。只要步长 \(\eta\) 充分小，一阶泰勒展开下没有任何一个目标函数值会升高。 - 尺度线性性（Linear under scaling）：对于任意对角正定缩放矩阵 \(\text{diag}(\bm{c})\)（\(\bm{c} \succ \bm{0}\)），映射 \(\bm{c} \mapsto \mathcal{A}(\text{diag}(\bm{c}) J)\) 必须关于 \(\bm{c}\) 是线性的。该性质保证每个目标的梯度贡献严格与其模长成正比，既不会让极小梯度主导或扼杀更新，也不会让某一大模长梯度被完全忽略。 - 加权性（Weighted）：更新方向必须完全落在雅可比矩阵行向量张成的子空间内，即存在权重向量 \(\bm{w} \in \mathbb{R}^m\) 满足 \(\mathcal{A}(J) = J^\top \bm{w}\)。这保证算法不会向一阶导数恒为零的无效正交方向浪费更新能量。
论文严密证明了现有多任务学习聚合器（MGDA、DualProj、PCGrad、CAGrad、Aligned-MTL 等）均无法同时满足这三项性质，而所提方法是首个兼具三者的聚合器。
2. 无冲突梯度投影聚合器 \(\mathcal{A}_{\text{UPGrad}}\)：对偶锥正交投影及其二次规划对偶
聚合器 \(\mathcal{A}_{\text{UPGrad}}\) 的几何直觉是将每个目标的梯度向量单独正交投影到所有梯度构成的对偶锥上，再对这 \(m\) 个投影向量取均值： $\(\mathcal{A}_{\text{UPGrad}}(J) = \frac{1}{m} \sum_{i \in [m]} \pi_J(J^\top \bm{e}_i)\)$ 其中 \(\pi_J(\bm{x}) = \arg\min_{\bm{y} \in \mathbb{R}^n: \bm{0} \mathrel{≼} J\bm{y}} \|\bm{y} - \bm{x}\|^2\)。由于凸锥在正线性组合下封闭，均值必然严格位于对偶锥内，天然满足非冲突性；由于对偶锥投影算子满足严格的正齐次性，进一步确保了尺度线性性。
当参数维度 \(n\) 极大（如神经网络具有数百万参数）时，在原空间求解投影不可行。论文通过 Lagrange 对偶理论推导出其完全与参数维度 \(n\) 无关的对偶二次规划（QP）形式： $\(\pi_J(J^\top \bm{u}) = J^\top \bm{w}, \quad \text{其中 } \bm{w} \in \arg\min_{\bm{v} \in \mathbb{R}^m: \bm{u} \mathrel{≼} \bm{v}} \bm{v}^\top (J J^\top) \bm{v}\)$ 只需分别令 \(\bm{u} = \bm{e}_i\)（\(i=1,\dots,m\)），求解 \(m\) 个 \(m\) 维的简单凸二次规划问题，即可得到闭式加权形式： $\(\mathcal{A}_{\text{UPGrad}}(J) = J^\top \left( \frac{1}{m} \sum_{i=1}^m \bm{w}_i \right)\)$
3. 平滑凸多目标优化向 Pareto 前沿收敛定理：超越弱驻点的坚实理论保证
针对 \(\beta\)-光滑且 \(\mathrel{≼}\)-凸的向量目标函数 \(\bm{f}\)，在 Pareto 前沿紧致且任意点均被 Pareto 最优解支配的标准假设下，设步长 \(\eta = \frac{1}{\beta \sqrt{m}}\)。若更新权重序列 \(\bm{w}_t\) 有界，定理 1 保证序列 \(\bm{f}(\bm{x}_t)\) 必收敛至某个 Pareto 最优解 \(\bm{f}(\bm{x}^*)\)，即： $\(\lim_{t \to \infty} \bm{f}(\bm{x}_t) \in \bm{f}(X^*)\)$ 相比文献中广泛采用的弱 Pareto 驻点（即仅存在非零非负权重使梯度的凸组合为零，容易停滞在非最优边界），收敛至真实 Pareto 前沿构成了实质性增强。进一步推导还给出了平均收敛速度界： $\(\frac{1}{T} \sum_{t=1}^T \bm{w}_t^\top \big(\bm{f}(\bm{x}_t) - \bm{f}(\bm{x}^*)\big) \le \frac{\sqrt{m}}{T} \left( \|\bm{f}(\bm{x}_1) - \bm{f}(\bm{x}^*)\| + \frac{\beta}{2} \|\bm{x}_1 - \bm{x}^*\|^2 \right) = \mathcal{O}\left(\frac{1}{T}\right)\)$
4. 格拉姆反向累加（Gramian Reverse Accumulation）：打破 \(m \times n\) 内存屏障
存储完整雅可比矩阵需要 \(\mathcal{O}(m \times n)\) 显存，这在深度神经网络中通常因显存爆炸而无法运行。作者发现，任何加权聚合器的权重计算实际仅依赖内积矩阵——即半正定的格拉姆矩阵 \(G = J J^\top \in \mathbb{R}^{m \times m}\)。一阶泰勒展开表明： $\(\bm{f}(\bm{x} + \bm{y}) = \bm{f}(\bm{x}) - \eta \mathcal{G}\bm{f}(\bm{x}) \cdot \bm{w} + o\left(\eta \sqrt{\bm{w}^\top \mathcal{G}\bm{f}(\bm{x}) \bm{w}}\right)\)$ 且梯度更新可转化为标量函数的单次反向传播：\(\mathcal{A}(\mathcal{J}\bm{f}(\bm{x})) = \nabla \big( \bm{w}^\top \bm{f} \big)(\bm{x})\)。
基于多层复合函数的格拉姆链式法则与可加性： $\(\mathcal{G}(\bm{f} \circ \bm{g})(\bm{x}) = \mathcal{J}\bm{f}(\bm{g}(\bm{x})) \cdot \mathcal{G}\bm{g}(\bm{x}) \cdot \mathcal{J}\bm{f}(\bm{g}(\bm{x}))^\top, \quad \mathcal{G}\bm{f}(\bm{x}_1, \dots, \bm{x}_k) = \sum_{i=1}^k \mathcal{G}_{\bm{x}_i}\bm{f}\)$ 算法仅需自顶向下反向传播输出关于各层激活值的雅可比切片 \(J_x \in \mathbb{R}^{m \times n_i}\)，在每一层计算该层参数的局部雅可比 \(J_p = J_x \cdot \mathcal{J}_{\bm{p}_i}\bm{f}_i\)，并就地将其格拉姆外积累加至 \(G \leftarrow G + J_p J_p^\top\)，随后立即释放 \(J_p\)。显存开销从 \(\mathcal{O}(m \times n)\) 彻底锐减至 \(\mathcal{O}(n + m^2)\)。
损失函数 / 训练策略¶
在样本级风险最小化（Instance-Wise Risk Minimization, IWRM）中，训练集每个样本的交叉熵损失被建模为独立的单一目标： $\(\bm{f}(\bm{\theta}) = [f_1(\bm{\theta}), f_2(\bm{\theta}), \dots, f_m(\bm{\theta})]^\top, \quad f_i(\bm{\theta}) = \mathcal{L}_{\text{CE}}(\bm{g}(\bm{x}_i; \bm{\theta}), \bm{y}_i)\)$ 在实际训练中，采用随机子雅可比下降（SSJD），每个迭代步抽取批大小为 \(B\)（例如 \(B=32\)）的样本计算对应的子雅可比矩阵。当使用平均聚合器 \(\mathcal{A}_{\text{Mean}}\) 时，算法精确退化为标准经验风险最小化（ERM）配合 SGD。实验采用恒定学习率策略（通过精密网格与细分搜索确定每个算法的最优固定学习率），并在附录中验证了与 Adam 动量更新的良好兼容性。
实验关键数据¶
主实验¶
主实验在 6 个图像分类数据集的子集（固定 1024 张图像以平衡总计 43,776 次训练的庞大计算预算）上，评估了 CNN 架构在 IWRM 设定下不同聚合器的测试集分类准确率（8 次独立随机试验取平均）：
| 聚合器 (Aggregator) | SVHN (%) | CIFAR-10 (%) | EuroSAT (%) | MNIST (%) | Fashion-MNIST (%) | Kuzushiji-MNIST (%) |
|---|---|---|---|---|---|---|
| \(\mathcal{A}_{\text{Mean}}\) (SGD 基线) | 69.6 | 47.7 | 68.5 | 96.3 | 80.7 | 76.9 |
| \(\mathcal{A}_{\text{UPGrad}}\) (本文方法) | 69.7 | 48.1 | 73.0 | 96.0 | 80.8 | 77.1 |
| \(\mathcal{A}_{\text{DualProj}}\) | 69.3 | 47.2 | 72.7 | 96.1 | 80.8 | 76.4 |
| \(\mathcal{A}_{\text{GradDrop}}\) | 68.9 | 46.8 | 71.9 | 96.3 | 80.9 | 76.9 |
| \(\mathcal{A}_{\text{PCGrad}}\) | 49.1 | 48.1 | 59.3 | 85.9 | 69.3 | 54.7 |
| \(\mathcal{A}_{\text{RGW}}\) | 66.2 | 39.5 | 67.9 | 94.0 | 77.1 | 73.6 |
| \(\mathcal{A}_{\text{Aligned-MTL}}\) | 61.3 | 47.3 | 65.3 | 88.3 | 78.5 | 67.9 |
| \(\mathcal{A}_{\text{CAGrad}}\) | 60.7 | 42.4 | 68.1 | 90.3 | 78.1 | 65.3 |
| \(\mathcal{A}_{\text{IMTL-G}}\) | 58.7 | 44.8 | 65.2 | 87.2 | 75.0 | 65.9 |
| \(\mathcal{A}_{\text{MGDA}}\) | 41.7 | 32.9 | 62.5 | 85.2 | 73.4 | 53.7 |
| \(\mathcal{A}_{\text{Nash-MTL}}\) | 39.7 | 31.3 | 60.8 | 87.6 | 72.0 | 52.7 |
在 EuroSAT 数据集上，\(\mathcal{A}_{\text{UPGrad}}\) 相比标准 SGD 基线展现出了显著优势（73.0% vs 68.5%，绝对提升 +4.5%）。
显存与计算时间消融/性能对比¶
在批大小为 32 的设置下，对比标准 ERM (SGD)、经典雅可比下降 (SSJD) 以及本文提出的格拉姆反向累加雅可比下降 (Gramian-based SSJD) 的峰值显存占用与单轮训练耗时（NVIDIA L4 GPU）：
| 训练目标 / 实现配置 | SVHN 显存 | CIFAR-10 显存 | EuroSAT 显存 | MNIST 显存 | CIFAR-10 单轮耗时 | 耗时相对原生SSJD降低 |
|---|---|---|---|---|---|---|
| ERM (SGD) | 32.0 MB | 44.6 MB | 72.6 MB | 54.2 MB | 0.44 s | - |
| IWRM (SSJD - 原生实现) | 315.4 MB | 619.3 MB | 1447.5 MB | 786.0 MB | 1.78 s | - |
| IWRM (SSJD - Gramian累加) | 39.8 MB | 75.3 MB | 104.2 MB | 65.6 MB | 1.12 s | 37.1% (提速 59%) |
格拉姆累加实现使峰值显存相较于原生雅可比存储减少了 87% ~ 93%，几乎压缩至与标准 SGD 处于同一量级。
关键发现¶
- 训练收敛速度反常反超基线：尽管 \(\mathcal{A}_{\text{Mean}}\) 直接对标量经验风险进行优化，但在所有数据集的训练损失曲线上，\(\mathcal{A}_{\text{UPGrad}}\) 均展现出更快的训练集损失下降速率。分析更新方向与 SGD 梯度的余弦相似度发现，在训练初期两者相似度偏低，表明 \(\mathcal{A}_{\text{UPGrad}}\) 成功阻止了困难样本的微弱梯度被简单样本的大梯度掩盖，从而提早拟合了高难度样本特征。
- 批大小（Batch Size）的对偶锥紧致效应：在 SSJD 中增大批大小不仅能降低方差，更能向对偶锥引入更多行向量约束，使对偶锥更加收紧，从而使每一次参数更新在更全局的多目标空间内保持无冲突。
- 尺度非线性导致传统算法坍塌：\(\mathcal{A}_{\text{MGDA}}\) 在几乎所有数据集上性能惨烈垫底（如 CIFAR-10 仅 32.9%），直接印证了其不满足尺度线性性的理论缺陷——一旦某一样本或目标的梯度范数接近零，MGDA 的聚合步长将趋于停滞；而 PCGrad 由于投影顺序的随机性以及高维下无法杜绝冲突，表现出明显的优化不稳定。
亮点与洞察¶
- 格拉姆链式法则打破内存诅咒：巧妙地将聚合更新表示为对内积矩阵的函数，并推导出仅依赖激活值雅可比与层局部格拉姆矩阵就地累加的链式法则，用数学等价变换将原本无法实用化的 \(\mathcal{O}(m \times n)\) 显存降至与 SGD 几乎相当的 \(\mathcal{O}(n + m^2)\)。
- IWRM 学习范式革新：传统观点通常将分类任务视为单一目标优化，而本文通过将样本损失视作相互制约的向量维度，展示了将单目标问题升维至多目标流形进行抗冲突求解可带来泛化与收敛双重收益。
- 对偶二次规划的高效求解转化：将高维参数空间中的对偶锥几何投影，严格转化为 \(m\) 个低维二次规划（维度仅为样本批大小 \(m\)），规避了在网络参数空间做凸集投影的高维灾难。
局限与展望¶
- 经验验证受限于小规模设定：由于网格搜索最优学习率需要完成数万次独立训练，当前实验仅在每个数据集 1024 张图像的微型子集和浅层 CNN 上进行，尚未在大规模数据集（如完整 ImageNet）和现代骨干模型（如 Transformer/LLM）上进行完整验证。
- 批内样本块对角稀疏性尚未利用：在 IWRM 中，样本损失对网络激活值的导数呈现天然的块对角稀疏结构，当前的通用反向传播实现未对此进行算子级稀疏加速，仍存在计算延迟。
- QP 求解器的硬件瓶颈：每个迭代步求解 \(m\) 个 QP 问题主要依赖 CPU 算法库（qpsolvers），尚未完全移至 GPU 端并行，当批大小 \(m\) 扩大时求解耗时会成为新的系统瓶颈。
相关工作与启发¶
- vs MGDA (Désidéri, 2012)：MGDA 通过极小化最差目标提升寻找下降方向，虽然满足非冲突性，但缺少尺度线性性，极易收敛至无意义的弱 Pareto 驻点并在小梯度处早停；本文的 \(\mathcal{A}_{\text{UPGrad}}\) 保持尺度线性且严格收敛到 Pareto 前沿。
- vs PCGrad (Yu et al., 2020)：PCGrad 采用贪心、随机顺序的成对正交投影，当目标数 \(m > 2\) 时数学上无法消除冲突；本文的 \(\mathcal{A}_{\text{UPGrad}}\) 联合投影到全局对偶锥，在任意维度下均严格保证非冲突。
- vs CAGrad (Liu et al., 2021a)：CAGrad 在平均梯度周围设定半径为 \(c\) 的局部球内权衡，当 \(c < 1\) 时不具备非冲突性；\(\mathcal{A}_{\text{UPGrad}}\) 则无须手工调节保守性超参数 \(c\)。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次系统性形式化雅可比下降，提出首个收敛至 Pareto 前沿的无冲突聚合器与格拉姆反向累加架构。
- 实验充分度: ⭐⭐⭐⭐ 理论证明严丝合缝，但深度学习实证部分受限于微型数据集规模，稍显不足。
- 写作质量: ⭐⭐⭐⭐⭐ 数学推导严谨流畅，对偶形式与几何直觉结合紧密，附录证明完备。
- 价值: ⭐⭐⭐⭐⭐ 开源的高性能 PyTorch 库 TorchJD 与 IWRM 学习范式，为多目标学习乃至基础模型训练提供了极具潜力的底层工具。