深入解析AI Transformer原理：从数学基础到高效实现

1次阅读

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

Transformer 架构自从 2017 年由 Vaswani 等人在论文《Attention Is All You Need》中提出后，已经彻底改变了自然语言处理（NLP）和计算机视觉（CV）领域。它摒弃了传统的循环神经网络（RNN）和卷积神经网络（CNN），完全依赖自注意力机制（Self-Attention）来捕捉序列中的长距离依赖关系。这种架构在机器翻译、文本生成、图像识别等任务中表现出色，成为现代 AI 模型（如 BERT、GPT）的核心组件。

然而，许多开发者在实现 Transformer 时常常遇到以下难点：

长序列处理 ：随着序列长度的增加，自注意力的计算复杂度呈平方级增长，导致内存和计算资源消耗剧增。
内存消耗 ：训练大型 Transformer 模型需要存储大量中间结果，显存不足成为常见瓶颈。
收敛困难 ：模型训练过程中可能出现梯度消失或爆炸，需要精细调参。

自注意力机制的核心是通过计算查询（Query）、键（Key）和值（Value）向量之间的相关性，来决定每个位置应该关注序列中的哪些部分。其数学表达如下：

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

其中，$Q$, $K$, $V$ 分别是通过线性变换从输入序列得到的矩阵，$d_k$ 是键向量的维度。分母 $\sqrt{d_k}$ 的作用是缩放点积，防止梯度消失。

Query（查询）：当前正在处理的位置，用于“询问”其他位置的相关性。
Key（键）：序列中其他位置的表示，用于与 Query 匹配。
Value（值）：实际用于加权求和的信息。

由于 Transformer 没有内置的顺序信息，需要通过位置编码为输入序列注入位置信息。常用的三角函数位置编码公式如下：

$$
PE_{(pos, 2i)} = \sin\left(\frac{pos}{10000^{2i/d_{\text{model}}}}\right) \
PE_{(pos, 2i+1)} = \cos\left(\frac{pos}{10000^{2i/d_{\text{model}}}}\right)
$$

其中，$pos$ 是位置，$i$ 是维度索引，$d_{\text{model}}$ 是模型的隐藏层维度。

以下是 Transformer 编码器层（Encoder Layer）的简化实现，包含自注意力和前馈网络（Feed-Forward Network）：

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().__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, x, mask=None):
        batch_size = x.size(0)

        # 线性变换得到 Q、K、V
        q = self.q_linear(x)
        k = self.k_linear(x)
        v = self.v_linear(x)

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

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

        # 应用 mask（如填充 mask 或未来 mask）if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)

        # softmax 归一化
        attention = F.softmax(scores, dim=-1)

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

        # 合并多头
        output = output.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)

        # 输出线性变换
        output = self.out_linear(output)
        return output

class TransformerEncoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, ff_dim, dropout=0.1):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, num_heads)
        self.ffn = nn.Sequential(nn.Linear(d_model, ff_dim),
            nn.ReLU(),
            nn.Linear(ff_dim, d_model)
        )
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        # 自注意力子层
        attn_output = self.self_attn(x, mask)
        x = x + self.dropout(attn_output)
        x = self.norm1(x)

        # 前馈网络子层
        ffn_output = self.ffn(x)
        x = x + self.dropout(ffn_output)
        x = self.norm2(x)
        return x