cos大壮 2025-02-28 14:32 浙江
手撕Query-Key-Value矩阵。
今儿咱们聊一下，Transformer中的Query-Key-Value 矩阵计算与注意力权重生成，还想了一个自认为比较美的标题：手撕Query-Key-Value矩阵的数学之美。
好了，下面咱们聊聊 Transformer 中 Query-Key-Value 矩阵计算与注意力权重生成的原理，以及一个手算例子。
在一次聚会中，每个人都有三张卡片：
Query 卡片：代表“我想知道和谁聊得来”，相当于提出问题；
Key 卡片：代表“我是谁”，用来标记每个人的特性；
Value 卡片：代表“我的真实信息或观点”。
在 Transformer 中，每个词（或 token）都会经过三个不同的线性变换，分别得到 Query、Key 和 Value 向量。
接下来，我们拿某个词的 Query 去与所有词的 Key 进行匹配（计算点积），匹配得分越高，说明这两个词的“特性”越接近。
为了让这些得分更稳定，还会除以一个缩放因子（Key 向量的维度平方根），然后再经过 softmax 得到一组概率，也就是注意力权重。
最后，用这些权重对所有词的 Value 进行加权求和，就得到了这个词的新表示——它综合了其他词对它的影响。
细节原理
输入及线性变换
设输入矩阵为（每一行对应一个词的嵌入向量），通过三个不同的线性变换得到：
其中、和是模型学习到的参数矩阵。
注意力权重的计算
相似度计算
对于任意一个词的 Query 向量，计算它与所有词的 Key 向量的点积：
缩放
为了防止点积值过大（尤其在高维空间中），将结果除以缩放因子：
其中是 Key 向量的维度。
归一化（softmax）
将缩放后的得分经过 softmax 函数，得到注意力权重：
这表示词对词的关注程度。
加权求和生成输出
用注意力权重对所有词的 Value 向量加权求和，得到词的输出表示：
用矩阵形式表达整个注意力层：
推理总结
输入转换：将输入嵌入转换为 Query、Key、Value 三个向量。
相似度计算与缩放：用 Query 与 Key 做点积并除以。
归一化：通过 softmax 得到归一化的注意力权重。
信息聚合：用这些权重对 Value 向量进行加权求和，输出综合上下文信息的表示。
手撕计算
假设我们有两个词，令它们的嵌入向量维度为 2，为了简化计算，我们令所有线性变换矩阵均为单位矩阵（即直接让）。
设定输入
令两个词的输入向量为：
因此：
设 Key 向量维度，则缩放因子为。
计算 Q 与 K 的点积
计算得：
缩放后的得分矩阵为：
计算注意力权重（softmax）
对每一行分别计算 softmax：
对于第一行：
计算指数：
；
。
求和：
。
得到注意力权重：
对于第二行（过程类似）：
得到注意力权重：
加权求和生成输出
用注意力权重对 Value 向量求和：
第一个词的输出：
其中,：
第二个词的输出：
完整案例
这里，咱们全流程展示 Transformer 中 Query-Key-Value 矩阵计算与注意力权重生成的应用。
整个案例包含了用虚拟数据集训练模型的 PyTorch 代码，并绘制了包含至少四个数据分析图形的综合图（图中包含训练损失变化、预测曲线、注意力权重热力图以及误差分布图等）。
我们构造一个简单的时序预测任务。假设生成的虚拟数据是一条带噪声的正弦曲线，模型的任务是：给定前 N 个时刻的数值，预测下一个时刻的数值。
这种任务可以模拟时间序列预测问题，同时能直观展示模型关注序列中哪些位置的信息（即注意力权重分布）。
大概3个步骤：
数据生成：利用正弦函数生成平滑曲线，再加上随机噪声。
窗口切分：设定固定的窗口大小（例如 10），将数据按滑动窗口切分为输入序列与预测目标。
数据归一化：对输入数据进行归一化，使模型训练更加稳定。
虚拟数据集的构造代码将在后面的 PyTorch 实现中给出，数据集会同时返回序列数据和目标值。
模型构建
我们构造一个简化版的 Transformer Encoder 模块，仅保留自注意力层（Self-Attention）部分，用于对输入序列进行建模。模型主要包含以下模块：
嵌入层：将输入序列转换为高维表示（在本例中我们直接使用原始数值构成的向量，也可引入位置编码）。
自注意力层：实现 Query-Key-Value 的计算过程，对输入序列进行信息聚合。
前馈网络：经过自注意力层后，通过一层简单的全连接层输出预测值。
import math
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt
from matplotlib import gridspec
# 固定随机种子，保证结果可复现
torch.manual_seed(42)
np.random.seed(42)
# 3.1 定义自注意力层（Attention Layer）
class SelfAttention(nn.Module):
def __init__(self, embed_dim):
super(SelfAttention, self).__init__()
self.embed_dim = embed_dim
# 定义 Query、Key、Value 的线性变换
self.W_Q = nn.Linear(embed_dim, embed_dim, bias=False)
self.W_K = nn.Linear(embed_dim, embed_dim, bias=False)
self.W_V = nn.Linear(embed_dim, embed_dim, bias=False)
def forward(self, x):
# x: [batch_size, seq_len, embed_dim]
Q = self.W_Q(x) # [batch_size, seq_len, embed_dim]
K = self.W_K(x)
V = self.W_V(x)
# 计算 Q 与 K 的点积
scores = torch.matmul(Q, K.transpose(-2, -1)) # [batch_size, seq_len, seq_len]
# 缩放因子：除以 sqrt(d_k)
scores = scores / math.sqrt(self.embed_dim)
# 应用 softmax 得到注意力权重
attn_weights = torch.softmax(scores, dim=-1) # [batch_size, seq_len, seq_len]
# 对 V 加权求和
output = torch.matmul(attn_weights, V) # [batch_size, seq_len, embed_dim]
return output, attn_weights
# 3.2 定义 Transformer 模型（仅包含一层自注意力和简单前馈）
class TransformerPredictor(nn.Module):
def __init__(self, seq_len, embed_dim, hidden_dim):
super(TransformerPredictor, self).__init__()
self.seq_len = seq_len
self.embed_dim = embed_dim
# 输入嵌入层（本例中直接将单个数值映射到高维空间）
self.embedding = nn.Linear(1, embed_dim)
# 位置编码（简化版：使用正弦、余弦编码）
self.positional_encoding = self._generate_positional_encoding(seq_len, embed_dim)
# 自注意力层
self.attention = SelfAttention(embed_dim)
# 前馈层，将注意力层的输出映射到预测值
self.fc = nn.Sequential(
nn.Linear(embed_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, 1)
)
def _generate_positional_encoding(self, seq_len, embed_dim):
# 生成位置编码矩阵，[seq_len, embed_dim]
pe = np.zeros((seq_len, embed_dim))
position = np.arange(0, seq_len)[:, np.newaxis]
div_term = np.exp(np.arange(0, embed_dim, 2) * -(np.log(10000.0) / embed_dim))
pe[:, 0::2] = np.sin(position * div_term)
pe[:, 1::2] = np.cos(position * div_term)
pe = torch.tensor(pe, dtype=torch.float)
return pe.unsqueeze(0) # shape: [1, seq_len, embed_dim]
def forward(self, x):
# x: [batch_size, seq_len, 1]
batch_size = x.size(0)
# 嵌入
x_embed = self.embedding(x) # [batch_size, seq_len, embed_dim]
# 加入位置编码
x_embed = x_embed + self.positional_encoding.to(x.device)
# 自注意力层
attn_output, attn_weights = self.attention(x_embed) # attn_output: [batch_size, seq_len, embed_dim]
# 对序列最后一个时间步的输出进行预测（也可以采用平均池化）
final_feature = attn_output[:, -1, :] # [batch_size, embed_dim]
prediction = self.fc(final_feature) # [batch_size, 1]
return prediction, attn_weights
数据集构造与预处理
生成一条带噪声的正弦曲线，并利用滑动窗口切分为多个样本。
生成正弦曲线，并在每个点加上随机噪声；
设定窗口大小（例如 10），将序列切分成输入序列和对应的目标值（下一个时刻的值）；
将数据集划分为训练集和测试集。
def generate_sin_data(seq_length=200, noise_std=0.1):
# 生成正弦曲线数据
x = np.linspace(0, 4 * np.pi, seq_length)
y = np.sin(x) + np.random.normal(scale=noise_std, size=seq_length)
return y
def create_dataset(data, window_size):
X, Y = [], []
for i in range(len(data) - window_size):
X.append(data[i:i+window_size])
Y.append(data[i+window_size])
X = np.array(X)
Y = np.array(Y)
return X, Y
# 生成数据
data = generate_sin_data(seq_length=300, noise_std=0.15)
window_size = 10
X_data, Y_data = create_dataset(data, window_size)
# 划分训练集和测试集
split_ratio = 0.8
split_index = int(len(X_data) * split_ratio)
X_train = X_data[:split_index]
Y_train = Y_data[:split_index]
X_test = X_data[split_index:]
Y_test = Y_data[split_index:]
# 转换为 PyTorch 张量，并调整形状：[batch_size, seq_len, 1]
X_train_tensor = torch.tensor(X_train, dtype=torch.float).unsqueeze(-1)
Y_train_tensor = torch.tensor(Y_train, dtype=torch.float).unsqueeze(-1)
X_test_tensor = torch.tensor(X_test, dtype=torch.float).unsqueeze(-1)
Y_test_tensor = torch.tensor(Y_test, dtype=torch.float).unsqueeze(-1)
我们构造了一个简单的回归问题数据集，每个样本输入为长度为 10 的时序数据，目标为下一个时刻的数值。
模型训练与结果记录
接下来定义训练函数、损失函数和优化器，训练模型，并记录训练过程中的损失变化、预测结果以及注意力权重信息，以便后续进行可视化分析。
# 超参数设置
embed_dim = 32
hidden_dim = 64
num_epochs = 300
learning_rate = 0.005
# 实例化模型
model = TransformerPredictor(seq_len=window_size, embed_dim=embed_dim, hidden_dim=hidden_dim)
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
# 记录训练过程中的损失和注意力权重
train_losses = []
all_attn_weights = None# 用于保存最后一批的注意力权重
model.train()
for epoch in range(num_epochs):
optimizer.zero_grad()
output, attn_weights = model(X_train_tensor)
loss = criterion(output, Y_train_tensor)
loss.backward()
optimizer.step()
train_losses.append(loss.item())
if (epoch+1) % 50 == 0:
print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}")
# 保存最后一次的注意力权重（仅取一批样本的第一个）
all_attn_weights = attn_weights[0].detach().cpu().numpy()
# 预测测试集
model.eval()
with torch.no_grad():
pred_test, _ = model(X_test_tensor)
pred_test = pred_test.squeeze().cpu().numpy()
Y_test_np = Y_test_tensor.squeeze().cpu().numpy()
在训练过程中，我们记录了每个 epoch 的训练损失，同时每隔一定周期打印当前损失值。最后保存了最后一批样本的注意力权重，用于后续分析。
数据可视化与分析
为了全面展示模型的训练和预测效果，绘制了：
训练损失曲线：展示模型训练过程中损失函数值的下降趋势，反映模型的学习情况。
测试集预测曲线：将测试集的真实值与预测值进行对比，直观展示模型的预测效果。
注意力权重热力图：可视化自注意力层中注意力权重的分布，揭示模型在不同时间步之间的关注程度。
误差分布图：展示预测值与真实值之间的残差分布，帮助判断模型是否存在系统性偏差或噪声问题。
# 绘图：构造一个包含 4 个子图的综合图形
plt.figure(figsize=(16, 12))
gs = gridspec.GridSpec(2, 2)
# 子图1：训练损失曲线
ax1 = plt.subplot(gs[0, 0])
ax1.plot(range(num_epochs), train_losses, color='mediumseagreen', linewidth=2)
ax1.set_title("Training Loss Curve", fontsize=14, color='navy')
ax1.set_xlabel("Epoch")
ax1.set_ylabel("MSE Loss")
ax1.grid(True, linestyle='--', alpha=0.6)
ax1.text(0.5, 0.9, "Reflects the trend of loss reduction during model training", transform=ax1.transAxes, fontsize=12, color='purple')
# 子图2：测试集预测曲线对比
ax2 = plt.subplot(gs[0, 1])
# 绘制真实值曲线
ax2.plot(Y_test_np, label="True Values", color='crimson', linewidth=2, marker='o')
# 绘制预测值曲线
ax2.plot(pred_test, label="Predicted Values", color='dodgerblue', linewidth=2, linestyle='--', marker='x')
ax2.set_title("Test Set Prediction Comparison", fontsize=14, color='navy')
ax2.set_xlabel("Sample Index")
ax2.set_ylabel("Value")
ax2.legend(fontsize=12)
ax2.grid(True, linestyle='--', alpha=0.6)
ax2.text(0.5, 0.9, "Shows the model's prediction performance on the test set", transform=ax2.transAxes, fontsize=12, color='darkgreen')
# 子图3：注意力权重热力图（以最后一次训练得到的注意力权重为例）
ax3 = plt.subplot(gs[1, 0])
if all_attn_weights isnotNone:
cax = ax3.imshow(all_attn_weights, cmap='plasma', aspect='auto')
plt.colorbar(cax, ax=ax3)
ax3.set_title("Self-Attention Weight Heatmap", fontsize=14, color='navy')
ax3.set_xlabel("Sequence Position (Key)")
ax3.set_ylabel("Sequence Position (Query)")
ax3.text(0.5, 0.9, "Reflects the model's attention distribution across time steps", transform=ax3.transAxes, fontsize=12, color='brown')
else:
ax3.text(0.5, 0.5, "Attention weights not saved", ha='center', va='center')
# 子图4：预测误差分布图（残差直方图）
ax4 = plt.subplot(gs[1, 1])
residuals = pred_test - Y_test_np
ax4.hist(residuals, bins=20, color='mediumorchid', edgecolor='black', alpha=0.8)
ax4.set_title("Prediction Residual Distribution", fontsize=14, color='navy')
ax4.set_xlabel("Residual")
ax4.set_ylabel("Frequency")
ax4.grid(True, linestyle='--', alpha=0.6)
ax4.text(0.5, 0.9, "Used to evaluate the distribution of prediction errors", transform=ax4.transAxes, fontsize=12, color='darkred')
plt.suptitle("Transformer Query-Key-Value Attention Mechanism Application Case Data Analysis", fontsize=18, color='darkblue')
plt.tight_layout(rect=[0, 0.03, 1, 0.95])
plt.show()
训练损失曲线：随着 epoch 的增加，损失值逐渐降低，表明模型在不断学习数据的特征；若曲线平稳下降且趋于收敛，则说明学习率和模型容量设置合理。
测试集预测曲线：通过比较真实值和预测值，观察模型是否在新数据上具有较好的泛化能力；若二者重合程度高，说明模型拟合效果较好。
注意力权重热力图：热力图中每个数值代表某一 Query 与相应 Key 的关注程度，颜色越深代表权重越大；从图中可以直观观察模型在序列中不同位置之间的信息传递情况。
预测残差分布图：直方图展示了预测误差的分布情况，若误差呈正态分布且均值接近 0，则表明模型没有明显的系统偏差；若存在偏移或尖峰，则需关注数据噪声或模型欠拟合/过拟合问题。
算法优化点与调参流程
模型结构改进
多头注意力（Multi-Head Attention）：当前模型仅使用单头自注意力，多头注意力能够让模型从不同子空间捕捉信息，提升表达能力。可在后续扩展时引入多个注意力头，并对每个头的输出进行拼接或加权求和。
位置编码优化：除了简单的正弦和余弦位置编码，还可以尝试使用可学习的位置编码或相对位置编码，这对序列建模任务往往能带来更好效果。
前馈网络加深：现有前馈网络较简单，可考虑增加层数、采用残差连接和层归一化（Layer Normalization），使得模型更深、学习到更复杂的特征。
数据预处理与增强
数据归一化：对输入数据做标准化或归一化处理，能帮助模型更快收敛。
数据扩充：对于时序数据，可考虑采用平移、缩放、加噪声等方式进行数据扩充，提升模型的泛化能力。
损失函数与正则化
自适应损失函数：针对不同任务，可尝试 MAE、Huber Loss 等鲁棒性更好的损失函数。
正则化方法：在模型中加入 Dropout、权重衰减（L2 正则化）等方法，防止过拟合。
Early Stopping：根据验证集损失监控训练过程，提前终止训练，防止模型在训练集上过拟合。
学习率调度
学习率衰减：采用学习率衰减策略（如 StepLR、CosineAnnealingLR），在训练后期逐渐降低学习率，有助于模型在局部最优处更稳定收敛。
Warm-up 策略：在训练初期使用较小的学习率，并逐步升高至设定值，能够稳定训练过程，尤其是在深度 Transformer 模型中尤为重要。
调参流程
初始参数设定：选择一个相对较大的模型规模（embed_dim、hidden_dim）和较小的学习率，先验证模型能否收敛。例如：embed_dim=32、hidden_dim=64、初始学习率 0.005、batch size 32。
固定基础参数，先调学习率：观察训练损失曲线与验证集表现，尝试调整学习率（例如 0.001、0.005、0.01 等），观察不同学习率下的损失下降速度与稳定性。可采用学习率衰减或 Warm-up 策略，验证哪种方式更适合当前任务。
调整模型容量：通过增大 embed_dim、hidden_dim、增加注意力头数等方式，测试模型容量对预测效果的影响。注意模型过大可能导致过拟合，因此需同时观察验证集误差，必要时加入正则化方法（Dropout、L2正则化）。
优化模型结构：尝试引入残差连接、层归一化以及更深的前馈网络，观察模型稳定性和训练速度。对比单头注意力与多头注意力在任务中的表现，选择最优结构。
数据预处理与扩充：试验数据归一化的不同方法（如 min-max scaling、z-score 标准化），观察其对模型训练收敛性的影响。利用数据扩充技术（如加噪声、时间平移），提升模型鲁棒性，并分析扩充前后模型在测试集上的表现差异。
以上面代码为例，可按照以下调参步骤逐步优化模型：
Step1：初始模型训练：使用 embed_dim=32，hidden_dim=64，学习率=0.005，训练 300 个 epoch，观察训练损失和测试预测效果。
Step2：学习率调优：若训练损失下降较慢，可尝试将学习率调大到 0.01，或采用学习率 warm-up 策略；若出现震荡，则降低学习率至 0.001。
Step3：模型容量调整：试验将 embed_dim 提升到 64，hidden_dim 调整到 128，并引入多头注意力（例如 4 个头），观察模型在验证集上的泛化效果。
Step4：正则化加入：加入 Dropout 层（如 dropout rate=0.2~0.5）以及 L2 正则化，监控训练与验证损失变化，防止过拟合。
Step5：融合多种改进策略：综合多头注意力、位置编码优化、前馈网络加深、正则化及学习率调度等改进，找到最佳参数组合。
实际调参过程中，需多次实验并结合领域知识和数据特点进行调整，确保模型既有较强表达能力，又能保持良好的泛化性能。
最后
希望可以帮助大家深入理解 Transformer 中 Query-Key-Value 矩阵计算与注意力权重生成的原理，当然在实际项目中，大家需要灵活应用这些方法，取得更好的效果。
大家有问题可以直接在评论区留言即可~
喜欢本文的朋友可以收藏、点赞、转发起来！
