GPT 的“思考引擎”：多头因果自注意力与前馈网络的精妙配合

今天，深入到 GPT 的“思考引擎”内部，聚焦最核心的两个部件：多头因果自注意力和前馈神经网络，看看它们是如何精妙配合，赋予 AI 理解上下文、预测未来的强大能力。

依然以 “The cat sat on the” 这个例句为例，模型需要预测下一个词。在将这句话转化为“数字暗号”并穿上“小马甲”（词嵌入和位置嵌入）后，数据就进入了 GPT 的核心处理单元—— Transformer Block。 而 Transformer Block 的魔力，就藏在多头因果自注意力和前馈神经网络这两大“核心引擎”中。

def transformer_block(x, mlp, attn, ln_1, ln_2, n_head):  # [n_seq, n_embd] -> [n_seq, n_embd]
    # multi-head causal self attention
    x = x + mha(layer_norm(x, **ln_1), **attn, n_head=n_head)  # [n_seq, n_embd] -> [n_seq, n_embd]

    # position-wise feed forward network
    x = x + ffn(layer_norm(x, **ln_2), **mlp)  # [n_seq, n_embd] -> [n_seq, n_embd]

    return x

多头因果自注意力：词语间的“单向交流”与“多重视角”

想象一下，句子中的每个词语都是一个信息节点，它们需要相互了解才能理解句子的整体含义。多头因果自注意力机制，正是构建起这张精密的“交流网络”，让每个词语都能“看到”并“关注”其之前的词语，从而捕捉上下文信息。

在 picoGPT 的代码中，mha 函数是这一机制的具体体现。让我们深入理解“因果”与“多头”的含义：

def mha(x, c_attn, c_proj, n_head):  # [n_seq, n_embd] -> [n_seq, n_embd]
    # qkv projection
    x = linear(x, **c_attn)  # [n_seq, n_embd] -> [n_seq, 3*n_embd]

    # split into qkv
    qkv = np.split(x, 3, axis=-1)  # [n_seq, 3*n_embd] -> [3, n_seq, n_embd]

    # split into heads
    qkv_heads = list(map(lambda x: np.split(x, n_head, axis=-1), qkv))  # [3, n_seq, n_embd] -> [3, n_head, n_seq, n_embd/n_head]

    # causal mask to hide future inputs from being attended to
    causal_mask = (1 - np.tri(x.shape[0], dtype=x.dtype)) * -1e10  # [n_seq, n_seq]

    # perform attention over each head
    out_heads = [attention(q, k, v, causal_mask) for q, k, v in zip(*qkv_heads)]  # [3, n_head, n_seq, n_embd/n_head] -> [n_head, n_seq, n_embd/n_head]

    # merge heads
    x = np.hstack(out_heads)  # [n_head, n_seq, n_embd/n_head] -> [n_seq, n_embd]

    # out projection
    x = linear(x, **c_proj)  # [n_seq, n_embd] -> [n_seq, n_embd]

    return x

1、“因果”二字非常重要！

这意味着在预测序列中的某个词时，模型只能依赖于它前面的信息，而不能“偷看”后面的内容。这就像我们阅读理解一样，理解当前内容只能基于已经读过的部分。

这种“因果性”是通过在计算注意力权重时使用 Mask（掩码） 实现的。 对于序列中的每个位置，Mask 会屏蔽掉其之后的位置，确保在计算注意力时，只能考虑到前面的词语。

    # causal mask to hide future inputs from being attended to
    causal_mask = (1 - np.tri(x.shape[0], dtype=x.dtype)) * -1e10  # [n_seq, n_seq]

“因果性”对于文本生成任务至关重要，因为它保证了模型是按照顺序逐步生成文本的，符合人类的写作习惯。

2、“多头”又是什么意思呢？

单单一次注意力计算可能只能捕捉到词语之间的一种关系。为了更全面地理解上下文，模型采用了“多头”策略，也就是同时进行多次独立的注意力计算。

• 独立学习： 每个“头”都拥有自己独立的 Query、Key、Value 变换矩阵，这意味着每个“头”可以学习到不同的注意力模式，关注到词语之间不同类型的关联。

• 并行处理： 这些不同的注意力计算是并行进行的，提高了模型的处理效率。

• 捕捉多样性： 例如，一个“头”可能关注语法关系（主谓宾），另一个“头”可能关注语义关系（实体指代），还有一个“头”可能关注更长距离的依赖关系。 通过组合多个“头”的结果，模型就能更全面、更鲁棒地理解上下文。

3、小结

总结来说，多头因果自注意力机制让模型能够：

• 理解语序： 由于“因果性”，模型能够区分 “猫追老鼠” 和 “老鼠追猫” 的不同含义。

• 捕捉多重关系： 通过“多头”，模型能够从不同角度理解词语间的复杂关联。

• 关注重要信息： 注意力机制使得模型能够自动聚焦于对当前预测最有帮助的词语。

前馈神经网络：每个词的“独立思考与加工”

经过多头因果自注意力机制的“交流”与“审视”后，句子中的每个词语都融入了上下文信息。 接下来，就轮到 前馈神经网络（Position-wise Feed-Forward Network） 发挥作用了。

picoGPT 的代码如下：

def ffn(x, c_fc, c_proj):  # [n_seq, n_embd] -> [n_seq, n_embd]
    # project up
    a = gelu(linear(x, **c_fc))  # [n_seq, n_embd] -> [n_seq, 4*n_embd]

    # project back down
    x = linear(a, **c_proj)  # [n_seq, 4*n_embd] -> [n_seq, n_embd]

    return x

理解前馈神经网络关键在于 “Position-wise”， 这意味着前馈神经网络是独立地、相同地应用于序列中的每个位置的词语表示的。

• 独立性： 与自注意力机制不同，前馈神经网络在处理每个位置的词语时，不会直接与其他位置的词语进行交互。每个位置的词语表示都被单独地输入到同一个前馈网络中进行处理。

• 相同性： 对于序列中的所有位置，使用的是同一个前馈神经网络（共享权重）。这意味着模型在不同的位置应用相同的非线性变换，学习通用的模式。

• 作用： 前馈神经网络的作用是对每个位置的词语表示进行更深层次的非线性变换，提取更高级别的特征，增强模型的表达能力。它可以被看作是对自注意力机制输出的进一步“加工”和“提纯”。

可以这样理解：

• 多头因果自注意力让句子中的词语“互相了解”，建立了联系；
• 而前馈神经网络则像是为每个词语提供了一个独立的“加工车间”，让它们在吸收了上下文信息后，进行“独立思考”，提取更深层次的特征。

总结：GPT 的“思考引擎” —— 分工合作，各司其职

多头因果自注意力和前馈神经网络是 GPT 模型中至关重要的组成部分，它们分工合作，共同塑造了模型的理解和生成能力：

• 多头因果自注意力： 负责捕捉输入序列中不同位置词语之间的依赖关系，理解上下文信息，是模型理解语言的关键。其“因果性”保证了生成的合理顺序，“多头”性提升了理解的全面性。

• 前馈神经网络： 负责对每个位置的词语表示进行独立的非线性变换，提取更高级别的特征，增强模型的表达能力。其 “Position-wise” 的特性使其能够独立地优化每个词语的表示。

理解了这两部分的工作原理，你就对 GPT 模型的“思考方式”有了更深入的认识。后续的层归一化和投影到词汇表等步骤，都是在这些核心机制的基础上进行的，用于最终的概率预测。