Masked Self-Attention保姆级教程：手把手实现Transformer文本生成

从零构建文本生成引擎：深入Masked Self-Attention的工程实践

如果你正在尝试构建自己的GPT风格模型，或者对Transformer解码器的内部运作机制感到好奇，那么你很可能已经听说过Masked Self-Attention。这个名字听起来有些神秘，但它实际上是现代文本生成模型（如GPT系列）能够“思考”并创造出连贯文字的核心秘密。与编码器中“全知全能”的注意力机制不同，解码器的注意力是“有纪律”的——它只能回顾过去，不能窥探未来。这种看似简单的限制，恰恰是生成式人工智能能够有序工作的基石。本文将抛开复杂的理论堆砌，直接进入代码层面，手把手带你用PyTorch实现一个完整的、可用于文本生成的Masked Self-Attention模块，并深入探讨其在训练和推理中的微妙差异。无论你是希望为自己的项目添加文本生成能力，还是想彻底理解Transformer解码器的工作原理，这里都将提供一条清晰的实践路径。

1. 理解Masked Self-Attention：为何“看不见未来”如此重要

在开始写代码之前，我们必须先弄清楚一个根本问题：为什么解码器需要“掩码”？想象一下，你正在写一篇小说。当你写下第一个句子时，你并不知道整篇故事的结局。你只能基于已经写下的文字，去构思下一句。这就是自回归生成的本质——每一步的输出都依赖于之前所有步骤的输出。如果解码器在生成第一个词时就能“看到”最后一个词，那无异于考试时提前知道了答案，模型将无法学会如何根据上下文进行合理的预测。因此，Masked Self-Attention的核心任务就是施加这种“时间因果性”约束。

从数学和计算图的角度看，普通的Self-Attention会计算序列中每个位置与其他所有位置（包括未来位置）的关联度。而在解码器中，我们需要一个下三角掩码矩阵，将未来位置的注意力权重设置为一个极大的负值（如 -1e9），这样在经过Softmax函数后，这些位置的权重就会趋近于零。这就确保了在计算位置 i 的表示时，模型只能聚合位置 0 到 i 的信息。

注意：这种掩码机制不仅适用于文本生成，也广泛应用于任何需要保持时间或顺序因果关系的序列生成任务中，如代码补全、音乐生成和语音合成。

理解这一点后，我们可以将其与编码器的注意力做一个快速对比：

这种设计上的差异，直接导致了训练和推理流程的根本不同，这也是我们后续需要仔细处理的部分。

2. 搭建基础：实现Scaled Dot-Product Attention与掩码

让我们从最基础的注意力函数开始实现。我们将遵循原始Transformer论文的设计，但会特别关注掩码的集成。

首先，我们实现核心的缩放点积注意力函数。这个函数接收查询（Q）、键（K）、值（V）三个张量，以及可选的掩码。

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

def scaled_dot_product_attention(q, k, v, mask=None):
    """
    计算缩放点积注意力。
    参数:
        q: 查询张量，形状为 (batch_size, num_heads, seq_len, d_k)
        k: 键张量，形状为 (batch_size, num_heads, seq_len, d_k)
        v: 值张量，形状为 (batch_size, num_heads, seq_len, d_v)
        mask: 掩码张量，形状为 (batch_size, 1, seq_len, seq_len) 或 (batch_size, seq_len, seq_len)
    返回:
        注意力加权的值张量，形状为 (batch_size, num_heads, seq_len, d_v)
        注意力权重张量，形状为 (batch_size, num_heads, seq_len, seq_len)
    """
    d_k = q