Taking the Road Less Scheduled with Adaptive Polyak Steps¶
会议: NeurIPS 2026
arXiv: 2511.07767
领域: 最优化
关键词: 最优化算法, Schedule-Free, Polyak步长, 自适应学习率, 大模型预训练
一句话总结¶
本文针对 Schedule-Free 优化器仍需人工网格搜索基准学习率的痛点,通过最小化参数距最优解距离的上界,推导了适配 Schedule-Free 迭代平均机制的自适应 Polyak 步长,彻底摆脱了预设训练步长与学习率调优的负担。
研究背景与动机¶
深度神经网络训练在传统实践中极其依赖预先设定的学习率调度策略(如余弦退火、线性衰减或阶梯下降)。这类调度方案的核心缺陷在于必须提前固定总训练步数 \(T\);一旦训练需要提前终止或后续追加训练步数,原有的调度曲线便不再最优。为了打破这一限制,Defazio 等人在 2024 年提出了 Schedule-Free 优化器,通过精巧的评估序列插值与迭代平均机制,在无需预设训练时长 \(T\) 的前提下达到了与精调余弦退火相当的最优收敛速率。
然而,现有 Schedule-Free 方法并未实现真正意义上的「免调参」。理论上其最优常数学习率 \(\gamma = D / (G\sqrt{T})\) 仍然显式依赖未知的初始距离 \(D = \|\bm{x}_0 - \bm{x}_\star\|\)、目标函数 Lipschitz 常数 \(G\),乃至方法本意试图消除的训练时长 \(T\)。在工程落地中,从业者不得不针对每个任务、模型架构和权重衰减配置运行大规模的基准学习率网格搜索,仅在预训练大语言模型时就会消耗数百 GPU 小时的算力成本。
经典的 Polyak 步长机制通过当前采样的函数值与梯度模长在每一步自适应计算标量步长,完全避免了依赖问题未知常数。然而,经典 Polyak 步长针对标准 SGD 设计,而 Schedule-Free 优化器维持了三个交织序列(梯度评估序列 \(\bm{y}_t\)、更新基序列 \(\bm{z}_t\) 和平均输出序列 \(\bm{x}_t\)),梯度求值点与参数更新点分离,直接套用经典公式会破坏理论收敛保证。本文的核心 idea 是:将参数距离最优解的动态递推分解到评估序列 \(\bm{y}_t\) 上,利用凸性与动量插值推导包含动量修正项的封闭解 Polyak 步长,并通过安全防护机制与 Adam 对角预条件化,实现无需预设训练周期、无需网格搜索基础学习率的端到端自适应优化。
方法详解¶
整体框架¶
算法围绕 Schedule-Free 的三序列架构展开:在第 \(t\) 步,算法首先结合历史基序列 \(\bm{z}_{t-1}\) 与迭代平均序列 \(\bm{x}_t\) 生成前向与梯度评估点 \(\bm{y}_t\);随后采集随机批次 \(\bm{\zeta}_t\),计算损失值 \(f_{\bm{\zeta}_t}(\bm{y}_t)\) 与随机梯度 \(\nabla f_{\bm{\zeta}_t}(\bm{y}_t)\);核心模块依据当前损失间隙与动量修正内积自适应计算标量步长 \(\gamma_t\)(在 Adam 变体中结合对角预条件子与动态梯度模长上限);最后利用该步长更新基序列 \(\bm{z}_t\),并通过权重系数 \(c_{t+1}\) 递推更新最终输出序列 \(\bm{x}_{t+1}\)。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["当前迭代状态<br/>基序列 z_{t-1} 与输出平均 x_t"] --> B["动量插值评估点构建<br/>y_t = (1-β)z_{t-1} + β x_t"]
B --> C["随机批次前向与梯度反传<br/>采样 ζ_t 计算 f_{ζ_t}(y_t) 与 ∇f_{ζ_t}(y_t)"]
C --> D["自适应动量修正 Polyak 步长计算<br/>分子融合损失间隙与动量内积项"]
D --> E["分母安全防护约束与预条件加权<br/>max(||g_t||_{D_t^{-1}}^2, M_t)"]
E --> F["基序列投影更新<br/>z_t = z_{t-1} - γ_t D_t^{-1} ∇f_{ζ_t}(y_t)"]
F --> G["输出序列无调度迭代平均<br/>x_{t+1} = (1-c_{t+1}) x_t + c_{t+1} z_t"]
关键设计¶
1. 动量解耦与上界极小化:推导适配 Schedule-Free 机制的闭式 Polyak 步长
在 Schedule-Free SGD 中,参数更新发生在基序列 \(\bm{z}_t = \bm{z}_{t-1} - \gamma_t \nabla f_{\bm{\zeta}_t}(\bm{y}_t)\),但梯度的评估位置却是 \(\bm{y}_t\)。展开基序列至最优解 \(\bm{x}_\star\) 的欧氏距离平方,交叉项为 \(-2\gamma_t \langle \nabla f_{\bm{\zeta}_t}(\bm{y}_t), \bm{z}_{t-1} - \bm{x}_\star \rangle\)。若直接对其优化,极小化步长将显式包含未知的全局最优解 \(\bm{x}_\star\)。
本文通过插值关系 \(\bm{y}_t = (1-\beta)\bm{z}_{t-1} + \beta\bm{x}_t\) 将其严谨解耦,重写出 \(\bm{z}_{t-1} = \bm{y}_t + \beta(\bm{z}_{t-1} - \bm{x}_t)\)。代入内积后,第一部分通过 \(f_{\bm{\zeta}_t}\) 的凸性放缩为下界 \(f_{\bm{\zeta}_t}(\bm{y}_t) - f_{\bm{\zeta}_t}(\bm{x}_\star)\),第二部分则转化为完全已知的动量修正项。对距离上界关于 \(\gamma_t \ge 0\) 进行二次极小化,由此导出 Oracle 版本的 Schedule-Free Polyak 步长: $\(\gamma_t = \frac{\left[ f_{\bm{\zeta}_t}(\bm{y}_t) - f_{\bm{\zeta}_t}(\bm{x}_\star) + \beta \langle \nabla f_{\bm{\zeta}_t}(\bm{y}_t), \bm{z}_{t-1} - \bm{x}_t \rangle \right]_+}{\|\nabla f_{\bm{\zeta}_t}(\bm{y}_t)\|^2}\)$ 在工程实现中,利用恒等式 \(\beta(\bm{z}_{t-1} - \bm{x}_t) = \bm{z}_{t-1} - \bm{y}_t\),分子中的动量修正内积可直接简化为 \(\langle \nabla f_{\bm{\zeta}_t}(\bm{y}_t), \bm{z}_{t-1} - \bm{y}_t \rangle\),计算时完全无需额外保留或加载 \(\bm{x}_t\) 显存。当 \(\beta = 0\) 时,\(\bm{y}_t = \bm{z}_{t-1}\),该式无缝退化为经典 SGD 的 SPS+ 规则。
2. 安全防护与已知下界替换(SF-SPS-safe):摆脱样本最优值依赖与步长数值爆炸
实际深度学习任务通常无法预先获知样本级最优质 \(f_{\bm{\zeta}_t}(\bm{x}_\star)\)。针对该痛点,受无平滑优化中安全防护步长启发,设计引入两重改进:第一,将未知项 \(f_{\bm{\zeta}_t}(\bm{x}_\star)\) 替换为任意可获取的理论下界 \(\ell_{\bm{\zeta}_t}^*\)(对于交叉熵与 MSE 等非负损失直接取 0);第二,在分母中引入常数阈值 \(M > 0\),采用 \(\max\{\|\nabla f_{\bm{\zeta}_t}(\bm{y}_t)\|^2, M\}\) 防止在梯度范数极小的平坦区出现步长数值爆炸: $\(\gamma_t = \frac{\left[ f_{\bm{\zeta}_t}(\bm{y}_t) - \ell_{\bm{\zeta}_t}^* + \beta \langle \nabla f_{\bm{\zeta}_t}(\bm{y}_t), \bm{z}_{t-1} - \bm{x}_t \rangle \right]_+}{\max\left\{ \|\nabla f_{\bm{\zeta}_t}(\bm{y}_t)\|^2, M \right\}}\)$ 与通过截断总步长 \(\min\{\gamma_t, \gamma_{\max}\}\) 的经典方法不同,该防护仅作用于分母,保留了分子损失差距主导的动态调节特性。理论证明在凸 Lipschitz 条件下,该步长能达到 \(\mathcal{O}(1/\sqrt{t} + \sigma^2)\) 的近最优收敛域,其中 \(\sigma^2\) 为下界松弛方差,在现代过参数化插值区间(\(\sigma^2 = 0\))可实现精准收敛。
3. 对角预条件矩阵扩展(SF-Adam-SPS):无缝兼容大语言模型异构曲率优化
Transformer 等现代架构参数曲率各向异性极强,纯欧氏几何下的梯度更新易受病态曲率牵制。为此,作者将框架推广至正定预条件矩阵 \(\bm{D}_t \succ 0\) 诱导的度量空间,在基序列更新中采用 \(\bm{z}_t = \bm{z}_{t-1} - \gamma_t \bm{D}_t^{-1} \nabla f_{\bm{\zeta}_t}(\bm{y}_t)\)。在能量上界推导中,将欧氏范数替换为加权诱导范数 \(\|\bm{v}\|_{\bm{D}_t}^2 = \bm{v}^\top \bm{D}_t \bm{v}\)。
在极小化求解后,步长分母的梯度模长自然变为预条件逆范数 \(\|\nabla f_{\bm{\zeta}_t}(\bm{y}_t)\|_{\bm{D}_t^{-1}}^2\)。实例化为 Adam 的二阶矩估计时,\(\bm{D}_t = \operatorname{diag}(\sqrt{\bm{v}_t / (1-\beta_2^{t+1})} + \epsilon)\),即可得到开箱即用的 SF-Adam-SPS 形式: $\(\gamma_t = \frac{\left[ f_{\bm{\zeta}_t}(\bm{y}_t) - \ell_{\bm{\zeta}_t}^* + \beta \langle \nabla f_{\bm{\zeta}_t}(\bm{y}_t), \bm{z}_{t-1} - \bm{x}_t \rangle \right]_+}{\max\left\{ \|\nabla f_{\bm{\zeta}_t}(\bm{y}_t)\|_{\bm{D}_t^{-1}}^2, M \right\}}\)$ 此时梯度更新方向和步长分母尺度均经过一阶、二阶曲率矫正,无需改动模型反向传播流。
4. 基于指数移动平均的自适应防护(EMA Safeguard):消除固定防护阈值的人工配置
固定阈值 \(M\) 虽然保证了数值安全性,但若 \(M\) 选取得过大,分母会退化为常数主导;选取得过小,则在初期震荡时无法有效抑制瞬态过大步长。为了消除该超参数的敏感性,本文提出基于梯度二阶范数的指数移动平均(EMA)自适应跟踪策略: $\(M_t = \beta_M M_{t-1} + (1 - \beta_M) \|\nabla f_{\bm{\zeta}_t}(\bm{y}_t)\|^2\)$ 初始值设为首步梯度范数平方 \(M_0 = \|\nabla f_{\bm{\zeta}_0}(\bm{y}_0)\|^2\),衰减系数推荐采用极度平稳的通用默认值 \(\beta_M = 0.99\)。这一机制使分母安全界动态跟随训练各阶段的梯度尺度自动伸缩,在训练初期梯度剧烈变化时提供柔性约束,而在后期平缓收敛时保持自适应缩放。
一个完整示例¶
假设我们使用 SF-Adam-SPS-safe 训练单层变换器模型,动量系数设为 \(\beta = 0.9\),下界 \(\ell^* = 0\)。在某一步 \(t\): 1. 生成评估状态:当前基参数 \(\bm{z}_{t-1}\) 与平均参数 \(\bm{x}_t\) 按照比例加权,得到前向计算点 \(\bm{y}_t = 0.1 \bm{z}_{t-1} + 0.9 \bm{x}_t\)。 2. 前向反传计算:在批次样本上计算得到交叉熵损失 \(f_{\bm{\zeta}_t}(\bm{y}_t) = 3.50\);反向传播得到当前梯度 \(\bm{g}_t = \nabla f_{\bm{\zeta}_t}(\bm{y}_t)\)。 3. 计算动量修正内积:利用向量差 \(\bm{z}_{t-1} - \bm{y}_t\),计算内积项 \(\langle \bm{g}_t, \bm{z}_{t-1} - \bm{y}_t \rangle\)。若梯度方向与基参数调整方向一致,该项为正(如 \(+0.12\)),有效分子项扩充为 \(3.50 + 0.12 = 3.62\);若方向相悖导致该和为负,ReLU 操作直接将其截断归零以确保稳定性。 4. 计算步长分母:获取当前 Adam 预条件矩阵 \(\bm{D}_t\),计算预条件梯度能量 \(\|\bm{g}_t\|_{\bm{D}_t^{-1}}^2 = 0.45\)。当前自适应安全防护值 \(M_t = 0.50\),分母取最大值 \(\max\{0.45, 0.50\} = 0.50\),从而自适应产出标量步长 \(\gamma_t = 3.62 / 0.50 = 7.24\)。 5. 双序列更新:基参数更新为 \(\bm{z}_t = \bm{z}_{t-1} - 7.24 \bm{D}_t^{-1} \bm{g}_t\);随后按权重 \(c_{t+1}\) 将 \(\bm{z}_t\) 合并进输出参数 \(\bm{x}_{t+1}\)。
损失函数 / 训练策略¶
在预训练工程实践中,为了完全对齐业界大模型训练规范,作者采用如下策略组合: - 权重均值更新策略:将理论证明采用的调和平均权重 \(c_{t+1} = 1/(t+1)\),替换为 Defazio 建议的高效工程启发式权重 \(c_{t+1} = \gamma_t^2 / \sum_{i=1}^t \gamma_i^2\),有效增强了训练后期优质迭代点的表征比重。 - 预热与衰减机制:搭配 2,000 步的 Warmup 策略;随后保持 Schedule-Free 恒定运行模式,无需预设任何退火周期。 - 超参数配置:梯度裁剪对于 Schedule-Free 与 Polyak 变体均全程禁用(仅 AdamW 基线保留 1.0 的裁剪阈值);权重衰减设为 0.1,Dropout 设为 0.0。
实验关键数据¶
主实验¶
在 FineWeb 10B 数据集上,针对 GPT-2 Small (124M) 与 Medium (355M) 进行了严谨的大规模自回归语言建模预训练测试(总训练标记量达 10.25B,全局批大小 480 个序列共 491,520 tokens,优化步数 20,856 步)。
| 优化器方案 | 学习率搜索空间 / 策略 | 调优所需 GPU 耗时 (4×H100) | 学习率敏感度表现 | 最终验证损失 (GPT2-small) | 最终验证损失 (GPT2-medium) |
|---|---|---|---|---|---|
| AdamW (Cosine Decay) | 6点网格搜索 \(\{5\text{e-}6 \sim 5\text{e-}1\}\) | ~220 GPU 小时 | 极敏感,最优范围窄 | 3.32 | 3.12 |
| Schedule-Free AdamW | 6点网格搜索 \(\{5\text{e-}6 \sim 5\text{e-}1\}\) | ~195 GPU 小时 | 敏感,偏离最优值迅速发散 | 3.29 | 3.08 |
| SF-Adam-SPS-safe (固定 M=10) | 无需搜索(单次运行) | ~21 GPU 小时 (4×A100) | 完全平坦,免调参 | 3.26 | 3.05 |
| SF-Adam-SPS-safe (EMA M) | 无需搜索(单次运行) | ~21 GPU 小时 (4×A100) | 完全平坦,免调参 | 3.26 | 3.05 |
消融实验¶
在模型知识蒸馏任务(将 gpt2-medium 345M 教师模型蒸馏至 67.7M 4层变换器学生模型,测试集 tiny_shakespeare,bfloat16 + FlashAttention)与预训练超参数消融中,对比了不同组件配置的影响:
| 实验配置 / 优化器设置 | 核心超参数配置 | 验证集最佳 Loss | 学习率稳健性 (鲁棒区间) | 机制说明 |
|---|---|---|---|---|
| SGD with Momentum | 扫基准学习率 | ~4.85 | 窄 (\(\pm 0.5\) 数量级) | 纯欧氏一阶动量更新 |
| Standard Adam | 扫基准学习率 | ~4.12 | 较窄 (\(\pm 1\) 数量级) | 标准矩估计对角预条件 |
| Schedule-Free Adam | 扫基准学习率 | ~4.05 | 中等 | 依赖人工调参基准 \(\gamma\) |
| Iams-Adam (Gower 2025) | 扫最大截断 \(\gamma_{\max}\) | ~3.98 | 宽 | 理想化随机 Polyak 步长 |
| SF-Adam-SPS+ (Ours Oracle) | 自适应步长 + 扫 \(\gamma_{\max}\) | 3.92 | 极宽(最优损失显著领先) | 教师损失近似 \(f(x_\star)\),充分释放动量优势 |
| SF-Adam-SPS-safe (Const M) | \(M = 10, \ell^* = 0\) | 3.26 (GPT2-S) | 免调参 | 常数分母防护抑制小梯度震荡 |
| SF-Adam-SPS-safe (EMA M) | \(\beta_M = 0.99, \ell^* = 0\) | 3.26 (GPT2-S) | 免调参 | 动态追踪逆梯度模长,行为更加精细自适应 |
关键发现¶
- 算力成本呈数量级削减:在标准预训练流程中,AdamW 与原版 Schedule-Free Adam 必须执行 6 组学习率遍历(各耗费约 220 与 195 GPU-hours);而 SF-Adam-SPS-safe 无需任何网格搜索,单次运行即收敛,整体算力开销缩减近 90%。
- 泛化性能反超精调基线:在 GPT2-small 与 GPT2-medium 预训练中,自适应步长产出的最终验证损失(3.26 与 3.05)均优于精心扫选得到的 Schedule-Free AdamW(3.29 与 3.08)以及标准 AdamW(3.32 与 3.12),显示出自适应步长在逃离局部极小与平坦极小值探索上的优势。
- 动态曲线自发展现退火特征:追踪实际步长轨迹发现,SF-Adam-SPS-safe 在训练初期随着损失差距快速维持大步长更新,在模型收敛后期随损失下降平稳自动衰减,在没有任何外部余弦/衰减调度器控制的前提下,天然演化出了优雅的降温退火轨迹。
亮点与洞察¶
- 严谨优美的数学闭式解:巧妙利用 Schedule-Free 中 \(\bm{z}_{t-1} = \bm{y}_t + \beta(\bm{z}_{t-1} - \bm{x}_t)\) 的线性代数等式,将不可计算的最优点内积转化为凸性下界与已知的动量修正项,将古典优化理论的 Polyak 优雅延展至现代前沿架构。
- 一石二鸟的工程实现精简:分子中的动量内积项被恒等化简为 \(\langle \nabla f(\bm{y}_t), \bm{z}_{t-1} - \bm{y}_t \rangle\),从而完全不需要在步长计算步骤中从内存加载或缓存加权历史输出 \(\bm{x}_t\),零显存开销即可实现无缝嵌入。
- 从根本上解决调度生命周期绑定问题:传统余弦退火强绑定训练寿命 \(T\),Schedule-Free 解耦了 \(T\) 却仍绑定常数 \(\gamma\),本文将两者同时解放,使大模型持续训练(Continual Pretraining)与动态弹性扩训无需重置学习率。
局限与展望¶
- 理论收敛假设限定于凸条件:当前 Theorems 3.1 和 3.2 的严格证明均建立在凸 Lipschitz 目标函数假设之上,对于深度神经网络所处的非凸复杂损失曲面,理论分析尚未涵盖全局极值收敛性。
- 非插值区域存在稳态误差邻域:对于存在不可消除随机方差的数据分布(\(\sigma^2 > 0\)),理论保证存在由 \(\mathcal{O}(\sigma^2)\) 控制的极限邻域,虽然在大模型强过参数化插值情形下表现良好,但在强噪声小模型场景中可能需要更精细的方差缩减配合。
相关工作与启发¶
- vs Schedule-Free SGD (Defazio et al., 2024):原版方法通过平均机制消除了 \(T\) 的依赖,但保留了最难调的基准学习率 \(\gamma\)。本文将 Polyak 步长嵌入该框架,填补了其自适应理论闭环,彻底实现完全免调参。
- vs SPS+ / Safeguarded SPS (Garrigos et al., 2023; Oikonomou & Loizou, 2026):既有随机 Polyak 步长均建立在标准单序列 SGD 或重球动量上,无法适配 Schedule-Free 评估点与更新点分离的特殊机制;本文建立了普适的映射理论,并在 \(\beta = 0\) 时严谨包含前者为特例。
- vs Iams-Adam (Gower et al., 2025):Iams 属于需要精调截断阈值 \(\gamma_{\max}\) 的理想化 Polyak 动量模型;本文给出了无需截断、完全基于损失下界与分母安全防护的自洽方案。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次将 Polyak 自适应步长与 Schedule-Free 架构融合,推导了严格的动量修正项并给出随时随地(anytime)收敛界。
- 实验充分度: ⭐⭐⭐⭐☆ 覆盖黑盒知识蒸馏与十亿级别真实预训练(FineWeb 10B),完整汇报了敏感性与算力开销。
- 写作质量: ⭐⭐⭐⭐⭐ 理论推导清晰完备,从 Oracle 到 Safeguard、从 SGD 到 Adam 的层层演进逻辑极具说服力。
- 价值: ⭐⭐⭐⭐⭐ 大幅降低大模型训练调参算力消耗,推动无调度自适应优化迈向实用化。