解密ChatGPT的Transformer自注意力机制：从数学原理到工程实现

1次阅读

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

自注意力机制是 Transformer 架构的核心组件，也是 ChatGPT 等大型语言模型能够高效处理序列数据的关键。传统的 RNN 和 LSTM 模型在处理长序列时存在梯度消失和并行计算困难的问题，而自注意力机制通过计算序列中每个元素与其他元素的关系权重，实现了高效的并行计算和长距离依赖建模。

自注意力机制的核心是 QKV（Query-Key-Value）矩阵计算和缩放点积注意力。下面是关键公式：

QKV 矩阵计算 :
$$
Q = XW_Q, \quad K = XW_K, \quad V = XW_V
$$
其中，$X$ 是输入序列，$W_Q$, $W_K$, $W_V$ 是可学习的权重矩阵。
缩放点积注意力 :
$$
\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V
$$
其中，$d_k$ 是 Key 向量的维度，缩放因子 $\sqrt{d_k}$ 用于防止点积结果过大导致梯度消失。

以下是一个基于 PyTorch 的自注意力层实现，包含多头注意力并行计算优化：

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

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

        self.W_Q = nn.Linear(d_model, d_model)
        self.W_K = nn.Linear(d_model, d_model)
        self.W_V = nn.Linear(d_model, d_model)
        self.W_O = nn.Linear(d_model, d_model)

    def split_heads(self, x):
        batch_size, seq_len, _ = x.size()
        return x.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)

    def forward(self, Q, K, V, mask=None):
        # 线性变换并分头
        Q = self.split_heads(self.W_Q(Q))
        K = self.split_heads(self.W_K(K))
        V = self.split_heads(self.W_V(V))

        # 计算缩放点积注意力
        scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.d_k, dtype=torch.float32))
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)
        attn_weights = F.softmax(scores, dim=-1)
        output = torch.matmul(attn_weights, V)

        # 合并多头输出
        output = output.transpose(1, 2).contiguous().view(output.size(0), -1, self.d_model)
        return self.W_O(output)

自注意力机制的计算复杂度为 $O(n^2)$，其中 $n$ 是序列长度。对于长序列，可以采用以下优化策略：