深入解析BART自注意力机制架构图：从原理到高效实现

1次阅读

没有评论

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

自注意力机制（Self-Attention）是 Transformer 架构的核心组件，它通过计算序列中每个位置与其他位置的关联权重，实现全局依赖建模。与传统 RNN 的串行计算相比，自注意力具有三大优势：

并行化能力 ：可同时计算所有位置的注意力权重
长程依赖捕获 ：直接建模任意距离的 token 关系
可解释性 ：注意力权重可视化展示模型关注点

BART（Bidirectional and Auto-Regressive Transformers）作为改进的 Transformer 变体，在自注意力机制上做出重要调整：

编码器 - 解码器不对称设计
编码器使用双向注意力（类似 BERT）
解码器采用带掩码的自回归注意力（类似 GPT）
注意力掩码策略优化
编码器：全可见注意力矩阵
解码器：严格左下三角掩码矩阵
位置编码改进
采用学习式位置编码而非固定公式
在嵌入层合并 token 和位置信息

下图展示 BART 的注意力计算流程（此处描述架构图内容）：

graph TD
    A[输入嵌入] --> B[加位置编码]
    B --> C{编码器 / 解码器?}
    C -->| 编码器 | D[双向自注意力]
    C -->| 解码器 | E[掩码自注意力]
    D --> F[前馈网络]
    E --> G[编码器 - 解码器注意力]

关键组件说明：

QKV 生成 ：同一线性层产生查询、键、值向量
缩放点积 ：得分除以√d_k 防止梯度消失
多头机制 ：并行多个注意力头提升表征能力

import torch
import torch.nn as nn
import math

class BartAttention(nn.Module):
    def __init__(self, embed_dim, num_heads):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads

        # 合并的 QKV 投影矩阵
        self.qkv_proj = nn.Linear(embed_dim, 3 * embed_dim)
        self.out_proj = nn.Linear(embed_dim, embed_dim)

    def forward(self, x, mask=None):
        batch_size, seq_len, _ = x.shape

        # 生成 QKV [batch, seq_len, 3*embed_dim]
        qkv = self.qkv_proj(x)

        # 拆分为多头 [batch, num_heads, seq_len, head_dim]
        qkv = qkv.reshape(batch_size, seq_len, 3, self.num_heads, self.head_dim)
        q, k, v = qkv.unbind(2)  # 拆解 QKV

        # 缩放点积注意力
        attn_weights = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)

        # 应用注意力掩码（解码器需要）if mask is not None:
            attn_weights = attn_weights.masked_fill(mask == 0, float('-inf'))

        attn_probs = torch.softmax(attn_weights, dim=-1)
        attn_output = torch.matmul(attn_probs, v)

        # 合并多头输出
        attn_output = attn_output.transpose(1, 2).reshape(batch_size, seq_len, self.embed_dim)
        return self.out_proj(attn_output)