Transformer电路逆向工程练习题集
Exercises
本练习集旨在通过动手编写注意力头的具体权重矩阵，从参数层面精确理解Transformer工作机制。内容涵盖：详解注意力头中W_Q、W_K、W_V、W_out矩阵的作用；分析读写子空间的控制矩阵及其乘积意义；探讨如何用两个矩阵等效表示注意力头及其秩的含义；研究跨层注意力头如何通过矩阵运算传递信息。并通过具体数值示例，演示多个“前词注意力头”如何协作实现“查看前两个词”的虚拟功能，以及手动构建实现“归纳头”的“指针算法”步骤。
帮助开发者亲手拆解Transformer内部机制，提升可解释性研究能力。
本练习集是我们用于逆向工程 Transformer 的数学框架的补充材料。这些练习通过编写注意力头中各个权重来实现算法。我们发现，这有助于加深我们自身的理解，确保我们真正掌握 Transformer 直至单个参数的完整机制原理，而不会模糊任何困惑之处。
下方提供了参考答案。
练习题
热身
- 从高层级描述 Transformer 架构
- 详细描述单个注意力头的工作原理，涉及矩阵 W_Q、W_K、W_V 和 W_{out}。（注意力头的方程和代码通常是将一层中所有注意力头拼接在一起编写的。这种实现计算效率更高，但更难理解，因此我们希望描述单个注意力头。）
- 注意力头将信息从一个 token 的残差流子空间移动到另一个 token 残差流中的不同子空间。哪个矩阵控制被读取的子空间，哪个矩阵控制被写入的子空间？它们的乘积意味着什么？
- 注意力头关注哪些 token 仅由定义该注意力头的四个矩阵中的两个控制。这两个矩阵是哪个？
- 注意力头可以用两个矩阵（W_Q^T · W_k 和 W_{out} · W_v）而不是四个矩阵来表示。在前两个问题中，你已经对这些矩阵给出了解释。现在，仅参考这两个矩阵来写出一个注意力头。
- 这些矩阵的秩是多少？
- 你想了解一个注意力头是否正在读取前一个注意力头的输出。W_V^2 · W_{out}^1 能告诉你什么？奇异值能告诉你什么？
练习 1 —— 构建一个简单的虚拟注意力头
小型 Transformer 通常有多个关注前一个 token 的注意力头，但没有关注前两个 token 的注意力头。在本练习中，我们将了解两个前一个 token 头如何实现一个“虚拟注意力头”来关注两个 token 之前的内容，而无需为此牺牲一个完整的注意力头。
让我们考虑两个注意力头，头 1 和头 2，它们都关注前一个 token。头 1 在第一层，头 2 在第二层。为了便于写出显式矩阵，我们设两个头的 k、q 和 v 向量均为 4 维，残差流为 16 维。
- (a) 写出头 1 的 W_V^1 和 W_{out}^1，使得该头将其输入的维度 0–3 复制到其输出的维度 8–11。
- (b) 写出头 2 的 W_V^2 和 W_{out}^2，使得它复制前一个 token 的另外 3 个维度，以及两个 token 之前的一个维度（使用前一个头写入的维度）。
- (c) 展开 W_{\text{net}}^1 = W_{out}^1 \cdot W_V^1 和 W_{\text{net}}^2 = W_{out}^2 \cdot W_V^2。这些矩阵告诉了你什么？
- (d) 展开以下矩阵：两个 token 复制：W_{\text{net}}^2 \cdot W_{\text{net}}^1。一个 token 复制：W_{\text{net}}^2 \cdot \text{Id} ~+~ \text{Id} \cdot W_{\text{net}}^1。
- 观察：当我们通常考虑一个注意力头时，它们需要将所有容量专用于一项任务。在这个例子中，两个头将 7/8 的容量用于一项任务，将 1/8 的容量用于另一项任务。
练习 2 — 使用归纳头复制文本（指针算术版本）
神经网络进行的最简单的上下文内元学习是增加在此上下文中之前见过的序列的概率。这是通过一个“归纳头”完成的，它查看我们上次看到某个 token 之后发生了什么。
实现归纳头至少有两条算法。在本练习中，你将手动构建“指针算术”算法。
- (a) 设 u^{\text{cont}}_0, ~~ u^{\text{cont}}_1, ~~ \ldots ~~ u^{\text{cont}}_n 为内容嵌入向量的主成分。用 W_Q 和 W_K 表示一个注意力头（查询和键的维度为 4），该注意力头会选择与当前 token 内容相似的 token，包括当前 token 自身。
- (b) 设 u^{\cos}_0, ~ u^{\sin}_0, ~ u^{\cos}_1, ~ u^{\sin}_1, ~ ... 为一个描述位置嵌入的基，该基由编码 token 位置正弦和余弦嵌入的向量（例如 \lambda\cos(\alpha_0 n_{tok})）组成，且幅值递减。用 W_Q 和 W_K 表示一个注意力头（查询和键的维度为 4），该注意力头会自选择当前 token 位置。
- (c) 使用 (b) 中描述的位置嵌入基，用 W_Q 和 W_K 表示一个注意力头（查询和键的维度为 4），该注意力头会自选择*上一个* token 位置。提示：考虑一个二维旋转矩阵。
- (d) 用 W_Q 和 W_K 表示一个注意力头（查询和键的维度为 8），该注意力头会选择与当前 token 内容相似的 token，但倾向于不关注自身。提示：参考 (b)，并为键和查询使用额外的 4 个维度。
- (e) 为你 (d) 中描述的注意力头写出 W_V 和 W_{out}，使其从所关注的 token 中提取位置嵌入的最大 8 个维度，并将它们写入向量 v_0, v_1, ... 中。
- (f) 用 W_Q 和 W_K 表示一个注意力头，该注意力头会关注当前 token 上一个副本之后的那个 token。提示：使用 (e) 中注意力头的输出以及你在 (c) 中使用的策略。
练习 3 — 使用归纳头复制文本（前一个 token K 组合版本）
某些位置编码机制（例如旋转位置注意力）不会将位置信息暴露给 W_V 矩阵乘法。使用这些机制训练的 Transformer 无法采用上一练习中 (e) 和 (f) 的策略来操控位置编码向量。
对于这些 Transformer 架构，我们观察到了一种替代机制：第一个注意力头将关于前一个 token 的信息复制到一个子空间中，第二个注意力头则利用该子空间来构建查询和键。假设采用与上述相同的位置编码机制，请写出实现该算法的一对注意力头的 \(W^1_Q\)、\(W^1_K\)、\(W^1_V\)、\(W^1_O\) 以及 \(W^2_Q\) 和 \(W^2_K\)。
解答
热身练习
关于所有这些问题的讨论，请参阅我们关于 Transformer 电路（可解释性研究）的论文。
练习 1 — 构建一个简单的虚拟注意力头
(1)(a) 写出注意力头 1 的 \(W_V^1\) 和 \(W_{out}^1\)，使得该头将其输入的维度 0–3 复制到其输出的维度 8–11。
(1)(b) 写出注意力头 2 的 \(W_V^2\) 和 \(W_{out}^2\)，使得该头复制前一个 token 的另外 3 个维度，以及两个 token 前的 token 的一个维度（利用前一个头写入的维度）。
请注意，这里有许多正确答案。一个关键的性质是，\(W_V^2\) 在某一列中有一个 1，该列对应于 \(W_O^1\) 的某一行；同时，在 \(W_V^1\) 和 \(W_O^1\) 未触及的列中有三个 1。
(1)(c) 展开计算 \(W_{\text{net}}^1 = W_{out}^1 \cdot W_V^1\) 和 \(W_{\text{net}}^2 = W_{out}^2 \cdot W_V^2\)。这些矩阵告诉了你什么？
这些矩阵描述了当一个注意力头将信息从一个位置（被关注的位置）的残差流移动到另一个位置时的完整操作：
\(h_i(x) = \sum_j A_{ij}W_{net}x_{j}\)
(1)(d) 展开计算以下矩阵：两个 token 复制：\(W_{\text{net}}^2 \cdot W_{\text{net}}^1\)。一个 token 复制：\(W_{\text{net}}^2 \cdot \text{Id} ~+~ \text{Id} \cdot W_{\text{net}}^1\)。
练习 2 — 使用归纳头复制文本（指针算术版本）
(2)(a) 设 \(u^{\text{cont}}_0, ~~ u^{\text{cont}}_1, ~~ \ldots ~~ u^{\text{cont}}_n\) 为内容嵌入的主成分。请写出一个注意力头（具有 4 维查询和键）的 \(W_Q\) 和 \(W_K\)，该头用于选择与当前 token 内容相似的 token，包括当前 token 本身。
\(W_Q = W_K = \begin{pmatrix}u_0^{cont} \\u_1^{cont} \\u_2^{cont} \\u_3^{cont} \end{pmatrix}\)
我们只需将内容维度投影到两个注意力头上。请注意，在本示例及后续所有示例中，我们都可以将任意旋转组合到两个矩阵中，从而得到一个功能等价的注意力头。我们展示的是我们认为这些矩阵最直接的版本。
(2)(b) 设 \( u^{\cos}_0, ~ u^{\sin}_0, ~ u^{\cos}_1, ~ u^{\sin}_1, ~ ... \) 为一个基，该基用向量描述位置嵌入，这些向量编码了 token 位置的正弦和余弦嵌入（例如 \(\lambda\cos(\alpha_0 n_{tok})\)），且幅值递减。请为一个自选择当前 token 位置的注意力头（查询和键的维度为 4）写出 \(W_Q\) 和 \(W_K\)。
\( W_Q = W_K = \begin{pmatrix} u_0^{cos} \\ u_0^{sin} \\ u_1^{cos} \\ u_1^{sin} \end{pmatrix} \)
(2)(c) 使用 (b) 中描述的位置嵌入基，请为一个自选择*前一个* token 位置的注意力头（查询和键的维度为 4）写出 \(W_Q\) 和 \(W_K\)。提示：考虑一个二维旋转矩阵。
请注意，根据定义，\( u^{\cos}_0, ~ u^{\sin}_0, … \) 将 token 索引 \(n\) 编码为余弦或正弦波。如果我们将对应的余弦和正弦对放在一起考虑，可以将其视为一个二维点：
\( (u_0^{\cos}, ~u_0^{\sin})\cdot x^0_n ~=~ \lambda_0(\cos(\alpha_0n),~ \sin(\alpha_0n)) \)
如果我们对 token \(n-1\) 做同样的处理，会发现它是 token \(n\) 的点旋转了 \(-\alpha_0\) 之后的结果：
\( \begin{aligned} (u_0^{\cos}, u_0^{\sin}) \cdot x^0_{n-1} &~=~ \lambda_0(\cos(\alpha_0(n-1)),~ \sin(\alpha_0(n-1))) \\ &~=~ \lambda_0(\cos(\alpha_0n-\alpha_0),~ \sin(\alpha_0n-\alpha_0))\\ &~=~ \lambda_0R_{-\alpha_0}(\cos(\alpha_0n),~ \sin(\alpha_0n) \end{aligned} \)
因此，我们想要获取位置嵌入向量，将正弦和余弦分量配对，并执行一个二维旋转。一个绕原点的二维旋转为：
\( R_\theta=\begin{pmatrix} \cos\theta & -\sin\theta \\ \sin\theta & \cos\theta \end{pmatrix} \)
我们将两个这样的旋转与位置基的投影组合起来，得到：
\( W_K = \textrm{(与 (b) 部分相同)} \)
W_Q = \begin{pmatrix} \cos\alpha_0 & \sin\alpha_0 & & \\-\sin\alpha_0 & \cos\alpha_0 & & \\& & \cos\alpha_1 & \sin\alpha_1 \\& & -\sin\alpha_1 & \cos\alpha_1 \end{pmatrix} \begin{pmatrix} u_0^{cos} \\ u_0^{sin} \\ u_1^{cos} \\ u_1^{sin} \end{pmatrix}
(2)(d) 写出一个注意力头（查询和键的维度为 8）的 W_Q 和 W_K，该注意力头会选择与当前 token 内容相似的 token，但倾向于不关注自身。提示：参考 (b)，并为键和查询额外使用 4 个维度。
W_K=\begin{pmatrix} u_0^{cont} \\ u_1^{cont} \\ u_2^{cont} \\ u_3^{cont} \\ u_0^{cos} \\ u_0^{sin} \\ u_1^{cos} \\ u_1^{sin} \end{pmatrix} ~~~~W_Q=\begin{pmatrix} u_0^{cont} \\ u_1^{cont} \\ u_2^{cont} \\ u_3^{cont} \\ -\beta{}u_0^{cos} \\ -\beta{}u_0^{sin} \\ -\beta{}u_1^{cos} \\ -\beta{}u_1^{sin} \end{pmatrix}
β 是一个参数，允许我们调整计算中“相同 token”部分和“非当前位置”部分的相对权重。
(2)(e) 写出你在 (d) 中描述的注意力头的 W_V 和 W_{out}，使其从所关注的 token 中提取位置嵌入的最大 8 个维度，并将它们写入向量 v_0, v_1, ...。
W_V=\begin{pmatrix} u_0^{cos} \\ u_0^{sin} \\ u_1^{cos} \\ u_1^{sin} \\ u_2^{cos} \\ u_2^{sin} \\ u_3^{cos} \\ u_3^{sin} \end{pmatrix} ~~~~W_O = \begin{pmatrix}v_0^T&v_1^T&v_2^T&\ldots{}\end{pmatrix}
(2)(f) 写出一个注意力头的 W_Q 和 W_K，该注意力头会关注当前 token 上一个副本之后的那个 token。提示：使用 (e) 中头的输出以及你在 (c) 中使用的策略。
(d) 和 (e) 中描述的头将“该 token 上一个实例的位置”放入由 v_0, v_1, ... 定义的子空间中。我们投影出该子空间，并以与 (c) 相同的方式旋转它，只是方向是向前而非向后：
W_K = \begin{pmatrix}u_0^{cos} \\ u_0^{sin} \\ u_1^{cos} \\ u_1^{sin}\end{pmatrix} ~~~~ W_Q = \begin{pmatrix} \cos\alpha_0 & -\sin\alpha_0 & & \\ \sin\alpha_0 & \cos\alpha_0 & & \\ & & \cos\alpha_1 & -\sin\alpha_1 \\ & & \sin\alpha_1 & \cos\alpha_1 \end{pmatrix} \begin{pmatrix} v_0\\ v_1\\ v_2\\ v_3 \end{pmatrix}
练习 3 — 使用归纳头复制文本（前一个 token 的 K 组合版本）
第一个头将前一个 token 的“内容”子空间复制到当前位置的 v_0, v_1, \ldots{} 子空间中：
W_K = \begin{pmatrix}u_0^{cos} \\ u_0^{sin} \\ u_1^{cos} \\ u_1^{sin}\end{pmatrix} ~~~~ W_Q = \begin{pmatrix} \cos\alpha_0 & \sin\alpha_0 & & \\ -\sin\alpha_0 & \cos\alpha_0 & & \\ & & \cos\alpha_1 & \sin\alpha_1 \\ & & -\sin\alpha_1 & \cos\alpha_1 \end{pmatrix} \begin{pmatrix}u_0^{cos} \\ u_0^{sin} \\ u_1^{cos} \\ u_1^{sin}\end{pmatrix}
W_V = \begin{pmatrix} u_0^{cont} \\ u_1^{cont} \\ u_2^{cont} \\ u_3^{cont} \end{pmatrix} ~~~~ W_O = \begin{pmatrix}v_0^T&v_1^T&v_2^T&\ldots{}\end{pmatrix}
然后，第二个头可以在其 key 投影中使用该子空间，并将其与当前 token 的内容进行比较：
W_K = \begin{pmatrix}v_0\\v_1\\v_2\\v_3\end{pmatrix} ~~~~W_Q = \begin{pmatrix} u_0^{cont} \\ u_1^{cont} \\ u_2^{cont} \\ u_3^{cont} \end{pmatrix}
关于
本文是《Transformer 电路的数学框架》的配套习题集。
致谢
我们感谢 Tuomas Oikarinen 和 Callum Canavan 发现了这些练习原版中的几处笔误。
来源：Anthropic：Transformer Circuits（可解释性研究） · transformer-circuits.pub