Transformer基础入门：从词嵌入到注意力机制的完整实现指南

1次阅读

没有评论

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

在自然语言处理（NLP）领域，传统的 RNN 和 LSTM 模型在处理长序列时存在两个主要问题：

梯度消失 / 爆炸：随着序列长度增加，反向传播时梯度会指数级衰减或增长，导致模型难以学习长期依赖关系。
顺序计算限制：必须按时间步逐个处理序列，无法充分利用 GPU 的并行计算能力。

而 Transformer 通过完全基于注意力机制的架构解决了这些问题。下面我们拆解它的核心组件。

词嵌入将离散的单词映射到连续的向量空间。假设词汇表大小为 V，嵌入维度为 d，则嵌入矩阵为：

$$ W_e \in \mathbb{R}^{V \times d} $$

PyTorch 实现示例：

import torch.nn as nn

embedding = nn.Embedding(num_embeddings=10000, embedding_dim=512)
# 输入形状: (batch_size, seq_len)
input_ids = torch.LongTensor([[1, 23, 456], [78, 901, 2]]) 
# 输出形状: (batch_size, seq_len, d)
embedded = embedding(input_ids)

由于 Transformer 没有循环结构，需要通过位置编码注入序列顺序信息。计算公式为：

$$
PE_{(pos,2i)} = \sin(pos/10000^{2i/d}) \
PE_{(pos,2i+1)} = \cos(pos/10000^{2i/d})
$$

实现时要注意避免数值溢出（常见错误见后文避坑指南）：

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000):
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        self.register_buffer('pe', pe)

    def forward(self, x):
        return x + self.pe[:x.size(1)]  # 形状匹配: (seq_len, d)

计算公式：

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

关键实现细节：

sqrt(d_k)缩放避免 softmax 梯度消失
注意 mask 处理（用于 decoder 或 padding）

def scaled_dot_product_attention(q, k, v, mask=None):
    # 输入形状: q (..., seq_len_q, d_k)
    #          k,v (..., seq_len_k, d_k)
    matmul_qk = torch.matmul(q, k.transpose(-2, -1))
    dk = k.size(-1)
    scaled_attention = matmul_qk / math.sqrt(dk)

    if mask is not None:
        scaled_attention += (mask * -1e9)  # 负无穷掩码

    attention_weights = F.softmax(scaled_attention, dim=-1)
    return torch.matmul(attention_weights, v)

将 Q、K、V 投影到 h 个子空间并行计算注意力，最后拼接结果：

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        assert d_model % num_heads == 0
        self.d_k = d_model // num_heads
        self.proj = nn.Linear(d_model, d_model)
        self.linear_layers = nn.ModuleList([nn.Linear(d_model, d_model) for _ in range(3)
        ])

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

        # 线性投影 -> (batch, seq_len, d_model)
        q,k,v = [l(x) for l,x in zip(self.linear_layers, (q,k,v))]

        # 分头 -> (batch, num_heads, seq_len, d_k)
        q = q.view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        k = k.view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        v = v.view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)

        # 计算注意力并拼接
        scaled_attention = scaled_dot_product_attention(q, k, v, mask)
        concat = scaled_attention.transpose(1, 2).contiguous() \
                 .view(batch_size, -1, self.d_model)

        return self.proj(concat)

单头注意力：$O(n^2 \cdot d)$
h 头注意力：$O(n^2 \cdot (d/h) \cdot h) = O(n^2 \cdot d)$（理论复杂度相同）
实际运行时间：多头会增加并行开销，但能更好利用 GPU

$$
显存(Bytes) \approx 4 \times (3 \times n \times d + n^2 \times h)
$$

其中 4 表示 float32 占用的字节数。

位置编码溢出 ：当pos/10000^(2i/d) 太小时，sin/cos计算可能丢失精度。解决方案是像前文代码那样用 exp(log) 形式计算。
注意力掩码错误：decoder 的因果掩码应使用上三角矩阵，padding 掩码需要广播到正确形状。
维度不匹配 ：分头操作后务必检查张量形状，常见错误是view 和transpose顺序不对。