多头注意力
Transformer 中的多头注意力机制
Section titled “Transformer 中的多头注意力机制”在Transformer架构中，输入嵌入（input embeddings）与位置编码（Positional Encoding）结合后，得到的向量被送入第一个核心组件：多头注意力（multi-head attention）机制。这一机制是相对于简单注意力方法的一项重大改进，它使模型能够更有效地处理信息。
多头注意力机制是Transformer模型的基础，对于其在处理序列数据（如自然语言处理Natural Language Processing, NLP 中的文本）方面表现出色至关重要。它使模型能够同时关注不同表示子空间中不同位置的信息。
在本章中，我们将探讨多头注意力机制的架构，讨论其优势，并提供一个概念性的使用NumPy实现的Python代码示例，以说明其工作原理。
理解自注意力机制（缩放点积注意力）
Section titled “理解自注意力机制（缩放点积注意力）”在深入了解多头注意力之前，理解其核心组件——自注意力（self-attention，通常称为缩放点积注意力scaled dot-product attention）至关重要。这种机制使模型在生成每个token的表示时，能够衡量输入序列中不同token相互之间的重要性。
自注意力机制的工作原理
Section titled “自注意力机制的工作原理”自注意力过程包括以下关键步骤：
- 创建查询（Query）、键（Key）和值（Value）向量：对于输入序列中的每个token（例如，词嵌入word embeddings），通过将其嵌入向量乘以学习到的权重矩阵，生成三个向量：查询向量（Q）、键向量（K）和值向量（V）。这些向量的维度分别为d_q、d_k和d_v（通常d_q = d_k）。
- 计算注意力分数：通过计算特定token的查询向量（Query）与序列中所有token（包括自身）的键向量（Key）的点积来计算注意力分数。这个分数决定了在处理当前token时，应将多少注意力放在输入序列的其他部分上。
- 缩放分数：然后，通过除以键向量维度（d_k）的平方根（√d_k）来缩放分数。这种缩放可以防止点积变得过大，从而避免在Softmax函数中出现梯度消失。
- 应用Softmax函数：将Softmax函数应用于缩放后的分数，以获得注意力权重（attention weights）。这些权重是正数且总和为1，代表输入序列上注意力的概率分布。
- 计算值的加权总和：最后，每个token的输出是序列中所有值向量（Value vectors）的加权总和，权重就是刚刚计算出的注意力权重。
数学上，缩放点积注意力定义为：
$$\mathrm{Attention(Q, K, V)} = \mathrm{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$
其中Q、K和V是分别包含查询、键和值向量的矩阵。
什么是多头注意力机制？
Section titled “什么是多头注意力机制？”多头注意力（Multi-Head Attention）通过并行执行多个自注意力操作来增强自注意力机制，每个并行操作称为一个“注意力头”（attention head）。
多头注意力不使用一组单一的查询、键和值权重矩阵，而是使用多组矩阵，允许模型同时关注来自不同表示子空间中不同位置的信息。
这可以形象地理解为模型学习同时关注输入序列的不同方面。例如，一个注意力头可能关注句法关系，而另一个则关注语义相似性。
从架构上看，多头注意力层由多个并行的注意力头组成。每个注意力头独立地执行缩放点积注意力。然后将这些注意力头的输出连接（concatenated）起来，并通过一个最终的线性层（linear layer）来生成多头注意力块的输出。这种结构是Transformer编码器（encoder）和解码器（decoder）层的关键部分。
多头注意力机制的步骤
Section titled “多头注意力机制的步骤”多头注意力的过程包含以下步骤：
- 每个注意力头的线性投影：输入嵌入（或前一层的输出）使用每个注意力头 ‘i’ 不同的学习到的权重矩阵 (W_i^Q, W_i^K, W_i^V) 被线性投影到多组查询 (Q_i)、键 (K_i) 和值 (V_i) 中。通常，如果模型维度是d_model并且有 ‘h’ 个注意力头，则每个注意力头处理的向量维度是 d_k = d_v = d_model / h。
- 并行缩放点积注意力：对每个注意力头，使用其各自的 Q_i、K_i 和 V_i 投影并行执行自注意力（缩放点积注意力）。这会产生 ‘h’ 个独立的输出矩阵 (head_i)。
- 连接：将所有注意力头的输出矩阵沿着特征维度连接（concatenate）起来。
- 最终线性投影：然后将连接后的输出通过最终的线性变换（乘以另一个学习到的权重矩阵 W^O）来产生多头注意力层的最终输出。这种投影通常将连接后的特征映射回原始的 d_model 维度。
数学上，多头注意力可以表示为：
$$\mathrm{MultiHead(Q, K, V)} = \mathrm{Concat(head_1, \dots, head_h)}W^O$$
其中每个注意力头的计算公式为：
$$\mathrm{head_i} = \mathrm{Attention(QW_i^Q, KW_i^K, VW_i^V)}$$
其中 W_i^Q、W_i^K、W_i^V 是注意力头 ‘i’ 的参数矩阵，而 W^O 是最终线性投影的参数矩阵。
多头注意力机制的优势
Section titled “多头注意力机制的优势”- 增强的表示学习：通过从不同的表示子空间关注输入序列的不同部分，模型可以捕捉数据中更丰富的关系和细微差别。
- 改进的聚焦能力：不同的注意力头可以学习关注不同类型的信息（例如，短程依赖 vs. 长程依赖，不同的句法或语义特征）。
- 并行化：每个注意力头的计算可以并行执行，使得该机制在GPU等现代硬件上高效运行，与处理长序列的循环模型相比，有助于加快训练时间。
使用NumPy的概念性Python实现
Section titled “使用NumPy的概念性Python实现”以下Python脚本提供了使用NumPy实现多头注意力机制的概念性示例。这有助于理解其底层工作原理。在实际应用中，像PyTorch或TensorFlow这样的深度学习框架提供了经过优化的实现。
import numpy as np
class MultiHeadAttention: def __init__(self, d_model: int, num_heads: int): """初始化 MultiHeadAttention 层。
Args: d_model (int): 输入和输出的维度（嵌入维度）。 num_heads (int): 注意力头的数量。 """ assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
self.d_model = d_model self.num_heads = num_heads self.depth = d_model // num_heads # 每个注意力头的键、查询、值的维度
# 所有注意力头的 Q、K、V 权重矩阵（可以看作是连接起来的 W_i^Q 等） # 在实际实现中，这些是可学习的参数。 # Xavier/Glorot 或 He 初始化是常用的。 self.W_q = np.random.randn(d_model, d_model) * np.sqrt(1 / d_model) self.W_k = np.random.randn(d_model, d_model) * np.sqrt(1 / d_model) self.W_v = np.random.randn(d_model, d_model) * np.sqrt(1 / d_model)
# 最终线性投影的权重矩阵 self.W_o = np.random.randn(d_model, d_model) * np.sqrt(1 / d_model)
def _scaled_dot_product_attention(self, q_h, k_h, v_h, mask=None): """计算单个注意力头的缩放点积注意力。
Args: q_h (np.ndarray): 注意力头的查询张量 (batch_size, seq_len_q, depth)。 k_h (np.ndarray): 注意力头的键张量 (batch_size, seq_len_k, depth)。 v_h (np.ndarray): 注意力头的值张量 (batch_size, seq_len_v, depth)。 (seq_len_k 通常等于 seq_len_v) mask (np.ndarray, optional): 应用于注意力分数的掩码。
Returns: tuple: (输出, 注意力权重) output (np.ndarray): 上下文表示 (batch_size, seq_len_q, depth)。 attention_weights (np.ndarray): 注意力分布 (batch_size, seq_len_q, seq_len_k)。 """ # 矩阵乘法 Q_h 和 K_h.T: (batch_size, seq_len_q, depth) @ (batch_size, depth, seq_len_k) # -> (batch_size, seq_len_q, seq_len_k) matmul_qk = np.matmul(q_h, k_h.transpose(0, 2, 1))
# 缩放 matmul_qk dk = k_h.shape[-1] scaled_attention_logits = matmul_qk / np.sqrt(dk)
# 如果提供了掩码，则应用（例如，用于填充或解码器中的未来 token 掩码） if mask is not None: # 在 mask 为 1 的位置添加一个很大的负数 scaled_attention_logits += (mask * -1e9)
# Softmax 应用于最后一个轴 (seq_len_k) 以获得权重 # 数值稳定 Softmax: scaled_attention_logits_stabilized = scaled_attention_logits - np.max(scaled_attention_logits, axis=-1, keepdims=True) attention_weights = np.exp(scaled_attention_logits_stabilized) attention_weights /= np.sum(attention_weights, axis=-1, keepdims=True)
# 矩阵乘法 attention_weights 和 V_h: (batch_size, seq_len_q, seq_len_k) @ (batch_size, seq_len_v, depth) # -> (batch_size, seq_len_q, depth) output = np.matmul(attention_weights, v_h) return output, attention_weights
def _split_heads(self, x: np.ndarray, batch_size: int): """将最后一个维度分割为 (num_heads, depth)。 转置结果为形状 (batch_size, num_heads, seq_len, depth)。 Args: x (np.ndarray): 输入张量 (batch_size, seq_len, d_model)。 batch_size (int): Batch 大小。 Returns: np.ndarray: 重塑后的张量 (batch_size, num_heads, seq_len, depth)。 """ seq_len = x.shape[1] x = np.reshape(x, (batch_size, seq_len, self.num_heads, self.depth)) # (batch_size, num_heads, seq_len, depth) return x.transpose(0, 2, 1, 3)
def call(self, q_in: np.ndarray, k_in: np.ndarray, v_in: np.ndarray, mask=None): """MultiHeadAttention 的前向传播。
Args: q_in (np.ndarray): 查询输入 (batch_size, seq_len_q, d_model)。 k_in (np.ndarray): 键输入 (batch_size, seq_len_k, d_model)。 v_in (np.ndarray): 值输入 (batch_size, seq_len_v, d_model)。 mask (np.ndarray, optional): 在注意力计算期间应用的掩码。
Returns: tuple: (输出, 所有注意力头的注意力权重) output (np.ndarray): 最终输出 (batch_size, seq_len_q, d_model)。 attention_weights_all_heads (np.ndarray): 所有注意力头的注意力权重 (batch_size, num_heads, seq_len_q, seq_len_k)。 """ batch_size = q_in.shape[0]
# 1. Q、K、V 的线性投影 # (batch_size, seq_len, d_model) @ (d_model, d_model) -> (batch_size, seq_len, d_model) q = np.dot(q_in, self.W_q) k = np.dot(k_in, self.W_k) v = np.dot(v_in, self.W_v)
# 2. 分割注意力头 # (batch_size, num_heads, seq_len, depth) q_split = self._split_heads(q, batch_size) k_split = self._split_heads(k, batch_size) v_split = self._split_heads(v, batch_size)
# 3. 每个注意力头的缩放点积注意力 # 我们将存储注意力权重以供检查，尽管通常只有输出会被传递。 attention_heads_outputs = [] attention_weights_all_heads_list = []
# 在向量化实现中，可以通过在 num_heads 维度上操作来避免这个循环。 # 为了清晰，我们在这里迭代。但为了效率，首选在所有注意力头之间进行批处理矩阵乘法。 # 现代 NumPy 支持批处理，因此 _scaled_dot_product_attention 可以处理像 # (batch_size, num_heads, seq_len, depth) 这样的张量，如果实现得足够仔细。对于本例，我们将保持简单。 # 对于真正的并行计算，可以将 q_split、k_split、v_split 视为并行注意力单元的输入。
# 在这里，我们演示一个简化版本，其中 scaled_dot_product_attention # 如果支持跨注意力头的批处理，则直接应用于已分割的 Q、K、V 张量。 # 我们当前的 _scaled_dot_product_attention 期望 (batch_size, seq_len, depth) 形状的输入， # 因此我们通常会循环或重塑来处理所有注意力头。 # 对于这个概念示例，我们假设 _scaled_dot_product_attention 可以处理注意力头维度： # (batch_size, num_heads, seq_len_q, depth) attention_output_concat, attention_weights_concat = self._scaled_dot_product_attention(q_split, k_split, v_split, mask) # 注意：上述假设 _scaled_dot_product_attention 已修改以处理 num_heads 维度。 # 如果没有，则需要一个遍历注意力头的循环： # for i in range(self.num_heads): # head_output, head_weights = self._scaled_dot_product_attention(q_split[:,i,:,:], k_split[:,i,:,:], v_split[:,i,:,:], mask) # attention_heads_outputs.append(head_output) # attention_weights_all_heads_list.append(head_weights) # attention_output_concat = np.stack(attention_heads_outputs, axis=1) # 如果循环 # attention_weights_concat = np.stack(attention_weights_all_heads_list, axis=1) # 如果循环
# 4. 连接注意力头 # 转置回 (batch_size, seq_len_q, num_heads, depth) attention_output_transposed = attention_output_concat.transpose(0, 2, 1, 3) # 重塑为 (batch_size, seq_len_q, d_model) seq_len_q = q_in.shape[1] concat_attention = np.reshape(attention_output_transposed, (batch_size, seq_len_q, self.d_model))
# 5. 最终线性投影 # (batch_size, seq_len_q, d_model) @ (d_model, d_model) -> (batch_size, seq_len_q, d_model) output = np.dot(concat_attention, self.W_o)
# attention_weights_concat 将是 (batch_size, num_heads, seq_len_q, seq_len_k) return output, attention_weights_concat
# 示例用法：d_model = 512 # 嵌入维度num_heads = 8 # 注意力头数量batch_size = 2 # 批处理中的序列数量seq_len = 10 # 每个序列的长度
# 创建 MultiHeadAttention 实例multi_head_attn = MultiHeadAttention(d_model, num_heads)
# 示例输入张量（用于演示的随机数据）# 在实际场景中，Q、K、V 可能是同一个张量（自注意力）# 或不同（例如，编码器-解码器注意力）。Q_input = np.random.randn(batch_size, seq_len, d_model)K_input = np.random.randn(batch_size, seq_len, d_model)V_input = np.random.randn(batch_size, seq_len, d_model)
# 执行多头注意力# 注意：此类中的概念性 _scaled_dot_product_attention 需要调整# 以正确处理来自 _split_heads 的 (batch_size, num_heads, seq_len, depth) 输入# 或者应该使用循环。为简洁起见，我们假设它已适配。# 让我们改进调用方式，以反映如果 _scaled_dot_product_attention 处理了批处理注意力头，它将如何工作：
# 简化的直接调用演示 - 假设内部已处理output, attention_weights = multi_head_attn.call(Q_input, K_input, V_input)
print("Input Query (Q_input shape):", Q_input.shape)print("Multi-Head Attention Output (shape):", output.shape)print("Attention Weights (shape):", attention_weights.shape)print("\nFirst batch, first token's output vector (first 5 elements):\n", output[0, 0, :5])print("\nAttention weights for first head, first batch, first query token (first 5 keys):\n", attention_weights[0, 0, 0, :5])预期输出结构
Section titled “预期输出结构”运行带有示例用法的脚本将生成输出张量（tensors）。其形状很重要：
Input Query (Q_input shape): (2, 10, 512)Multi-Head Attention Output (shape): (2, 10, 512)Attention Weights (shape): (2, 8, 10, 10)
First batch, first token's output vector (first 5 elements):\n [ ... some float values ... ]
Attention weights for first head, first batch, first query token (first 5 keys):\n [ ... some float values summing to <=1 ... ]输出张量的形状与输入查询相同，表示上下文嵌入（contextualized embeddings）。注意力权重张量显示了对于每个注意力头和每个token，它对序列中所有其他token的关注程度。
多头注意力机制是Transformer架构中强大而不可或缺的一部分。它通过同时从多个注意力视角处理信息，使模型能够有效地捕捉序列数据中的复杂依赖关系。
其学习多样化表示的能力以及对并行计算的适应性，使其成为现代深度学习模型的基石，尤其是在自然语言处理（NLP）领域。
尽管本章提供了概念性概述和基于NumPy的实现以增加清晰度，但实际生产模型利用PyTorch和TensorFlow等框架中的优化版本。理解这些基础知识对于领会当前先进AI系统背后的创新至关重要。
如需进一步学习，请查阅原始论文《Attention Is All You Need》以及流行深度学习库中的实现。