LLM Systems Assignment 3: Transformer Architecture

本次 homework 的任务是实现 Decoder-only 的 Transformer 模型，也就是 GPT-2，并在一个翻译任务数据集上进行模型训练。具体来说，包括以下几个子任务：

• 实现 CrossEntropy Loss
• 实现基础模块，包括 Linear/Dropout/LayerNorm1d/Embedding
• 实现 Transformer，包括：MultiHeadAttention/FeedForward/TransformerLayer/DecoderLM
• 实现 Transformer 模型 generate 接口，即贪婪解码

下面对实现这几个子任务时的要点进行介绍。

1. CrossEntropy Loss

Problem 1 的 TODOs 要求实现 logsumexp 和 softmax_loss 两个函数，乍一看它们似乎和 CrossEntropy Loss 的原始定义没有任何关系。但实际上这是工程上实现 CrossEntropy Loss 数值稳定计算的典型方式。下面进行详细推导。

1.1 数值稳定的 CrossEntropy Loss 计算方式推导

假设某个样本的原始 logits 为 \(\mathbf{z} = [z_1, z_2, \dots, z_C]\)，真实类别为 \(c\)（即 one-hot 向量中只有第 \(c\) 个位置为 1，其余为 0）。CrossEntropy 的定义为： \[ L = - \sum_{i=1}^{C} y_i \log(p_i) \] 因为只有真实类别 \(c\) 的 \(y_c = 1\)，其他 \(y_i = 0\)，公式直接简化为： \[ L = - \log(p_c) \] 其中，概率 \(p_c\) 由原始 logits \(z_i\) 经过 Softmax 计算得到： \[ p_c = \frac{\exp(z_c)}{\sum_{j=1}^{C} \exp(z_j)} \] 将 \(p_c\) 代入 \(- \log(p_c)\) 中： \[ L = - \log\left( \frac{\exp(z_c)}{\sum_{j=1}^{C} \exp(z_j)} \right) \] 上式整体是一个 \(\log{\frac{a}{b}}\) 的形式，根据对数计算规则，其等于 \(\log{a} - \log{b}\) ，即： \[ L = - \left[ \log(\exp(z_c)) - \log\left( \sum_{j=1}^{C} \exp(z_j) \right) \right] \] 由于 \(\log(\exp(z_c)) = z_c\)（自然对数和指数运算互为逆运算，直接抵消）： \[ L = - \left[ z_c - \log\left( \sum_{j=1}^{C} \exp(z_j) \right) \right] \] 把负号放到里面，得到 Problem 1 TODOs 形式： \[ L = \underbrace{\log\left( \sum_{j=1}^{C} \exp(z_j) \right)}_{\text{logsumexp\_res}} - \underbrace{z_c}_{(\text{logits} \times \text{labels}).\text{sum()}} \] 其中第一项就是需要实现的 logsumexp，第二项是 label 对应的 logits 值。

1.2 LogSumExp 数值稳定计算方式推导

在上述 CrossEntropy Loss 化简形式的第一项，就是 LogSumExp（LSE） 的原始定义： \[ \text{LSE}(\mathbf{x}) = \log \left( \sum_{i=1}^{n} \exp(x_i) \right) \] 其中，假设沿着某维度有 \(n\) 个数：\(\mathbf{x} = (x_1, x_2, \dots, x_n)\)。

设该维度上的最大值为 \(x^* = \max(x_1, x_2, \dots, x_n)\)。我们可以将每个 \(x_i\) 拆分为 \((x_i - x^*) + x^*\)： \[ \text{LSE}(\mathbf{x}) = \log \left( \sum_{i=1}^{n} \exp((x_i - x^*) + x^*) \right) \] 根据指数运算法则 \(e^{a+b} = e^a \cdot e^b\)，展开求和项： \[ \text{LSE}(\mathbf{x}) = \log \left( \sum_{i=1}^{n} \left[ \exp(x_i - x^*) \cdot \exp(x^*) \right] \right) \] 因为 \(\exp(x^*)\) 与求和索引 \(i\) 无关，可以将其提到求和号 \(\sum\) 外面： \[ \text{LSE}(\mathbf{x}) = \log \left( \exp(x^*) \cdot \sum_{i=1}^{n} \exp(x_i - x^*) \right) \] 然后根据对数运算法则 \(\log(a \cdot b) = \log(a) + \log(b)\)，将对数拆开： \[ \text{LSE}(\mathbf{x}) = \log(\exp(x^*)) + \log \left( \sum_{i=1}^{n} \exp(x_i - x^*) \right) \] 第一项简化最终得到： \[ \text{LSE}(\mathbf{x}) = x^* + \log \left( \sum_{i=1}^{n} \exp(x_i - x^*) \right) \]

从上述公式可以看出，通过减去原始 logits 最大值，可以避免指数运算溢出风险，让数值计算更稳定。这一形式也是 Problem 1 中 logsumexp 函数中我们需要实现的形式。

2. Basic Modules

Basic Modules 的实现，在有了 Homework 1 的经验之后整体流程非常清晰：

• 定义 Parameter
• 实现 forward

不过，由于 minitorch 框架本身的一些局限性，部分实现需要用一点比较 trick 的方式，下面简要介绍。

Dropout

Dropout 的实现主要要注意两点：

• 推荐采用一般框架实现的训练时 scale，推理时 no-op 的方式。即：在训练模式下，在得到 dropout 之后结果后，进行一个 scale 计算：result / (1.0 - self.p_dropout)，然后在推理模式下直接返回输入 tensor 即可。
• 实现训练模式下 forward 时，根据文档推荐使用以下形式：

  mask = np.random.binomial(1, p_keep, size=x.shape)
  mask = tensor_from_numpy(mask, x.backend)

• 需要注意的是，调用 tensor_from_numpy 函数从 numpy array 构造 minitorch tensor 时，记得传入 backend 参数。

LayerNorm1d

这个实现公式本文就不贴了，就是个简单的：\(\frac{(x - mean)}{\sqrt{var}} * weight +bias\) ，但是需要注意的是，forward 实现过程中可能需要调用 mean/var 等接口时，传入 dim 参数最好不要是负数 -1，除非你的 CUDA kernel 处理了传入的 reduce_dim 是负数的情况（泪的教训）。

Embedding

Embedding Layer 从概念上来说非常简单：就是一个 Lookup Table，但是在基于 minitorch 框架实现 forward 时，还是不太直观的。理论上来说，我们应该：

• forward 时从 weights embeddings 中 select 对应 index 位置的 embedding，并 gather 成需要输出的 tensor shape。
• backward 时，将上游传下来的 grads token_id 相同的位置累加起来，对于没出现的 token_id 梯度设置为 0。

不过这一实现方式需要我们自行实现 Embedding 层的 fwd/bwd Function，但是这显然不是本次 homework 的重点。（实际上，PyTorch 的 Embedding 层 fwd/bwd 实现要更加复杂，尤其是 bwd，这里我们不深入讨论。）

对于本次 homework 而言，推荐的 Embedding forward 实现方式是：

• 根据输入 token ids tensor 构造一个 one-hot 矩阵，形状为 (bs * seq_len, vocab_size)；
• 然后与 weights 矩阵进行矩阵乘法，得到 shape 为 (bs * seq_len, n_embd) 的结果矩阵；
• 最后将其 view 为目标 shape：(bs, seq_len, n_embd)。

3. TransformerLayer

MultiHeadAttention

Attention 的实现相对复杂一些，但是已有的代码框架和 homework 文档已经非常详细，跟着一步一步实现也比较清晰。这里简单列几个注意点：

• shape 变换需要调用 Tensor 的 permute 方法。注意调用该方法后如果需要 view，要先调用一次 contiguous 然后再调用 view。
• Dropout 要应用在 Attention Score 矩阵上。

TransformerLayer

实现这一模块主要注意点是：

• forward 时注意添加适当的 view 改变 tensor 形状，因为 LayerNorm1D/Attention 等模块预期的输入 tensor 维度不一样，部分要求 2D 输入部分要求 3D。

4. generate

Problem 4 的 generate 函数，从概念上来讲非常简单：

• 将 token_ids 转换为 tensor 输入模型得到 logits
• 取最后一个 token 位置的 logits，得到其中最大值对应的 index，作为预测的 next token
• 循环上述过程

不过，在代码实现时，由于 minitorch 框架接口的限制，推荐首选把 logits 转换为 numpy array，然后愉快的调用 slice/argmax 等接口实现。

5. 总结

本次 homework 主要收获点：

• CrossEntropy Loss 数值稳定实现方式。
• Decoder-only Transformer 实现。