注意力机制深度解析：从原理到Transformer实战

1次阅读

没有评论

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

注意力机制的核心思想是通过动态权重分配来捕捉输入序列中不同部分的重要性。其数学表达如下：

给定查询向量 Q、键向量 K 和值向量 V，注意力权重计算过程为：

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

其中：
– Q ∈ ℝ^{n×d_k}：查询矩阵
– K ∈ ℝ^{m×d_k}：键矩阵
– V ∈ ℝ^{m×d_v}：值矩阵
– d_k：键向量的维度

与传统 RNN 相比，注意力机制的优势在于：

解决了 RNN 的梯度消失问题：通过直接连接任意距离的 token
并行计算：所有注意力权重可以同时计算
可解释性：权重矩阵可视化了 token 间的关系

注意力类型	计算复杂度	优点	缺点	适用场景
Scaled Dot-Product	O(n²d)	计算简单，易于实现	对长序列内存消耗大	通用 NLP 任务
Additive	O(n²d²)	可学习非线性变换	计算资源消耗高	小规模数据集
Multi-Head	O(n²dk)	捕捉多维度特征	实现复杂度较高	需要丰富特征的任务

import torch
import torch.nn as nn
from einops import rearrange
from torch.utils.checkpoint import checkpoint

class MultiHeadAttention(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

        self.qkv = nn.Linear(embed_dim, embed_dim * 3)
        self.out = nn.Linear(embed_dim, embed_dim)

    def forward(self, x, mask=None):
        # 使用梯度检查点节省显存
        return checkpoint(self._forward, x, mask)

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

        # 生成 QKV [B, L, 3*D] -> [3, B, H, L, D/H]
        qkv = self.qkv(x)
        q, k, v = rearrange(qkv, 'b l (k h d) -> k b h l d', 
                           k=3, h=self.num_heads)

        # 注意力得分 [B, H, L, L]
        scores = torch.einsum('bhqd,bhkd->bhqk', q, k) / (self.head_dim ** 0.5)

        # 掩码处理
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))

        # 注意力权重
        attn = torch.softmax(scores, dim=-1)

        # 输出 [B, H, L, D/H] -> [B, L, D]
        out = torch.einsum('bhal,bhld->bhad', attn, v)
        out = rearrange(out, 'b h l d -> b l (h d)')

        return self.out(out)

Flash Attention 通过以下技术优化显存使用：

分块计算 ：将大矩阵分解为小块，减少中间结果存储
内存融合 ：将多个操作合并为单个 CUDA 内核
在线 softmax：避免存储完整的注意力矩阵

示例 CUDA 内核融合实现：

__global__ void flash_attention_kernel(
    float* Q, float* K, float* V,
    float* output, int seq_len, int d_model) {extern __shared__ float shared_mem[];

    // 分块加载 Q、K、V 到共享内存
    // ...

    // 在线计算 softmax
    float max_val = -INFINITY;
    float sum_exp = 0.0f;

    for (int i = 0; i < tile_size; ++i) {
        float score = 0.0f;
        for (int j = 0; j < d_model; ++j) {score += Q[tid * d_model + j] * K[i * d_model + j];
        }
        score /= sqrtf(d_model);

        max_val = fmaxf(max_val, score);
        sum_exp += expf(score - max_val);
    }

    // 计算结果并写入全局内存
    // ...
}