Scaled Dot-Product Attention

终于到了注意力了，伟大的assignment 1行程已经走过了一大半。

我们现在将实现缩放点积注意力（scaled dot-product attention）。作为一个前置步骤，Attention操作的定义将使用softmax。

softmax是一种操作，它接收一个未归一化的向量，并将其转换为一个归一化的分布：

$softmax(v)_i = \frac{\exp(v_i)}{\sum ^n_{j=1}\exp(v_j)}$

请注意，当 $v_i$ 很大的时候， $\exp(v_i)$ 可能会变成 inf（此时 inf/inf = NaN）。

可以通过“softmax对于给所有输入加上同一个常数c是不变的”的性质，从而避免这个问题。

可以利用这一性质来提高数值稳定性——通常我们会从向量v的所有元素中减去它的最大值，使得新的最大值变为0。

Experiment：softmax

编写一个函数，对张量应用softmax操作。你的函数应接受两个参数：一个张量和一个维度i，并在该张量的第i个维度上应用softmax。输出张量应与输入张量具有相同的形状，但它的第i个维度现在应是一个归一化的概率分布。

为了避免数值稳定性问题，需要使用这样一个技巧：在计算softmax之前，先从第i个维度的所有元素中减去该维度上的最大值。

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

1	uv run pytest -k test_softmax

直接给代码：

import torch

def softmax(x, dim):
    x_max = torch.max(x, dim = dim, keepdim = True)[0]
    x_exp = torch.exp(x - x_max)
    x_sum = torch.sum(x_exp, dim = dim, keepdim = True)
    return x_exp / x_sum

imagepng

现在就可以用数学公式描述注意力运算了：

$Attention(Q,K,V)=softmax\left(\frac{Q^TK}{\sqrt{d_k}}\right)V, \\ Q\in \mathbb{R}^{n\times d_k},K\in \mathbb{R}^{m\times d_k}, V \in \mathbb{R}^{m\times d_v}$

这里的Q、K、V都不是可训练的参数，如果你对为什么不用 $QK^T$ 有疑问，不妨看看上面的Mathematical Notation and Memory Ordering（数学表示与内存顺序）那一章。

Masking（掩码）：
有时对注意力操作的输出进行掩码是很方便的。一个掩码应具有形状 $M \in \{ \text{True}, \text{False} \}^{n \times m}$ ，这个布尔矩阵的每一行 i 表示查询 i 应该关注哪些键。

按照惯例（并且有点容易混淆），在位置 (i, j) 上：

值为 True 表示查询 i 会关注键 j；
值为 False 表示查询 i 不会关注键 j。

换句话说，值为 True 的 (i, j) 对应的位置是“信息流动”的位置。

例如，考虑一个形状为 $1 \times 3$ 的掩码矩阵 $[[ \text{True}, \text{True}, \text{False} ]]$ ，则这个单一的查询向量只会关注前两个键。

从计算的角度来看，使用掩码比对子序列分别计算注意力要高效得多。

我们可以通过在 softmax 之前的得分矩阵 $\left(\frac{Q^TK}{\sqrt{d_k}}\right)$ 中，把掩码矩阵中为 False 的位置加上一个 $-\infty$ ，来实现这一点。

Experiment : Scaled dot-product attention

请实现scaled dot-product attention函数。你的实现需要能够处理如下形状的输入：

keys 和 queries 的形状为 (batch_size, ..., seq_len, d_k)
values 的形状为 (batch_size, ..., seq_len, d_v)
这里的 ... 表示任意数量的、类似 batch 的额外维度（如果存在的话）。

该函数的输出应具有形状 (batch_size, ..., d_v)。

你的实现还应支持一个可选的、用户提供的布尔型掩码（mask），其形状为 (seq_len, seq_len)。
对于 mask 中值为True的位置，其对应的注意力概率应整体归一化为 1；
对于 mask 中值为False的位置，其对应的注意力概率应为 0。最后，在adapters.run_scaled_dot_product_attention中实现你的功能。

使用

1	uv run pytest -k test_scaled_dot_product_attention

会在三阶（3D）输入张量上测试你的实现；

使用

1	uv run pytest -k test_4d_scaled_dot_product_attention

会在四阶（4D）输入张量上测试你的实现

直接上代码：

from unittest import result

from jaxtyping import Float, Bool
from torch import  Tensor
from einops import einsum
from cs336_basics.Transformer.Softmax import softmax
import math

def get_scaled_dot_product_attention(query: Float[Tensor, " ... queries d_k"], key: Float[Tensor, " ... keys d_k"],
                 value: Float[Tensor, " ... values d_v"], mask: Bool[Tensor, " ... queries keys"] | None):
    q_k = einsum(query, key, "... queries d_k, ... keys d_k -> ... queries keys")
    d_k = query.shape[-1]
    score_matrix = q_k / math.sqrt(d_k)
    if mask is not None:
        score_matrix = score_matrix.masked_fill(~mask, float("-inf"))
    return einsum(softmax(score_matrix, dim = -1), value, "... queries keys, ... keys d_v -> ... queries d_v")

这个函数就是对公式的计算过程的还原。其中有几个要点，比如对mask的处理使用了masked_fill()函数，这个函数会根据输入的mask矩阵为True的位置，对相同形状的目标矩阵赋值。例如这里先将mask取反，这样会让里面为False的位置赋值-inf。

然后要注意使用einsum时对维度注释的要求。显然参数中queries和values的大小是相同的，但是在表达式中是不能把queries和values分开写的。就像最后return语句那样，为了让函数工作正常，必须把value的维度从原来的... values d_v改为相同维度的... keys d_v，这样才能让函数正常进行矩阵乘法。

最后考虑到使用einops会造成性能损失，这里顺便给出不用它的写法：

# q_k = einsum(query, key, "... queries d_k, ... keys d_k -> ... queries keys")
q_k = query @ key.transpose(-1, -2)
# return einsum(softmax(score_matrix, dim = -1), value, "... queries keys, ... keys d_v -> ... queries d_v")
return softmax(score_matrix, dim=-1) @ value

测试结果如下，对不同维度的张量都有效：

imagepng

Causal Multi-Head Self-Attention

因果多头自注意力（Causal Multi-Head Self-Attention）

我们将实现多头自注意力部分，多头自注意力能够从多方面对数据进行解读，效果比单一注意力更全面。从数学上讲，应用多头注意力的运算公式如下：

$\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \dots, \text{head}_h)$

其中 $\text{head}_i = \text{Attention}(Q_i, K_i, V_i)$

这里 $(Q_i, K_i, V_i)$ 分别是 Q、K、V 的第 $(i \in {1, \dots, h})$ 个切片，切片大小分别为 $(d_k)$ 或 $(d_v)$ 的嵌入维度。Attention就是上一小节的scaled dot-product attention操作。

由此，我们可以得到多头自注意力的运算公式：

$\text{MultiHeadSelfAttention}(x) = W_O \text{MultiHead}(W_Q x, W_K x, W_V x)$

这里的可学习参数为：

$W_Q \in \mathbb{R}^{h d_k \times d_\text{model}}, \quad W_K \in \mathbb{R}^{h d_k \times d_\text{model}}, \quad W_V \in \mathbb{R}^{h d_v \times d_\text{model}}, \quad W_O \in \mathbb{R}^{d_\text{model} \times h d_v}$

由于在多头注意力操作中 Q、K、V 会被切片，我们可以认为 $(W_Q, W_K, W_V)$ 沿输出维度为每个 head 分开处理。当这一过程实现后，你应该通过三个矩阵乘法就能计算出 key、value 和 query 的投影。

进阶要求：可以尝试将 key、query、value 的投影合并到一个单一的权重矩阵中，这样就只需要一次矩阵乘法即可完成投影操作

因果掩码（Causal Masking）

在实践中，需要防止模型访问序列中的未来 token。换句话说，如果模型得到一个 token 序列 $(t_1, \dots, t_n)$ ，并希望计算前缀 $(t_1, \dots, t_i)$ （其中 (i < n)）的下一词预测时，模型不能访问位置 $(t_{i+1}, \dots, t_n)$ 的 token 表示。原因是，在推理阶段生成文本时，这些未来token并不可见，同时访问它们会泄露真实下一词的信息，使语言建模预训练任务变得过于简单。

对于输入序列 $(t_1, \dots, t_n)$ ，一种直接方法是对每个前缀单独运行多头自注意力 n 次。但更高效的方法是使用因果注意力掩码（causal attention masking）：

对于 token (i)，它只能 attend 到所有 $(j \le i)$ 的位置。
可以使用torch.triu或广播索引比较来构造这个掩码。
你已经在缩放点积注意力实现中支持了attention mask，因此可以直接利用这一点。

RoPE（旋转位置编码）的应用

RoPE 应用于query和key向量，但不应用于 value 向量。
在多头注意力中，每个 head 是独立应用注意力的，因此head维度应当作为批量维度处理。
也就是说，对于每个 head，query 和 key 都要应用完全相同的 RoPE 旋转，以保证注意力计算的一致性。

Experiment: causal multi-head self-attention

接下来的实验可能是到现在为止最难的一次，我们要实现因果多头自注意力机制，并将其作为一个torch.nn.Module。

模块至少应该包含以下参数：

d_model int：Transformer块输入的维数
num_heads int：用于多头自注意力的头的数量。

此外，还需要设置 $d_k=d_v=d_{model}/h$ 。

在adapters.run_multihead_self_attention中完成模块的编写，并且使用下面的语句进行测试：

1	uv run pytest -k test_multihead_self_attention

如果你已经运行了测试样例就会发现，这个测试单元实际上包括两个小样例，分别是带RoPE和不带RoPE的。为了提高模块实用性，我们需要在模块内处理这两种情况。

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model: int, num_heads: int, rope: bool = False, theta: float | None = None, max_seq_len: int | None = None):
        super().__init__()
        self.num_heads = num_heads
        self.d_model = d_model
        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
        self.rope = rope

        if rope:
            # rope=True → 必须提供 theta 和 max_seq_len
            if theta is None or max_seq_len is None:
                raise ValueError("When using rope, theta and max_seq_len must be provided.")

            self.theta = theta
            self.max_seq_len = max_seq_len

            # 初始化ROPE
            self.roper = RotaryPositionalEmbedding(theta, d_model // num_heads, max_seq_len)
        else:
            self.theta = None
            self.max_seq_len = None

由于输入维度会被多个头划分，因此需要确保输入维度能被头数量整除。后面的self.roper也是一样的道理，由于RoPE是分别应用于各个头上的，因此输入维度是每个头的维度d_model // num_heads。

接下来先实现第一条公式 $\text{head}_i = \text{Attention}(Q_i, K_i, V_i)$

def _split_heads(self, query, key, value, token_positions : Int[Tensor, " ... sequence_length"] | None = None):
        batch, seq_len, d_k = query.shape
        head_dim = d_k // self.num_heads

        # 把最后的 d_k 切成 num_heads 个 head_dim
        # (batch, seq_len, D) → (batch, seq_len, H, head_dim)
        query = query.view(batch, seq_len, self.num_heads, head_dim)
        key = key.view(batch, seq_len, self.num_heads, head_dim)
        value = value.view(batch, seq_len, self.num_heads, head_dim)

_split_heads这个函数用于按头数量切分输入的Q、K、V张量。这里使用torch.view完成这个操作。

按照步骤，接下来需要对切分出来的每个头进行可选的RoPE，然后进行attention的计算。我们已经实现的attention支持batch，但是由于我们使用view扩展了Q、K、V张量维度，目前他们的维度为(batch, seq_len, H, head_dim),而接下来的attention操作入参为query: Float[Tensor, " ... queries d_k"], key: Float[Tensor, " ... keys d_k"], value: Float[Tensor, " ... values d_v"]。为了和attention函数的入参契合，我们需要对维度进行调换。

# 交换 seq_len 和 head 的位置，保证最后两个维度 = (seq_len, head_dim)
# (batch, seq_len, H, head_dim) → (batch, H, seq_len, head_dim)
query = query.transpose(1, 2)
key = key.transpose(1, 2)
value = value.transpose(1, 2)

这样我们就可以不需要手动批量进行attention操作了。

if self.rope:
            # 添加旋转位置编码
            query = self.roper(query, token_positions)
            key = self.roper(key, token_positions)

        # causal mask, shape (seq_len,seq_len)
        causal_mask = ~torch.triu(torch.ones(seq_len, seq_len, dtype=torch.bool, device = query.device), diagonal=1)
        heads = get_scaled_dot_product_attention(query, key, value, mask=causal_mask)
        return heads

最后为每个头加上RoPE和attention处理。这里的causal_mask就是因果注意力掩码矩阵，使用torch.triu配合diagonal = 1，能够实现一个不包含主对角线的上三角矩阵。

接下来实现第二条公式 $\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \dots, \text{head}_h)$ 。

def _multi_head(self, query, key, value, token_positions : Int[Tensor, " ... sequence_length"] | None = None):
    heads = self._split_heads(query, key, value, token_positions)
    # (B,H,T,d) → (B,T,H*d)
    batch, head, seq_len, head_dim = heads.shape
    multi = heads.transpose(1, 2).reshape(batch, seq_len, head * head_dim)
    return multi

_multi_head函数就是对之前切分头时进行的维度操作进行逆向操作。

最后来实现第三条公式 $\text{MultiHeadSelfAttention}(x) = W_O \text{MultiHead}(W_Q x, W_K x, W_V x)$

如果无脑一点的话，通过以下的操作就能很简单地获得三个权重的投影。

1
2
3

q_x = einsum(w_query, x, "d_k d_in, ... seq_len d_in -> ... seq_len d_k")
k_x = einsum(w_key, x, "d_k d_in, ... seq_len d_in -> ... seq_len d_k")
v_x = einsum(w_value, x, "d_v d_in, ... seq_len d_in -> ... seq_len d_v")

但是，重复的mutmal操作无疑会对运行效率造成负担。因此，可以参考扩展要求中提到的方法，将三个权重合并到一个矩阵中，进行计算后再分开，这样只需要进行一次矩阵乘法即可。

def forward(self, w_query, w_key, w_value, w_out, x, token_positions : Int[Tensor, " ... sequence_length"] | None = None):
    """
    q_proj_weight: Float[Tensor, " d_k d_in"],
    k_proj_weight: Float[Tensor, " d_k d_in"],
    v_proj_weight: Float[Tensor, " d_v d_in"],
     o_proj_weight: Float[Tensor, " d_model d_v"],
    in_features: Float[Tensor, " ... sequence_length d_in"],
    """
    assert w_query.size(1) == w_key.size(1) == w_value.size(1) == x.size(-1), "d_in of w_query, w_key, w_value and x are not same"
    w_qkv = torch.cat([w_query, w_key, w_value], dim = 0)
    qkv_x = x @ w_qkv.T
    q_x, k_x, v_x = torch.split(qkv_x, [w_query.size(0), w_key.size(0), w_value.size(0)], dim = -1)
    multi_head = self._multi_head(q_x, k_x, v_x, token_positions)
    return einsum(w_out, multi_head, "d_model d_v, ... seq_len d_v -> ... seq_len d_model")

测试结果如下：

$image\.png$

Experiment: Transformer Block

$image\.png$

在各个部分完成后，终于可以将它们组合起来了。一个Transformer block 主要分为两个部分：前半部分是因果多头自注意力机制（带RoPE），后半部分是逐位置前馈网络（我们已经实现了SwiGLU）。每个部分都需要先进行RMSNorm，再进行主要操作，最后加入残差连接。

具体来说，Transformer块的前半部分应该执行下面的一组更新，输入x然后输出y:

$y=x+MultiHeadSelfAttention(RMSNorm(x))$

模块至少应该包含以下参数：

d_model int：Transformer块输入的维数
num_heads int：用于多头自注意力的头的数量。
d_ff int: 逐位置前馈网络的维度

在adapters.run_transformer_block中完成模块的编写，并且使用下面的语句进行测试：

1	uv run pytest -k test_transformer_block

这个实验就很简单了，基本就像拼积木一样，直接给代码：

from torch import nn, Tensor
import torch

from cs336_basics.Transformer.Causal_multi_head_self_attention import MultiHeadAttention
from cs336_basics.Transformer.RMSNorm import RMSNorm
from cs336_basics.Transformer.SwiGLU import SwiGLU

class TransformerBlock(nn.Module):

    def  __init__(self, d_model: int, num_heads: int, d_ff: int, max_seq_len: int, theta: float, weights: dict[str, Tensor] | None = None):
        super().__init__()
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_ff = d_ff
        self.max_seq_len = max_seq_len
        self.theta = theta
        if weights is not None:
            # Initialize the parameters with the provided weights
            self.q_proj = nn.Parameter(weights['attn.q_proj.weight'].clone())
            self.k_proj = nn.Parameter(weights['attn.k_proj.weight'].clone())
            self.v_proj = nn.Parameter(weights['attn.v_proj.weight'].clone())
            self.o_proj = nn.Parameter(weights['attn.output_proj.weight'].clone())
            self.swiglu = SwiGLU(d_model= self.d_model, d_ff= self.d_ff,
                             w1= weights['ffn.w1.weight'],
                             w2= weights['ffn.w2.weight'],
                             w3= weights['ffn.w3.weight'])
            self.res_norm_first = RMSNorm(self.d_model)
            self.res_norm_second = RMSNorm(self.d_model)
            self.res_norm_first.load_state_dict({"weight": weights['ln1.weight']})
            self.res_norm_second.load_state_dict({"weight": weights['ln2.weight']})
        else:
            # Initialize the parameters randomly(using truncated normal initialization)
            self.q_proj = self._init_linear_weight(d_model, d_model)
            self.k_proj = self._init_linear_weight(d_model, d_model)
            self.v_proj = self._init_linear_weight(d_model, d_model)
            self.o_proj = self._init_linear_weight(d_model, d_model)
            w1 = self._init_linear_weight(d_ff, d_model)
            w2 = self._init_linear_weight(d_model, d_ff)
            w3 = self._init_linear_weight(d_ff, d_model)
            self.swiglu = SwiGLU(d_model= self.d_model, d_ff= self.d_ff,
                         w1= w1,
                         w2= w2,
                         w3= w3)
            self.res_norm_first = RMSNorm(self.d_model)
            self.res_norm_second = RMSNorm(self.d_model)
        self.mha = MultiHeadAttention(d_model= self.d_model, num_heads= self.num_heads, rope= True, max_seq_len= self.max_seq_len, theta= self.theta)

    def forward(self, x: Tensor) -> Tensor:
        """
        x (Float[Tensor, "batch sequence_length d_model"]):
                    Tensor to run your implementation on.
        """
        # save x
        x_old = x
        # RMSNorm(first)
        x = self.res_norm_first(x)
        # Causal Multi-Head Attention with RoPE
        token_positions = torch.arange(x.shape[1], device = x.device)
        x = self.mha(w_query= self.q_proj, w_key= self.k_proj, w_value= self.v_proj, w_out= self.o_proj, x= x, token_positions= token_positions)

        # y = x + MultiHeadSelfAttention(RMSNorm(x))
        x = x_old + x
        x_old = x
        # RMSNorm(second)
        x = self.res_norm_second(x)
        # SwiGLU Feed-Forward
        x = self.swiglu(x)
        x = x_old + x
        return x

代码中专门区分了有提供权重和没有提供权重下的初始化操作，这里使用的初始化方法是正态截断的随机初始化方法，和我们在Linear层的方法是一致的，具体方法如下：

@staticmethod
def _init_linear_weight(d_out: int, d_in: int) -> nn.Parameter:
        std = math.sqrt(2.0 / (d_in + d_out))
        w = torch.empty(d_out, d_in)
        nn.init.trunc_normal_(w, mean=0.0, std=std, a=-3 * std, b=3 * std)
        return nn.Parameter(w)

需要注意的是不要在forward()中做RMSNorm，Mhlti-Head Attention和SwiGLU的初始化，防止非必要的反复计算。同时，如果反复new模块，权重状态不会保存，会导致优化器无法更新参数，这使得训练将无法进行下去。

然后可以发现token_positions是需要自己计算的，这时候就要注意入参x的维数是动态获取的。

测试结果如下：

imagepng

Experiment: Transformer LM

$image\.png$

现在我们将这些transformer block放在一起，如图中所示。

按照我们对token嵌入的描述，将其输入到num_layers层transformer block中，然后在三个输出层间依次传递，最后就能获得词汇上的分布。

你的实现至少应该接受所有上述transformer block的构造参数，以及这些附加参数：

vocab_size int: 词表的大小，可以确定嵌入矩阵的维度
context_length int: 确定位置嵌入矩阵维度所需的最大上下文长度
num_layers int: 就是上面提到的Transformer block的块数

在adapters.run_transformer_lm中完成模块的编写，并且使用下面的语句进行测试：

1	uv run pytest -k test_transformer_lm

和上面一样，这里也是直接拼积木就行了，但是尤其要留意各个层级的权重有没有用对地方。特别需要提醒的是，结果返回的是没有归一化的结果，意思就是并不需要最后一层的softmax，这是最坑的地方。

@dataclass
class Transformer(nn.Module):
*    *vocab_size: int
    context_length: int
    d_model: int
    num_layers: int
    num_heads: int
    d_ff: int
    rope_theta: float
    weights: dict[str, Tensor]

    def __post_init__(self):

        super().__init__()
        self.token_embeddings = Embedding(self.vocab_size, self.d_model)
        self.token_embeddings.load_state_dict({"weight": self.weights['token_embeddings.weight']})
        self.res_norm = RMSNorm(self.d_model)
        self.res_norm.load_state_dict({"weight": self.weights['ln_final.weight']})
        self.linear = Linear(self.d_model, self.vocab_size)
        self.linear.load_state_dict({"weight": self.weights['lm_head.weight']})
        self.layers = nn.ModuleList()
        for i in range(self.num_layers):
            prefix = f"layers.{i}."
            layer_weight = {
                k[len(prefix):]: v
                for k, v in self.weights.items()
                if k.startswith(prefix)
            }
            self.layers.append(
                TransformerBlock(self.d_model, self.num_heads, self.d_ff, self.context_length, self.rope_theta,
                                 layer_weight)
            )

    def forward(self, x):
        """
        forward the transformer process
        Args:
            x (Int[Tensor, "batch_size sequence_length"]) :Tensor with input indices to run the language model on. Shape is (batch_size, sequence_length), where `sequence_length` is at most `context_length`.
        Returns:

            Float[Tensor, "batch_size sequence_length vocab_size"]: Tensor with the predicted unnormalized
            next-word distribution for each token.
        """
        x = self.token_embeddings(x)
        for transformer_block in self.layers:
            x = transformer_block(x)
        x = self.res_norm(x)
        x = self.linear(x)
        # x = softmax(x, dim=-1)
        return x

可以看到这次使用的是@dataclass的写法，这种方式会让入参和module的属性自动绑定（类似于C#），就不需要写很多重复的self.xx = xx了。

在__post_init__的时候就进行了各种模块的初始化，是因为如果在forward中初始化的话将会反复初始化严重拖慢性能。

最后是layer_weight的操作，主要用来将各层的权重进行分离，这样才能适配TransformerBlock的内部权重命名要求。

最终测试结果如下，再次注意softmax是不需要使用的。

imagepng

Conclusion

在这次课程和很多的实验中，我们成功从零搭建了一个基础的Transformer架构，并且成功通过了测试。但是可以发现，我们现在还缺少一些组件，让我们对训练过程中的参数和状态缺乏有效的观察手段。下节课将会补充这些组件，并开始我们第一次的Transformer训练。