课程简介
CS336 Lecture 11:Scaling Laws 进阶——从经典理论到大规模训练实践
学习目标
完成本讲学习后,你应该能够:
- 理解经典 scaling laws(Kaplan、Chinchilla)在真实开源模型训练中的适用性与局限性
- 掌握两种主流的超参数缩放策略:基于拟合 scaling laws 的方法(DeepSeek 路线)与基于参数化不变性的方法(μP 路线)
- 理解 warmup-stable-decay(WSD)学习率调度的动机、实现方式及其对 scaling 实验的价值
- 了解现代 LLM 优化器(尤其是 Muon)的基本原理、规模依赖特性及大规模迁移的挑战
- 能够批判性地评估小规模实验结论向大规模迁移的有效性
课程主线
本讲围绕"如何从经典 scaling laws 走向真实大规模训练"这一核心问题展开,分为三个主要部分:
第一部分:开源模型中的 scaling laws 实践——通过 MiniCPM、DeepSeek、Qwen、Kimi K2、MiniMax 等模型的技术报告,展示 scaling laws 如何被实际应用于超参数选择、架构决策和训练策略设计。
第二部分:优化器与超参数的规模依赖——讨论学习率、批大小、初始化等超参数如何随规模变化,以及现代优化器(Muon 等)在大规模训练中的表现与挑战。
第三部分:μP(Maximal Update Parametrization)的理论基础与实践——从数学推导角度理解 μP 如何通过初始化和学习率的特殊设置实现超参数规模不变性。
第一部分:开源模型中的 Scaling Laws 实践
1.1 经典 Scaling Laws 的回顾与现状
课程将此前的内容(Kaplan 2020、Chinchilla 2022)定位为"经典扩展定律"。这些工作建立了三个核心结论:
- 损失与计算量之间存在幂律关系
- 最优 token/参数比约为 20:1(Chinchilla 结论)
- 临界批大小随损失目标呈幂律缩放
然而,课程强调:从 2022 年至今,scaling laws 领域的高质量论文数量有所减少,且主要贡献来自中国开源社区。这一现象的原因在于——核心的 scaling 方法论(Chinchilla 分析、学习率扫描)已经成为"默认知识",不再值得在论文中大幅讨论。
1.2 MiniCPM:μP + WSD 的工程实践
MiniCPM(2024)代表了 1-2B 参数规模的高性能小模型训练实践。其核心贡献在于将两类关键技术落地:
μP(Maximal Update Parametrization)初始化
MiniCPM 采用的 μP 方案包含以下具体操作:
- Embedding 输出缩放:对 embedding 层的输出进行缩放
- 残差连接缩放:使用层数的平方根缩放残差连接
- 初始化缩放:所有矩阵形状张量的初始化按 fan-in 与 fan-out 的比例进行缩放
- 逐参数学习率缩放:不同张量使用不同的学习率(非全局统一)
- LM Head 缩放:对语言模型头进行额外缩放
μP 的核心目标是:当模型规模变化时,最优学习率保持恒定。MiniCPM 的实验结果非常干净——在不同规模模型上扫描学习率,最小损失几乎都落在
课程解释:μP 并非为了直接训练大规模模型,而是为了通过小规模实验"一次性"确定超参数,从而避免大规模训练的反复试错。
WSD(Warmup-Stable-Decay)学习率调度
WSD 的学习率曲线形似一个大梯形:
- Warmup 阶段:固定步数(而非训练总步数的比例)内线性上升
- Stable 阶段:长时间保持恒定学习率(占训练的大部分)
- Decay 阶段:快速衰减至 0,通常占总训练时长的 10%-20%
WSD 的核心优势在于训练可复用性:
- 从 stable 阶段结束时的任意 checkpoint 可以继续训练(只需回滚到最后一个 stable checkpoint)
- 若要增加训练数据量,无需从头重新训练,只需继续 stable 训练并在最后做一次 decay
- 这使得沿着数据维度的 scaling 扫描变得极其高效——每次扫描只需 10% 左右的重训练成本
WSD 与余弦调度的性能对比:WSD 在 decay 阶段之前看起来"落后",但在衰减后能追平甚至偶尔超越精心调优的余弦调度。课程指出,许多实践者认为余弦调度在多数情况下略优,但 WSD 接近同等水平且具有更好的复用性。
MiniCPM 的 Scaling 实验
MiniCPM 复现了 Chinchilla 的方法 1 和方法 3(课程认为这两种方法最不可靠),并获得了合理的结果。值得注意的是,其联合拟合得到的指数与 Chinchilla 原文差异较大,作者据此声称应该用远多于 Chinchilla 比例的 token 训练——课程对此结论持保留态度。
MiniCPM 要点总结:
- μP 是稳定学习率随规模变化的实用技巧
- WSD 学习率是基础知识点,应被所有 LLM 从业者掌握
1.3 DeepSeek LLM:拟合 Scaling Laws 的路线
DeepSeek LLM(初代,非 MoE 版本)代表了与 MiniCPM 截然不同的策略:不尝试稳定超参数,而是直接拟合超参数的缩放规律。
方法:网格搜索 + 幂律拟合
- 在不同模型规模下,对学习率和批大小做密集网格搜索
- 固定计算量(FLOPs),观察最终损失
- 找到每个规模下的最优学习率和批大小
- 在多个规模下重复,拟合幂律关系
实验结果显示:最优批大小随非嵌入 FLOPs 呈清晰的线性(对数-对数坐标)缩放;最优学习率随 FLOPs 增大而降低。
课程指出了该方法的一个潜在问题:网格搜索存在量化误差,只有当网格点位置恰当时,拟合出的"线性关系"才可靠。学习率的拟合质量不如批大小理想,但最终模型训练成功,说明该方法的实用性。
Chinchilla 复现与预测
DeepSeek 同样使用 WSD 学习率进行 Chinchilla 风格分析(略有变体:使用两个 decay 阶段,这一做法后续未流行)。其 IsoFLOPs 扫描的曲线非常干净,成功复现了 Chinchilla 的数据-模型规模权衡,并获得了对实际训练模型的准确损失预测。
DeepSeek 路线要点:处理敏感超参数有两种思路——要么稳定它们(μP),要么拟合它们的缩放规律并跟随。DeepSeek 代表后者。
1.4 其他开源模型的 Scaling 实践
Qwen 2.5 / Qwen 3
Qwen 2.5 基本复用了 DeepSeek 的分析方法:做 scaling 实验找出最优批大小和学习率,并用 scaling law 进行预测。该方法对 MoE 模型同样适用。到 Qwen 3 时,这一流程已完全标准化——直接照搬上一代的做法。
Kimi K2
Kimi K2(课程系列中讨论的最新模型)聚焦于 MoE scaling laws。核心问题是:给定 FLOPs,不同稀疏度水平(激活参数占比)如何影响验证损失?结论是:网络越稀疏,给定 FLOPs 下的验证损失越好(符合直觉)。但关键价值在于定量化这一关系,从而理性选择稀疏度。Kimi K2 最终选择稀疏度 48,因为在该点 scaling 收益进入边际递减。
Hunyuan
另一开源 MoE 模型,做了类似分析——固定稀疏度水平,得到每个激活参数 96 个数据点的比例。这反映了 MoE 模型必须做"激活参数量 vs. 损失"的 scaling 分析这一趋势。
LLaMA 3
LLaMA 3 的 scaling 工作包含两个维度:
- 左侧:常规 IsoFLOP 分析(token/参数比与 Chinchilla 略有不同但相近)
- 右侧:将训练损失(log loss)与下游 benchmark 准确率联系起来,拟合 sigmoid 曲线
课程指出后者虽有系统性偏差,但基本结论成立:从 loss 到 accuracy 的映射与 sigmoid 形式的缩放关系一致——这为"优化 log loss 即优化下游性能"提供了实证基础。
MiniMax-01:用 Scaling Laws 做架构决策
MiniMax-01 通过对比三种注意力架构的 scaling 表现来做架构选择:
- Lightning attention(线性注意力)
- Softmax attention(完全注意力)
- 混合架构(两者结合)
结论是三种架构基本可相提并论,需要相似规模的模型,因此选择混合架构用于实际部署。这是"基于 scaling laws 做架构决策"的典型案例。
1.5 后训练时代的 Scaling Laws
课程回答了一个关键问题:当后训练(post-training)成为标准流程时,scaling laws 如何调整?
目前没有好的答案。核心困难在于:后训练有时会改变预训练的本质判断。有初步工作研究预训练中的覆盖度/多样性概念以预测后训练效果,但这仍处于早期阶段。缺少将后训练纳入完全整合的 scaling 分析的成熟框架。
第二部分:优化器与超参数的规模依赖
2.1 学习率与批大小的缩放规律
课程重点介绍 StepFun 团队的超参数调优预印本——目前最全面的相关研究之一。
不同缩放规律的分歧
| 来源 | 批大小缩放 | 学习率缩放 |
|---|---|---|
| OpenAI (Kaplan) | 基于最终损失(临界批大小) | — |
| DeepSeek | 基于计算量 | 基于计算量 |
| MiniCPM | 基于数据量 | 基于数据量 |
| StepFun | 基于数据量 | 基于数据量、模型规模 |
课程强调:这些公式的输入变量甚至各不相同,说明超参数缩放的本质尚无共识。因此 StepFun 的结论也不应视为金科玉律——但实验现象本身值得学习。
StepFun 的关键发现
通过高分辨率网格搜索(1B 参数、100B token 规模),StepFun 发现:
超参数景观是光滑凸函数:固定数据量,损失作为批大小和学习率的函数,每个切片都是近似光滑凸的——这保证网格搜索的可靠性。
最优批大小只依赖总数据量:不同模型规模的点(不同颜色)沿同一条趋势线排列,服从以数据量为变量的幂律。
最优学习率的依赖关系反直觉:
- 更大的模型 → 更小的最优学习率
- 更大的数据量 → 更大的最优学习率
课程标注此结果可能脆弱(其他论文认为数据依赖方向相反),但在该实验中趋势清晰。
学习率本身相当稳健:
到 通常是合理起点——这得到了大量经验支持。 缩放规律可迁移至 MoE:对 MoE 的最优超参数预测与直接搜索的结果接近,只需控制激活参数规模。
最终的简化缩放规律(StepFun Law)大致为:
- 批大小 ≈ 常数 × √数据量
- 学习率随批大小或数据量增大而增大,随模型规模增大而减小
注意:当按照 Chinchilla 方式缩放时(模型和数据量都由计算量驱动),两个效应部分抵消,净效应是学习率随计算量降低、批大小随计算量增大——这与 DeepSeek Law 方向一致,但指数不同。
2.2 优化器:Muon 的兴起与规模之谜
小规模的成功:NanoGPT Speedrun
在 NanoGPT speedrun 基准(作业 1 的灵感来源)上,muon 优化器相对 AdamW 带来了显著加速。Muon 不仅效果更好,而且速度比 AdamW 快得多。
然而课程提出关键问题:这种小规模优势能否迁移到大规模?
优化器比较的困境
课程引用 Kaiyue、Tengyu、Percy 和 David 的大规模优化器比较研究,揭示了许多方法论陷阱:
超参数差异:不同优化器的最优超参数不同。若 AdamW 未调优,其表现会显著劣于调优过的替代品——这会错误地放大替代品的优势。
权重衰减的影响:权重衰减未最优调整时,优化器之间的相对排名可能改变。
规模依赖的两个关键轴:
- 计算量轴:muon 在小规模下优势显著,但随规模增大优势衰减
- Chinchilla 比例轴(数据量/模型参数比):某些算法在过参数化(比例小)时表现好,另一些在大比例时表现好
Scaling 实验可能失败:即使是精心设计的实验,也可能在大规模时崩溃。课程展示了一个案例:某组超参数在小规模时呈现完美线性缩放,但继续扩大后突然偏离并崩溃。解决方法是换用更精细的 μP 风格参数化和不同优化器。
Muon 算法原理
Muon 的核心洞见:梯度对所有参数一视同仁,但网络中的参数并非同质。
- 向量参数(RMS 归一化层)需要常规更新
- 矩阵参数(注意力、MLP)可以享受谱空间操作
Muon 的具体操作:
- 计算梯度,做标准动量更新(得到矩阵 B_t)
- 对 B_t 应用 Newton-Schultz 5 迭代,将其正交化(近似 SVD 的效果,但仅用矩阵乘法,GPU 高效)
- 使用正交化后的矩阵作为参数更新
本质:在 Adagrad/Adam 中,按坐标除以梯度大小使每个坐标单位化;在 Muon 中,改为在谱范数下操作,使每个方向单位化。
工程细节:Newton-Schultz 不是 SVD——SVD 在 GPU 上不快,而 Newton-Schultz 是仅含矩阵乘法的有限迭代,系统实现高效。
Muon 通常只应用于矩阵参数;向量参数仍使用 AdamW。
Muon 的大规模验证:Kimi K2
课程讲述了 Muon 故事的完整弧线:
- 小规模(NanoGPT speedrun)表现惊艳
- 中间规模研究表明优势随规模衰减——一度被认为"故事结束"
- Kimi K2 使用纯 Muon 训练整个模型(加上防发散技巧),获得优秀性能
Kimi K2 没有提供消融实验,因此无法确定该规模下 Muon 是否优于 AdamW,但至少证明了 Muon 在大规模下可行。
核心教训:判断某个方法在大规模下是否有效极其困难。小规模实验迁移仍是最常用的科学方法,但要意识到这是有风险的外推。
2.3 回答学生问题
Q:是否应直接采用已有缩放规律,而非自己做网格搜索?
A:取决于计算预算和精度需求。StepFun Law 在相近计算区间内是很好的默认值。但大规模训练中常有细微差异(如权重衰减的大小),这些因素是否同方式缩放并不明确。这就是为什么许多团队仍重新做 Chinchilla 分析——有足够多的细微差异需要验证。
Q:只调整学习率和批大小,不担心与其他超参数相互作用吗?
A:超参数空间是指数级的,不可能全部网格搜索。常见做法是针对最敏感的参数(学习率)做网格搜索,其他参数(如权重衰减)做单变量扫描。
第三部分:μP 的理论基础与实践
3.1 μP 的目标
μP 的目标可以用一张图概括:随着模型规模(宽度)增大,最佳学习率通常会移动;μP 希望让最优学习率不随规模变化。
实现手段:
- 逐层调整初始化
- 逐参数调整学习率
- 根据模型规模缩放残差连接
3.2 μP 的实证验证
Cerebras(芯片公司但有 LLM 训练部门,由 Hestnes 领导)在 0.1B 到 13B 规模上对比了标准 Chinchilla 配方和 μP 变体:
- 使用 μP 参数化时,scaling law 拟合更稳定
- μP 的损失预测与实际模型几乎完全吻合
- 非 μP 模型的损失预测波动更大
这为 μP 在大规模下的有效性提供了重要证据。
3.3 μP 的数学推导
课程推荐 Greg Yang 的 tensor programs 系列(较难)、Jeremy Bernstein 的综述(较好懂),以及一些物理学家写的解释文章。
两个核心假设(不变性)
A1(初始化激活值不变性):初始化后,输入随机数据,激活值的规模应大致保持为网络规模的恒定函数——不应随宽度爆炸或消失。
A2(特征学习不变性):一步梯度更新后,激活值的变化应为 O(1) 量级——这是"特征学习"的含义。对比神经正切核(NTK):NTK 中激活值变化随宽度趋零,这不是我们想要的。
推导:深度线性网络的初始化
对深度线性网络,设权重初始化为高斯分布,方差为
当 fan-out 不比 fan-in 大太多时,层的输出范数 ≈ 输入范数 × 算子范数。
通过归纳法,若第 L-1 层激活范数为
结论:初始化方差应包含两项——传统的
推导:学习率缩放
对一步 SGD 更新,权重变化是秩 1 更新(损失梯度 × 激活值的外积)。激活值变化有三项:
- 上一层激活值变化的直接传播项
- 权重变化直接导致的交叉项(两项)
为了使每项的量级均为
A3(损失更新单位量级):损失在一步更新中的变化为 O(1)——即模型在取得非平凡进步。
在此假设下,通过一阶泰勒展开和秩 1 性质,可以解出所需学习率为:
即每层学习率正比于 fan-out / fan-in。
对 Adam 的变体:Adam 的更新规则不同,推导结果变为
3.4 μP 的实践规则
初始化:
当 fan-out/fan-in = 1 时退化为标准初始化。
学习率:
- SGD:
- Adam:
这带来有趣的层自适应学习率——这是标准参数化不提供的。
3.5 μP 的局限性
独立研究者对 μP 的压力测试发现:
- 兼容的情况:SwiGLU 激活、RMS norm、不同初始化方式——多数情况下 μP 仍能工作
- 破坏 μP 的情况:
- 在 RMS norm 中学习增益项(可移除而不损伤性能)
- 使用基于符号梯度的优化器(如 Lion)——符号梯度与 Muon 精神相近,但破坏 μP
- 较大的解耦权重衰减——最严重的失效场景
3.6 μP 的定位总结
μP 是工具库中控制超参数随规模漂移的工具之一。Scaling laws 在真实应用中"更像艺术而非精确科学":
- 可用于选择架构、优化器、超参数
- 但外推始终存在不确定性
- 需要主动采取合理措施最大化成功机会
- μP 和拟合缩放规律是两种互补的路线,目前无万能药
课程作业关联
本讲内容直接关联作业 1(NanoGPT speedrun 灵感来源)和后续大规模训练项目:
- WSD 学习率:作业中可能作为默认调度的候选
- Muon 优化器:NanoGPT speedrun 中 Muon 是最优选项之一
- μP 初始化:如果遇到超参数迁移问题,考虑 μP 可能有益
- Scaling 实验设计:作业中做小规模实验时,应关注规模依赖的可能
常见误区
将某个 Scaling Law 视为普适真理:不同工作的公式输入变量和指数都不一致,具体结论高度依赖实验设置和数据。
小规模优化器优势直接外推:Muon 的故事证明小规模优势可能随规模衰减,也可能在大规模下保持——需要实证验证。
忽略 Chinchilla 比例轴:算法优劣可能取决于数据/参数比,而不仅仅是计算量。研究缩放时需同时关注两个轴。
假设所有超参数同等敏感:学习率是最值得调的超参数,批大小次之,其他参数做单变量扫描即可。
使用 WSD 时忽视衰减阶段的重要性:最终退火阶段的训练极其关键——性能增益主要在衰减阶段获得。
μP 与任意修改兼容:学习增益项、符号梯度优化器、大权重衰减都可能破坏 μP 的不变性。
网格搜索没有量化误差意识:网格分辨率不足时,拟合出的"线性关系"可能不可靠。
复习清单
- [ ] 能解释 μP 的 A1、A2 假设及其意义
- [ ] 能推导深度线性网络中 μP 的初始化方差公式
- [ ] 能说明 SGD 和 Adam 在 μP 下的学习率缩放差异及原因
- [ ] 能描述 WSD 学习率的三个阶段及其核心优势
- [ ] 能解释为什么 WSD 使数据规模的 scaling 扫描成本大幅降低
- [ ] 能列出 MiniCPM 和 DeepSeek 在超参数缩放策略上的根本差异
- [ ] 能说明 StepFun 关于批大小和学习率依赖关系的主要发现
- [ ] 能解释 Muon 的核心思想(谱空间正交化)及其与 Adam 的区别
- [ ] 能分析优化器比较中"Chinchilla 比例"作为混淆变量的作用
- [ ] 能列举破坏 μP 有效性的几种情况
- [ ] 能说明 MoE scaling laws 要解决的核心问题(稀疏度 vs. 损失)
- [ ] 能讨论 scaling laws 在实际应用中的"艺术性"来源
缩放定律的“现实世界”:从Chinchilla到muP,我们如何驯服大模型训练中的不确定性
导语:当缩放定律从论文走向工程
在CS336这门课程中,我们已经花了不少时间讨论缩放定律(Scaling Laws)的理论基础——从Kaplan等人的经典工作,到Chinchilla提出的数据与模型规模的最优配比。这些理论告诉我们:随着计算量的增加,模型的损失会按照幂律(power law)可预测地下降,而最优的模型规模与数据量之间存在一种微妙的平衡。这些结论干净、优雅,仿佛只要我们照着公式操作,就能“速通”到大模型时代的彼岸。
但现实远比公式复杂。当我们真正试图去训练一个几十亿甚至上千亿参数的开源模型时,会遇到一连串理论没有完全回答的问题:最优学习率到底是多少?批量大小(batch size)应该怎么设置?为什么有些方法在小规模实验上效果惊艳,一放大就崩盘?初始化方式真的有那么重要吗?更重要的是——那些在论文中看起来理所当然的假设,在实际训练中是否站得住脚?
本讲(CS336 Lecture 11)正是为了回答这些问题而设计的。它把我们带回“扩展前沿”(scaling frontier)的真实场景,考察过去几年中开源社区(从MiniCPM到DeepSeek,再到Kimi K2)是如何在实际操作中处理这些棘手问题的。我们还将深入讨论两个核心主题:一是如何通过缩放定律本身来预测最优超参数(DeepSeek路线),二是如何通过重新参数化(如muP)来让超参数对规模不敏感(MiniCPM路线)。最后,我们会讨论优化器(optimizer)这个往往被忽视、却对规模极其敏感的组件,看看muon这类新方法是如何在“小规模惊艳、大规模存疑”的困境中挣扎前行的。
这篇文章将带你走过这段旅程:从Chinchilla的经典理论出发,看看人们在实践中如何复现它、修正它、甚至绕开它;然后深入muP的数学推导,理解其背后的物理直觉;最后讨论优化器选择和超参数调整中的种种陷阱。我们希望读完这篇文章后,你不仅知道“怎么做”,更能理解“为什么这样做”,以及“什么时候这些方法可能会失效”。
第一部分:经典缩放定律的现代复现——MiniCPM与DeepSeek的实践
从Chinchilla到“开源复现潮”
在进入具体细节之前,我们需要回顾一个背景。Chinchilla论文在2022年给出了一个简洁而深刻的结论:在给定的计算预算下,模型参数数量与训练token数量应按大致相等的比例增长。这个结论指导了此后几乎所有大模型的训练配置。但正如本讲所指出的,Chinchilla分析本身存在一些模糊之处——它提供了三种不同的拟合方法,结果有细微差异。
更有趣的是,当开源社区开始大规模复现这些分析时,出现了新的问题。不同团队可能采用不同的参数化方式(parameterization)、不同的学习率调度,甚至不同的优化器——这些因素都会影响最优超参数的选择。因此,当我们说“在某个规模下,最优学习率是多少”时,这个答案可能是条件性的,取决于你采用的具体设置。
这正是本讲第一部分要探讨的内容。我们将跟随讲师的目光,审视两篇在开源社区中产生广泛影响的论文:MiniCPM和DeepSeek的初始报告。它们代表了两种截然不同的处理超参数缩放问题的哲学。
MiniCPM:利用muP稳定学习率
MiniCPM是一个高性能的小语言模型系列(1-2B参数规模),发布于2024年。它的特别之处在于,团队在训练过程中非常注重对超参数的系统研究。他们面临一个典型的缩放困境:最终模型有约1.5B参数,但直接对这么大的模型做超参数搜索是不现实的。于是他们采用了一种“扩展阶梯”(scaling ladder)策略:训练一组小得多的模型(最大规模大约比最终模型小5倍),从中学习超参数如何随规模变化,再外推到目标规模。
这种方法的关键在于:如何确保从模型上学习到的学习率规律能够可靠地迁移到最终的大模型?MiniCPM团队采用了一类特殊的初始化方法——muP(Maximal Update Parameterization)。muP的核心思想是:通过仔细调整初始化和学习率的缩放方式,使得最优学习率不随模型规模的变化而剧烈改变。具体来说,他们采用了以下操作清单:
- Embedding输出缩放:对embedding层的输出进行缩放。
- 残差连接缩放:用层数的平方根缩放残差连接。
- 矩阵张量的初始化缩放:根据fan-in和fan-out的比例调整初始化方差。
- 张量学习率缩放:对不同参数的梯度更新设置不同的学习率系数。
- LM头缩放:对输出层(LM head)应用类似的规则。
这些操作听起来相当技术化,但目的很简单:让最优学习率保持稳定。MiniCPM的实验结果显示,这个方法确实有效——在不同规模的模型上扫描学习率时,最低损失点几乎都落在同一个值上(约10⁻²)。这意味着,只要使用了muP,你就不再需要为每个规模单独调优学习率。这是一个极其实用的结果,极大地降低了超参数搜索的成本。
WSD学习率:一种更高效的计算范式
除了muP,MiniCPM还介绍了一种学习率调度的变体——Warmup-Stable-Decay(WSD)调度。这是一种“大梯形”形状的调度:先预热一小段(固定步数),然后在一个较长的稳定阶段保持恒定学习率,最后在训练末期快速衰减到零。
WSD的妙处在于它的可重入性(re-entrant)。想象你正在训练一个模型,已经用WSD跑了一段时间(处于稳定阶段)。现在你决定要让模型多看一些数据集。为了做到这一点,你不需要重新开始整个训练。你只需回滚到最新的稳定检查点(checkpoint),继续用稳定学习率训练,然后在最后再做一次衰减。相比之下,余弦学习率调度(cosine schedule)要求你在训练开始前就知道总步数,因而无法灵活延长训练。
MiniCPM的实验表明,精心调优的余弦调度可能略优于WSD(这一观察在社区中普遍存在),但WSD的性能差距非常小,而其灵活性带来的便利却是巨大的。因此,WSD很快成为开源社区的标准做法之一。
Chinchilla复现与数据规模预测
最后,MiniCPM团队还做了一次Chinchilla分析。他们的方法并非最可靠的那种(他们选择了方法1和方法3,而非更流行的IsoFLOPs方法),但结果依然展示了合理的缩放趋势。有趣的是,他们通过联合拟合得到的指数与Chinchilla论文有较大差异,使其声称应该用比Chinchilla多得多的文本量来训练模型。讲师对此持保留态度——这可能反映了他们的拟合方式或实验设置中的一些奇怪之处,而非真正的规律差异。
无论如何,MiniCPM展示了处理超参数缩放的第一种路线:通过重新参数化(muP)稳定学习率,通过灵活的WSD调度降低实验成本,然后在此基础上进行更干净的缩放分析。
DeepSeek:用缩放定律直接拟合超参数
与MiniCPM不同,DeepSeek团队(当时尚未发布MoE模型)采用了另一种策略:不依赖muP这样的特殊参数化,而是试图直接用缩放定律来拟合最优学习率和批量大小随计算量变化的关系。
他们的做法如下:在不同规模的模型上,进行网格搜索(grid search)——在不同的学习率和批量大小组合下训练模型,并观察在给定token数量下的最终损失。固定计算量,变化这两个参数,找到损失最低的点(即最小损失)。然后,在不同计算量水平上重复这个过程。最终,他们得到了两种缩放关系:
- 最优批量大小随FLOPs的幂律增长;
- 最优学习率随FLOPs的幂律下降。
这种方法的优点在于,一旦你拟合出这些缩放定律,你就可以直接预测大规模训练的最佳超参数,而无需在大规模上做任何搜索。DeepSeek的拟合结果(至少对于批量大小)看起来是相当清晰的线性关系(在对数坐标下)。学习率的拟合稍欠完美,但模型最终仍然训练得很好。
DeepSeek还用类似的方法复现了Chinchilla分析(同样使用WSD学习率,并做了一个小的变体:两次衰减阶段),得到了不错的缩放曲线。最令人印象深刻的是,他们将缩放定律外推到实际训练的模型规模时,预测的损失与实际达到的损失非常接近(误差很小)。这验证了缩放定律在实际中的预测能力。
两种路线的对比:稳定 vs. 拟合
至此,我们看到了处理敏感超参数的两种相反哲学:
- 稳定化路线(MiniCPM):通过muP这样的重新参数化,使得最优学习率(甚至其他超参数)对规模不敏感。这样,你只需要在小模型上调好超参数,然后直接迁移到大模型。
- 拟合路线(DeepSeek):接受超参数随规模变化的现实,但试图用幂律关系来刻画这种变化,从而预测大规模下的最优值。
这两种方法各有优劣。muP的承诺很诱人,但它对模型架构和优化器有额外的要求(我们后面会看到它的局限性)。而缩放定律拟合虽然直接,但需要足够多的实验点才能得到可靠的拟合,且拟合结果可能对实验设置高度敏感。
在后续的迭代中(Qwen 2.5/3、Kimi K2、Hunyuan),我们看到了这些路线的融合。例如,Qwen团队采用了类似DeepSeek的方法,但将分析扩展到了MoE模型(关注激活参数量的缩放)。Kimi K2则聚焦于稀疏度(sparsity)与损失的关系,指导其选择了48的稀疏度。Hunyuan做了固定稀疏度下的Chinchilla分析。MiniMax-01则用缩放定律比较了线性注意力(lightning attention)和全注意力架构的缩放表现。这些例子共同表明,缩放定律已经成为开源模型训练的基础工具,但正如讲师所强调的,它们已经从“精确的科学”变成了“一门艺术”——需要在具体情境中谨慎应用。
第二部分:优化器与规模——为什么Adam不是终点
在讨论了学习率和批量大小的缩放之后,我们转向另一个对规模极为敏感、却常被忽略的组件:优化器(optimizer)。讲师指出,在过去的几年中,LLM优化器领域出现了一些令人兴奋的进展,而理解它们需要同时注意两个缩放轴:计算量(FLOPs)和“Chinchilla比例”(token与模型参数之比)。
小规模惊喜,大规模存疑:muon的案例
本讲引入了一个引人入胜的案例:muon优化器。这个故事始于NanoGPT speedrun——一个在很短的时间内将小模型训练到特定损失值的基准测试。在这个基准上,muon相对于AdamW展现出了巨大的优势(性能提升显著,而且每个step的速度还更快)。这让很多人兴奋不已,但随即而来的问题是:它在更大规模下真的能用吗?
一些早期的扩展研究(scaling studies)似乎给出了不太乐观的答案:muon的优势随着模型规模的增大而逐渐减弱。这个结果一度让社区对muon的未来感到悲观。然而,故事的转折出现在Kimi K2上——据报道,Kimi K2模型完全使用muon训练(外加一些防止发散的小技巧),而Kimi K2的性能非常出色,训练曲线也相当合理。虽然Kimi没有提供正式消融实验来证明muon优于AdamW,但至少证明muon在超大模型上是可行的。
这个案例深刻地说明了算法开发中的困境:小规模实验发现有希望的候选方法,大规模验证又可能推翻或支持初始结论,而在这个过程中,我们需要小心翼翼地控制各种变量。一个关键的教训是:不同优化器的超参数是不同的,甚至可能有不同的缩放指数。如果你用调好AdamW的那套规则去套muon,很可能会得到糟糕的结果。反之亦然。
从两个维度审视优化器:FLOPs与Chinchilla比例
讲师特别强调,在比较优化器时,我们不能只看模型规模(或FLOPs)这一变量。另一个同样重要的维度是token与参数之比(即“Chinchilla比例”)。这个比例对优化器的相对表现有显著影响。
想象你在一个过参数化(overparameterized)状态下训练模型(即模型参数很多,数据相对较少)。在这种状态下,某些算法可能因为其隐式正则化特性而表现得更好。而在接近Chinchilla最优比例(即token与参数大致相等)时,其他算法可能因为数据/参数效率更高而胜出。因此,如果不对这个变量进行控制,你的“缩放结论”可能只是某个特定Chinchilla比例下的偶然现象。
尽管有些论文声称做了良好的缩放实验,但往往忽略了Chinchilla比例的控制——可能是因为没有足够的算力去覆盖这个维度。讲师提醒我们,无论是做研究还是做工程,我们都应该始终留意这两个缩放因子。
超参数调优的残酷现实:一个失败的缩放曲线
为了让这一论点更加深刻,讲师展示了一张令人印象深刻的“崩溃”缩放图。在这种图中,研究人员进行了标准的Chinchilla分析,在小规模到中等规模(大约跨越两个数量级)的数据点上都观察到了一条非常平滑的缩放曲线。看起来万事大吉。然而,当把规模进一步扩大时,预测的损失突然偏离了曲线,然后性能急剧恶化——用讲师的话说,“它就这样彻底崩掉了”。
这个案例来自Will Held在Marin(Percy Liang的开源LLM项目)合作的工作,试图验证一组超参数(如cautious AdamC、平方根批大小缩放等)。尽管他们在中等规模上的结果看起来完美,但继续外推时遭遇了灾难性的失败。最终,解决方案是转向更精细的muP风格参数化,或者更换优化器,才获得了跨更多数量级的平滑缩放曲线。
这个经验教训是深刻的:即使是最标准的Chinchilla风格实验,也不能保证外推的成功。我们的缩放曲线总是建造在某种“信任”之上——信任实验室配置与目标场景的相似性,信任没有未察觉的变量在暗中作祟。而在现实中,这些信任经常会被打破。
muon:一个为矩阵参数而生的优化器
在讨论完这些困境后,讲师深入介绍了muon的原理。传统的梯度下降(或动量SGD)对所有参数一视同仁:对梯度做一次动量更新,然后直接用它更新参数。但 muon 的核心思想是:语言模型中的参数并非同质的。有些是向量参数(如RMS归一化层的增益),有些是矩阵参数(如注意力或MLP的权重)。矩阵参数具有谱结构(奇异值),这可以被利用。
Muon的具体操作是:对梯度做动量更新后,将更新矩阵(记为B)进行正交化——即通过奇异值分解(SVD)将所有奇异值设置为1,得到一个新的更新矩阵(U V^T)。这个操作的效果是:无论原始梯度的大小如何,更新矩阵的谱范数(即最大奇异值)都保持为单位长度。这类似于Adam按坐标归一化梯度(每个坐标除以梯度的二阶矩),但muon是在谱域中进行归一化。
Muon的实现并不需要精确的SVD,而是通过Newton-Schultz迭代来近似正交化。这个迭代只涉及矩阵乘法,因此在GPU上是高效的。
为什么muon在小规模上表现这么好?一种可能的解释是,它将每个矩阵参数视为一个独立的结构单元,通过强制更新保持“各向同性”(即所有方向的步长一致),可能有助于避免优化路径的扭曲。但这仍然是一个未完全理解的问题。
muP的边界:什么会破坏它?
在讨论muP时,我们需要认识到它的理论框架是有限的。muP建立在一系列关于激活值规模、梯度更新的假设之上,这些假设对许多先进架构来说并不完全成立。例如,SwiGLU激活函数、RMS归一化、初始化变体等都不在原始muP的框架内。但有趣的是,实践表明这些架构通常仍与muP兼容。
然而,有些因素确实会破坏muP的有效性。讲师列举了三类情况:
- 在RMS归一化中学习增益项(gain)——这会破坏muP。
- 使用基于符号梯度的优化器(如Lion)——这与muP不兼容。
- 使用较大的解耦权重衰减(decoupled weight decay)——这可能是最严重的破坏因素。
这意味着,muP虽是一个有用的工具,但它并非万能药。我们必须理解其适用范围和边界条件。
第三部分:muP的数学推导与直觉
本讲最技术性的部分,是对muP数学原理的推导。讲师从两个核心的不变性假设出发,用“物理学家式”的数学推导,为我们展示了muP缩放规则是如何涌现的。
两个核心不变性
设想我们有一个深度线性网络,其宽度(hidden size)为n,层数为L。我们做以下两个断言:
- 初始化时激活值的规模应该与宽度无关。也就是说,无论宽度是1024还是8192,初始化后各层激活值的绝对值都不应爆炸或消失。
- 经过一步梯度更新后,激活值的变化量应该是O(1)(即与宽度无关)。这被称为“特征学习” (feature learning)。在神经正切核(NTK)极限下,激活值的变化在宽度趋向无穷时会趋于零,这意味着网络无法学习新特征——这是我们不希望看到的。
这两个不变性意味着:在宽度变化时,网络的前向传播和后向梯度更新都保持在一个特定的规范化状态。
初始化的推导
基于这些不变性,我们可以推导出初始化方差的设置。对于每一层,我们假设权重初始化为高斯分布,方差为σ²。如果我们希望激活值在各层之间保持相同的数量级(即每层的激活范数与(隐藏单元数量)^(1/2)成正比),那么我们可以解出σ的值。
推导过程大致是:
- 算子的谱范数(operator norm)与扇入(fan-in)和扇出(fan-out)有关。
- 如果我们设置σ ∝ 1/√(fan-out) ,那么经过各层之后,激活值范数与隐藏单元数量的平方根成正比。
最终,我们获得了muP的初始化规则:权重初始化方差为1/(fan-out)(而非标准参数化中的1/(fan-in))。
学习率缩放的推导
更复杂的是学习率的设置。我们需要考虑一步梯度更新后激活值的变化量。对于SGD,权重更新是秩-1更新(由损失对层权重的梯度和该层的激活值组成)。
通过追踪各项的数量级,我们可以解出学习率η需要满足的条件,使得激活值的变化保持在恒定的数量级,同时保证损失值确实在下降(即损失在更新后按O(1)减少)。这个推导过程涉及多个假设(如没有抵消效应等),但最终得到的结果是简洁的:
- 第L层的SGD学习率应正比于(fan-out / fan-in)。
- 对于Adam,学习率应正比于1/fan-in,因为Adam对梯度的缩放改变了其中一项的数量级。
这些规则,正是在现代Transformer模型中应用muP时需要遵循的。
muP的实用价值与局限
讲师强调,muP不仅仅是技术细节,更是一种极具启发性的思维方式——它展示了如何通过仔细追踪数量级来设计超参数缩放规则,而非盲目依赖实验搜索。然而,它也有其局限性,特别是当我们使用更复杂的优化器(如muon)或未涵盖在理论框架内的架构组件时,这些规则可能需要重新调整。
结语:缩放的艺术与科学
在本讲的结尾,讲师回到了一个更宏大的视角:缩放定律的实践应用是“棘手的”(messy)。最初,缩放定律可能给人一种“精确科学”的感觉——画一条直线,按步骤操作,就能准确预测大规模的结果。但现实远比这复杂。
我们已经看到,即使是训练良好的DeepSeek模型,其缩放定律拟合也并非完美;MiniCPM的Chinchilla复现也存在可疑之处;muon在小规模上的惊艳优势在大规模上可能减弱或消失;而一个看似平滑的缩放曲线,也可能在某个未知的比例下突然崩溃。
然而,这并不意味着我们应该放弃缩放定律。相反,缩放定律已经成为开源社区训练模型时不可或缺的工具。人们用它来选择架构(如Kimi K2的稀疏度决策)、优化器(如muon的采用)、超参数(如批大小和学习率)。这些都说明,缩放定律在实践中确实很有用。
但我们必须记住,它是一种“艺术”而非精确科学。我们必须意识到可能的外推风险,采取合理措施(如使用muP稳定超参数,或进行多尺度网格搜索)来最大化成功的机会。同时,我们也应该保持开放心态,研究新的方法——无论是muP、WSD、muon还是其他尚未发明的技术——来更好应对规模带来的挑战。
这就是缩放定律在“野外”的现实:既有科学的美感,又有艺界的微妙。而作为一名语言模型的构建者,理解这种双重性是至关重要的。
本讲在课程中的位置
在CS336的课程体系中,Lecture 11是“规模化探索之旅”中承上启下的关键一环。前几讲(尤其是关于Scaling Laws的讲座)建立了理论基础——Kaplan、Chinchilla等经典论文定义了最优规模与数据配比。而本讲则深入到工程实践,探讨了这些理论在真实开源模型(如MiniCPM、DeepSeek)中的复现与修正,并引入了优化器(muon)和参数化(muP)等高级主题。
此外,本讲与课程中的作业和实践紧密相连:讲师提到NanoGPT speedrun启发了作业1,而WSD学习率、muP初始化等技术也是学员实际训练模型时可能采用的工具。因此,本讲不仅是对基础知识的检验,更是为学员即将进行的实践(如训练更大规模的模型)提供了必要的方法论。
官方资源
- 视频:CS336 Lecture 11 (YouTube)
- 讲义:CS336课程官网(可访问课程主页获取相关Slides与讲义笔记)
本文基于CS336 Lecture 11的中文逐字稿编写,旨在提供一篇独立阅读的技术博客。文中提到的论文和结果均来自课程内容,未添加任何外部引用。
CS336 2026 Lecture 11:完整中文逐字稿
- 视频:官方视频
- 来源:英文人工字幕提取后翻译为中文;保留段落起始时间戳,只修正明显的断句、重复和语序。
- 说明:这是逐字稿,不是讲义摘要;相邻且语义连续的片段已合并为阅读段落;术语保留必要英文。
正文
[00:05.080] 好的,我们将继续我们的规模化探索之旅。今天的主题比较庞杂。我的意思是,我们实际上要讨论扩展定律中一些更进阶的细节,以及我们如何扩大模型规模。在此过程中,我会覆盖一些我认为在实际扩大语言模型规模时非常重要的细节。所以我们可能会想问一些问题,比如两节课前你们讨论的扩展定律,在真正构建开源模型规模的实践中真的有效吗?如果我们能对其中一些进行优化就好了。人们在训练中会使用很多学习率技巧和其他方法,我觉得作为了解如何构建语言模型的人,你们应该知道这些。所以我们会讨论一些这些优化细节。然后,我想讨论的最后一点是,当我们扩大规模时,有些东西的调参会变得有点棘手。我们知道优化器及其行为实际上对规模非常敏感。
[01:08.060] 正因为如此,像初始化、学习率、批大小这些,我们需要根据规模仔细思考。所以我们会讨论如何仔细地思考这些问题。到目前为止,我们讨论的大部分内容是我所说的经典扩展定律经典内容。我们讨论了Kaplan。我们还讨论了,哦,甚至Hestnes,然后是Chinchilla论文。所以我们跟上了,大概是2022年左右的知识。从那时起,有一些训练大模型的人发表了相当多的扩展论文。不过,我认为这类论文近年来数量有所减少,而且大部分都来自中国的开源社区。所以今天,我想谈一谈我策划和挑选的一些论文,让你们了解,在扩展前沿——至少是开放前沿——是什么样子?这些人关心或担忧的是什么?
[02:05.570] 以及你需要做什么才能从Chinchilla一路“速通”到,比如说,Kimi K2——那是最近发布的包含一些扩展细节的模型。这就是本节课的第一部分。我们会过几篇论文,详细讨论其中的细节。第二部分将讨论优化器和初始化。记住,我两节课前花了不少时间讨论批大小的选择。这将是一个重要因素。而且在思考如何调整学习率等参数时,这也很重要。这些参数非常敏感,对正确设置至关重要。事实证明,人们已经投入了大量精力来把这些事情做对。我们会讨论一些相关技术。这实际上甚至可能对你们日常训练各种语言模型的工作有用,因为这些技术会让你们了解到如何在扩大或缩小规模时调整你的超参数和实验。所以我会花不少时间重点讨论两篇论文。它们不是最新的论文,也不一定是最——实际上,DeepSeek那篇可能是整体上最详尽的扩展论文之一。
[03:14.950] 但我认为这达到了一个最佳平衡点。这两篇论文都来自构建了性能相当不错的模型的人,是严肃的扩展研究。重要的是,对于我们今天要提出的一些关键问题,它们采取了截然不同的方法。所以我认为它们是引出本节课其余内容的一个好方式。我会讨论这两篇。第一篇是MiniCPM。我记得是在2024年,这是一个高性能的小语言模型。还有DeepSeek——你们现在都很熟悉了。我不再需要解释为什么我要谈DeepSeek,这很好。不过我们先从MiniCPM开始讲。我确实很喜欢这篇论文,尤其是它刚出来的时候。它给人一种非常新颖、前沿的感觉。即使在2026年,我仍然觉得CPM团队做的很多事情对我们很有启发。
[04:03.710] 我稍后会讲到更新的论文时再谈这个。但我认为,很多经典的扩展知识,比如如何设置学习率、如何做Chinchilla分析,这些东西现在被认为是理所当然的了。大家都假设你知道怎么做这些事情。所以人们在论文中不再那么频繁地讨论这些了,正因如此,讨论这些较老的论文是非常好的。MiniCPM是中国开源社区里一群产学研合作者共同打造高性能、较小规模语言模型的成果。在它发布的时候,它基本上代表了1到2B参数规模区间的最高水平。如今,Gemma和其他模型——Jemma——以及2D或3D规模的模型性能会更高。但这些都是有竞争力的、训练良好的模型,至少在2024年的标准下是这样。
[04:58.520] 那你怎么训练这些模型呢?正如我之前所说,我们必须把一些扩展参数调对。我今天要讲并会相当强调的一点是,初始化非常重要。初始化之所以重要,是因为它们能让某些超参数不那么依赖规模。所以我想从MiniCPM论文中讲的第一个技术是一类特殊的初始化方法,叫做muP。muP,这种方法的核心观点是,我想确保当我扩大或缩小规模时,我的最优学习率保持不变。我在两节课前非常简短地提到过这一点,希望你们还记得一些。但这是我们今天要讨论的第一个实际应用案例。具体是什么样子的呢?MiniCPM论文给出了一个很好的操作清单,他们缩放embedding的输出。
[05:55.960] 他们用层数的平方根来缩放残差连接。他们根据fan-in和fan-out的比例来缩放所有矩阵形状张量的初始化。他们缩放张量的学习率。你们应该注意这一点,因为如果你不习惯按参数设置学习率,这会比较特别。然后他们还缩放LM头。我们会讨论为什么要做这些事情,或者如何推导出这一套特定的缩放规则。可以说,你只需要稍微改变一下初始化和学习率的做法,以期让学习率更稳定一些。当然,我们在这里训练的是小模型,但我们仍然想要一个好的扩展策略。这项研究的部分意义不在于你要暴力训练一个1.5B的模型,比如训练几百个。而是你希望第一次就把超参数完全调对。
[06:45.885] 所以你要训练一堆小得多的模型,然后建立一个扩展阶梯,一直延伸到目标规模的5倍左右。他们训练阶梯上最大的模型与他们最终发布的实际模型之间大约有5倍的差距。他们的目标是确定,我认为,大多数人认为最敏感的参数,比如最优批大小、学习率。然后是我们之前在Chinchilla讨论过的token与模型规模的比例。他们也想复现一次Chinchilla分析,这对我们现在来说不那么重要,可能对你们训练模型时也一样。但无论如何我都要讲讲它,因为他们的做法中有一个重要之处,我认为这很有启发性和实用性。正如我之前所说,采用这种奇怪的初始化方式,或者说非标准初始化方式,即 muP,其全部意义在于试图让最优学习率保持稳定。
[07:38.800] 是这样吗?至少在 MiniCPM 的案例中,答案是肯定的。事实上,在他们的实验中,这个结果非常非常干净。他们在各种不同规模的模型上扫描学习率,每条线对应一种模型规模。最小损失基本上正好落在 10 的负 2 次方上。我想最小值对于最小的模型来说稍有偏移。但即便如此,它也非常接近最小值。所以这是 muP 方法的一个很好的成功案例。我们稍后会讨论其他一些使用 muP 的模型和论文。这不是唯一的方法,但这是其中一种有效的方法。如果你能得到这样的结果,那么你就不再需要调优学习率了。但当然,即使你有一个很好的参数化方法,能让你的学习率固定在一个最优值上,你的最优批大小仍然不会是一样的。所以批大小实际上会随着数据集大小和模型规模两者的变化而变化。
[08:34.010] 所以在他们的案例中,他们运行了这些模型训练实验。你在这些图上看到的每一列点都是一个训练实验。他们有一个固定的批大小,然后他们增加处理的 token 数量,我想是从下到上递增的。然后他们问,好,我们已经做了很多这样的训练实验。什么是最优批大小?他们将二次曲线拟合到这里的每一条线上,也就是这条等损失等高线。他们发现的是这种最优批大小随处理的 token 数量变化的缩放规律,而且他们是在不同的模型规模下做的。考虑到我们在 Kaplan 论文中看到的关于临界批大小的那些漂亮缩放规律,你可能对此并不感到惊讶。但最优批大小确实与目标损失之间存在幂律关系。所以我们看到的正是 Kaplan 风格的分析:如果你的损失在左侧更低,你的批大小就应该更大,这就是临界批大小的缩放规律。
[09:31.950] 那么现在,给定一个特定的损失目标——你可以从缩放定律中反推出这个目标——你基本上就可以设定最优的或者说临界的批大小,它会给你最佳的权衡。所以这里有一个很清晰干净的趋势。这将让你非常容易且精确地设定批大小。所以这两个——批大小和学习率缩放——会一次又一次地出现。它们可以说是做缩放定律分析时最重要和最敏感的参数。但我还要讲讲 MiniCPM 论文中的最后一件事情,我认为它很重要。请记住,要做 Chinchilla 分析,我们需要训练很多不同的模型。所以我们要,比如说,做 IsoFLOPs 分析。这意味着我要固定一个 FLOPs 预算,然后在 token 数量和模型规模之间进行权衡。为了做到这一点,我要,比如说,增加我训练的数据量,也就是增加序列的数量。
[10:24.310] 但随着我增加序列数量,请记住每一个训练都会有不同的学习率调度。你们都在用余弦学习率训练。余弦学习率看起来是不同的。你需要在训练之前就知道最终的总预算,才能制定余弦调度。而这导致了一个恼人的事实,那就是你只能不断地重新训练你的模型。如果你想训练一个 800 万序列的模型,你不能从 400 万序列模型的终点重新开始。你只能从头开始,因为你的学习率调度从根本上就是不同的。这非常烦人。从某种意义上说,这就像做这件事的二次方成本,因为你不得不一直从头开始重新运行。这看起来非常浪费,我想很多其他人也这么认为。我认为 MiniCPM 论文中的一个很好的解决方案就是现在所谓的,我想,warmup stable decay 学习率。
[11:17.370] 或者说这是流行叫法,它基本上就是一个大梯形。你有一个预热阶段。预热阶段通常定义为固定数量的步数,而不是你训练目标的一个比例。所以这与你的训练时长无关。你有一段相当长的时期保持学习率恒定。那是稳定阶段。然后在最后,你会相当快地衰减到 0。那是衰减阶段。一般来说,衰减阶段大约是你总训练时长的 10% 到 20%。所以你在训练的大部分时间里都处于稳定状态,然后在最后非常快速地衰减,通常衰减到最大学习率的 10% 左右。这就是非常典型的 warmup、stable、decay 学习率。那我们为什么这样做呢?
[12:04.980] 这样做的意义何在?我们使用 WSD 是因为现在我们可以从稳定阶段结束时的任何时间点重新开始一个训练。如果我想复用一个训练并让它训练得更久,我只需要把训练回滚到最后一个稳定检查点。然后我继续用稳定的学习率向前训练,最后把模型衰减到 0。所以这为我们提供了一种非常简便的方法来进行数据规模的实验,而无需不断地重新运行实验。你每次确实需要重新衰减,但那只占总成本的 10% 左右。反复衰减比反复重新运行预训练要划算得多。这是一个在很多地方都使用的非常常见的技巧。我接下来会展示一个性能对比。这种形式的图已经有很多了。我想,最近做这个的可能是——也许是 Martin Jaggi 和其他人,我记得是苏黎世联邦理工学院的,他们对 WSD 性能做了广泛的研究。
[13:06.160] 但我认为很多论文都展示了类似的结果:你有一条调优得很好的余弦学习曲线。这是这里的橙色线。然后你有各种 WSD 曲线。WSD 曲线看起来总有点奇怪,因为看起来你远远落后,直到你进入衰减阶段,突然你重新获得了所有的收益。然后你就能匹配甚至有时超过余弦学习率的性能。据传闻,如果你和一群训练网络的人聊聊,我想很多人会说在许多情况下余弦稍好一些。但 WSD 在其他许多情况下基本上也差不多好。这是一个非常多用途的默认学习率。它不一定需要重新运行如果你只是想要继续训练的话。但如果你以前没有摆弄过学习率衰减,看到学习率衰减有多大影响,那真的很酷。衰减学习率的这一部分对于最终退火阶段的训练来说极其重要。
[14:06.000] 好的,这就是那里的对比。那么现在我们已经配备了 WSD 风格的 SGD 学习率,Chinchilla 风格的分析实际上就非常容易了。你要做的本质上是运行一个长的训练,然后你只需回退检查点并反复衰减。那会给你沿着数据维度的扫描结果。然后你还会沿着模型维度做扫描。你可以做 FLOPs 或者任何你喜欢的其他分析。MiniCPM 的作者们,不管出于什么原因,选择尝试复制方法 1 和方法 3,而在我看来,这两种是 Chinchilla 方法中最不可靠的。但他们确实得到了看起来相当不错的缩放曲线。如果你看MiniCPM的Chinchilla方法1,他们看到了漂亮的下包络线,在这些点上看起来基本是线性的,这很好。然后方法3,他们在计算量与非嵌入参数之间的权衡上,大部分都得到了平滑的缩放结果。
[15:03.110] 所以看起来,他们在大多数情况下都得到了合理的、符合Chinchilla风格的拟合结果。我要指出的一点是,他们整体联合拟合得到的指数与Chinchilla论文相比有相当大的差异。他们在论文中声称这意味着他们应该用比Chinchilla多得多的文本量来训练。我完全不确定这是否真的成立,还是说他们的Chinchilla拟合相对于Chinchilla论文本身来说有些奇怪。好,我在这里暂停一下,以防有人对MiniCPM有问题。也许需要记住的几个要点是,第一,使用muP是一个有用的技巧,我猜是为了在规模变化时稳定学习率的变化方式。
[15:46.398] 第二件要记住的事情是WSD学习率。这是一个非常常见的技巧。你们都应该了解这种学习率。这是基础知识点。好,没问题。那么接下来我继续讲DeepSeek。这是DeepSeek最初的论文,在他们开始讨论MoE和其他那些很酷的东西之前。我认为,即使在DeepSeek的第一篇论文中,你就能看出这些人非常认真,因为我觉得DeepSeek LLM,即使是现在,也是开源世界里做得最漂亮的缩放分析之一。而且我觉得他们在实验设计上很有品味。所以即使在R1等成果出来之前,就挺清楚这些人会有大作为的。好的,DeepSeek在优化超参数方面采取了与MiniCPM和其他几个模型非常不同的缩放策略。他们不会真正使用muP。
[16:44.990] 他们不会做这些有趣的缩放调整。相反,他们会尝试拟合一个缩放定律来同时估计最优批次大小和最优学习率。我们知道学习率和批次大小实际上应该随规模变化。但如果它们变化的方式是可预测的,就像缩放定律那样,那我们就没问题了。我们不需要担心太多。所以他们做的就是在不同规模下运行大量的网格搜索,并试图找到最优值。这是他们在做的事情的一个例子。他们扫描学习率,扫描批次大小。然后他们观察在处理特定数量token时的最终损失。所以他们固定计算量,然后变化这两个参数。他们会在多个不同的规模下做这件事,然后基本上试图确定最低损失在哪里,在这种情况下就是最浅的颜色。所以他们会放一个小星星在那里,然后他们在不同规模下重复这个过程。如果你在足够多的规模下这样做,你会得到一个看起来有点像这样的图。
[17:48.800] 这个图说的是,好的,我要在某个范围内变化我的非嵌入训练FLOPs。你可以看到这里的灰点。当我变化这个时,我的最优批次大小会改变。记住,最优批次大小是通过查看这些大网格并找到最小值点来得到的。然后我会得到看起来像一条线的东西。最优批次大小应该随我的非嵌入FLOPs缩放。类似地,对于我的最优学习率,我会在我的非嵌入FLOPs范围内得到一堆灰点。然后我可以拟合一条线来估计我的最优学习率应该是什么样子。这相当不错。我想指出的一点是,他们的完整运行在星星上。这些是他们选择的运行。在更高的FLOPs下,他们有更大的批次大小和更低的学习率。星星是他们实际用于大规模运行的东西。
[18:52.000] 这些拟合对于批次大小来说看起来是合理的。这确实看起来像一条线。学习率那个——我不太确定这是不是我见过的最好的线性拟合,但我觉得它行得通。可能没问题。模型确实训练出来了。模型确实得到了合理的结果。我觉得这是——我想说,这是学习率方法的一个缺点。不太清楚如果你尝试这种网格搜索的话,除非你的网格位置放对了,否则你会有很多量化误差。然后你得不到很好的缩放定律拟合,最终得到的东西看起来像一条任意的线性关系。好的,哦,是的,抱歉。为什么在相同训练步数的情况下,你会得到多个数据点,且最优值范围很大?为什么在不同的训练FLOPs下,你会得到很大的最优学习率范围?你问的是这里的y轴吗?
[19:49.660] 是的,所以我认为原因在于模型大小之类的变化。2024年回到了那个时代,我认为那时每个开源模型构建者都在复现整个缩放定律体系。所以DeepSeek也用基本上是Chinchilla复现的方式来补充他们的学习率分析。为了做到这一点,他们也会做和MiniCPM相同的事情。他们会采用WSD风格的学习率。他们有一个奇怪的小变化,就是做两个衰减阶段而不是一个。我不太确定这是为什么。这之后并没有真正流行起来,但确实是他们做的一件事。不管怎样,这里的重点是说明WSD是一种有效的方法。你可以用它来复现Chinchilla。而且DeepSeek比MiniCPM有更好的缩放曲线,因为他们选择了IsoFLOPs方法。所以他们在固定的FLOPs范围内得到了非常漂亮、干净的扫描结果。
[20:59.050] 然后他们也得到了相当干净漂亮的Chinchilla风格权衡曲线,和之前一样。这让他们基本上能够复现Chinchilla的结果,并得到类似的数据模型缩放定律。我展示这些的重点只是想说明,看,Chinchilla定律已经被其他人在相当大的计算量范围内复现了。所以那个分析总体上是可靠的。最后,我觉得这是——我有点遗憾这类图在开源发布论文中现在有些罕见了,但这确实是很多缩放定律工作的最终核心结论。你在较小的计算量范围内训练一堆精心挑选的小模型,然后绘制出幂律。最后,如果你缩放得足够远,你会得到非常接近现实的东西。所以两颗星星是他们实际训练的模型。灰点是他们的缩放定律拟合。他们在星星上得到了相当好的预测。我的意思是,本可以更好,但已经相当不错了。
[21:59.650] 所以他们能够在实际训练的开源模型上得到一些相当准确的缩放定律预测。好,我就在这里停一下。如果有人对DeepSeek的东西有疑问,尽管问。我觉得我的一行总结是,这是处理敏感超参数的另一种方法。你要么尝试稳定它们,要么尝试直接用缩放定律拟合它们,然后跟随缩放定律走。DeepSeek代表了这种问题处理上更偏向于拟合缩放定律的方法。好,没问题。那么还有其他模型,一些更新的模型也有不错的缩放定律,他们会讨论。不过,正如我之前所说,我认为这方面的大量工作现在——怎么说呢——大家都已经知道了。所以没有人真的会把这类内容作为大篇幅章节写进论文里。比如你看 Qwen 2.5,他们大概会说,好吧,我们需要调一些最优超参数,比如 batch size 和学习率。
[23:02.245] 所以我们做了一堆 scaling 实验来找出最优值。然后我们有一个 scaling law,我们用这个来预测最优值。而且即使是对于 MoE 模型,我们也能做到这一点。所以这基本上就是 Qwen 2.5 里提到的 DeepSeek 那套分析方法的翻版。Qwen 3 也是同样的做法。到了 Qwen 3 的时候,他们说,行吧,我们在 Qwen 2.5 里已经搞清楚该怎么做了。我们就是完全照做一遍。所以这已经成为一套非常标准的流程,用来确定最优学习率和 batch size。Kimi K2 可能是我要讲的这一系列论文里最新的。到这儿,我就打算快速过几篇论文,让大家了解一下在开源模型发布中,人们在 scaling 工作方面都在做些什么。Kimi K2 和其他几篇论文基本上都聚焦于 MoE 的 scaling laws。
[23:57.440] 现在已经 2026 年了。大多数人都知道怎么做 Chinchilla scaling law,甚至怎么做学习率 scaling law。那什么才是新的东西呢?嗯,很多公司现在都在转向 MoE,完全转向 MoE。如果你要这么做,你就很想知道:我训练的激活参数量对不对?于是他们就做这样的事:好,我们要改变 FLOPs,然后观察损失的变化。而且我会针对不同的稀疏度水平来做这件事,这样我就能搞清楚稀疏度对损失到底有多少帮助或者损害。他们的结论当然是,如果看不同稀疏度水平,网络越稀疏,验证损失越好——
[24:33.640] 在给定 FLOPs 的情况下——这并不意外。但关键的是,他们现在能定量地知道稀疏度和损失之间的关系,这意味着他们可以在模型中放入多少稀疏度这件事上做出更理性的决策。他们看着数据然后说,好,我们要用的做法是采用 48 的稀疏度,因为在那个点上,我们的 scaling 损失已经开始进入边际收益递减的阶段了。所以你可以看到,这种分析是如何被用来为不同开源发布中的不同架构决策提供依据的。Hunyuan——这是几年前另一个开源的 MoE 模型——他们做了非常类似的事情。我觉得每家公司一旦转向 MoE,就都得做这种 MoE scaling law,来搞清楚激活参数数量或者说稀疏度大小是如何影响训练损失的。
[25:19.620] 然后他们会重做很多这类分析。比如说,在 Hunyuan 的例子中,他们最终得到的是每个激活参数对应 96 个数据点的比例。在他们的情况里,他们是固定稀疏度水平的。所以稀疏度不是他们的搜索项。LLaMA 3 也有一些不错的 scaling laws。不过可惜的是,这些并不是特别有用的那种。他们做了 IsoFLOP 式的 scaling 实验,在 token 与模型参数的比例上得到了略有不同但相近的数值。在 LLaMA 3 的 scaling 工作里,他们做的最有意思的事情是右边的这张图——左边那栏是他们常规的 scaling law。
[26:01.340] 这个图的意思是,如果我拥有更多算力,那么模型在某个 benchmark 上回答的 logprob 就会持续下降。而在右边那栏,他们展示的是,好的,如果你有更好的 log loss,那么下游任务的准确率也会跟着提升。所以他们的基本意思是,我们可以直接在那些紫色点上拟合一条 sigmoid 曲线。这可能会让我们对不同模型的好坏有一个感知。你可以看到这些点与曲线的系统性偏差。所以我真不觉得我会完全相信这条曲线就是唯一真相。但你确实可以看出,从 loss 到 accuracy 的映射与一个 sigmoid 形式的 scaling 关系是一致的。我觉得这非常有用,因为我们一直在讲 log loss。
[26:39.730] 但在某些情况下,或者说在很多情况下,log loss 和下游准确率之间有着紧密的耦合关系。然后关于 scaling 我想讨论的最后一篇论文,是去年的 MiniMax-01,里面有一组很好的 scaling 实验,他们做的正是那种——我在两讲之前所提到的 scaling 的用处,也就是架构变体。他们想搞清楚:我有 lightning attention,这是一种线性注意力架构;我还有一个 softmax 版本,也就是完全注意力;另外我还有一个混合版本,把两者结合在一起。这些东西的 scaling 表现如何?是不是存在这种情况:如果我扩大规模,其中某个会比另一个差很多?所以你可以看到,他们基本上是在复现 Kaplan 风格的图——随着我增加 FLOPs,学习曲线会怎么变?对于这些不同的架构,在不同算力水平下,我需要的参数量是不是不一样?
[27:46.970] 而答案大体上是否定的,混合架构、lightning 和 softmax 基本上是可以相提并论的。它们需要相似规模的模型。这组图被用来为他们在实际部署系统中使用混合架构的做法提供依据。这正是那种你基于 scaling law 来做决策的方式。所以这基本上就是我想覆盖的关于各家公司近期 scaling law 方案的内容。就像我之前说的,对于 DeepSeek 的方案,有两点。真正的关键是针对 batch 和学习率的 scaling 分析。这是基于 scaling 的方法。而 MiniCPM 走的是更像 muP 风格的路线,试图获得各种不变性。然后我们看到一些趋势——我们看到 Kimi、Qwen 以及其他人,或者说 Hunyuan,在转向 MoE 的过程中做了这些 MoE scaling 实验。
[28:39.730] 我们看到 MiniMax 做了一些架构方面的工作,也看到其他几家公司复现了 IsoFLOPs 分析。回到我们之前说的那个观点,关键在于,我觉得很多公司现在写得没那么详细了,因为核心的 scaling 机制,比如做 Chinchilla 或者学习率 scaling,现在大家都已经很清楚了。所以对于其中一些模型发布来说,展示这些实验未必有额外价值。或者至少这是我对为什么这些发布中这类实验少了很多的猜测。好的,我花一分钟时间,然后停一下,让大家消化一下,问一两个问题,以防你们对这部分有疑问。好,所以所有这些都覆盖了预训练阶段。那我的问题是,当后训练成为模型开发的标准流程时,什么东西会变,什么不会变?啊,好的,这个问题问的是——当你进行——或者我猜广义来说是当你考虑后训练时,这个东西会有什么变化?
[29:47.990] 我觉得这仍然是一个很大的开放性问题。目前没有一种很好的方式把后训练纳入进来做完全整合的 scaling 分析,因为后训练有时候真的会改变你对预训练应该做什么的本质判断。我能想到的最接近考虑后训练的方式是,有一些工作在研究预训练中的覆盖度或者多样性概念,这有助于预测后训练中会发生什么,但这仍然是很初步的工作。我觉得对于“我们如何与后训练形成协同”这个问题,目前没有任何好的答案。好。行,那我们现在来讲另一块内容。所以我认为本讲剩余的两部分正好对应这两个例子。我将讨论一种方法,或者说一系列方法,即关于学习率、批大小和优化器的缩放定律。这对应的是 DeepSeek 的思路。
[30:50.080] 然后我会讨论 muP 的内容。这是 MiniCPM 的思路。我们能否重新参数化一切,让我们不必担心规模问题?现在我们开始讨论 DeepSeek 的方法。如何正确地缩放学习率?这部分会涉及超参数选择和优化器选择。我会大量引用 StepFun 团队最近发布的这篇预印本。我不认为目前有一项真正全面、可靠、稳健的关于超参数(尤其是学习率)的研究。但这篇可能已经相当接近了。StepFun 团队训练了可信的大模型并发布出来。他们在这项缩放定律研究中,耗费了大量算力对超参数空间进行网格搜索。所以我会过一遍这篇论文,并向你们展示一些关于学习率、批大小以及这些因素如何影响下游损失的有趣经验现象。
[31:51.300] 很多人讨论过各种不同的缩放方式,学习率和批大小大致作为计算量的函数。而 StepFun 的这篇综述——不是综述论文,是超参数调优论文——有一张非常好的表格,涵盖了其中的很多内容。我会指出几件不同的事情。其中一些我们已经知道了。最上面一行,这个 OpenAI 缩放定律——这来自 Kaplan 的论文。如果你看那里的批大小,注意它写的是 2eE18,大写的 L 的指数。那是临界批大小,对吧。它基本上是说,你应该根据最终损失来缩放批大小。好的。DeepSeek 我们刚刚看过。这是基于计算量的缩放定律,同时涵盖学习率和批大小。还有另外几组人提出了各种不同的缩放方案,包括我们之前讨论过的 MiniCPM。
[32:47.840] 我想指出的另一件重要事情是,这些学习率和批大小的公式,它们的输入变量甚至都不一样。人们似乎都不太一致,哪些变量应该影响这些学习率?比如 OpenAI 那个是最终损失。而 StepFun 团队声称,正确的批大小缩放是数据的函数。所以关于超参数缩放的本质,存在很多分歧。我想在讲的过程中说的一点是,即便是最后一行也不要把当作金科玉律。这些东西在很多方面都相当脆弱,但实验和现象本身是你可以很好学习的内容。所以我会这样做。这篇论文的设计与 DeepSeek 那篇非常相似。他们就是训练大量模型,对学习率空间和批大小空间进行网格搜索,并在一系列模型和数据集规模上重复这个过程。
[33:44.810] 通过这样做,他们希望基本上覆盖一组小模型的超参数空间,然后确定最优值在哪里。和 DeepSeek 的设计类似,他们在网格化空间之后——在这种情况下,他们有一个等高线图。但底层的对象就是,他们跑了一堆模型。他们的目标是尽量最小化损失。这些等高线图相当不错,我觉得它们提供了很好的可视化效果。他们跑了分辨率相当高的网格。正因为如此,我们能看到的其中一件事就是,他们实际上可以绘制出相当高分辨率的等高线图,展示超参数空间的样子。在这种情况下,这是一个 10 亿参数和 1000 亿 token 的切片,我记得是这样。你可以看到右侧,那个曲面就是超参数景观。作为批大小(一个轴)和学习率(另一个轴)的函数,如果你固定模型的总数据量,损失如何变化?
[34:46.710] 我们看到,每个切片,如果你固定批大小或固定学习率——每个切片都是相当好的凸函数。甚至几乎是平滑的。这非常好,因为这意味着你可以相当好地搜索最优超参数,并且可以确信最小值不会离得太远。如果这个面非常崎岖不平,你可能会对这种网格化空间的方案是否可行产生严重怀疑。另外,这也许能给你一些关于超参数空间整体平滑度的良好直觉。而且我认为实际上有一个非常有趣的总体缩放趋势。看看其他与这个主题相关的论文,我觉得,我不确定我是否见过任何与这个观察相矛盾的证据。所以我认为一个看起来非常有趣、也是这项工作的关键要点之一的事情,基本上就是,如果你看最优批大小,它似乎唯一取决于的是你训练的总数据量。
[35:46.650] 所以看右边这个图。横轴是你训练的数据量。不同颜色的点——那些是不同的模型。注意不同颜色的点似乎大致都沿着同一条趋势线,相当可预测的一条线。这是对数-对数坐标,所以你应该预期某种以 D 为变量的幂律。在左边这个面板中,我们看到的是学习率的依赖关系。我们看到这是一个非常不同的行为。如果你从蓝色一直看到粉红色,更大的模型有更小的最优学习率。而更大的数据量有更高的最优学习率,这实际上是反直觉的。所以这可能有点脆弱。有其他类型的论文认为对 D 的依赖应该是相反的。但这里的趋势,至少在这次实验中,是相当清晰的。
[36:45.120] 我想指出的最后一点是,学习率实际上相当稳健。我认为这得到了很多人的经验和实验的支持。我觉得人们有一种感觉——其实作为一个小调查,大家知道一个好的超参数是什么吗?如果让你明天训练一个语言模型,会是 1e 负 3,还是 1e 负 4?我觉得我们大致有一个感觉,实际上,学习率的这个范围通常是一个不错的起点。而且这通常是成立的。所以他们得到的缩放定律实际上可以迁移到,比如说,MoE 上,在一定程度上。黄色星星——可能有点难看清——是他们的缩放定律给出的预测。红色叉号是直接在 MoE 上推导出的最优点。所以当你切换到 MoE,只要你大致控制住激活参数,你最终不会得到非常不同的缩放定律。唯一看起来可能有点棘手的是,当你改变训练数据时,最优学习率和最优批大小确实会有所移动,这意味着所有这些现象都很可能是条件性的。
[37:57.020] 这些缩放定律中的具体数值等,可能更多地取决于你训练的数据,而非其他因素。所以在我继续之前,我想把这行再显示一次。你最终会得到一个简单的缩放定律,它说,好的,批大小也许应该大约是数据量的平方根乘以某个常数。你的学习率应该随着批量大小的增加而向上调整,或者随着数据量的增加而向上调整,但随着模型规模的增大则应向下调整。这是一个令人惊讶的效果,但也许并不算疯狂,因为这些实验在很多情况下都假设了类似 Chinchilla 的缩放规律,在那种情况下,这两个效应会相互抵消。随着计算量的增加,你的学习率应该下降。希望之前关于 Chinchilla 的评论足够清晰。我的意思是,当你按照 Chinchilla 方式进行缩放时,模型规模和数据量都由计算量驱动。随着计算量的增加,你预期学习率会下降。
[38:51.808] 所以你最终会得到类似于 DeepSeek Law 的结果,即学习率随计算量增加而降低,批量大小随计算量增加而增大,尽管指数与 DeepSeek Law 不同。好的。所以在某种意义上,这是 DeepSeek 分析的最新且最大规模的版本。并且可以看到一种方式,如果你采用缩放规律式的方法,你会如何处理拟合学习率这个问题。这就是超参数的故事。如果有人有问题,随时可以打断我。然后接下来,紧接着,我将继续讨论优化器。我认为这是架构讲座中没有时间讲到的内容之一。那场讲座已经相当充实了。但我认为如果我不讲优化器,那将是一种失职。而且优化器是依赖于规模的,所以这正好是一个合适的时机来讨论它们,并谈谈优化器领域一些最新、最重大的进展。
[39:52.630] 以及它们与规模的关系。嗯,是吗?我想我是在试图弄清楚如何解读这一点。这是否意味着我们应该直接采用现有的缩放规律并加以应用,而不是自己进行网格搜索?是的,我的意思是,我认为这取决于你所处的具体情境。问题是——哦,抱歉,我需要为线上听课的同学重复一遍——我们是否应该直接采用缩放规律并运行它们,还是应该自己做网格搜索?这确实取决于你拥有的计算预算以及你需要多精确地命中目标。我确实认为,例如,StepFun Law——如果你处于正确的计算区间,比如接近他们运行网格搜索时的条件,这些可能是非常好的默认值。比起我凭空想出来的东西,我更愿意押注在这些上面。话虽如此,如果你正在运行自己大规模预训练,你很可能会遇到一些细微的差异。
[40:44.730] 有人可能会说,我确实相信我们应该使用大的权重衰减。我的架构就使用了大的权重衰减。我们不一定知道这些因素是否会以相同的方式缩放。所以你可能会想要重新做很多这方面的工作。这也是为什么我认为我给你们展示的许多报告都在重新做 Chinchilla、重新做所有这些东西,因为有足够多的细微差异,人们想要检查并确保他们与这些结果在大致上是一阶正确的。我认为在接下来的优化器部分,这个问题会再次出现。我之前也和[MUTED]讨论过这个问题,那就是缩放规律有一种非常科学的感觉。就像是,没错,拟合这些线并外推它们。但归根结底,缩放规律的很大一部分仍然是基于感觉的。就像,我真的相信这份 StepFun 调查的实验设置与我的足够相似,以至于它们可以迁移吗?我们不可能完全确定。
[41:34.210] 总是会有微小的差异。好的,那么现在我要开始讲优化器了。我认为在过去一两年里,LLM 优化器,甚至可能整个深度学习优化器领域,都有一些非常令人兴奋的进展。所以现在正是讨论它们的好时机。我想开始讲的一个故事,也许应该用来引出这个小节的内容,就是这边的这个对比,或者我们可能从左边这个面板开始。这来自 NanoGPT speedrun,它启发了我们的作业1。在那里,你有一个非常非常小的模型,你在相当短的时间内训练它,并且你想尽快达到某个,我忘了,3点几的损失值。很多人在这个基准测试上一直在进行爬山式优化。我认为在这个基准测试上发生的最引人注目的事情之一是,在这里用紫色表示的叫做 muon 的基准出现了,并且相对于用蓝色表示的 Adam 获得了非常显著的提升。
[42:36.460] 它就在那边的绿线正下方。所以你可以看看绿色和紫色线以及它们之间的差距。然后你会说,哇,仅仅通过将优化器换成这个 muon 东西就带来了巨大的收益。而且它甚至不比 Adam 慢,比 Adam 快得多。这太棒了。但是接下来有一个问题是,这个东西在大规模下真的表现好吗?然后你去做一些缩放研究,也许结果就没那么好了。这是我们如何进行研究中一个非常棘手和重要的部分。我们有一个小规模的实验。我们有一些依赖于规模的东西,比如优化器肯定就是。然后我们在小规模下看到了非常大的收益。对于大规模我们该怎么办呢?除了提供一个可靠的、具体的处方之外,我希望我们几乎可以提供一个故事,这样你就能了解我们能做什么,以及哪些事情我们不一定能知道。
[43:27.300] 所以实际上,在我进入下一张幻灯片之前,这种事情经常发生,你会在小规模下想出很棒的东西。但它可能在大规模下并不能成功。那么区别是什么呢?我将花几张幻灯片来讨论 Kaiyue、Tengyu、Percy 和 David 的论文,基本上是关于进行一系列大规模实验来比较优化器,并尝试让这种比较做得更好。从中我们可以学到很多关于缩放以及其他方面的东西,但也有很多是标准的深度学习内容。有一些非常基础的东西。比如,不同模型的超参数是不同的。所以如果你要扩大规模,也许它们甚至有不同的缩放指数。所以你必须把 Adam 调得非常好。如果你调得不好,你会得到这条更深的棕色线。然后你会说,哇,这个东西比我左边这边的所有这些其他算法都要差得多。
[44:18.275] 但是然后你稍微改变一下学习率。然后你会说,哇,实际上我所有的收益都消失了。类似地,在右边,你可能会遇到你已经把学习率调好了的情况。但是现在你有其他东西,比如权重衰减,没有被最优地调整。所以如果你使用相同的权重衰减,那么对于这条绿色线,你肯定会得到更差的结果。如果你选择这个小的权重衰减,相比于选择最优的权重衰减,你会为这条绿色线得到更差的结果。所以有各种各样相当棘手的事情在发生,从严格意义上说这不是缩放相关的事情,但这是我想指出的一个非常重要的通用实证机器学习问题。但也许与当前课堂讨论更相关的是这些事物的规模依赖性。我认为这是一个很好的案例研究,也是一个思考如何真正将算法开发视为规模的函数的例子,因为对我来说,在思考规模时,有两个变量总是很重要的。
[45:11.970] 这篇论文在很大程度上揭示了这两个变量。一个轴是计算量。如果你固定模型规模与数据的比例,比如这两者耦合在一起,那么这里 x 轴本质上就是计算量轴。所以当你进行研究时,本质上就是扩大计算量并观察作为规模函数的趋势。我们在这里可以看到,如果观察Adam(按定义为1.0)与其他所有替代优化器之间的相对加速比,会发现muon在小规模下表现非常好。而随着规模增大,这种优势会逐渐减弱。这不一定是你希望看到的情况。如果你在做算法开发,还有另一个维度你需要始终保持警惕,那就是Chinchilla比例。Chinchilla比例可以这样理解:它是你拥有的数据量与模型规模之间的比率。
[46:12.930] 这是一个非常非常非常重要的参数,因为有些算法在过参数化的情况下表现良好,也就是模型规模远大于数据量的时候。也许你的算法引入了某种隐式正则化,或者它数据效率不高,但因为计算原因仍然表现良好。所以有些算法可能在较小的Chinchilla比例下表现良好。相反,其他算法可能在较大的Chinchilla比例下表现良好,即在每个参数对应大量数据的情况下,因为你可以更高效地将知识压缩进参数中。所以这是一个很大的混淆因素。即使是那些我认为在沿模型规模轴做缩放实验方面有着良好规范的论文,也常常忽略或只是没有足够的算力去研究Chinchilla的token与参数比。而事实证明,这在很多时候出奇地重要。
[47:06.010] 但在这种情况下,并非如此。在这里,你可以看到优化器的收益是相当一致的。在不同的Chinchilla比例下,各条线之间的比值是相同的。但这并不总是如此。你总是需要对这两个缩放因素保持警惕。这一点非常重要。然后我还要指出,如果你有一个新算法或新设置,建立缩放规律是相当不平凡的。我这里的例子取自——Will Held当时在和Marin合作,那是Percy的开源语言模型训练项目。在那项工作中,他们基本上是在尝试验证一组特定的超参数。这甚至算是新算法。这是cautious AdamC,它是一种Adam变体、平方根批大小缩放,以及一系列其他看起来标准的做法。
[47:55.880] 他们做的本质上就是你们在其他论文中见过的Chinchilla分析。但因为这不一定是论文,这是一群人在做实验,你不仅仅发布成功的实验。他们很好地把失败的运行结果也发布了出来。这里的例子是,他们做的基本上是标准的Chinchilla实验,看起来非常棒。如果你走到那条位于10的20多次方的虚线竖线,这里的缩放曲线看起来非常漂亮。然后,随着你继续扩大规模,结果就——嗯,稍微差一点了。差了很多。然后,哇,那个运行结果已经不太好了。所以它就这样彻底崩掉了。这里的要点是,即使是看起来在多个数量级上都表现良好的缩放趋势,也可能突然给你一个意外。它可能偏离线性关系,然后立刻崩溃。
[48:44.760] 在他们的情况中,解决方案就是选择更精细的muP风格参数化。优化器的改动。他们换了另一种优化器,然后得到了在更多数量级上更平滑的缩放曲线。但这里的重点是,当我回过头来说“好吧,你们都应该做缩放实验,你们应该缩放所有的学习率”————这并不容易。你会有相当大一部分时间最终得到像这样子的图。而且全都糟糕透顶,你也不太确定哪里出了问题。好,接下来。然后,我最后想谈的是关于muon的一些内容,以及muon故事的收尾。我觉得我不得不比我之前讲的更详细地谈一下muon,因为它开始被整合进大规模模型训练中了。这就是这类课中许多新研究主题的纳入标准:它是否被用于大规模训练中?
[49:48.300] 而muon现在已经达到了这个标准。而且它本身也是一个有趣的算法,所以我会提一下,讲几分钟,然后继续往下。在梯度下降中,我们通常有一个相当简单的过程。暂时忽略第5行。然后你大概会得到一个标准的基于动量的优化器。我取一个梯度,把它加到B的t上。这是我的动量步骤,所以我以参数mu继续之前的动量。然后我基于这个动量梯度更新我的参数。非常标准的动量梯度方法。但我觉得真正有趣、至少对我来说有些跳出框架的想法是,梯度有一个特性,它对待所有参数都一样。但我们知道,语言模型中,甚至深度神经网络中,并非所有参数都是相同的。
[50:40.920] 有些参数是向量参数,它们属于RMS归一化层。有些参数是矩阵参数,比如注意力或MLP中的参数。矩阵参数可能与向量参数不同。这就是这个想法的出发点,我觉得这确实是一个非常有趣的想法。对于矩阵参数可以做的一件事是,你可以开始观察它们的谱。它们是矩阵。矩阵有特征值和奇异值。所以你可以对奇异值做一些操作。Muon中的做法就是这样一个相当简单的想法。你计算一个梯度,做一次动量更新。然后我取我的更新量B的t。而不是直接在第6行将它应用到我的参数上,我要做的是执行这个NewtonSchultz操作,具体来说是NewtonSchultz 5。这个操作的作用是将B的t矩阵正交化。
[51:33.567] 我在底部写下了这个运算。假设我取B的t,把它写成奇异值分解,USV转置。然后我把所有的奇异值都设为1。所以如果我的奇异值很大,我会把它们缩小回单位大小。如果它们很小,我会把它们扩大回单位大小。然后我取UV转置,这就是我要应用的新更新。如果你想要直观理解,而你已经有对adagrad或其他优化器、甚至Adam的直觉——在adagrad或Adam中,你会除以梯度的大小。所以每个坐标大致是相同的大小。在muon中,你改在谱范数下操作。所以每种方向都是单位大小。注意,这只对矩阵才有意义。我只能对矩阵值的B的t定义奇异值分解这个概念。所以muon的有趣之处在于它是一个矩阵优化器,你将它应用于这些矩阵值的参数。而对于向量值的参数,你可能仍然使用类似AdamW的东西,因为那些与NewtonSchultz项没有真正的交互。
[52:40.990] 另外,NewtonSchultz是对那个操作的一种近似,它只使用矩阵乘法。所以它在系统上是高效的,但并非精确的正交化。但是好吧,我为什么讲了这么多关于muon的内容?首先,它对NanoGPT speedrun这类任务非常强大。这一点是在这一系列事物上被发现的。所以你可能会觉得它效果很好。正如我在上一张幻灯片中提到的,有一些扩展性研究认为,muon 的优势会随着规模的增大而减小。但我想把这个故事讲完,我认为,最初这个故事是这样的。我当时想,哇,muon 真有趣。我想知道会不会有人把它扩展到更大规模。然后[听不清]和其他人把 muon 扩展上去了。他们说,嗯,也许它效果不是特别好。我当时以为这个故事到这里就结束了。我想,好吧,也许这事就到这儿了。
[53:33.110] 没有人会投入巨大的算力把这家伙扩展到完整的预训练规模。然后 Kimi K2 出来了,那个整个模型都是用 muon 训练的,完全用 muon,外加一些防止 muon 发散的小技巧。他们发现了大量不稳定性问题。但 Kimi K2 是一个非常出色的模型。它是一个非常好的模型,而且训练曲线看起来也很合理。所以这实际上表明的是,嗯,muon 在大规模下是可行的。它比 Adam 更好吗?Kimi K2 显然没有任何消融实验,所以在这个规模上我们无从得知。但它绝对是一个有效且可用的方案。那么这个故事的要点是什么?是要说明,判断某个东西在大规模下是否有效真的很难。非常非常难。但迄今为止,大量小规模实验并将它们迁移到大规模,一直是我们做科学的方式。
[54:25.030] 所以你仍然要重视像 NanoGPT speedrun 这样的东西。我们从中学到了不少,获得了好的想法。然后最终,它们会进入大模型。我们会拭目以待,看人们是否会继续在大规模下做消融实验,真正搞清楚 muon 在大规模下是否比 AdamW 更好,等等。这些都是我们尚未有答案的有趣问题。好的,我在这里暂停一下。这是我的优化器部分的结尾。我想确保大家都跟上了,特别是因为我刚刚讲了 muon。那是一个全新的、我以前没怎么讲过的话题。嗯?[听不清]的字母[听不清]在 GPU 上更快?好的,[听不清]。很好。是的,问题是,SVD 在 GPU 上快吗?SVD 在 GPU 上不快,但他们做的不是 SVD。他们做的是 NewtonSchultz。NewtonSchultz 是一种仅包含矩阵乘法的有限迭代算法,用于近似矩阵正交化的过程。
[55:18.940] 所以这里有几个巧妙的部分。其中之一是把它当作一个独立的矩阵来处理的想法。然后是这个正交化的想法,然后是使用 NewtonSchultz 来实现它的系统工程部分。所以这些都是重要且有趣的部分?嗯,后面那位?超参数对于非平凡的[听不清]来说是不同的吗?这里的超参数——不,它们是不同的,是的。我认为有好几个人一直在研究这种 muon 风格的工作方向。其中一个,我想,比如 Jeremy Bernstein,我觉得他一直在说,实际上每一层都应该有自己的学习率。我认为这的终极极限是,你从——甚至,抱歉,每一层都应该有自己的优化器。这个方向的终极极限是说,transformer 中的每一个参数都是不同的。为什么不为每一个参数都设置一个不同的优化器呢?
[56:08.660] 那么你就需要为每一个参数调整超参数,而这是我不想做的。但这也许就是优化的未来。嗯?像——关于学习率和批量大小,你只考虑这两个东西的有效性是什么?然后你从中选出最好的,然后就可以用它来调整其他东西。但这些事情会相互作用吗?是的,或者也许我所理解的问题版本是,很多这些人只是调整学习率和批量大小。而且他们是在网格上做的。但可以想见,这些东西会与一大堆其他超参数相互作用。你不担心这个吗?我认为,总的来说,是的,你确实担心超参数的相互作用。但现实是,它们是指数增长的。这不是你可以把整个空间都铺成网格来做的事情。所以我认为人们经常做的是,他们针对最令人担忧、最敏感的东西做网格搜索。而学习率非常重要。我觉得,如果你要调整超参数,学习率是最值得调的东西。
[57:08.880] 所以他们会非常激进地在这方面下功夫。然后像权重衰减这类东西,他们之后可能会做一个单变量扫描,以使这些参数达到局部最优。好的。行。好了,这节课的最后一部分,我想谈谈 muP。muP 是一个有点神秘的东西。尽管有很多论文写过它,实际上甚至有很多实现,但我不认为它们对底层数学或底层实现应该是什么样都有一致的看法。但有一组核心的想法是大家共享的,而且它似乎确实通常有效。所以我将谈谈我对整个计划的理解,姑且这么叫吧。目标希望是非常清晰的。我认为这张图就是目标,这张图是说,如果我们增加模型的规模,在这种情况下,仅仅是宽度,我们最优的学习率通常会移动。
[58:03.590] 但理想情况下,我们希望最优学习率无论如何都是相同的。这就是我们想要玩的游戏。而我愿意用来实现这一点的旋钮是,我愿意在逐层的基础上调整初始化。我也愿意在逐参数的基础上调整学习率。然后我有时会根据模型的大小来缩放或调整残差连接之类的东西。这就是我想要玩的游戏。那它在实践中效果如何?有几篇论文已经相当广泛地研究了这些问题。Cerebras,它其实是一家芯片公司,也有一个语言模型训练部门,由以缩放定律论文闻名的 Hestnes 负责。我觉得他们做了非常有趣的事情,因为他们现在写了很多关于语言模型 muP 缩放的论文,而且他们几年前有 Cerebras GPT,他们用通常的 Chinchilla 配方训练 0.1 到 13b 的模型,
[59:02.970] 但也包括一个 muP 变体,他们用它来尝试更有效地调整超参数。我认为其中一个好处是,他们发现如果使用 muP 参数化,那么他们的缩放定律拟合通常会更稳定。他们预测的 muP 趋势几乎与他们的实际模型完全吻合。而我认为,他们的非 muP 模型在损失预测方面的波动要大得多。所以他们声称,好吧。muP 也许真的是个好主意。所以这在大规模下得到了验证。你看到了 MiniCPM 和其他一些。很多人用 muP 变体训练过模型。所以我也可以讲很多关于用 muP 训练的例子或模型。但我认为现在对我来说更有趣的事情,实际上是谈谈 muP 的概念基础和数学原理。
[01:00:01.735] 如果你对此感兴趣,并且想读一篇相对易读的论文,我会推荐这一篇。我认为这篇论文对我来说是最清晰的。有几篇论文。Greg Yang,他开创了这个研究计划,有一系列论文,我会称之为 tensor programs。我觉得它们有些难以理解。Jeremy Bernstein,也就是那位同样在研究muon(一种优化方法)相关工作的学者,写了一篇不错的——不是总结性论文,而是一篇综述论文,把它放在了一个不同的框架里,我觉得那个框架非常好懂。另外还有几篇物理学家写的文章,在解释muP到底是什么方面也很有意思。Cerebras(芯粒公司)的人其实也做了类似的一些工作。不过不管怎样,回到这篇论文上来,muP的核心思想基于两个断言。这是一种非常物理学家的思维方式。所以如果我把网络变大——比方说我现在只改变宽度。
[01:00:53.280] 我把网络的宽度变大。我要断言两个不变性,它们应该在我做这个缩放极限的过程中始终保持。首先,初始化时的激活值——也就是说,如果我初始化模型,然后输入一些随机数据,激活值的规模应该大致保持为网络规模的某个不变函数。所以如果它们随着网络规模变大而爆炸,或者随着网络规模变大而收缩到0,那就说明我选择了错误的参数化方式。这是讲得通的。我的激活值应该大致保持在同一量级,不管我的网络规模有多大。其次,在经过一步梯度更新之后,我的激活值的变化量应该是O(1)量级的。这就是人们所说的特征学习。特征学习的意思是,如果我走一步梯度,我网络里的东西会发生显著的变化,而且这个变化量应该是网络宽度的一个固定函数。如果在座有搞理论的,你们可能熟悉神经正切核。
[01:01:46.920] 神经正切核的行为会非常不同。那种情形下,激活值的变化实际上会随着网络宽度趋近于零。而你并不希望这样,因为你确实希望网络能够改变、学习特征、做一些事情,即使在大型神经网络的极限情况下也是如此。所以这两个就是不变性。它们是有道理的。它们基本上是在说,激活值以及激活值的变化量,应该随着网络宽度的变化而大致保持相似。然后注意,我有时候会在单个激活值(也就是这个)和范数之间来回切换。然后还有范数——如果每个激活值大致是O(1)量级的,那么整个激活向量的范数就应该是该层隐藏单元数量的平方根。好,这就是主要内容。muP有意思的地方在于,它从一组非常简单且易于接受的假设出发。
[01:02:39.343] 你们可能都会同意这是合理的事情。然后我们会再做一系列额外的强假设。但是到最后,我们会从中得到一组非常简单的缩放规则。所以记住,我们要做的是确保初始化时的激活值大致为单位量级。那我们怎么做呢?好,如果我们有一个简单的深度线性网络,并且我们把所有权重初始化为服从sigma平方缩放的高斯分布,那么这个矩阵的谱范数会是多少呢?你可以把它写出来。我不要求你们去验证,但这是关于随机高斯矩阵算子范数的标准矩阵集中论证。然后现在注意到,在某些范围里,当扇出(fan-out)——我想,不比扇入(fan-in)大太多的时候,你会得到该层的输出大致等于该层输入乘以算子范数。这总是一个上界,但在某些高维情况下这是一个近似。
[01:03:39.130] 现在,如果这是成立的,那我们能做什么呢?我们可以解出sigma,试着让我的激活值在所有层之间、在所有宽度上保持相同。所以想象在第L减1层,我有平方根L减1大小的激活值。记住,这个符号与我们之前的表示一致,至少对于整个向量来说,我应该有n的L层的平方根大小。所以这是我们的归纳基础情形。然后我们的归纳步骤说,记住我在第L层有这个算子范数,它是关于扇入和扇出的函数。所以如果我选择一个特定的sigma,这就会等于这个。所以假设这是一个ansatz(试探解)。我假设这就是答案,然后我们看看归纳会怎样展开。如果我选择这个作为假设,并把它代入sigma,注意到我会得到算子范数的比值是这个,也就是扇出比值的平方根除以——抱歉,是扇出的平方根除以扇入的平方根。
[01:04:35.110] 然后我把它代回这里。然后我会得到输出的范数等于我这一层单元数量的平方根,加上一些我不打算追踪的低阶项。所以这个sigma是从哪里来的——我的意思是,这像是从帽子里变出来的。但我们可以验证,如果我把这个噪声初始化代入我的公式,那么基本上在每一层,我的激活值的大小大致等于该层隐藏单元数量的平方根。这里做了一些假设,比如这个约等号。但总的来说这是一组合法的假设。条件A1很简单。条件A2说实话稍微棘手一点。所以我们确实需要处理梯度更新了。而且我只会考虑单一样本在批量中SGD的情况,针对线性神经网络,也就是深度线性网络。
[01:05:36.212] 那这会是什么样子呢?好,我某一特定层L的权重变化——如果你想想这会如何表现,它会是我损失函数梯度和激活值之间的一个秩1更新。这是成立的,仅仅因为反向传播的行为方式。前向传播是一个秩1乘法。反向传播则会是一个秩1外积。所以这就是权重的变化。那么那一层激活值的变化是什么呢?好,激活值的变化——我只需要取h等于W乘以hl减1,然后把它展开。我把这个展开成这个公式。这恰好就是第L层激活值的变化。它会有两项,一项是来自L减1层激活值变化的直接项,另一项在某种意义上是一个交叉项,来自我的权重发生了变化,而不仅仅是因为前一层激活值变了。
[01:06:31.487] 所以这个更新有两项。这就是第L层激活值的变化。现在我要做的是,逐一看这三个项。我希望这三个项中的每一项都有相同的量级,并且我希望那个量级是n的L层的平方根。为什么?记住,如果我回到两页幻灯片之前,我对A2的条件是,经过一步梯度更新后,激活值的变化应该是O(1)量级的。所以如果我想要这个成立,那么那一层的范数应该是nl的平方根。要做到这一点,我希望这三个项的每一项都有那个量级。如果这是真的,并且没有相互抵消的情况,那么这就是正确的量级。你们可以看到这些粗略的假设在累积。我们假设了诸如没有抵消之类的条件,但这给了我们一些有趣的方式来思考这些对象中每一个的量级。
[01:07:25.740] 所以此时,正确的理解方式并不是“Tatsu 坐在那里做严谨的数学推导”。我在这里做的事情,本质上是在追踪所有项的數量级。而我要说的是,在给定假设的前提下,这些东西大致应该有多大?这是物理学家的数学。无意冒犯物理学家。这种方法最初就是由物理学家这样发展出来的。好,那现在真的——抱歉,往回翻一页幻灯片让这一切更清楚。第一项,这几乎是直接来自假设的,是归纳假设。我们假设 delta h of l minus 1 具有正确的數量级。所以这一项是白来的。另外这两项——就更微妙一些。我们有这些项,它们将权重变化乘以算子范数再乘以根号 L minus 1。因为这两项的數量级分别是这两个量级的乘积,所以我们知道这个 delta L 需要有多大。
[01:08:21.245] 我们可以解出这个,因为我们知道左侧需要是根号 L 的量级。然后我们有这一项,所以我们解出这个算子范数。而这正好就是这个量。所以这就是下一张幻灯片的目标。我想试着弄清楚什么样的学习率能让我得到大致按扇出平方根除以扇入平方根进行缩放的更新。这就是目标。如我所说,在我选定一个学习率之后,这应该大致成立。那我们怎么做到呢?好,我们要再做另一个假设。而这可能是所有假设中最让人难以接受的一个,至少在我看来是这样,那就是我们假设损失更新以单位量级缩放。换句话说,我们假设我们的模型确实在取得一些可观的、非平凡的进步,并且这种进步作为规模的函数是大致相似的。我不太确定这个假设对我而言是否可接受,但这又是另一个假设。
[01:09:14.740] 所以我们把这个假设放下去。但这个假设的意义在于,如果这是真的,剩下的数学推导就会非常干净地展开,因为损失的变化等于我的梯度乘以我的权重变化量。这是一阶泰勒近似。然后因为这是秩 1 的,我可以把这个写成 Frobenius 范数和算子范数。所以右侧现在就让我们得到了我们想要的东西。这是 O(1)。这一项我们希望是根号 n 除以根号 l minus 1。这就是右边这里的这一项。最后,我有这一项 L 的梯度项。这将会是根号 n_l 减 1 除以根号 n_l 的平方根量级。那现在我需要什么样的 eta 呢?好,我需要的 eta 就是解出这个量,也就是这两个量的比值,即 n_l 除以 n_l minus 1。所以我们做了一个假设,即左边的量级是 O(1)。然后我们解出了所有量,最终得到了每一层学习率的扇出除以扇入的比值。
[01:10:13.770] 我要做的一个附带说明是,对于 Adam,由于其中一些项的变化,你会得到一个 1 除以 n_l minus 1 的缩放,这点我会在这里提一下。所以我已经逐一讲解了希望看起来比较直观的各类缩放推导。但最终,也许你感兴趣的是,这意味着什么?这对我有什么用?这意味着我们要稍微改变设置初始化和学习率的方式。所以我们的初始化会变成这样——对 muP 来说,蓝色框里是我们的做法。我们的初始化将是 1 除以根号 n_l minus 1,再乘以一个 1 除以根号(扇出除以扇入) 的分量。所以如果这个比值是 1,那么这就等同于通常的初始化。这是标准参数化。但如果这个比值不小于 1,那你得到的基本上就是不同的东西。
[01:11:08.790] 你得到的是更接近根号(扇出除以扇入) 的结果。学习率通常你会设为模型规模的常数函数,但你会想把它设为扇出和扇入的比值。也许更有趣的是,对于 Adam,整个东西会有些不同。我甚至不会去碰那个推导,那里你会去掉这个 1 除以 n_l 的项,得到每层学习率为 1 除以扇入。所以基本上,在 Adam 中扇入较大的层会得到较小的学习率。这让你能够拥有这些相当有趣的层自适应学习率。所以这就是 muP 的推导。我讲了很多细节。我确实想再谈谈这其中的高层要点。高层要点是,muP 作为一个论证过程非常有趣。不仅其实际内容很有用。你得到这个算法的方式本身就很有趣,因为这是一种非常不同的数学推理方式,其中你在说,好,我要对我的网络取一个缩放极限。
[01:12:06.110] 然后我断言我的网络具有某种不变性,然后我再做额外的假设。然后我要试图弄清楚我的超参数上的约束是什么。而这些约束本质上也给了你超参数上的某种缩放极限。所以我认为,这是提出不同类型算法或超参数缩放规则的一个通用原则,这些可能和你习惯的 CS 或机器学习训练中的做法非常不同。好,这些是 muP 的推导。我想以一些实证内容来结束 muP 这部分。有一篇由独立研究者写的有趣论文,基本上是在用各种不同方式对 muP 进行压力测试。而 muP 可以被视为一种超参数调节的程序。因此取决于你如何做事,Adam 会根据网络的不同组件有不同的学习率函数。所以你有嵌入参数、注意力参数、输入/输出 MLP 或 softmax 线性层。所有这些都是按 muP 或非 muP 的缩放规则具有不同的缩放方式。
[01:13:18.540] 这里有趣的是,那篇特定论文复现了一批在不同规模下训练的语言模型的 muP 结果。也许不令人惊讶的是,他发现了几乎与 MiniCPM 和其他工作一样的头条结论:如果你正确使用了 muP,并且只以非常受控的方式进行缩放,你就能得到学习率最优性不变性。无论是基线 muP——或者说只看 muP 的一个跟踪注意力投影偏差的变体,两者都能很好地迁移。但你可能会问,现实世界中有很多东西是 muP 不处理的。严格来说,SwiGLU 并不适合 muP 的框架,初始化变体、初始化方式或 RMS norm 也都不适合。
[01:14:12.600] 或者如果你真的乱动优化器,那也会破坏它。所以你可能会问,这些理论所不允许的东西中,哪些真的会破坏 muP?实际上,在大多数情况下,这些东西大多还是和 muP 兼容的。但一些不兼容的情况包括:如果你在 RMS norm 中学习增益项,这被证明会破坏 muP。尽管在许多情况下这些可以被移除而不一定伤害到你的性能,所以也许这不是个坏主意。但这确实看起来会破坏 muP。如果你使用更奇特的优化器,比如基本上依赖梯度符号的 Lion,它似乎也会被破坏。而符号梯度(sign gradient)在精神上与 Muon 及其他一些想法相似。因此,可能还有其他有趣的情况会打破 μP(muP)的适用性。最后,看起来最令人担忧的一点是,如果你使用较大的解耦权重衰减(decoupled weight decay),最终会导致严重的 μP 失效。
[01:15:06.920] 所以这或许是它确实无法通过的一项压力测试。但为了结束这最后一节,你最终也许会问,μP 到底有没有用?如果你真的想稳定学习率,它其实相当有用。有很多这类实验表明,当你增加宽度时,最佳学习率会随规模变化而非常可预测地、但幅度很大地发生偏移。而使用 μP 时,学习率则不会这样偏移。因此,μP 是工具库中众多工具之一,用来尝试控制超参数随规模变化的漂移。我认为这是一个有趣的开放研究领域。这件事还没有完全定论,也不是只有一种正确做法。但 μP 风格的初始化方案有很大的前景。同样,单纯拟合缩放定律(scaling laws)也有很大的前景。最后我想以这一点来收尾:实际应用中的缩放(scaling in the wild)是非常棘手的。正如我在之前的回复中提到的那样,我认为缩放定律最初的表述方式让它听起来像一门精确的科学。
[01:16:04.550] 就好像是说,没错,画这样一条线,按这些步骤操作,你就能确切知道大规模下会发生什么。但我认为,现实情况远比这更加混乱、更加未知。人们在实际中确实会用缩放定律来做很多事情,比如选择架构、选择优化器、选择超参数。这些都是人们用缩放定律去做的事情,但在某种意义上,这确实是一门艺术,因为你并不真正确定它是否能一直外推下去。你需要采取合理的措施来最大化成功的机会。所以人们提出了很多解决方案。你可以使用 μP。你也可以在所有其他这类方面搜索最佳学习率。这些方法都是控制超参数漂移的手段。但目前还没有万能药。也许明年会有一个模块说,我们解决了这个问题,但现在还没有。
[01:16:53.390] 好的,非常感谢。之后我可以回答一些问题。
CS336 2026 · Lecture 11 · Scaling laws
- 视频:官方视频
- 字幕:en-US · manual captions
- 处理:连续字幕合并为自然段;每段保留首条字幕的起始时间戳
- 状态:原语言字幕稿,待校对
正文
[00:05.080] OK, we are going to continue our scaling journey. Today's a bit of a grab bag of stuff. I mean, we're really going to talk about some of the more advanced details of scaling laws and how we scale up models. And in doing so, we're going to cover what I feel like are some important details in how you would scale up a language model in practice. So we might want to ask some questions, like the scaling law stuff that you talked about two lectures ago, does that actually work at real building open-source model scales? It would be nice if we could optimize some of this. There's a lot of learning rate tricks and other things that people do that I think you should know as people that are informed about how to build language models. So we'll talk about some of these optimization details. And then I think the last thing I want to talk about is, when we scale up, there are certain things that are going to be a little bit trickier to tune. We know that optimizers and optimizer behavior is going to be actually pretty scale sensitive.
[01:08.060] And because of that, things like initializations, learning rates, batch size, those we need to carefully think about as a function of scale. So we will talk about how to think about those carefully. So up until now, most of our discussions is what I would call the classical scaling laws cannon. We talked about Kaplan. We talked about, oh, even Hestnes and then the Chinchilla paper. So we're up to date with, I don't know, 2022 or something. Since then, there have been a number of scaling papers that are published by people who train big models. Although, I think the number of such papers have decreased in recent years, and most of them have come out of the open-source community in China. And so today, I want to talk about a number of these that I've curated and selected to give you a sense of, what does scaling at the frontier, at least open frontier, look like? What are the things that these people care about or worry about?
[02:05.570] And what would you need to do to speedrun your way from Chinchilla all the way to, let's say, Kimi K2, the most recent release that has some scaling details. So that's going to be the first part of this lecture. We're going to go through a couple of papers, and we're going to go through the details. The second part is going to be a discussion of optimizers and initializations. And remember, I spent a nontrivial amount of time two lectures ago on the selection of batch sizes. That's going to be an important factor. And it's also going to be important thinking about how we're going to tune things like learning rates. Those are very sensitive, very important to get right. And so it turns out that a lot of effort has been invested into getting those things right. And we will talk about some of the techniques to do so. And this may, in fact, even be useful in your day to day lives of training various language models because they'll give you a sense of how to change your hyperparameters and your experiments as you scale up or down. So I'm going to spend a nontrivial amount of time talking about two papers in particular. They're not the newest papers. They're not necessarily the most-- actually, DeepSeek one might be one of the more detailed ones-- most detailed scaling paper in general.
[03:14.950] But I think this hits a sweet spot. So both of these papers are serious scaling investigations from people who build reasonably high-performance models. And importantly, they take pretty different approaches to some of the key questions that we will ask today. And so I think they're a good way to motivate a lot of the rest of the lecture. So I'm going to talk about these two. First one is MiniCPM. This was, I think, back in 2024, high-performance small language model. And DeepSeek-- you all are now familiar with this. I no longer have to justify why I'm talking about DeepSeek, which is great. We'll start by talking about MiniCPM, though. I do like this paper, especially when it first came out. It felt very fresh and new and at the cutting edge. Even in 2026, I feel a lot of the things that the CPM folks did has been informative to us.
[04:03.710] I'll talk about this later as I get to the newer papers. But I think it's becoming the case that a lot of this classic scaling stuff, how you set learning rates, How do you do Chinchilla? those kinds of things I think are now taken for granted. Everyone just assumes that you how to do those things. And so people don't really talk about those as much anymore in the paper, so it's really nice to talk about some of these older papers for that reason. So MiniCPM is an effort by a bunch of industry academic partnerships folks over in the Chinese open-source community to come up with a high-performance, small-ish language model. And at the time when this came out, it was basically state of the art for the 1 to 2 billion parameter bracket. Nowadays, Gemma and other models-- Jemma-- and other models at 2D or 3D are going to be higher performance. But these are competitive, well-trained models, at least as at 2024 standards.
[04:58.520] So how do you train these models? As I said before, we have to get some of these scaling parameters right. And one of the things that I'll talk about today that I'll emphasize quite a bit is, initializations are very important. Initializations are important because they allow you to make certain hyperparameters less scale dependent. So the first technique that I want to talk about from this MiniCPM paper is a special class of initializations called muP. So muP, the whole point of this approach is to say, I want to make sure that my optimal learning rate is the same as I scale up and down. I talked about this very briefly in the two lectures ago, so you should hopefully remember that a little bit. But this is the first real instantiation that we're going to talk about. What does that look like specifically? The MiniCPM paper has a nice list of things to do, so they scale the embedding output.
[05:55.960] They scale the residual connections by the square root of the number of layers. They scale the initialization of all of their matrix-shaped tensors by the fan-in, fan-out ratio. They scale the learning rates of tensors. And you should pay attention to this one because this is exotic if you aren't used to per parameter learning rates. And then they also scale the LM heads. We'll talk about why you do these things, or how you derive this particular set of scalings. Suffice to say that you're just going to change up how you're going to do your initializations and your learning rates a little bit in the hopes of stabilizing your learning rates a little bit. And of course, we're training a small model here, but we're still wanting to do a nice scaling strategy. Part of the point of this study is not that you want to brute force a 1.5 model, like training hundreds of them. You want to get the hyperparameters exactly right the first time.
[06:45.885] So you're going to train a bunch much smaller models and then get a scaling ladder up all the way to your target size of 5x or so. And so there's about a 5x gap between the biggest model they train on their ladder and the actual model that they're going to finally release. And their goal is going to nail down, I think, what most people consider to be the sensitive parameters, like optimal batch size, learning rate. And then the token-to-size ratio that we talked about at Chinchilla. They also want to replicate a Chinchilla analysis, much less important to us now and probably to you even if you're training. But I'm going to talk about it regardless because there's one important thing in the way they do that and that I think has been informative and useful. So as I said earlier, the whole point of going through this weird initialization or nonstandard initialization called muP is to try to make the optimal learning rate stable.
[07:38.800] Is it? At least in the MiniCPM case, the answer is yes. In fact, it's very, very clean in their experiments. They sweep learning rate across a variety of models, so each line is a model size. The minimal loss is basically right on 10 to the negative 2. I guess the minima slightly shifted for the smallest model. But even then, it's quite close to the minimum. So this is a nice success case of the muP approach. We will talk about several other models and other papers that use muP. This is not the only way to do it, but this is one effective way to do it. And if you can get something like this now you've removed the need to tune the learning rate. But of course, even if you have a nice parameterization that fixes your learning rate to a single of optimum, your optimal batch size is still not going to be the same. So the batch size actually varies as a function of both data set size and model size.
[08:34.010] So in their case, they run these model training runs. So each column of points that you see on these plots is a training run. So they have a fixed batch size, and they increase the number of tokens processed, I guess, from bottom to top. And then they ask, OK, we've done many of these training runs. What is the optimal batch size fitting these quadratic curves onto each of these lines over here, this equal loss contour? What they find is this kind of optimal batch size scaling as a function of the number of tokens processed, and they do this across different model sizes. Maybe you would not be surprised by this, given the nice scalings that we saw in the Kaplan paper for critical batch sizes. But the optimal batch size does have a power law structure with respect to the target loss. And so we see exactly the Kaplan-style analysis of, if your loss is lower on the left, your batch size should be bigger, critical batch size scaling.
[09:31.950] And so now given a particular loss target, which you can back out from the scaling laws, you can basically set what the optimal or critical batch sizes that would give you the best trade-offs. So there's a nice clean trend here. This will let you set batch size very easily and precisely. So those two, the batch size and the learning rate scaling, they're going to come back over and over again. Those are arguably the most important and sensitive parameters to do scaling laws on. But I'll talk about one last thing in the MiniCPM paper that I think is important. So remember, to do Chinchilla, we need to train a lot of different models. So we're going to, let's say, do IsoFLOPs. That means I'm going to fix a FLOPs budget, and I'm going to vary my trade-offs between number of tokens and model size. And to do that, I'm going to, let's say, increase the amount of data I train on, so increase the number of sequences.
[10:24.310] But as I increase the number of sequences, remember that each one of these gets a different learning rate schedule. You all training on cosine learning rates. Cosine learning rates look different. You need to know the terminal total budget before you train in order to have a cosine schedule. And so this leads to an annoying fact, which is that you just need to keep retraining your models. If you want to train a 8 million sequence model, you can't restart from the end of a 4-million sequence model. You have to just restart from scratch because your learning rate schedule is just fundamentally different. This is just very annoying. It's like a quadratic cost in some sense of doing this because you have to keep restarting your runs just from the start. And this seems very wasteful, and I think lots of other people did, too. And I think one of the nice solutions from the MiniCPM paper is what's now called, I think, the warmup stable decay learning rate.
[11:17.370] Or this is the popular term for this, and it's basically a big trapezoid. So you've got a warmup phase. The warmup phase is usually defined as a constant number of steps, not as a fraction of your training target. So this is independent of your training horizon. You've got a fairly long run where you're holding the learning rate constant. That's the stable phase. And then at the end, you're going to decay fairly rapidly down to 0. That's the decay phase. And generally, the decay phase is something like 10% to 20% of your total training run. So you're going to be stable for the vast majority of your training run and then do a very rapid decay down to usually about 10% of your maximum learning rate at the very end here. And that's a very typical warmup, stable, decay learning rate. And why do we do this?
[12:04.980] What's the point of doing this? Well, we do WSD because now we can restart a run at any point at the end of the stable phase. So if I want to reuse a run and train it for longer, I'll just roll back my training to the last stable checkpoint. And then I'll just run the stable learning rates forward and then decay the model down to 0. So this provides a very easy way for us to essentially run scale in data experiments without constantly rerunning your experiments. So you do have to redecay every time, but that's maybe 10% of the total cost. It's a much better situation to be repeatedly decaying than to repeatedly rerun your pre-training runs. And this is a very common trick that's used in many places. I'll show, I guess, a performance comparison. There have been many plots of this form. I think, who was the most recent folks that were-- maybe it was Martin Jaggi and others, I think, over at ETH had an extensive study of WSD performance.
[13:06.160] But I think a lot of papers have shown similar things of, you have a well-tuned cosine learning curve. This is the orange line over here. And then you've got various WSD curves. And WSD curves always look a little bit funky because it looks like you're way underperforming until you hit the decay phase, where suddenly you reclaim all of your gains. And then you match or even sometimes exceed the cosine learning rate performance. Anecdotally, I think if you talk to a bunch of people training networks, I think a lot of people will say the cosine is slightly better in many cases. But WSD is basically as good in many other situations. It is a very versatile default learning rate. It doesn't necessarily suffer from having to rerun if you want to just continue runs. It is very cool though, to see just how much of an impact learning rate decay has if you haven't played with that before. This part of decaying the learning rate is just incredibly important to getting the final annealing parts of your training together.
[14:06.000] OK, so that's the comparison there. So now that we're equipped with WSD SGD style learning rates, Chinchilla style analyzes are actually very easy. What you're going to do is you're going to essentially run one long run, and then you'll just rewind the checkpoints and then repeatedly decay. And that will give you sweeps along the data dimension. And then you'll also do sweeps along the model dimension. You can do flops or any other kind of analysis that you would like. The MiniCPM authors, for whatever reason, choose to try to replicate methods 1 and 3, which, to me, are the least reliable of the Chinchilla methods. But they do get fairly nice looking scaling curves. If you look at the MiniCPM Chinchilla method 1, they see nice, lower envelopes that look mostly linear across these, which is good. And then method 3, they get nice, smooth scaling for the most part in both the compute to nonembedding parameter trade-off.
[15:03.110] So it seems like, for the most part, they get reasonable Chinchilla-looking fits. One thing I will point out is that they get pretty different exponents for their overall joint fit compared to Chinchilla. They claim in the paper that this means that they should be training with way more pieces of text than Chinchilla. It is highly unclear to me whether this is a really true or whether their Chinchilla fits are a little bit strange relative to the Chinchilla paper. OK, I'll stop here for a moment in case someone has questions about MiniCPM. Maybe the main couple things to take away is, one is the use of muP is a useful trick, I guess, to try to stabilize how your learning rate changes over scales.
[15:46.398] And then the second thing maybe to remember is the WSD learning rate. That's just a very common trick. You should all know about this learning rate. That's a basic thing to know. OK, good. All right, so then I'll move on to DeepSeek. This is the original DeepSeek paper, before they started talking about MOEs and all these other very cool things. And I think, even at the very first DeepSeek paper, you could tell that these were very serious people because I think DeepSeek LLM, even now, is, I think, one of the more nicely executed scaling analyses in the open world. And I feel like they have lots of great taste in the experiments that they run. So even before R1 and so on came out, it was kind of clear that these people were going places. All right, so DeepSeek takes a very different scaling strategy than MiniCPM and several others in trying to optimize their hyperparameters. They're not going to really do muP.
[16:44.990] They're not going to do this interesting scaling adjustments. Instead, they're going to try to fit a scaling law to estimate both the optimal batch size and the optimal learning rate. So we that the learning rate and batch sizes should actually change with scale. But if the way in which they change are predictable, just like a scaling law, then we're good to go. We don't really need to worry about very much. So what they do is they basically run extensive grid searches at various different scales, and they try to find the optimum. So this is an example of the kind of thing they're doing. They're sweeping learning rates. They're sweeping batch sizes. And they're looking at the terminal loss if they're processing a particular number of tokens. So they fixed the compute, and then they vary these two. They'll do this across several different scales, and then they'll basically try to identify what the lowest loss is, so in this case, the lightest color. So they'll put a little star there, and then they'll do this across scales. So if you do this across sufficiently many scales, you will get a plot that looks a little like this.
[17:48.800] And this plot is saying, OK, I'm going to vary my nonembedding trading flops across some range. And you can see these gray dots over here. And as I vary this, my optimal batch size will change. Remember, the optimal batch size is derived by looking at these big grids and finding the minimizer. And I will get something that looks like a line. The optimal batch size should scale with my nonembedding FLOPs. And similarly for my optimal learning rate, I'll get a bunch of gray dots in my nonembedding FLOPs range. And then I can fit a line to that to try to estimate what my optimal learning rate should look like. So this is fairly nice. And I guess one thing I'll point out is, their full runs are on the star. These are the runs that they chose. At higher FLOPs, they have much higher batch sizes and much lower learning rates. The stars are the things they actually ran for their big runs.
[18:52.000] These fits look reasonable for the batches. This does look like a line. The one on the learning rate-- I'm not quite sure that constitutes the best linear fits that I've seen in my life, but I guess it works. It's probably fine. The models do train. The models do get reasonable things. I think this is the one, I would say, like drawback of the learning rates approach. It's not really clear that if you try to do this grid search kind of thing, unless your grid is in the right place, you've got a lot of quantization error. And then you're not getting very good scaling law fits, and you end up with something that looks like an arbitrary linear thing going on. OK, oh, yeah, sorry. How come for the same number of training swaps there, you get multiple data points with a large range of optimal vertices? Why is it that you get a large range of optimal learning rates for different training FLOPs? Is that the question on the y-axis over here?
[19:49.660] Yeah, so the reason I think there is because of the changes to the model sizes and things like this. So 2024 was back in the era where I think every open model builder was replicating the entire scaling law stack. And so DeepSeek also complements the learning rate analysis with basically a Chinchilla replication. And to do that, they're also going to do the same thing as MiniCPM. They're going to follow the WSD-style learning rate. They have a weird little variation where they do two decay phases instead of one. I am not really sure why this is. It hasn't really been a thing that has taken off afterwards, but it is a thing that they do. Regardless, the whole point here is to say, WSD is an effective approach. And you can use that to replicate Chinchilla. And DeepSeek has much nicer scaling curves than MiniCPM because of their choice to go IsoFLOPs. So they get these very nice, clean sweeps across the fixed flops ranges.
[20:59.050] And then they also get fairly nice and clean Chinchilla-style trade-offs as before. So this allows them to essentially replicate the Chinchilla results and get a similar kind of data model scaling law. Really, the point that I'm showing with this is just to say, look, the Chinchilla laws have been replicated by other people at large-ish compute ranges. So that analysis is generally sound. Finally, I think, this is the-- I'm a little bit sad that these kinds of plots are now somewhat rare in open-source release papers, but this is really the final punchline of a lot of scaling law work. You train a bunch of small, curated models at smaller compute ranges, and then you draw out the power law. And then in the end, you get something pretty close to reality if you scale out enough. So the two stars are the models that they actually trained. The gray dots are their scaling law fits. And they're getting pretty good predictions on their star. I mean, it could be better, but it's not bad at all.
[21:59.650] And so they're able to get some pretty accurate scaling law predictions on real open-source models that they've trained. OK, I'm actually I'll stop here. If people have questions about the DeepSeek stuff, feel free to ask them. I think my one-line summary here is to say, this is the other approach to dealing with the sensitive hyperparameters. You can either attempt to stabilize them, or you can attempt to just scaling law fit them and then just go with the scaling law. And DeepSeek represents the more fit the scaling law approach to this problem. Good, OK. So there are other models, more recent ones that have nice scaling laws that they talk about. As I said before, though, I think a lot of this work is now-- what's the-- it's all known to everybody. And so no one is really putting a lot of this stuff into big sections into the paper. So if you look at Qwen 2.5, they say something like, OK, well, we need to tune some optimal hyperparameters, like batch sizes and learning rates.
[23:02.245] So we do a bunch of scaling experiments to figure out what the optimum are. And then we have a scaling law, and we predict the optimum. And then we're able to do that even for MOEs. So this is basically the same thing as the DeepSeek analysis that they're talking about in Qwen 2.5. Same thing with Qwen 3. By the time it's Qwen 3, they say, all right, we've already figured out what to do in Qwen 2.5. We're just doing exactly the same thing. So this has become part of a very standard recipe that you use in order to nail down the optimal learning rates and batches. Kimi K2 is probably the newest paper that I'll talk about of the series. At this point, I'm just going to quickly go through several papers maybe to give you an overview of what people are doing in terms of scaling work in the open-source model releases. So Kimi K2 and a few others have focused on essentially MoE scaling laws.
[23:57.440] So it's 2026. Most people know how to do a Chinchilla scaling law or even a learning rate scaling law. What things are new? Well, a lot of these companies are now switching to MoEs, fully switching to MoEs. And if you're doing that, you really want to know, am I training the right amount of active parameters? And so they're doing things, like, OK, we're going to vary the flops and look at my losses. And I'm going to do so for different levels of sparsity so that I can figure out how much sparsity helps or hurts my losses. And their conclusion, of course, is if you look at the different sparsity levels, the more sparse your network, the better your validation loss,
[24:33.640] given a FLOP-- unsurprising. But now the fact that they can quantitatively know what the relationship is between sparsity and loss means that they can make more rational decisions about how much sparsity to put into their model. And they look at this, and they say, OK, what we're going to do is we're going to use sparsity of 48, because at that point we're hitting diminishing returns on our scaling loss. So you can see how this kind of analysis is used to justify different architectural decisions across different open-source releases. Hunyuan, which is another open-source MoE from a few years ago now, they do a very similar thing. I think every company, once they switch to the MoE, has to do this kind of MoE scaling law to figure out how the number of activated
[25:19.620] parameters or the amount of sparsity city changes training loss. And then they redo a lot of these analyses. So for example, in Hunyuan's case, they end up with a 96 data point per active parameter ratio. In their case, they're fixing the sparsity level. And so that's not a thing that they're searching over. LLaMA 3 also has some nice scaling laws. Unfortunately, they're not particularly useful ones. They do IsoFLOP-style scaling and get slightly different but similar ratios for token-to-model. The most interesting thing that they do in terms of scaling in LLaMA 3 is the plot on the right here, which is, they have the usual scaling law, which is the left panel.
[26:01.340] That's saying, if I have more compute, then the logprob on the answer for a benchmark decays consistently. And on the right side, they show, OK, well, if you have better log loss, then the downstream accuracy also goes up. So they're basically saying, well, we can just fit a sigmoid on these purple points. And that might give us a sense of how good these different models are. You can see their systematic deviations from the curve. So I don't really know if I would really buy that this curve is like the one truth that's going on. But you can certainly see that it is consistent with a sigmoidal scaling mapping from the loss to the accuracy, which I think
[26:39.730] is very useful because we've been talking only about log losses. But there is a tight coupling in some cases, or in many cases, between log losses and downstream accuracy. And then for one last paper that I want to talk about for scaling, MiniMax-01 from last year had this nice set of scaling experiments where they were doing the exact kind of thing I was saying is a useful affordance of scaling from two lectures ago, and that is architecture variations. And so they wanted to understand, OK, I have lightning attention, which is some kind of a linear attention architecture. I have a softmax version like full attention, and I have a hybrid version that mixes the two. How do these things scale? Is it the case that if I scale up, one of them will get much worse than the other? And so you can see, they're replicating essentially the Kaplan plots of, as I increase FLOPs, how do the learning curves change? And is the number of parameters that I need different for different levels of compute for these different architectures?
[27:46.970] And the answer is, for the most part, no, that hybrid and lightning and softmax are all comparable for the most part. And they require similar model sizes. And this set of plots is used to justify their use of the hybrid architecture in the actual deployed system. So this is exactly the kind of decision making that you would do based upon a scaling law. So that's basically what I wanted to cover for the recent scaling law recipes from the various companies. So the two things, as I said before, for the DeepSeek recipe. Really, the key thing is the scaling analysis on the batch and learning rate. This is the scaling-based approach. And MiniCPM follows a more of a muP-style approach to try to get invariances. And then we see some trends. We see that Kimi and Qwen and others, or in Hunyuan, are doing these MoE scaling experiments as they switch to MoEs.
[28:39.730] We see some architecture work by MiniMax, and we also see some IsoFLOPs analysis being replicated by several others. And to return to our previous point, the point here is, I think a lot of companies are now less detailed because the core scaling machinery, like doing Chinchilla or doing learning rate scaling, is now fairly well understood by everybody. So there's no necessarily extra value in showing those experiments for some of these releases. Or at least that's my guess for why we have much fewer of these kinds of experiments in the release. OK, I'll take one minute and then pause and let you digest and ask a question or two, in case you have questions for this part. OK, so all these cover the proofing stage. So my question is, what changes and what do not change when no post training is standardized process for model development? Yeah, so the question was about, what changes about this when you do-- or when you think about post training, I guess, broadly?
[29:47.990] I think that's still a big, open question. There's no great way to do fully integrated scaling with post-training accounted for, because post-training can sometimes really change the nature of what pre-training you should do. The closest things that I can think of to thinking about post-training is there's been some work on understanding coverage or diversity notions in pre-training that would help predict what happens in post-training, but it's still pretty nascent work. I don't think there's any good answer to the question of, what do we do about synergy with post-training? Cool. OK, so now we're going to talk about the other piece. So I think the two remaining parts of this lecture map on cleanly to the two examples. So I'm going to talk about one approach, one set of approaches, which is about scaling laws for learning rates, batch sizes, and optimizers. That corresponds to the DeepSeek approach to the world.
[30:50.080] And then I'm going to talk about the muP stuff. That's the MiniCPM approach to the world. Can we just reparameterize everything to make it so that we don't have to worry about scale? So we're going to now talk about the DeepSeek approach. How do we scale learning rate properly? And so there's going to be a hyperparameter part and an optimizer choice part. And I'm going to draw quite a bit from this pretty recent preprint from the StepFun folks. I don't think there's a great, really, truly reliable, robust study on hyperparameters, especially learning rates. But this is probably getting somewhat close. The StepFun folks train credible large models and release them. And they burned a ton of compute in grid searching the space of hyperparameters in this scaling law study. So I'll go through this and show you some intriguing empirical phenomena about how learning rates, batch size, and all of those affect downstream losses.
[31:51.300] So lots of people have talked about lots of different ways of scaling, learning rate and batch sizes as roughly as a function of compute. And the StepFun survey-- not survey paper, the hyperparameter tuning paper has this very nice table that captures a lot of this. And I'll point out a couple of different things. Some of this we already know. So the very top row, this OpenAI scaling law-- this is from Kaplan. If you look at the batch size over there, notice how it says 2eE18, capital L to the exponent. That's critical batch sizes right. It's basically saying, you should scale your batch size as a function of your terminal loss. Great. DeepSeek we just saw. This is the compute-based scaling loss for both learning rate and batch size. And there's been a couple of other folks that have proposed various kinds of scaling, including MiniCPM, which we talked about as well.
[32:47.840] Another important thing that I want to point out is that these formulas for learning rate and batch sizes, they don't even have the same inputs. People don't even seem to agree, what are the variables that should affect these learning rates? Like, the OpenAI one is terminal loss. Whereas, the StepFun folks, they claim that the right batch size scaling is a function of the data. So there's a lot of different disagreements about the nature of scaling for hyperparameters. And I think one thing I'll say as I go through this is not to take even this last row as gospel. A lot of these things are pretty brittle in many ways, but the experiments and the phenomena are things that you can learn well from. So I'll do that. So the design of this paper is very similar to the DeepSeek one. They're just going to train a ton of models to grid out both the learning rate space and the batch size space and do them across a range of models and range
[33:44.810] of data set sizes. And by doing so, they hope to essentially saturate the hyperparameter space for a set of small models and then nail down what the optimum is. And much like in the DeepSeek design, they're going to be-- after they grid up the space-- in this case, they have a contour plot. But the underlying objects are just, they run a bunch of models. They aim to try to minimize the losses. And the contour plots are quite nice. I feel like they give a good visualization. They ran a pretty high resolution grid. And thanks to that, one of the things that we can see is that they can actually plot out a fairly high-resolution contour of what the hyperparameter space looks like. In this case, this is a slice at 1b parameters and 100b tokens, I believe. And you can see on the right side, that surface is the hyperparameter landscape. As a function of the batch size, which is one axis, and learning rate, which is the other axis, how does the loss change if you fix the total amount of data in the model?
[34:46.710] Now, we see that each slice, if you fix a batch size or you fix a learning rate-- each slice is kind of nice and convex. And it's almost even smooth. And this is quite nice because it means that you can search for the optimal hyperparameter fairly nicely, and you can be confident that the minima is not too far off. This was very jaggedy. You might have some very serious doubts about whether this kind of program of gridding the space is even viable. Also, this should maybe give you some nice intuition about the smoothness of the hyperparameter space in general. And I think there's actually one really interesting general scaling trend. And looking at other related papers to this topic, I think, I don't know if I've seen any contradicting evidence to this observation. So I think the one thing that seems very interesting and one of the key takeaways of this work is, essentially, if you're looking at the optimal batch size, the only thing that seems to depend on is the total amount of data that you're training on.
[35:46.650] So this right-hand plot is the one to look at. So x-axis here is the amount of data that you're training on. The different colored dots-- those are different models. Notice how the different colored dots seem to follow roughly the same trend line here, fairly predictable one. And this is in log-log scale, so you should expect some of a power law that's a function of D. On the left panel, we have the dependence of learning rate. And we see that this is a very different behavior. So if you go from blue all the way to this pink color right, bigger models have smaller optimal learning rate. And larger amounts of data have higher optimal learning rate, which is counterintuitive, actually. And so this might be a little fragile. There are other kinds of papers that argue that the dependence on D should be reversed. But the trends here are at least on this experiment, are quite clear.
[36:45.120] And one final thing that I want to point out is, the learning rates are actually pretty robust. And I think this is backed up by a lot of experiences and experiments that people have. I think people have a sense of-- actually as a poll, do people know what a good hyperparameter is? If you were told to train a language model tomorrow, would it be 1e negative 3, 1e negative 4? I feel like we roughly have a sense of, actually, this range of hyperparameters for learning rates is generally a good starting point. And that is generally true. So the scaling laws that they got actually transfer over to, for example, MoEs to a certain extent. The yellow stars, which might be a little harder to see, are the predictions from their scaling laws. The red x's are the optima that are derived directly on the MoEs. And so you switch to an MoE. You don't end up with a very different scaling law, as long as you're roughly controlling for active parameters. The one thing that does seem a little bit more potentially dicey is, as you shift the training data, the optimal learning rates and the optimal batch sizes do shift a little bit, which means that all of these phenomena are likely contingent.
[37:57.020] The specific numbers that are in these scaling laws and so on are potentially contingent to the data that you train on more so than others. And so I want to flash this back up before I move on. So you end up with some of simple scaling law at the end, which says, OK, batch size maybe should be roughly square root of the number of data times some constant. Your learning rate should scale upwards with batch size or upwards with data but downwards with the model size. And that's a surprising effect but not maybe crazy, because these experiments often, in many cases, are assuming Chinchilla-like scaling, in which case, those two would cancel. And as you increase compute, your learning rate should go down. And hopefully, that Chinchilla comment was clear. What I mean is when you're scaling by Chinchilla, N and D are both driven by compute. And as compute goes up, you expect your learning rate to go down.
[38:51.808] And so you will end up with something like the DeepSeek Law, where the learning rate decreases with compute, and the batch size increases with compute, albeit with different exponents than the DeepSeek Law. Cool. OK, so in some sense, this is the newest and most large-scale version of the DeepSeek analysis. And visit one way that, if you were doing the scaling-law-style approach, you would approach this problem of fitting a learning rate. So that's the hyperparameter story. If someone has questions, feel free to stop me. And then next, right after, I'm going to move on to talking about optimizers. I think this is one thing that didn't have time to talk about in the architectures lecture. That was already a pretty packed lecture. But I think it would be a disservice to if I don't talk about optimizers. And optimizers are scale dependent, so this is a good place to put them and talk about some of the latest and greatest developments in optimizers
[39:52.630] and their relationship with scale. Oh, yeah? So I guess I'm trying to figure out how to interpret this. Does that mean we should just take the scaling laws that we have and then apply them, rather than running our own grid search? Yeah, I mean, I think it depends what regime you're in. The question is like-- oh, sorry, I have to repeat it for the online folks-- is should we just take the scaling law and run them, or should we just do our own grid search? It does depend on the amount of compute budget that you have and how much you need to nail it. I do think, for example, maybe the StepFun law-- if you're in the right compute regime, like close to where they were running the grid, these are probably really good defaults. I would bet on these over something I pull out of my hat. That said, if you were doing your own big pre-training run, you're probably going to have minor differences.
[40:44.730] Someone might say, I really believe we should have a big weight decay. My architecture has a big weight decay. We don't know necessarily that those things are going to scale the same. And so you might want to actually redo a lot of this. And this is why I think a lot of the reports I showed you were redoing Chinchilla, redoing all these things, because there's enough minor differences that people want to check and make sure that they're roughly first-order correct with these results. I think this will come up again in this next optimizer part. And I was talking about this with [MUTED] earlier, which is, scaling laws have this very scientific feel to them. It's like, yes, fit these lines and extrapolate them. But ultimately, a big part of scaling laws is still vibes. It's like, do I really believe that the experiment setting of this StepFun survey is similar enough to mine that they will transfer? We can't possibly right.
[41:34.210] There are always tiny differences. Good, OK, so now I get to talk about optimizers. And I think there's been some really exciting developments in LLM optimizers, and maybe even deep learning optimizers in general over the past year or two. And so this is an exciting time to talk about them. And so a story that I want to start with and one that should maybe motivate this subsection is this comparison over here, or maybe we'll start on this left panel. So this is from NanoGPT speedrun, which inspired our assignment 1. And there you have a really teeny, tiny model that you're training on a fairly short amount of time, and you want to get the, I forget, 3 point-some loss as quickly as possible. And many people have been hill climbing on this benchmark. And I think one of the more remarkable things that have happened on the benchmark is as a benchmark called muon in purple over here came and got really significant gains over Adam, which is blue.
[42:36.460] It's right underneath the green line over there. So you can just look at the green and the purple and that gap. And you say, wow, lots of big gains from just changing the optimizer to this muon thing. And it's not even slower than Adam, much slower than Adam. That's great. But then there's a question of, well, is this thing actually good at scale? And then you go and run some scaling studies, and maybe they're not quite as good. This is a really tricky and important part of how we do research. We have a small-scale experiment. We have something that's scale-dependent, for sure, like an optimizer. And then we see really big gains in the small scale. What do we do about the large scale? And more than providing a solid, concrete prescription, I want us to maybe almost even provide a story here so that you can get a sense of what we can do and what things we can't necessarily know.
[43:27.300] So actually, before I move on to the next slide, and this kind of thing happens all the time, where you're like, come up with a great thing small scale. And it just doesn't maybe pan out large scale. So what's the difference? I'm going to spend a few slides talking about Kaiyue and Tengyu and Percy and David's paper on basically doing a large-scale set of experiments on comparing optimizers and trying to do that comparison better. And there's a lot to learn about scaling and other things from that, but also just standard deep learning stuff. There's really basic things. Like, the hyperparameters for different models are different. So if you're scaling things up, maybe they even have different scaling exponents. So you have to tune Adam really well. So if you tune it badly you get this darker, brown line. And you're like, wow, this thing is much worse than all these other algorithms that I have over here on the left.
[44:18.275] But then you change the learning rate a little bit. And you're like, wow, actually, all my gains went away. Similarly on the right, you can have cases where maybe you've tuned your learning rate well. But now you've got other stuff like weight decay that is not optimally tuned. And so if you use the same weight decay, then you're certainly going to get a much worse result for this green line. If you pick this small weight decay, you'll get a much worse result for this green line than if you pick the optimal weight decay. So there's all sorts of pretty tricky things going on, technically not a scaling thing, but a very important general empirical machine learning thing that I want to point out. But maybe more relevant to the discussion in this class right now is the scale dependence of these things. And I think this is a good case study and example of how to really think about algorithm development as a function of scale, because to me, there are two variables that are always important when
[45:11.970] thinking about scale. And this paper exposes, to most extent, both of them. So one axis is compute. And if you're fixing the model size to data ratio, like the two are coupled together, then this is essentially a compute axis on the x-axis here. So whenever you're doing a study essentially scale up the compute and look at the trends as a function of scale. And we see right here that if we look at the relative speedup between Adam, which is at 1.0 by definition, and all these other alternatives, we see that muons are really good at small scale. And that gain diminishes down as we increase in scale. That's not really necessarily the picture that you want to see. There's another axis that you should always be vigilant of if you're doing algorithm development, and the other one that you should be vigilant of is the Chinchilla ratio. So the Chinchilla ratio can be thought of. It is the ratio between how much data you have and the amount of model you have.
[46:12.930] And this is a really, really, really important parameter because some algorithms work well in the overparameterized case where you have a way bigger model than you have lots of data. Maybe your algorithm adds some implicit regularization. Or maybe it is data inefficient, but it works well because of compute reasons. So some algorithms might work well at the small Chinchilla ratio case. In contrast, other algorithms might work well in the large Chinchilla ratio case, where you have lots of data per unit parameter because you can pack knowledge more efficiently into the parameters. So this is a big, big big confounder. Even papers that I think have pretty good hygiene in doing scaling experiments on the model size axis often neglect to or just don't have the compute to study the Chinchilla token to parameter ratio. And this turns out to be surprisingly important a lot of time.
[47:06.010] In this case, it's not. In this case, you see that the optimizer gains are pretty consistent. The ratio is between all the lines are the same at different Chinchilla ratios. But this is not always the case. You always want to be vigilant about both of these scaling factors. That's quite important. And then I'm going to also point out, if you have a new algorithm or a new setting, establishing scaling is quite nontrivial. So I've taken this from one of-- Will Held was working with Marin, which was Percy's open-source language model training project. And there they were basically trying to prove out a particular set of hyperparameters. This is even like a new algorithm. So this is cautious AdamC, which is like an Adam variant, square root, batch size scaling and a bunch of other standard-looking things.
[47:55.880] And they're doing essentially the same Chinchilla analysis that you've seen all the other papers do. But because this is not necessarily a paper, this is a bunch of people doing experiments, you don't just publish the successful one. They have the niceness of publishing their failed runs. And the example here is that they're doing basically standard Chinchilla stuff, and it's looking great. If you go up to this dashed vertical line at 10 to the 20 something, this is a really beautiful-looking scaling loft here. And then, as you scale up, it's like, well, it's a little bit worse. It's a lot worse. And then, wow, now that run is not so good anymore. So it's totally blown up on you. And the point here is that even good-looking scaling trends for many orders of magnitude can suddenly bite you. It can deviate from linear, and then it can then blow up immediately.
[48:44.760] And in their case, the fix here is just picking more careful muP-style parameterizations. Optimizer changes. They switch to a different kind of optimizer, and they get nicer scaling across more orders of magnitude. But the point here is to say, when I go back here, and I say, OK, well, you should all do scaling experiments, you should scale all your learning rates and so on, that's not easy. You end up fairly nontrivial fraction of the time with plots that look like this. And it's all horrible, and you're not quite sure what has gone wrong. Good. OK. And then, so the last thing I want to end with is just talking a little bit about muon and the closing of the muon story. So I feel like I have to talk about muon in a little bit more detail than I already have, just because it's beginning to be integrated into a large scale model training runs. That's the inclusion criteria for a lot of new research topics for the class. Has it been in a big training run?
[49:48.300] And muon now checks that mark. And it's also just an interesting algorithm, so I'll mention it and talk about it for a few minutes and then move on. So in gradient descent, we usually have a fairly simple thing that's happening. Ignore line for the moment, sorry, line 5 for the moment. And then you're going to get roughly a standard momentum-based optimizer. I take a gradient. I add it to B of t. This is my momentum step, so I continue my momentum from before with parameter mu. And then I update my parameters based on this momentum gradient. Very standard momentum gradient stuff. But then I think the really interesting and, to me at least, a bit of an out of the box set of ideas is to say, well, gradients have this property that they treat, all the parameters are the same. But we know that not all the parameters in a language model, or even a deep neural network, are the same.
[50:40.920] Some parameters are like vector parameters. They belong to the RMS norm. Some parameters are matrix parameters, like for the attention or for MLPs. And matrix parameters are probably different than vector parameters. And so that's the starting point for this idea and I think a really interesting one. And so the one thing that you could do about matrix parameters is you can start to look at their spectra. They are matrices. Matrices have eigenvalues and singular values. And so you can do things to the singular values. And so what you do here in Muon is this fairly simple idea. You compute a gradient. You take a momentum step. And now I take my update B of t. And instead of directly applying it to my parameters in line 6, what I'm going to do is I'm going to do this NewtonSchultz thing, NewtonSchultz 5 to be specific. And what that is it's going to orthogonalize the B of t matrix.
[51:33.567] And so I've written the operation down there on the bottom. Let's say I take B of t. I write it as singular value decomposition, USV transpose. And then I make all of my singular values 1. So if my singular values are really big, I'll shrink them back to unit. If they're really small, I'll expand them back up to being unit. And then I take UV transpose, and that's my new update that I'm going to apply. If you want intuition for this, and you already have intuition for something like adagrad or these other, or Adam even. In adagrad or Adam, you're going to divide by the size of the gradient. So every coordinate is roughly going to be the same size. In muon, you are operating in the spectral norm instead. So every kind of direction is going to be unit size. And notice that this only makes sense for a matrix. I can only define this notion of a singular value decomposition for a matrix valued B of t. So the interesting thing about muon is it's a matrix optimizer, so you apply it to these matrix valued parameters. And for the vector-valued ones, you might still use something like AdamW because those don't really have any interaction with the NewtonSchultz term.
[52:40.990] Also, NewtonSchultz is an approximation to that operation that only uses matrix multiplies. So it's systems efficient and not exactly the orthogonalization. But OK, why have I spent all this time talking about muon? Well, first of all, it's very powerful for things like the NanoGPT speedrun. It was discovered for this set of things. And so you might think that it works well. As I said in the previous slide, there have been some scaling studies that have argued that the gain for muon decreases as scale goes up. But then I want to close off this story, which, I think, initially this came out. And I was like, wow, muon is interesting. I wonder if anyone will scale it up. And then [INAUDIBLE] and others scaled up muon. And they said, well, maybe it doesn't work super well. And I thought the story was closed here. I was like, well, maybe this is done.
[53:33.110] No one will spend the big compute to scale this guy up to the full pre-training size. And then Kimi K2 comes out, and that whole model is trained with muon, fully with muon, with a few bells and whistles to prevent muon from blowing up on you. They found a whole bunch of instability issues. But Kimi K2 is an outstandingly good model. It's a very good model, and the training curves look reasonable. And so really, what this suggests is like, well, muon worked at scale. Is it better than Adam? Kimi K2 certainly doesn't have any ablation, so we don't know at that scale. But it is certainly a valid and workable solution. And so what's the point of this story? It's to say, It's really hard to whether something works at scale. It's very, very hard. But thus far, a lot of small scale experiments and transferring them up to large scale has been the way by which we do science.
[54:25.030] And so you still want to respect things like the NanoGPT speedrun. We learn quite a bit from that and get good ideas. And then eventually, they make their way to the big models. We'll see if people continue to do ablations at big scale and really nail down, is muon better than AdamW at scale and so on? Those are interesting questions to which we don't yet have the answers. OK, I'll pause here for a moment. This is the end of my optimizer section. I want to make sure everyone's good, also especially because I was talking about muon. And that's a whole new can of worms that I haven't talked about much before. Yeah? [INAUDIBLE] letter [INAUDIBLE] faster on GPU? OK, [INAUDIBLE]. Great. Yeah, the question was, is SVD fast on GPUs? SVD is not fast on GPUs, but they're not doing SVD. What they're doing is NewtonSchultz. NewtonSchultz is a matrix multiply only finite iteration algorithm that approximates the process of orthogonalizing a matrix.
[55:18.940] So there's several clever pieces here. One of them is the idea of treating it as a matrix on its own. And then there's this orthogonalization idea, and then there's the systems piece of doing it with NewtonSchultz. So those are all important, interesting pieces? Yeah, back there? Are hyperparameters for nontrivial [INAUDIBLE] ones? Are the hyperparameters here-- no, they're different, yeah. I think there are several people that have been working on this muon-style line of work. And one of them, I think, like Jeremy Bernstein, I think he's been saying, actually every layer should have its own learning rate. I think the ultimate limit of this is like, you go from-- or even, sorry, every layer should have its own optimizer. The ultimate limit of this is to say, every parameter in a transformer is different. Why not have a different optimizer for each one of them?
[56:08.660] And so you would have to tune hyperparameters for each one of them as well, which I do not want to do. But that's maybe the future of optimization. Yes? What's the validity of like-- like, on the learning rate and batch size, you just consider these two things. And then you pick the best out of those, and then you can use it to tune other things. But do these things interact with each other? Yeah, or maybe the version of the question that I interpreted was, a lot of these people are just tuning learning rate and batch size. And they do this on the grid. But presumably, these things interact with a whole bunch of other hyperparameters. Aren't you worried about that? And I think, in general, yes, you are worried about hyperparameter interactions. But realistically, they grow exponentially. This is not a thing where you can grid out the whole space. And so I think the thing that people often do is they grid the stuff that's most concerning and most sensitive. And learning rates are very important. I think, if you ever tune stuff, learning rates are the most important thing to tune.
[57:08.880] So they very aggressively go on that. And then stuff like weight decay, they might do a univariate sweep afterwards to be locally optimal with respect to those parameters. Good. OK. All right, so the final part of this lecture is I want to talk about muP. muP is a bit of a mysterious object. Even with many papers written about it and, in fact, even many implementations, I don't think all of them even agree on what the underlying math is or what the underlying implementation should be. But there's a core set of ideas that are shared, and it does seem generally effective. And so I'll talk about my interpretation of the whole program, let's call it. So the goal hopefully is very clear. And I think this picture is the goal, the picture being, well, if we increase the size of our model, in this case, just width, our optimal learning rate normally shifts.
[58:03.590] But ideally, we want the optimal learning rate to be the same, no matter what. That is the game that we would like to play. And the knobs that I'm willing to tune to do this is I'm willing to mess with the initializations at a per-layer basis. I'm willing to mess with the learning rates on a per-parameter basis as well. And then I'm sometimes willing to scale up or down things like residual connections based on the size of my model. So that's the game that I'd like to play. And how well does it work in practice? There's been several papers that have studied these fairly extensively. Cerebras which is really a chip company, also has a language model training arm run by Hestnes of the scaling law paper fame. And I feel like they do very interesting things because they now are writing lots of papers about muP scaling of language models, and they have Cerebras GPT from a couple of years ago, where they were training the 0.1 to 13b models using the usual
[59:02.970] Chinchilla recipe, but also including a muP variant, which they use to try to tune the hyperparameters more effectively. And I think one of the nice things about this is they find that if they use the muP parameterization, then their scaling law fits are generally more stable. Their projected muP trends are almost right on the money for their actual models. Whereas, I think their non-muP models fluctuated much more wildly in terms of the loss predictions that they had. And so their claim was, OK. MuP is maybe a really good idea. So this is validated at scale. You saw MiniCPM on others. Many people have trained models with muP variants. So I could also talk a lot about training or examples of models trained with muP. But I think maybe the more interesting thing for me to do right now is actually to talk about the conceptual foundations and the math for muP.
[01:00:01.735] And if you're interested in this, and you want to read a somewhat accessible paper, I would recommend this one. This is the paper that I think, for me, was the clearest. There are several papers. Greg Yang, who started this research program, has a series of papers I'll call tensor programs. I find them somewhat inscrutable. Jeremy Bernstein, who was also working on the muon stuff, has written this nice-- not recap paper, like a review paper, putting it into a different framework, which I think is very accessible. And there are several others by physicists that are also interesting in terms of explaining what muP is. The Cerebras people actually do a bit of that as well. But anyway, returning to this paper, the core idea of muP is based off of two kind of assertions. And this is a very physicist way of thinking about this. So if I make my network big-- let's say I'm only changing the width for now.
[01:00:53.280] I'm making the width of my network big. I'm going to assert two invariants that should hold as I do the scaling limit. First of all, the activations at initialization-- so if I initialize my model, and I put in some random data, the activations should remain the same size roughly as a function of the network. So if they blow up with the network size, or they shrink to 0 with the network size, I have chosen a wrong parameterization. And that makes sense. My activation should be roughly the same scale, regardless of the size of my network. Now, after one gradient step, the change in my activation should be O of 1. So this is what people call feature learning. So feature learning means if I take a gradient step, things in my network change a significant amount and that this amount should be a fixed function of the width of the network. In case some of you are theorists, you might be familiar with neural tangent kernels.
[01:01:46.920] Neural tangent kernels would have a very different behavior. Those would be the kinds of things where the change in activations actually vanish as a function of the network width. And you don't want that, because you do want the network to change and learn features and do things even in the limit of your big neural networks. So these two are the invariants. They make sense. They're basically saying, activations and the change in the activations should remain roughly similar as a function of the width of the network. And then note that I'm going to sometimes go back and forth between individual activations, which is this. And then the norm, which-- if each activation is roughly O of 1, the norm of the whole vector of activation should be square root of the number of hidden units in that layer. OK, so that's the main thing. Now, the thing about muP is it's interesting because it starts with this very simple and palatable set of assumptions.
[01:02:39.343] You might all agree that this is a reasonable thing to do. And then we're going to make a number of additional strong assumptions. But then, in the end, we're going to get a very simple set of scaling rules that drop out of this. So remember, what we want to do is we want to make sure that the activation is at init are roughly unit. So what do we do? Well, if we have a simple deep linear network, and we initialize all of our weights as a sigma squared scaled Gaussian, then what is going to be the spectral norm of this matrix? Well, you can write that down. I'm not going to ask you to verify that, but that's standard matrix concentration arguments for the operator norm of a random Gaussian matrix. And then now note that in certain regimes where the fan-out is, I think, not much bigger than the fan-in, you're going to get that the output of the layer is similar to the product of the input of the layer times the operator norm. This is always an upper bound, but this is approximate in certain high-dimensional regimes.
[01:03:39.130] Now, if this is true, now what can we do? Well, we can solve for sigma to try to keep our activations the same across all the layers and across all the width. So imagine that at layer L minus 1, I have square root of L minus 1 activation. And this is, remember, consistent with this notation that, at least for this whole vector, I should have square root of n of L. So this is our base case for induction. And then our inductive case says, well, remember that I have this operator norm at layer L, which is this function of the fan in and the fan out. And so this is going to be equal to this if I pick a particular sigma. So let's say this is the ansatz. I'm going to assume that this is the answer, and then we'll see what happens to the induction. If I pick this as a hypothesis, and I plug this into the sigma, notice that I get that the ratio of the operator norm is this, the ratio of the fan out over the-- sorry, square root of the fan-out over the square root of the fan-in.
[01:04:35.110] And then I plug this back in here. And then I'm going to get that the norm of the output is going to be equal to the square root of the number of units in my layer plus some lower-order terms that I'm not going to bother to keep track of. So where this sigma came out of-- I mean, this is drawing it out of a hat. But we can verify that if I plug this noise initialization into my formula, then, essentially, at every layer, the size of my activations is roughly equal to the square root of the number of hidden units at that layer. There are some assumptions made here, like this approximate equals. But this is mostly a legitimate set of assumptions. So condition A1 is pretty simple. Condition A2 is a little bit hairier, to be honest. So really, we need to deal with the updates now. And I'm only going to consider the case of a single example in a batch for SGD for linear neural networks, so deep linear network.
[01:05:36.212] So what is this going to look like? Well, the change in my weights for a particular layer L-- if you think about how this is going to behave, it's going to be a rank 1 update between my gradient for the loss and the activations. This is true just because of the behavior of the backwards. The forwards is a rank 1 multiplication. The backwards is going to be a rank 1 outer product. So now this is the change in the weights. And now what is the change in the activations at that layer? Well, the change in the activations-- I just take h equals W times hl minus 1, and I expand this out. And I expand this out into this formula. This is going to be exactly the change in the activation at layer L. And it's going to have two terms, one that is the direct term from changes in the activation at L minus 1 and the cross term in some sense of what happens because my weights have changed, not just because the activations of the previous layer have changed.
[01:06:31.487] So there's two terms to this update. And so this is the change in the activation at layer L. And now what I'm going to do is I'm going to look through each of these three terms. I want each of these three terms to have the same order of magnitude, and I want that order of magnitude to be square root of n of L. Why is that? Remember that, if I go back two slides, the condition that I want for A2 is that after one gradient step, the change in activation should be O of 1 or omega of 1. And so if I want that to be true, then the norm of that layer should be square root of nl. And so to do that, what I want is for each of these three terms to have that order of magnitude. And if that's true, and there's no cancelations, then this is going to be the right order of magnitude. You can see the sketchy pieces adding up. We're assuming things like no cancelations and so on, but this is giving us some interesting ways of thinking about the magnitude of each of these objects.
[01:07:25.740] So at this point, the right way to think about this is not like, OK, Tatsu is sitting here doing rigorous mathematics. What I am doing here is I'm basically keeping track of all the orders of magnitude of all the terms. And I am saying that, given our assumptions, how big should these roughly be? It's physicists math. No offense to physicists. This was developed by physicists in this way. OK, so now really-- sorry, go back one slide just to make all this clear. So the first term, this comes by almost assumption. It's the inductive assumption. We assume that delta h of l minus 1 has the right order of magnitude. So this comes for free. These other two terms-- these are trickier. We have these terms that multiply the change in my weights in operator norm times square root of L minus 1. And so because these two are the order of magnitude of both of these terms, we know how big this delta L needs to be.
[01:08:21.245] We can solve for this because we know that this left-hand side needs to be square root and L. And then we have this term, so we solve for this operator norm. And that's exactly this quantity. So this is going to be the goal for the next slide. I want to try to figure out what kinds of learning rates are going to give me an update that's roughly scaled as the square root of the fan-out over the square root of the fan-in. So that's the goal. As I said, after I pick a learning rate, this should roughly be true. So how can we do that? Well, we're going to make yet another assumption. And this is probably the least palatable of the assumptions, in my opinion, which is that we're going to assume that the loss update is going to be scaled by unit. In other words, we're going to assume that our model is actually making some appreciable, nontrivial progress as a function of scale and that that's roughly similar. Not quite sure that's palatable to me, but that is yet another assumption.
[01:09:14.740] So we throw that down. But now the point of this assumption is that if this is true, the rest of the math works out quite cleanly, because the change in the loss is the product of my gradient plus how much my weights changed. This is Taylor approximation. And then because this is rank 1, I can rewrite this thing as a Frobenius norm and as an operator norm. And so this right-hand side now allows us to get what we want. This is O of 1. This term we want to be square root n over square root l minus 1. That's this term on the right here. And then finally, I have this gradient of L term. And this is going to be square root of NL minus 1 over square root nl. And then so now what eta do I need? Well, the eta that I need is going to be solving for this quantity, which is the ratio of these two guys, which is nl over nl minus 1. So we made this one assumption that the left-hand side of this is O of 1. And then we've solved for all the quantities, and we've ended up with a fan-out over fan-in ratio for a learning rate for a layer.
[01:10:13.770] One side note that I'll make is that for Adam, because of the changes in some of this, you're going to get a 1 over nl minus 1 scaling, which I'll mention here. So I've walked through a bunch of hopefully intuitive-looking scaling arguments for various components. But in the end, maybe you are interested in, what does this mean? What does this do for me? What this means is that we're going to change how we set our initializations and learning rates a little bit. So our initialization is going to look like-- so for muP, the blue box is what we do. Our initialization is going to be 1 over square root nl minus 1 times a component that is 1 over the square root fanout over fanin. So if this is 1, then this is equivalent to the usual initialization. This is the standard parameterization. But if this ratio is not smaller than 1, then you get basically something different.
[01:11:08.790] You get something that's closer to the square root fanout over fanin. The learning rate normally you would set to be a constant as a function of your model size, but you would want to set this to the ratio of fanout and fanin. And maybe more interestingly, for Adam, this whole thing is a little bit more different. I'm not going to even touch the derivation for that, where you drop this 1 over nl term, and you get 1 over the fanin as the layer specific learning rate. And so basically, layers that have big fanins for Adam get smaller learning rates. It allows you to have these layer-adaptive learning rates that are quite interesting. So that was the derivation of muP. I went through a lot of details. I do want to maybe talk about what the high-level point of that was. And the high-level point of that is, muP as an argument is very interesting. Not only is there actual useful stuff. The way in which you get to the algorithm is itself interesting because it's this very different kind
[01:12:06.110] of math in which you're saying, OK, I'm going to take a scaling limit of my network. And then I assert some invariance about my network, and then I'm going to make additional assumptions. And then I'm going to try to figure out what the constraints on my hyperparameters are. And the constraints essentially give you certain kinds of scaling limits on your hyperparameters as well. So this is a general principle for coming up with different kinds of algorithms or hyperparameter scalings, I think, which are probably very different than what you're used to in CS or ML training. OK, those are the muP derivation. I want to close out on the muP part with just a few empirical things. There's an interesting paper by an independent researcher on basically stress testing muP in lots of different ways. And muP can be thought of as a procedure for tuning hyperparameters. And so depending on how you do things, Adam is going to have a certain learning rate as a function of different components of the network. So you have an embedding parameter, attention parameters, input/output MLPs or softmax linears. All of those are going to have different kinds of scalings under muP, as well as not muP scaling rules.
[01:13:18.540] And the interesting thing here is that particular paper replicated a bunch of the muP results for language models trained at different scales. And maybe unsurprisingly, he found almost the headline result from MiniCPM and others, which is that if you get the muP right, and you're only scaling with, in this case, in a very controlled way, you get exactly the learning rate optimality invariance that baseline muP-- or just looking at a variant of muP that tracks projection biases for attention, both of those transfer very nicely. But then you might ask, basically, there's a lot of things in the real world that muP does not handle. Technically, SwiGLU does not fit into muP, neither does initialize variations, initialization or RMS norm.
[01:14:12.600] Or if you really mess with your optimizer, that will mess that up, too. And so you might ask, like, which of these things that are not allowed by the theory really break muP? And actually, for the most part, these mostly work with muP. But some of the things that don't work is if you learn the gain term in RMS norm, this turns out to break muP. Although, these can be removed in many cases without necessarily hurting you, and so maybe it's not such a bad idea. But this does seem to break muP. It also seems to break if you use more exotic optimizers like Lion, which basically rely on the sign of the gradient. And sign gradient is like a spiritually similar to muon and other ideas. And so there might be other interesting things that break with muP as well. Finally, the thing that seems most concerning is that if you do large decoupled weight decay, you end up with significant muP failure.
[01:15:06.920] So this is maybe the one stress test that it does fail. But to close out this final section, you might, in the end, ask, like, is muP useful? And if you really want to stabilize your learning rates, it is actually pretty useful. There's a lot of experiments of this form, where if you increase the width, your optimal learning rate shifts very predictably, but very greatly as a function of scale. Whereas, with muP it doesn't. And so it's one of the many tools in the toolkit to try to control hyperparameter drift as a function of scale. I think it's an interesting open area of research. It's not really done and dusted, and this is the one right thing to do. But there's a lot of promise in the muP-style initialization program. There's also a lot of promise in just fitting scaling laws as well. So final thing I want to close with, scaling in the wild is very tricky. As I mentioned in one of my replies. I think the initial presentation of scaling laws makes it sound like a science.
[01:16:04.550] It's like, yes, draw this line, do these procedures, and you will know, for sure, what will happen at scale. But I think, in reality, it's a lot more messy and a lot more unknown than that. There are lots of things that people do use in practice scaling laws for, like picking architectures, picking optimizers, picking hyperparameters. All of these things are things that people do with scaling laws, but there is really an art to it in the sense of, you don't really if it's truly going to extrapolate forever. You want to do reasonable things to maximize the chances of success. And so there's lots of solutions that people have. You can use muP. You can search for the optimal learning rate in all of these other kinds of things. And these are ways of controlling hyperparameter drift. But there's no silver bullet yet. Maybe next year, there will be a module that's like, we solved it, but not quite yet.
[01:16:53.390] Great Thanks a lot. I can take questions afterwards.
