在 Transformer 中,Self-Attention(自注意力机制)是其核心组件之一,用于建模序列中各个 token 之间的关系。它让模型在处理某个词时,可以“关注”同一句中其他所有词,从而捕捉上下文信息。
一、Self-Attention 的作用
传统的语言模型(如 RNN、CNN)只能看到固定窗口或前后邻居的信息。而 Self-Attention 可以实现:
-
全局建模:每个 token 能看到整个输入序列;
-
动态权重分配:不同上下文下,同一个词关注的对象可能不同;
-
并行计算:相比 RNN 的逐个处理,Self-Attention 支持并行训练。
二、Self-Attention 的核心思想
输入是一个序列:
例如一句话被编码成向量矩阵:
输入:x = [x₁, x₂, ..., xₙ] ∈ ℝ^{n×d_model}
我们将每个 token 映射为 3 个向量:
-
Query(查询)
-
Key(键)
-
Value(值)
这些是通过权重矩阵线性变换得到的:
三、Self-Attention 的公式
Self-Attention 的输出为:
其中:
-
-
-
-
:计算每个 token 对其他 token 的注意力打分(注意缩放)
-
softmax 保证每个 token 对其他 token 的权重之和为 1
为什么要除以
?
为了防止点积值过大导致 softmax 输出梯度消失。否则随着维度增加,分数可能变得非常大,导致 softmax 输出变得极端。
四、Self-Attention 的运行示意(图形化描述)
假设输入有 3 个 token:
["I", "love", "you"]
计算过程如下:
-
每个词 → 计算 Query、Key、Value 向量;
-
对每个词(如 “love”):
-
用它的 Query 去与所有 Key 计算相似度;
-
通过 softmax 得到关注权重;
-
用这些权重加权 Value 向量 → 输出结果。
-
五、代码实现(PyTorch 简化版)
import torch
import torch.nn as nn
import torch.nn.functional as F
class SelfAttention(nn.Module):
def __init__(self, embed_dim):
super().__init__()
self.q_linear = nn.Linear(embed_dim, embed_dim)
self.k_linear = nn.Linear(embed_dim, embed_dim)
self.v_linear = nn.Linear(embed_dim, embed_dim)
self.scale = embed_dim ** 0.5
def forward(self, x):
Q = self.q_linear(x)
K = self.k_linear(x)
V = self.v_linear(x)
scores = torch.matmul(Q, K.transpose(-2, -1)) / self.scale
weights = F.softmax(scores, dim=-1)
output = torch.matmul(weights, V)
return output
使用示例:
x = torch.rand(2, 10, 512) # batch_size=2, seq_len=10, embed_dim=512
sa = SelfAttention(embed_dim=512)
out = sa(x) # output: (2, 10, 512)
六、与 Multi-Head Attention 的关系
Self-Attention 是 Multi-Head Attention 的一个“单头版本”。Multi-Head 就是将 Self-Attention 复制多份(通常是 8 或 12),分别使用不同权重做注意力计算,然后拼接起来。
七、优点总结
特性 | 说明 |
---|---|
全局上下文 | 每个词都能与所有其他词交互 |
位置无关 | 不依赖相对/绝对位置,可与位置编码组合使用 |
并行计算 | 所有 token 同时处理,速度快 |
灵活建模 | 根据不同输入动态学习注意力分布 |
八、结语
Self-Attention 是 Transformer 模型的核心灵魂,既实现了强大的表达能力,又兼顾了效率。掌握了它,就为理解整个 Transformer 奠定了基础。