起因#
CS336 的 training loop 里，loss.backward() 和 optimizer.step() 挨在一起。前者计算 gradient，后者更新 parameter；但 gradient 存在哪里、什么时候清零、optimizer 还维护什么状态，需要分开看。于是结合 MIT 6.S191 的 Lecture 1，把一次训练迭代重新过了一遍。
训练到底在做什么#
一个 neuron：
$$ z = Wx + b $$$$ h = f(z) $$\(W\) 和 \(b\) 是 parameter，\(f\) 是 activation function。堆很多层就有了 neural network。
这里有个容易跳过的细节：如果没有 non-linear activation，连续堆叠多层——
$$ W_3(W_2(W_1 x)) $$仍然只是一个线性变换，等价于某个 \(W'x\)。深度本身不带来表达能力。Non-linearity 才让多层 composition 有意义。
有了 architecture 以后，所谓"训练"就是在解一个优化问题：
$$ \theta^* = \arg\min_\theta L(\theta) $$说白了，architecture 决定了网络能表达哪些函数，loss 定义了目标，training 就是在 parameter space 里找一个尽量好的点。
Backward 不更新模型#
其实 backward 和 optimizer step 是两件不同的事。
Backward：
$$ L \rightarrow \nabla_\theta L $$回答的是 loss 对每个 parameter 有多敏感。
Optimizer step：
$$ \nabla_\theta L \rightarrow \Delta\theta $$决定 parameter 往哪走、走多远。
最终：
$$ \theta_{t+1} = \theta_t + \Delta\theta $$Backward 只算 gradient。更新 parameter 是 optimizer 的事。
对应 PyTorch：
1optimizer.zero_grad()
2
3pred = model(x)
4loss = loss_fn(pred, y)
5
6loss.backward()
7optimizer.step()model(x) 和 loss_fn(pred, y) 是 forward。loss.backward() 沿 computation graph 反向算出每个 parameter 的 gradient，存进 .grad。optimizer.step() 根据 .grad 和 optimizer 自己的逻辑更新 parameter。
zero_grad() 是因为 PyTorch 的 .grad 默认累加——不清零的话，这轮的 gradient 会叠到上轮的上面。这不是 bug，gradient accumulation 就靠这个行为工作：累加多个 mini-batch 的 gradient 再做一次 update，模拟更大的 batch size。正常 training loop 里每步独立，所以每轮开头必须清零。
Chain Rule 和 Backpropagation#
Backward 的核心是 chain rule。
最简单的情况：
$$ x \rightarrow y \rightarrow L $$$$ \frac{\partial L}{\partial x} = \frac{\partial L}{\partial y} \cdot \frac{\partial y}{\partial x} $$也就是说：
$$ \text{upstream gradient} \times \text{local derivative} = \text{gradient passed backward} $$Backpropagation 就是在 computation graph 上反向地、逐层地执行这个过程。
有一个容易忘掉的细节。考虑某一层：
$$ y = f(x, W) $$Backward 要算两个东西。
Parameter gradient：
$$ \frac{\partial L}{\partial W} $$这个用来更新 \(W\)，没问题。
但还有 input gradient：
$$ \frac{\partial L}{\partial x} $$\(x\) 不是可学习的 parameter，它是上一层的输出。为什么要算它的 gradient？
因为上一层的 backward 需要它。上一层要算自己的 \(\frac{\partial L}{\partial W_{\text{prev}}}\)，就得知道 \(\frac{\partial L}{\partial x}\)——也就是当前层传回去的那个值。
对 input 求 gradient 不是为了更新它，而是为了让 gradient 能继续往回传。
这也是为什么 backward 的内存行为和 forward 很不一样。推理的时候，一层算完，input 可以丢。训练时 backward 算 parameter gradient 需要 forward 阶段的输入，所以这些中间结果必须一直留着，直到对应的 backward 跑完。
Optimizer#
SGD 最直接——沿当前 gradient 走一步：
$$ \theta_{t+1} = \theta_t - \eta g_t $$其中 \(g_t = \nabla_\theta L\)。
但 gradient 只是当前这一个 mini-batch、这一个点的 local information。它可能很 noisy，可能在不同方向上 scale 差得很远。直接跟着走，trajectory 经常很抖。
Momentum#
Momentum 的想法是引入历史方向：
$$ v_t = \beta v_{t-1} + g_t $$$$ \theta_{t+1} = \theta_t - \eta v_t $$如果某个方向上的 gradient 长期一致，就在这个方向上积累速度。偶尔出现的噪声会被平滑掉。
Adaptive Learning Rate#
SGD 和 Momentum 对所有 parameter 用同一个 learning rate。但不同 parameter 的 gradient scale 可能差好几个数量级。
AdaGrad、RMSProp 这类方法引入 gradient 二阶矩的历史——对梯度经常很大的 parameter 自动缩小步长，对梯度经常很小的 parameter 自动放大。
Adam#
Adam 大致可以理解成 momentum 的方向估计加上 adaptive 的 scale 估计。
它为每个 parameter 维护两个状态：
$$ m_t = \beta_1 m_{t-1} + (1 - \beta_1) g_t $$$$ v_t = \beta_2 v_{t-1} + (1 - \beta_2) g_t^2 $$\(m_t\) 是 gradient 一阶矩估计，\(v_t\) 是二阶原始矩估计。Update 大致是 \(m_t / \sqrt{v_t}\) 再乘 learning rate（还有 bias correction，这里不展开）。
Optimizer State 占的是真实的显存#
从 Infra 的角度看，这些数学设计直接变成了 GPU 上的 memory footprint。
Momentum 要为每个 parameter 存一份 \(v_t\)。Adam 要存两份：\(m_t\) 和 \(v_t\)。
粗估一下。FP32 Adam 下，每个 parameter 对应 parameter 本身、gradient、first moment、second moment，四份 float32 tensor。一个 1B 参数的模型，光这些就大约 16 GB。
Optimizer 的数学选择直接决定了训练时的一部分 memory model。 Mixed precision、ZeRO 的 optimizer-state sharding、master weights——归根结底都在处理这个问题。
Optimization 和 Generalization 是两件事#
Training loop 本身已经完整了。但有一个根本问题。
训练时优化的是 \(L_{\text{train}}\)，真正关心的是 \(L_{\text{test}}\)。
$$ \min L_{\text{train}} \nRightarrow \min L_{\text{test}} $$Train loss 可以一直降，validation performance 可能早就在变差——模型把 training data 里的 noise 和 pattern 一起学进去了。
所以光把 optimization 做好不够。
Dropout#
Dropout 在训练时随机 mask 一部分 activation：
$$ h' = \frac{m \odot h}{1 - p} $$\(m\) 是随机 binary mask，\(p\) 是 dropout rate。
注意不是"删除神经元"。Neuron 还在，只是这次 forward 它的输出被置零了。下次 forward 重新抽 mask，被 mask 掉的是另一批。
分母 \(\frac{1}{1-p}\) 是 inverted dropout 的关键：训练时 scale up 保证 \(E[h'] = h\)，这样 inference 时直接关掉 dropout 就行，不用额外缩放。
Dropout 的作用不只是笼统的"防过拟合"。更准确地说，它打破了模型对固定 feature 组合的 co-adaptation——每个 hidden unit 不能依赖某些特定的 unit 一定在，被迫学到更鲁棒的 representation。
有人把它理解成"同时训练很多 subnetwork"。有一定道理，但只是 intuition，不是严格的 ensemble。
Early Stopping#
另一种思路：不改变模型怎么学，而是控制学多久。
训练过程中定期在 validation set 上评估，performance 不再改善就停下来。
一个细节：early stopping 不一定在 epoch 结束时才发生。大规模训练里一个 epoch 可能包含几十亿 token，epoch 不一定是最自然的 evaluation 单位。常见的做法是 every N steps 或 every N tokens 评估一次。
对比一下：Dropout 改变模型怎么学，early stopping 控制学多久。 一个在每步 forward 里引入随机性，一个在 optimization trajectory 层面做截断。
重新看这条 Pipeline#
回到整体。一次 training iteration：
1data → forward → prediction → loss → backward → gradients → optimizer → updated parameters → next iteration在实际训练系统里，这条线展开以后会变成很多东西：forward 时 activation 必须保留给 backward；autograd 自动追踪 computation graph；gradient 可能累加多个 micro-batch 再 update；optimizer state 额外占显存；定期 checkpoint；distributed training 要跨卡同步 gradient；mixed precision 在不同阶段用不同 dtype；activation checkpointing 丢掉中间结果再重算..
但每一个都能在上面那条简单 pipeline 里找到它的位置。
Infra 层面的很多复杂性，追到源头就是这条 training loop 在规模变大以后的工程展开。
参考资料#
- MIT 6.S191: Introduction to Deep Learning — https://introtodeeplearning.com/