从零理解c3tr模块中的多头自注意力机制：原理剖析与实战指南

1次阅读

没有评论

共计 1920 个字符，预计需要花费 5 分钟才能阅读完成。

自注意力机制是 Transformer 架构的核心，它让模型能够动态地关注输入序列中不同位置的信息。在自然语言处理（NLP）任务中，这一机制尤为重要，因为它能够捕捉长距离依赖关系，解决了传统 RNN 模型在处理长序列时的梯度消失问题。

自注意力机制广泛应用于机器翻译、文本摘要、问答系统等任务中。通过计算每个词与其他词的相关性，模型能够更好地理解上下文，从而生成更准确的输出。

自注意力机制的核心是 Q（Query）、K（Key）、V（Value）三个矩阵。简单来说：

Query：表示当前需要关注的词
Key：表示其他词的标识
Value：表示其他词的实际内容

通过计算 Query 和 Key 的相似度，我们可以得到每个词对其他词的注意力权重，然后用这个权重加权 Value，得到最终的输出。

注意力权重的计算公式如下：

Attention(Q, K, V) = softmax(QK^T / sqrt(d_k)) V

其中，d_k是 Key 的维度。除以 sqrt(d_k) 是为了防止点积结果过大，导致 softmax 函数的梯度消失。

以下是 c3tr 模块中多头自注意力机制的 PyTorch 实现代码：

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

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super(MultiHeadAttention, self).__init__()
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads

        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)

    def forward(self, x):
        batch_size = x.size(0)

        # 线性变换得到 Q, K, V
        Q = self.W_q(x)
        K = self.W_k(x)
        V = self.W_v(x)

        # 分割多头
        Q = Q.view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        K = K.view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        V = V.view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)

        # 计算缩放点积注意力
        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.d_k ** 0.5)
        attn_weights = F.softmax(scores, dim=-1)
        output = torch.matmul(attn_weights, V)

        # 合并多头
        output = output.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)

        # 最终线性变换
        output = self.W_o(output)
        return output

单头注意力只能学习到一种注意力模式，而多头注意力可以让模型同时关注不同位置的不同特征。例如，一个头可能关注句子的主语，另一个头关注谓语，这样模型就能更全面地理解句子结构。