在本文，我们将从数学的角度出发，了解最基本的 pre-training 的数学原理。这是理解当先 LLM 大规模 pretraining 和 scaling 的第一步。

Abstractions: Auto-regressive token

假设输入文本经过 Tokenizer 后，得到一个长度为 TT 的 Token 序列：

x=(x1,x2,…,xT) x=(x_1,x_2,\ldots,x_T)

LLM 可以被建模为自回归的概率模型：

Pθ(x1,…,xT)=∏t=1TPθ(xt∣x<t) P_\theta(x_1,\ldots,x_T) = \prod_{t=1}^T P_\theta(x_t\mid x_{<t})

在这里，我们不考虑 LLM 内部复杂的模型结构，只考虑模型的输入输出, 模型输入历史的 tokens，输出对当前输出 token 的 logits:

zt=fθ(x≤t)∈RV z_t=f_\theta(x_{\le t}) \in \mathbb{R}^{V}

因此，输出的 Logits 为一个词表长度的向量:

zt=[zt,1zt,2⋮zt,V] z_t= \begin{bmatrix} z_{t,1}\\ z_{t,2}\\ \vdots\\ z_{t,V} \end{bmatrix}

最终，输出 logits 经过 softmax 归一化后，在推理侧通过适当的采样方式，得到最终的输出 tokens

pt,k=exp⁡(zt,k)∑j=1Vexp⁡(zt,j) p_{t,k}= \frac{\exp(z_{t,k})} {\sum_{j=1}^{V}\exp(z_{t,j})}

Training Objectives

在 Pre-training 阶段，我们希望模型能够根据历史 Tokens，尽可能准确地预测下一个 Token。

假设训练语料来自真实数据分布 PdataP_{\mathrm{data}}，模型参数化的条件概率分布为 PθP_\theta。对于输入前缀 x≤tx_{\le t}，模型经过 Softmax 输出词表上的概率分布：

pt,k=Pθ(xt+1=k∣x≤t)=exp⁡(zt,k)∑j=1Vexp⁡(zt,j) p_{t,k}=P_\theta(x_{t+1}=k\mid x_{\le t}) =\frac{\exp(z_{t,k})}{\sum_{j=1}^{V}\exp(z_{t,j})}

Maximum Likelihood Estimation (MLE)

我们首先来看最大似然估计的基本数学定义，给定训练数据集 D=x(1),x(2),…,x(N)\mathcal D={x^{(1)},x^{(2)},\ldots,x^{(N)}}, 模型为一段序列分配对应 token 的概率为:

Pθ(x(i))=∏t=1TiPθ(xt(i)∣x<t(i)) P_\theta(x^{(i)}) =\prod_{t=1}^{T_i} P_\theta(x_t^{(i)}\mid x_{<t}^{(i)})

PreTraining 的目标就是最大化这个概率 θMLE=arg⁡max⁡θ∏i=1NPθ(x(i))\theta_{\mathrm{MLE}} =\arg\max_\theta \prod_{i=1}^{N}P_\theta(x^{(i)}), 去对数后可以得到导出的训练目标:

θMLE=arg⁡max⁡θ∑i=1N∑t=1Tilog⁡Pθ(xt(i)∣x<t(i)) \theta_{\mathrm{MLE}} =\arg\max_\theta \sum_{i=1}^{N} \sum_{t=1}^{T_i} \log P_\theta(x_t^{(i)}\mid x_{<t}^{(i)})

Negative Log-Likelihood (NLL)

由于深度学习通常采用 Gradient Descent 最小化 Loss，我们可以对 Maximum Log-Likelihood 取负号，将最大化问题转换为最小化问题。我们定义 Negative Log-Likelihood:

LNLL,t=−log⁡Pθ(y∣x≤t)=−log⁡pt,y=−log⁡(exp⁡(zt,y)∑k=1Vexp⁡(zt,k))=−zt,y+log⁡∑k=1Vexp⁡(zt,k) \begin{aligned} \mathcal L_{\mathrm{NLL},t} &=-\log P_\theta(y\mid x_{\le t}) =-\log p_{t,y}\\ &=-\log\left( \frac{\exp(z_{t,y})} {\sum_{k=1}^{V}\exp(z_{t,k})} \right)\\ &=-z_{t,y}+\log\sum_{k=1}^{V}\exp(z_{t,k}) \end{aligned}

在 LLM-Pretraining 中，优化目标就是下一个输出的 token，可以被看作一个 one-hot 的 cross-entropy loss, 两者在数学上等价：

H(qt,pt)=−∑k=1Vqt,klog⁡pt,k=−log⁡pt,y=LNLL,t H(q_t,p_t) =-\sum_{k=1}^{V}q_{t,k}\log p_{t,k} = -\log p_{t,y} = \mathcal L_{\mathrm{NLL},t}

其中 qt,kq_{t,k} 是一个独热向量，只有 t=yt = y 的时候是 1, 其余时刻都是 0.

这个本质上也对 pretraining-loss 提出了一种全新的形式，从概率分布上，模型分布是在 Cross-Entropy 的监督下，去拟合 pretraining-data 的数据分布。

因此，我们最终导出了 LLM 训练的 Pretraining Loss (sequence-level):

LPT(θ)=1T−1NLLseq(θ)=1T−1∑t=1T−1NLLt(θ)=−1T−1∑t=1T−1log⁡Pθ(xt+1∣x≤t) \mathcal L_{\mathrm{PT}}(\theta) = \frac{1}{T-1} \mathrm{NLL}_{\mathrm{seq}}(\theta) = \frac{1}{T-1} \sum_{t=1}^{T-1}\mathrm{NLL}_t(\theta) =-\frac{1}{T-1} \sum_{t=1}^{T-1} \log P_\theta(x_{t+1}\mid x_{\le t})

Training Gradients

我们暂时不考虑时间步，只考虑当前 token 的 auto-regressive 输出:

L=NLLt(θ)=log⁡∑j=1Vezj−zy \mathcal L= \mathrm{NLL}_t(\theta) = \log\sum_{j=1}^{V}e^{z_j}-z_y

我们的重点是计算当前 loss 对模型输出 logits 的梯度 ∂L∂zk\frac{\partial\mathcal L}{\partial z_k}:

zkz_k -> pkp_k -> L\mathrm{L}

∂L∂zk=∂∂zklog⁡∑j=1Vezj−∂zy∂zk=ezk∑j=1Vezj−∂zy∂zk=pk−I[k=y] \frac{\partial\mathcal L}{\partial z_k} = \frac{\partial}{\partial z_k} \log\sum_{j=1}^{V}e^{z_j} - \frac{\partial z_y}{\partial z_k} = \frac{e^{z_k}}{\sum_{j=1}^{V}e^{z_j}} - \frac{\partial z_y}{\partial z_k} = p_k-\mathbb I[k=y]

我们可以考虑如下的梯度下降过程:

  • 当 kk 不等于 yy 时，indicator 为 0，此时梯度就是对应的输出 token 的概率。输出 token 的概率越大，梯度下降的越狠，更新后的 logits
  • 当 k=yk = y 的时候，此时 indicator 为 1，此时梯度下降的更新值为负数，对应的输出 logits 会被放大。

因此，从外部的 training-objective 出发，我们从理论证明了数学上 NLL-loss 可以保证对应 logits 的输出概率会被正确的校准到训练数据的分布概率（one-hot-objective），至于具体内部的模型参数更新，就依赖反向传播来进行就可以了。

KL Divergence

我们在 Agentic RL 的章节介绍过 KL 散度，重新引用一下定义：

假设有两个概率分布：

  • P(x)P(x)：目标分布（Target Distribution），希望模型学习的分布。
  • Qθ(x)Q_\theta(x)：模型分布（Model Distribution），可以通过参数 θ\theta 优化。

KL Divergence（相对熵）定义为：

DKL(P∥Q)=Ex∼P[log⁡P(x)Q(x)] D_{\mathrm{KL}}(P\|Q)=\mathbb E_{x\sim P}\left[\log\frac{P(x)}{Q(x)}\right]

对于离散分布：

DKL(P∥Q)=∑xP(x)log⁡P(x)Q(x) D_{\mathrm{KL}}(P\|Q)=\sum_x P(x)\log\frac{P(x)}{Q(x)}

对于连续分布：

DKL(P∥Q)=∫P(x)log⁡P(x)Q(x)dx D_{\mathrm{KL}}(P\|Q)=\int P(x)\log\frac{P(x)}{Q(x)}dx

KL 分布有两个重要的数学性质:

  • DKL(P∥Q)≥0D_{\mathrm{KL}}(P\|Q)\geq 0
  • DKL(P∥Q)=0  ⟺  P=QD_{\mathrm{KL}}(P\|Q)=0\iff P=Q

但是，KL 散度不是对称的，如果交换 Q 和 P 的位置，会产生不同的优化目标: DKL(P∥Q)≠DKL(Q∥P)D_{\mathrm{KL}}(P\|Q)\neq D_{\mathrm{KL}}(Q\|P)

在 Pretraining 中，KL 散度是 Forward-KL 的形式:

Forward KL 对当前参数分布中进行采样，检验该分布在目标分布中出现的比例。

DKL(P∥Qθ)=Ex∼P[log⁡P(x)Qθ(x)]=Ex∼P[log⁡P(x)]−Ex∼P[log⁡Qθ(x)] D_{\mathrm{KL}}(P\|Q_{\theta})=\mathbb E_{x\sim P}\left[\log\frac{P(x)}{Q_{\theta}(x)}\right] = E_{x\sim P}\left[\log P(x)\right] - E_{x\sim P}\left[\log Q_{\theta}(x)\right]

第一项是一个常数，和具体的模型参数无关，第二项就是 Pretrain/SFT 的训练目标。