Transformer注意力机制深度剖析
注意力机制的起源与直觉 Transformer的核心创新在于抛弃了RNN的序列依赖,转而使用注意力机制直接建模序列中任意两个位置之间的关联。这一思想源于人类视觉的注意力选择性——在处理信息时,我们会自然而然地将更多注意力分配给相关度高的部分。 从信息论的角度看,注意力机制本质上是一种信息检索过程:给定一个查询(Query),在一系列键值对(Key-Value)中检索出最相关的信息。这种检索不是硬性的离散选择,而是通过软权重实现连续可微的加权聚合。 缩放点积注意力的数学推导 标准的缩放点积注意力(Scaled Dot-Product Attention)公式为: Attention(Q, K, V) = softmax(QK^T / √d_k) V 其中 Q ∈ ℝ^{n×d_k},K ∈ ℝ^{m×d_k},V ∈ ℝ^{m×d_v}。 为什么要除以√d_k? 这是一个容易被忽视但至关重要的细节。当 d_k 较大时,QK^T 的元素值会随之增大。假设 Q 和 K 的每个元素都是均值为0、方差为1的独立随机变量,那么 QK^T 的每个元素的方差为 d_k。当 d_k = 512 时,点积值的量级可能达到 ±22 左右,这会使得 softmax 函数进入梯度饱和区。 除以 √d_k 将方差归一化为1,确保 softmax 的梯度保持健康。这个看似微小的缩放因子,在深层 Transformer 训练中起到了不可替代的稳定作用。 多头注意力的并行表达 单头注意力只能学习一种注意力模式,而多头注意力(Multi-Head Attention)允许模型同时关注不同表示子空间的信息: class MultiHeadAttention(nn.Module): def __init__(self, d_model=512, n_heads=8): super().__init__() self.d_model = d_model self.n_heads = n_heads self.d_k = d_model // n_heads self.W_q = nn.Linear(d_model, d_model) self.W_k = nn.Linear(d_model, d_model) self.W_v = nn.Linear(d_model, d_model) self.W_o = nn.Linear(d_model, d_model) def forward(self, x, mask=None): batch_size, seq_len, _ = x.shape Q = self.W_q(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2) K = self.W_k(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2) V = self.W_v(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2) scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k) if mask is not None: scores = scores.masked_fill(mask == 0, float('-inf')) attn = F.softmax(scores, dim=-1) out = torch.matmul(attn, V) out = out.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model) return self.W_o(out) 每个头独立学习不同的注意力模式——有的头关注语法依赖,有的关注语义相似度,有的捕捉长距离位置关系。经验研究表明,8头注意力中不同头确实呈现出了明确的功能分化。 ...

