注意力机制深度解析：从原理到常见实现及自注意力优势

1次阅读

没有评论

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

在深度学习领域，如何让模型有效地关注输入数据中的关键信息一直是个核心挑战。传统的序列模型（如 RNN、LSTM）虽然能够处理序列数据，但存在两个主要问题：

长距离依赖问题 ：随着序列长度的增加，早期信息在传递过程中会逐渐衰减
固定编码问题 ：无论输入如何变化，模型都以相同的方式处理所有位置的输入

注意力机制的提出完美解决了这些问题，它允许模型动态地关注输入的不同部分，根据当前任务需要灵活调整关注重点。

原理：通过一个单层前馈网络计算注意力分数
公式：e_ij = v^T tanh(W_1h_i + W_2h_j)
优点：适用于查询和键维度不同的情况
缺点：计算复杂度较高

原理：直接计算查询和键的点积作为注意力分数
公式：e_ij = q_i^T k_j
优点：计算效率高，实现简单
缺点：当维度较高时，点积结果可能过大，导致 softmax 梯度很小

改进：对点积结果进行缩放，解决梯度消失问题
公式：e_ij = q_i^T k_j / √d_k
优势：Transformer 采用的基础注意力机制

自注意力是注意力机制的特例，其中查询、键和值都来自同一输入序列。它的核心优势在于：

全局依赖建模 ：直接建立序列中任意两个位置的关联
并行计算 ：所有位置的注意力可以同时计算
灵活性强 ：不依赖序列顺序，适合处理各种结构化数据

线性变换 ：将输入 X 分别通过三个权重矩阵投影得到 Q、K、V
注意力分数 ：计算 Q 和 K 的点积并缩放
softmax 归一化 ：得到注意力权重
加权求和 ：用注意力权重对 V 进行加权

为了捕获不同子空间的特征，通常会并行多个自注意力层：

将 Q、K、V 分割到 h 个头
在每个头上分别计算注意力
拼接所有头的输出
通过线性层融合结果

import torch
import torch.nn as nn
import math

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

        self.q_linear = nn.Linear(d_model, d_model)
        self.k_linear = nn.Linear(d_model, d_model)
        self.v_linear = nn.Linear(d_model, d_model)
        self.out_linear = nn.Linear(d_model, d_model)

    def forward(self, q, k, v, mask=None):
        batch_size = q.size(0)

        # 线性投影
        q = self.q_linear(q).view(batch_size, -1, self.num_heads, self.head_dim)
        k = self.k_linear(k).view(batch_size, -1, self.num_heads, self.head_dim)
        v = self.v_linear(v).view(batch_size, -1, self.num_heads, self.head_dim)

        # 转置以准备矩阵乘法
        q = q.transpose(1, 2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)

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

        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)

        attn_weights = torch.softmax(scores, dim=-1)
        output = torch.matmul(attn_weights, v)

        # 拼接多头输出
        output = output.transpose(1, 2).contiguous()
        output = output.view(batch_size, -1, self.d_model)

        return self.out_linear(output)