多头自注意力机制（Multi-Head Self-Attention）是深度学习中一种用于处理序列数据的重要机制，广泛应用于自然语言处理（NLP）和计算机视觉等领域。它最早出现在 Transformer 模型中。

1. Self Attention 机制图示

计算过程如下：

1. 先将输入 Token 进行词嵌入计算；
2. 为每个输入 Token 分别初始化注意力权重参数：w_q、w_k、w_v，输入维度为词嵌入维度，输出维度可指定；
3. 将每个 Token 与各自的 w_q、w_k、w_v 分别计算，产生 q、k、v 三个张量；
4. 以计算第一个字 “我” 的注意力张量为例：
   1. 将 “我” 的 q 张量分别与：“我”、“爱”、“你” 三个的 k 张量进行计算，得出三个 att_score
   2. 将计算得到的 3 个 att_score 经过 softmax 计算，得到三个 att_weight，表示当前 Token 对不同的 Token 的关注程度，也表示不同的 Token 对当前 Token 语义表示的贡献；
   3. 将 3 个 att_weight 与各自的对应的 v 进行计算，此时会得到 3 个向量，将这 3 个向量加起来或者拼接起来作为 “我” 对其他 Toekn 的注意力张量；
   4. 最后，将注意力张量加到 Token 上，即表示当前 Token 结合了上下文的语义表示的向量表示。

计算案例：

import torch
import torch.nn.functional as F
import torch.nn as nn


# 1. 注意力机制基本理解
def test01():

    # 输入 shape=(2, 3), 表示2个词, 每个词3维
    inputs = torch.tensor([[1, 0, 0],
                           [0, 2, 2]], dtype=torch.float)

    # 初始化注意权重, shape=(3, 2)
    w_q = torch.tensor([[1, 0], [1, 0], [0, 1]], dtype=torch.float)
    w_k = torch.tensor([[0, 1], [1, 0], [1, 0]], dtype=torch.float)
    w_v = torch.tensor([[2, 0], [3, 0], [0, 3]], dtype=torch.float)

    # 计算得到每个词的 q、k、v, shape=(2, 2)
    # 第一个词的 v 为: [0, 1]
    # 第二个词的 v 为: [4, 0]
    q = inputs @ w_q  # tensor([[1, 0], [2, 2]])
    k = inputs @ w_k  # tensor([[0, 1], [4, 0]])
    v = inputs @ w_v  # tensor([[2, 0], [6, 6]])

    # 计算注意力分数, shape=(2, 2)
    # 第一个词 v: [2, 0] 对第一个词注意力分数为: 0, 对第二个词的注意力分数为: 4
    # 第二个词 v: [6, 6] 对第一个词注意力分数为: 2, 对第二个词的注意力分数为: 8
    att_score = q @ k.T  # tensor([[0, 4], [2, 8]])

    # 计算得到注意力权重
    # 第一个 v: [0.0180, 0.9820] 对第一个词注意力权重为: 0.0180, 对第二个词的注意力权重为: 0.9820
    # 第二个 v: [0.0025, 0.9975] 对第一个词注意力权重为: 0.0025, 对第二个词的注意力权重为: 0.9975
    attn_weight = F.softmax(att_score, dim=1)  # tensor([[0.0180, 0.9820], [0.0025, 0.9975]])

    # 第一个词的注意力值
    #  [2, 0]、 [6, 6]
    #  0.0180、 0.9820
    # 第一个词对词1和词2的注意力值：[0.036, 0]、[5.8919, 5.8919]

    # 第二个词的注意力值
    #  [2, 0]、 [6, 6]
    #  0.0025、 0.9975
    #  # 第二个词对词1和词2的注意力值：[0.005, 0]、[5.985, 5.985]

    # 最终注意力值
    # 把对每个 v 的注意力加起来, 构成这个词的注意力值
    # [0.036, 0] + [5.892, 5.892]
    # [0.005, 0] + [5.985, 5.985]
    # 得到
    # [5.928, 5.892]
    # [5.990, 5.985]

    # tensor([[2, 0],
    #         [6, 6]])
    # tensor([[0.0180, 0.9820],
    #         [0.0025, 0.9975]])

    # (2, 1, 2)
    # (2, 2, 2)
    attn_values = v[:, None] * attn_weight.T[:, :, None]
    print(attn_values.shape)
    print(attn_values.sum(dim=0))

    # v[:, None]
    # tensor([[[2., 0.],
    #          [2., 0.]],
    #
    #         [[6., 6.],
    #          [6., 6.]]])

    # attn_weight.T[:, :, None]
    # tensor([[[0.0180, 0.0180],
    #          [0.0025, 0.0025]],
    #
    #         [[0.9820, 0.9820],
    #          [0.9975, 0.9975]]])

    # 两个矩阵相加，维度不同需要先进行广播
    # 扩展成相同维度，再进行点乘

    # tensor([[[2 * 0.0180, 0 * 0.0180],
    #          [2 * 0.0025, 0 * 0.0025]],
    #
    #         [[6 * 0.9820, 6 * 0.9820],
    #          [6 * 0.9975, 6 * 0.9975]]])

    # tensor([[[0.036, 0],
    #          [0.005, 0]],
    #
    #         [[5.8919, 5.8919],
    #          [5.985, 5.985]]])

    print(attn_values[0])
    print(attn_values[1])

    # tensor([[2, 0], [6, 6]])
    # tensor([[0.0180, 0.9820], [0.0025, 0.9975]])


# 2. 注意力机制的维度控制
def test02():
    # 输入
    input_shape = (2, 3)
    inputs = torch.randint(0, 5, input_shape, dtype=torch.float)

    # 初始化注意权重
    # 设置自注意力张量维度
    attn_dim = 8

    # 计算得到每个词的 q、k、v
    q_linear = nn.Linear(input_shape[1], attn_dim)
    k_linear = nn.Linear(input_shape[1], attn_dim)
    v_linear = nn.Linear(input_shape[1], attn_dim)

    q = q_linear(inputs)
    k = k_linear(inputs)
    v = v_linear(inputs)

    # 计算注意力分数
    att_score = q @ k.T

    # 计算得到注意力权重
    attn_weight = F.softmax(att_score, dim=1)
    print(attn_weight, attn_weight.shape)

    attn_values = v[:, None] * attn_weight.T[:, :, None]
    print(attn_values.sum(dim=0).shape)


if __name__ == '__main__':
    test01()

在 test02 中，可以通过 attn_dim 变量来控制最终计算的注意力张量的维度.

官网也给出了一段关于自主力机制的封装函数：

import torch
import math
import torch.nn.functional as F
import torch.nn as nn


# 1. 接受的输入为已经计算得出的 query、key、value 矩阵
# 2. query、key、value 的 shape = (batch_size, multi_head, seq_length, dim)
# 3. query、key、value 的 shape = (批量, 头数量, 序列长度, 维度)
# 4. mask :
# 4.1 当使用普通注意力机制时, 需要使用 mask 掩盖未来词, 避免提前看到预测的词
# 4.2 当使用自注意力机制时, 需要对部分为 0 的位置进行特殊处理
# 5. dropout : 对计算的注意力权重进行随机丢弃
def attention_calculate(query, key, value, mask=None, dropout=None):

    # 获得 q、k、v 维度
    query_dim = query.shape[-1]
    # 计算注意力分数
    attn_scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(query_dim)
    # 将 mask == 0 的位置的分数设置极小值, 计算 softmax 时使得该位置的权重接近 0
    if mask is not None:
        attn_scores = attn_scores.masked_fill(mask == 0, -1e-9)
    # 使用 softmax 函数将注意力分数转换为注意力权重
    attn_weight = F.softmax(attn_scores, dim=-1)
    # 对注意力权重值进行随机丢弃
    if dropout is not None:
        attn_weight = dropout(attn_weight)
    # 计算注意力值
    attn_tensor = torch.matmul(attn_weight, value)

    return attn_tensor, attn_weight


if __name__ == '__main__':

    # (批量, 头数量, 序列长度, 维度)
    # inputs = ['我爱你', '我恨你']
    # 2 表示该批次共输入 2 个句子
    # 8 表示对于批次中的每个句子每个字或者词使用8个注意力头
    # 3 表示输入句子长度
    # 256 表示维度

    q = torch.randn(2, 8, 3, 256)
    k = torch.randn(2, 8, 3, 256)
    v = torch.randn(2, 8, 3, 256)

    attn_tensor, attn_weight = attention_calculate(q, k, v, mask=None, dropout=None)
    print('注意力张量:', attn_tensor.shape)
    print('注意力权重:', attn_weight.shape)

2. 多头自注意力机制

多头注意力机制扩展了模型专注不同位置的能力，或者说专注了从不同的角度关注的能力，最终将将多个不同角度关注到的信息（张量）拼接起来。

假设：我们有 3 head，则：每个输入将会初始化 3 套 (q1, k1, v1)、(q2, k2, v2)、(q3, k3, v3)，当计算第一个 Token 的注意张量时，会用第一 Token 的 q1、q2、q3 分别与其他 Token 的 k1、k2、k3 进行缩放点积计算，得到分数之后，softmax 得到注意力权重，在分别乘以其他 Token 的 v1、v2、v3，得到 3 个向量，这 3 个向量可以理解为从不同角度对其他 Token 关联程度的理解。

示例代码：

import torch
import torch.nn.functional as F
import torch.nn as nn

# 1. 单独计算每个注意力头张量
def test01():

    # 固定随机数种子
    torch.random.manual_seed(0)

    # 输入 shape=(2, 3), 表示2个词, 每个词3维
    inputs = torch.randint(1, 5, (2, 3), dtype=torch.float)

    # 初始化注意权重
    head_number = 2
    # 一个头的的注意力权重: shape=(2, 3)
    # 多个头的的注意力权重: shape=(2, 3 * head_number)

    W_q = torch.randint(1, 5, (3, 3 * head_number), dtype=torch.float)
    W_k = torch.randint(1, 5, (3, 3 * head_number), dtype=torch.float)
    W_v = torch.randint(1, 5, (3, 3 * head_number), dtype=torch.float)

    w_q1 = W_q[:, :3]
    w_k1 = W_k[:, :3
    w_v1 = W_v[:, :3]

    w_k2 = W_k[:, 3:]
    w_q2 = W_q[:, 3:]
    w_v2 = W_v[:, 3:]


    # 计算得到每个词的 q、k、v, shape=(2, 2)
    q1 = inputs @ w_q1
    k1 = inputs @ w_k1
    v1 = inputs @ w_v1

    q2 = inputs @ w_q2
    k2 = inputs @ w_k2
    v2 = inputs @ w_v2


    # 计算注意力得分
    att_score1 = q1 @ k1.T
    att_score2 = q2 @ k2.T

    # 计算注意力权重
    attn_weight1 = F.softmax(att_score1, dim=1)
    attn_weight2 = F.softmax(att_score2, dim=1)

    attn_values1 = (v1[:, None] * attn_weight1.T[:, :, None]).sum(dim=0)
    attn_values2 = (v2[:, None] * attn_weight2.T[:, :, None]).sum(dim=0)

    # 将两头注意力张量拼接到一起
    multi_head_attn = torch.cat([attn_values1, attn_values2], dim=1)
    print(multi_head_attn)


# 2. 矩阵计算每个注意力头张量
def test02():

    # 固定随机数种子
    torch.random.manual_seed(0)

    # 输入 shape=(2, 3), 表示2个词, 每个词3维
    inputs = torch.randint(1, 5, (2, 3), dtype=torch.float)

    # 初始化注意权重
    head_number = 2
    # 一个头的的注意力权重: shape=(2, 3)
    # 多个头的的注意力权重: shape=(2, 3 * head_number)

    W_q = torch.randint(1, 5, (3, 3 * head_number), dtype=torch.float)
    W_k = torch.randint(1, 5, (3, 3 * head_number), dtype=torch.float)
    W_v = torch.randint(1, 5, (3, 3 * head_number), dtype=torch.float)

    # 计算得到每个词的 q、k、v, shape=(2, 2)
    Q = inputs @ W_q
    K = inputs @ W_k
    V = inputs @ W_v

    # 计算注意力得分
    att_score = Q @ K.T

    # 计算注意力权重
    attn_weight = F.softmax(att_score, dim=1)
    attn_values = (V[:, None] * attn_weight.T[:, :, None]).sum(dim=0)
    print(attn_values)


if __name__ == '__main__':
    test01()
    print('-' * 30)
    test02()