CS336 第一章③ Transformer语言模型

一个语言模型将一批序列的整数token ID作为输入（即形状为(batch_size, sequence_length)的torch.Tensor），并输出一个（同样按批次的）在词表上的归一化概率分布（即形状为(batch_size, sequence_length, vocab_size)的torch.Tensor）。这个输出中的每个位置给出了对下一个token的预测分布。

在训练语言模型时，我们使用这些“下一个词”的预测，与真实的下一个token计算交叉熵损失。

在推理（生成文本）时，我们会取模型在最后一个时间步（也就是输入序列最后一个token的位置）输出的预测分布，基于它生成下一个token（例如取概率最大的token，或从分布中采样）。然后将新生成的token加到输入序列末尾，再重复这一过程。

在这次课程中将从零开始构建这个Transformer语言模型。我们会先给出模型的高层结构，然后再逐步介绍各个组成部分。

Transformer Language Model

在上节课我们编写的BPE Tokenizer最后能够输出一个token id序列，这个序列将作为transformer语言模型的输入。Transformer语言模型首先使用输入嵌入（input embedding）将token ID转换为稠密向量，然后将这些嵌入后的token依次传入num_layers个Transformer块，最后通过一个可学习的线性投影（称为“输出嵌入”或“LM head”）来生成对下一个token的预测logits。

logits就是模型在做softmax前的原始分数

接下来分节详细介绍各个部分。

Token Embeddings（Token 嵌入）

在一开始Transformer会将序列化的token id嵌入到向量序列中，并保留token的标识信息。这实际上和我们之前将字符串文本转化为token id是一样的作用，都是在转换形式的同时保留原始信息。

从程序的角度上来说，每个嵌入层接收一个形状为(batch_size, sequence_length)的整数张量，并输出一个形状为(batch_size, sequence_length, d_model)的向量序列。

预归一化（Pre-norm）Transformer块
在完成嵌入之后，激活值会被送入若干结构完全相同的神经网络层中处理。一个标准的只有解码器（decoder-only）的Transformer语言模型由num_layers个相同的层组成（通常称为Transformer块）。
每个Transformer块的输入形状为(batch_size, sequence_length, d_model)，输出的形状也是(batch_size, sequence_length, d_model)。
每个块都会在序列范围内聚合信息（通过自注意力机制self-attention），并对信息进行非线性变换（通过前馈网络feed-forward layers）。

Output Normalization & Embedding（输出归一化与输出嵌入）

在经过num_layers个Transformer块之后，我们会将最终的激活值转换为对整个词表的概率分布。由于我们实现的是“预归一化”（pre-norm）Transformer块，这要求在最后一个Transformer块之后再进行一次层归一化，以确保输出的尺度被正确调整。

在完成这一步归一化后，我们会使用一个标准的可学习线性变换，将Transformer块的输出转换为预测的下一个token的logits。

Transformer Block

从上图就可以看出，在整个Transformer中，我们会对许多“类batch”的输入执行相同的计算。以下是几个例子：
• batch的各元素：我们在每个batch元素上执行完全相同的Transformer前向计算。
• 序列长度：像RMSNorm和前馈网络这样的“逐位置（position-wise）”操作，会在序列的每个位置上以相同方式执行。
• 注意力头：在多头注意力（multi-head attention）中，注意力操作会在多个注意力头之间进行批处理。

因此，我们需要一种便捷的方式来对这些额外的批次维度进行操作，从而充分利用GPU，同时代码也容易阅读和理解。许多PyTorch操作可以接受张量开头的额外“类batch维度”，并在这些维度上自动重复或广播运算。

举个例子：

假设我们正在进行一个“逐位置的、带batch的”操作。

我们有一个形状为(batch_size, sequence_length, d_model)的“数据张量”D，并希望将它与一个形状为(d_model, d_model)的矩阵A进行批量向量-矩阵乘法。

在这种情况下，D @ A 会执行一个批量矩阵乘法，而(batch_size, sequence_length)会被视为批次维度并被自动处理。这是PyTorch中一个非常高效的底层操作。

正因为如此，你的函数通常应该假设输入张量可能带有额外的“类batch”维度，并且应该保持这些维度位于PyTorch张量形状的最前面。

为了将张量整理成适合这种批处理的格式，通常需要多次使用view、reshape和transpose来调整形状。这可能会变得麻烦，而且让代码不容易看懂，也不容易知道每一步的张量形状到底是什么。

一个更方便的方式是使用torch.einsum中的einsum记法，或者进一步使用与框架无关的库，例如einops或einx。这类库的两个核心操作是：

• einsum：用于在任意维度上进行张量收缩（tensor contraction）。

• rearrange：用于重新排序、拼接、或切分任意维度。

事实证明，机器学习中几乎所有操作都可以看作：维度操作 + 张量收缩 + 偶尔的逐点非线性函数。

这意味着如果使用einsum记法，你的代码往往能变得更清晰、更灵活，强烈推荐在这门课中学习并使用einsum记法。

• 对从未接触过einsum的朋友，建议先用einops。

• 对已经熟悉einops的朋友，则建议进一步学习更通用的einx。
  这两个库都已预装在官方提供的环境中。

我知道你读到这里已经很懵了，没事其实我也很懵。其实上面讲了这么一大坨，你只要理解transformer的结构（token embedding - 若干个transformer blocks- normalize - output embedding(linear) - softmax)和各层的作用即可。

至于上面提到的einops和einx，下面是我自己的理解

Einops & Einx

虽然Transformer里面存在大量不明所以的操作，但归根结底都是矩阵操作。但是人类并不是肉编器，有时候无法想象出高维张量的维度操作过程。而Einops和Einx就是用形象化的操作来拯救人类的。

直接看实际操作的例子，你会对这个操作印象更深：

import torch
from einops import einsum
# C = A @ B 
torch.einsum("ij, jk->ik", A, B)

# multi-head Q·KT calculate
# assume Q and K shape like (batch, heads, seq, dim)
scores = torch.einsum("bhqd, bhkd -> bhqk", q, k)

相信你现在已经了解einops的基本用法，我们不需要用reshape transpose等复杂的维度操作，而是直接用清晰直观的矩阵表达式来描述我们的要求。例如第一个例子的A、B矩阵相乘，直接用线性代数课上学过的表达式ij, jk->ik来表示。第二个例子的多头注意力$Q·K^T$直接用bhqd, bhkd -> bhqk表示了在batch和head没有变化的情况下进行seq和dim的相乘。

通过使用einops，我们能对张量做可读的自然语言化描述。下面会给出一些einsum记法的使用示例。这些示例是对einops文档的补充，应当先阅读它们的官方文档再继续下面的阅读。

矩阵乘

import torch
from einops import rearrange, einsum

# Basic implementation
# 普通矩阵乘法，难以看出输入输出的形状
Y = D @ A.T

# Einsum 更清晰、健壮
# 明确 D 和 A 的维度对应关系
# D A -> Y
Y = einsum(D, A, "batch sequence d_in, d_out d_in -> batch sequence d_out")

# 更通用的 batched 版本
# D 可以有任意前导 batch-like 维度（用...表示），但 A 维度固定
Y = einsum(D, A, "... d_in, d_out d_in -> ... d_out")

维度变换

import torch
from einops import rearrange, einsum

# 假设有一个 batch 的图片
images = torch.randn(64, 128, 128, 3)  # (batch, height, width, channel)

# 生成 10 个不同亮度/缩放因子
dim_by = torch.linspace(start=0.0, end=1.0, steps=10)  # dimming factor

# 方法 1：使用 rearrange 分别扩展维度再相乘
# 扩展 dim_by 的维度以匹配 images
dim_value = rearrange(dim_by, "dim_value -> 1 dim_value 1 1 1")
# 扩展 images 的维度以匹配 dim_value
images_rearr = rearrange(images, "b height width channel -> b 1 height width channel")
# 相乘得到 dimmed images
dimmed_images = images_rearr * dim_value
# dimmed_images.shape -> (64, 10, 128, 128, 3)

# 方法 2：使用 einsum 一步到位
dimmed_images = einsum(
    images, dim_by,
    "batch height width channel, dim_value -> batch dim_value height width channel"
)
# dimmed_images.shape -> (64, 10, 128, 128, 3)

Einsum 的动作规则：

• 所有在两个输入中都出现的维度 → 会发生点乘或广播

• 只出现一次的维度 → 被自动广播扩展

• 输出中出现的维度 → 保留

• 输出中不出现的维度 → 会被求和减少 (sum-reduction)

广播就是让两个张量形状自动对齐的机制，使它们能进行逐元素的乘法（或加减等）运算，而无需手动扩展维度或复制数据

图片像素混合

# Suppose we have a batch of images represented as a tensor of shape
# (batch, height, width, channel), and we want to perform a linear
# transformation across all pixels of the image, but this transformation
# should happen independently for each channel. Our linear transformation is
# represented as a matrix B of shape (height × width, height × width).

channels_last = torch.randn(64, 32, 32, 3)   # (batch, height, width, channel)
B = torch.randn(32 * 32, 32 * 32)

# 1. Manual way using view + transpose

# Rearrange an image tensor for mixing across all pixels
channels_last_flat = channels_last.view(
    -1,
    channels_last.size(1) * channels_last.size(2),
    channels_last.size(3)
)

channels_first_flat = channels_last_flat.transpose(1, 2)
channels_first_flat_transformed = channels_first_flat @ B.T
channels_last_flat_transformed = channels_first_flat_transformed.transpose(1, 2)

channels_last_transformed = channels_last_flat_transformed.view(
    *channels_last.shape
)

# 2. Using einops

height = width = 32

# Rearrange replaces clunky torch view + transpose
channels_first = rearrange(
    channels_last,
    "batch height width channel -> batch channel (height width)"
)

channels_first_transformed = einsum(
    channels_first,
    B,
    "batch channel pixel_in, pixel_out pixel_in -> batch channel pixel_out"
)

channels_last_transformed = rearrange(
    channels_first_transformed,
    "batch channel (height width) -> batch height width channel",
    height=height,
    width=width
)

# 3. One-liner using einx.dot (equivalent of einops.einsum)

height = width = 32

channels_last_transformed = einx.dot(
    "batch row_in col_in channel, (row_out col_out) (row_in col_in)"
    "-> batch row_out col_out channel",
    channels_last,
    B,
    col_in=width,
    col_out=width
)

# Summary:
# The first implementation here could be improved by placing comments
# before and after to indicate what the input and output shapes are,
# but this is clunky and susceptible to bugs.
# With einsum notation, documentation is implementation!

Einsum表示法可以处理任意输入批处理维度，同时也具有自说明文件（功能一目了然）的关键优势。在使用einsum符号的代码中，输入和输出张量的相关形状要清楚得多。对于剩下的张量，可以考虑使用Tensor类型提示，例如使用jaxtyping库(不是Jax所特有的)。我们将更多地讨论赋值2中使用einsum符号的性能影响，但现在只需要知道它们几乎总是优于其他方案。

Mathematical Notation and Memory Ordering（数学表示与内存顺序）

很多的机器学习论文都用行向量作为主要的计算向量，这样运算的表达就像下面这样：

$$y=xW^T$$

这也是NumPy和PyTorch默认使用的计算顺序。

但是一般的线性代数表达通常使用列向量，这样线性变换运算的表达是这样：

$$y = Wx$$

在这次课程中我们使用列向量表示法，因为这样更容易贴近数学上的表示方法。要切记在PyTorch中是使用行向量计算的，不过如果你用了einsum进行矩阵计算的话这些都不是问题。

Experiment：Linear Module & Embedding Module

说了那么多，现在我们就来自己实现一个Transformer吧。首先是最基础的两个模块：线性运算模块（Linear Module）和嵌入模块（Embedding Module）。

参数初始化

高效地训练神经网络通常需要认真地初始化模型参数，一个水平较差的初始化容易导致模型表现不可接受，如发生梯度消失或爆炸问题。预归一化转换器对初始化非常耐操，但仍然会对训练速度和收敛性产生重大影响。由于这次课程内容已经很长，我们将这部分的处理细节等到assignment 3再另外展示，而现在你只需要使用以下的初始化参数就可以了，这些参数在大部分情况下表现都很好：

• Linear weights：$\mathcal{N} \left(\mu = 0, \sigma^2 = \frac{2}{d_{in} + d_{out}} \right) \ truncated\ at [-3\sigma, 3\sigma]$

• Embedding：$\mathcal{N} \left(\mu = 0, \sigma^2 = 1 \right) \ truncated\ at [-3, 3]$

• RMSNorm：1

要使用这些参数，需要使用torch.nn.init.trunc_normal_初始化截断正态分布的权重。

线性运算模块（Linear Module）

线性运算层通常是transformer和神经网络的基本组成单元。首先，你将实现自己的继承于torch.nn的Linear类，能够进行线性变换 $y = Wx$

跟大多数现代LLM一样，我们不使用偏置量bias。

偏置量bias的作用就是在进行线性运算的时候偏移基准，使模型能拟合非零均值的数据，提高表达自由度，同时使大部分激活函数能够被激活。从数学角度上来表示就是$y=Wx + b$。但是现代的LLM都可以设置bias = False，因为transformer的LayerNorm（对每个样本的特征维度的归一化层）已经提供平移项

接下来是实际的实验要求：

• 继承nn.Module

• 记得super().__init__()

• 以$W$（而不是$W^T$）的形式创建和存储你的参数，最后用nn.Parameter的形式保存

• 不要使用nn.Linear或者nn.functional.linear

你可能需要使用Module.load_state_dict来加载权重，并使用torch.nn.init.trunc_normal_初始化权重。

最后，在[adapters.run_linear]中实现你的功能，并用以下的语句测试：

uv run pytest -k test_linear

老样子，我们还是将这部分的实现独立出来，创建文件Transformer/Linear.py，然后填写模板：

import torch
from torch import nn


class Linear(nn.Module):
    def __init__(
            self,
            in_features: int,
            out_features: int,
            device: torch.device | None = None,
            dtype: torch.dtype | None = None,
    ):
        """
        Construct a linear transformation module.

        Args:
            in_features : int
                The dimensionality of the input feature vector.

            out_features : int
                The dimensionality of the output feature vector.

            device : torch.device or None, optional
                Device on which to allocate parameters (e.g., 'cuda', 'cpu').

            dtype : torch.dtype or None, optional
                Data type used to store the parameters (e.g., torch.float32).
        """
        super().__init__()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Apply the linear transformation to the input.
        """

要将维度为in_features的张量$x$转换为out_features维度的张量$y$输出，实际上就是要进行以下的运算：

$$y = Wx$$

在einops的写法就应该是下面这样：

y = einsum(x, self.weight, "... d_in, d_out d_in -> ... d_out")

那么其实我们要做的就是构造这个$W$，也就是self.weight。按照上文所说，这个权重的初始化需要截断正态分布赋值。那么按照上文的公式就是这样的：

self.weight = nn.Parameter(torch.empty(out_features, in_features, **kwargs))
std = math.sqrt(2 / (self.in_features + self.out_features))
nn.init.trunc_normal_(self.weight,
                      mean=0,
                      std=std,
                      a=-3 * std,
                      b=3 * std)

尤其注意一开始nn.Parameter的声明张量维度顺序，是out在前，in在后，这和我们的公式计算顺序是息息相关的，这样结果的y才会以out作为其中一维的维数，同时才能在张量乘的时候与维度为in的x正确计算。

其他的就没什么好说的了，完整的Linear.py如下：

import math

import torch
from torch import nn
from einops import einsum


class Linear(nn.Module):
    def __init__(
            self,
            in_features: int,
            out_features: int,
            device: torch.device | None = None,
            dtype: torch.dtype | None = None,
    ):
        """
        Construct a linear transformation module.

        Args:
            in_features : int
                The dimensionality of the input feature vector.

            out_features : int
                The dimensionality of the output feature vector.

            device : torch.device or None, optional
                Device on which to allocate parameters (e.g., 'cuda', 'cpu').

            dtype : torch.dtype or None, optional
                Data type used to store the parameters (e.g., torch.float32).
        """
        super().__init__()

        self.in_features = in_features
        self.out_features = out_features
        kwargs = {"device": device, "dtype": dtype}
        self.weight = nn.Parameter(torch.empty(out_features, in_features, **kwargs))
        std = math.sqrt(2 / (self.in_features + self.out_features))
        nn.init.trunc_normal_(self.weight,
                              mean=0,
                              std=std,
                              a=-3 * std,
                              b=3 * std)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Apply the linear transformation to the input.
        """
        # y = Wx
        y = einsum(x, self.weight, "... d_in, d_out d_in -> ... d_out")
        # y = self.weight @ x
        return y

最后，我们在adapters.py内的run_linear()函数调用我们写好的类就可以进行测试了

def run_linear(
    d_in: int,
    d_out: int,
    weights: Float[Tensor, " d_out d_in"],
    in_features: Float[Tensor, " ... d_in"],
) -> Float[Tensor, " ... d_out"]:
    """
    Given the weights of a Linear layer, compute the transformation of a batched input.

    Args: 
        in_dim (int): The size of the input dimension
        out_dim (int): The size of the output dimension
        weights (Float[Tensor, "d_out d_in"]): The linear weights to use
        in_features (Float[Tensor, "... d_in"]): The output tensor to apply the function to

    Returns:
        Float[Tensor, "... d_out"]: The transformed output of your linear module.
    """
    module = Linear(d_in, d_out)
    module.load_state_dict({'weight' : weights})
    return module.forward(in_features)

如果你理清代码逻辑的话可能会觉得很奇怪，我们在Linear.__init__()里面对weight进行了初始化并以截断正态分布赋值，但是在调用的时候还使用了module.load_state_dict({'weight' : weights})来强行对weight重新赋值，那前面的操作不都是做无用功？

其实这是课程为了能够进行测试而特意设计的，测试数据已经进行符合截断正态分布的赋值，不需要我们对其再做一次。课程的核心教学目标是让我们掌握einops的使用方法，而不是处理数据。

当然，在正常使用时，如果我们需要能够在类内部对权重进行相关处理，可以在初始化的时候传入原始权重，再进行处理即可。

测试结果如下，其实只有一个点要测，如果不通过的话大概率是权重重复计算或者是矩阵相乘时的位置关系没有正确处理：

语义嵌入模块（Embedding Module）

上文提到语义嵌入模块位于transformer入口第一层，作用是将token ids映射到维度为d_model的向量空间中。接下来我们要实现一个Embedding类，需要继承自 torch.nn.Module（因此nn.Embedding被ban了）。
forward 方法应当通过索引一个形状为(vocab_size, d_model) 的嵌入矩阵，来为每个token ID选取对应的嵌入向量。输入的token ID是一个torch.LongTensor，形状为(batch_size, sequence_length)。

同样地，使用torch.nn.init.trunc_normal_初始化权重。

最后，在[adapters.run_embedding]中实现你的功能，并用以下的语句测试：

uv run pytest -k test_embedding

创建文件Transformer/Embedding.py，填写模板：

import torch

class Embedding(torch.nn.Module):
    def __init__(self, num_embeddings, embedding_dim, device=None, dtype=None):
        """
        Construct an embedding module.
    
        Args:
            num_embeddings (int):
                Size of the vocabulary.

            embedding_dim (int):
                Dimension of the embedding vectors, i.e., d_model.
  
            device (torch.device | None, optional):
                Device to store the parameters on.

            dtype (torch.dtype | None, optional):
                Data type of the parameters.
        """
        super().__init__()
    
    
    def forward(self, token_ids: torch.Tensor) -> torch.Tensor:
        """
        Lookup the embedding vectors for the given token IDs.
        """

然后把上文给的权重参数放进代码里，就像这样：

def __init__(self, num_embeddings, embedding_dim, device=None, dtype=None):
        """
        Construct an embedding module.

        Args:
            num_embeddings (int):
                Size of the vocabulary.

            embedding_dim (int):
                Dimension of the embedding vectors, i.e., d_model.

            device (torch.device | None, optional):
                Device to store the parameters on.

            dtype (torch.dtype | None, optional):
                Data type of the parameters.
        """
        super().__init__()
        self.num_embeddings = num_embeddings
        self.embedding_dim = embedding_dim
        kwargs = {'device': device, 'dtype': dtype}
        self.weight = torch.nn.Parameter(torch.empty(num_embeddings, embedding_dim, **kwargs))
        nn.init.trunc_normal_(self.weight,mean=0,std=1,a=-3,b=3)

现在做的和Linear类初始化过程的思路是一致的，都在初始化时设定了权重的大小。

然后在forward()函数中实际上我们只需要这样实现就可以了：

def forward(self, token_ids: torch.Tensor) -> torch.Tensor:
        """
        Lookup the embedding vectors for the given token IDs.
        """
        return self.weight[token_ids]

What??为什么这样一句话就能实现语义嵌入功能了？这就要回到嵌入的原理上了。

一言而概之，语义嵌入实际上就是查表。

也就是说，self.weight其实就是一个大小为 num_embeddings × embedding_dim的可学习查找表，而对于每个token ID，embedding的工作就是：在一个矩阵中找到这个 token 对应的一行，并把它作为语义向量返回。

为什么这就完成了语义嵌入？因为Embedding的本质就是**把离散 ID → 映射成连续的向量空间，**这样在训练时，反向传播会自动更新self.weight[id]对应的那一行，从而学习语义信息。

而这一句话会让PyTorch 做多维索引，把token_ids中的每一个整数ID当成索引到weight的第 0 维（行）里 取出对应行的embedding vector。

效果是：

token_ids: (batch_size, seq_len)
↓
self.weight[token_ids]: (batch_size, seq_len, embedding_dim)

举个具体例子：

token_ids = tensor([[5, 3, 8],
                    [2, 7, 1]])   # shape: (2, 3)

"""
self.weight 的第 5 行是词 5 的向量
第 3 行是词 3 的向量
第 8 行是词 8 的向量
...
"""
# 返回的 tensor 就是：
[
  [ embedding(5), embedding(3), embedding(8) ],
  [ embedding(2), embedding(7), embedding(1) ]
]

最后，我们在adapters.py内的run_embedding()函数调用我们写好的类就可以进行测试了

def run_embedding(
    vocab_size: int,
    d_model: int,
    weights: Float[Tensor, " vocab_size d_model"],
    token_ids: Int[Tensor, " ..."],
) -> Float[Tensor, " ... d_model"]:
    """
    Given the weights of an Embedding layer, get the embeddings for a batch of token ids.

    Args:
        vocab_size (int): The number of embeddings in the vocabulary
        d_model (int): The size of the embedding dimension
        weights (Float[Tensor, "vocab_size d_model"]): The embedding vectors to fetch from
        token_ids (Int[Tensor, "..."]): The set of token ids to fetch from the Embedding layer

    Returns:
        Float[Tensor, "... d_model"]: Batch of embeddings returned by your Embedding layer.
    """
    module = Embedding(vocab_size, d_model)
    module.load_state_dict({'weight': weights})
    return module.forward(token_ids)

测试结果如下，同样只有一个点要测