3 Transformer Language Model Architecture#

Token Embeddings#

第一步中，Transformer 将（batch 化的）token IDs 序列嵌入为一系列包含 token 身份信息（token identity）的向量序列（Figure 1 中的红色块）。 更具体地说，给定一个 token IDs 序列，Transformer 语言模型使用一个 Token Embedding 层来生成一系列向量。每个 Embedding 层接收一个 shape 为 (batch_size, sequence_length) 的整数张量，并生成一个形状为 (batch_size, sequence_length, d_model) 的向量序列。

Pre-norm Transformer Block#

在 embedding 之后，激活值由几个结构完全相同的神经网络层处理。

标准的仅解码器（decoder-only）型 Transformer 语言模型由 num_layers 个相同的层组成（通常称为 Transformer “blocks”）。每个 Transformer 块接收 shape 为 (batch_size, sequence_length, d_model) 的输入，并返回 shape 为 (batch_size, sequence_length, d_model) 的输出。

每个块都会在序列上聚合信息（通过自注意力）并对其进行非线性变换（通过前馈层）。经过 num_layers 个 Transformer 块之后，我们将取最终的激活值，并将其转换为词汇表上的分布。

我们将采用“预规范”型（Pre-norm）的 Transformer 块（详见第 3.4 节），此外，在最后一个 Transformer 块之后还需要使用层归一化（详情见下文），以确保其输出得到恰当的缩放。

经过这种规范处理后，我们将使用标准的学习型线性变换将 Transformer 块的输出转换为预测的下一个标记的对数概率。

Example (einstein_example1): Batched matrix multiplication with einops.einsum#

import torch
from einops import rearrange, einsum
## Basic implementation
Y = D @ A.T
# Hard to tell the input and output shapes and what they mean.
# What shapes can D and A have, and do any of these have unexpected behavior?

## Einsum is self-documenting and robust
#                          D                A     ->          Y
Y = einsum(D, A, "batch sequence d_in, d_out d_in -> batch sequence d_out")
## Or, a batched version where D can have any leading dimensions but A is constrained.
Y = einsum(D, A, "... d_in, d_out d_in -> ... d_out")

Example (einstein_example2): Broadcasted operations with einops.rearrange#

# 我们有一批图像，对于每一张图像，我们希望根据某个缩放系数生成 10 个 dimmed 的版本：

images = torch.randn(64, 128, 128, 3) # (batch, height, width, channel)
dim_by = torch.linspace(start=0.0, end=1.0, steps=10)

## Reshape and multiply
dim_value = rearrange(dim_by, "dim_value -> 1 dim_value 1 1 1")
images_rearr = rearrange(images, "b height width channel -> b 1 height width channel")
dimmed_images = images_rearr * dim_value

## Or in one go:
dimmed_images = einsum(
	images, dim_by,
	"batch height width channel, dim_value -> batch dim_value height width channel"
)

Example (einstein_example3): Pixel mixing with einops.rearrange#

'''
假设我们有一组图像，以形状为(batch, height, width,channel)的张量形式表示，并且我们想要对图像的所有像素进行线性变换，但这种变换应该针对每个通道分别独立进行。我们的线性变换表示为一个形状为(height * width, height * width)的矩阵 𝐵
'''

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

## 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)

# Instead, 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
)
# Or, if you’re feeling crazy: all in one go using einx.dot (einx 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
)
'''
这里的第一种实现方式可以通过在前后添加注释来说明输入和输出的形状，这样会更清晰易懂，但这种方式显得较为笨拙，还容易出现错误。而使用 einsum 表示法，文档信息就包含在实现代码中了！
'''

Einsum notation 能够处理任意的输入批处理维度，而且还有一个关键的优点，那就是它具有自解释性（self-documenting）。在使用 Einsum notation 编写的代码中，更清晰地能够看出输入和输出张量的相关形状。对于其余的张量，您可以考虑使用 Tensor 类型提示，例如使用 jaxtyping 库（并非专门针对 JAX）。

我们将在 assignment 2 中进一步讨论使用 einsum notation 时的性能影响，但目前要知道，它们几乎总是优于其他选择！

3.2.1 数学符号与内存顺序(Mathematical Notation and Memory Ordering)#

许多机器学习论文在其表述中使用行向量，这使得所得到的表示形式与 NumPy 和 PyTorch 默认采用的“行优先”内存顺序相契合。使用行向量时，线性变换的形式看起来是

其中 $𝑊∈ℝ^{𝑑_{out}×𝑑_{in}}$ ，行向量 $𝑥∈ℝ^{1×𝑑_{in}}$ .请注意，这使我们能够通过增加 $x$ 的最外层维度来对输入进行批量处理，这意味着我们可以用向量输入 $x$ 来替代矩阵输入$X ∈ ℝ^{batch×d_{in}}$.

在线性代数中，通常更常用的是使用列向量，此时线性变换的表现形式为

其中 $𝑊∈ℝ^{𝑑_{out}×𝑑_{in}}$，列向量 $𝑥∈ℝ^{𝑑_{in}}$ .在这种情况下对输入进行分批处理时，$x$ 的 batch 维度应放在最后位置，因此 $x$ 应替换为一个矩阵 $X ∈ ℝ^{d_{in}×batch}$.

在本次作业中，我们将主要使用列向量来进行数学符号表示，因为数学通常遵循这种表示方式。您需要记住，如果您想使用普通的矩阵乘法符号，那么您就必须像 Equation (1) 中的行向量惯例那样，将带有转置的矩阵进行相乘，因为 PyTorch 使用的是行优先的内存排列方式。如果您使用 einsum 进行线性代数运算，只要您正确标注轴，这通常就不会成为问题。顺便说一下，值得注意的是，像 Matlab、Julia 和 Fortran 这样的其他语言/线性代数包都使用列优先的内存排列方式，这意味着批处理维度在最后，但 Python 和相关包已经采用了 C 标准的行优先排列方式。

3.3.1 Parameter Initialization#

高效地训练神经网络通常需要对模型参数进行精心的初始化——不恰当的初始化可能会导致诸如梯度消失或爆炸等不良现象。预规范变换器 (Pre-norm transformers) 对初始化的适应性异常强，但仍会对训练速度和收敛性产生显著影响。鉴于本次作业已经较长，我们将把具体细节留到 assignment 3 中讨论，而在此先为您提供一些大致的初始化值，这些值在大多数情况下应该都能很好地发挥作用。

目前，请使用：

• Linear weights： $𝒩︀(𝜇=0, 𝜎^2=\frac{2}{𝑑_{in}+𝑑_{out}})$ truncated at $[−3𝜎, 3𝜎]$ .

• Embedding： $𝒩︀(𝜇=0, 𝜎^2=1)$ truncated at $[−3, 3]$

• RMSNorm： $\mathbb{1}$

您应当使用 torch.nn.init.trunc_normal_ 来初始化截断正态分布的权重。

3.3.2 Linear Module#

线性层是 Transformer 模型以及一般神经网络中的一个基本构建模块。首先，您将实现自己的 Linear 类，该类继承自 torch.nn.Module 并执行线性变换：

请注意，我们未加入偏差项（bias），这与大多数现代语言模型的做法一致。

def __init__(self, in_features, out_features, device=None, dtype=None):
	'''
    构建一个线性变换模块。此函数应接受以下参数:
    in_features: int  输入数据的维度
    out_features: int  输出数据的维度
    device: torch.device | None = None  用于存储参数的 device
    dtype: torch.dtype | None = None  参数的数据类型
    '''
def forward(self, x: torch.Tensor) -> torch.Tensor:
    '''对输入应用线性变换'''

3.3.3 Embedding Module#

如上所述，Transformer 的第一层是一个 embedding 层，它将整数 token IDs 转换为维度为 d_model 的向量空间。我们将实现一个自定义的 Embedding 类，该类继承自 torch.nn.Module（因此您不应使用 nn.Embedding）。forward 方法应通过使用一个 shape 为 (batch_size， sequence_length) 的 token IDs 的 torch.LongTensor（其中 token IDs 的值为整数）从一个 shape 为 (vocab_size， d_model) 的 embedding 矩阵中选取每个 token ID 的 embedding 向量。

def __init__(self, num_embeddings, embedding_dim, device=None, dtype=None):
    '''
        Construct an embedding module.
        参数列表:
            num_embeddings: int, vocabulary 大小
            embedding_dim: int, embedding 向量维度
            device: torch.device | None = None
            dtype: torch.dtype | None = None
    '''
    pass
def forward(self, token_ids: torch.Tensor) -> torch.Tensor:
    '''
    	查找给定 token ID 的 embedding 向量。
    '''
	pass

3.4.1 Root Mean Square Layer Normalization#

A. Vaswani等人的原始 Transformer 实现使用了层归一化来对激活值进行归一化。遵循H. Touvron等人的做法，我们将使用**均方根层归一化（RMSNorm, 公式 4）**进行层归一化。

给定一个激活值的向量$𝑎∈\mathbb{R}^{𝑑_{model}}$，RMSNorm 将按如下方式对每个激活值 $a_i$ 进行重新缩放：

其中，$RMS(a) = \sqrt{ \frac{1}{d_{model}} \sum^{d_{model}}_{i=1} a^2_i + 𝜀 }$，在这里，$g_i$是一个可学习的“增益”参数（总共有 d_model 个这样的参数），而 𝜀 是一个超参数，通常固定为 1e-5。

你应该将输入类型向上转换到 torch.float32，以防止在对输入进行平方运算时发生溢出。总体来说，你的 forward 方法应该如下所示：

in_dtype = x.dtype
x = x.to(torch.float32)
# Your code here performing RMSNorm
...
result = ...
# Return the result in the original dtype
return result.to(in_dtype)

def __init__(self, d_model: int, eps: float = 1e-5, device=None, dtype=None): 
	'''
		Construct the RMSNorm module.
		接收参数：
            d_model: int Hidden dimension of the model
            eps: float = 1e-5 Epsilon value for numerical stability
            device: torch.device | None = None Device to store the parameters on
            dtype: torch.dtype | None = None Data type of the parameters
	'''
	pass
def forward(self, x: torch.Tensor) -> torch.Tensor:
    '''
    	处理输入形状为 (batch_size, sequence_length, d_model) 的输入 x，并且返回一个相同形状的 Tensor
    '''
	pass

3.4.2 Position-Wise Feed-Forward Network（逐位置前馈网络）#

在原始的 Transformer 论文中，Transformer 的前馈网络由两个线性变换组成，并且在两者之间有一个 ReLU 激活函数（ReLU($x$) = max(0, $x$)）。在最初的架构中，内部前馈层的维度通常是输入维度的 4 倍。

然而，与这种原始设计相比，现代语言模型通常会引入两个主要的变化：使用另一种激活函数，并采用一种门控机制。具体来说，我们将实现 “SwiGLU” 激活函数，该函数在像 Llama 3 和 Qwen 2.5 这样的大型语言模型中被采用，它将 SiLU（通常称为 Swish）激活函数与一种称为门控线性单元（GLU）的门控机制相结合。我们还将省略线性层中有时使用的偏置项，这与大多数现代大型语言模型 PaLM 和 LLaMA 的做法一致。

SiLU 激活函数的定义为：

如 Figure 3 所示，SiLU 激活函数与 ReLU 激活函数类似，但在零点处是平滑的。

门控线性单元（GLUs）最初由 Y. N. 戴芬等定义为通过一个经过 sigmoid 函数处理的线性变换和另一个线性变换得到的元素级乘积：

其中 $\odot$ 表示 element-wise multiplicaton（逐元素乘法，也称 Hadamard 乘积）。门控线性单元被建议用于“通过为梯度提供一条线性路径来缓解深度架构中的梯度消失问题，同时保留非线性能力。”

将 SiLU(Swish) 和 GLU 结合在一起，我们就得到了 SwiGLU，我们将会把它用于我们的前馈网络：

其中$x \in \mathbb{R}^{d_{model}}$，$W_1, W_3 \in \mathbb{R}^{d_{ff} \times d_{model}}$，$W_2 \in \mathbb{R}^{d_{model} \times d_{ff}}$，并且按照常规，$d_{ff}=\frac{8}{3}d_{model}$. 对于具体的实现方式，将这个值四舍五入到接近 64 的倍数以提高硬件效率是可行的。

N. Shazeer 首先提出了将 SiLU/Swish 激活函数与 GLU 结合，并进行了实验，结果表明 SwiGLU 在语言建模任务中比诸如 ReLU 和 SiLU（不带门控）这样的基准模型表现更好。在本次作业的后续部分，你将对 SwiGLU 和 SiLU 进行比较。尽管我们已经提及了这些组件的一些启发式论点，但保持实证视角是很有好处的：Shazeer 论文中的一句著名引述是：

“We offer no explanation as to why these architectures seem to work; we attribute their success, as all else, to divine benevolence.”

”对于这些架构为何如此有效，我们无法解释；如同其他一切事物，我们将它们的成功归因于神圣的眷顾。“

3.4.3 Relative Positional Embeddings（相对位置编码）#

为了将位置信息注入模型中，我们将采用 “Rotary Position Embeddings” 方法，该方法常被称为 “RoPE”（旋转位置嵌入）。

对于位于第 $i$ 个位置的 query token $q^{(i)}=W_qx^{(i)} \in \mathbb{R}^d$ ，我们将应用一个成对旋转矩阵（pairwise rotation matrix） $R^i$ ，从而得到 $q'^{(i)}=R^iq^{(i)}=R^iW_qx^{(i)}$。

在此，矩阵 $R^i$ 会将一对对的 embedding 元素 $q^{(i)}_{2k-1:2k}$ 旋转成一个 2d 向量，旋转角度为 $\theta_{i,k} = \frac{i}{\Theta^{(2k-2)/d}}$ ，其中 $k \in \{1,...,d/2\}$，$\Theta$ 是常数。因此，我们可以认为矩阵 $R^i$ 是一个大小为 $d \times d$ 的块对角矩阵，块 $R^i_k$ 中有 $k \in \{1,...,\frac{d}{2}\}$ ，并且具有以下形式：

因此，我们就得到了完整的旋转矩阵（pairwise rotation matrix）

其中 0 代表 2×2 的零矩阵。虽然可以构建完整的 $d \times d$ 矩阵，但一个好的解决方案应当利用该矩阵的特性来更高效地实现转换。由于我们只关心给定序列中 tokens 的相对旋转，因此我们可以重复使用 $cos(\theta_{i,k})$ 和 $sin(\theta_{i,k})$ 计算出的值，在不同 layer 和不同 batch 中都可以使用。如果您想要对其进行优化，可以使用由所有层引用的单个 RoPE 模块，并且它可以在初始化时通过 self.register_buffer(persistent=False) 创建一个 2d 预计算的正弦和余弦值缓冲区，而不是使用 nn.Parameter（因为我们不想去记忆这些固定的余弦和正弦值）。然后，我们对 $k^{(j)}$ 所执行的完全相同的旋转过程与我们对 $q^{(i)}$ 所执行的过程完全相同，只是旋转角度为相应的 $R^j$。需要注意的是，这一层没有可学习的参数。

def __init__(self, theta: float, d_k: int, max_seq_len: int, device=None):
	'''
        构建RoPE模块, 并根据需要创建缓冲区
        theta: float,       RoPE 的 theta 值
        d_k: int,           query 向量和 key 向量的维度
        max_seq_len: int,   输入的最大序列长度
        device: torch.device | None = None Device to store the buffer on
    '''
    pass
def forward(self, x: torch.Tensor, token_positions: torch.Tensor) -> torch.Tensor:
	'''
        处理一个形状为 (..., seq_len, d_k) 的输入张量，并返回一个相同形状的张量。
        请注意，你应该能够处理具有任意数量的批量维度的 x。
        你应该假设 token 位置是一个形状为 (..., seq_len) 的 Tensor，用于指定 x 在序列维度上的标记位置。
    '''
	pass