AdaLRS: Loss-Guided Adaptive Learning Rate Search for Efficient Foundation Model Pretraining¶
会议: NeurIPS 2025
arXiv: 2506.13274
代码: 无
领域: 优化 / 基础模型预训练
关键词: learning rate search, online optimization, loss velocity, foundation model pretraining, cosine scheduler
一句话总结¶
提出AdaLRS，一种即插即用的在线学习率搜索算法，通过监控损失下降速度（loss velocity）来自适应调整学习率，将学习率超参搜索的成本从多次独立训练降低到单次训练，实现~50%的训练成本节省。
背景与动机¶
学习率是基础模型预训练中最关键的超参数，但寻找最优学习率通常需要大量代理模型实验或多次独立训练。现有方法要么将搜索限制在小型代理模型上（可能与大模型不一致），要么需要许多独立运行。对于成本高昂的基础模型预训练（LLM/VLM），需要一种能在单次训练过程中自动找到最优学习率的方法。
核心问题¶
如何在单次训练运行中，在线地找到接近最优的学习率，且兼容现代学习率调度器（cosine、WSD）？
方法详解¶
整体框架¶
监控训练损失曲线的斜率（下降速度），当斜率衰减时上调学习率，通过k步窗口最小二乘法估计斜率，并设置验证机制防止过度调整。
关键设计¶
- 核心理论洞察：
- 证明训练损失 \(L(\eta)\) 和损失下降速度 \(V(\eta)\) 都是关于学习率的凸函数，且共享同一最优学习率 \(\eta^*\)
- 理论推导：SGD 下 \(\mathbb{E}[L_{t+1}-L_t] \approx -\eta\|\nabla L_t\|^2 + \frac{C_{Lip}}{2}\eta^2\|\nabla L_t\|^2\)，最优 \(\eta^* = 1/C_{Lip}\)
-
这意味着可以通过优化可在线估计的斜率来间接优化损失
-
在线调整算法：
- 每 k 步窗口用最小二乘法估计损失斜率 \(v(\eta)\)
- 调整规则：试探性放大 \(\alpha'\eta\)，比较 \(v(\alpha'\eta)\) 与 \(v(\eta)+2e\) 的关系决定保持/放大/缩小
- 缩放因子递减：\(\alpha' = \max(\lambda^t\alpha, 1)\), \(\beta' = 1/\max(\lambda^t\beta, 1)\)，\(\lambda=0.99\) 默认
-
搜索仅在训练的 [0.1, 0.4] 阶段进行，后期交还给标准调度器
-
稳定性机制：
- 回退策略（backtracking）：放大失败时恢复模型和优化器状态，防止参数被破坏性更新污染
- 早停：若损失超过历史最大值，停止放大
-
边界条件：连续2窗口损失上升则缩小学习率
-
收敛保证：
- Theorem 2.1: \(\lim_{t\to\infty}\mathbb{P}(|\eta_t-\eta^*|<e)=1\)（几乎必然收敛到 \(e\)-邻域）
- Theorem 2.4: 几何误差衰减 \(|\eta_{t+k}-\eta^*|\leq\gamma|\eta_t-\eta^*|\)，复杂度 \(O(\log R)\)
实验关键数据¶
主实验：LLM 预训练¶
| 场景 | 模型 | 训练损失 (baseline→AdaLRS) | PPL (val) |
|---|---|---|---|
| 小LR起步 | Qwen2.5-1.5B | 2.56→更优 | 12.66→改善 |
| 大LR起步 | Qwen2.5-1.5B | 5.21→改善 | 183.94→改善 |
| 小LR起步 | Qwen2.5-7B | 2.38→更优 | 10.61→改善 |
- 小 LR 起步：~50% 训练成本节省到达基线损失
- 大 LR 起步：>30% 训练成本节省
- 训练预算：120B-160B tokens, 10000-20000 910B NPU hours
VLM 预训练（2B SAIL-VL, 7个基准）¶
| LR 设置 | 平均分 (基线/AdaLRS) |
|---|---|
| Fit LR | 56.16 / 55.80 |
| Small LR | 57.34 / 53.77 |
| Large LR | 48.96 / 47.67 |
消融实验¶
| 配置 | 说明 |
|---|---|
| 有/无回退策略 | 回退策略关键：无回退时大LR损失在5.0-5.2高位不降 |
| 超参 (α,β,λ) 组合 | (3/2,0.99), (2/1.67,0.99), (1.5/1.43,0.99) 等均有效，鲁棒性强 |
| Cosine vs WSD 调度器 | 两者均兼容，验证即插即用特性 |
| 继续预训练 | 小LR有效 (0.8851→0.8286)，大LR受限于灾难性遗忘 |
亮点¶
- 理论优雅：损失和损失斜率共享最优LR的凸性证明，为方法提供了坚实基础
- 实用性强：即插即用，兼容主流调度器，无需修改训练pipeline
- 多场景验证：LLM/VLM、不同规模、不同初始LR、预训练/继续训练
- 显著节省：50%训练成本节省对大模型预训练意义重大
局限性 / 可改进方向¶
- 初始 LR 过大时方法失效——参数已被破坏性更新污染，仅靠降 LR 无法修复
- 仅保证收敛到 \(e\)-邻域而非精确最优，\(e\) 的大小取决于估计质量
- 继续预训练中大 LR 下无法消除灾难性遗忘
- "Fit LR" 基线通过 grid search 确定，可能引入比较偏差
- 对 batch size 变化的适应性分析不足
相关工作与启发¶
- vs Chinchilla-style LR 搜索：需要数百次独立训练建立 scaling law；AdaLRS 单次即可
- vs μP/μTransfer：从代理模型转移超参，但代理搜索本身仍昂贵；AdaLRS 完全不需要代理模型
- vs 标准 cosine 调度：AdaLRS 在 cosine 基础上叠加自适应调整，是互补而非替代关系
- 对实际大模型训练有直接工程价值：减少因 LR 选择不当导致的计算浪费
评分¶
- 新颖性: ⭐⭐⭐⭐ 损失斜率凸性的理论洞察新颖
- 实验充分度: ⭐⭐⭐⭐⭐ 多模型、多场景、消融充分
- 写作质量: ⭐⭐⭐⭐ 理论-实践结合清晰
- 价值: ⭐⭐⭐⭐⭐ 对大模型预训练实践有重要影响