跳转至

Low-Rank Curvature for Zeroth-Order Optimization in LLM Fine-Tuning

会议: AAAI 2026
arXiv: 2511.07971
代码: GitHub
领域: LLM微调优化
关键词: 零阶优化, 曲率感知, 低秩预条件, 方差缩减, 内存高效

一句话总结

提出 LOREN，一种曲率感知的零阶优化方法，通过低秩块对角预条件器捕获损失景观的各向异性曲率，并结合 REINFORCE Leave-One-Out 方差缩减技术，在 LLM 微调中实现了更高精度和更快收敛，同时相比 MeZO-Adam 节省高达 27.3% 的峰值内存。

研究背景与动机

  1. 领域现状：零阶（ZO）优化方法仅通过前向传播估计梯度，无需反向传播和中间激活存储，使得 LLM 微调的内存需求大幅降低。MeZO 开创了这一范式，后续工作如 HiZOO、LOZO 等进一步改进。
  2. 现有痛点：现有 ZO 方法存在两个根本缺陷——（1）有限差分梯度估计在高维空间中方差极大，导致参数更新不稳定；（2）对损失景观的各向异性曲率视而不见，在高曲率方向振荡、低曲率方向停滞，甚至收敛到鞍点。
  3. 核心矛盾：引入曲率信息（如 Hessian）通常需要额外内存和计算，这与 ZO 方法追求极低内存的初衷直接矛盾。如何在几乎不增加内存的前提下注入曲率感知能力？
  4. 本文要解决什么？ 同时解决 ZO 优化中的高方差和曲率无关问题，在保持内存高效的同时显著提升收敛速度和微调精度。
  5. 切入角度：将梯度预条件问题重新表述为自适应估计扰动分布的协方差矩阵，借助自然演化策略（NES）框架和 Kronecker 分解将协方差参数化为低秩结构。
  6. 核心idea一句话：ZO 优化中的梯度预条件等价于从各向异性高斯分布中采样扰动向量，用低秩 Kronecker 分解的协方差矩阵近似 Hessian 逆，结合 RLOO 方差缩减实现高效曲率感知更新。

方法详解

整体框架

LOREN 在每次迭代中：（1）从参数化的各向异性高斯分布采样 K 个扰动向量；（2）用扰动后的参数做 K 次前向传播算损失；（3）用 RLOO 估计器计算低方差的梯度估计和协方差参数梯度；（4）同时更新模型参数 x 和协方差参数 a。

关键设计

  1. 预条件 = 各向异性扰动分布：标准 ZO-SGD 从各向同性高斯采样扰动 u ~ N(0, I)，忽略了曲率差异。LOREN 的核心洞察在于：预条件梯度更新 H̃⁻¹∇f(x) 等价于从 N(0, H̃⁻¹) 采样扰动后做有限差分估计。因此，学习一个合适的协方差矩阵 Σ = H̃⁻¹ 就自然实现了曲率感知的预条件。

  2. 低秩块对角 Hessian 近似：对每个网络层参数矩阵 X ∈ R^{m×n}，将 Hessian 近似为 H̃ = I_m ⊗ (ρI_n + aa^T)，其中 ρ 是阻尼因子，a ∈ R^n 是可学习向量。该近似具有闭式逆和闭式平方根逆，额外内存仅为 O(n)（存储向量 a），远低于 MeZO-Adam 的 O(mn)。这是首个在 ZO 优化中使用块对角 Hessian 近似的方法。

  3. REINFORCE Leave-One-Out 方差缩减：不同于传统 ZO 方法使用 SPSA（两点有限差分），LOREN 采用 RLOO 估计器——对 K 个扰动样本中的每一个，用其余 K-1 个样本的损失均值作为基线减去，有效降低方差。实验显示 MSE 降低显著。

  4. 自然演化策略（NES）框架：将优化目标转化为最小化搜索分布下的期望损失 J(θ) = E_{z~p(z;θ)}[f(z)]，利用 score function estimator 计算协方差参数 a 的梯度，无需额外前向传播。参数 x 和 a 在每步中同时更新。

损失函数 / 训练策略

全参数微调，不使用 prompt。每步做 K=6 次前向传播（扰动评估），RLOO 估计器无需反向传播。使用 early stopping 防止过拟合。heavyball 动量加速收敛。收敛率为 O(1/√T)。

实验关键数据

主实验

GPT-2-XL (1.5B) 在 GLUE 上的准确率（%）：

方法 MNLI QNLI SST-2 CoLA 平均
MeZO 39.1 58.8 73.8 65.4 59.3
MeZO-Adam 50.9 72.3 91.2 71.6 71.5
HiZOO 48.6 66.3 89.6 71.5 69.0
LOREN 51.2 74.6 89.8 72.0 71.9

OPT-13B (SuperGLUE)：LOREN 在 CB 上达 73.7%，全面超越所有 ZO baseline。

消融实验

峰值 GPU 内存对比（OPT-13B，BF16）：

方法 内存 (GB) 相对 MeZO
MeZO 32.9 1.00×
MeZO-Adam 76.0 2.31×
HiZOO 59.6 1.81×
LOREN ~41 ~1.25×

LOREN 的额外内存开销仅为 O(n)，远低于 MeZO-Adam 的 O(mn)。

关键发现

  • LOREN 在 DistilBERT (66M) 到 OPT-13B 的不同规模模型上一致优于或持平所有 ZO baseline
  • 在 RoBERTa-large 上，LOREN 平均准确率 70.1%，大幅超越 MeZO 的 58.4% 和 HiZOO 的 64.2%
  • RLOO 方差缩减在 1000 维测试函数上梯度估计 MSE 持续低于标准 ZO-SGD
  • 在 monkey saddle 函数上，LOREN 是唯一成功逃出鞍点区域的 ZO 方法

亮点与洞察

  • 将 "预条件" 和 "扰动分布学习" 统一起来是非常优雅的理论贡献
  • 低秩 Kronecker 分解将协方差参数化为单个向量 a，内存开销极小
  • 首次在 ZO 优化中同时实现曲率感知和方差缩减
  • 提供了 O(1/√T) 的收敛性保证

局限性 / 可改进方向

  • 全参数微调设定下与 LoRA 等参数高效方法的对比缺失
  • 每步需要 K=6 次前向传播，总计算量高于仅需 2 次前向传播的 MeZO
  • 块对角近似假设层间参数独立，忽略了跨层曲率相关性
  • 在 LLaMA-3-8B 上提升相比小模型没那么显著，对更大模型的可扩展性有待验证

相关工作与启发

  • vs MeZO: LOREN 在 MeZO 基础上同时引入曲率和方差缩减，GPT-2-XL 平均准确率从 59.3 提升到 71.9
  • vs HiZOO: HiZOO 只估计 Hessian 对角线且需额外前向传播，LOREN 用块对角近似捕获更丰富曲率，内存更低
  • vs LOZO: LOZO 做低秩梯度估计，LOREN 做低秩预条件器，思路互补
  • vs MeZO-SVRG: SVRG 需要全批次参考梯度存储，内存高；RLOO 无此需求

评分

  • 新颖性: ⭐⭐⭐⭐ NES+低秩 Kronecker+RLOO 的组合首次出现，preconditioning-as-distribution-learning 视角新颖
  • 实验充分度: ⭐⭐⭐⭐ 覆盖 66M 到 13B 多种模型、GLUE/SuperGLUE 多个任务，但缺少生成任务评估
  • 写作质量: ⭐⭐⭐⭐ 数学推导严谨完整，但符号较多，读起来有一定门槛
  • 价值: ⭐⭐⭐⭐ 在内存受限场景下提升 ZO 微调性能有实际意义，但随着 LoRA 的普及，ZO 微调的应用窗口在收窄