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

在AI领域，Transformer模型及其核心组件——多头注意力机制（MHA）——已经成为一种强大的工具，本文将使用Python的NumPy库来实现这一机制。

什么是多头注意力机制？

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

MHA的关键优势

1. 表示多样性：每个头可以学习输入的不同特征，增加了模型的表示能力。
2. 并行处理：多头可以同时工作，提高了模型的计算效率。
3. 灵活性和适应性：头的数量和大小可以根据不同的任务进行调整。

MHA的简化实现

1. 定义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)

2. 实现缩放点积注意力

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

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

3. 多头注意力的整合

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

def multi_head_attention(queries, keys, values, num_heads, head_dim, mask=None):
    batch_size = queries.shape[0]
    seq_len = queries.shape[1]
    features = queries.shape[2]

    q = queries.reshape(batch_size, seq_len, num_heads, head_dim)
    k = keys.reshape(batch_size, seq_len, num_heads, head_dim)
    v = values.reshape(batch_size, seq_len, num_heads, head_dim)

    scale = head_dim ** 0.5
    outputs, _ = scaled_dot_product_attention(q, k, v, scale, mask)

    outputs = outputs.reshape(batch_size, seq_len, num_heads * head_dim)
    return outputs

示例代码

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

# 示例数据
queries = np.random.rand(2, 4, 16)  # (batch_size, seq_len, features)
keys = np.random.rand(2, 4, 16)     # (batch_size, seq_len, features)
values = np.random.rand(2, 4, 16)   # (batch_size, seq_len, features)

# 定义参数
num_heads = 4
head_dim = 4

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

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

Code

代码已上传至：Machine Learning and Deep Learning Algorithms with NumPy
此项目包含更多机器学习和深度学习算法的NumPy实现，供大家学习参考使用，欢迎star~

备注

个人水平有限，有问题随时交流~