深入解析BART自注意力机制中的QKV图：从原理到实践

1次阅读

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

自注意力机制是 Transformer 架构的核心组件，它允许模型在处理序列数据时动态地关注不同位置的元素。BART（Bidirectional and Auto-Regressive Transformers）是一种基于 Transformer 的预训练语言模型，结合了双向编码和自回归解码的优点，在文本生成和理解任务中表现优异。

自注意力机制通过 Query、Key、Value（QKV）三元组来计算输入序列中各个元素的重要性权重。理解 QKV 图是掌握自注意力机制的关键，也是初学者常见的难点。

自注意力机制中的 QKV 图可以形象地表示输入序列中元素之间的关系。Query、Key 和 Value 都是从输入序列的嵌入表示通过线性变换得到的：

Query（Q）：表示当前需要计算注意力的位置
Key（K）：表示所有可能被关注的位置
Value（V）：包含每个位置的实际信息内容

计算过程分为四步：

计算 Q 和 K 的点积，得到注意力分数
对注意力分数进行缩放（除以√d_k，d_k 是 Key 的维度）
应用 softmax 函数得到归一化的注意力权重
用注意力权重加权求和 Value，得到最终的注意力输出

下面是使用 PyTorch 实现 BART 自注意力机制的简化代码：

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

class SelfAttention(nn.Module):
    def __init__(self, embed_size, heads):
        super(SelfAttention, self).__init__()
        self.embed_size = embed_size
        self.heads = heads
        self.head_dim = embed_size // heads

        # QKV 投影矩阵
        self.values = nn.Linear(self.head_dim, self.head_dim, bias=False)
        self.keys = nn.Linear(self.head_dim, self.head_dim, bias=False)
        self.queries = nn.Linear(self.head_dim, self.head_dim, bias=False)
        self.fc_out = nn.Linear(heads * self.head_dim, embed_size)

    def forward(self, values, keys, queries, mask):
        N = queries.shape[0]  # 批大小
        value_len, key_len, query_len = values.shape[1], keys.shape[1], queries.shape[1]

        # 分割嵌入到多个头
        values = values.reshape(N, value_len, self.heads, self.head_dim)
        keys = keys.reshape(N, key_len, self.heads, self.head_dim)
        queries = queries.reshape(N, query_len, self.heads, self.head_dim)

        # 计算注意力分数
        energy = torch.einsum("nqhd,nkhd->nhqk", [queries, keys])
        if mask is not None:
            energy = energy.masked_fill(mask == 0, float("-1e20"))

        # 计算注意力权重
        attention = torch.softmax(energy / (self.embed_size ** (1/2)), dim=3)

        # 应用注意力权重
        out = torch.einsum("nhql,nlhd->nqhd", [attention, values])
        out = out.reshape(N, query_len, self.heads * self.head_dim)

        # 最终线性变换
        out = self.fc_out(out)
        return out

在实际 NLP 任务中，BART 的自注意力机制可以应用于多种场景：