多头注意力机制（MHA）的NumPy实现

type

status

date

slug

summary

什么是多头注意力机制？

MHA是Transformer架构中的关键组成部分，它允许模型在多个位置并行地捕捉输入序列的不同方面。与传统的注意力机制相比，MHA通过将输入数据分割成多个头，每个头学习不同的表示，从而增强了模型的表达能力。

MHA的关键优势

表示多样性：每个头可以学习输入的不同特征，增加了模型的表示能力。

并行处理：多头可以同时工作，提高了模型的计算效率。

灵活性和适应性：头的数量和大小可以根据不同的任务进行调整。

MHA的简化实现

定义Softmax函数首先，我们需要一个Softmax函数来计算注意力权重：


import numpy as np
def softmax(x, axis=-1):
    exps = np.exp(x - np.max(x, axis=axis, keepdims=True))
    return exps / exps.sum(axis=axis, keepdims=True)

实现缩放点积注意力

接着，我们实现计算缩放点积注意力的函数：


def attention(q,k,v,scale,mask=None):
    scores = np.matmul(q,k.transpose(0,1,3,2))/scale
    if mask is not None:
        scores+=(mask==0)*-1e9 
    attention_w = softmax(scores,axis=-1)
    output = np.matmul(attention_w,v)
    return output

多头注意力的整合

最后，我们将实现多头注意力机制，包括分割、并行处理和合并头：


def multi_head_attention(query,key,value,num_heads,head_dim,mask=None):
    bs = query.shape[0]
    q = query.reshape(bs,-1,num_heads,head_dim).transpose(0,2,1,3) 
    k = key.reshape(bs,-1,num_heads,head_dim).transpose(0,2,1,3)
    v = value.reshape(bs,-1,num_heads,head_dim).transpose(0,2,1,3)

    scale = head_dim**0.5 
    attention_out = attention(q,k,v,scale,mask) 
    out = attention_out.transpose(0,2,1,3).reshape(bs,-1,num_heads*head_dim) 
    return out

示例代码

假设我们有一些查询、键和值的示例数据，并希望应用多头注意力机制：


# 示例数据
batch_size = 2
seq_len = 5
d_model = 16  # 总的特征维度
num_heads = 4
head_dim = d_model // num_heads  # 每个头的特征维度

queries = np.random.rand(batch_size, seq_len, d_model)
keys = np.random.rand(batch_size, seq_len, d_model)
values = np.random.rand(batch_size, seq_len, d_model)

# 应用多头注意力机制
output = multi_head_attention(queries, keys, values, num_heads, head_dim)

print("输入形状:", queries.shape)
print("输出形状:", output.shape)