Fast Gauss-Newton for Multiclass Cross-Entropy¶
会议: NeurIPS 2026
arXiv: 2605.06081
领域: 优化/理论
关键词: Gauss-Newton 优化, 多分类交叉熵, 真实与其余边际, 共轭梯度法, 行空间求解器
一句话总结¶
针对多分类 Softmax 交叉熵广义高斯牛顿（GGN）曲率在类别数增大时耦合严重的计算瓶颈，本文提出 Fast Gauss-Newton（FGN），通过严谨的代数分解将 GGN 拆分为秩一的“真实对其余”边际项与竞争类间协方差项，在保留前者并精确维持原损失与梯度的前提下，导出了规模仅取决于批大小的无矩阵共轭梯度行空间求解器。
研究背景与动机¶
在深度学习和机器学习的二阶优化方法中，曲率模型的结构与计算复杂度直接决定了算法的可扩展性。对于分类任务中的 Softmax 交叉熵损失，广义高斯牛顿（GGN）矩阵因其半正定性（PSD）和对 Hessian 矩阵的优良近似，成为了自然梯度法、随机高斯牛顿法（SGN）以及二阶预条件子设计的核心支柱。然而，标准多分类 Softmax GGN 的输出空间曲率继承了全类别 Softmax 协方差结构 \(\mathrm{diag}(\bm{p}) - \bm{p}\bm{p}^\top\)。在无矩阵（matrix-free）隐式曲率向量积计算中，该协方差项将所有类别的 Logit 紧密耦合在一起，导致曲率计算的开销随类别数 \(C\) 的增长而急剧增加，尤其在数千类的大规模分类或词表场景下成为严重的性能瓶颈。
现有优化文献在面对这一挑战时，大多将焦点放在参数空间的曲率近似上（例如 K-FAC 的分层 Kronecker 分解、Shampoo 的张量分解，或 AdaHessian、Sophia 的对角/低秩近似），但这些方法依然是在完整的 Softmax 输出空间曲率几何之上进行参数侧的压缩。另一类做法是通过修改训练目标本身来降低开销，例如采用负采样、采样 Softmax 或一对多（one-vs-rest）二分类折约，但这类手段改变了原始的多分类交叉熵目标与精确梯度方向。因此，如何在严格保留原始多分类 Softmax 交叉熵损失及其真实一阶梯度的前提下，从本质上简化输出空间的曲率耦合，成为了二阶优化中悬而未决的核心矛盾。
本文的核心思路是回归多分类交叉熵本身的标量代数结构：通过将 Softmax 交叉熵精确重写为真类对竞争类 Log-Sum-Exp 聚合项的标量边际形式，揭示出标准多分类 GGN 实际上存在一个代数意义上的精确正交分解。核心 idea：将多分类 Softmax GGN 精确分解为“真类对其余”（True-vs-Rest）边际曲率与非负半正定的竞争类内部协方差残差项，FGN 仅保留真类对其余项并舍弃竞争类内部协方差，从而在保持原始损失和梯度完全不变的前提下，构建出一种保半正定、二分类下严格等价、且仅依赖批大小维度的无矩阵行空间二阶求解器。
方法详解¶
整体框架¶
FGN 方法的核心逻辑建立在“标量边际重参数化 \(\to\) GGN 曲率代数精确分解 \(\to\) 白化行空间系统构建 \(\to\) 无矩阵共轭梯度求解”这一完整闭环上。对于任意单一样本，算法首先利用竞争类 Logit 的 Log-Sum-Exp 聚合将多分类损失等价表示为关于标量边际的 Softplus 复合函数；随后，基于该标量复合形式定义的 FGN 曲率在输出空间展现出纯粹的秩一结构，从而与标准 GGN 之间剥离出一个纯粹位于竞争类子空间内的半正定协方差矩阵；最后，在小批量批次下，由于保留的曲率矩阵由每个样本单一的标量边际梯度外积构成，算法将原本位于 \(d\) 维参数空间的阻尼线性系统精确规约到规模仅为 \(b\)（批大小）维的对称正定行空间系统中，利用矩阵无关（matrix-free）的 JVP 和 VJP 进行高效 CG 迭代。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["输入样本特征与标签"] --> B["标量边际与精确损失重表示<br/>构建真类与聚合竞争类标量边际 s"]
B --> C["GGN 代数精确分解<br/>分离真类对其余曲率与竞争类内协方差"]
C --> D["白化行空间线性系统转换<br/>将 d 维参数系统规约至 b 维对称正定系统"]
D --> E["无矩阵共轭梯度求解<br/>基于标量边际 JVP/VJP 迭代更新"]
E --> F["输出：阻尼高斯牛顿参数更新方向 d"]关键设计¶
1. 标量边际重表示与精确梯度保持：在不改变损失和梯度的前提下重构输出几何 传统多分类 Softmax 交叉熵直接以 \(C\) 维 Logit 向量 \(\bm{z} = (z_1, \dots, z_C)^\top\) 为自变量。设真实类别为 \(\star\)，竞争类别集合为 \(-\star\)。本文定义聚合竞争类 Logit 为 \(z_\dagger = \log \sum_{j \neq \star} e^{z_j}\)，由此得到竞争类的总概率质量 \(p_\dagger = \sum_{j \neq \star} p_j = 1 - p_\star\)。在此基础上，定义标量真类对其余边际为 \(s = z_\dagger - z_\star\)。利用对数指数恒等变换，原始的多分类损失可精确表示为标量 Softplus 函数： $\(\ell(\bm{w}; \bm{x}, \bm{y}) = -\log p_\star(\bm{w}) = \log(1 + e^{s(\bm{w})}) = \phi(s(\bm{w}))\)$ 通过链式法则可知，标量函数的高阶导数满足 \(\phi'(s) = p_\dagger\) 以及 \(\phi''(s) = p_\star p_\dagger\)。此时标量边际关于参数的 Jacobian 行为 \(\bm{J}_s = \nabla_{\bm{w}} s^\top = -\nabla_{\bm{w}} z_\star^\top + \sum_{j \neq \star} \rho_j \nabla_{\bm{w}} z_j^\top\)，其中 \(\rho_j = p_j / p_\dagger\) 为在竞争类上的条件概率分布。直接计算参数梯度可得 \(\nabla_{\bm{w}} \ell = \phi'(s) \nabla_{\bm{w}} s = p_\dagger \bm{J}_s^\top = \bm{J}_z^\top (\bm{p} - \bm{y})\)。这一关键恒等式表明：该标量边际转换不仅保持了函数值严格不变，而且其诱导出的梯度与标准 Softmax 交叉熵梯度在代数上完全恒等，所有的近似均被严格限制在曲率层面。
2. GGN 曲率代数精确分解与 FGN 算子：剥离竞争类内协方差残差 标准的多分类 GGN 曲率定义为 \(\bm{H}^{\mathrm{GGN}} = \bm{J}_z^\top (\mathrm{diag}(\bm{p}) - \bm{p}\bm{p}^\top) \bm{J}_z\)。若直接对标量边际复合形式 \(\ell = \phi(s)\) 采用高斯牛顿构造（忽略外层标量映射的二阶项 \(\nabla_{\bm{w}}^2 s\)），则诱导出 FGN 曲率矩阵： $\(\bm{H}^{\mathrm{FGN}} = \phi''(s) \bm{J}_s^\top \bm{J}_s = p_\star p_\dagger \bm{J}_s^\top \bm{J}_s\)$ 定理证明，在 Logit 坐标排布 \((z_\star, \bm{z}_{-\star})\) 下，Softmax 协方差矩阵存在唯一的正交块分解： $\(\mathrm{diag}(\bm{p}) - \bm{p}\bm{p}^\top = p_\star p_\dagger \nabla_{\bm{z}} s \nabla_{\bm{z}} s^\top + p_\dagger \begin{bmatrix} 0 & 0 \\ 0 & \mathrm{diag}(\bm{\rho}) - \bm{\rho}\bm{\rho}^\top \end{bmatrix}\)$ 投影到参数空间后，二者之间的差值恰好为 \(\bm{H}^{\mathrm{GGN}} - \bm{H}^{\mathrm{FGN}} = p_\dagger \sum_{j=1}^{C-1} \rho_j (\bm{J}_j - \bar{\bm{J}}_\rho)^\top (\bm{J}_j - \bar{\bm{J}}_\rho)\)，即竞争类各 Jacobian 行围绕其条件加权均值 \(\bar{\bm{J}}_\rho\) 的加权协方差。由于该残差矩阵天然是半正定的（\(\bm{R}^{\mathrm{comp}} \succeq 0\)），因此 FGN 是多分类 GGN 的半正定欠逼近（\(0 \preceq \bm{H}^{\mathrm{FGN}} \preceq \bm{H}^{\mathrm{GGN}}\)）。特别地，在二分类情形下 \(C=2\)，竞争类只有一个，条件分布 \(\bm{\rho}=[1]\) 导致类内协方差精确为 0，此时 \(\bm{H}^{\mathrm{FGN}} = \bm{H}^{\mathrm{GGN}}\) 严格等价。通过定义竞争类分散度指标 \(\xi = 1 - \|\bm{\rho}\|_2^2\)，论文明确指出了残差项在 Logit 空间的迹为 \(p_\dagger \xi\)：当样本预测置信度极高（\(p_\dagger \approx 0\)）或存在单一占主导地位的强竞争负类（\(\xi \approx 0\)）时，FGN 与完整 GGN 的差距趋近于零。
3. 白化行空间系统转换与无矩阵共轭梯度：将计算复杂度与类别数解耦 在小批量 \(\mathcal{B} = \{(\bm{x}_i, \bm{y}_i)\}_{i=1}^b\) 下，批级 FGN 曲率展现出独特的对角加权低秩形式：\(\bm{H}^{\mathrm{FGN}}(\bm{w}; \mathcal{B}) = \frac{1}{b} \bm{J}^\top \bm{Q} \bm{J}\)，其中 \(\bm{J} \in \mathbb{R}^{b \times d}\) 由各样本的标量边际梯度行堆叠而成，\(\bm{Q} = \mathrm{diag}(q_1, \dots, q_b)\) 为严格正对角阵（\(q_i = p_\star^{(i)} p_\dagger^{(i)} \in (0, 1/4]\)）。标准的阻尼高斯牛顿方程为 \((\bm{H}^{\mathrm{FGN}} + \lambda \bm{I}_d)\bm{d} = -\bm{g}\)。由于梯度 \(\bm{g} = \frac{1}{b}\bm{J}^\top \bm{r}\)（其中 \(\bm{r}_i = p_\dagger^{(i)}\)）完全位于 \(\bm{J}^\top\) 的值域内，解向量必可参数化为 \(\bm{d} = -\bm{J}^\top \bm{\delta}\)。代入后可将其转化为未白化的行空间线性系统 \((\bm{Q}\bm{J}\bm{J}^\top + b\lambda \bm{I}_b)\bm{\delta} = \bm{r}\)。 为消除系统中的非对称性以应用标准的共轭梯度（CG）法，定义白化雅可比 \(\tilde{\bm{J}} = \bm{Q}^{1/2} \bm{J}\) 与白化残差 \(\tilde{\bm{r}} = \bm{Q}^{-1/2} \bm{r}\)，进而导出等价的对称正定白化行空间系统： $\((\tilde{\bm{J}}\tilde{\bm{J}}^\top + b\lambda \bm{I}_b)\bm{u} = \tilde{\bm{r}}, \quad \text{其中} \quad \bm{d} = -\tilde{\bm{J}}^\top \bm{u}\)$ 在实际算法执行中，算子 \(\bm{B}\bm{v} = \tilde{\bm{J}}(\tilde{\bm{J}}^\top \bm{v}) + b\lambda \bm{v}\) 的计算完全无需显式构建任何 \(d \times d\) 或 \(b \times b\) 的矩阵。每一次 CG 迭代只需调用一次针对标量边际映射的向量-雅可比积（VJP）与一次雅可比-向量积（JVP），外加 \(\mathcal{O}(b)\) 的向量运算。相比于标准 SGN 需要在每次曲率积中承受 \(C\) 维 Softmax 协方差矩阵的全量乘法，FGN 的曲率求解完全脱离了类别维度 \(C\) 的耦合，极大地降低了显存与时间开销。
损失函数 / 训练策略¶
在算法实现与优化超参数配置方面，FGN 与基线优化器统一遵循严谨的二阶优化设置： - 目标函数：针对参数 \(\bm{w}\) 的非正则化小批量多分类交叉熵损失 \(\mathcal{L}(\bm{w}; \mathcal{B}) = \frac{1}{b} \sum_{i=1}^b \ell(\bm{w}; \bm{x}_i, \bm{y}_i)\)。 - 阻尼策略：引入固定阻尼参数 \(\lambda_0 = 1.0\)，既保证了行空间算子严格正定可逆，又在理论上压制了被舍弃的半正定残差项对步长方向的影响。 - 内层 CG 求解：设置共轭梯度法最大迭代步数 \(N_{\mathrm{CG}} = 5\)，收敛容差阈值 \(\epsilon = 10^{-5}\)。FGN 直接在白化行空间执行无热启动的 CG 求解，每次迭代仅执行 1 次 JVP 与 1 次 VJP。 - 外层更新：获得参数空间方向 \(\bm{d}_t = -\tilde{\bm{J}}_t^\top \bm{u}_t\) 后，采用学习率 \(\alpha = 0.1\) 进行确定性参数步长更新 \(\bm{w}_{t+1} = \bm{w}_t + \alpha \bm{d}_t\)。
实验关键数据¶
主实验¶
论文在细粒度 196 类车辆分类基准 Cars196 上进行了冻结特征的线性分类头优化评估。骨干网络采用在 ImageNet 上预训练的 ResNet-50，提取 2048 维全局均值池化特征并固定。待优化的仿射分类头参数量为 \(2048 \times 196 + 196 = 401,604\)。由于分类头对参数是严格线性的，且不施加权重衰减，全量 Softmax GGN 严格等价于该分类头目标函数的精确 Hessian 矩阵，因此 SGN 构成了最直接的全曲率参照基准。实验在单张配备 8GB 显存的 GPU 上进行，记录 10 个随机种子下的测试准确率、交叉熵损失及达到目标精度的挂钟时间。
| 优化算法 | 3.65s 处测试准确率 (%) | 3.65s 处测试交叉熵 (CE) | 达到 42% 精度耗时 (s) | 达到 43% 精度耗时 (s) | 达到 44% 精度耗时 (s) |
|---|---|---|---|---|---|
| Adam (一阶基线) | 42.24 ± 0.09 | 2.8449 ± 0.0302 | 0.693 | – (未达到) | – (未达到) |
| SGN (全量 GGN) | 44.16 ± 0.09 | 2.3603 ± 0.0006 | 0.774 | 1.320 | 3.017 |
| FGN (本文) | 44.29 ± 0.07 | 2.3519 ± 0.0008 | 0.561 | 0.937 | 2.209 |
消融实验与机制检验¶
为了验证理论分解的各个环节，论文针对类别数扩展性、输出空间迹分解以及阻尼二次下降比率进行了三组机制性检验与消融分析。
| 评估维度 / 设定 | 评估指标 | FGN (本文) | SGN (全量 GGN) | 对比结论与机制验证 |
|---|---|---|---|---|
| 稳态更新耗时 (\(C=4096\)) | 单步更新时间 (ms) | 显著低于 SGN | 随 \(C\) 呈显著线性陡增 | 证明 FGN 规避 \(C \times C\) 协方差乘积后，在大类别数下展现出优异的计算扩展性 |
| 输出空间迹分解 (\(p_\star=0.6, C=10\)) | 遗漏残差迹 \(\tau_{\mathrm{drop}} = p_\dagger \xi\) | 保留迹 \(\tau_{\mathrm{ret}} = p_\star p_\dagger (2-\xi)\) | 全量迹 \(\tau_{\mathrm{full}} = 2p_\star p_\dagger + p_\dagger^2 \xi\) | 严格验证了理论分解恒等式，残差随竞争类分散度 \(\xi\) 从 0 增大到最大值而单调上升 |
| 步长二次下降保持率 (\(s_\lambda=1.0\), 强阻尼) | \(R_{\lambda_s}(\xi) = \frac{\Delta^{\mathrm{GGN}}(\bm{d}_{\mathrm{FGN}})}{\Delta^{\mathrm{GGN}}(\bm{d}_{\mathrm{GGN}})}\) | 接近 1.0 (全区间) | 基准基准值 1.0 | 证明在合理阻尼下，被舍弃的半正定类内协方差对局部最优更新方向的影响可忽略不计 |
| 步长二次下降保持率 (\(s_\lambda=0.01\), 弱阻尼) | \(R_{\lambda_s}(\xi)\) (高分散度 \(\xi\)) | 下降至约 0.7 左右 | 基准基准值 1.0 | 证实了当阻尼极小且竞争类极度分散时，类内协方差的丢失会导致步长质量出现一定偏差 |
关键发现¶
- 大类别场景下的速度优势：稳态单步耗时随类别数 \(C\) 的增长实验（\(C \in [2, 4096]\)）表明，随着 \(C\) 增加，SGN 的更新时间急剧飙升，而 FGN 则保持极低的计算开销，证明避免类别空间显式协方差乘法能有效打破多分类二阶优化的算力瓶颈。
- 精度完全媲美全量 Hessian：在 Cars196 线性头优化中，FGN 最终测试精度达到 44.29%，略优于全量 SGN（44.16%），且达到 44% 精度的优化时间从 SGN 的 3.017 秒缩短至 2.209 秒（提速约 26.8%），同时显著优于一阶优化器 Adam（42.24%）。
- 阻尼抑制效应与有效工作区间：实验严格印证了定理推论——FGN 的近似质量取决于竞争类分布的集中程度（\(\xi\)）以及阻尼因子 \(\lambda\)。在实际深度学习训练常用的阻尼设置下，FGN 能够以极小的代数近似代价换取二阶步长的精确度。
亮点与洞察¶
- 严密精确的代数分解视界：不同于以往对高斯牛顿矩阵进行经验性粗暴剪枝或层级低秩分解，FGN 通过标量边际变换揭示了 Softmax GGN 内部天然存在的“真类 vs 竞争类内”正交分解，在保持一阶损失和梯度绝对精确的前提下，将近似精准局限在非负半正定的类内残差上。
- 行空间系统的高效白化与矩阵无关性：借助对角正定加权阵 \(\bm{Q}\) 的可逆性，将非对称的行空间系统转化为无显式矩阵存储的对称正定白化系统，完美适配现代自动微分框架中的 JVP/VJP 闭包复用。
- 二阶优化向高维类别任务迁移的通用构件：该设计思路可无缝迁移至极端多分类（如大规模物体识别）、语言模型最后的 Vocab 线性输出层，或者基于对比学习的密集检索任务中，为海量类别下的二阶曲率加速提供了即插即用的新范式。
局限与展望¶
- 作者承认的局限：FGN 的近似效果具有明显的机制依赖性（regime-dependent），当竞争类概率极其平摊（如未充分训练的初期阶段，模型输出接近均匀分布且无明显最难负样本）且阻尼较小时，被忽略的类内协方差会导致曲率方向产生不可忽视的偏离。
- 自身发现的局限：实验评估主要集中在固定特征提取器的线性分类头上（此场景下 GGN 等价于精确 Hessian），对于完整多层非线性神经网络端到端微调时，非线性特征层带来的外在 Hessian 项与 FGN 曲率的交互影响仍需更大规模的实证检验。
- 未来改进思路：可以探索将“真类对其余”主导项与竞争类内协方差的低秩近似（如仅保留 Top-\(k\) 竞争负类的秩一外积修正）相结合，进一步缩小高分散场景下的理论曲率残差。
相关工作与启发¶
- vs SGN (Stochastic Generalized Gauss-Newton): SGN 在参数空间或样本空间直接使用全量 Softmax GGN，在每次 CG 迭代中都必须承受 \(C \times C\) 维度的 Softmax 协方差乘法；本文 FGN 通过剔除半正定的类内协方差，将算子精简为纯粹的标量边际积，算力显著节约且在二分类下与 SGN 严格一致。
- vs 经验 Fisher / 对角预条件子 (如 Sophia / AdaHessian): Sophia 和 AdaHessian 等方法是对参数空间对角或低秩曲率的经验估计，仍依赖全量 Softmax 梯度的外积或采样；FGN 则是从损失函数底层代数结构出发重新定义输出曲率，与参数侧的结构化压缩形成正交互补。
- vs 负采样 / 采样 Softmax: 负采样等技术为了提速改变了优化目标和梯度期望；FGN 严格保持原多分类 Softmax 交叉熵损失及梯度的精确性，仅在二阶曲率构造时进行保半正定松弛。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次从标量边际角度给出多分类 GGN 的精确代数分解，兼具代数美感与理论完备性。
- 实验充分度: ⭐⭐⭐⭐ 机制验证极其透彻，闭环验证了理论推论；若能补充端到端大模型完整微调则更佳。
- 写作质量: ⭐⭐⭐⭐⭐ 概念阐述极其清晰，数学推导严丝合缝，符号系统与结构展开典范级水准。
- 价值: ⭐⭐⭐⭐⭐ 为困扰二阶优化多年的多类别计算瓶颈提供了全新的低开销解决路径，具有极高的理论与工程启发性。