Transformer QKV

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
登录后复制