BART自注意力机制架构图解析：从原理到新手实践指南

1次阅读

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

在自然语言处理领域，传统的 RNN 和 LSTM 模型在处理长序列时存在梯度消失和并行计算困难的问题。而 Transformer 架构通过自注意力机制彻底改变了这一局面，成为现代 NLP 模型的基石。BART 作为一种基于 Transformer 的序列到序列模型，其核心正是自注意力机制。

自注意力机制允许模型在处理每个词时，直接关注到输入序列中的所有其他词，动态计算它们之间的相关性权重。这种机制相比 RNN/LSTM 具有三大优势：

更强的长距离依赖建模能力
更高的并行计算效率
更灵活的特征提取方式

BART 的自注意力层遵循标准的 Transformer 架构，但作为编解码模型，它包含编码器和解码器两套注意力机制。让我们一步步拆解其数据流：

输入表示：每个词经过嵌入层得到维度为 d_model 的向量，加上位置编码后作为自注意力层的输入
Q/K/ V 矩阵计算 ：通过三个不同的线性变换将输入转换为查询(Query)、键(Key) 和值 (Value) 矩阵
缩放点积注意力 ：计算 Q 和 K 的点积，除以 sqrt(d_k) 进行缩放，然后应用 softmax 得到注意力权重
多头注意力：将上述过程在多个子空间并行执行，最后拼接各头的输出并通过线性变换
残差连接和层归一化：注意力输出与原始输入相加后进行层归一化

解码器还额外包含编码器 - 解码器注意力层，其 K / V 来自编码器输出，Q 来自解码器上一层的输出。

以下是 PyTorch 实现 BART 自注意力层的核心代码片段：

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

class BartSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.num_heads = config.num_attention_heads
        self.head_dim = config.hidden_size // config.num_attention_heads

        # Q/K/ V 投影矩阵
        self.query = nn.Linear(config.hidden_size, config.hidden_size)
        self.key = nn.Linear(config.hidden_size, config.hidden_size)
        self.value = nn.Linear(config.hidden_size, config.hidden_size)

        # 输出投影
        self.dense = nn.Linear(config.hidden_size, config.hidden_size)

    def forward(self, hidden_states, attention_mask=None):
        batch_size, seq_length, hidden_size = hidden_states.size()

        # 计算 Q /K/V [batch_size, seq_len, hidden_size]
        q = self.query(hidden_states)
        k = self.key(hidden_states)
        v = self.value(hidden_states)

        # 重塑为多头 [batch_size, num_heads, seq_len, head_dim]
        q = q.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2)
        k = k.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2)
        v = v.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2)

        # 计算注意力分数 [batch_size, num_heads, seq_len, seq_len]
        scores = torch.matmul(q, k.transpose(-1, -2)) / math.sqrt(self.head_dim)

        # 应用 attention mask(解码器使用)
        if attention_mask is not None:
            scores = scores + attention_mask

        # softmax 归一化
        attn_weights = nn.functional.softmax(scores, dim=-1)

        # 注意力加权求和
        attn_output = torch.matmul(attn_weights, v)

        # 重塑回原始形状 [batch_size, seq_len, hidden_size]
        attn_output = attn_output.transpose(1, 2).contiguous()
        attn_output = attn_output.view(batch_size, seq_length, hidden_size)

        # 最终投影
        attn_output = self.dense(attn_output)

        return attn_output, attn_weights

使用 BertViz 工具可以直观地观察 BART 的注意力模式：

安装 BertViz：pip install bertviz
加载预训练 BART 模型
准备输入文本
调用 bertviz 的注意力可视化函数

示例代码：

from transformers import BartModel, BartTokenizer
from bertviz import head_view

model = BartModel.from_pretrained('facebook/bart-base')
tokenizer = BartTokenizer.from_pretrained('facebook/bart-base')

inputs = tokenizer("The cat sat on the mat", return_tensors="pt")
outputs = model(**inputs, output_attentions=True)

# 可视化最后一层的注意力
attention = outputs.attentions[-1]  
head_view(attention, tokenizer.convert_ids_to_tokens(inputs['input_ids'][0]))

可视化结果将显示不同注意力头在各个位置的关注强度，有助于理解模型如何构建词与词之间的关系。

影响 BART 自注意力机制效果的三个关键超参数：