2D多头自注意力计算流程图解：从零实现Transformer核心模块

1次阅读

没有评论

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

在传统的 RNN 和 LSTM 中，序列处理是逐步进行的，这导致两个主要问题：

难以捕捉长距离依赖关系
无法并行计算，训练速度慢

Transformer 通过自注意力机制完美解决了这些问题。它的核心思想是：让序列中的每个元素都能直接 ” 看到 ” 其他所有元素，并通过学习到的权重来决定关注哪些部分。

假设我们有一个序列输入 $X \in \mathbb{R}^{n \times d_{model}}$，其中 n 是序列长度，$d_{model}$ 是嵌入维度。

通过三个不同的线性变换得到查询 (Query)、键(Key) 和值(Value)：

$$
Q = XW^Q, \quad K = XW^K, \quad V = XW^V
$$

其中 $W^Q, W^K, W^V \in \mathbb{R}^{d_{model} \times d_k}$

使用 einops 库将 Q,K,V 拆分为 h 个头：

from einops import rearrange

# 假设 h 是头数，d_k 是每个头的维度
q = rearrange(q, 'b n (h d) -> b h n d', h=h)
k = rearrange(k, 'b n (h d) -> b h n d', h=h)
v = rearrange(v, 'b n (h d) -> b h n d', h=h)

计算注意力权重并应用 softmax：

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

Python 实现：

attn = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
attn = attn.softmax(dim=-1)
out = attn @ v

将多个头的输出拼接起来：

out = rearrange(out, 'b h n d -> b n (h d)')

import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange
import math

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, h, dropout=0.1):
        super().__init__()
        assert d_model % h == 0
        self.d_k = d_model // h
        self.h = h

        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.dropout = nn.Dropout(dropout)
        self.out = nn.Linear(d_model, d_model)

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

        # 线性变换
        q = self.q_linear(q)
        k = self.k_linear(k)
        v = self.v_linear(v)

        # 拆分为多头
        q = rearrange(q, 'b n (h d) -> b h n d', h=self.h)
        k = rearrange(k, 'b n (h d) -> b h n d', h=self.h)
        v = rearrange(v, 'b n (h d) -> b h n d', h=self.h)

        # 计算注意力
        scores = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(self.d_k))

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

        attn = F.softmax(scores, dim=-1)
        attn = self.dropout(attn)

        # 应用注意力权重
        out = attn @ v

        # 拼接多头
        out = rearrange(out, 'b h n d -> b n (h d)')

        # 最终线性变换
        return self.out(out)

自注意力的计算复杂度是 $O(n^2d)$，其中 n 是序列长度，d 是嵌入维度
当序列长度增加时，显存占用会平方级增长

使用梯度检查点(gradient checkpointing)
混合精度训练
适当减少头数(h)，例如从 8 头改为 4 头

梯度爆炸：确保除以 $\sqrt{d_k}$ 进行缩放

变长序列处理：

# 创建 mask 矩阵
mask = torch.ones(bs, seq_len)
for i in range(bs):
    mask[i, lengths[i]:] = 0
mask = mask.unsqueeze(1).unsqueeze(1)  # 形状变为(b,1,1,seq_len)