掌握Transformer张量维度：从嵌入到语言模型头的完整指南
引言
在构建生成式AI模型时，理解张量维度的变化是掌握模型内部机制的关键。许多初学者在阅读Transformer相关代码时，常常对张量形状的变换感到困惑。本文将以一个简单的文本生成模型为例，逐步剖析从输入到输出的整个过程中，张量维度是如何演变的。通过具体的数值示例和矩阵乘法计算，你将清晰地看到每个步骤的维度变化，从而更深入地理解注意力机制和前馈网络的工作原理。
前置知识
在深入之前，你需要对矩阵乘法的维度规则有扎实的理解。例如，两个矩阵相乘时，第一个矩阵的列数必须等于第二个矩阵的行数，结果矩阵的行数等于第一个矩阵的行数，列数等于第二个矩阵的列数。如果你对此还不熟悉，建议先复习相关概念，因为本文中的许多计算都依赖于这一基本规则。
模型设置
大多数生成式AI模型采用解码器-only架构。我们以一个简单的文本生成模型为例，输入句子“Hello world !”经过分词后，得到三个词元：Hello、world和!。此外，我们添加两个特殊词元：<bos>（句子开始）和<eos>（句子结束），以确保输入的正确移位。分词后的输入张量表示为[12, 15496, 2159, 5145]，其中每个数字代表一个词元的ID。当以批次输入时，会添加一个批次维度，形状变为[[12, 15496, 2159, 5145]]。为了简化，我们关注张量的维度，将输入表示为[1, 4]，其中1是批次大小，4是序列长度。
嵌入层
输入张量首先经过嵌入层，其形状变为[1, 4, 768]，其中768是嵌入维度。嵌入层有两个重要作用：
- 维度传播：嵌入维度贯穿整个网络，并在注意力层中被大量使用。
- 语义表示：它将词元转换为高维向量，捕捉词之间的语义关系。例如，虽然“king”和“man”在数值上可能不同（如8848和9584），但它们在向量空间中的表示会显示出有意义的相似性。
位置编码
位置编码层不改变张量的维度，但会向输入中注入位置信息。由于后续的并行计算，模型需要知道词元的顺序，位置编码确保了这一点。
解码器层
生成模型通常包含多个连续的解码器层，每一层包含：
- 掩码多头注意力层
- 加和归一化操作
- 前馈网络
掩码多头注意力
多头注意力层使模型能够关注输入的不同部分，为每个词元在序列上下文中加权表示。掩码多头注意力则限制每个词元只能关注自身及之前的词元，通过掩码未来词元来实现。
数据首先通过三个并行的线性层，每个线性层为nn.Linear(768, 768)，保持输入形状不变。输出分别称为查询（Q）、键（K）和值（V），每个张量形状为[1, 4, 768]。
接下来，嵌入维度被分割为8个头，每个头的大小为96（因为768 = 8 * 96）。这会将张量重塑为[1, 4, 8, 96]。为了进行矩阵乘法，我们需要将序列长度和头大小维度对齐，因此转置后得到形状[1, 8, 4, 96]，分别对应Q、K、V。
注意力机制的计算公式为：
Attention(Q, K, V) = softmax(QK^T / sqrt(d_k)) V
首先计算QK^T：
- K的转置形状为
[1, 8, 96, 4]。 - QK^T的结果形状为
[1, 8, 4, 4]，计算过程为[1, 8, 4, 96] × [1, 8, 96, 4]。
掩码：掩码确保每个词元只关注自身及之前的词元，防止模型在生成时访问未来词元。
注意力权重：使用softmax(QK^T / sqrt(d_k))计算，其中d_k = 96，sqrt(d_k) ≈ 9.8。缩放防止值差异过大，softmax确保每个词元的权重和为1，并将-inf值置零，但只改变值，不改变形状。
计算注意力输出：将注意力权重与V相乘：
[1, 8, 4, 4] × [1, 8, 4, 96] = [1, 8, 4, 96]
拼接：将头数恢复，通过转置：
[1, 8, 4, 96] → [1, 4, 8, 96] → [1, 4, 768]
投影层：使用nn.Linear(768, 768)保持形状为[1, 4, 768]。
观察：张量形状在注意力机制后恢复为
[1, 4, 768]，与输入形状一致，这确保了与后续层的兼容性。
加和归一化
这一步包括跳跃连接，将注意力层前后的张量相加并归一化。加法确保张量值被更新而非替换，归一化防止值指数增长。这些操作在每一层后应用，以保持张量的原始特性。
前馈网络
前馈网络通常由两个连续的线性层组成：一个扩展张量，另一个收缩张量，并常伴有dropout层进行正则化。这些层引入非线性变换，使模型能够捕捉嵌入维度内的更复杂模式。
扩展因子通常为1→3，收缩因子为3→1，因此线性层结构为：
nn.Linear(768, 3 * 768)
nn.Linear(3 * 768, 768)最终输出形状回到输入形状[1, 4, 768]。保持形状允许在前馈网络后应用加和归一化层。此外，由于解码器的最终形状与输入形状相同，可以无缝堆叠多个解码器层。
语言模型头
经过一系列解码器层后，张量到达最后的线性层，将嵌入维度转换为词汇表大小。
假设词汇表大小为9735，则输出张量形状为[1, 4, 9735]，其中：
- 1：批次大小
- 4：序列长度
- 9735：词汇表大小
应用softmax函数并计算模型输出与真实标签之间的损失，优化器使用该损失更新模型权重。
在掩码多头注意力层中，每个词元的注意力仅使用当前输入词元及其前驱计算。由于输入是右移的，生成新词元时，模型需要考虑所有先前词元的向量表示及其与前驱的关系。这一机制是生成式AI模型工作的基础。
Transformer与交叉注意力
Transformer架构通常包含编码器-解码器结构，常用于上下文和输出不直接相关的任务，如翻译。然而，如今解码器-only架构更受青睐。
在编码器层中，张量形状的传播与解码器类似，但关键区别在于注意力层不应用掩码，允许每个词元关注序列中的所有其他词元，包括前后。
例如，考虑上下文“I am at home”和目标“
- 上下文：
[1, 4] - 目标：
[1, 6]
输入和目标序列长度不同，这由交叉注意力层处理。编码器输出形状为[1, 4, 768]，而解码器的掩码多头注意力层输出为[1, 6, 768]。在交叉注意力中，键（K）和值（V）来自编码器，查询（Q）来自解码器。
建议手动计算张量维度以加深理解，但下面提供了解答。
注意力公式：Attention(Q, K, V) = softmax(QK^T / sqrt(d_k)) V
- 查询：
[1, 6, 768] - 键和值：
[1, 4, 768]
分割并转置维度：
- 查询：
[1, 8, 6, 96] - 键和值：
[1, 8, 4, 96]
计算注意力：
- K^T：
[1, 8, 96, 4] - QK^T：
[1, 8, 6, 96] * [1, 8, 96, 4] = [1, 8, 6, 4] - softmax(QK^T / sqrt(d_k))：
[1, 8, 6, 4] - softmax(QK^T / sqrt(d_k)) V：
[1, 8, 6, 4] * [1, 8, 4, 96] = [1, 8, 6, 96]
拼接：
[1, 8, 6, 96] → [1, 6, 8, 96] → [1, 6, 768]
观察：张量形状恢复为
[1, 6, 768]，与掩码多头注意力层的输出形状一致，确保与后续层的兼容性。
总结
本文旨在帮助读者更清晰地理解注意力机制的工作原理以及张量形状在Transformer架构中的传播过程。通过逐步分析每个组件的维度变化，我们看到了模型如何保持形状一致性，从而支持深层堆叠和高效计算。掌握这些概念对于深入研究和开发生成式AI模型至关重要。
如果你觉得这篇文章有帮助，欢迎点赞和分享。对于任何反馈或问题，请随时通过我的个人网站联系我：https://not-lain.github.io/。