跳转至

Trust the Direction, Search the Step: Zero-and-First-Order Methods for LLM Fine-Tuning

会议: NeurIPS 2026
arXiv: 2610.02190
代码: https://github.com/nizswan/Zeroth-First-Order-Framework
领域: LLM 其他
关键词: 零阶优化, 一阶优化, 步长自适应, 局部曲率建模, 大模型微调

一句话总结

提出零阶与一阶混合优化框架(ZFO),将高维参数更新解耦为一阶优化器确定搜索方向与沿该方向一维子空间内的低开销零阶前向探测,仅需两次对称探测即可构建二阶/三阶泰勒或帕德局部曲率模型并闭式选取最佳步长,显著改善大模型微调的稳定性和收敛质量。

研究背景与动机

基于梯度的一阶(First-Order, FO)优化算法(如 SGD、AdamW 和 Muon)是现代深度学习与大语言模型(LLM)微调的标准基石。在这类算法中,梯度反向传播能够提供高质量的高维参数更新方向,然而沿着该方向究竟应该更新多远——即学习率或步长的选取,长期以来却只能依赖人工调优、全局学习率调度表或静态经验规则。在大语言模型微调(特别是基于强化学习可验证奖励 RLVR / GRPO 或监督微调 SFT)场景下,损失曲面的局部几何结构极其复杂且在训练过程中剧烈剧变;保守的步长设定会导致收敛极其缓慢、无法充分利用批次中的样本信息,而过于激进的步长则会诱发训练发散与灾难性性能崩塌。传统数值优化中的经典线搜索方法需要多次重算目标函数,在十亿级参数模型上成本过高;纯无梯度的零阶(Zeroth-Order, ZO)优化(如 MeZO)虽免去了梯度图反向传播,但在极高参数维度下沿随机方向差分会面临极大的方差与维数灾难。

这一困境背后的核心矛盾在于:高维参数空间中方向探索的有效性高度依赖一阶梯度的精确引导,但高维优化在确定方向后本质上即退化为一条一维射线上的标量步长选取问题,而现有一阶方法却完全缺失了对该一维局部几何曲率的动态感知机制。与此同时,现代 LLM 微调在每个迭代步中都会构建一个固定的样本小批次;在固定批次(Common Random Numbers, CRN)的条件下,在该点附近沿给定方向进行扰动评估仅需要重复执行轻量的前向推理,而无需重新采样提示词、生成响应或重构计算图,计算与显存开销极为低廉。

因此,与其用高方差的零阶探测盲目猜测高维更新方向,不如让一阶优化器专精于提供可靠的方向向量,而将廉价的零阶前向探测限制在这一维子空间内,用极少的函数调用恢复局部的二阶甚至三阶曲率信息,自适应决定当前步的位移大小。核心 idea:提出零阶与一阶混合优化框架(ZFO),在一阶优化器选定更新方向后,仅复用当前批次进行两次对称的前向扰动探测,估计二阶与三阶方向导数并构建局部泰勒或帕德近似模型,在有界搜索区间内闭式求解最优步长。

方法详解

整体框架

ZFO 将每次迭代的优化决策解耦为两步:首先由底层一阶优化器(如 AdamW)在当前参数 \(\theta_t\) 和小批次 \(\mathcal{S}_t\) 上计算梯度并提议参数位移向量 \(g_t\),归一化为单位方向 \(z_t = g_t / \|g_t\|\),从而将高维目标函数限制在一条一维线路上 \(\phi_t(s) = F_t(\theta_t + s z_t)\);随后,ZFO 在有界区间 \(\mathcal{I}_t = [0, R_t]\)(其中 \(R_t = \beta \alpha \|g_t\|\),\(\beta\) 为搜索半径倍率)内通过零阶评估选择标量步长 \(s_t\)。前向计算给出起点函数值 \(\phi_t(0)\),反向传播顺带提供一阶方向导数 \(\phi_t'(0) = \langle \nabla F_t(\theta_t), z_t \rangle\);在此基础上,ZFO 仅引入两次在固定小批次上的对称前向探测 \(\phi_t(\epsilon) = F_t(\theta_t + \epsilon z_t)\) 与 \(\phi_t(-\epsilon) = F_t(\theta_t - \epsilon z_t)\),利用二阶和三阶有限差分快速拟合一维局部代理模型 \(M(s)\),并在区间内解析求解最大化步长。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["当前参数与固定小批次<br/>θ_t, S_t"] --> B["一阶优化器提议更新方向<br/>z_t = g_t / ||g_t||"]
    B --> C["共享批次对称零阶探测<br/>评估 phi(+epsilon) 与 phi(-epsilon)"]
    C --> D["多阶局部显式建模<br/>构建 T2/T3 或 P2/P3 代理模型"]
    D --> E["有界区间极值求解与安全回退<br/>解出 s_t in [0, R_t],极点退回 T2"]
    E --> F["参数更新<br/>θ_t+1 = θ_t + s_t z_t"]

关键设计

1. 共享批次对称零阶探测:以低廉开销捕获高阶方向导数
传统的数值精确线搜索动辄需要 5-10 次全模型前向评估,而直接在高维空间估计 Hessian 矩阵则面临内存爆炸。为了在几乎不增加推理负担的前提下获取可信的局部曲率,ZFO 提出在完全相同的输入小批次 \(\mathcal{S}_t\) 上复用随机数(Common Random Numbers, CRN),仅沿一阶方向做正负微小扰动 \(\epsilon\)。通过常规前向传播得到 \(\phi(0)\),反向传播内积得到 \(a = \phi'(0) = \langle \nabla F(\theta), z \rangle\)。二阶中心有限差分估计为: $\(b = \widehat{\phi''}(0) = \frac{\phi(\epsilon) - 2\phi(0) + \phi(-\epsilon)}{\epsilon^2}\)$ 进一步利用已知精确的一阶方向导数 \(a\),通过对称探针的反对称分量消去四阶偶数阶误差项,推导得到三阶导数的闭式差分估计: $\(c = \widehat{\phi'''}(0) = \frac{3}{\epsilon^3}\left(\phi(\epsilon) - \phi(-\epsilon) - 2\epsilon a\right)\)$ 理论证明表明,该三阶估计的确定性差分截断误差仅为 \(\mathcal{O}(\epsilon^2)\)。由于所有探针均在固定批次上计算,条件随机性被消除,有限样本估计误差服从次高斯集中不等式,以高概率收敛于总体真实曲率。

2. 多阶多族局部显式建模:泰勒多项式与帕德有理逼近
获得了局部的一阶、二阶与三阶导数估计 \(\{a, b, c\}\) 后,必须构建能够向外推演至整个搜索区间的解析模型。ZFO 统一提供了多项式族与有理函数族共 4 个代表性实例: - 二阶泰勒模型(\(T_2\)):\(T_2(s) = \phi(0) + as + \frac{1}{2} b s^2\)。若 \(b < 0\)(局部凹),其驻点为 \(s = -a/b\);结合边界点比较即可求得区间极值。 - 三阶泰勒模型(\(T_3\)):\(T_3(s) = \phi(0) + as + \frac{1}{2} b s^2 + \frac{1}{6} c s^3\)。导数方程 \(a + bs + \frac{1}{2} c s^2 = 0\) 是一元二次方程,通过求根公式可在常数时间内解析找出落于搜索区间内部的候选极值点。 - 二阶帕德模型(\(P_2\)):\(P_2(s) = \phi(0) + \frac{as}{1 + q_2 s}\),其中 \(q_2 = -b / (2a)\),匹配至二阶导数。 - 三阶帕德模型(\(P_3\)):\(P_3(s) = \phi(0) + \frac{as + \rho s^2}{1 + q_3 s}\),其中 \(q_3 = -c / (3b), \rho = \frac{b}{2} + a q_3\),匹配至三阶导数。 有理帕德模型能够比单纯的多项式更好地捕获神经网络目标函数沿更新方向的平缓饱和特性与非多项式曲率变化,从而在更大的搜索半径内给出更平滑的步长决策。

3. 有界搜索区间与极点安全回退:规避外推发散与稳定性保障
有理模型在提供高阶拟合能力的同时,其分母多项式存在在搜索区间内部产生奇点(极点导致值趋于无穷)的潜在风险;此外,局部有限差分在远距离外推时截断误差会急剧累积。为此,ZFO 设定了严格的双重防护:首先,将搜索区间严格限制在 \(\mathcal{I}_t = [0, R_t]\) 内,其中有效搜索半径 \(R_t = \beta \alpha \|g_t\|\),由超参数 \(\beta\) 控制相对于标称一阶步长的放缩上限;其次,当所选的帕德模型(\(P_2\) 或 \(P_3\))在区间 \([0, R_t]\) 内出现分母等于零的极点或系数未定义时,算法自动触发安全机制,回退至二阶泰勒模型 \(T_2\) 求解。在理论层面,全局收敛性定理证明了即使存在局部建模截断误差 \(\tilde{\Delta}_T\),ZFO 迭代序列在非凸平滑目标下仍能以 \(\mathcal{O}(1/T)\) 速率严格收敛到由平均步长选择误差控制的稳定点邻域中。

损失函数 / 训练策略

在 LLM 的强化学习可验证奖励(RLVR)微调中,底层优化器采用 AdamW,策略梯度采用广义相对策略优化(GRPO)。每次迭代中,模型针对提示词采样一组候选响应并评估确定性奖励,由此构建替代目标函数: $\(F_t(\theta) = \mathbb{E}_{(x, y) \sim \mathcal{S}_t}\left[ \min\left(r_t(\theta) A_t, \text{clip}(r_t(\theta), 1-\epsilon_{\text{clip}}, 1+\epsilon_{\text{clip}}) A_t\right) - \beta^{\star} D_{\text{KL}}(\pi_\theta \| \pi_{\text{ref}}) \right]\)$ 其中 \(r_t(\theta) = \frac{\pi_\theta(y \mid x)}{\pi_{\theta_{\text{old}}}(y \mid x)}\)。在两次零阶对称探测时,冻结已生成的响应序列 \(y\) 与奖励优势值 \(A_t\),仅在参数 \(\theta_t \pm \epsilon z_t\) 下重新计算模型前向概率和 KL 惩罚,从而在常数级时间内完成目标函数值评测。探测步长尺度统一固定为 \(\epsilon = 10^{-3}\),基线学习率 \(\alpha \in \{10^{-6}, 5\times 10^{-6}, 10^{-5}\}\),搜索半径参数 \(\beta \in \{3, 5, 10\}\)。

实验关键数据

主实验

在 Qwen-2.5-Math-1.5B、Phi-2、Gemma-2-2B 和 Llama-3.2-1B 上,针对 GSM8K、MATH、SVAMP、AsDiv 和 OpenBookQA 多个数学与常识推理基准评测 ZFO 各变体相对于一阶 AdamW 和零阶 MeZO 基线的表现(3 个随机种子均值与标准差):

模型 数据集 AdamW (FO) MeZO (ZO) ZFO-Taylor2 ZFO-Taylor3 ZFO-Padé2 ZFO-Padé3 (本文最佳)
Qwen-2.5-Math-1.5B GSM8K 78.24 ± 0.59 0.00 81.00 ± 0.11 75.56 ± 0.24 80.95 ± 0.49 80.97 ± 0.73
Qwen-2.5-Math-1.5B MATH 31.77 ± 1.63 22.01 ± 3.93 40.10 ± 1.19 20.18 ± 3.03 37.50 ± 1.79 39.58 ± 0.45
Qwen-2.5-Math-1.5B SVAMP 88.11 ± 0.69 68.67 ± 10.33 90.78 ± 0.96 89.11 ± 0.51 90.89 ± 1.02 90.33 ± 2.19
Qwen-2.5-Math-1.5B AsDiv 91.45 ± 0.36 64.72 ± 3.51 94.09 ± 0.29 88.27 ± 0.35 93.04 ± 0.22 93.46 ± 0.18
Qwen-2.5-Math-1.5B OpenBookQA 26.07 ± 2.53 8.00 ± 10.05 37.80 ± 17.85 39.73 ± 20.99 60.67 ± 0.95 64.07 ± 1.22
Phi-2 (2.7B) SVAMP 55.44 ± 4.91 22.44 ± 3.89 55.67 ± 7.86 39.33 ± 0.88 64.22 ± 0.19 65.11 ± 5.39
Phi-2 (2.7B) AsDiv 56.51 ± 18.76 10.15 ± 4.88 65.42 ± 2.59 55.85 ± 2.55 65.58 ± 3.74 69.27 ± 1.92
Phi-2 (2.7B) OpenBookQA 51.20 ± 1.11 20.33 ± 10.90 58.27 ± 13.62 51.27 ± 0.95 49.40 ± 0.20 52.07 ± 0.23
Gemma-2-2B SVAMP 52.00 ± 1.76 14.44 ± 5.93 59.56 ± 0.51 47.22 ± 2.01 59.00 ± 1.76 58.22 ± 3.69
Llama-3.2-1B GSM8K 4.85 ± 5.47 1.30 ± 0.98 4.02 ± 3.98 9.00 ± 6.68 5.05 ± 6.94 1.72 ± 0.88
Llama-3.2-1B AsDiv 16.60 ± 28.12 3.33 ± 0.48 13.72 ± 23.17 44.77 ± 1.42 25.23 ± 21.97 52.16 ± 4.75
Llama-3.2-1B OpenBookQA 25.53 ± 0.31 0.00 34.40 ± 1.56 31.53 ± 0.81 34.60 ± 1.04 25.64 ± 0.23

消融实验

为了检验 ZFO 的性能增益是否单纯来源于扩大了搜索半径(\(\beta > 1\)),在 Qwen-2.5-Math-1.5B 上进行 \(\beta=1\) 的严格控制消融(此时搜索区间强制截断在标称一阶步长内 \([0, s_{\text{FO}}]\),ZFO 仅能缩短或重新选点,绝不能前移超过默认步长):

数据集 基础学习率 AdamW (FO) ZFO-Taylor2 ZFO-Taylor3 ZFO-Padé2 ZFO-Padé3
ARC-Challenge \(1\times 10^{-5}\) 31.71 ± 8.11 56.77 ± 0.61 56.51 ± 0.36 55.49 ± 0.98 49.86 ± 0.27
StrategyQA \(1\times 10^{-5}\) 59.24 ± 10.14 70.69 ± 5.10 74.87 ± 0.37 67.25 ± 13.64 48.18 ± 0.44
ASDiv \(5\times 10^{-6}\) 82.08 ± 0.08 81.82 ± 0.26 78.05 ± 1.41 81.87 ± 0.26 81.92 ± 0.46

此外,在与其他步长自适应策略(带余弦退火的 AdamW 及无参数自适应优化器 Prodigy)对比中,ZFO 最优变体在 Qwen 上所有 8 个推理任务上均显著胜出(例如在 GSM8K 上达到 81.00%,而 AdamW+Cosine 为 72.88%,Prodigy 为 64.32%)。

关键发现

  • 普适性显著提升:在全部 12 个主测试配置中,至少有一种 ZFO 变体性能显著优于 AdamW 基线;在更广泛的 16 个评测设置中,Padé3 和 Taylor2 分别在 14 个设置中取得净胜(胜率达 87.5%)。
  • 缓解高方差与优化崩塌:当基线一阶学习率略大导致不稳定时(如 AdamW 在 OpenBookQA 和 AsDiv 上的极高方差),ZFO 能迅速通过负曲率探测将步长回调至稳定区间,大幅平抑跨随机种子的方差。
  • 开销微不足道:实测显示,由于强化学习微调的主要时间开销在轨迹 Rollout 阶段,ZFO 增加的两次前向探测在 Qwen-1.5B 上仅增加约 5.2% 的墙钟时间(Wall-clock time),GPU 峰值显存增幅仅在 11.0% - 19.1% 之间,未引入新的显存瓶颈。

亮点与洞察

  • 解耦维度的巧思:巧妙避开了零阶优化在高维空间中“方向随机、方差爆炸”的致命缺陷,将零阶前向探测严格约束在一阶梯度指向的 1D 子空间内,以极低代价换取了珍贵的局部曲率信息。
  • 重用梯度的三次差分技巧:利用反向传播已经给出的精确一阶方向导数 \(\phi'(0)\),通过两次对称探测即可同时消除奇数阶和偶数阶截断误差,闭式恢复三阶导数 \(c\),使得构建三阶泰勒与有理帕德模型成为可能。
  • 通用包装器特性:ZFO 本质上是一个优化器包装器(Wrapper),底层不仅可以搭载 AdamW,还能直接兼容 Muon 或标准 SGD,甚至可拓展至经典非凸数值优化(在 Rosenbrock 和 Beale 测试函数上同样显著加速收敛)。

局限与展望

  • 作者承认的局限:ZFO 的步长选择严重依赖于底层一阶优化器提议的方向是否准确;若当前微调小批次的梯度方向本身与总体目标方向严重偏离,沿该方向无论如何搜索步长都无法避免次优更新。
  • 深层局限与假设约束:算法引入了额外的超参数 \(\beta\)(搜索半径倍率);实验显示若初始学习率 \(\alpha\) 已经过大,盲目搭配过大的 \(\beta\) 仍会导致外推失效。此外,对于纯预训练或海量 SFT,每次更新增加 2 次前向传播虽然相对微调开销小,但若吞吐受限仍会有轻微累计开销。
  • 改进方向:探索动态自适应调节搜索半径 \(\beta\) 的元策略;进一步结合小批次梯度的方差估计,自适应决定何时触发高阶帕德回退。

相关工作与启发

  • vs MeZO [Malladi et al., 2023]: MeZO 采用随机高维扰动估计梯度以节省显存,但在复杂任务上由于方差巨大导致收敛缓慢;ZFO 保留了一阶梯度的方向引导,仅用零阶探测做 1D 步长校准,性能大幅压制 MeZO。
  • vs GeN [Bu & Xu, 2025]: GeN 采用二次拟合进行广义牛顿步长选择;ZFO 证明了 GeN 是该框架下的一个特定二次实例,并进一步拓展到了三阶泰勒与有理帕德模型,提供了更统一的集中性与平稳点收敛性分析。
  • vs 传统线搜索 (Armijo, Barzilai-Borwein, PLS): 经典线搜索需要多次试探性回溯或概率高斯过程建模,在大模型上调用开销无法承受;ZFO 是固定预算(仅 2 次探针)的封闭解析式求解,在保持极高精度的同时计算开销确定且极低。

评分

  • 新颖性: ⭐⭐⭐⭐☆(将高维一阶方向与 1D 零阶高阶曲率拟合解耦,设计紧凑而兼具理论完备性)
  • 实验充分度: ⭐⭐⭐⭐⭐(横跨 4 个不同架构大模型、十余个推理任务、6 个经典非凸优化问题,含详尽的消融与计算开销实测)
  • 写作质量: ⭐⭐⭐⭐⭐(概念清晰、理论推导严谨自洽、逻辑线索极佳)
  • 价值: ⭐⭐⭐⭐☆(即插即用的大模型微调优化器加速组件,具有很高的实用工程参考价值)