缩放点积注意力

Transformer 架构使用注意力作为其核心机制，取代了传统的序列处理方法。注意力机制 (attention mechanism)不依赖于逐步传递的隐藏状态，而是让模型在处理特定部分时，能够直接衡量输入序列不同部分的权重 (weight)。此项操作的基本单元是缩放点积注意力。

想象一下，你正在句子“The river bank was eroding.”（河岸正在被侵蚀）中理解“bank”这个词的含义。为了辨别“bank”的含义，你会很自然地更关注“river”（河流），而不是“eroding”（侵蚀）或“was”（是）。自注意力 (self-attention)机制将这种直觉形式化。对于序列中的每个元素（例如，每个词的嵌入 (embedding)向量 (vector)），我们希望计算一个表示。这个表示结合了来自其他元素的信息，并根据它们的关联性进行加权。

为了做到这一点，缩放点积注意力基于从输入序列嵌入中得到的三个输入进行操作：

1. 查询 ($Q$)：一组向量，表示当前正在获取信息的元素。可以把它想象成元素在问：“序列的哪些部分与我有关联？”
2. 键 ($K$)：一组与序列中每个元素关联的对应向量。它们就像元素的标签或标识符。通过将它们与查询进行比较来确定关联性。
3. 值 ($V$)：另一组与每个元素关联的向量。这些向量包含元素的实际信息或表示。一旦关联性确定（通过查询-键交互），关联元素的“值”向量就会被聚合。

在实际操作中，$Q$、$K$和$V$通常是通过带有学习权重的独立线性层对输入嵌入（或前一层的输出）进行投影来生成的。设输入序列嵌入的维度为 $d_{model}$，键和查询向量的维度为 $d_k$。值向量的维度为 $d_v$（通常 $d_k = d_v$，但并非一定如此）。

计算过程分为几个步骤：

1. 计算相似度得分

第一步是衡量每个查询与所有键之间的兼容性或相似度。这通过点积完成。对于单个查询 $q$ 和所有键 $K$，我们对每个键 $k_i$ 计算 $q \cdot k_i$。对于完整的矩阵 $Q$ 和 $K$，这可以有效地通过矩阵乘法计算：

$$Scores = Q K^T$$

结果矩阵包含原始得分；$Scores_{ij}$ 表示查询 $i$ 和键 $j$ 之间的相似度。点积越高，表示查询和键之间的关联性越大。

2. 缩放得分

点积的数值可能变得很大，特别是对于较大的维度 $d_k$。输入到softmax函数（下一步）中的较大值可能导致梯度极小，从而阻碍学习。为解决此问题，得分会除以键维度的平方根，即 $\sqrt{d_k}$ 进行缩放：

$$缩放得分 = \frac{Q K^T}{\sqrt{d_k}}$$

这种缩放有助于稳定梯度并使训练更具可靠性。选择 $\sqrt{d_k}$ 是基于以下假设：$Q$ 和 $K$ 的分量是均值为零、方差为一的独立随机变量。在此假设下，点积 $q \cdot k = \sum_{i=1}^{d_k} q_i k_i$ 的均值为 0，方差为 $d_k$。通过 $\sqrt{d_k}$ 进行缩放可将方差恢复到 1，使 softmax 的输入保持在合理的范围内。

3. 计算注意力权重 (weight) (Softmax)

为了将缩放后的得分转换为表示注意力权重的概率分布，对缩放得分矩阵的每一行应用 softmax 函数：

$$注意力权重 = softmax(\frac{Q K^T}{\sqrt{d_k}})$$

注意力权重矩阵的每行现在总和为 1，并且每个元素 $Weights_{ij}$ 表明查询 $i$ 应该对值 $j$ 投入多少注意力。

4. 计算加权值

最后，注意力权重 (weight)用于计算值向量 (vector)的加权和。这意味着将 注意力权重 矩阵乘以 值 矩阵 $V$：

$$输出 = 注意力权重 \cdot V$$

得到的 输出 矩阵包含注意力加权表示。每行 $Output_i$ 是一个向量，它是 $V$ 中所有值向量的加权组合，其中权重由查询 $i$ 与所有键的相似度决定。这个输出向量有效地融入了整个序列的上下文 (context)，并根据关联性进行了加权。

因此，缩放点积注意力的完整公式为：

$$注意力(Q, K, V) = softmax(\frac{QK^T}{\sqrt{d_k}})V$$

我们来绘制数据流图：

缩放点积注意力机制 (attention mechanism)的数据流图。输入嵌入 (embedding)被投影到Q、K、V矩阵中，然后经过矩阵乘法、缩放和softmax处理，生成加权输出表示。

以下是使用 PyTorch 实现的核心计算的简化示例：

import torch
import torch.nn.functional as F
import math

def scaled_dot_product_attention(query, key, value, mask=None):
    """
    计算缩放点积注意力。

    参数：
        query: 查询张量 (批量, 查询序列长度, 键维度)
        key: 键张量 (批量, 键值序列长度, 键维度)
        value: 值张量 (批量, 键值序列长度, 值维度)
        mask: 可选的掩码张量 (批量, 1, 查询序列长度, 键值序列长度)

    返回：
        输出张量 (批量, 查询序列长度, 值维度),
        注意力权重 (批量, 查询序列长度, 键值序列长度)
    """
    dim_k = query.size(-1)
    # 矩阵乘法 QK^T: (批量, 查询序列长度, 键维度) x (批量, 键维度, 键值序列长度)
    # -> (批量, 查询序列长度, 键值序列长度)
    scores = torch.matmul(query, key.transpose(-2, -1))

    # 缩放
    scaled_scores = scores / math.sqrt(dim_k)

    # 可选掩码 (例如，用于填充或阻止解码器中的未来信息查看)
    if mask is not None:
        # 当mask为True（或0）时，应用一个很大的负值
        scaled_scores = scaled_scores.masked_fill(mask == 0, -1e9)
        # 使用 -1e9 保持数值稳定性

    # Softmax
    # Softmax 在最后一个维度 (键值序列长度) 上操作
    attention_weights = F.softmax(scaled_scores, dim=-1)

    # 矩阵乘法 权重 * V: (批量, 查询序列长度, 键值序列长度) x
    # (批量, 键值序列长度, 值维度) -> (批量, 查询序列长度, 值维度)
    output = torch.matmul(attention_weights, value)

    return output, attention_weights

# 示例用法（简化维度）
batch_size = 1
seq_len_q = 3 # 查询序列长度
seq_len_kv = 5 # 键/值序列长度
dim_k = 8
dim_v = 10

# 虚拟输入张量
q = torch.randn(batch_size, seq_len_q, dim_k)
k = torch.randn(batch_size, seq_len_kv, dim_k)
v = torch.randn(batch_size, seq_len_kv, dim_v)

# 计算注意力
output, weights = scaled_dot_product_attention(q, k, v)

print("输出形状:", output.shape)
print("注意力权重形状:", weights.shape)

# Example output:
# Output shape: torch.Size([1, 3, 10])
# Attention weights shape: torch.Size([1, 3, 5])

在这段代码中：

• 我们执行矩阵乘法 $Q K^T$。
• 我们通过 math.sqrt(dim_k) 对结果进行缩放。
• 代码中包含一个可选的 mask 参数 (parameter)。掩码在Transformer中非常重要，例如，它们可以防止模型关注填充标记 (token)，或者在解码器中防止关注未来的标记（前瞻掩码）。我们通过将掩码位置设置为一个非常大的负数，在softmax 之前 应用掩码，以确保它们在softmax后得到接近零的概率。
• F.softmax 计算注意力权重。
• 最后，我们将权重乘以 值 张量以获得输出。

这种机制构成了Transformer能够捕获序列中不限距离的依赖关系的核心能力，与RNN中固有的序列瓶颈相比，这是一个显著的优势。然而，单次注意力计算可能只关注一种类型的关系。为了同时捕获多种关系，Transformer采用了多头注意力 (multi-head attention)，我们将在接下来进行审视。