Pretrain 的训练目标和 Loss

这一节先回答 Pretrain 中最核心的理论问题：

模型到底在学什么，loss 是怎么定义的?
为什么 MiniMind 的 pretrain 用的是 Cross Entropy Loss?
LLM真的使用了交叉熵Loss吗？如何严格的推导损失函数的形式？

以及一个非常现实,但却能极大的方便理解的问题:

实际的训练数据集是什么样子的?

Q1: Pretrain 阶段模型到底在预测什么?

以 GPT 类自回归语言模型为例，pretrain 的核心任务是 next token prediction，也就是根据前面的 token 预测下一个 token。

假设一段文本经过 tokenizer 之后得到 token 序列：

\[ x_1, x_2, \cdots, x_T \]

其中，$x_t$ 表示第 $t$ 个 token id，$T$ 表示序列长度。

自回归语言模型会把整段文本的联合概率拆成一连串条件概率：

\[ p(x_1, x_2, \cdots, x_T) = \prod_{t=1}^{T} p(x_t \mid x_1, x_2, \cdots, x_{t-1}) \]

也就是说，模型不是一次性预测整段文本，而是在每个位置上预测“下一个 token 应该是什么”。这是一个自回归(AutoRegressive)的过程,所以只需要有任何无标注的自然文本(只需要把文本往后错一位,作为预测的ground truth即可),就可以进行自监督学习(Self‑Supervised Learning).这是模型可以scaling up的基石.

但是需要注意的是,对于一个句子,模型并不是依次预测 $x_1$ (根据上下文预测第一个 token)，$x_2$ (根据上下文预测第二个 token)，一直到 $x_T$ (根据上下文预测最后一个 token),而是并行预测所有token作为loss整体计算,也就是说模型不但学习了如何预测下一个词,同时也是被隐式约束了如何预测整个句子.

Q1.1 实际训练数据的样子?

这是Minimind中预训练数据集的示例,实际上预训练数据集非常大,即使是纯文本数据,也有几个GB的规模,并且高质量数据集是模型效果好的重要条件.

{"text": "如何才能摆脱拖延症？治愈拖延症并不容易，但以下建议可能有所帮助。"}
{"text": "清晨的阳光透过窗帘洒进房间，桌上的书页被风轻轻翻动。"}
{"text": "Transformer 通过自注意力机制建模上下文关系，是现代大语言模型的重要基础结构。"}

在代码里，这个思想通常会被实现成输入和标签错开一位：

# src/minimind_learning/dataset/lm_dataset.py
input_ids = encoding.input_ids.squeeze()
loss_mask = (input_ids != self.tokenizer.pad_token_id)

X = torch.tensor(input_ids[:-1], dtype=torch.long)
Y = torch.tensor(input_ids[1:], dtype=torch.long)
loss_mask = torch.tensor(loss_mask[1:], dtype=torch.long)
return X, Y, loss_mask

这里的 X 是去掉最后一个 token 的输入序列，Y 是去掉第一个 token 的目标序列。于是模型在位置 $t$ 看到的是 $x_1, \cdots, x_t$，要预测的是 $x_{t+1}$。

Q2: 为什么 Pretrain 用 Cross Entropy Loss?

根据上文思想,模型是在一个离散的Vocab Size的词表上进行多分类,并且输出的是每个分类的上面的概率,即输出了一个概率分布.我们希望最大化正确分类的的概率,所以这相当于需要优化整个分布,所以这个Loss中有 Entropy这个名称.虽然他看起来好像只是在最优化选择正确的概率,但实际上这只是这个分布式离散的onehot分布产生的错觉. 下面给出 Cross Entropy Loss 的定义。

假设真实分布是 $p$，模型预测分布是 $q$，词表大小为 $V$。对于一个离散分类问题，交叉熵定义为：

\[ H(p, q)=-\sum_{i=1}^{V}p_i \log q_i \]

其中，$p_i$ 表示真实分布中第 $i$ 个类别的概率，$q_i$ 表示模型预测第 $i$ 个类别的概率。

在 pretrain 的 next token prediction 里，真实标签通常是一个 one-hot 分布。也就是说，真实的下一个 token 只有一个正确类别。假设真实 token id 是 $y$，那么：

\[ p_i = \begin{cases} 1, & i = y \\ 0, & i \ne y \end{cases} \]

代入交叉熵公式后，只有正确类别 $y$ 那一项会保留下来：

\[ H(p, q)=-\log q_y \]

所以从表面上看，Cross Entropy Loss 好像只是在最大化正确 token 的概率；但更准确地说，它是在比较两个离散分布：真实的 one-hot 分布和模型预测出来的词表概率分布。只是因为真实分布是 one-hot，公式最后才简化成了 $-\log q_y$。

具体来说在每一个预测位置上，模型最初的输出的不是一个单独的概率，而是对整个词表的打分。

假设词表大小为 $V$，模型在第 $t$ 个位置输出 logits：

\[ z_t \in \mathbb{R}^{V} \]

其中，$z_{t,i}$ 表示模型认为第 $i$ 个 token 作为下一个 token 的原始分数。经过 softmax 之后，可以得到模型预测的概率分布：

\[ q_\theta(i \mid x_{\le t}) = \frac{\exp(z_{t,i})}{\sum_{j=1}^{V}\exp(z_{t,j})} \]

其中，$q_\theta(i \mid x_{\le t})$ 表示参数为 $\theta$ 的模型，在看到前文 $x_{\le t}$ 后，认为下一个 token 是 $i$ 的概率。

真实标签只有一个 token id。假设真实的下一个 token 是 $y_t$，那么这个位置的交叉熵损失就是：

\[ \ell_t = -\log q_\theta(y_t \mid x_{\le t}) \]

这句话很关键：Cross Entropy Loss 会惩罚模型没有把足够高的概率分给真实 token。

MiniMind 的训练代码里对应的是：

# src/minimind_learning/trainer/train_pretrain.py
loss_fct = nn.CrossEntropyLoss(reduction='none')

with autocast_ctx:
    # X: [batch_size, seq_len]
    res = model(X)
    # res.logits: [batch_size, seq_len, vocab_size]
    # 每一个 token 位置都会输出一个 vocab_size 维的 logits，
    # 表示模型对“下一个 token 属于词表中每个类别”的原始打分。
    loss = loss_fct(
        # [batch_size, seq_len, vocab_size]
        #   -> [batch_size * seq_len, vocab_size]
        # 把 batch 维和 seq_len 维合并，把每个 token 位置都看成一个分类样本。
        res.logits.view(-1, res.logits.size(-1)),

        # Y: [batch_size, seq_len]
        #   -> [batch_size * seq_len]
        # 每个元素是对应位置的真实 next token id。
        Y.view(-1)
    ).view(Y.size())
    # loss: [batch_size * seq_len] -> [batch_size, seq_len]
    # reduction='none' 表示暂时不求平均，保留每个 token 位置自己的 loss，
    # 后面还要用 loss_mask 去掉 padding token。

这里 res.logits 的形状是 [batch_size, seq_len, vocab_size]。也就是说，batch 里的每一个样本、每一个 token 位置，都会输出一个覆盖整个词表的分类分布。

PyTorch 这里容易让人迷惑的一点是：nn.CrossEntropyLoss 当然支持 batch，但它默认把第 2 个维度当成类别维。常见输入形状是 [N, C]，其中 $N$ 是样本数，$C$ 是类别数；如果是更高维输入，也通常要求形状类似 [N, C, d1, d2, ...]。

但语言模型的 logits 通常是 [batch_size, seq_len, vocab_size]，类别维 vocab_size 在最后一维。为了让 CrossEntropyLoss 正确理解“每个 token 位置是一个分类样本，类别数是 vocab_size”，MiniMind 这里选择把前两维展平：

\[ [B, L, V] \rightarrow [B \times L, V] \]

其中，$B$ 表示 batch size，$L$ 表示序列长度，$V$ 表示词表大小。

对应的标签也要从：

\[ [B, L] \rightarrow [B \times L] \]

这样展平之后，CrossEntropyLoss 看到的就是一个标准多分类问题：一共有 $B \times L$ 个样本，每个样本有 $V$ 个类别，标签是一个真实 token id。算完后再 .view(Y.size()) 还原回 [batch_size, seq_len]，方便后面和 loss_mask 对齐。

MiniMind3 升級的Shifted LM Loss

旧版训练代码中，Dataset 返回 (X, Y, loss_mask)，训练脚本手动计算交叉熵：

loss = loss_fct(
    # B batch size, L sequence length, V vocab size
    # [B,L,V] -> [B*L,V]
    res.logits.view(-1, res.logits.size(-1)),
    Y.view(-1)
).view(Y.size())

loss = (loss * loss_mask).sum() / loss_mask.sum()

MiniMind 3 更接近 HuggingFace 模型的接口习惯：Dataset 返回 input_ids 和 labels，模型内部完成 shift 和 ignore_index=-100 的 loss 计算。

对应代码在 src/minimind_learning/model/model_minimind.py：

loss = None
if labels is not None:
    # shifted LM loss
    # B batch size, L sequence length, V vocab size
    # logits : [B,L,V] -> [B,L-1,V] 
    # labels: [B,L] -> [B,L-1]

    x = logits[..., :-1, :].contiguous()
    y = labels[..., 1:].contiguous()
    loss = F.cross_entropy(x.view(-1, x.size(-1)), y.view(-1), ignore_index=-100)

这个变化看起来只是接口变化，但实际影响很大：pretrain、SFT、DPO 都可以围绕同一种模型输出结构来组织，训练脚本不再需要各自手写一套语言模型 loss。

Q3: BCE、Cross Entropy 和 KL 散度有什么关系?

这里先把几个常见 loss 的数学形式放在一起。它们不是完全割裂的东西，而是在不同任务设定下对“分布差异”的不同写法。

KL 散度

假设真实分布是 $p$，模型预测分布是 $q$，离散类别数为 $V$。KL 散度定义为：

\[ D_{\mathrm{KL}}(p \Vert q)=\sum_{i=1}^{V}p_i \log \frac{p_i}{q_i} \]

其中，$p_i$ 表示真实分布中第 $i$ 个类别的概率，$q_i$ 表示模型预测第 $i$ 个类别的概率。

KL 散度衡量的是：如果真实分布是 $p$，但我们用 $q$ 去近似它，会多付出多少信息代价。

!注意 KL 散度不是一个对称的距离，也就是说 $D_{\mathrm{KL}}(p \Vert q)$ 和 $D_{\mathrm{KL}}(q \Vert p)$ 是不一样的。

Cross Entropy

交叉熵定义为：

\[ H(p, q)=-\sum_{i=1}^{V}p_i \log q_i \]

\[ H(p, q)=-\mathbb{E}_{x \sim p}[\log q(x)] \]

它和 KL 散度的关系可以直接展开：

\[ D_{\mathrm{KL}}(p \Vert q)=\sum_{i=1}^{V}p_i \log \frac{p_i}{q_i} =\sum_{i=1}^{V}p_i \log p_i-\sum_{i=1}^{V}p_i \log q_i \]

而真实分布 $p$ 的熵为：

\[ H(p)=-\sum_{i=1}^{V}p_i \log p_i \]

所以：

\[ D_{\mathrm{KL}}(p \Vert q)=H(p, q)-H(p) \]

也就是：

\[ H(p, q)=H(p)+D_{\mathrm{KL}}(p \Vert q) \]

在训练时，真实数据分布 $p$ 是固定的，$H(p)$ 不依赖模型参数 $\theta$。因此最小化 Cross Entropy $H(p, q_\theta)$，等价于最小化 KL 散度 $D_{\mathrm{KL}}(p \Vert q_\theta)$。

one-hot 标签下的 CE

在 next token prediction 中，真实标签通常是 one-hot 分布。假设真实类别是 $y$，则：

\[ p_i = \begin{cases} 1, & i = y \\ 0, & i \ne y \end{cases} \]

代入交叉熵：

\[ H(p, q)=-\sum_{i=1}^{V}p_i \log q_i=-\log q_y \]

所以 PyTorch 里的多分类 CrossEntropyLoss，在 one-hot 标签场景下，最终就变成了对真实类别概率取负对数。(它其实是个distribution loss,而不是最大化某个概率)

这也是语言模型 pretrain 的形式：

\[ \ell_t=-\log q_\theta(y_t \mid x_{\le t}) \]

其中，$y_t$ 是第 $t$ 个位置的真实 next token。

BCE

BCE 是 Binary Cross Entropy，也就是二分类交叉熵。二分类时，标签 $y \in {0, 1}$，模型预测正类概率为 $\hat{y}$。

二分类可以看成一个两类分布：

\[ p = [y, 1-y] \]

\[ q = [\hat{y}, 1-\hat{y}] \]

把它代入交叉熵公式：

\[ H(p, q)=-\sum_{i=1}^{2}p_i \log q_i \]

得到：

\[ \mathrm{BCE}(y, \hat{y})=-\left[y \log \hat{y}+(1-y)\log(1-\hat{y})\right] \]

所以 BCE 不是另一种完全独立的 loss，而是 Cross Entropy 在二分类问题下的特例。

小结

这几者的关系可以简化成：

KL 散度衡量两个分布之间的差异。
Cross Entropy 和 KL 散度只差一个不依赖模型参数的 $H(p)$，所以优化 CE 等价于优化 KL。
one-hot 多分类下，CE 简化为 $-\log q_y$。
二分类下，CE 就是 BCE。

回到 pretrain，next token prediction 是在整个词表上做多分类，而不是二分类。因此代码中使用的是多分类 Cross Entropy Loss。

Q4: LLM真的使用了交叉熵Loss吗?

LLM真的是交叉熵loss吗? 交叉熵loss其实有点奇怪,一个loss从理论上到论文里实际落地的过程中,一般是先写出模型,明白模型在做什么,然后理论上推导loss的期望形式,最后在实际落地的时候,loss的样本期望又被用样本平均近似,但是仔细推导一下交叉熵loss的这个过程,会觉得比较奇怪. 因为如果直接把交叉熵当作loss,会发现写出样本期望并不是很直观.

我们尝试按照刚才的流程,写模型的优化目标:

假设文本序列来自真实数据分布 $p_{\text{data}}(x_{1:T})$。其中，$x_{1:T}$ 表示一段长度为 $T$ 的 token 序列，$x_{<t}$ 表示第 $t$ 个 token 之前的所有上下文，模型预测分布是 $q_\theta(x_t \mid x_{<t})$。

如果进一步展开到“给定上下文后预测下一个 token”的条件分布,然后求CE Loss，整个模型的优化目标可以写成:

(输入的样本是 $ x_{<t} $ 随机变量序列)

$\mathcal{L}(\theta)=\mathbb{E}_{x_{<t}\sim p_{\text{data}}} \left[H\left(p_{\text{data}}(\cdot \mid x_{<t}), q_\theta(\cdot \mid x_{<t})\right)\right]$

这里的 $p_{\text{data}}(\cdot \mid x_{<t})$ 表示真实数据中“给定上下文 $x_{<t}$ 后，下一个 token 的条件分布”，$q_\theta(\cdot \mid x_{<t})$ 表示模型预测的条件分布。里面的交叉熵实际是一个条件交叉熵.这时候最好把交叉熵协程期望的形式,在写成条件期望的形式,也就是:

$H\left(p_{\text{data}}(\cdot \mid x_{<t}), q_\theta(\cdot \mid x_{<t})\right) =\mathbb{E}_{y\sim p_{\text{data}}(\cdot \mid x_{<t})} \left[-\log q_\theta(y \mid x_{<t})\right]$

把里面的交叉熵继续写成离散形式，就是：

$H\left(p_{\text{data}}, q_\theta\right) =-\sum_{y\in \mathcal{V}}p_{\text{data}}(y \mid x_{<t})\log q_\theta(y \mid x_{<t})$

其中，$\mathcal{V}$ 表示词表。这个求和是在所有可能的 next token 上求和。

但真实的数据分布(next token的分布,即y的分布)我们并不知道，也不可能真的对所有上下文和所有 next token 做完整期望。训练时，我们只有数据集里的有限样本。所以实际训练会用样本平均来近似这个期望。

这就是为什么CE损失函数的绕人的地方,是因为他做了两次样本期望的近似,一次是Batch 中的多条序列样本的期望,另一次是每条序列中next token的样本的期望.所以看起来好像是“期望套期望“.这确实不trival.

具体来说假设一个 batch 中有 $B$ 条样本，每条样本有 $L$ 个预测位置，第 $b$ 条样本第 $t$ 个位置的真实 next token 是 $y_{b,t}$，上下文是 $x_{b,<t}$，那么离散化后的训练 loss 可以写成：

$\mathcal{L}_{\text{batch}}(\theta) =-\frac{1}{B L}\sum_{b=1}^{B}\sum_{t=1}^{L}\log q_\theta(y_{b,t} \mid x_{b,<t})$

如果把 batch 里的所有有效 token 位置都重新编号为 $n=1,\cdots,N$，也可以写成更紧凑的形式：

$\mathcal{L}_{\text{batch}}(\theta) =-\frac{1}{N}\sum_{n=1}^{N}\log q_\theta(y^{(n)} \mid c^{(n)})$

其中，$N$ 表示参与统计的有效 token 数量，$c^{(n)}$ 表示第 $n$ 个训练位置对应的上下文，$y^{(n)}$ 表示这个位置对应的真实 next token。

另一点更绕的是,从期望形式写出两层Loss时,里面的条件期望并不是上面那种给定 $x_{<t} $ 后的条件期望, 序列本身也变成了随机变量,更准确的写法是$X_{<t} $,即这个期望是$X_{<t} $ 这个随机变量序列的函数:

那么Loss就变成了

$\mathcal{L}(\theta) = \mathbb{E}_{x_{<t}\sim p_{\text{data}}} \left[ \mathbb{E}_{y\sim p_{\text{data}}(\cdot \mid X_{<t})} \left[-\log q_\theta(y \mid X_{<t})\right] \right]$

这并不是多此一举，只有这样才能使用全期望公式把外层的期望和内层的期望合成一个期望： \[ E[X] = E[E[X|Y]] \\ E[X|Y] = E[E[X|Y,Z]|Y] \\ \]

$\mathcal{L}(\theta) = \mathbb{E}_{y ,x_{<t} \sim p_{\text{data}}} \left[-\log q_\theta(y \mid x_{<t})\right]$

这里的两个下下标实际暗示进行了两次采样，一次对序列，来自batch，构造context，一次在序列中对nexttoken

返璞归真

很多时候自回归语言模型的理论训练目标可以写成，下面这个简单形式，它的下标只表示了对不同的序列Batch进行了采样，实际来自已经样本化的CE loss.这个式子并不向它表面那么trival：

$\mathcal{L}(\theta)=\mathbb{E}_{x_{1:T}\sim p_{\text{data}}} \left[-\sum_{t=1}^{T}\log q_\theta(x_t \mid x_{<t})\right]$

这个公式的意思是：从真实数据分布中采样一段文本，然后让模型在每个位置都预测当前 token，并把所有位置的负对数概率加起来。训练的目标就是让这个期望尽可能小。

但是如果我们忽略掉之前所有的推导，尤其是CE的部分，我们会发现，实际我们模型建模的是整个序列的联合分布，即 \[ p_\theta(x_{1:T}) = \prod_{t=1}^{T} q_\theta(x_t \mid x_{<t}) \] 我们的样本是一系列序列，所以优化目标实际应该写成：

$\mathcal{L}(\theta) = \mathbb{E}_{x_{1:T} \sim p_{\text{data}}} \left[ p_\theta(x_{1:T})\right]$

如果要最大化这个目标函数，可以等效于最小化$-\log p_\theta(x_{1:T})$，也就是：

$\mathcal{L}(\theta)=-\mathbb{E}_{x_{1:T}\sim p_{\text{data}}} \left[\log \prod_{t=1}^{T} q_\theta(x_t \mid x_{<t})\right]$ $\mathcal{L}(\theta)= - \mathbb{E}_{x_{1:T}\sim p_{\text{data}}} \left[\sum_{t=1}^{T}\log q_\theta(x_t \mid x_{<t})\right]$

正好就是我们之前写的那个公式。也就是说，虽然我们在推导过程中引入了交叉熵、条件分布、期望等概念，但最终的训练目标其实就是最大化模型对真实数据序列的联合概率。这个视角比使用交叉熵清晰的多得多得多，也许代码里使用交叉熵损失函数，只是这个api适合通过token index选择对应的概率，这功能恰巧和交叉熵损失函数一致了。

但这确实是两个完全不同的视角，前者是从“分布差异”的角度出发，强调模型预测分布和真实分布之间的关系；后者是从“概率最大化”的角度出发，强调模型对真实数据的拟合程度虽然它们最终等价，但前者更适合理解为什么使用交叉熵损失，而后者更适合理解模型的最终优化目标。

分布拟合视角（Cross‑Entropy / KL）
最大似然视角（Maximize joint probability）

更深层的角度 📌 最大似然、KL 散度与交叉熵的统一视角

在统计学中有一个非常重要的普适定理：

$\textbf{最大似然} \;\Longleftrightarrow\; \textbf{最小化 } KL(p_{\text{data}}\parallel p_\theta) \;\Longleftrightarrow\; \textbf{最小化交叉熵}$

这并不是语言模型特有的性质，而是所有概率模型都成立的基本等价关系。

KL 散度与最大似然的等价性

对任意真实分布 $p_{\text{data}}$ 和模型分布 $p_\theta$，KL 散度定义为：

$KL(p_{\text{data}}\parallel p_\theta) = \mathbb{E}_{x\sim p_{\text{data}}} \left[\log p_{\text{data}}(x) - \log p_\theta(x)\right]$

注意到：

第一项 $\mathbb{E}[\log p_{\text{data}}(x)]$ 与模型参数 $\theta$ 无关
因此优化 KL 散度只影响第二项

于是：

$\min_\theta KL(p_{\text{data}}\parallel p_\theta) \quad\Longleftrightarrow\quad \max_\theta \mathbb{E}_{x\sim p_{\text{data}}} \left[\log p_\theta(x)\right]$

右侧正是最大似然（MLE）的目标函数。

结论

最小化 KL(p‖q) 与最大化模型似然完全等价。
这条等价关系在所有概率模型中成立，不限于语言模型。

Q5: `loss_mask` 在 loss 里起什么作用?

训练时并不是所有位置都应该参与 loss。比如 padding token 只是为了把不同长度的样本补齐到同一个长度，它不应该影响模型训练。

MiniMind 里先让 CrossEntropyLoss(reduction='none') 返回每个位置的 loss，然后再用 loss_mask 过滤：

# src/minimind_learning/trainer/train_pretrain.py
loss = loss_fct(
    res.logits.view(-1, res.logits.size(-1)),
    Y.view(-1)
).view(Y.size())

# loss 和 loss_mask 的shape都是 [batch_size, seq_len]
loss = (loss * loss_mask).sum() / loss_mask.sum()

这里 loss_mask 中为 1 的位置参与损失计算，为 0 的位置被忽略。最后除以 loss_mask.sum()，表示只对有效 token 的 loss 取平均。

用公式写出来，假设 $m_t$ 表示第 $t$ 个位置的 mask，$m_t = 1$ 表示参与训练，$m_t = 0$ 表示忽略，那么最终 loss 是：

\[ \mathcal{L}=\frac{\sum_{t=1}^{T}m_t \ell_t}{\sum_{t=1}^{T}m_t} \]

其中，$\ell_t$ 表示第 $t$ 个位置的交叉熵损失。

这也是为什么 loss_mask 看起来只是一个工程细节，但实际上它决定了哪些 token 真正参与了模型学习。

Keyboard shortcuts

MiniMind 学习笔记