精读笔记

Problem Setting

这篇论文解决的不是一般的 learning rate schedule tuning，而是大模型预训练中更实际的 post hoc horizon extension：训练到某个预算后，发现还值得继续训，但不希望重新 warmup、重调 peak LR，或丢掉大量已有计算。

关键矛盾是：continuation 需要 schedule 在中途 checkpoint 上保持“还活着”，而 endpoint model 又需要低学习率收尾以获得好的 last-iterate loss。Cosine 把整条曲线绑定到预设终点，训练到末端后 LR 基本耗尽；WSD 通过 stable phase 保留 continuation anchor，但 stable phase 的 peak LR 仍是按原 horizon 调的。也就是说，WSD 只解决了“从哪里继续”的问题，没有解决“以什么 LR 继续仍然合适”的问题。

真正困难点在于 horizon 不确定性和 last-iterate optimality 之间有理论冲突：完全 horizon-free 的 last iterate 会付出额外 log / polylog 代价；但完全 known-horizon 的 schedule 又不适合训练中途延长。WSqD 的问题设定正是夹在这两者之间。

Motivation

已有路线缺的不是又一个 decay shape，而是一个能把“训练可延长性”和“最终模型质量”解耦的 schedule。Cosine 的 endpoint dependence 太强；WSD 的 flat stable phase 太粗糙，隐含假设是同一个 peak LR 可以长时间保持最优；power-law scheduler 虽然接近本文方向，但更多来自 empirical scaling fit；schedule-free / averaging 类方法则改变优化器或训练流程，不再是一个纯 LR schedule。

作者的核心观察是：如果 WSD 的问题在于 peak LR 随 horizon 漂移，那么 base phase 就不该是常数。一个随时间缓慢下降、但不依赖最终 T 的 base LR，可以自然吸收“长训练需要更小 LR”的趋势。stochastic convex optimization 中的 1/sqrt(t) 正好提供了一个非经验拟合的候选。缺口是：纯 1/sqrt(t) 对最后迭代点不是最优，因此还需要末端 linear decay。

Core Idea

WSqD 的真正想法是：把 schedule 分成“可继续的 base trajectory”和“提交模型前的 terminal refinement”。base trajectory 用 shifted inverse-square-root，因此每个 pre-decay checkpoint 都在同一条 horizon-free 曲线上；当决定要停时，再用 linear cooldown 把 last iterate 拉到更好的 endpoint。

和 WSD 的本质区别不是多了一个 decay，而是 stable phase 的 inductive bias 变了。WSD 假设长期恒定 LR 是好的探索 / 训练状态；WSqD 假设训练越久，有效 LR 应随步数下降，并且这种下降规律可以跨 horizon 复用。这让 LR tuning 从“为每个总 token budget 找 peak”变成“为一条 anytime base curve 找 scale”。

这个设计也绕开了严格 anytime lower bound：WSqD 不声称整个序列都 horizon-free，它只要求 base phase horizon-free；最后 cooldown 明确使用 horizon 信息。这个折中很实用，也比纯理论上的 anytime 更贴近预训练 workflow。

Method

1. Shifted inverse-square-root base：解决 WSD flat phase 在长 horizon 下 peak LR 过大的问题。它让 LR 以 1/sqrt(t+T0) 下降，核心变化是把 horizon-dependent retuning 变成 time-dependent decay。

2. Shift parameter T0：解决早期 LR 过高和曲线过陡的问题。它不是决定性机制，更像把 1/sqrt(t) 平滑成一个可用于 LLM pretraining 的工程化版本。经验上 T0 太小会早期衰减过快；中等以上后不太敏感。

3. Final linear cooldown：解决 base schedule 的 last-iterate 缺陷。纯 inverse-square-root 对平均迭代或 suffix 可能好，但最后一个点会有 log overhead；linear cooldown 的作用是把 suffix 中已有的好 checkpoint 质量传播到最终 iterate。

4. Continuation protocol：真正重要的是训练延长时不重新调 base LR，而是沿同一条 base trajectory 继续。只有当需要产出模型时才启动 cooldown。这个 protocol 才是 WSqD claim 的核心测试条件。

Key Insight / Why It Works

最核心的 insight 是：大模型 continued training 的 LR 问题不是“衰减或不衰减”，而是“何时引入 horizon 信息”。Cosine 从一开始就引入 horizon，所以不适合事后延长；WSD 把 horizon 信息推迟到 cooldown，但 peak LR 仍被短 horizon tuning 污染；WSqD 进一步把 base phase 的 scale dynamics 设计成 horizon-free，只在最后阶段使用 T。

方法可能有效的主要原因是 better inductive bias，而不是新优化器能力。1/sqrt base 与经验 power-law LR scaling 非常接近，天然表达了“训练预算越长，合适 LR 越低”的规律；这比 WSD 的 constant LR 更符合长程预训练中的更新尺度需求。这里的贡献更像 schedule-level scaling bias，而不是新的 representation learning 机制。

linear cooldown 是必要但偏辅助：它保证 endpoint model 好，理论上消除 last-iterate log overhead；但 empirical gain 相比 WSD 更可能主要来自 base phase 的逐渐降 LR，而不是 cooldown 本身，因为 WSD 也有 cooldown。真正有价值的是 base LR optimum 在多个 horizon 上更稳定。

需要直接说：这不是证明了 LLM 训练等价于 convex optimization。理论提供的是设计启发和 last-iterate 形式化解释，不是对 AdamW Transformer 的严格解释。Appendix B 关于 Adam/Muon 的 mirror-descent view 也只是说明 omitted scale factor 在实验中较稳定；文中未充分说明为什么这些 scale factor 应该稳定。

Relation To Prior Work

最接近的路线是 WSD、power scheduler、schedule-free / averaging 方法，以及 convex last-iterate SGD theory。

相对 WSD，WSqD 的实质创新是替换 constant stable phase 的假设。WSD 的 continuation anchor 是结构上的，但 peak LR 仍 horizon-specific；WSqD 让 base LR 本身沿 horizon-free curve 下降，因此 continuation 不只是“能继续”，而是“继续时 LR 尺度更合理”。

相对 Shen et al. power scheduler，WSqD 并不完全新颖：两者都在 WSD skeleton 内加入 decreasing base，且 exponent 接近 -1/2。差别在于 WSqD 给出 convex optimization 推导，并用 shift 而不是 cap 处理早期大 LR。实质上，WSqD 可以看成给近似 power-law schedule 提供了一个理论化、简化版解释。

相对 schedule-free optimizer / averaging / checkpoint merging，WSqD 更保守：不改 optimizer，不引入 averaging state，不要求 merge checkpoint。它属于 schedule-shape 这条技术谱系，而不是 optimizer-level anytime training。

相对 last-iterate SGD theory，本文的新意是把 known-horizon cooldown 和 anytime base 拼在一起，并明确服务于 LLM continuation workflow。理论部分的结构本身不算特别激进，更多是把已有 suffix averaging、block propagation、linear decay last-iterate 分析重组到 WSqD schedule 上。

Dataset / Evaluation

实验主要覆盖 LLM pretraining continuation，而不是下游任务泛化。主设置是 213M LLaMA-style model，SlimPajama，AdamW，短上下文；附录补了 OpenWebText2、三 seed、124M 模型。这些实验足以支持一个有限 claim：在小中规模 controlled pretraining 中，WSqD 的 base LR 比 WSD 更适合跨 horizon 复用。

但 evaluation 还没有充分验证更强的 claim。第一，规模偏小，距离真正 large model training 仍远；第二，数据场景主要是通用预训练 corpus，没有覆盖真实 mid-training 中常见的数据分布切换、质量重加权、domain continuation；第三，比较指标主要是 validation loss，没有显示 downstream 或 long-context impact；第四，power scheduler 在主文比较中略优于 WSqD，说明 WSqD 的 empirical dominance 并不绝对。

因此实验支持“WSqD 是比 WSD 更稳的 continuation schedule”，但还不能支持“它是大规模训练中普遍最优的 horizon-free schedule”。

Limitation

理论上，WSqD 的 guarantee 依赖 convex、bounded subgradient、bounded domain、mirror descent、unbiased oracle。这些假设与 AdamW 训练 Transformer 的实际机制差距很大。文中承认这一点，但这意味着理论主要是 design rationale，而不是实际训练的解释闭环。

方法上，WSqD 并非严格 horizon-free。最终 cooldown 仍要知道 T；如果用户不断延长，就需要丢弃之前 cooldown 部分，从 pre-decay checkpoint 继续。它把 horizon dependence 从整条曲线转移到了末端，而不是完全消除。

增益归因不完全清晰。因为 WSD 和 WSqD 都有 linear decay，差异主要来自 base phase；但尚不清楚收益是 inverse-square-root 的理论结构，还是简单地因为逐步降低 LR 比 constant LR 更匹配这些实验的 token budget。可能主要来自 scaling-shape calibration，而不是更深的优化原理。

泛化上限未确定。T0 推荐为 pilot horizon 是经验规则；在更大模型、更长训练、更大 batch、不同 tokenizer / context / optimizer 下是否稳定，文中未充分说明。对 Muon、AdamW adaptive state、momentum 和 weight decay 的交互也没有严格分析。

此外，evaluation 没有覆盖真实 production continuation 的复杂因素：数据分布改变、annealing data、quality upsampling、checkpoint selection、multi-stage objective。WSqD 在这些场景中可能仍有用，但当前证据不足。

Takeaway

• 1. 最值得记住的是“base phase horizon-free，terminal phase horizon-aware”这个拆分。
• 它比争论 cosine / WSD / power-law 的具体形状更重要。
• 2. WSD 的 flat stable phase 可能不是一个好的默认 continuation prior。
• 长训练中，LR 逐步下降可能是更自然的 scaling bias。

一句话总结

WSqD 是一篇把 WSD-style continuation 与 inverse-square-root anytime base 结合起来的 learning-rate schedule 论文，真正贡献在于把 horizon dependence 限制到最终 cooldown，从而给 continued pretraining 提供了一个更可迁移的 LR scaling bias。