Transformer的注意力和多头注意力结构（含Pytorch代码实现）

原创已于 2025-10-17 10:23:29 修改 · 1.2k 阅读

CC 4.0 BY-SA版权

文章标签：

于 2025-10-14 15:44:51 首次发布

注意力机制是Transformer的关键机制。传统的循环神经网络（RNN），卷积神经网络（CNN）的一个主要问题是长距离依赖的难以捕捉。RNN只能串行输入，长距离的依赖必须经过多步才能在层级之间传递。而CNN的卷积核的感受野从设计结构上就只能捕捉局部信息。注意力机制通过加权求和的方式，让输入的每个token都能了解其他token的重要性，从而获取全局依赖关系。

三个关键向量：q，k，v

Transformer模型的输入叫做token序列，也就是一串token。每个token也是一个1维向量，元素个数为 $d_{model}$ 。对于每个token，我们有3个全连接的线性层，和token进行线性变换之后，就得到了q，k，v向量。这里q，k，v的具体含义暂时不解释，毕竟为其赋予意义是论文写作者的工作。在这一步，q，k，v的形状和token相同，都为 $d_{model}$ 。

从计算注意力权重到注意力向量

上面我们提到，输入是token的序列。也就是说，对于输入长度为L的token序列，每个token都有对应的q，k，v向量。单个token的注意力权重是要和其他所有token一起计算的。

现在我们计算第i个token的注意力权重。用这个token的q向量q_i和另外一个token的k向量k_j做点乘，得到一个值，我们将其称作注意力分数：

$\text{score}_{ij} = \mathbf{q}_i \mathbf{k}_j^T$

令j从1取到L，就得到了单个token的注意力分数向量。我们可以用矩阵相乘的方式获得这个向量，按行排列所有token的k向量，得到矩阵K，转置后和q向量相乘，得到注意力分数向量。通常要对其做一个缩放处理，对所有元素除以一个常数d_k的平方根：

$$ \mathbf{score}_{i} = \frac{\mathbf{q}_i \mathbf{K}^T}{\sqrt{d_k}} $$

将得到的向量做softmax处理，就得到了注意力权重向量。

$$ \mathbf{\alpha}_i = softmax(\mathbf{score}_i) $$

最后将所有token的v值和注意力权重向量做加权求和，得到这个token的注意力向量：

$\mathbf{z}i = \sum_j \alpha_{ij}\mathbf{v}_j$

多头注意力

卷积神经网络中，卷积核可有多个，多个卷积核能够改变特征图的通道数。同理，注意力机制中可以有多个线性变换，算出不同的注意力向量。在多个注意力向量计算完毕后，将其拼接就成为token的注意力向量。拼接后的向量要和原先的向量大小相同，所以需要在token阶段先进行切分成h份，然后独立进行线性变换（h个线性变换的参数独立），计算完毕后拼接。

在多头注意力中，理论上是token的不同部分进行切分，然后进行线性变换后得到q,k,v向量，然后分别计算注意力后拼接，实际上因为线性变换可统一进行，所以是得到Q，K，V矩阵后再切分。这个在后续代码中会再次讲解。

矩阵形式

单个token的注意力向量计算我们已经掌握了，我们可用矩阵运算简化上述公式。在计算注意力分数向量的公式中，我们把q向量按行排布成矩阵Q，得到一系列的注意力分数向量\alpha，得到注意力分数矩阵A。最后将A按行进行softmax操作得到注意力权重矩阵，其中每一行代表一个token的注意力权重向量，最后将注意力权重矩阵A和v向量构成的矩阵V进行矩阵乘法得到注意力矩阵：

$$ \mathbf{A} = softmax\left(\frac{\mathbf{Q} \mathbf{K}^T}{\sqrt{d_k}}\right) $$

$$ \mathbf{Attention} = \mathbf{A}\mathbf{V} $$

Pytorch实现

在Pytorch中已经有了现成的注意力模块，我们只需要使用

mha = nn.MultiheadAttention(embed_dim=512, num_heads=8,
                            dropout=0.1, batch_first=True)

就可以创建一个token大小为512，注意力头数为8，包含dropout层（禁用概率为0.1），支持batch操作的注意力模块。

下面我们手动实现一个多头注意力模块，首先给出完整代码，然后一步步讲解：

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

class Attention(nn.Module):
    def __init__(self, embed_dimension, num_heads):
        super(Attention, self).__init__()
        assert embed_dimension % num_heads == 0 , "token维度必须能被num_heads整除"

        self.embed_dimension = embed_dimension
        self.num_heads = num_heads
        self.head_dimension = embed_dimension // num_heads # 每个头的q,k,v长度

        # 生成Q，K，V的线性层
        self.W_q = nn.Linear(embed_dimension, embed_dimension)
        self.W_k = nn.Linear(embed_dimension, embed_dimension)
        self.W_v = nn.Linear(embed_dimension, embed_dimension)

        # 输出投影层
        self.W_o = nn.Linear(embed_dimension,embed_dimension)
        self.dropout = nn.Dropout(0.1)
    def forward(self, x, mask=None):
        """
        x: (batch_size, seq_len, embed_dimension)
        mask: (batch_size, 1, 1, seq_len)
        """
        batch_size, seq_len, _ = x.size()

        # 计算Q，K，V
        Q = self.W_q(x)
        K = self.W_k(x)
        V = self.W_v(x)

        # 拆分
        Q = Q.view(batch_size, seq_len, self.num_heads, self.head_dimension).transpose(1, 2)
        K = K.view(batch_size, seq_len, self.num_heads, self.head_dimension).transpose(1, 2)
        V = V.view(batch_size, seq_len, self.num_heads, self.head_dimension).transpose(1, 2)
        # (B, h, L, d_k)

        # 计算注意力分数矩阵
        attn_scores = Q @ K.transpose(-2,-1) / (self.head_dimension ** 0.5)
        # (B, h, L, L)

        if mask is not None:
            attn_scores = attn_scores.masked_fill(mask==0, float('-inf'))

        # softmax
        attn_weights = F.softmax(attn_scores, dim=-1) # (B, h, L, L)
        attn_weights = self.dropout(attn_weights)

        # 计算注意力
        attn_output = attn_weights @ V # (B, h, L, d_k)

        # 超级拼装
        attn_output = attn_output.transpose(1,2).contiguous().view(batch_size, seq_len, self.embed_dimension)

        # 最后线性层结束
        output = self.W_o(attn_output)

        return output, attn_weights

我们首先对输入进行了三个线性变换得到了Q,K,V矩阵。Pytorch的线性层和矩阵乘法都是只针对最后两个维度的，也就是说在更高的维度上可以看作按元素操作，因此我们只需要把batch保持在最前面，头数h放在第二个就可以很方便的应用矩阵乘法了：

Q = self.W_q(x)
K = self.W_k(x)
V = self.W_v(x)

在计算完毕之后，我们根据头数将其拆分成多个，每个token的长度是d_k：

Q = Q.view(batch_size, seq_len, self.num_heads, self.head_dimension).transpose(1, 2)
K = K.view(batch_size, seq_len, self.num_heads, self.head_dimension).transpose(1, 2)
V = V.view(batch_size, seq_len, self.num_heads, self.head_dimension).transpose(1, 2)
# (B, h, L, d_k)

计算注意力分数矩阵：

attn_scores = Q @ K.transpose(-2,-1) / (self.head_dimension ** 0.5)
        # (B, h, L, L)

这里我们添加了一个mask掩码，目的是将值为0的空白填充从注意力分数去掉，我们将其设置为负无穷大，在softmax之后的值为0。

if mask is not None:
      attn_scores = attn_scores.masked_fill(mask==0, float('-inf'))

最后进行softmax操作，乘以V矩阵后将其拼回即可：

 # softmax
        attn_weights = F.softmax(attn_scores, dim=-1) # (B, h, L, L)
        attn_weights = self.dropout(attn_weights)

        # 计算注意力
        attn_output = attn_weights @ V # (B, h, L, d_k)

        # 超级拼装
        attn_output = attn_output.transpose(1,2).contiguous().view(batch_size, seq_len, self.embed_dimension)

        # 最后线性层结束
        output = self.W_o(attn_output)

        return output, attn_weights