1.1 Softmax的本质

Softmax函数将一个向量映射为概率分布。对于输入向量 $z = (z_1, z_2, ..., z_n)$，其输出为：

$\text{softmax}(z_i) = \frac{e^{z_i}}{\sum_{j=1}^{n} e^{z_j}}$

这个公式有两个关键特性：

1. 输出总和为1：$\sum_i \text{softmax}(z_i) = 1$
2. 保持单调性：$z_i$ 越大，输出概率越大

在注意力机制中，Softmax应用于注意力分数（Q与K的点积除以$\sqrt{d_k}$），将分数转换为“关注权重”——权重越大，表示模型越关注对应的Value向量。

1.2 数值溢出问题

上述公式在实现中存在一个严重的数值隐患：指数函数 $e^{z_i}$ 增长极快。当 $z_i$ 取值较大时（如大于100），$e^{100}$ 会导致浮点数溢出（overflow），在Python中表现为 inf；而当 $z_i$ 取值较小（如小于-100），$e^{-100}$ 会下溢（underflow）为 0。

在注意力中，$z_i$ 是Q和K的点积除以$\sqrt{d_k}$，其绝对值可能很大，尤其是在训练初期或模型参数尺度较大时。数值溢出会导致梯度变为NaN，训练崩溃。

1.3 防止溢出的技巧

解决溢出的方案为在指数运算前，从每个输入中减去其最大值。

对于任意向量 $z$，令 $z_{\max} = \max_i z_i$。我们对 $z_i - z_{\max}$ 应用Softmax：

$\text{softmax}(z_i) = \frac{e^{z_i - z_{\max}}}{\sum_{j=1}^n e^{z_j - z_{\max}}}$

这个变换是数学上恒等的，因为：

$\frac{e^{z_i - z_{\max}}}{\sum_j e^{z_j - z_{\max}}} = \frac{e^{z_i} / e^{z_{\max}}}{\sum_j e^{z_j} / e^{z_{\max}}} = \frac{e^{z_i}}{\sum_j e^{z_j}}$

现在，所有指数函数的输入都 $\leq 0$，最大值为0。这完美地防止了溢出（最大值也不会超过$e^0=1$），同时保留下溢问题不影响梯度。

代码实现极为简洁：

1
def softmax(x: torch.Tensor, dim: int = -1) -> torch.Tensor:
2
    x_max = torch.max(x, dim=dim, keepdim=True)[0]
3
    x_exp = torch.exp(x - x_max)
4
    x_softmax = x_exp / torch.sum(x_exp, dim=dim, keepdim=True)
5
    return x_softmax

关键点：keepdim=True 保留了原始张量的维度结构，使得广播减法可以正确执行（将每个元素减去对应行的最大值）。

1.4 在注意力中的应用

在缩放点积注意力中，通常直接使用torch.softmax，它内部已经实现了数值稳定技巧。我们自定义的softmax函数则展示了其底层原理。在注意力分数张量中，dim=-1对应于key_len维度 —— 即对每个query，将其对所有key的注意力分数归一化为概率分布。

2.1 注意力机制的直觉理解

在深入数学之前，先建立直觉：注意力机制本质上是一个“查询-键值”检索系统。

想象你在图书馆找一本书：

• Query（查询）：你要找的书名或主题
• Key（键）：每本书的标签或摘要
• Value（值）：书本身的内容

你通过比对Query和所有Key的相似度，找到最相关的几本书，然后根据相似度加权组合它们的Value（内容），得到最终的信息。

在Transformer中，这个过程被形式化为：

$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{Q K^T}{\sqrt{d_k}}\right) V$

2.2 公式的逐行实现

1
def Scaled_dot_product_attention(Q, K, V, mask=None):
2
    d_k = Q.size(-1)
3
    scores = torch.einsum('... q d, ... k d -> ... q k', Q, K) / (d_k ** 0.5)
4
    if mask is not None:
5
        scores = scores.masked_fill(mask == False, float('-inf'))
6
    attention_weights = torch.softmax(scores, dim=-1)
7
    output = torch.einsum('... q k, ... k v -> ... q v', attention_weights, V)
8
    return output

第1步：计算点积分数

1
scores = torch.einsum('... q d, ... k d -> ... q k', Q, K) / (d_k ** 0.5)

einsum表达式 '... q d, ... k d -> ... q k' 表示（... 表示前导维度（如batch_size和num_heads）被保留）：

• 输入Q的形状为 (..., q, d_k)（q是query数量，d_k是维度）
• 输入K的形状为 (..., k, d_k)（k是key数量）
• 输出形状为 (..., q, k)，其中每个元素是Q中第i个向量与K中第j个向量的点积

关键点：除以 $\sqrt{d_k}$ 的原因 —— 这是缩放点积注意力中“缩放”二字的来源。

• 当 $d_k$ 较大时，点积 $Q \cdot K$ 的方差也会很大（约等于 $d_k$），导致softmax的输入值分布在较宽的范围，梯度进入饱和区（softmax输出的概率接近0或1），从而梯度很小，训练缓慢。
• 除以 $\sqrt{d_k}$ 将方差拉回到1左右，使softmax的梯度保持在一个健康的范围。这是Transformer作者的一个关键洞见。

第2步：应用掩码

1
if mask is not None:
2
    scores = scores.masked_fill(mask == False, float('-inf'))

掩码是一个布尔张量，形状与scores兼容（通常为 (batch, 1, q, k)）。值为False的位置表示“不允许关注”，我们将对应的分数设为 -inf。这样在softmax中，$e^{-\infty}=0$，该位置的权重变为0。

第3步：Softmax归一化

1
attention_weights = torch.softmax(scores, dim=-1)

对每个query（最后一维是key维度），将其分数转换为概率分布。

第4步：加权求和

1
output = torch.einsum('... q k, ... k v -> ... q v', attention_weights, V)

将权重张量 (..., q, k) 与V (..., k, v) 进行张量乘法，得到输出 (..., q, v)——即每个query聚合了所有value的加权和。

2.3 注意力的几何意义

从几何视角看，注意力机制将每个查询向量映射到由所有键值对构成的凸组合（因为softmax输出非负且和为1）。输出向量位于由V的行向量张成的凸包中。这种“软选择”机制使得模型可以根据输入动态调整信息流，这正是注意力的强大之处。

3.1 为什么需要因果性？

语言模型的任务是自回归生成：给定已生成的token序列，预测下一个token。这意味着在预测第 $t$ 个token时，模型只能看到位置 $1, 2, ..., t-1$ 的token，不能看到未来位置 $t, t+1, ...$ 的信息。

如果不施加因果约束，模型在训练时可以通过“偷看”未来位置来降低损失，但推理时未来位置不存在，这将导致训练与推理的不一致（称为“信息泄露”）。因此，我们必须在注意力计算中引入因果掩码（Causal Mask）。

3.2 下三角掩码矩阵

因果掩码是一个下三角矩阵（Lower Triangular Matrix）：

$\text{mask}_{ij} = \begin{cases} 1 & \text{if } i \geq j \\ 0 & \text{if } i < j \end{cases}$

其中 $i$ 是query的位置索引，$j$ 是key的位置索引。当 $i \geq j$ 时，表示query $i$ 可以关注到key $j$（即位置 $j$ 是“过去”或“当前”的）；当 $i < j$ 时，表示query $i$ 不能关注到未来的key $j$。

在代码中，我们通过torch.tril生成下三角矩阵：

1
mask = torch.tril(torch.ones(seq_len, seq_len, device=x.device)).bool()
2
mask = mask.unsqueeze(0).unsqueeze(0)  # (1, 1, seq_len, seq_len)

tril（Lower Triangle）保留矩阵的下三角部分（包括对角线），其余置0。然后转换为布尔型，并扩展两个维度以匹配注意力分数的形状：(1, 1, seq_len, seq_len)，其中第一个1是batch维度（可广播），第二个1是head维度（可广播）。

3.3 掩码的应用

在Scaled_dot_product_attention中，掩码通过masked_fill应用：

1
scores = scores.masked_fill(mask == False, float('-inf'))

所有False位置（即上三角，未来位置）被设为-inf，softmax后权重为0。

注意点：在实际实现中，mask的维度可能包含head维度，也可能不包含，取决于具体用法。代码中采用了广播机制，使得 (1, 1, seq_len, seq_len)的掩码可以自动匹配(batch_size, num_heads, seq_len, seq_len)的分数张量。

4.1 单头注意力的局限

单头注意力（Single-head Attention） 只能从一个角度计算query和key的相似度。但在自然语言中，一个词与上下文的关系往往是多维度的 —— 例如，“苹果”既可以指水果，也可以指公司，不同头可以分别关注这些不同的语义方面。

多头注意力（Multi-Head Attention） 的设计思路是：使用多组独立的投影矩阵，将 $d_{\text{model}}$ 维的Q、K、V投影到多个子空间，每个子空间分别计算注意力，然后将结果拼接起来。

数学上：

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

其中：

$\text{head}_i = \text{Attention}(Q W_i^Q, K W_i^K, V W_i^V)$

4.2 多头的拆分与合并

在代码中，没有为每个头单独创建投影矩阵，而是用一个大的投影矩阵一次性投影出所有头的Q、K、V，然后用einops.rearrange进行维度拆分。

1
q = self.wq(x)  # (batch, seq_len, d_model)
2
k = self.wk(x)
3
v = self.wv(x)
4

5
q = rearrange(q, 'b s (n_h d_k) -> b n_h s d_k', n_h=self.num_heads)
6
k = rearrange(k, 'b s (n_h d_k) -> b n_h s d_k', n_h=self.num_heads)
7
v = rearrange(v, 'b s (n_h d_k) -> b n_h s d_k', n_h=self.num_heads)

rearrange 的表达式 'b s (n_h d_k) -> b n_h s d_k' 的含义是：

• 左侧 'b s (n_h d_k)'：输入形状为 (batch, seq_len, d_model)，其中最后一维d_model 被解释为 n_h(n_heads) 乘以 d_k
• 右侧 'b n_h s d_k'：输出形状为 (batch, n_heads, seq_len, d_k)

这种写法比手动view或reshape更清晰、更不易出错。einops库让张量维度的变换变得具有可读性。

4.3 多头注意力的效率优势

多头机制的核心优势在于并行化。在单头注意力中，所有计算是顺序的；而在多头中，所有头可以独立并行计算，因为它们在batch和head维度上是独立的（注意力的softmax在每个头内部进行，头与头之间没有数据依赖）。

这使得多头注意力在GPU上可以充分发挥并行计算能力，同时通过拆分维度降低了每个头的d_k，从而减少了点积计算的复杂度（$O(h \cdot d_k^2)$ vs $O((h \cdot d_k)^2)$）。

4.4 输出投影

注意力计算完成后，我们得到形状为 (batch, n_heads, seq_len, d_k) 的张量，需要将其合并回 (batch, seq_len, d_model)：

1
attn_output = rearrange(attn_output, 'b n_h s d_k -> b s (n_h d_k)', n_h=self.num_heads)
2
out = self.wo(attn_output)

最后的 wo 是一个线性层，将拼接后的向量投影回 d_model 维，作为多头注意力的最终输出。

5.1 位置编码的不同范式

在第二篇博客中，我们介绍了 RoPE（旋转位置编码）。它与传统位置编码（如正弦位置编码或可学习位置编码）有一个根本区别：

• 传统方式：在输入层的词嵌入上添加位置向量（加法融合）
• RoPE方式：在注意力计算的Q和K上施加旋转（乘法融合）

RoPE的设计使得位置信息直接嵌入到注意力分数中，而不是作为额外的特征向量。这有两个好处：

1. 位置编码不占用模型的“表达容量”（词嵌入维度不变）
2. 注意力分数天然地包含了相对位置信息

5.2 在多头注意力中集成RoPE

在CausalMultiHeadAttention类中，我们通过use_rope参数控制是否启用RoPE：

1
class CausalMultiHeadAttention(nn.Module):
2
    def __init__(self, d_model, n_heads, use_rope=False, max_seq_len=None, theta=10000.0):
3
        # ...
4
        if use_rope:
5
            self.rope = RotaryPositionalEmbedding(theta, self.d_k, max_seq_len=max_seq_len)

在前向传播中，我们在将Q和K拆分为多头之后，对每个头分别应用RoPE：

1
if self.use_rope:
2
    assert token_positions is not None
3
    q_rope = []
4
    k_rope = []
5
    for head in range(self.num_heads):
6
        q_head = q[:, head, :, :]  # (batch, seq_len, d_k)
7
        k_head = k[:, head, :, :]
8
        q_head_rope = self.rope(q_head, token_positions)
9
        k_head_rope = self.rope(k_head, token_positions)
10
        q_rope.append(q_head_rope.unsqueeze(1))
11
        k_rope.append(k_head_rope.unsqueeze(1))
12
    q = torch.cat(q_rope, dim=1)
13
    k = torch.cat(k_rope, dim=1)

• 对每个头分别应用的原因：RoPE的d_k是每个头的维度，而整个d_model可能很大。每个头拥有独立的d_k维子空间，旋转操作在每个头的子空间内独立进行。
• 只对Q和K施加RoPE的原因：位置信息通过Q·K的点积影响注意力权重，而V只是被权重加权聚合，不需要携带位置信息。这已被理论和实验证明是有效的。

5.3 token_positions参数的传入

在forward中，接受token_positions参数（形状为 (batch, seq_len)），它指定了每个token在序列中的绝对位置（从0开始）。这些位置索引用于从预计算的cos/sin表中取值。

在训练中，如果我们直接处理连续的序列片段，token_positions通常就是0, 1, 2, ..., seq_len-1。但在推理过程中，我们可能增量生成token，此时位置偏移需要正确设置。

5.4 因果掩码与RoPE的协同

在带RoPE的因果多头注意力中，因果掩码和RoPE共同工作：

• RoPE —— 在注意力分数中隐式地编码了相对位置信息
• 因果掩码 —— 显式地阻止未来位置的注意力

两者相辅相成：RoPE告诉模型位置差为d的两个词应该如何交互，而因果掩码强制这种交互只能是单向的（从过去到未来）。