跳转至

Zeroth-Order Sharpness-Aware Learning with Exponential Tilting

会议: NeurIPS 2026
arXiv: 2510.16157
代码: https://github.com/xuchengo/ZEST
领域: 优化与学习理论
关键词: 零阶优化, 锐度感知最小化, 指数倾斜, 大语言模型微调, 泛化性

一句话总结

本文提出基于指数倾斜的零阶锐度感知优化算法 ZEST,通过倾斜参数在平均损失平滑与极值损失之间构建连续过渡,不仅保持了经典零阶优化极低的显存开销,还在理论与实证上成功规避高曲率尖锐极小值,显著提升了大语言模型下游任务的泛化性能。

研究背景与动机

在超大参数量模型微调、黑盒对抗攻击、差分隐私以及科学计算等场景中,一阶梯度往往不可用、不可靠或显存开销极其高昂。以 MeZO 为代表的零阶(Zeroth-Order, ZO)优化方法通过前向扰动的双点函数值估计梯度,仅需常数级的显存开销即可微调数十亿参数的语言模型。经典理论表明,高维对称双点估计本质上优化的是原目标函数在参数摄动邻域内的平滑版本(期望损失)。在泰勒展开下,该期望损失等价于经验风险加上与 Hessian 矩阵迹成正比的正则项,因而无意间带来了一种隐式偏差——促使局部极小值在平均意义上变得更平缓。

然而,平均平滑与经典的锐度感知最小化(SAM)存在本质分歧。标准 SAM 及其变体采用极小-极大(min-max)形式,针对参数邻域内的最坏情况(最大损失)进行惩罚,从而有效压制 Hessian 矩阵的最大特征值 \(\lambda_{\max}\)。相比之下,经典零阶方法仅对特征值平均值敏感,极易被“迹相同但某个方向曲率极高”的伪平坦极小值误导,导致在关键方向上泛化恶化;同时,由于零阶优化缺乏梯度信息,无法直接求解非光滑的 min-max 鞍点问题。

为了打破零阶优化只能优化平均损失的局限,本文引入指数倾斜(Exponential Tilting)技术,在平均损失与最坏极值损失之间架起了一座可微的光滑过渡桥梁。核心 idea:通过引入倾斜参数 \(t\) 构建软 SAM 目标函数(\(t\)-SAM),利用高斯 Stein 引理与超球面散度定理推导出仅依赖函数值采样的无偏零阶梯度估计,并设计免显存存储的伪随机数重构方案,实现保持经典零阶优化极低显存的同时自适应惩罚高曲率方向的 ZEST 算法。

方法详解

整体框架

ZEST(Zeroth-Order Exponential-Tilted Sharpness-Aware Training)的核心目标是求解带有指数倾斜的软 SAM 目标函数(\(t\)-SAM): $\(F_t(x) = \frac{1}{t} \log \mathbb{E}_{\mu(\epsilon)} \left[ e^{t f(x+\epsilon)} \right]\)$ 其中 \(\mu(\epsilon)\) 代表摄动分布(如高斯分布 \(\mathcal{N}(0, I_d)\) 或超球体 \(\sqrt{d}\mathbb{B}^d\) 上的均匀分布),\(t > 0\) 为倾斜超参数。当 \(t \to 0\) 时,\(F_t(x) \to \mathbb{E}[f(x+\epsilon)]\),退化为传统零阶优化所对应的平均损失目标;当 \(t \to \infty\) 时,\(F_t(x) \to \max_{\|\epsilon\| \le \rho} f(x+\epsilon)\),逼近标准 SAM 的最坏情况损失。

在优化流程上,ZEST 在每轮迭代中采样微批次数据与 \(k\) 个随机种子;第一阶段通过固定随机种子就地(in-place)生成正负摄动扰动并计算前向损失,求得指数倾斜权重后立即释放摄动向量;第二阶段基于选定的比率估计器计算加权系数,再次利用相同种子重构摄动并就地更新模型参数。整个过程避免了梯度回传和摄动存储,显存消耗与经典 MeZO 完全一致(仅需 \(2|x|\))。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入当前参数 x 与数据批次 D"] --> B["阶段 1:伪随机生成正负摄动与前向评估<br/>仅记录指数损失,不保留高维扰动"]
    B --> C["阶段 2:期望比率加权与偏差校正<br/>Naive 或 Bias-Corrected 归一化权重"]
    C --> D["阶段 3:原位重构摄动并执行参数更新<br/>同种子重生成扰动向量,显存开销仅 2|x|"]
    D --> E["输出更新后的平坦极小值参数"]

关键设计

1. 倾斜零阶梯度推导:借助积分几何将梯度积分转化为纯函数值差分

求解 \(F_t(x)\) 的首要难题在于其一阶梯度定义为 \(\nabla_x F_t(x) = \frac{\mathbb{E}_{\mu(\epsilon)} [e^{t f(x+\epsilon)} \nabla f(x+\epsilon)]}{\mathbb{E}_{\mu(\epsilon)} [e^{t f(x+\epsilon)}]}\),分子中显式包含参数梯度 \(\nabla f(x+\epsilon)\)。为了在无反向传播的前提下消除该梯度,本文针对不同摄动分布采用积分几何变换: 对于高斯摄动 \(v \sim \mathcal{N}(0, I_d)\),利用高维 Stein 引理 \(\mathbb{E}_v [g(v) v] = \mathbb{E}_v [\nabla_v g(v)]\),将分子中的梯度转化为被摄动变量加权后的函数评估;结合对称化构造,推导出精确等价形式: $\(\nabla_x F_t(x) = \frac{1}{t\rho} \frac{\mathbb{E}_{v \sim \mathcal{N}} \left[ \left(e^{t f(x+\rho v)} - e^{t f(x-\rho v)}\right) v \right]}{\mathbb{E}_{v \sim \mathcal{N}} \left[ e^{t f(x+\rho v)} + e^{t f(x-\rho v)} \right]}\)$ 对于半径为 \(\sqrt{d}\) 的超球体均匀摄动 \(v \sim \mathcal{B} = \mathcal{U}(\sqrt{d}\mathbb{B}^d)\),利用微分流形的高维散度定理(Divergence Theorem),将球体内部的梯度积分转化为球面边界 \(\mathcal{S} = \mathcal{U}(\sqrt{d}\mathbb{S}^{d-1})\) 上的通量积分。更关键的是,基于高维测度集中现象(超球体内几乎所有体积都聚集在外层超薄球壳上,\(\text{Var}(\|v\|) \approx \frac{1}{3d}\)),将分母对球体的积分直接近似为超球面上的期望,从而允许在一次采样中同时复用分子与分母的摄动向量,消除额外前向计算开销。

2. 期望比率估计与二阶偏差校正:解决有限采样下的非线性统计偏差

在实际优化中,每步只能采样有限个(\(k\) 个)随机摄动方向,而目标梯度的形式为“期望之比”(Ratio-of-Expectations)\(\frac{\mathbb{E}[A]}{\mathbb{E}[B]}\),其中 \(A_i = (a_i^+ - a_i^-) v_i\),\(B_i = a_i^+ + a_i^-\),\(a_i^\pm = e^{t f(x \pm \rho v_i)}\)。直接采用样本均值之比进行代入(Naive Plug-In,\(G_{\text{N}}^k\))会因 Jensen 不等式引入 \(O(1/k)\) 的渐近偏差。 为此,本文基于统计学中的二阶泰勒展开推导了带偏差校正的估计器(Bias-Corrected Plug-In,\(G_{\text{BC}}^k\))。通过减去协方差与方差产生的一阶偏倚项,将估计偏差的收敛阶提升至 \(O(1/k^2)\): $\(G_{\text{BC}}^k = \frac{1}{t\rho} \sum_{i=1}^k \left\{ 1 + \frac{k}{k-1} \left[ \bar{a}_i^+ + \bar{a}_i^- - \sum_{j=1}^k (\bar{a}_j^+ + \bar{a}_j^-)^2 \right] \right\} (\bar{a}_i^+ - \bar{a}_i^-) v_i\)$ 其中 \(\bar{a}_i^\pm = a_i^\pm / Z\) 为归一化指数损失,\(Z = \sum_{i=1}^k (a_i^+ + a_i^-)\)。该设计在不增加任何前向评估次数与显存消耗的前提下,显著稳定了小样本扰动下的更新方向。

3. 锐度偏置解析:从迹正则化到最大特征值的主动规避

理论分析揭示了 \(F_t(x)\) 在不同 \(t\) 取值下的隐式正则化项 \(R_t(x)\) 机制(定义 \(F_t(x) = f(x) + R_t(x) + O(\rho^2 d)\))。引入锐度敏感度概念 \(\phi_i(t) = \frac{\partial R_t}{\partial \lambda_i}\)(其中 \(\lambda_i\) 为 Hessian 矩阵特征值): 在高斯摄动下,\(R_t = \frac{1}{2t} \sum_{i=1}^d \left[ \frac{(t\rho g_i)^2}{1 - t\rho^2 \lambda_i} - \log(1 - t\rho^2 \lambda_i) \right]\)。当 \(t \to 0\) 时,敏感度 \(\phi_i \to \rho^2 / 2\) 为常数,所有特征方向受同等惩罚(退化为 \(\text{Tr}(\nabla^2 f(x))\));随着 \(t\) 增大,\(\phi_i\) 成为 \(\lambda_i\) 的单调递增函数,对大特征值施加不成比例的重罚。 在球体摄动与 \(t \to \infty\) 的极限情形下,\(R_\infty = \max_{\|u\| \le \sqrt{d}} \left( \rho g^\top u + \frac{\rho^2}{2} u^\top \Lambda u \right)\),同时包含斜率惩罚(反向抑制大梯度方向)和曲率惩罚。在驻点区域,该惩罚严格退化为 \(\sqrt{d} \max(\lambda_1, 0)\),仅由最大特征值 \(\lambda_{\max}\) 主导。这证明了 ZEST 具有在优化过程中主动识别并规避任意方向高曲率尖锐极值点的内在动力学特性。

一个完整示例

以参数维度为 2 的驻点优化场景为例:目标函数在两处极小值 \((1,0)\) 与 \((-1,0)\) 处的经验损失均为 0。在 \((1,0)\) 处,Hessian 特征值为 \(\{12/5, 2/5\}\);在 \((-1,0)\) 处,特征值为 \(\{10/5, 4/5\}\)。 1. 传统零阶方法(如 MeZO)评估:两者特征值之和均为 \(14/5\)(迹完全相同)。由于 MeZO 仅对平均特征值施加惩罚,优化器无法区分两者的平坦度差异,在特定初始轨迹下容易落入 \((1,0)\)。然而 \((1,0)\) 处的 \(\lambda_{\max} = 2.4\) 明显高于 \((-1,0)\) 处的 \(\lambda_{\max} = 2.0\),意味着沿第一特征向量方向受到参数微扰时,损失将剧烈攀升。 2. ZEST 优化迭代:设 \(t=1\)。在采样摄动方向时,落入高曲率方向的扰动会因指数函数 \(e^{t f(x+\rho v)}\) 的放大效应,在该方向赋予极高的归一化权重 \(\bar{a}_i^+\)。 3. 参数逃逸与收敛:加权梯度产生排斥高曲率方向的强推力,使更新轨迹避开 \((1,0)\),成功收敛到最大特征值更小、全局更平坦的理想极小值 \((-1,0)\)。

损失函数 / 训练策略

在具体实现中,ZEST 算法的超参数主要包括摄动步长 \(\rho\)、采样扰动数 \(k\)、倾斜参数 \(t\) 以及学习率 \(\eta\): - 摄动生成与显存控制:摄动向量通过伪随机数种子实时生成。前向传播计算正向损失 \(f(x+\rho v_i)\) 与负向损失 \(f(x-\rho v_i)\),计算得到指数项标量后立即销毁 \(v_i\);待加权系数确定后,二次生成相同 \(v_i\) 执行模型权重原位更新,无需额外保存激活值或累积扰动张量。 - 超参数设定:语言模型实验中,摄动规模通常固定为 \(\rho = 0.002\)(在 RoBERTa 与 OPT 上),每次迭代采样 \(k=5\) 条扰动方向;学习率在 RoBERTa 上设置为 \(\eta \in \{5\text{e}-6, 1\text{e}-5, 2\text{e}-5\}\),batch size 固定为 128。倾斜参数推荐默认选择 \(t=1\)。

实验关键数据

主实验

主实验分别在掩码语言模型(RoBERTa-Base,512-shot 分类)与自回归模型(OPT-1.3B,1000-shot 多选及生成)上展开,全面对比了一阶基线(SGD、ESAM、SAM、TSAM)与零阶基线(MeZO)。

表 1:RoBERTa-Base 在 GLUE 基准上的分类精度(干净标签与 30% 噪声标签)

方法类型 算法 SST-2 (干净/噪声) SST-5 (干净/噪声) QQP (干净/噪声) MRPC (干净/噪声) MNLI (干净/噪声) RTE (干净/噪声)
一阶方法 SGD 92.8 / 89.2 56.2 / 53.7 84.0 / 73.8 88.2 / 77.0 78.4 / 73.8 78.3 / 66.1
一阶方法 ESAM 93.0 / 89.9 56.4 / 54.6 84.3 / 79.5 88.5 / 77.5 78.4 / 75.4 79.4 / 66.8
一阶方法 SAM 93.2 / 91.1 56.4 / 55.2 84.8 / 80.2 90.0 / 78.9 79.3 / 76.9 80.1 / 68.6
一阶方法 TSAM 93.5 / 91.5 57.5 / 55.2 85.0 / 81.0 89.2 / 77.7 79.5 / 76.5 80.5 / 67.5
零阶基线 MeZO 92.1 / 89.0 48.6 / 44.7 71.4 / 62.4 81.9 / 67.2 71.8 / 60.3 72.9 / 59.9
零阶本文 \(\text{ZEST}_{\text{N}}\) 92.2 / 89.4 49.4 / 46.2 71.6 / 68.3 83.6 / 68.6 73.6 / 63.4 73.3 / 61.4
零阶本文 \(\text{ZEST}_{\text{BC}}\) 92.0 / 88.2 49.7 / 44.7 72.6 / 62.7 81.6 / 68.9 73.8 / 63.4 72.9 / 61.7

表 2:OPT-1.3B 在多选问答与文本生成任务上的性能对比(测试集准确率 / F1 分数)

方法类型 优化算法 COPA (多选 Acc) ReCoRD (多选 Acc) SQuAD (生成 F1) DROP (生成 F1) 显存复杂度
一阶方法 SGD 75.0 72.2 83.4 29.7 \(\sum_l \max(\|a_l\|, \|x_l\|) + \|x\|\)
一阶方法 ESAM 76.0 72.5 83.7 31.2 \(\sum_l \max(\|a_l\|, \|x_l\|) + 2\|x\|\)
一阶方法 SAM 77.0 72.7 84.3 31.8 \(\sum_l \max(\|a_l\|, \|x_l\|) + 2\|x\|\)
一阶方法 TSAM 77.0 72.1 84.6 31.3 \(\sum_l \max(\|a_l\|, \|x_l\|) + (k+1)\|x\|\)
零阶基线 MeZO 74.0 72.4 78.8 25.2 \(2\|x\|\)
零阶本文 \(\text{ZEST}_{\text{N}}\) 78.0 72.3 79.4 25.5 \(2\|x\|\)
零阶本文 \(\text{ZEST}_{\text{BC}}\) 77.0 72.5 79.0 25.7 \(2\|x\|\)

消融与分析

表 3:不同倾斜参数 \(t\) 下 RoBERTa-Base 的验证集准确率对比(消融分析)

设置 (\(k=5\)) SST-2 (干净) SST-5 (干净) MRPC (干净) MNLI (30% 噪声) QQP (30% 噪声) 说明
\(t=0\) (退化为 MeZO) 90.9 48.0 81.6 58.7 62.3 仅惩罚 Hessian 平均迹
\(t=1\) (ZEST 推荐值) 92.6 (+1.7) 47.9 (-0.1) 82.4 (+0.8) 61.6 (+2.9) 68.0 (+5.7) 显著惩罚大特征值,提升稳健性
\(t=5\) 91.8 48.4 82.1 60.5 67.2 强倾向于极值损失惩罚
\(t=20\) 91.5 48.1 81.9 60.2 66.8 逼近最大损失,方差略有上升

关键发现

  • 泛化表现全面超越传统零阶优化:在干净数据集上,ZEST 相较 MeZO 准确率最高提升 1.7%;在包含 30% 标签噪声的严苛测试下,提升幅度更是高达 5.9%(QQP 上从 62.4% 跃升至 68.3%),证明倾斜 SAM 带来的平坦极小值具有极强的抗噪鲁棒性。
  • 多选任务中比肩甚至反超一阶 SAM:在 OPT-1.3B 的 COPA 任务中,\(\text{ZEST}_{\text{N}}\) 取得了 78.0% 的准确率,不仅比 MeZO 高出 4.0%,还超越了一阶 SGD(75.0%)与一阶 SAM(77.0%),且显存开销不到一阶方法的几分之一。
  • 真实平坦度定量验证:对微调收敛点的邻域采样与 Hessian 计算显示,ZEST 找到的极小值在不同摄动半径下的平均领域损失均显著低于 MeZO,且 Hessian 矩阵前 5 个最大特征值全面小于 MeZO(如在 MRPC 上前 5 大特征值显著下移),实证支撑了理论对 \(\lambda_{\max}\) 敏感性的预测。
  • 超参数 \(t\) 具备极高宽容度:实验表明 \(t=1\) 是开箱即用的稳健选择,在 7/8 的任务设置中均战胜或持平 \(t=0\),无需繁琐的超参搜索。

亮点与洞察

  • 巧用积分几何实现无梯度 SAM 求解:通过 Stein 引理与超球面散度定理,将原本必须依赖反向传播的指数梯度期望化简为函数评估差分之比,解决了零阶优化长期无法求解 min-max 类最坏情况目标函数的难题。
  • 高维球壳体积集中性质的巧妙复用:借助高维空间中球体体积集中在表面的测度特性,将分母的球面积分直接近似为超球面期望,使正负摄动评估结果得以同时服务于分子与分母,避免了翻倍的前向评估开销。
  • 基于种子的原地重构消除了扰动存储显存:通过两阶段重构技巧,将摄动生成解耦为“种子采样 \(\to\) 标量损失记录 \(\to\) 原地更新”,实现了与经典一阶 SAM 相同的平坦解偏置,却仅需经典零阶优化的基准显存。

局限与展望

  • 小采样数下偏差校正增益受限:虽然理论上 \(G_{\text{BC}}^k\) 将估计偏差从 \(O(1/k)\) 降低至 \(O(1/k^2)\),但在实际受限于算力而选择 \(k=5\) 的极小扰动数时,两者的经验性能差异并不明显,更高级的比率估计器仍有待探索。
  • 依赖双倍前向评估:尽管无需反向传播梯度,但每次迭代仍需进行 \(2k\) 次前向传播(正反方向各 \(k\) 次),虽然显存开销极小且单次前向极快,但在总体训练轮数上相较一阶方法仍需更多步数收敛。
  • 更广泛非凸场景的理论边界:当前的曲率敏感性推导建立在二次泰勒展开附近,在模型参数高度病态或非光滑激活区域中的理论保证仍可进一步拓展。

相关工作与启发

  • vs MeZO (Malladi et al., 2023): MeZO 采用基础双点估计微调 LLM,本质上优化的是等权重惩罚所有 Hessian 特征值的迹正则化;ZEST 通过引入指数倾斜参数 \(t\),平滑插值到最大特征值敏感惩罚,有效规避了尖锐极小值。
  • vs SAM (Foret et al., 2020) & TSAM (Li et al., 2024): 一阶 SAM 及其倾斜变体必须依赖梯度计算扰动方向并反向传播更新,显存消耗大且需维护多个模型副本;ZEST 实现了倾斜 SAM 的零阶等价体,在完全无梯度的前提下保持 \(2|x|\) 显存上限,为极端受限显存下的模型微调提供了新范式。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次建立零阶优化与指数倾斜 SAM 之间的精确数学映射,理论优雅实用。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 RoBERTa 到 OPT-1.3B、分类到生成、干净到加噪,并附带完备的曲率谱分析。
  • 写作质量: ⭐⭐⭐⭐⭐ 数学推导严谨精炼,结构层次清晰,实验结论自洽。
  • 价值: ⭐⭐⭐⭐⭐ 为超大模型黑盒/低显存微调开辟了兼顾平坦泛化与计算可行性的有效路径。