Transformer QKV
时间: 2025-02-25 10:44:49 浏览: 19
Transformer 中 QKV 的概念
在神经网络中的注意力机制里,Query (查询), Key (键), 和 Value (值) 是三个核心组件。这三个向量用于计算不同位置之间的关系强度。
- Query: 表示当前词项对于其他所有词项的兴趣程度。
- Key: 对应于每个词项的表示形式,用来和其他 Query 进行匹配比较。
- Value: 当某个特定 Query 与 Key 成功配对时所返回的内容或信息[^2]。
具体来说,在多头注意力层内部:
- 输入序列会被映射成相同长度但维度不同的三组向量——Queries, Keys 及 Values;
- 使用缩放点积注意函数来衡量 Queries 和 Keys 之间相似度得分;
- 得分经过 Softmax 归一化处理后作为权重系数加权求和对应的 Values 向量得到最终输出;
这种设计允许模型在同一时间步上关注来自不同位置的信息,并且可以并行执行而无需等待前馈操作完成。
import torch
import torch.nn as nn
class MultiHeadAttention(nn.Module):
def __init__(self, d_model, num_heads):
super(MultiHeadAttention, self).__init__()
assert d_model % num_heads == 0
self.d_k = d_model // num_heads
self.num_heads = num_heads
# Linear layers to project queries, keys and values from input embeddings.
self.q_linear = nn.Linear(d_model, d_model)
self.v_linear = nn.Linear(d_model, d_model)
self.k_linear = nn.Linear(d_model, d_model)
def forward(self, q, k, v, mask=None):
batch_size = q.size(0)
# Perform linear operation and split into heads.
k = self.k_linear(k).view(batch_size, -1, self.num_heads, self.d_k)
q = self.q_linear(q).view(batch_size, -1, self.num_heads, self.d_k)
v = self.v_linear(v).view(batch_size, -1, self.num_heads, self.d_k)
# Transpose to get dimensions [batch_size, num_heads, seq_length, d_k].
k = k.transpose(1, 2)
q = q.transpose(1, 2)
v = v.transpose(1, 2)
# Calculate attention using function we will define next.
scores = ...
output = ...
return output
相关推荐


















