NLP学习记录十：多头注意力

news2025/5/15 5:08:39

一、单头注意力

单头注意力的大致流程如下：

① 查询编码向量、键编码向量和值编码向量分别经过自己的全连接层（Wq、Wk、Wv）后得到查询Q、键K和值V；

② 查询Q和键K经过注意力评分函数（如：缩放点积运算）得到值权重矩阵；

③ 权重矩阵与值向量相乘，得到输出结果。

图1 单头注意力模型

二、多头注意力

2.1 使用多头注意力的意义

看了一些对多头注意力机制解释的视频，我自己的浅显理解是：在实践中，我们会希望查询Q能够从给定内容中尽可能多地匹配到与自己相关的语义信息，从而得到更准确的预测输出。而多头注意力将查询、键和值分成不同的子空间表示（representation subspaces）（有点类似于子特征？），使得匹配过程更加细化。

2.2 代码实现

也许直接看代码能更快地理解这个过程：

import torch
from torch import nn
from attentionScore import DotProductAttention

# 多头注意力模型
class MultiHeadAttention(nn.Module):
    def __init__(self, key_size, query_size, value_size, num_hiddens,
                 num_heads, dropout, bias=False, **kwargs):
        super(MultiHeadAttention, self).__init__(**kwargs)
        self.num_heads = num_heads
        self.attention = DotProductAttention(dropout)
        self.W_q = nn.Linear(query_size, num_hiddens, bias=bias)
        self.W_k = nn.Linear(key_size, num_hiddens, bias=bias)
        self.W_v = nn.Linear(value_size, num_hiddens, bias=bias)
        self.W_o = nn.Linear(num_hiddens, num_hiddens, bias=bias)

    # queries：(batch_size，查询的个数，query_size)
    # keys：(batch_size，“键－值”对的个数，key_size)
    # values：(batch_size，“键－值”对的个数，value_size)
    def forward(self, queries, keys, values, valid_lens):
        # queries，keys，values的形状：(batch_size，查询或者“键－值”对的个数，num_hiddens)
        queries = self.W_q(queries)
        keys = self.W_k(keys)
        values = self.W_v(values)

        # 经过变换后，输出的queries，keys，values的形状：(batch_size*num_heads，查询或者“键－值”对的个数，num_hiddens/num_heads)
        queries = transpose_qkv(queries, self.num_heads)
        keys = transpose_qkv(keys, self.num_heads)
        values = transpose_qkv(values, self.num_heads)

        # valid_lens的形状：(batch_size，)或(batch_size，查询的个数)
        if valid_lens is not None:
            # 在轴0，将第一项（标量或者矢量）复制num_heads次，然后如此复制第二项，然后诸如此类。
            valid_lens = torch.repeat_interleave(valid_lens, repeats=self.num_heads, dim=0)

        # output的形状：(batch_size*num_heads，查询的个数，num_hiddens/num_heads)
        output = self.attention(queries, keys, values, valid_lens)

        # output_concat的形状：(batch_size，查询的个数，num_hiddens)
        output_concat = transpose_output(output, self.num_heads)
        return self.W_o(output_concat)

# 为了多注意力头的并行计算而变换形状
def transpose_qkv(X, num_heads):
    # 输入X的形状:(batch_size，查询或者“键－值”对的个数，num_hiddens)
    # 输出X的形状:(batch_size，查询或者“键－值”对的个数，num_heads，num_hiddens/num_heads)
    X = X.reshape(X.shape[0], X.shape[1], num_heads, -1)

    # 输出X的形状:(batch_size，num_heads，查询或者“键－值”对的个数, num_hiddens/num_heads)
    X = X.permute(0, 2, 1, 3)

    # 最终输出的形状:(batch_size*num_heads,查询或者“键－值”对的个数, num_hiddens/num_heads)
    return X.reshape(-1, X.shape[2], X.shape[3])

# 逆转transpose_qkv函数的操作
def transpose_output(X, num_heads):
    X = X.reshape(-1, num_heads, X.shape[1], X.shape[2])
    X = X.permute(0, 2, 1, 3)
    return X.reshape(X.shape[0], X.shape[1], -1)

可以发现，前面的处理流程和单头注意力的第①步是一样的，都是使用全连接层计算查询Q、键K、值V。但在进行点积运算之前，模型使用transpose_qkv函数对QKV进行了切割变换，下图可以帮助理解这个过程：