同时优化模型平坦性和梯度对齐,提升多分布学习泛化能力
Flatness and Gradient Alignment Are Both Necessary: Spectral-Aware Gradient-Aligned Exploration for Multi-Distribution Learning

- 基于损失曲面分解,同时关注曲率与梯度对齐两项几何特性
- 在五个域泛化和两个多任务学习基准上达到新最优性能
- 适合追求鲁棒泛化的深度学习研究者和工程实践者
尖锐性感知与梯度对齐方法虽能提升泛化性能,但各自仅关注损失曲面的单一几何特性。本文揭示,在多分布学习中,忽略任一特性都存在结构性缺陷。我们推导出过拟合风险的分解式,包含两项主导项:(i) 由平均海森矩阵 \bar{H} 与梯度协方差 \Sigma_g 的迹控制的对齐项 \tr(\bar{H}^{-1}\Sigma_g),(ii) 由 \bar{H} 控制的曲率项。值得注意的是,\bar{H} 在两项中分别以逆和非逆形式出现。通过反例证明二者互不蕴含,任何仅优化其一的方法无法保证低过拟合风险。为此提出 SAGE 方法,同时优化两项:曲率部分用牛顿-舒尔茨迭代计算梯度矩阵极分解,实现各方向等幅上升;对齐部分在下降步注入与跨分布梯度差异成比例的各向同性噪声。在五项域泛化与两项多任务学习基准测试中,SAGE 在 DomainBed 上达新最优,并可作为通用改进模块显著提升基础多任务学习器表现。
原文摘要 · Abstract (English)
Sharpness-aware and gradient-alignment methods have been shown to improve generalization, however each family of methods targets a single geometric property of the loss landscape, while ignoring the other. In this paper, we show that this omission is structurally unavoidable and that both flatness and gradient alignment should be considered in multi-distribution learning settings. Specifically, we derive an excess-risk decomposition that yields two additive leading-order terms: (i) an alignment term, controlled by the trace of $\bar{H}^{-1}Σ_g$ and (ii) a curvature term, controlled by $\bar{H}$, where $\bar{H}$ is the average Hessian and $Σ_g$ is the covariance of the gradient across distributions. Notably, $\bar{H}$ appears inverted in one and non-inverted in the other. We further show, via a counterexample, that neither quantity bounds the other in general, so no algorithm targeting only one term can guarantee low excess risk. Motivated by this decomposition, we propose SAGE (Spectral-Aware Gradient-Aligned Exploration) that targets both terms. The curvature component replaces SAM's gradient-scaled perturbation with the polar factor of each layer's gradient matrix, computed via Newton-Schulz iteration, so that the ascent step probes all directions with similar magnitude. On the other hand, the alignment component injects isotropic noise at the descent step, the magnitude of which scales with cross-distribution gradient disagreement. Experiments on five domain-generalization and two multi-task learning benchmarks show that the proposed method establishes a new state-of-the-art on DomainBed and acts as a general-purpose improvement to base MTL solvers, remaining competitive with, or even surpassing, state-of-the-art methods.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。