优化器、学习率和数据设置

这一节我们进入训练中更偏实践的部分,即训练的优化器和数据部分.这部分非常重要,直接决定了模型能否稳定训练,大模型的训练不是简单的“用个 Adam 就行了”,而是需要对优化器、学习率调度、数据设置等多个方面进行细致调整和理解.才可以保证训练的稳定性. 这一节实际上有很多相关的问题:

1. 优化器本身,各种优化器的原理和性质?
2. 优化过程中的参数和细节: 学习率、权重初始化、weight decay、warmup,等等,很多这些参数还是变化的.
3. 数据的重要性,从数据集如何构建,到训练过程中batch size、epoch,这些数据相关参数的设置.
4. 数据视角更进一步的 scaling law,数据量和模型规模的关系,数据质量相关的内容.

这些问题看起来分散，但它们都是决定训练过程是否成功的重要环节.

这一节按照如下顺序组织：

1. Warmup、学习率调度、权重初始化、weight decay 这些训练技巧分别在解决什么问题？
2. Adam 和 AdamW 优化器到底是什么？为什么 MiniMind 里用 AdamW？
3. Lion、Muon 这类新型优化器大概是什么方向？
4. Pretrain 数据集需要注意什么？数据质量、数据分布和 scaling law 为什么重要？
5. Batch size、epoch、tokens seen 这些概念应该怎么理解？

Q1: 基本的优化模型?

模型训练的最基本更新形式是：

\[ \theta_{t+1}=\theta_t-\eta \nabla_\theta \mathcal{L}(\theta_t) \]

其中，\(\theta_t\) 是第 \(t\) 步的模型参数，\(\eta\) 是学习率，\(\nabla_\theta \mathcal{L}(\theta_t)\) 是 loss 对参数的梯度。

这个公式看起来很简单，但实际训练中会马上遇到很多问题：

• 参数一开始应该怎么初始化？
• 训练刚开始时梯度不稳定怎么办？
• 学习率应该一直不变，还是随着训练变化？
• 参数会不会变得太大？
• 梯度更新要不要考虑历史方向和历史尺度？
• 一个 batch 里应该放多少数据？
• 数据集本身是否足够大、足够干净、分布是否合理？

正如开头所说的,这些因素大致可以被归结为

1. 优化训练方面: 如使用什么优化器? Warmup、学习率调度、weight decay、AdamW,数据视角的 batch size、epoch
2. 数据层面: 基本的数据集分布

对于后续的问题,我们也会按照训练的流程来组织,从最开始Warmup和权重初始化开始.

Q2: 权重初始化为什么会影响训练?

权重初始化不是“随便随机一下”。它会影响前向传播时激活值的尺度，也会影响反向传播时梯度的尺度。

如果初始化太大，激活值和梯度可能迅速变大，训练容易不稳定。
如果初始化太小，信号在层与层之间传播时可能变得很弱，模型学习速度会变慢。

经典初始化方法，比如 Xavier initialization 和 Kaiming initialization，本质上都是在尝试维持不同层之间的方差稳定。

先看最简单的一类：小方差正态初始化。它会让权重从一个均值为 0、标准差较小的正态分布中采样：

\[ W_{ij}\sim \mathcal{N}(0,\sigma^2) \]

其中，\(W_{ij}\) 表示权重矩阵中的一个元素，\(\sigma\) 是初始化标准差。很多 Transformer 实现会使用类似 \(\sigma=0.02\) 的小标准差初始化。它的直觉很简单：让模型一开始不要输出过大的激活值，避免训练刚开始就数值不稳定。

另一类是 Xavier initialization，也叫 Glorot initialization。它主要希望前向传播和反向传播中的方差都尽量保持稳定。设某一层输入维度为 \(d_{\text{in}}\)，输出维度为 \(d_{\text{out}}\)，则 Xavier uniform 常写成：

\[ W_{ij}\sim U\left(-\sqrt{\frac{6}{d_{\text{in}}+d_{\text{out}}}},\sqrt{\frac{6}{d_{\text{in}}+d_{\text{out}}}}\right) \]

对应的 Xavier normal 可以写成：

\[ W_{ij}\sim \mathcal{N}\left(0,\frac{2}{d_{\text{in}}+d_{\text{out}}}\right) \]

其中，\(U(a,b)\) 表示在区间 \([a,b]\) 上均匀采样。Xavier 初始化常用于 tanh、sigmoid 或较早期的全连接网络设置，它的核心目标是让信号在层与层之间传播时不要快速放大或缩小。

再看 Kaiming initialization，也叫 He initialization。它主要是为 ReLU 这类激活函数设计的。因为 ReLU 会把一部分负值截断为 0，所以需要稍微调整方差。Kaiming normal 常写成：

\[ W_{ij}\sim \mathcal{N}\left(0,\frac{2}{d_{\text{in}}}\right) \]

对应的 Kaiming uniform 可以写成：

\[ W_{ij}\sim U\left(-\sqrt{\frac{6}{d_{\text{in}}}},\sqrt{\frac{6}{d_{\text{in}}}}\right) \]

它的直觉是：在使用 ReLU 类激活函数时，仍然尽量保持每一层输出的方差稳定。

需要注意的是，上面这些初始化方法并不是“所有参数都这样初始化”。它们主要针对带权重矩阵的可学习层，最典型的是线性层 Linear，也包括卷积层 Conv 这类本质上也是线性变换的层。

在 LLM 里，情况会更复杂一些。Transformer 里大量使用残差连接、LayerNorm、Attention 和 MLP，很多初始化策略还会考虑模型深度。例如有些实现会对残差分支相关参数使用更小的初始化尺度，避免很多层残差叠加后激活值过大。

更具体地说：

• Xavier / Glorot 初始化主要适合 tanh、sigmoid 或比较对称的激活函数场景。它同时考虑 fan_in 和 fan_out，希望前向传播和反向传播的方差都比较稳定。
• Kaiming / He 初始化主要适合 ReLU、LeakyReLU、GELU 这类会改变激活分布的非线性函数。它最初是为 ReLU 设计的，更强调根据 fan_in 保持前向激活方差稳定。
• 小方差正态初始化在 Transformer / LLM 里很常见，常用于 embedding、attention projection、MLP projection 等参数。它不一定严格绑定某个激活函数，而是作为一种经验上稳定的初始化尺度。

但有些参数的初始化规则不同：

• bias 通常初始化为 0。
• LayerNorm 的 scale / weight 通常初始化为 1，bias 初始化为 0。
• embedding 矩阵通常也会用正态分布初始化，但它的角色和普通线性层不完全一样。
• RoPE 这类位置编码缓存不是可训练参数，一般不涉及权重初始化。

所以，初始化方法需要结合“参数属于哪一类层”和“后面接什么激活函数”一起理解。对于 Transformer 来说，还要额外考虑这个参数是否处在 attention、MLP、embedding、output head 或残差分支中。

这部分在 MiniMind 的训练脚本里不显眼，因为模型初始化通常封装在模型定义里。但从训练角度看，初始化决定了优化开始时的位置。一个好的初始化不保证模型训练成功，但一个不合适的初始化很容易让训练一开始就出问题。

Q2.1: 为什么不用全 0 初始化?

关于小方差初始化,一个很自然的问题是,如果方差已经很小了,为什么不直接初始化为全 0 呢?

两者还是有关键区别的,如果把所有权重都初始化为 0，会出现一个严重问题：对称性无法打破。

假设一层里有多个神经元，第 \(i\) 个神经元可以写成：

\[ h_i=f(W_i x+b_i) \]

其中，\(W_i\) 是第 \(i\) 个神经元对应的权重，\(b_i\) 是 bias，\(f\) 是激活函数。

如果所有 \(W_i\) 都是 0，所有 \(b_i\) 也一样，那么这些神经元一开始的输出完全一样。反向传播时，它们收到的梯度也会完全一样，于是更新后仍然一样。

也就是说，这一层里虽然有很多神经元，但它们行为上像是同一个神经元的复制品，模型容量被浪费了。

小方差随机初始化则不同：

\[ W_{ij}\sim \mathcal{N}(0,\sigma^2) \]

它的均值仍然是 0，但每个权重会有一点点随机差异。这样不同神经元一开始就不完全一样，反向传播时收到的梯度也会逐渐不同，模型才能学到不同特征。

所以小方差初始化是在折中：

• 随机：打破对称性。
• 均值接近 0：避免整体偏移。
• 方差较小：避免一开始数值爆炸。

bias 倒是经常可以初始化为 0，因为只要权重已经随机了，神经元之间就已经被区分开了。

Q3: Warmup 是什么? 为什么训练初期需要它?

Warmup 指的是：训练刚开始时，不直接使用目标学习率，而是先从一个很小的学习率逐步升高到目标学习率。

第一种最常见的方法是从 0 线性升高到目标学习率。假设目标学习率是 \(\eta_{\max}\)，warmup 总步数是 \(T_{\text{warmup}}\)，那么第 \(t\) 步的学习率可以写成：

\[ \eta_t=\eta_{\max}\cdot \frac{t}{T_{\text{warmup}}} \]

其中，\(0 \le t \le T_{\text{warmup}}\)。当 \(t=0\) 时，学习率是 0；当 \(t=T_{\text{warmup}}\) 时，学习率升到 \(\eta_{\max}\)。

类似的方法是从一个较小的起始学习率升高到目标学习率。有时我们不希望学习率真的从 0 开始，而是从 \(\eta_{\min}\) 开始逐步升高到 \(\eta_{\max}\)。这时可以写成：

\[ \eta_t=\eta_{\min}+(\eta_{\max}-\eta_{\min})\cdot\frac{t}{T_{\text{warmup}}} \]

其中，\(\eta_{\min}\) 表示 warmup 起始学习率，\(\eta_{\max}\) 表示 warmup 结束后达到的目标学习率。这个形式比从 0 开始更一般：当 \(\eta_{\min}=0\) 时，它就退化成第一种线性 warmup。

实际训练里，warmup 后面通常会接一个学习率衰减策略，比如 cosine decay。于是完整学习率调度会变成两段：前面先 warmup，后面再 decay。

为什么需要它？因为训练刚开始时，模型参数还很随机，激活值和梯度的尺度也可能不稳定。如果一开始就用较大的学习率，参数更新可能过猛，导致 loss 剧烈震荡，甚至直接出现 NaN。

Warmup 相当于给训练一个缓冲期：先小步走，等模型进入相对稳定的区域后，再使用正常学习率。

MiniMind 当前的 pretrain 脚本没有单独实现 warmup，但大模型训练里 warmup 非常常见。是否需要 warmup，通常和模型规模、batch size、初始化方式、优化器和学习率大小有关。

Q4: Learning Rate 调整? Cosine Decay 在做什么?

Learning rate 决定每一步参数更新走多远。我们之前已经提到了Warmup,它是训练初期的一个缓冲机制。Warmup 之后，通常还需要一个学习率衰减策略，让学习率随着训练进度逐渐降低。

如果学习率太大，参数更新可能越过较好的区域，导致 loss 震荡甚至发散。
如果学习率太小，训练虽然稳定，但下降太慢，在有限训练时间内学不到足够的东西。

因此，学习率通常不会固定不变，而是随着训练进度调整。MiniMind 使用的是一个简单的 cosine decay：

# src/minimind_learning/trainer/trainer_utils.py
def get_lr(current_step, total_steps, lr):
    return lr / 10 + 0.5 * lr * (1 + math.cos(math.pi * current_step / total_steps))

训练循环中，每一步都会重新设置 optimizer 的学习率：

# src/minimind_learning/trainer/train_pretrain.py
lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
for param_group in optimizer.param_groups:
    param_group['lr'] = lr

用公式写出来，MiniMind 这里的学习率大致是：

\[ \eta_t=\frac{\eta_0}{10}+\frac{1}{2}\eta_0\left(1+\cos\left(\pi\frac{t}{T}\right)\right) \]

其中，\(\eta_t\) 是第 \(t\) 步学习率，\(\eta_0\) 是初始学习率，\(T\) 是总训练步数。

这个公式的直觉是：训练前期学习率较高，帮助模型快速学习；训练后期学习率逐渐降低，让参数更新变得更细。

这里还有一个小细节：MiniMind 这个实现不是从 \(\eta_0\) 严格衰减到 0，而是从大约 \(1.1\eta_0\) 衰减到 \(0.1\eta_0\)。这不影响理解它作为 cosine schedule 的作用，但读代码时需要注意。

Q4.1: 还有哪些常见的学习率调整方式?

学习率调整方式很多，本质上都是在回答同一个问题：训练不同阶段应该用多大的步长更新参数。

1. Constant learning rate

最简单的方法是不调整学习率，整个训练过程都使用同一个值：

\[ \eta_t=\eta_0 \]

这种方式简单，但对于大模型 pretrain 往往不够灵活。训练初期可能需要 warmup，训练后期通常也希望学习率逐渐下降。

PyTorch 中如果不设置 scheduler，其实就是这种形式：

optimizer = torch.optim.AdamW(model.parameters(), lr=5e-4)

2. Step decay

Step decay 会每隔固定步数把学习率乘上一个系数：

\[ \eta_t=\eta_0\cdot \gamma^{\left\lfloor t/S \right\rfloor} \]

其中，\(S\) 是衰减间隔，\(\gamma\) 是衰减系数。

scheduler = torch.optim.lr_scheduler.StepLR(
    optimizer,
    step_size=1000,
    gamma=0.1,
)

它的特点是学习率呈阶梯状下降，简单直接，但下降点会比较突兀。

3. Exponential decay

Exponential decay 会让学习率按指数形式连续下降：

\[ \eta_t=\eta_0\cdot \gamma^t \]

其中，\(0<\gamma<1\)。

scheduler = torch.optim.lr_scheduler.ExponentialLR(
    optimizer,
    gamma=0.99,
)

这种方式比 step decay 更平滑，但如果 \(\gamma\) 设置不合适，学习率可能下降得太快或太慢。

4. Linear decay

Linear decay 会让学习率从初始值线性下降到某个最小值。假设总训练步数为 \(T\)，最终学习率为 \(\eta_{\min}\)，可以写成：

\[ \eta_t=\eta_{\min}+(\eta_0-\eta_{\min})\left(1-\frac{t}{T}\right) \]

def linear_decay_lr(step, total_steps, lr0, lr_min=0.0):
    ratio = min(step / total_steps, 1.0)
    return lr_min + (lr0 - lr_min) * (1.0 - ratio)

很多训练设置会使用 warmup + linear decay，也就是前面升高，后面线性下降。

5. Cosine decay

Cosine decay 使用余弦曲线让学习率平滑下降：

\[ \eta_t=\eta_{\min}+\frac{1}{2}(\eta_0-\eta_{\min})\left(1+\cos\left(\pi\frac{t}{T}\right)\right) \]

PyTorch 里可以直接使用：

scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
    optimizer,
    T_max=total_steps,
    eta_min=1e-5,
)

它的优点是下降平滑，训练后期会自然变成小步更新，所以在 Transformer 和 LLM 训练中很常见。

6. Warmup + decay

实际大模型训练里，常见的不是单独一种 decay，而是 warmup 和 decay 组合：

\[ \eta_t= \begin{cases} \eta_{\max}\cdot \frac{t}{T_{\text{warmup}}}, & t<T_{\text{warmup}} \\ \eta_{\text{decay}}(t), & t\ge T_{\text{warmup}} \end{cases} \]

其中，\(\eta_{\text{decay}}(t)\) 可以是 linear decay、cosine decay 或其他衰减函数。

一个简单的 warmup + cosine decay 代码可以写成：

import math

def warmup_cosine_lr(step, total_steps, warmup_steps, lr_max, lr_min=0.0):
    if step < warmup_steps:
        return lr_max * step / warmup_steps

    progress = (step - warmup_steps) / (total_steps - warmup_steps)
    progress = min(max(progress, 0.0), 1.0)
    return lr_min + 0.5 * (lr_max - lr_min) * (1 + math.cos(math.pi * progress))

这个形式很接近许多大模型训练里常用的学习率 schedule。

7. ReduceLROnPlateau

还有一类方法不是按固定 step 变化，而是根据验证集指标变化来调整学习率。如果 validation loss 长时间不下降，就降低学习率：

scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
    optimizer,
    mode="min",
    factor=0.5,
    patience=3,
)

# 每次验证后调用
scheduler.step(val_loss)

这种方式在传统深度学习任务中很常见。不过在大规模 pretrain 里，训练通常更依赖预先设计好的 schedule，因为验证成本高，训练步数也非常大。

Q4.2: 如何判断学习率是否合适?

这是一个很玄学的事情,并且这应该归结到Eval章节,但我还是决定在这里先提一下.后续再更多的展开.

最直接的观察对象仍然是 loss 曲线，但 loss 曲线只能给信号，不能单独证明某个学习率是最优的。

如果学习率太大，常见现象是：

• loss 一开始不下降，或者下降一点后剧烈震荡。
• loss 曲线有很多尖峰，甚至越来越高。
• 出现 NaN / Inf。
• 同样配置多跑几次，训练差异很大。
• 如果记录了 grad norm，通常会看到明显尖峰。

直觉上，这是因为每一步迈得太大，参数更新越过了比较好的区域，甚至破坏了已经形成的结构。

如果学习率太小，常见现象是：

• loss 很稳定，但下降非常慢。
• 曲线长时间几乎没有明显改善。
• 训练看起来很安全，但单位时间或单位 token 的学习效率很低。
• 提大学习率后，如果 loss 仍然稳定下降且下降更快，通常说明原来的学习率偏小。

比较合适的学习率通常表现为：warmup 后 loss 能比较快地下行，曲线允许有 batch-level 抖动，但整体趋势稳定下降；没有频繁的大尖峰，更不会出现 NaN。

一个常用的低成本方法是 LR range test。它不是严格控制变量实验，而是一个预实验，用来粗略测试当前训练系统能承受多大的学习率。

具体做法是：

1. 从一个很小的学习率开始，比如 \(10^{-7}\) 或 \(10^{-6}\)。
2. 每个 step 或每几个 step，把学习率按指数方式增大。
3. 记录每一步的 smoothed loss。
4. 一旦 loss 明显发散、剧烈震荡或出现 NaN，就停止。
5. 选择 loss 开始快速下降之后、明显发散之前的一段作为候选学习率区间。

指数增长可以写成：

\[ \eta_t=\eta_{\min}\left(\frac{\eta_{\max}}{\eta_{\min}}\right)^{t/T} \]

其中，\(\eta_{\min}\) 是起始学习率，\(\eta_{\max}\) 是测试上限，\(T\) 是测试总步数。

为什么说它不是严格实验？因为模型每一步都在更新，所以不同学习率对应的 loss 并不是在同一个模型状态下测出来的。严格来说，如果要控制变量，应该从同一个 checkpoint 出发，同时用多个学习率分别训练一小段再比较。但这样资源消耗会高很多。

所以实际做法通常是折中：

1. 先用 LR range test 快速排除明显不合适的学习率。
2. 再从同一个 checkpoint 开 2 到 4 个短跑实验，分别用几个候选学习率训练几百到几千 step。
3. 最后选择曲线更稳、下降更快、没有发散迹象的学习率。

因此，LR range test 更准确的理解是：在正式训练前，先用一小段训练过程扫描学习率，粗略摸清这个模型、数据、batch size、optimizer、混合精度配置下的稳定区间。但是这些参数变化后,学习率区间也会变化.

它主要回答的不是“最优学习率是多少”，而是：

• 学习率太小时，loss 是不是几乎不动？
• 学习率进入哪个区间后，loss 开始明显下降？
• 学习率大到什么程度后，loss 开始震荡、上升或 NaN？
• 正式训练的初始学习率应该避开哪些明显危险区域？

Q5: Adam 和 AdamW 优化器是什么? 它们是如何被构造出来的?

Q5.1: GD 族优化器

在介绍 Adam 之前，可以先从最基础的 SGD 看起。

SGD 全称是 Stochastic Gradient Descent，也就是随机梯度下降。它的基本形式是：

\[ \theta_{t+1}=\theta_t-\eta g_t \]

其中，\(\theta_t\) 是第 \(t\) 步的参数，\(\eta\) 是学习率，\(g_t=\nabla_\theta \mathcal{L}_t(\theta_t)\) 是当前 mini-batch 上估计出来的梯度。

需要注意, 深度学习语境下的 SGD 实际上通常指的是 mini-batch SGD. 其本质都是使用不同的样本数来估计梯度的期望.在优化算法里,本质上都是 Gradient Descent（梯度下降）的一族变种:

SGD 的特点是简单直接：当前 batch 给出一个梯度，就沿着负梯度方向走一步。但它也有两个问题：

• 当前 batch 的梯度可能噪声很大，更新方向会抖动。
• 所有参数使用同一个学习率，不会根据不同参数的梯度尺度自适应调整。

为了缓解第一个问题，可以加入 momentum。Momentum SGD 会维护一个速度项 \(v_t\)：

\[ v_t=\mu v_{t-1}+g_t \]

\[ \theta_{t+1}=\theta_t-\eta v_t \]

其中，\(\mu\) 是 momentum 系数。直觉上，momentum 会累积过去梯度的方向，让更新方向更平滑，减少 batch 噪声带来的来回震荡。

Q5.2: Adam

Adam 可以看成是在 Momentum SGD 的方向上继续发展：它不只关心梯度方向的滑动平均，还关心梯度平方的滑动平均。也就是说，SGD 主要看当前梯度，Momentum SGD 看历史梯度方向，而 Adam 同时维护梯度的一阶矩和二阶矩。

Adam 的全称是 Adaptive Moment Estimation。它对每个参数都维护两类状态：

• 一阶矩 \(m_t\)：梯度的指数滑动平均，可以理解为“带 momentum 的梯度方向”。
• 二阶矩 \(v_t\)：梯度平方的指数滑动平均，可以理解为“梯度尺度”的估计。

初始化时通常令：

\[ m_0=0,\quad v_0=0 \]

假设第 \(t\) 步梯度是：

Adam 会先更新一阶矩和二阶矩：

\[ m_t=\beta_1m_{t-1}+(1-\beta_1)g_t \]

\[ v_t=\beta_2v_{t-1}+(1-\beta_2)g_t^2 \]

其中，\(\beta_1\) 和 \(\beta_2\) 控制历史信息保留多少。常见默认值是 \(\beta_1=0.9\)，\(\beta_2=0.999\)。(看起来是个线性插值/平滑滤波)

因为 \(m_0\) 和 \(v_0\) 都初始化为 0，所以训练初期的 \(m_t\)、\(v_t\) 会偏向 0。Adam 会使用 bias correction 修正这个偏差：

\[ \hat{m}_t=\frac{m_t}{1-\beta_1^t} \]

\[ \hat{v}_t=\frac{v_t}{1-\beta_2^t} \]

然后用修正后的一阶矩和二阶矩更新参数：

\[ \theta_t=\theta_{t-1}-\eta\frac{\hat{m}_t}{\sqrt{\hat{v}_t}+\epsilon} \]

其中，\(\eta\) 是学习率，\(\epsilon\) 是一个很小的常数，通常类似 \(10^{-8}\)，用于避免除以 0。

这个更新式的直觉是：

• \(\hat{m}_t\) 决定主要更新方向。
• \(\sqrt{\hat{v}_t}\) 根据历史梯度尺度对更新做归一化。
• 梯度长期较大的参数，更新会被缩小。
• 梯度长期较小的参数，更新相对不会被压得太厉害。

所以 Adam 会根据每个参数自己的梯度历史，自适应地调整更新尺度。梯度波动大的参数，更新会更谨慎；梯度相对稳定的参数，更新会更顺滑。

Adam 也有一些需要注意的点：

• Adam 的状态开销比 SGD 大，因为每个参数都要额外保存 \(m_t\) 和 \(v_t\)。
• Adam 对学习率仍然敏感，不是用了 Adam 就不需要调学习率。
• Adam 里的 \(\epsilon\)、\(\beta_1\)、\(\beta_2\) 通常使用默认值，但在极大 batch、混合精度或特殊任务里也可能需要调整。
• 如果直接把 L2 regularization 加进 Adam 的梯度里，它会和自适应缩放混在一起，这也是 AdamW 要解决的问题。

Q5.2.1: \(v_t\) 是严格意义上的二阶矩吗?

这个问题很关键,它揭示了Adam本质上还是一阶优化算法,而不是严格意义上的二阶优化算法(没有使用二阶Hessian矩阵).

严格说，Adam 里的 \(v_t\) 不是完整意义上的二阶矩矩阵，也不是协方差矩阵。它更准确地说是逐元素梯度平方的指数滑动平均，也就是对每个参数维度的 second raw moment estimate。

如果参数 \(\theta\) 的形状是：

\[ \theta \in \mathbb{R}^{d} \]

那么梯度 \(g_t\)、一阶矩 \(m_t\)、二阶矩估计 \(v_t\) 的 shape 都和参数一样：

\[ g_t,m_t,v_t \in \mathbb{R}^{d} \]

Adam 中的二阶矩更新是逐元素的：

\[ v_t=\beta_2v_{t-1}+(1-\beta_2)(g_t\odot g_t) \]

其中，\(\odot\) 表示逐元素乘法。

所以它不是下面这种完整矩阵：

\[ g_tg_t^\top \]

完整二阶矩矩阵或协方差矩阵会包含不同参数维度之间的相关性，计算和存储代价都非常高。Adam 只保留每个参数位置自己的梯度平方历史，所以它更像是一种对角近似：估计每个参数维度自己的梯度尺度，但不建模不同参数维度之间的相关性。

Q5.3: AdamW

AdamW 可以理解为 Adam 加上 decoupled weight decay，也就是把 weight decay 从 Adam 的梯度自适应更新中解耦出来。

先看普通 Adam 如果加入 L2 regularization，常见做法是把正则项加进 loss：

这样梯度会变成：

如果把 \(g’t\) 送进 Adam，那么 \(\lambda\theta{t-1}\) 也会进入一阶矩、二阶矩，并被 Adam 的自适应缩放处理。这样 weight decay 的效果会和参数自己的梯度尺度纠缠在一起。

AdamW 的做法是：Adam 的自适应梯度更新仍然只根据 \(g_t\) 计算，然后在参数更新时单独加上 weight decay。

AdamW 的状态初始化和 Adam 一样：

\[ m_0=0,\quad v_0=0 \]

先计算梯度：

再按 Adam 的方式更新一阶矩、二阶矩和 bias correction：

\[ m_t=\beta_1m_{t-1}+(1-\beta_1)g_t \]

\[ v_t=\beta_2v_{t-1}+(1-\beta_2)g_t^2 \]

\[ \hat{m}_t=\frac{m_t}{1-\beta_1^t},\quad \hat{v}_t=\frac{v_t}{1-\beta_2^t} \]

最后更新参数时，把 Adam 更新和 weight decay 分开：

也可以写成：

这里 \(\lambda\) 是 weight decay 系数。

这就是 AdamW 里 decoupled weight decay 的含义：参数衰减不再作为梯度的一部分进入 Adam 的自适应矩估计，而是在最后更新参数时单独作用。

MiniMind 的 pretrain 代码里使用的是 AdamW：

# src/minimind_learning/trainer/train_pretrain.py
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)

在 Transformer / LLM 训练里，AdamW 是非常常见的默认选择。不过使用 AdamW 时也要注意：

• learning rate 仍然需要调，AdamW 不是免调参优化器。
• weight decay 通常不会施加到所有参数上。很多训练设置会排除 bias、LayerNorm weight、embedding 等参数。
• weight decay 和 learning rate 会共同决定参数衰减强度，因为衰减项里有 \(\eta\lambda\)。
• AdamW 的状态开销仍然较大，因为它和 Adam 一样要保存 \(m_t\) 和 \(v_t\)。

Q6: Weight Decay 是什么?

有一个很容易和 Learning Rate Decay 混淆的概念叫 Weight Decay。它们虽然名字里都有 decay，但其实是两个完全不同的东西。Weight decay 来自 L2 正则项:

\[ \mathcal{L}’(\theta)=\mathcal{L}(\theta)+\frac{\lambda}{2}|\theta|^2 \]

其中，\(\lambda\) 是正则强度，也就是 weight decay 系数。

对这个新的 loss 求梯度，可以得到：

\[ \nabla_\theta \mathcal{L}’(\theta)=\nabla_\theta \mathcal{L}(\theta)+\lambda\theta \]

如果记 \(g_t=\nabla_\theta \mathcal{L}(\theta_t)\)，那么在 SGD 中更新就是：

\[ \theta_{t+1}=\theta_t-\eta(g_t+\lambda\theta_t) \]

展开后：

\[ \theta_{t+1}=\theta_t-\eta g_t-\eta\lambda\theta_t \]

也可以写成：

\[ \theta_{t+1}=(1-\eta\lambda)\theta_t-\eta g_t \]

这里的 \(-\eta\lambda\theta_t\) 就是 weight decay 项。它的作用是每一步都把参数按比例往 0 的方向拉一点，避免参数规模无限变大。

所以在普通 SGD 中，可以说 L2 正则和 weight decay 是等价的：L2 正则项的梯度自然产生了参数衰减项。

但是在 Adam 中，情况会变得更微妙。因为 Adam 会把梯度送进一阶矩 \(m_t\) 和二阶矩估计 \(v_t\)，并做自适应缩放。如果直接把 L2 正则项加到 loss 里，那么 \(\lambda\theta_t\) 也会进入 Adam 的自适应缩放过程。这样一来，参数衰减就不再是简单的“按比例往 0 拉”，而是会受到每个参数历史梯度尺度的影响。

AdamW 的意义就是把这件事解耦：Adam 仍然只根据原始 loss 的梯度 \(g_t\) 计算自适应更新，而 weight decay 在最后参数更新时单独作用：

也就是把“根据梯度更新参数”和“按比例衰减参数”分开处理。

关于 \(\lambda\) 是否变化，通常来说 weight decay 系数本身是固定超参数。常见取值可能是 0.01、0.1 这一类，具体要看模型规模、数据量、训练步数和是否对某些参数排除 weight decay。

不过要注意，实际每一步的衰减强度里还有学习率 \(\eta_t\)：

\[ -\eta_t\lambda\theta_t \]

所以即使 \(\lambda\) 固定，只要 learning rate schedule 在变化，实际每一步的参数衰减幅度也会跟着变化。

在 Transformer / LLM 训练里，还常常不会对所有参数都使用 weight decay。比如 bias、LayerNorm 的 weight、某些 embedding 参数，可能会被排除在 weight decay 之外。这样做的原因是这些参数本身的尺度意义和普通线性层权重不完全一样，直接衰减未必总是有益。

Q7: Lion 和 Muon 这类新型优化器是什么?

AdamW 很常见，但它不是优化器发展的终点。近年来也出现了一些针对深度网络和大模型训练的新优化器，比如 Lion 和 Muon。

7.1 Lion

Lion 来自论文 Symbolic Discovery of Optimization Algorithms。它的名字是 EvoLved Sign Momentum。

和 Adam 相比，Lion 更轻量。Adam 通常需要保存一阶矩和二阶矩两个状态，而 Lion 主要保存 momentum，因此内存开销更小。Lion 的更新还使用了 sign operation，也就是参数更新方向更像是由符号决定，而不是直接使用连续梯度值。

Lion 的状态初始化通常是：

\[ m_0=0 \]

其中，\(m_t\) 是 momentum 状态，shape 和参数 \(\theta_t\) 一样。

假设第 \(t\) 步梯度是：

Lion 会先用 \(\beta_1\) 构造当前更新方向：

\[ c_t=\beta_1m_{t-1}+(1-\beta_1)g_t \]

然后用 sign update 更新参数。如果带 decoupled weight decay，可以写成：

\[ \theta_t=\theta_{t-1}-\eta\left(\operatorname{sign}(c_t)+\lambda\theta_{t-1}\right) \]

最后再用 \(\beta_2\) 更新 momentum 状态：

\[ m_t=\beta_2m_{t-1}+(1-\beta_2)g_t \]

其中，\(\eta\) 是学习率，\(\lambda\) 是 weight decay 系数。Lion 常见默认 \(\beta\) 设置类似 \((0.9, 0.99)\)。

Lion 和 Adam 的一个重要区别是：Adam 用 \(\sqrt{\hat{v}_t}\) 做逐元素尺度归一化，而 Lion 不维护二阶矩估计。它只保留 momentum，并用 \(\operatorname{sign}(c_t)\) 决定每个参数位置更新的方向。

一个简化理解是：

• AdamW：用一阶矩和二阶矩做自适应更新。
• Lion：用 momentum 加 sign update，状态更少，更新形式更简单。

不过 Lion 不是“无脑替换 AdamW”。因为 sign update 的更新范数通常更大，Lion 往往需要比 AdamW 更小的学习率；为了保持 \(\eta\lambda\) 的衰减强度，weight decay 有时也会相应调大。它在不同任务上的收益并不总是稳定，所以更适合作为了解现代优化器方向的例子。

7.2 Muon

Muon 的主实现来自 KellerJordan/Muon，全称可以理解为 MomentUm Orthogonalized by Newton-Schulz。

Muon 的核心思想是：对神经网络隐藏层中的 2D 权重矩阵，先用 momentum 得到更新方向，再对这个更新矩阵做近似正交化。实现上常用 Newton-Schulz iteration 来高效近似这个操作。

为了避免符号混乱，先把这一节里的符号列出来：

• \(W_t\)：第 \(t\) 步的 2D 权重矩阵，也就是 Muon 要优化的参数。
• \(G_t\)：第 \(t\) 步的梯度矩阵，shape 和 \(W_t\) 相同。
• \(M_t\)：momentum 状态，用来累积历史梯度方向。
• \(U_t\)：正交化之前的候选更新矩阵，也就是 update matrix。
• \(O_t\)：对 \(U_t\) 做正交化之后得到的更新矩阵。
• \(\mu\)：momentum 系数。
• \(\eta\)：学习率。
• \(\lambda\)：weight decay 系数。

也就是说，Muon 并不是直接用梯度 \(G_t\) 或 momentum \(M_t\) 更新参数，而是先得到候选更新 \(U_t\)，再把它正交化为 \(O_t\)，最后用 \(O_t\) 更新权重。

对于一个 2D 参数矩阵 \(W_t\)，Muon 的 momentum 状态可以初始化为：

\[ M_0=0 \]

假设当前梯度是：

Muon 先做类似 SGD momentum 的更新：

\[ M_t=\mu M_{t-1}+G_t \]

有些实现还会使用 Nesterov 形式：

\[ U_t=\mu M_t+G_t \]

如果不使用 Nesterov，也可以简单理解为：

\[ U_t=M_t \]

接下来，Muon 会对更新矩阵 \(U_t\) 做正交化：

\[ O_t=\operatorname{Ortho}(U_t) \]

理想情况下，如果 \(U_t=A\Sigma B^\top\) 是奇异值分解，那么正交化结果可以理解为：

\[ \operatorname{Ortho}(U_t)=AB^\top \]

也就是说，它保留更新矩阵的“方向结构”，但把奇异值压到接近 1。实际实现中不会真的每一步做完整 SVD，而是使用 Newton-Schulz iteration 近似这个正交化操作。

最后更新参数：

\[ W_t=W_{t-1}-\eta O_t \]

如果带 AdamW-style weight decay，也可以写成：

\[ W_t=W_{t-1}-\eta\left(O_t+\lambda W_{t-1}\right) \]

它的使用方式也和 AdamW 不完全一样。Muon 通常只用于隐藏层里的矩阵参数；embedding、输出层、bias、gain 等参数仍然建议使用 AdamW 之类的标准优化器。

所以 Muon 更像是一种“针对矩阵参数结构的优化器”，而不是简单的通用 AdamW 替代品。它利用了 2D 权重矩阵的结构，而 AdamW / Lion 这类优化器更多是逐元素更新。

对于这类新优化器，我现在更倾向于先把它们当作了解方向：它们说明优化器设计正在从“通用自适应更新”继续往“利用参数结构、减少状态开销、提高训练效率”的方向发展。但在 MiniMind 这种学习项目里，AdamW 仍然是最稳定、最容易理解的选择。

Q8: Pretrain 数据集需要注意什么?

Pretrain 的本质是拟合数据分布。模型学到什么，很大程度上取决于训练数据长什么样。这是一个非常复杂的问题,我们会在后续的章节展开说.但在这里先给出一些直观的要点,和一些实现细节的注意点：

数据集至少要关注几件事：

• 数据质量：低质量文本、乱码、重复模板会直接被模型学进去。
• 数据分布：代码、百科、小说、问答、网页文本的比例不同，模型能力也会偏向不同方向。
• 重复数据：大量重复样本会让模型浪费训练步数，也可能增加记忆风险。
• 数据污染：如果 eval 数据混进训练集，会让评估结果失真。
• 领域覆盖：如果训练数据过窄，模型在其他领域的泛化会比较弱。

MiniMind 的 pretrain dataset 读取 JSONL 文件中的 text 字段：

# src/minimind_learning/dataset/lm_dataset.py
def load_data(self, path):
    samples = []
    with open(path, 'r', encoding='utf-8') as f:
        for line_num, line in enumerate(f, 1):
            data = json.loads(line.strip())
            samples.append(data)
    return samples

随后每条样本会被 tokenizer 编码、截断和 padding：

encoding = self.tokenizer(
    str(sample['text']),
    max_length=self.max_length,
    padding='max_length',
    truncation=True,
    return_tensors='pt'
)

这意味着 max_seq_len、padding 和 truncation 都会影响模型实际看到的数据。

max_seq_len 不只是显存参数。它决定了单条样本最多保留多少 token，也影响模型能从训练中学习多长范围内的上下文依赖。较短序列训练更快、更省显存，但模型很难从中学习长文本结构；较长序列信息更完整，但训练成本更高。

需要根据模型的能力以及数据集的状态进行权衡,比如数据集缺乏长句子时，过长的 max_seq_len 可能会浪费显存；数据集里有很多长文本时，过短的 max_seq_len 又会丢失重要信息。所以在数据集构造的时候,就需要结合模型能力做考量.

Q8.1: 模型在训练和推理的时候能接受多长的句子?

模型能接受多长的句子，通常受几类因素共同限制：

1. 模型结构里的最大位置长度。
2. 训练时见过的序列长度。
3. attention 计算和显存成本。
4. 位置编码或位置嵌入的外推能力。

如果使用传统的 learned absolute position embedding，输入 token embedding 和位置 embedding 通常会相加：

\[ x_t=e_t+p_t \]

其中，\(e_t\) 是第 \(t\) 个 token 的 embedding，\(p_t\) 是第 \(t\) 个位置的位置 embedding。

这时位置 embedding 通常是一张固定大小的表：

\[ P\in\mathbb{R}^{L_{\max}\times d} \]

其中，\(L_{\max}\) 是最大位置长度，\(d\) 是 hidden size。如果训练时设置 \(L_{\max}=1024\)，那么这张表只有 1024 个位置。推理时输入超过 1024，就没有对应的 \(p_t\) 可以查。因此这种模型的上下文长度会被 position embedding 表硬限制住。

即使强行扩展这张表，比如随机初始化新的位置 embedding，效果通常也不可靠。因为模型训练时没有学过这些新位置。

后来很多 LLM 使用 RoPE、ALiBi 或其他相对位置编码方式。以 RoPE 为例，它不是查一个固定长度的位置表，而是根据位置 \(t\) 对 query 和 key 做旋转：

\[ \operatorname{RoPE}(x_t,t) \]

理论上，\(t\) 可以继续变大，所以 RoPE 不像 learned absolute embedding 那样被固定表硬卡住。但这不代表模型可以无限使用长上下文。因为模型训练时可能只见过 1024、2048 或 4096 长度的序列，它学到的 attention 行为和长距离依赖方式，主要仍然来自训练长度范围。

所以 RoPE 解决的是“位置编码能不能算到更远”的问题，但没有完全解决“模型会不会用这么远的信息”的问题。

如果一个模型只用 \(1024\) token 的长度训练，那么推理时能到多长要分情况：

• 如果是 learned absolute position embedding，硬限制通常就是 1024。
• 如果是 RoPE / ALiBi，技术上可能能喂超过 1024，但这属于外推，效果不保证。
• 即使推理框架和显存允许更长输入，模型也未必真的能利用远处信息。

所以训练长度、位置编码和推理长度之间不是简单的等号关系。位置编码决定“能不能表示位置”，训练长度决定“模型有没有学过这种长度范围内的行为”，推理资源决定“实际能不能跑得动”。

Q8.2: 支持长上下文的模型一般怎么训练?

支持长上下文的模型，通常不能只靠推理时把长度硬拉长。更常见的做法是让模型在训练中见过足够长的上下文，或者至少经过长上下文继续训练。

常见做法有几类。

第一类是从预训练阶段就使用长序列。比如训练时直接使用 4k、8k、32k token 的 sequence length。这样模型从一开始就能见到长文档结构、长距离依赖和长 attention pattern。

但这样非常贵，因为标准 attention 的计算和显存复杂度大致是：

\[ O(L^2) \]

其中，\(L\) 是序列长度。序列越长，attention 矩阵越大。

第二类是先短上下文预训练，再做长上下文继续训练。这是很常见的路线。先用较短长度训练基础语言能力，比如 2k 或 4k；之后再用更长 sequence length 继续训练，比如 8k、16k 或 32k。

这样做更省，因为很多基础语言能力不一定需要特别长的上下文。长上下文阶段主要补的是：

• 更长距离的位置泛化。
• 长文档结构。
• 跨段落依赖。
• 长上下文检索能力。
• 注意力在长序列里的稳定性。

第三类是使用 RoPE scaling / position interpolation 后继续训练。对于 RoPE 模型，可以通过缩放位置编码，把原来的位置分布映射到更长范围。例如从 4k 扩到 32k。但通常仍然需要继续训练或微调，让模型适应新的位置编码尺度。

换句话说，RoPE scaling 让模型“能算”，长上下文训练让模型“会用”。

长上下文训练还常常配合一些工程和数据策略：

• FlashAttention：降低 attention 显存占用并加速。
• gradient checkpointing：减少训练时保存的激活。
• sequence packing：提高 token 利用率。
• document-level 数据：让训练样本保留真实长文档结构。
• 长文档 QA / retrieval-style 数据：让模型学习使用远处信息。

训练时的显存压力通常比推理更大。因为训练不仅要保存参数，还要保存：

• forward 激活，用于 backward。
• 梯度。
• optimizer state，比如 AdamW 的 \(m_t\)、\(v_t\)。
• 混合精度训练中的额外状态。

而推理不需要 backward，主要额外保存的是 KV cache。所以同样 32k context，推理可能还能跑，训练会贵很多。

因此，长上下文能力可以拆成三层：

1. 位置编码是否支持更长位置。
2. 模型是否通过训练学会使用长距离信息。
3. 训练和推理资源是否支持这么长的序列。

一句话总结：位置编码扩展让模型“能处理更长位置”，长上下文训练让模型“真的会用长上下文”。

Q8.3: 下一个 token 的预测是否受限于最后一个 hidden state?

这里还有一个容易被忽略的问题：对于 causal language model 来说，模型预测下一个 token 时，通常确实只使用最后一层、最后一个位置的 hidden state。

假设输入序列是：

\[ x_1,x_2,\dots,x_t \]

经过 \(L\) 层 Transformer 之后，第 \(t\) 个位置会得到最后一层的 hidden state：

\[ h_t^{(L)}\in \mathbb{R}^{d_{\text{model}}} \]

其中，\(h_t^{(L)}\) 表示最后一层第 \(t\) 个位置的表示，\(d_{\text{model}}\) 是模型的 hidden size。预测下一个 token 时，模型通常会把它送入 LM Head：

然后再经过 softmax 得到下一个 token 的概率分布：

\[ p(x_{t+1}\mid x_{\le t})=\operatorname{softmax}(\text{logits}_{t+1}) \]

从这个角度看，这种直觉是对的：虽然前面的 token 可以在多层 attention 里不断组合、交换和变换信息，但最终用于预测 \(x_{t+1}\) 的，确实是一个固定维度的向量 \(h_t^{(L)}\)。也就是说，整个历史上下文中与当前预测有关的信息，最后都需要被压缩到这个 \(d_{\text{model}}\) 维表示里。

这可以看成一种隐含的信息瓶颈。上下文长度可以变长，attention 可以看到更多 token，但最后参与下一个 token 预测的表示维度并不会随着上下文长度线性增长。如果序列非常长，模型就必须学会选择性地保留、压缩和聚合信息，而不是把所有上下文无损地塞进最后一个向量。

不过这里也需要稍微更精确一点：\(h_t^{(L)}\) 并不是一个静态的“全文摘要”。它是每一层 attention 根据当前位置的预测需求逐层构造出来的表示。模型不需要完整记住前面所有 token 的每个细节，而是需要保留对预测下一个 token 有用的信息。

比如在预测一句话的下一个词时，模型可能只需要最近几个 token 的局部语法信息；但在长文档问答、代码补全、跨段落引用这类任务里，模型可能需要从很远的位置取回关键信息。这个时候，难点就不只是“上下文窗口够不够长”，还包括：

• 最后一个位置的 hidden state 能不能有效聚合远距离信息。
• attention 是否真的学会在长序列中找到关键 token。
• 固定的 \(d_{\text{model}}\) 维表示是否足够承载当前预测需要的信息。
• 训练数据里是否有足够多需要长距离依赖的样本。

所以长上下文能力并不只是 position embedding 或显存问题。即使模型结构上能接收很长的输入，它仍然要通过 attention 和 hidden state，把“长上下文里真正有用的部分”压缩成当前预测所需的表示。这个压缩过程做得好不好，也是长上下文模型能力的重要限制之一。

Q9: Scaling Law 是什么?它说明了什么?

Scaling law 讨论的是模型性能如何随着模型参数量、训练数据量和计算量变化。它不是一个严格的理论定理，而是从大量训练实验里拟合出来的经验规律。

先定义几个符号：

• \(L\)：模型的 cross entropy loss，也可以理解为平均 negative log-likelihood。
• \(N\)：模型参数量。Kaplan 论文里通常指不包含 embedding 的参数量。
• \(D\)：训练数据量，也就是训练 token 数。
• \(C\)：训练 compute，通常用 FLOPs 衡量。

对于 dense Transformer language model，一个常用的训练 compute 近似是：

\[ C\approx 6ND \]

其中，\(N\) 是参数量，\(D\) 是训练 token 数。这个公式的直觉是：每个 token 的 forward 和 backward 大致都要对模型参数做若干次计算，所以总 compute 会同时随参数量和 token 数增长。

OpenAI 的 Scaling Laws for Neural Language Models 观察到，当其他因素不成为瓶颈时，loss 会随着模型参数量、数据量、训练 compute 呈现幂律下降。

当主要瓶颈是模型参数量 \(N\) 时，论文写成：

\[ L(N)=\left(\frac{N_c}{N}\right)^{\alpha_N} \]

其中，\(N_c\) 是拟合得到的常数，\(\alpha_N\) 是参数量方向的 scaling exponent。论文给出的拟合结果大约是：

\[ \alpha_N\approx 0.076 \]

当主要瓶颈是训练数据量 \(D\) 时：

\[ L(D)=\left(\frac{D_c}{D}\right)^{\alpha_D} \]

其中，\(D_c\) 是拟合得到的常数，\(\alpha_D\) 是数据量方向的 scaling exponent。论文给出的拟合结果大约是：

\[ \alpha_D\approx 0.095 \]

当主要瓶颈是 compute，并且模型大小和训练过程都接近 compute-optimal 时：

\[ L(C_{\min})=\left(\frac{C_c^{\min}}{C_{\min}}\right)^{\alpha_C^{\min}} \]

其中，\(C_{\min}\) 表示为了达到某个 loss 所需的最小 compute，\(C_c^{\min}\) 是拟合常数，论文给出的指数大约是：

\[ \alpha_C^{\min}\approx 0.050 \]

Kaplan 论文还给了一个同时考虑模型大小和数据量的形式：

\[ L(N,D)= \left[ \left(\frac{N_c}{N}\right)^{\frac{\alpha_N}{\alpha_D}} + \frac{D_c}{D} \right]^{\alpha_D} \]

这个公式表达的是：模型太小会带来损失，数据太少也会带来损失；最终 loss 可以看成这两类限制共同作用的结果。

根据这套拟合关系，Kaplan 论文得到的 compute-optimal 趋势大致是：

\[ N_{\text{opt}}\propto C^{0.73} \]

\[ D_{\text{opt}}\propto C^{0.27} \]

也就是说，如果训练 compute 增加，Kaplan 版本的建议更偏向于优先增大模型参数量，而训练 token 数增长得相对慢一些。

后来的 Chinchilla 工作 Training Compute-Optimal Large Language Models 重新研究了这个问题。它的核心问题是：在固定 compute 预算 \(C\) 下，应该把预算更多分给模型参数量 \(N\)，还是训练 token 数 \(D\)？

Chinchilla 论文使用的参数化 loss 形式是：

\[ L(N,D)=E+\frac{A}{N^\alpha}+\frac{B}{D^\beta} \]

其中：

• \(E\)：理想生成过程在数据分布上的不可约 loss，可以理解为数据本身的熵下界。
• \(\frac{A}{N^\alpha}\)：模型参数量不够带来的额外 loss。
• \(\frac{B}{D^\beta}\)：训练 token 数不够带来的额外 loss。

论文拟合得到：

\[ E=1.69,\qquad A=406.4,\qquad B=410.7 \]

\[ \alpha=0.34,\qquad \beta=0.28 \]

然后在 compute 约束下做优化：

也就是在 \(6ND=C\) 这个预算约束下，寻找能让 loss 最小的 \(N\) 和 \(D\)。

根据 Chinchilla 的参数化形式，可以得到 compute-optimal frontier：

\[ N_{\text{opt}}(C)=G\left(\frac{C}{6}\right)^a \]

\[ D_{\text{opt}}(C)=G^{-1}\left(\frac{C}{6}\right)^b \]

其中：

\[ G=\left(\frac{\alpha A}{\beta B}\right)^{\frac{1}{\alpha+\beta}} \]

\[ a=\frac{\beta}{\alpha+\beta},\qquad b=\frac{\alpha}{\alpha+\beta} \]

把 \(\alpha=0.34\)、\(\beta=0.28\) 代进去：

\[ a=\frac{0.28}{0.34+0.28}\approx 0.45 \]

\[ b=\frac{0.34}{0.34+0.28}\approx 0.55 \]

这说明 Chinchilla 的结论和 Kaplan 不太一样。它认为当 compute 增加时，模型参数量和训练 token 数应该接近同比例增长，而不是主要把预算用来增大模型。

论文里用三种方法估计这个关系，最后得到的趋势大致可以概括为：

\[ N_{\text{opt}}\propto C^{0.5},\qquad D_{\text{opt}}\propto C^{0.5} \]

这就是常说的 Chinchilla scaling：模型变大时，训练 token 数也应该同步变多。

这个结论的一个重要后果是：很多早期大模型其实是 undertrained 的。它们参数量很大，但训练 token 数不够多，所以并没有在给定 compute 下达到最优。Chinchilla 本身就是一个典型例子：它比 Gopher 小很多，但训练 token 更多，在相近 compute 下效果更好。

对 MiniMind 这样的学习项目来说，我们不需要直接套这些工业级公式。因为这些公式是在特定数据集、特定模型结构和大规模训练设置下拟合出来的，小模型、小数据集不一定满足相同规律。但它们给了一个非常重要的直觉：

• 参数量、数据量、计算量要一起看。
• 只增大模型、不增加训练 token，模型可能 undertrained。
• 只增加数据、不增加模型容量，模型也可能吃不下更复杂的分布。
• pretrain 里不能只看 epoch，更应该关注 tokens seen。
• 数据质量和数据规模本身就是 pretrain 的核心变量。

这也是为什么 MiniMind 的意义不是直接复制工业级训练结论，而是帮助我们在小规模实验中看清这些变量之间的关系。

Q10: Batch Size、Epoch 和 Tokens Seen 应该怎么理解?

这一节其实是在回答一个训练计量问题：一次参数更新到底看了多少数据？整个训练过程一共看了多少 token？这些量又和显存、epoch、compute 有什么关系？

Q10.1: Batch size、梯度累积和分布式训练是什么关系?

在代码里看到的 batch_size，通常不是整个训练系统的总 batch size，而是每张卡、每次 forward/backward 实际处理的样本数。这个值也常被叫作 micro-batch size。

MiniMind 默认设置里：

parser.add_argument("--batch_size", type=int, default=32, help="batch size")
parser.add_argument("--accumulation_steps", type=int, default=8, help="梯度累积步数")

如果显存不够，最直接的办法就是把每张卡上的 batch_size 设小一点。这样每次 forward/backward 处理的样本少，激活占用也会少一些，显存压力会下降。

但是 batch size 太小，单次梯度估计会比较嘈杂。于是训练里经常使用 gradient accumulation。它的意思是：连续做多次 forward/backward，把梯度累积在参数上，暂时不执行 optimizer.step()；等累积了指定次数之后，再统一做一次参数更新。

可以把它理解成：

1. 第 1 个 micro-batch 计算梯度，先不更新参数。
2. 第 2 个 micro-batch 继续计算梯度，把梯度加到前面的梯度上。
3. 重复 \(A\) 次。
4. 最后做一次 optimizer.step()，完成一次真正的参数更新。

这里的 \(A\) 就是 accumulation_steps。

分布式训练里还会多一个维度：多张 GPU 同时处理不同的数据。每张卡先在自己的 micro-batch 上计算梯度，然后通过 all-reduce 把不同 GPU 上的梯度求平均。这样就相当于多张卡一起组成了一个更大的 batch。

所以 effective batch size 可以写成：

\[ B_{\text{eff}}=B_{\text{micro}}\times A\times W \]

其中：

• \(B_{\text{micro}}\)：每张卡单次 forward/backward 的 batch size。
• \(A\)：gradient accumulation steps。
• \(W\)：world size，也就是参与训练的 GPU 数量。

例如单卡训练时，MiniMind 默认配置下：

\[ B_{\text{eff}}=32\times 8\times 1=256 \]

如果是 4 张 GPU，并且每张卡仍然使用 batch_size=32、accumulation_steps=8，那么：

\[ B_{\text{eff}}=32\times 8\times 4=1024 \]

如果每条样本长度是 \(L\)，也可以进一步估计每次参数更新处理了多少 token：

\[ T_{\text{update}}=B_{\text{eff}}\times L \]

其中，\(T_{\text{update}}\) 表示每次 optimizer.step() 对应的 token 数。对于 LLM 训练来说，这个量往往比“多少条样本”更直观，因为不同样本的长度可能不同。

那么 effective batch size 是不是越大越好？不一定。

更大的 batch size 通常有几个好处：

• 梯度估计更平滑，训练曲线可能更稳定。
• 多卡训练时吞吐更高，硬件利用率可能更好。
• 每次参数更新看过的数据更多，噪声更小。

但它也有代价：

• 在总 token 数固定时，batch 越大，参数更新次数越少。
• batch 变大后，学习率、warmup 和 weight decay 往往都需要重新调。
• batch 太大时，梯度噪声太小，泛化不一定更好。
• 继续增大 batch 到某个程度后，收益会明显变小。

所以 batch size 不是单纯追求越大越好，而是在显存、吞吐、训练稳定性、更新次数之间做平衡。LLM pretrain 通常会使用较大的 effective batch size，但它会配合学习率调度、warmup 和稳定性监控一起调。

Q10.2: Epoch、tokens seen 和 training compute 应该怎么理解?

Epoch 表示完整遍历训练集的次数。在传统深度学习里，数据集通常没有那么大，所以经常会训练很多 epoch。比如图像分类里，一个数据集可能被反复看几十遍甚至上百遍。

但是 LLM pretrain 的情况很不一样。语言模型的数据集往往非常大，大到完整训练一个 epoch 就已经需要巨大的计算量。所以在 LLM 里，大家更常用的训练计量单位不是 epoch，而是 tokens seen，也就是训练过程中模型实际处理过的 token 总数。

假设去重后的训练数据集一共有 \(D_{\text{dataset}}\) 个 token，训练了 \(E\) 个 epoch，那么：

\[ D_{\text{seen}}=E\cdot D_{\text{dataset}} \]

其中，\(D_{\text{seen}}\) 是训练过程中实际喂给模型的 token 总数。

Scaling law 里的训练 compute 近似可以写成：

\[ C\approx 6ND_{\text{seen}} \]

其中，\(N\) 是模型参数量，\(D_{\text{seen}}\) 是训练过程中实际处理过的 token 数。这里的 \(D_{\text{seen}}\) 不是说“数据集必须只过一遍”，而是模型训练期间总共看过多少 token。

如果数据集固定，就可以把 compute 和 epoch 联系起来：

\[ C\approx 6NE D_{\text{dataset}} \]

也就是说，在模型大小 \(N\) 和数据集大小 \(D_{\text{dataset}}\) 固定时，训练更多 epoch 会线性增加训练 compute。

反过来，如果有固定 compute 预算 \(C\)，可以粗略估计最多训练多少个 epoch：

\[ E\approx \frac{C}{6ND_{\text{dataset}}} \]

如果不知道 compute \(C\) 是多少, Chinchilla 论文里的一个经验值是,总token数大约为模型参数量的20倍：

\[ D_{\text{seen}} \approx 20 N \]

这样我们可以得到epoch: \[ E \approx \frac{D_{\text{seen}}}{D_{\text{dataset}}} \approx \frac{20 N}{D_{\text{dataset}}} \]

不过这里有一个容易混淆的问题：如果数据有限，第二遍、第三遍重复同一批数据，这到底算不算“增加数据”？

从训练计数上看，它确实增加了 tokens seen：

\[ D_{\text{seen}}=E\cdot D_{\text{dataset}} \]

也确实增加了 compute：

\[ C\propto D_{\text{seen}} \]

但它不等价于增加了新的数据多样性。更准确地说，我们可以区分两个概念：

\[ D_{\text{seen}}=\text{训练过程中实际处理过的 token 总数} \]

\[ D_{\text{unique}}=\text{去重后真正不同的 token 或文档规模} \]

如果同一个样本被看了三遍，那么 \(D_{\text{seen}}\) 会按三倍计算，但 \(D_{\text{unique}}\) 并没有变成三倍。第二遍、第三遍仍然能帮助优化，因为模型会继续从这些样本上更新参数；但它带来的信息增量通常会下降，并且可能增加过拟合和记忆风险。

这也是 LLM pretrain 和传统小数据深度学习的一个重要区别。传统深度学习经常默认“一个 epoch 太少”，因为数据集相对小，模型需要反复拟合这些样本。而 LLM pretrain 面对的是巨大的 token 流，很多时候一个 epoch 或少数几个 epoch 就已经足够昂贵。相比“把同一批数据刷很多遍”，更有价值的往往是扩大高质量、去重后的数据覆盖。

所以可以这样理解：

• 增加 epoch：增加 tokens seen，也增加 training compute。
• 增加 epoch 但不增加新样本：增加优化步数，但不增加数据多样性。
• 增加去重后的高质量数据：既增加可训练 token，也增加数据覆盖。
• 在 LLM pretrain 中，tokens seen 通常比 epoch 更适合作为训练进度指标。

因此，对于 pretrain，更自然的观察单位通常是：

• steps：优化器更新了多少次。
• tokens seen：模型总共处理了多少 token。
• unique tokens / unique documents：训练数据本身有多少去重后的信息。
• effective batch size：每次参数更新大约用了多少样本或 token。
• learning rate schedule：这些更新发生在什么学习率下。

Q10.3: 如何决定 effective batch size?

这个问题麻烦的地方在于：effective batch size 和 learning rate 是强耦合的。它们都会影响 loss 曲线的稳定性，而且有时候会造成相似的现象。所以实践里通常不是先拍脑袋决定一个 batch size，再单独决定 learning rate，而是把它们当成一组训练配置一起调。

参数更新可以粗略写成：

\[ \theta_{t+1}=\theta_t-\eta \hat{g}_t \]

其中，\(\eta\) 是 learning rate，\(\hat{g}_t\) 是当前 batch 估计出来的梯度。batch size 影响的是 \(\hat{g}_t\) 的噪声大小，learning rate 影响的是沿着这个梯度方向走多远。

如果 batch size 变大，\(\hat{g}_t\) 通常会更接近真实平均梯度，噪声更小：

但如果 learning rate 太大，即使梯度估计更平滑，参数也可能一步走太远，导致 loss spike 甚至 NaN。反过来，如果 batch size 很小，loss 也可能抖动很大，但这是梯度估计噪声造成的，不一定代表学习率已经过大。

不过实际操作时，仍然需要一个清晰的顺序。一个比较稳妥的流程是：

第一步：先确定硬件能承受的 micro-batch。

单卡显存主要受模型大小、序列长度和 micro-batch size 影响。直观上，可以理解为：

\[ \text{memory}\uparrow \quad \text{when}\quad B_{\text{micro}}\uparrow,\ L\uparrow,\ d_{\text{model}}\uparrow,\ \text{layers}\uparrow \]

其中，\(B_{\text{micro}}\) 是每张卡单次 forward/backward 的样本数，\(L\) 是序列长度。通常先固定模型和 seq_len，然后从 batch_size=1,2,4,8... 往上试，找到不会 OOM、吞吐也还可以的 \(B_{\text{micro}}\)。这个阶段主要受显存约束，不要急着讨论理论最优。

第二步：通过梯度累积和多卡得到几个候选 effective batch size。

确定 \(B_{\text{micro}}\) 后，再用：

\[ B_{\text{eff}}=B_{\text{micro}}\times A\times W \]

得到实际每次参数更新对应的 batch size。这里 \(A\) 是 accumulation_steps，\(W\) 是 GPU 数。

对于 LLM，更建议同时看每次 optimizer step 处理多少 token：

\[ T_{\text{update}}=B_{\text{eff}}\times L \]

其中，\(L\) 是序列长度。相比单纯看 \(B_{\text{eff}}\)，\(T_{\text{update}}\) 更适合 LLM，因为 seq_len=512 和 seq_len=4096 时，同样的 batch size 对应的 token 数完全不同。

对于学习项目或小模型，可以先从比较温和的范围试起。一个粗略参考是：

• 如果 \(L=512\)，可以先试 \(B_{\text{eff}}=64,128,256,512\)，对应每次更新大约 \(3.2\times 10^4\) 到 \(2.6\times 10^5\) tokens。
• 如果 \(L=1024\)，可以先试 \(B_{\text{eff}}=32,64,128,256\)，对应每次更新大约 \(3.2\times 10^4\) 到 \(2.6\times 10^5\) tokens。
• 如果是更大的训练，常见做法会继续把 \(T_{\text{update}}\) 提到 \(10^5\sim 10^6\) 甚至更高，但这通常需要更仔细的 LR、warmup 和稳定性调参。

这些不是理论最优值，只是起步参考。真正的选择还要看模型大小、数据质量、显存、吞吐和 loss 曲线。

第三步：在候选 \(B_{\text{eff}}\) 上重新寻找 learning rate。

这一步很重要：当 \(B_{\text{eff}}\) 改变时，learning rate 通常也要跟着改。一个常见经验是：

\[ \eta_{\text{new}}\approx \eta_{\text{old}}\left(\frac{B_{\text{new}}}{B_{\text{old}}}\right)^p \]

其中，\(p\) 通常可以先在 \([0.5,1.0]\) 之间理解。\(p=1\) 接近 linear scaling rule，\(p=0.5\) 接近 square-root scaling。对 LLM 的 AdamW 训练来说，这只是调参起点，不是定律。batch 变大时，通常还需要更长的 warmup。

所以更实际的做法是：不要只试一个 learning rate。对于一个候选 \(B_{\text{eff}}\)，至少试两三个 learning rate，比如原始学习率、稍小一点、稍大一点。这样才能判断问题到底来自 batch size，还是来自 LR 不匹配。

第四步：用 loss 曲线区分 batch 问题和 LR 问题。

从 loss 曲线上，可以大致这样判断：

• batch 太小：单 step loss 抖动很明显，但滑动平均 loss 仍然能稳定下降；梯度范数可能有噪声，但不一定持续变大。
• learning rate 太大：loss 不只是抖动，而是出现明显 spike、持续升高，甚至 NaN；梯度范数也可能突然变大。
• batch 很大但 LR 偏小：loss 曲线很平滑，但按 tokens seen 比较时下降很慢，训练像是在“慢慢挪”。
• batch 变大但 LR / warmup 没跟上：前期 loss spike，或者 warmup 结束附近突然不稳定。

这里最容易误判的是第一种和第二种：它们都可能表现为 loss 抖动。但 batch 太小通常是“有噪声地下降”，而 LR 太大更像是“更新走过头”，会出现 spike、发散或者 NaN。

第五步：比较实验时按 tokens seen，而不是只按 step。

batch size 改变后，每个 step 处理的 token 数也会改变。如果只看 step，batch 大的实验每一步看了更多 token，比较会不公平。所以更合理的是固定模型、数据、seq_len 和总 tokens seen，然后做少量网格：

如果只能做很少实验，我更倾向于先选一个不会让更新次数太少的中等 \(B_{\text{eff}}\)，再围绕它调 learning rate。因为 batch 太大以后，虽然曲线好看，但在固定 tokens seen 下参数更新次数会明显减少，小模型实验里不一定划算。

所以总结成一句话：先由显存确定 micro-batch 的可行范围，再用梯度累积和多卡得到适中的候选 \(B_{\text{eff}}\)，然后根据候选 batch 重新调 LR, 如果能work最好,但是如果不行的话就需要重新选择 batch 再选择LR ，最后评估的时候,除了要看 step 级别的曲线,还要按 tokens seen 的 loss 曲线判断模型是否更新的足够快。 需要始终记住的是，loss 抖动不一定只说明 batch 小，也可能是 LR 太大；这两个旋钮必须一起看。

Q11: 为什么不同实验的 loss 曲线不能随便比较?

这算是Eval章节的一个实操预演,我们会在Eval章节进行更多的展开.

比如，一个实验 padding 很多，另一个实验有效 token 更多，它们虽然都在算平均 loss，但实际参与统计的 token 分布并不一样。再比如，effective batch size 变大后，loss 曲线可能更平滑，但这不一定代表模型能力更强，因为它的更新次数、梯度噪声和 learning rate 适配关系都变了。

所以训练指标不是孤立数字。每一条 loss 曲线背后，都绑定着一整套数据和训练配置。

如果要比较不同实验，至少需要尽量检查下面这些参数。这里大致按照重要性排列：

1. 数据集是否一致：包括数据来源、数据清洗、去重、过滤规则、数据混合比例。
2. tokenizer 是否一致：不同 tokenizer 会改变 token 数、切分方式和 loss 的统计单位。
3. tokens seen 是否一致：不同实验最好按训练 token 数比较，而不是只按 step 或 epoch 比较。
4. max_seq_len 是否一致：序列长度会影响上下文范围、padding 比例和每步 token 数。
5. packing、padding、truncation 规则是否一致：这些会影响有效 token 比例。
6. loss mask 是否一致：padding token、特殊 token、prompt token 是否参与 loss，都会改变 loss 数值。
7. 模型结构是否一致：参数量、层数、hidden size、attention head、position encoding、是否 tie embedding。
8. optimizer 是否一致：AdamW、Lion、Muon 这类优化器的更新方式不同，loss 曲线形态也可能不同。
9. learning rate 和 schedule 是否一致：包括初始 LR、warmup、cosine decay、最小 LR。
10. effective batch size 是否一致：包括 batch_size、accumulation_steps、world size。
11. weight decay、gradient clipping 是否一致：这些会影响训练稳定性和参数范数。
12. mixed precision 设置是否一致：比如 fp16、bf16、loss scale、是否使用 gradient scaler。
13. 随机性是否一致：seed、shuffle、dataloader worker、分布式采样方式都会影响短期曲线。
14. eval 设置是否一致：eval 数据集、eval 间隔、eval batch size、是否固定 eval samples。

除了看一条 train loss 曲线，实际训练时还可以同时看几类曲线：

• train step loss：每一步训练 loss，最原始，但噪声最大。
• smoothed train loss：对 step loss 做滑动平均，更适合看整体下降趋势。
• train loss vs tokens seen：比 loss vs step 更适合比较不同 batch size 的实验。
• eval loss：在固定验证集上的 loss，更适合判断泛化能力。
• eval perplexity：由 eval loss 转换得到，公式是：

\[ \text{PPL}=\exp(L) \]

其中，\(L\) 是平均 cross entropy loss。

• learning rate curve：确认 warmup、decay 是否按预期执行。
• gradient norm curve：判断是否有梯度爆炸、LR 过大或训练不稳定。
• parameter norm / update norm：观察参数本身和每次更新幅度是否异常。
• loss scale / NaN count：混合精度训练时用于判断数值稳定性。
• tokens/sec 或 step time：观察吞吐，不然可能 loss 下降快但训练效率很差。
• train loss 和 eval loss 的 gap：如果 train loss 继续下降但 eval loss 不降，可能有过拟合、数据污染或训练分布不匹配的问题。

所以比较实验时，不要只说“第 N step 的 loss 更低”。尤其当 \(B_{\text{eff}}\)、seq_len 或 accumulation_steps 不同时，第 N step 背后对应的 tokens seen 可能完全不同。更稳妥的做法是：用 loss vs tokens seen 作为主轴，再结合 eval loss、gradient norm、learning rate curve 和吞吐一起判断。

Q12: 一个 CheckList

整个训练流程可以按照如下步骤检查规划:

1.根据目标和算力，确定模型参数量和模型规模：

• 影响 layers、hidden size、heads、训练预算。

2.根据目标上下文能力，设计数据集和序列长度：

• 影响 max_seq_len、packing、padding、truncation。

3.根据模型参数量和数据集规模，估算训练 token 数：

• 影响 tokens seen、epoch、是否需要重复数据。

4.根据 optimizer、weight decay、precision 等训练配置，估算训练状态开销：

• 影响显存、optimizer state、是否需要 mixed precision / checkpointing。

5.根据实际硬件显存

• 设置 batch_size、gradient_accumulation_steps 和分布式方式
• 影响 B_eff、tokens per update、吞吐。

6.观察训练稳定性和收敛速度

• 根据 B_eff 和 loss 曲线
• 调整 learning rate、warmup、decay schedule

7.训练过程中其他Eval指标的监控

• loss、eval loss、PPL, gradient norm、LR curve、tokens/sec。

8.实验管理

• 记录配置，保证实验可复现、可比较。
• 可视化清晰

参考资料

• Scaling Laws for Neural Language Models
• Training Compute-Optimal Large Language Models
• Symbolic Discovery of Optimization Algorithms
• Muon: An optimizer for hidden layers in neural networks