注意力机制全解析：从基础原理到Transformer实战应用

1次阅读

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

在自然语言处理（NLP）任务中，传统的循环神经网络（RNN）及其变种 LSTM、GRU 在处理序列数据时面临一个主要问题：长程依赖。随着序列长度的增加，RNN 模型难以有效地捕捉远处单词之间的关系。例如，在句子 ”The animal didn’t cross the street because it was too tired” 中，”it” 指代的是 ”animal”，但两者相隔较远，传统 RNN 可能无法准确建立这种联系。

这种局限性促使研究者寻找更有效的建模方式，注意力机制应运而生。注意力机制的核心思想是：在处理每个位置的信息时，动态地关注与当前任务最相关的其他位置的信息，而不是像 RNN 那样依赖固定的时间步传递。

注意力机制的计算过程可以分解为三个核心步骤：

投影阶段 ：将输入序列转换为 Query(Q)、Key(K) 和 Value(V)三个矩阵。这三个矩阵通常是通过不同的线性变换从输入序列得到的。
注意力分数计算：计算 Query 与每个 Key 的相似度得分，常用的相似度函数有点积、加性注意力等。
加权求和：使用 softmax 函数将注意力分数归一化为权重，然后用这些权重对 Value 进行加权求和。

数学表达式为：

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

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

并行计算：不像 RNN 需要顺序处理，注意力可以并行计算所有位置的注意力分数
长程依赖：直接建模任意两个位置的关系，不受距离限制
可解释性：注意力权重可以直观展示模型关注的重点

自注意力是一种特殊的注意力机制，其 Query、Key 和 Value 都来自同一个输入序列。它使模型能够同时关注输入序列的不同位置，从而学习丰富的内部表示。

在 Transformer 中，自注意力机制允许每个单词直接与句子中的所有其他单词建立联系，无论它们在序列中的位置如何。这种特性使得 Transformer 特别适合捕捉长距离依赖关系。

交叉注意力则涉及两个不同的序列：Query 来自一个序列，而 Key 和 Value 来自另一个序列。这种机制在编码器 - 解码器架构中特别有用，例如在机器翻译任务中，解码器可以关注编码器的输出。

多头注意力将 Query、Key 和 Value 分别投影到多个子空间（称为 ” 头 ”），在每个子空间中独立计算注意力，最后将结果拼接起来。这种方法允许模型在不同的表示子空间中关注不同的信息。

数学表达式为：

$$
\text{MultiHead}(Q,K,V) = \text{Concat}(head_1,…,head_h)W^O
$$

其中每个头的计算为：

$$
head_i = \text{Attention}(QW_i^Q,KW_i^K,VW_i^V)
$$

为了降低长序列上的计算开销，研究者提出了各种稀疏注意力变体，如：

局部注意力：只关注固定窗口内的邻近位置
步长注意力：以固定间隔采样关键点
轴向注意力：沿着不同的轴分别计算注意力

这些方法可以显著减少计算量，但可能会牺牲一些模型性能。

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

class AttentionLayer(nn.Module):
    def __init__(self, embed_dim, num_heads=8, dropout=0.1):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads

        assert self.head_dim * num_heads == embed_dim, "embed_dim must be divisible by num_heads"

        self.q_proj = nn.Linear(embed_dim, embed_dim)
        self.k_proj = nn.Linear(embed_dim, embed_dim)
        self.v_proj = nn.Linear(embed_dim, embed_dim)

        self.out_proj = nn.Linear(embed_dim, embed_dim)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        batch_size, seq_len, embed_dim = x.size()

        # 投影到 Q,K,V
        q = self.q_proj(x)  # [batch_size, seq_len, embed_dim]
        k = self.k_proj(x)  
        v = self.v_proj(x)

        # 分割多头
        q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        k = k.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        v = v.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)

        # 计算注意力分数
        attn_scores = torch.matmul(q, k.transpose(-2, -1)) / (self.head_dim ** 0.5)

        # 应用 mask（如果有）if mask is not None:
            attn_scores = attn_scores.masked_fill(mask == 0, float('-inf'))

        # 计算注意力权重
        attn_weights = F.softmax(attn_scores, dim=-1)
        attn_weights = self.dropout(attn_weights)

        # 加权求和
        output = torch.matmul(attn_weights, v)

        # 合并多头
        output = output.transpose(1, 2).contiguous().view(batch_size, seq_len, embed_dim)

        # 最终投影
        output = self.out_proj(output)

        return output, attn_weights

class TransformerEncoderLayer(nn.Module):
    def __init__(self, embed_dim, num_heads, ff_dim=2048, dropout=0.1):
        super().__init__()
        self.self_attn = AttentionLayer(embed_dim, num_heads, dropout)
        self.ffn = nn.Sequential(nn.Linear(embed_dim, ff_dim),
            nn.ReLU(),
            nn.Linear(ff_dim, embed_dim)
        )
        self.norm1 = nn.LayerNorm(embed_dim)
        self.norm2 = nn.LayerNorm(embed_dim)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)

    def forward(self, src, src_mask=None):
        # 自注意力
        src2, attn_weights = self.self_attn(src, src_mask)
        src = src + self.dropout1(src2)
        src = self.norm1(src)

        # 前馈网络
        src2 = self.ffn(src)
        src = src + self.dropout2(src2)
        src = self.norm2(src)

        return src, attn_weights