掌握 Transformer 中的张量维度

先决条件

对矩阵乘法的有效形状和维度有扎实的理解至关重要。强烈建议在继续之前熟悉此主题。

设置

大多数生成式 AI 模型都采用仅解码器架构。在这篇博文中，我们将探讨一个简单的文本生成模型，如下图所示。

我们从一个输入示例开始，以供参考。
句子 Hello world ! 可以被分词为三个部分：Hello、world 和 !。此外，还附加了两个辅助标记，<bos> 和 <eos>，分别表示**句子的开头**和**句子的结尾**。这确保了输入正确地进行了位移。

分词后，输入成为一个张量，例如 [12, 15496, 2159, 5145]。当以批次形式传递给模型时，会添加一个额外的维度，变为 [[12, 15496, 2159, 5145]]。为简化起见，我们将关注张量维度，将输入表示为 $[1,4][1,\; 4]$，其中 1 是批次大小，4 是序列长度。

嵌入层

输入张量在每个层中都会进行转换，这会影响其值，在某些情况下还会影响其形状。当张量到达嵌入层时，其形状变为 $[1,4,768][1,\; 4,\; 768]$，其中 768 表示嵌入维度。

嵌入层对架构至关重要，原因有二：

• 嵌入维度在神经网络中传播，并在注意力层中被大量使用。
• 它将标记转换为高维向量，捕捉词语之间的语义关系。例如，虽然“国王”和“男人”等标记在数值上可能看似无关（例如，8848 和 9584），但它们在高维空间中的向量表示揭示了有意义的相似性。

位置编码

此层不改变张量维度，但会将位置信息**注入**输入。这至关重要，因为输入将在架构的后续部分进行并行计算，而位置编码确保模型保留了序列中标记顺序的信息。

解码器层

生成模型通常由多个连续的解码器层组成，每个层包含：

• 一个遮蔽多头注意力层
• 一个添加和归一化操作
• 一个前馈网络

遮蔽多头注意力

“多头注意力”层使模型能够关注输入的各个部分，并在整个序列的上下文中对每个标记进行加权和表示。
“遮蔽”多头注意力通过遮蔽未来标记的注意力权重，限制每个标记只能关注自身和之前的标记。

数据首先通过三个并行的线性层，每个层都是 nn.Linear(768, 768)，这保留了输入形状。输出被称为**查询（Q）**、**键（K）**和**值（V）**，它们的张量形状均为 $[1,4,768][1,\; 4,\; 768]$。

然后，嵌入维度根据 $768=8\ast 96768\; =\; 8\; *\; 96$ 进行拆分，其中 8 是头数，96 是头大小。这会将张量重塑为 $[1,4,8,96][1,\; 4,\; 8,\; 96]$。为了使维度与矩阵乘法对齐，序列长度和头大小进行了转置，从而得到 Q、K 和 V 的形状均为 $[1,8,4,96][1,\; 8,\; 4,\; 96]$。

注意力机制，如原始论文所定义，计算如下：
$$\text{Attention}(Q,K,V)=\text{softmax}\left(\frac{Q{K}^{\mathrm{\top }}}{\sqrt{d_k}}\right)V\backslash text\{Attention\}(Q,\; K,\; V)\; =\; \backslash text\{softmax\}\backslash left(\backslash frac\{QK^\backslash top\}\{\backslash sqrt\{d\_k\}\}\backslash right)\; V$$

首先，计算 $Q{K}^{\mathrm{\top }}QK^\backslash top$

• $K^{\mathrm{\top }}K^\backslash top$（转置后的 Key）的形状为 $[1,8,96,4][1,\; 8,\; 96,\; 4]$。
• $Q{K}^{\mathrm{\top }}QK^\backslash top$ 得到形状为 $[1,8,4,4][1,\; 8,\; 4,\; 4]$ 的张量，计算方式为 $[1,8,4,96]\times [1,8,96,4][1,\; 8,\; 4,\; 96]\; \backslash times\; [1,\; 8,\; 96,\; 4]$。

请参阅上面的矩阵乘法详情，以便更清晰地理解维度转换。

• 遮蔽
  遮蔽确保每个标记只关注自身和之前的标记，防止模型在生成过程中访问未来的标记。

• 注意力权重
  注意力权重使用以下公式计算：
  $$\text{softmax}\left(\frac{Q{K}^{\mathrm{\top }}}{\sqrt{d_k}}\right)\backslash text\{softmax\}\backslash left(\backslash frac\{QK^\backslash top\}\{\backslash sqrt\{d\_k\}\}\backslash right)$$
  其中 $\sqrt{d_k}=\sqrt{96}\backslash sqrt\{d\_k\}\; =\; \backslash sqrt\{96\}$。通过头大小进行缩放可以防止值的大幅差异，而 softmax 确保每个标记的向量表示之和为 1。这会将 -inf 值置零，同时保留正值，只影响张量值，不影响其形状。

• 计算注意力
  注意力权重乘以值：
  $$\text{softmax}\left(\frac{Q{K}^{\mathrm{\top }}}{\sqrt{d_k}}\right)V\backslash text\{softmax\}\backslash left(\backslash frac\{QK^\backslash top\}\{\backslash sqrt\{d\_k\}\}\backslash right)\; V$$
  $$[1,8,4,4]\times [1,8,4,96]=[1,8,4,96][1,\; 8,\; 4,\; 4]\; \backslash times\; [1,\; 8,\; 4,\; 96]\; =\; [1,\; 8,\; 4,\; 96]$$

• 拼接
  通过转置维度恢复头数：
  $$[1,8,4,96]\Rightarrow [1,4,8,96]\Rightarrow [1,4,768][1,\; 8,\; 4,\; 96]\; \backslash Rightarrow\; [1,\; 4,\; 8,\; 96]\; \backslash Rightarrow\; [1,\; 4,\; 768]$$

• 投影层
  一个带有 nn.Linear(768, 768) 的线性层将张量形状维持在 $[1,4,768][1,\; 4,\; 768]$。

观察
张量形状已恢复到 $[1,4,768][1,\; 4,\; 768]$，与注意力机制之前的初始形状匹配。这种一致性对于确保与模型后续层和计算的兼容性至关重要。

添加并归一化

此步骤涉及一个跳跃连接，其中注意力层之前和之后的张量被相加并归一化。相加确保张量值得到“更新”而不是替换，而归一化则防止值的指数增长。

这些添加和归一化操作在每个层之后应用，以保持张量的原始特性。

前馈网络

前馈网络通常由两个连续的线性层组成：一个扩展张量，另一个收缩张量，通常还伴随着一个用于正则化的 dropout 层。这些层引入非线性转换，使模型能够捕获嵌入维度（头数和头大小）中更丰富、更复杂的模式。

扩展因子通常为 $1\Rightarrow 31\; \backslash Rightarrow\; 3$，收缩因子为 $3\Rightarrow 13\; \backslash Rightarrow\; 1$。这导致线性层的结构如下：

nn.Linear(768, 3 * 768)  
nn.Linear(3 * 768, 768)  

最终输出形状返回到输入形状，即 $[1,4,768][1,\; 4,\; 768]$。保持此形状允许在前馈步骤后应用添加和归一化层。

此外，由于解码器的最终形状与其输入形状匹配，因此可以无缝堆叠多个连续的解码器层。

语言模型头

经过一系列解码器层后，张量到达一个最终的线性层，该层将 embed_dim $\Rightarrow \backslash Rightarrow$ 转换为 vocab_size。

这会产生一个形状为 $[1,4,9735][1,\; 4,\; 9735]$ 的张量（假设词汇表大小为 9735），其中：

• 1：批次大小
• 4：序列长度
• 9735：词汇表大小

应用 softmax 函数并计算模型输出与真实值之间的损失（或误差），优化器将使用该损失来更新模型权重。

在**遮蔽多头注意力层**中，每个标记的注意力仅使用当前**输入**标记及其前身来计算。由于输入向右偏移，生成新标记需要模型考虑所有先前标记的向量表示，包括它们与先前标记的关系。这种机制是生成式 AI 模型工作的基础，如下图所示：

Transformer 和交叉注意力

Transformer 架构通常由编码器-解码器结构组成，常用于上下文和输出不直接相关的任务，例如翻译。然而，如今仅解码器架构通常更受此类任务的青睐。

在编码器层中，张量形状的传播方式与解码器类似，但有一个关键区别：注意力层不应用遮蔽，允许每个标记关注序列中的所有其他标记，无论是在之前还是之后。

例如，考虑上下文 I am at home 和目标 <bos> je suis à la maison。张量形状将是：

• 上下文（I am at home）：$[1,4][1,\; 4]$
• 目标（<bos> je suis à la maison）：$[1,6][1,\; 6]$

输入和目标的序列长度不同，这由交叉注意力层处理。编码器输出的形状为 $[1,4,768][1,\; 4,\; 768]$，而解码器中遮蔽多头注意力层的输出为 $[1,6,768][1,\; 6,\; 768]$。在这里，**键（K）**和**值（V）**来自编码器，而**查询（Q）**来自解码器。

建议手动计算张量维度以求清晰，但下面提供了解决方案以供验证。

[点击此处查看解决方案]

注意力公式如下： $$\text{Attention}(Q,K,V)=\text{softmax}\left(\frac{Q{K}^{\mathrm{\top }}}{\sqrt{d_k}}\right)V\backslash text\{Attention\}(Q,\; K,\; V)\; =\; \backslash text\{softmax\}\backslash left(\backslash frac\{QK^\backslash top\}\{\backslash sqrt\{d\_k\}\}\backslash right)\; V$$

• 查询：$[1,6,768][1,\; 6,\; 768]$
• 键和值：$[1,4,768][1,\; 4,\; 768]$

拆分和转置维度 :

• 查询：$[1,8,6,96][1,\; 8,\; 6,\; 96]$
• 键和值：$[1,8,4,96][1,\; 8,\; 4,\; 96]$

计算注意力 :

• $K^{\mathrm{\top }}K^\backslash top$：$[1,8,96,4][1,\; 8,\; 96,\; 4]$
• $Q{K}^{\mathrm{\top }}QK^\backslash top$：$[1,8,6,96]\ast [1,8,96,4]=[1,8,6,4][1,\; 8,\; 6,\; 96]\; *\; [1,\; 8,\; 96,\; 4]\; =\; [1,\; 8,\; 6,\; 4]$
• $\text{softmax}\left(\frac{Q{K}^{\mathrm{\top }}}{\sqrt{d_k}}\right)=[1,8,6,4]\backslash text\{softmax\}\backslash left(\backslash frac\{QK^\backslash top\}\{\backslash sqrt\{d\_k\}\}\backslash right)\; =\; [1,\; 8,\; 6,\; 4]$
• $\text{softmax}\left(\frac{Q{K}^{\mathrm{\top }}}{\sqrt{d_k}}\right)V=[1,8,6,4]\ast [1,8,4,96]\Rightarrow [1,8,6,96]\backslash text\{softmax\}\backslash left(\backslash frac\{QK^\backslash top\}\{\backslash sqrt\{d\_k\}\}\backslash right)\; V\; =\; [1,\; 8,\; 6,\; 4]\; *\; [1,\; 8,\; 4,\; 96]\; \backslash Rightarrow\; [1,\; 8,\; 6,\; 96]$

拼接

• $[1,8,6,96]\Rightarrow [1,6,8,96]\Rightarrow [1,6,768][1,\; 8,\; 6,\; 96]\; \backslash Rightarrow\; [1,\; 6,\; 8,\; 96]\; \backslash Rightarrow\; [1,\; 6,\; 768]$

观察
张量形状已恢复到 $[1,6,768][1,\; 6,\; 768]$，与**遮蔽多头注意力层的输出形状**匹配。这种一致性确保了与后续层的兼容性。

结语

这篇博文旨在更清晰地解释注意力机制的工作原理以及张量形状如何在 Transformer 架构中传播。

如果您觉得这篇博文有帮助，请考虑点赞 🤗。

如有任何反馈或问题，请通过我的作品集上列出的任何联系方式与我联系：https://not-lain.github.io/。