缩放点积注意力机制和多头注意力机制

缩放点积注意力机制
在深度学习领域，缩放点积注意力机制是一种用于计算查询（Query）、键（Key）和值（Value）之间关系的方法。具体来说，在该机制下，输入的查询向量会与所有的键向量做点乘操作来衡量它们之间的相似度得分；这些得分会通过softmax函数转换成概率分布形式表示各个位置的重要性权重；最后利用得到的概率加权求和对应的价值向量获得最终输出[^1]。
为了防止大数值下的梯度消失问题以及稳定训练过程中的激活值范围，通常会对上述提到的点乘结果除以根号d_k(d_k代表模型维度)，即进行了所谓的“缩放”。
import numpy as np

def scaled_dot_product_attention(Q, K, V):
d_k = Q.shape[-1]
scores = np.matmul(Q, K.T) / np.sqrt(d_k)
p_attn = softmax(scores)
output = np.matmul(p_attn, V)
return output, p_attn

def softmax(x):
e_x = np.exp(x - np.max(x))
return e_x / e_x.sum(axis=-1, keepdims=True)
登录后复制
多头注意力机制
多头注意力允许模型在同一序列的不同表征子空间内共同关注不同位置的信息。简单而言就是将原始特征映射到多个低维空间里分别执行一次标准版或改进后的自注意层运算后再拼接起来作为整体输出[^2]。这种设计有助于捕捉更丰富的语义结构，并且可以显著提升性能表现特别是在处理长依赖性和复杂模式识别任务上效果明显优于单头版本。
每个头部都独立工作于各自变换过的Q,K,V矩阵之上从而能够平行化加速整个计算流程同时也增加了参数数量使得网络具备更强表达能力去理解多样化上下文环境变化情况。
class MultiHeadedAttention(nn.Module):
def __init__(self, h, d_model, dropout=0.1):
super().__init__()
assert d_model % h == 0

self.d_k = d_model // h
self.h = h
self.linears = clones(nn.Linear(d_model, d_model), 4)
self.attn = None
self.dropout = nn.Dropout(p=dropout)

def forward(self, query, key, value, mask=None):
if mask is not None:
mask = mask.unsqueeze(1)

nbatches = query.size(0)

query, key, value = \
[l(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)
for l, x in zip(self.linears, (query, key, value))]

x, self.attn = attention(query, key, value, mask=mask,
dropout=self.dropout)

x = x.transpose(1, 2).contiguous()\
.view(nbatches, -1, self.h * self.d_k)

del query
del key
del value

return self.linears[-1](x)

def clones(module, N):
"Produce N identical layers."
return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])

def attention(query, key, value, mask=None, dropout=None):
"Compute 'Scaled Dot Product Attention'"
d_k = query.size(-1)
scores = torch.matmul(query, key.transpose(-2, -1)) \
/ math.sqrt(d_k)

if mask is not None:
scores = scores.masked_fill(mask == 0, -1e9)

p_attn = F.softmax(scores, dim=-1)

if dropout is not None:
p_attn = dropout(p_attn)

return torch.matmul(p_attn, value), p_attn
登录后复制