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

1次阅读

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

Transformer 通过自注意力机制实现了序列数据的全局依赖建模，彻底解决了 RNN 系列模型无法并行计算的瓶颈。其恒定长度的路径距离（O(1)）让长距离依赖捕捉不再受序列长度限制，多头注意力机制则赋予模型同时关注不同位置信息的能力。这些特性使其成为机器翻译、文本生成等任务的绝对主流架构。

在处理长文本时，传统 RNN/LSTM 面临两个致命缺陷：

时间复杂度 ：RNN 需要 O(n) 次顺序计算才能建立 n 个 token 的关联，而 Transformer 通过自注意力机制一次性计算所有 token 关系（理论复杂度 O(n²)但可优化）
信息衰减：LSTM 虽然通过门控机制缓解梯度消失，但超过 100 步的序列仍会出现信息丢失，而 Transformer 的注意力权重直接建模任意两个 token 的关联

以 512 长度的文本为例：
– LSTM 需要串行执行 512 次计算
– Transformer 通过矩阵运算一次完成所有 token 的关联计算

词嵌入层将离散 token 转换为连续向量：

self.token_embed = nn.Embedding(vocab_size, d_model)

位置编码使用正弦函数生成，公式为：
$$PE_{(pos,2i)} = \sin(pos/10000^{2i/d_{\text{model}}})$$
$$PE_{(pos,2i+1)} = \cos(pos/10000^{2i/d_{\text{model}}})$$

实现代码：

def positional_encoding(max_len, d_model):
    position = torch.arange(max_len).unsqueeze(1)
    div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))
    pe = torch.zeros(max_len, d_model)
    pe[:, 0::2] = torch.sin(position * div_term)
    pe[:, 1::2] = torch.cos(position * div_term)
    return pe

注意力计算流程：
1. 通过线性变换生成 Query(Q)、Key(K)、Value(V)矩阵
2. 计算注意力分数：$\text{Attention}(Q,K,V) = \text{softmax}(\frac{QK^T}{\sqrt{d_k}})V$

核心代码：

class SelfAttention(nn.Module):
    def __init__(self, d_model, d_k):
        super().__init__()
        self.W_q = nn.Linear(d_model, d_k)
        self.W_k = nn.Linear(d_model, d_k)
        self.W_v = nn.Linear(d_model, d_k)
        self.scale = 1 / (d_k ** 0.5)

    def forward(self, x, mask=None):
        Q = self.W_q(x)  # [batch, seq_len, d_k]
        K = self.W_k(x)
        V = self.W_v(x)

        scores = torch.matmul(Q, K.transpose(-2, -1)) * self.scale
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)
        attn = torch.softmax(scores, dim=-1)
        return torch.matmul(attn, V)

将注意力机制扩展到多个子空间：

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.d_head = d_model // n_heads
        self.n_heads = n_heads
        self.attention = SelfAttention(d_model, self.d_head)
        self.fc_out = nn.Linear(d_model, d_model)

    def forward(self, x, mask=None):
        batch_size = x.size(0)
        # 分头处理 [batch, seq_len, n_heads, d_head]
        x = x.view(batch_size, -1, self.n_heads, self.d_head)
        x = x.transpose(1, 2)  # [batch, n_heads, seq_len, d_head]

        attended = self.attention(x, mask)
        attended = attended.transpose(1, 2).contiguous()
        attended = attended.view(batch_size, -1, self.n_heads * self.d_head)
        return self.fc_out(attended)

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

    def forward(self, x, mask=None):
        attn_out = self.attention(x, mask)
        x = self.norm1(x + self.dropout(attn_out))
        ffn_out = self.ffn(x)
        return self.norm2(x + self.dropout(ffn_out))