📚 课程目录第2章从零实现Transformer解码器

2. 从零实现Transformer解码器

任务说明

本案例要求从零实现Transformer解码器的一个最小完整版本，使其能够接收经过编码的输入序列，通过掩码自注意力与交叉注意力逐步生成目标序列。核心目标是：不依赖任何深度学习框架，仅用Python与NumPy实现解码器的前向传播逻辑，并验证其输出形状与数值可计算性。

解码器结构定义

解码器由三个核心子层堆叠而成：掩码多头自注意力（Masked Multi-Head Self-Attention）、交叉多头注意力（Cross Multi-Head Attention，以编码器输出为Key/Value）、以及逐位置前馈网络（Position-Wise Feed-Forward Network）。每个子层后均包含残差连接与层归一化（LayerNorm）。为保持实现聚焦，本案例实现单层解码器，但代码结构支持直接堆叠多层。

核心实现

以下代码实现了一个完整的单层解码器前向传播。所有矩阵运算均使用NumPy完成，权重采用小随机数初始化。掩码自注意力通过将上三角位置填充为负无穷大来实现因果约束，确保位置i只能关注位置≤i的输入。

import numpy as np

class LayerNorm:
    def __init__(self, d_model, eps=1e-6):
        self.gamma = np.ones(d_model)
        self.beta = np.zeros(d_model)
        self.eps = eps

    def forward(self, x):
        mean = x.mean(axis=-1, keepdims=True)
        var = x.var(axis=-1, keepdims=True)
        return self.gamma * (x - mean) / np.sqrt(var + self.eps) + self.beta

class MultiHeadAttention:
    def __init__(self, d_model, num_heads):
        assert d_model % num_heads == 0
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        # 合并的Q/K/V权重
        self.W_q = np.random.randn(d_model, d_model) * 0.02
        self.W_k = np.random.randn(d_model, d_model) * 0.02
        self.W_v = np.random.randn(d_model, d_model) * 0.02
        self.W_o = np.random.randn(d_model, d_model) * 0.02

    def _split_heads(self, x, batch_size):
        # x: (batch, seq_len, d_model) -> (batch, heads, seq_len, d_k)
        x = x.reshape(batch_size, -1, self.num_heads, self.d_k)
        return x.transpose(0, 2, 1, 3)

    def forward(self, query, key, value, mask=None):
        batch_size = query.shape[0]
        Q = query @ self.W_q
        K = key @ self.W_k
        V = value @ self.W_v

        Q = self._split_heads(Q, batch_size)
        K = self._split_heads(K, batch_size)
        V = self._split_heads(V, batch_size)

        scores = Q @ K.transpose(0, 1, 3, 2) / np.sqrt(self.d_k)
        if mask is not None:
            scores = scores + mask  # mask中0保留，-1e9屏蔽

        attn_weights = np.exp(scores - scores.max(axis=-1, keepdims=True))
        attn_weights /= attn_weights.sum(axis=-1, keepdims=True)

        context = attn_weights @ V  # (batch, heads, seq_len, d_k)
        context = context.transpose(0, 2, 1, 3).reshape(batch_size, -1, self.d_model)
        return context @ self.W_o

class FeedForward:
    def __init__(self, d_model, d_ff=2048):
        self.W1 = np.random.randn(d_model, d_ff) * 0.02
        self.b1 = np.zeros(d_ff)
        self.W2 = np.random.randn(d_ff, d_model) * 0.02
        self.b2 = np.zeros(d_model)

    def forward(self, x):
        return self.W2 @ np.maximum(0, x @ self.W1 + self.b1) + self.b2

class DecoderLayer:
    def __init__(self, d_model, num_heads, d_ff):
        self.self_attn = MultiHeadAttention(d_model, num_heads)
        self.cross_attn = MultiHeadAttention(d_model, num_heads)
        self.ffn = FeedForward(d_model, d_ff)
        self.norm1 = LayerNorm(d_model)
        self.norm2 = LayerNorm(d_model)
        self.norm3 = LayerNorm(d_model)

    def forward(self, tgt, memory, causal_mask):
        # 子层1：掩码自注意力 + 残差 + LayerNorm
        attn1 = self.self_attn.forward(tgt, tgt, tgt, causal_mask)
        tgt = self.norm1.forward(tgt + attn1)

        # 子层2：交叉注意力（Q来自tgt，K/V来自memory）
        attn2 = self.cross_attn.forward(tgt, memory, memory)
        tgt = self.norm2.forward(tgt + attn2)

        # 子层3：前馈网络 + 残差 + LayerNorm
        ffn_out = self.ffn.forward(tgt)
        return self.norm3.forward(tgt + ffn_out)

class TransformerDecoder:
    def __init__(self, d_model=64, num_heads=4, d_ff=128, num_layers=1):
        self.layers = [DecoderLayer(d_model, num_heads, d_ff) for _ in range(num_layers)]

    def forward(self, tgt, memory):
        # 构造因果掩码: (1, 1, tgt_len, tgt_len)
        tgt_len = tgt.shape[1]
        mask = np.triu(np.full((tgt_len, tgt_len), -1e9), k=1)
        mask = mask[np.newaxis, np.newaxis, :, :]

        for layer in self.layers:
            tgt = layer.forward(tgt, memory, mask)
        return tgt

# ========== 使用示例 ==========
if __name__ == "__main__":
    np.random.seed(42)
    batch_size, tgt_len, mem_len, d_model = 2, 5, 7, 64

    # 模拟输入：目标序列嵌入与编码器输出
    tgt_input = np.random.randn(batch_size, tgt_len, d_model).astype(np.float32)
    memory = np.random.randn(batch_size, mem_len, d_model).astype(np.float32)

    decoder = TransformerDecoder(d_model=64, num_heads=4, d_ff=128, num_layers=1)
    output = decoder.forward(tgt_input, memory)

    print("输出形状:", output.shape)  # 预期: (2, 5, 64)
    print("输出数值范围: [{:.4f}, {:.4f}]".format(output.min(), output.max()))

落地与验证

运行方式：将上述代码保存为decoder.py，在Python 3.8+环境中执行python decoder.py（仅需NumPy库）。

验证方法：

形状验证：程序输出的张量形状必须为(2, 5, 64)，对应(batch_size, tgt_len, d_model)。
因果性验证：修改代码，将第一个的查询向量与最后一个的键向量之间的注意力分数打印出来，确认其为-1e9（被屏蔽）。
数值稳定性：连续运行两次（不重置随机种子），输出应不同，证明权重初始化生效；但输出数值应保持在有限范围内（无NaN或无穷大）。

判定标准：若输出形状正确、因果掩码生效（未来位置分数为负无穷）、且无数值溢出，则本案例实现通过验证。