注意力机制的输入

flyfish

注意力机制用于确定序列中每个组成部分相对于其他部分的相对重要性。
绘图源码

import matplotlib.pyplot as plt
from matplotlib.patches import FancyArrowPatchplt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False
# 创建一个新的图形
fig, ax = plt.subplots()# 设置背景颜色
ax.set_facecolor('#ffffff')# 添加文本，使用白色和浅灰色来提高可读性和对比度
ax.text(0.5, 0.8, 'The Essence of Attention', fontsize=24, color='#ecf0f1', ha='center', fontweight='bold')
ax.text(0.2, 0.5, 'Attend to All', fontsize=18, color='#bdc3c7', ha='center', fontweight='semibold')
ax.text(0.8, 0.5, 'Focus on Key Points', fontsize=18, color='#e67e22', ha='center', fontweight='semibold')# 添加箭头，使用亮橙色使箭头更加突出
arrow = FancyArrowPatch((0.45, 0.5), (0.55, 0.5), arrowstyle='->', mutation_scale=20, lw=2, color='#e67e22')
ax.add_patch(arrow)# 设置图形的边界
ax.set_xlim(0, 1)
ax.set_ylim(0, 1)
ax.axis('on')  # 关闭坐标轴# 显示图形
plt.show()

注意力机制的输入

以输入句子 "The cat sat on the mat" 为例：

---

1. 输入的表示 (Tokenization & Embedding)

输入：

句子 "The cat sat on the mat" 通过词汇表转化为 token 索引：
$$\text{tokens}=[0,1,2,3,4,5]$$

嵌入层公式：

词嵌入矩阵 $W_e\in {\mathbb{R}}^{V\times d_{\text{model}}}$（其中 $V$是词汇表大小，$d_{\text{model}}$是嵌入维度）。
对于每个 token，将其映射为一个 $d_{\text{model}}$-维的嵌入向量：
Embedding ( x ) = W e [ x ] , x ∈ { 0 , 1 , 2 , … , V − 1 } \text{Embedding}(x) = W_e[x], \quad x \in \{0, 1, 2, \dots, V-1\} Embedding(x)=We​[x],x∈{0,1,2,…,V−1}

代码对应部分：

self.embedding = nn.Embedding(vocab_size, d_model)
x = self.embedding(src)

---

2. 位置编码 (Positional Encoding)

位置编码的公式如下：
$$P{E}_{\text{pos},2i}=\sin \left(\frac{\text{pos}}{10000^{\frac{2i}{d_{\text{model}}}}}\right)$$
$$P{E}_{\text{pos},2i+1}=\cos \left(\frac{\text{pos}}{10000^{\frac{2i}{d_{\text{model}}}}}\right)$$

• $\text{pos}$ 表示位置索引（例如，第 0 个单词、第 1 个单词）。
• $2i$ 和 $2i+1$ 表示奇偶维度。

位置编码加入到嵌入后：
$$x=\text{Embedding}(x)+PE$$

代码对应部分：

x = self.pos_encoding(self.embedding(src))

---

3. 多头注意力机制 (Multi-Head Attention)

输入：

输入张量 $x\in {\mathbb{R}}^{B\times L\times d_{\text{model}}}$，其中：

• $B$: 批量大小。
• $L$: 序列长度。
• $d_{\text{model}}$: 嵌入维度。

通过一个线性变换生成 Query (Q)、Key (K) 和 Value (V)：
$$Q=x{W}_Q,K=x{W}_K,V=x{W}_V$$
其中 $W_Q,W_K,W_V\in {\mathbb{R}}^{d_{\text{model}}\times d_{\text{model}}}$。

分头操作：
将 $Q,K,V$划分为 $h$个头（头的数量为 $\text{num\_heads}$），每个头的维度为 $d_k=\frac{d_{\text{model}}}{h}$：
$$Q\to Q^{\prime }\in {\mathbb{R}}^{B\times h\times L\times d_k},K\to K^{\prime },V\to V^{\prime }$$

计算注意力权重：

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

• $Q{K}^T$: 计算相似性分数。
• $\sqrt{d_k}$: 缩放因子，避免分数值过大。
• softmax: 将分数转换为概率分布。

多头合并：

将每个头的输出重新拼接，并通过线性变换：
$$\text{MultiHead}(Q,K,V)=\text{Concat}({\text{head}}_1,{\text{head}}_2,\dots ,{\text{head}}_h)W_O$$
其中 $W_O\in {\mathbb{R}}^{d_{\text{model}}\times d_{\text{model}}}$。

代码对应部分：

qkv = self.qkv_linear(x).reshape(B, L, 3, self.num_heads, self.d_k).permute(2, 0, 3, 1, 4)
Q, K, V = qkv[0], qkv[1], qkv[2]
scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
attn = torch.softmax(scores, dim=-1)
out = torch.matmul(attn, V)

---

4. 前馈网络 (Feed Forward Network, FFN)

每个输入通过两个线性变换，中间使用 ReLU 激活函数：
$$\text{FFN}(x)=\text{ReLU}(x{W}_1+b_1)W_2+b_2$$
其中：

• $W_1\in {\mathbb{R}}^{d_{\text{model}}\times d_{ff}}$，$W_2\in {\mathbb{R}}^{d_{ff}\times d_{\text{model}}}$。
  -$d_{ff}$: 前馈网络的中间层维度。

代码对应部分：

return self.linear2(torch.relu(self.linear1(x)))

---

5. 残差连接与层归一化 (Residual Connection & Layer Normalization)

每一层后加入残差连接和归一化：
$$\text{Output}=\text{LayerNorm}(x+\text{SubLayer}(x))$$

代码对应部分：

attn_out = self.dropout(self.attn(self.norm1(x))) + x
ffn_out = self.dropout(self.ffn(self.norm2(attn_out))) + attn_out

---

6. 编码器层与堆叠

每一层的处理：

输入依次经过多头注意力机制和前馈网络：
$$\text{LayerOutput}=\text{FFN}(\text{MultiHeadAttention}(x)+x)+\text{MultiHeadAttention}(x)$$

堆叠多层：

编码器由多个编码器层堆叠：
$$x\to {\text{Layer}}_1(x)\to {\text{Layer}}_2(x)\to \cdots \to {\text{Layer}}_N(x)$$

代码对应部分：

for layer in self.layers:x = layer(x)

---

完整的输入到输出的流程总结

1. 输入文本："The cat sat on the mat" 转为 token 序列：[0, 1, 2, 3, 4, 5]。
2. 嵌入层：将 token 映射为嵌入向量$x\in {\mathbb{R}}^{1\times 6\times 512}$。
3. 位置编码：加上位置信息，得到加权的嵌入表示。
4. 编码器层：
   • 多头注意力机制：计算每个 token 与其他 token 的关系。
   • 前馈网络：对每个 token 的嵌入进一步提取特征。
   • 残差连接与归一化：稳定训练。
5. 堆叠层：经过 6 层编码器层，输出最终表示$\text{Output}\in {\mathbb{R}}^{1\times 6\times 512}$。

---

输出：

print("输出表示形状:", output.shape)  # (1, 6, 512)

输入句子的每个 token 被转化为一个 512 维的上下文感知表示。

完整代码

import torch
import torch.nn as nn
import math# 定义位置编码模块
class PositionalEncoding(nn.Module):def __init__(self, d_model, max_len=5000):"""参数:- d_model: 每个单词嵌入向量的维度- max_len: 最大序列长度"""super(PositionalEncoding, self).__init__()# 初始化一个 (max_len, d_model) 的零张量，用于存储位置编码pe = torch.zeros(max_len, d_model)# 定义序列中的每个位置 [0, 1, 2, ..., max_len-1]position = torch.arange(0, max_len).unsqueeze(1)# 定义公式中 div_term 的部分，用于计算位置编码div_term = torch.exp(torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model))# 位置编码：偶数维度用 sin，奇数维度用 cospe[:, 0::2] = torch.sin(position * div_term)  # 偶数维度pe[:, 1::2] = torch.cos(position * div_term)  # 奇数维度# 注册为 buffer，表示模型内部固定参数，不参与训练self.register_buffer("pe", pe)def forward(self, x):"""前向传播:- 将位置编码与输入张量相加"""x = x + self.pe[:x.size(1)]  # 根据输入序列长度截取位置编码return x# 定义多头注意力机制模块
class MultiHeadAttention(nn.Module):def __init__(self, d_model, num_heads):"""参数:- d_model: 输入向量的维度- num_heads: 注意力头的数量"""super(MultiHeadAttention, self).__init__()self.d_model = d_modelself.num_heads = num_headsself.d_k = d_model // num_heads  # 每个注意力头的维度# 确保 d_model 可以被 num_heads 整除assert d_model % num_heads == 0, "d_model 必须能被 num_heads 整除"# 定义用于生成 Q、K、V 的线性层self.qkv_linear = nn.Linear(d_model, 3 * d_model)# 定义最终输出的线性层self.out_linear = nn.Linear(d_model, d_model)def forward(self, x):"""前向传播:- 输入: x, 形状 [B, L, d_model]"""B, L, _ = x.size()  # 获取批量大小 B 和序列长度 L# 通过线性层生成 Q、K、V，并 reshape 成多头格式qkv = self.qkv_linear(x).reshape(B, L, 3, self.num_heads, self.d_k).permute(2, 0, 3, 1, 4)Q, K, V = qkv[0], qkv[1], qkv[2]  # 分别提取 Q、K、V# 计算注意力得分scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)  # [B, num_heads, L, L]attn = torch.softmax(scores, dim=-1)  # 对最后一个维度进行 softmax 归一化# 通过注意力分数对 V 进行加权求和out = torch.matmul(attn, V)  # [B, num_heads, L, d_k]# 多头输出合并为 [B, L, d_model]out = out.transpose(1, 2).reshape(B, L, self.d_model)return self.out_linear(out)  # 返回线性变换后的结果# 定义前馈网络模块
class FeedForward(nn.Module):def __init__(self, d_model, d_ff):"""参数:- d_model: 输入维度- d_ff: 前馈层的中间层维度"""super(FeedForward, self).__init__()self.linear1 = nn.Linear(d_model, d_ff)  # 第一个全连接层self.linear2 = nn.Linear(d_ff, d_model)  # 第二个全连接层def forward(self, x):"""前向传播:- 使用 ReLU 激活函数"""return self.linear2(torch.relu(self.linear1(x)))# 定义编码器层模块
class EncoderLayer(nn.Module):def __init__(self, d_model, num_heads, d_ff, dropout=0.1):"""参数:- d_model: 输入维度- num_heads: 注意力头的数量- d_ff: 前馈层中间层维度- dropout: dropout 概率"""super(EncoderLayer, self).__init__()self.attn = MultiHeadAttention(d_model, num_heads)  # 多头注意力机制self.ffn = FeedForward(d_model, d_ff)  # 前馈网络self.norm1 = nn.LayerNorm(d_model)  # 层归一化 (注意力)self.norm2 = nn.LayerNorm(d_model)  # 层归一化 (前馈网络)self.dropout = nn.Dropout(dropout)  # dropoutdef forward(self, x):"""前向传播:- 输入 x: [B, L, d_model]"""attn_out = self.dropout(self.attn(self.norm1(x))) + x  # 多头注意力 + 残差连接ffn_out = self.dropout(self.ffn(self.norm2(attn_out))) + attn_out  # 前馈网络 + 残差连接return ffn_out# 定义编码器模块
class Encoder(nn.Module):def __init__(self, vocab_size, d_model, num_heads, d_ff, num_layers, max_len=5000):"""参数:- vocab_size: 词汇表大小- d_model: 嵌入向量维度- num_heads: 注意力头的数量- d_ff: 前馈网络中间层维度- num_layers: 编码器层的数量- max_len: 最大序列长度"""super(Encoder, self).__init__()self.embedding = nn.Embedding(vocab_size, d_model)  # 嵌入层self.pos_encoding = PositionalEncoding(d_model, max_len)  # 位置编码# 创建 num_layers 个编码器层self.layers = nn.ModuleList([EncoderLayer(d_model, num_heads, d_ff) for _ in range(num_layers)])self.norm = nn.LayerNorm(d_model)  # 最终层归一化def forward(self, src):"""前向传播:- 输入 src: [B, L]"""x = self.pos_encoding(self.embedding(src))  # 嵌入 + 位置编码for layer in self.layers:  # 依次通过编码器层x = layer(x)return self.norm(x)# 词汇表与输入句子
vocab = {"The": 0, "cat": 1, "sat": 2, "on": 3, "the": 4, "mat": 5}
sentence = "The cat sat on the mat"
tokens = [vocab[word] for word in sentence.split()]  # 将句子转化为索引列表: [0, 1, 2, 3, 4, 5]# 超参数
vocab_size = len(vocab)
d_model = 512
num_heads = 8
d_ff = 2048
num_layers = 6# 输入处理
src = torch.tensor([tokens])  # 输入形状: [B, L]# 构造编码器
encoder = Encoder(vocab_size, d_model, num_heads, d_ff, num_layers)# 运行编码器
output = encoder(src)# 输出结果
print("输出表示形状:", output.shape)  # 输出形状: (1, 6, 512)