深入解析BART自注意力机制中的QKV图：原理与实现优化

1次阅读

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

BART（Bidirectional and Auto-Regressive Transformers）是一种结合了双向编码和自回归解码的预训练模型，广泛应用于文本生成、摘要和翻译等 NLP 任务。与标准 Transformer 不同，BART 的自注意力机制在处理 QKV（Query-Key-Value）时增加了对双向上下文的支持，这使得其在处理长文本时表现更优。

自注意力机制的核心是通过 QKV 矩阵计算输入序列中各个位置的相关性。具体来说：

Query（Q）：表示当前需要关注的位置。
Key（K）：表示序列中其他位置的标识。
Value（V）：包含每个位置的实际信息。

计算公式如下：

$$
\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V
$$

其中，$d_k$ 是 Key 的维度，用于缩放点积结果，防止梯度消失。

与传统 Transformer 相比，BART 的 QKV 计算在编码器和解码器中有所不同：

编码器：完全双向，每个位置可以关注整个输入序列。
解码器：自回归，只能关注当前位置及之前的序列部分。

以下是使用 PyTorch 生成 QKV 矩阵的完整代码示例：

import torch
import torch.nn as nn
import numpy as np

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

        # 确保 embed_size 可以被 heads 整除
        assert self.head_dim * heads == embed_size, "Embed size needs to be divisible by 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, query, mask):
        N = query.shape[0]  # 批次大小
        value_len, key_len, query_len = values.shape[1], keys.shape[1], query.shape[1]

        # 拆分 embedding 到多个头
        values = values.reshape(N, value_len, self.heads, self.head_dim)
        keys = keys.reshape(N, key_len, self.heads, self.head_dim)
        queries = query.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)

        # 应用注意力权重到 Values
        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