多头注意力机制公式

多头注意力机制的数学公式
多头注意力机制的核心在于通过多个不同的线性变换来计算输入序列的不同表示形式，并将这些不同头部的结果连接起来形成最终输出。以下是其具体实现过程：
查询、键和值的定义
给定查询 ( Q )，键 ( K ) 和值 ( V )，它们分别经过线性变换得到各自对应的矩阵表示[^2]。
[
Q = XW_Q,\quad K = XW_K,\quad V = XW_V
]
其中，( W_Q, W_K, W_V ) 是可学习参数矩阵，用于映射原始特征到新的空间。
单头自注意力计算
单头自注意力可以通过缩放点积注意（Scaled Dot-Product Attention）来计算:
[
\text{Attention}(Q,K,V)=\text{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}\right)V
]
这里 ( d_k ) 表示键向量维度大小，用来防止点积过大而导致梯度消失或爆炸问题。
多头注意力组合
为了捕获更丰富的信息，在多头注意力中会并行执行多次上述操作，每次使用独立的学习权重集。最后将各头结果拼接并通过另一个投影层汇总成单一输出:
[
\text{MultiHead}(Q,K,V) = [\text{head}_1,...,\text{head}_h]\cdot W_O,
]
其中，
[
\text{head}_i=\text{Attention}(QW_i^Q ,KW_i^K ,VW_i^V )
]
这里的 ( h ) 表示头的数量；( W_I^{Q}, W_I^{K}, W_I^{V} ) 对应第 i 个头中的转换矩阵；而 ( W_O ) 则是用来融合所有头输出的一个额外全连接网络参数矩阵。
import torch
import math

def scaled_dot_product_attention(query, key, value):
dk = query.size()[-1]
scores = torch.matmul(query, key.transpose(-2,-1)) / math.sqrt(dk)
attn_weights = torch.nn.functional.softmax(scores,dim=-1)
context_vector = torch.matmul(attn_weights,value)
return context_vector,attn_weights

class MultiHeadedSelfAttention(torch.nn.Module):
def __init__(self,d_model,num_heads):
super(MultiHeadedSelfAttention,self).__init__()
assert d_model % num_heads ==0,"Model dimension must be divisible by number of heads."

self.d_head = int(d_model/num_heads)
self.num_heads=num_heads

self.query_projection=torch.nn.Linear(d_model,d_model,bias=False)
self.key_projection=torch.nn.Linear(d_model,d_model,bias=False)
self.value_projection=torch.nn.Linear(d_model,d_model,bias=False)

self.final_linear_layer=torch.nn.Linear(d_model,d_model,bias=False)

def forward(self,x):
batch_size=x.shape[0]

q=self.query_projection(x).view(batch_size,-1,self.num_heads,self.d_head).transpose(1,2)
k=self.key_projection(x).view(batch_size,-1,self.num_heads,self.d_head).transpose(1,2)
v=self.value_projection(x).view(batch_size,-1,self.num_heads,self.d_head).transpose(1,2)

outputs,_=scaled_dot_product_attention(q,k,v)
concatenated_outputs=outputs.transpose(1,2).contiguous().view(batch_size,-1,self.num_heads*self.d_head)
final_output=self.final_linear_layer(concatenated_outputs)
return final_output
登录后复制