A3 · Self-Attention and Transformers

做什么

题面来自官方材料（本地原件在 _local/，官方链接见「材料」）。 编号是原件的编号（notebook 是 cell 序号、PDF 是页码），方便跟官方题面对号。 本项目不提供答案（D53）。

第 1 页

CS 224N Winter 2026 Assignment 3

Self-Attention 与 Transformers

截止日期：February 5th, Thursday, 4:30 PM PST

本次作业是对 Transformer 的研究，Transformer 是当前前沿 LLM 所采用的主流架构。

本 problem set 包含三个问题：

• 第一个问题中，你将直观理解 Transformer 中的 self-attention 机制如何工作。
• 第二部分中，你将推导 positional encodings 的一些性质。
• 第三部分中，你将（几乎）从零开始编写一个 transformer，并在你的笔记本电脑上开始训练它。

请在 Gradescope 上正确标记题目，否则如果你不标记题目，助教将会扣分。

对于代码提交，运行：

bash create_submission.sh

它会将你的 model_solution.py、train.py 和 utils.py 打包成 zip。然后直接将生成的 submission.zip 上传到 Gradescope。

第 2 页

1. Attention 探索（Attention Exploration，14 points）

Multi-head self-attention 是 Transformer 的核心建模组件。本题中，我们将练习 self-attention 方程，并说明为什么 multi-headed self-attention 可能优于 single-headed self-attention。

回顾一下，attention 可以看作是对一个 query vector $q \in \mathbb{R}^d$、一组 value vectors $\{v_1,\dots,v_n\}$（$v_i \in \mathbb{R}^d$）以及一组 key vectors $\{k_1,\dots,k_n\}$（$k_i \in \mathbb{R}^d$）的操作，具体定义如下：

$$ c = \sum_{i=1}^{n} v_i \alpha_i \tag{1} $$

$$ \alpha_i = \frac{\exp(k_i^\top q)}{\sum_{j=1}^{n} \exp(k_j^\top q)} \tag{2} $$

其中 $\alpha = \{\alpha_1,\dots,\alpha_n\}$ 被称为“attention weights”。注意输出 $c \in \mathbb{R}^d$ 是 value vectors 的加权平均，权重为 $\alpha$。

(a) (3 points) 在 attention 中复制（Copying in attention）

Attention 的一个优点是：将某个 value vector“复制”到输出 $c$ 特别容易。本题中，我们将说明为什么会这样。

• i. (2 points) 分布 $\alpha$ 通常相对“弥散”；概率质量分散在许多不同的 $\alpha_i$ 之间。然而，情况并非总是如此。请描述（用一句话）在什么条件下，categorical distribution $\alpha$ 会把几乎全部权重放在某个 $\alpha_j$ 上，其中 $j \in \{1,\dots,n\}$（即 $\alpha_j \gg \sum_{i \ne j} \alpha_i$）。关于 query $q$ 和/或 keys $\{k_1,\dots,k_n\}$，必须满足什么条件？
• ii. (1 point) 在你于 (i) 中给出的条件下，描述输出 $c$。

(b) (2 points) 两个的平均（An average of two）

不是只关注一个向量 $v_j$，Transformer 模型可能希望纳入来自多个 source vectors 的信息。

考虑我们改为希望纳入来自两个向量 $v_a$ 和 $v_b$ 的信息，它们对应的 key vectors 为 $k_a$ 和 $k_b$。

假设 (1) 所有 key vectors 都正交，因此对所有 $i \ne j$ 有 $k_i^\top k_j = 0$；(2) 所有 key vectors 的范数都为 1。求一个 query vector $q$ 的表达式，使得 $c \approx \frac{1}{2}(v_a + v_b)$，并证明你的答案。[注 *]（回忆你在 (a) 中学到的内容。）

(c) (5 points) 单头 attention 的缺点（Drawbacks of single-headed attention）

在上一部分中，我们看到了 single-headed attention 如何可能同等地关注两个 values。同样的概念可以很容易地推广到任意 values 子集。本题中我们将看到为什么这不是一个实用的解决方案。

考虑一组 key vectors $\{k_1,\dots,k_n\}$，它们现在是随机采样得到的，$k_i \sim \mathcal{N}(\mu_i, \Sigma_i)$，其中均值 $\mu_i \in \mathbb{R}^d$ 已知，但协方差 $\Sigma_i$ 未知（除非题目中另有说明）。进一步假设均值 $\mu_i$ 全部垂直；若 $i \ne j$，则 $\mu_i^\top \mu_j = 0$，且范数为单位范数，$\|\mu_i\| = 1$。

• i. (2 points) 假设协方差矩阵为 $\Sigma_i = \alpha I$，$\forall i \in \{1,2,\dots,n\}$，其中 $\alpha$ 极小，可忽略。设计一个用 $\mu_i$ 表示的 query $q$，使得如前所述 $c \approx \frac{1}{2}(v_a + v_b)$，并简要论证它为什么有效。
• ii. (3 points) 尽管 single-headed attention 对 keys 中的小扰动具有抵抗性，某些类型的较大扰动可能会带来更大的问题。在某些情况下，一个 key vector $k_a$ 的范数可能比其他 key vectors 更大或更小，但仍然指向与 $\mu_a$ 相同的方向。[注 †]

例如，让我们考虑 item $a$ 的协方差为 $\Sigma_a = \alpha I + \frac{1}{2}(\mu_a \mu_a^\top)$，其中 $\alpha$ 极小，可忽略（如图 1 所示）。这会导致 $k_a$ 大致指向与 $\mu_a$ 相同的方向，但幅度上具有很大的方差。进一步，令所有 $i \ne a$ 有 $\Sigma_i = \alpha I$。

[注 *] 提示：虽然 softmax function 永远不会精确地平均这两个向量，但通过在表达式中使用一个很大的标量倍数，你可以接近这一结果。

[注 †] 与原始 Transformer 不同，一些较新的 Transformer 模型在 attention 之前应用 layer normalization。在这些 pre-layernorm 模型中，keys 的范数不能相差太大，这使得本题中的情况不太可能发生。

第 3 页

Figure 1: 向量 $\mu_a$（这里以 2D 为例展示），$k_a$ 可能取值的范围用红色表示。如前所述，$k_a$ 大致指向与 $\mu_a$ 相同的方向，但幅度可能更大或更小。

（接 (c) ii.）当你多次采样 $\{k_1,\dots,k_n\}$，并使用你在第 i. 部分定义的 $q$ 向量时，对于不同的样本，你预计向量 $c$ 在定性上会是什么样子？思考它与第 (i) 部分有何不同，以及 $c$ 的方差会受到怎样的影响。

(d) (3 points) 多头 attention 的优点（Benefits of multi-headed attention）

现在我们来看看 multi-headed attention 的一些能力。我们将考虑一个简单版本的 multi-headed attention，它与我们已介绍的 single-headed self-attention 相同，只是定义了两个 query vectors（$q_1$ 和 $q_2$），从而得到一对向量（$c_1$ 和 $c_2$），每个都是 single-headed attention 在给定各自 query vector 后的输出。Multi-headed attention 的最终输出是它们的平均值，$\frac{1}{2}(c_1 + c_2)$。

如在问题 1(c) 中一样，考虑一组随机采样的 key vectors $\{k_1,\dots,k_n\}$，$k_i \sim \mathcal{N}(\mu_i, \Sigma_i)$，其中均值 $\mu_i$ 已知，但协方差 $\Sigma_i$ 未知。也如前所述，假设均值 $\mu_i$ 两两正交；若 $i \ne j$，则 $\mu_i^\top \mu_j = 0$，且范数为单位范数，$\|\mu_i\| = 1$。

• i. (1 point) 假设协方差矩阵为 $\Sigma_i = \alpha I$，其中 $\alpha$ 极小，可忽略。设计用 $\mu_i$ 表示的 $q_1$ 和 $q_2$，使得 $c$ 大约等于 $\frac{1}{2}(v_a + v_b)$。注意 $q_1$ 和 $q_2$ 应该有不同的表达式。
• ii. (2 points) 假设协方差矩阵为 $\Sigma_a = \alpha I + \frac{1}{2}(\mu_a \mu_a^\top)$，其中 $\alpha$ 极小，可忽略；且对所有 $i \ne a$，$\Sigma_i = \alpha I$。取你在第 i. 部分设计的 query vectors $q_1$ 和 $q_2$。对于 key vectors 的不同样本，你预计输出 $c$ 在定性上会是什么样子？请用 $c_1$ 和 $c_2$ 的方差简要解释。你可以忽略 $k_a^\top q_i < 0$ 的情况。

(e) (1 point)

基于 (d) 部分，简要总结 multi-headed attention 如何克服你在 (c) 部分指出的 single-headed attention 的缺点。

第 4 页

2. Position Embeddings Exploration（6 分）

Position embeddings 是 Transformer 架构中的一个重要组成部分，使模型能够根据 token 在序列中的位置区分它们。在这个问题中，我们将探讨 Transformer 中 positional embeddings 的必要性，以及它们可以如何设计。

回顾一下，Transformer 架构的关键组件是 self-attention layer 和 feed-forward neural network layer。给定输入张量 $X \in \mathbb{R}^{T \times d}$，其中 $T$ 是序列长度，$d$ 是 hidden dimension，self-attention layer 计算如下：

$$ Q = XW_Q,\quad K = XW_K,\quad V = XW_V $$

$$ H = \mathrm{softmax}\left(\frac{QK^\top}{\sqrt{d}}\right)V $$

其中 $W_Q, W_K, W_V \in \mathbb{R}^{d \times d}$ 是权重矩阵，$H \in \mathbb{R}^{T \times d}$ 是输出。

接下来，feed-forward layer 应用以下变换：

$$ Z = \mathrm{ReLU}(HW_1 + \mathbf{1} \cdot b_1)W_2 + \mathbf{1} \cdot b_2 $$

其中 $W_1, W_2 \in \mathbb{R}^{d \times d}$，$b_1, b_2 \in \mathbb{R}^{1 \times d}$ 是权重和偏置；$\mathbf{1} \in \mathbb{R}^{T \times 1}$ 是全 1 向量[注 ‡]；$Z \in \mathbb{R}^{T \times d}$ 是最终输出。

（注意，为简单起见，我们省略了 Transformer 架构的一些细节。）

(a)（4 分）置换输入。

i.（3 分） 假设我们置换输入序列 $X$，使 token 被随机打乱。这可以表示为乘以一个 permutation matrix $P \in \mathbb{R}^{T \times T}$，即 $X_{\mathrm{perm}} = PX$。（关于 permutation matrices 的回顾，见 Wikipedia。）

证明对于置换后的输入 $X_{\mathrm{perm}}$，输出 $Z_{\mathrm{perm}}$ 将是 $Z_{\mathrm{perm}} = PZ$。

已知对于任意 permutation matrix $P$ 和任意矩阵 $A$，以下等式成立：

$$ \mathrm{softmax}(PAP^\top) = P\,\mathrm{softmax}(A)\,P^\top $$

以及

$$ \mathrm{ReLU}(PA) = P\,\mathrm{ReLU}(A). $$

ii.（1 分） 思考你在第 i 部分推导出的结果的含义。解释 Transformer 模型的这个性质在处理文本时为什么可能有问题。

(b)（2 分）Position embeddings 是编码每个 token 在序列中位置的向量。它们会在输入 word embeddings 被送入 Transformer 之前加到这些 word embeddings 上。一种方法是使用位置和 embedding 维度的固定函数来生成 position embedding。如果输入 word embeddings 是 $X \in \mathbb{R}^{T \times d}$，则 position embeddings $\Phi \in \mathbb{R}^{T \times d}$ 按如下方式生成：

$$ \Phi(t,2i) = \sin\left(t/10000^{2i/d}\right) $$

$$ \Phi(t,2i+1) = \cos\left(t/10000^{2i/d}\right) $$

其中 $t \in \{0, 1, \ldots, T-1\}$，$i \in \{0, 1, \ldots, d/2 - 1\}$[注 §]。

具体来说，position embeddings 会被加到输入 word embeddings 上：

$$ X_{\mathrm{pos}} = X + \Phi $$

i.（1 分） 你认为 position embeddings 会有助于解决你在第 (a) 部分指出的问题吗？如果有，请解释如何有帮助；如果没有，请解释为什么没有。

ii.（1 分） 输入序列中两个不同 token 的 position embeddings 可以相同吗？如果可以，请给出一个例子。如果不可以，请解释为什么不可以。

[注 ‡] 与 $\mathbf{1}$ 的外积表示 broadcasting 操作，并使 feed-forward network 的记号在数学上严谨。

[注 §] 这里假设 $d$ 为偶数，这在大多数模型中通常是这样的情况。

第 5 页

3. Coding a transformer from scratch（30 分）

在这个问题中，你将填写代码，实现一个 decoder only、GPT-2 风格的 transformer，以及一个简单的 training loop。

对于 (a) 部分，我们为每个子问题包含了 unit tests，可以在你的笔记本电脑上本地运行。如果你通过 unit test，该子问题就会获得满分。不要编辑 unit test 文件，因为我们会在你提交代码时单独运行这些测试。

在作业的这一部分，以下提示可能会有用：

• 添加 assert 语句，检查张量的 shape 是否与你想的一致。
• 考虑使用 Jaxtyping 包来为张量的 shape 添加 type hint。
• 考虑使用 einops 包来操作张量（einops.rearrange 尤其有用）。

这不仅会帮助你写出 bug 更少的代码，还会让你的代码可读性高得多。

(a)（20 分）

在问题的这一部分，我们将在 model.py 文件中实现一个 transformer。该文件包含许多不同的类，需要你来实现。最终，你将得到一个 Transformer 类的实现，其 forward 和 generate 方法可以正常工作。在 (b) 部分，我们将（开始）训练你实现的 Transformer。

i.（0 分） 熟悉 model.py 文件中的各个类。我们会要求你按 MLP、CausalAttention、DecoderBlock、最后 Transformer 的顺序实现它们。我们让你按这个顺序实现这些类，是因为这是依赖顺序。Transformer 依赖 DecoderBlock，而 DecoderBlock 又依赖 CausalAttention 和 MLP。

ii.（1 分） 实现 MLP.forward。检查你是否通过对应的 test。

iii.（6 分） 实现 CausalAttention.forward。检查你是否通过对应的 test。

iv.（2 分） 实现 DecoderBlock.forward。检查你是否通过对应的 test。

v.（6 分） 实现 Transformer.forward。检查你是否通过对应的 test。

vi.（5 分） 实现 Transformer.generate。检查你是否通过对应的 test。注意：这个函数应当实现 greedy decoding。

(b)（10 分）

完成 (a) 部分后，你现在有一个可以工作的 Transformer 模型。如果你查看 Transformer.__init__，可以看到当你创建 Transformer 类的实例时，我们会根据 Transformer._init_weights 方法用随机权重初始化模型。在问题的这一部分，你将实现一个 training loop，并开始在笔记本电脑上本地训练一个小模型。

i.（0 分） 查看 train.py 文件，并熟悉 training loop。我们将运行这段代码来训练模型。

ii.（7 分） 首先，实现 Transformer.get_loss_on_batch。这个函数将一批 token 映射为单个 loss 值。我们在 train.py 中使用这个函数来获取一个 batch 上的 loss。检查你是否通过对应的 test。

iii.（3 分） 运行 train.py。这会使用你实现的 Transformer.get_loss_on_batch 在 100 个 batch 的数据上训练模型。训练结束时，它会把训练过程中的 training loss 和 gradient norm 图保存到 losses_and_grad_norms.png，并在下面附上该图片。如果一切正确，你应该会看到 loss 曲线下降。

(c)（9 分）（Bonus）

在这个可选的 bonus 问题中，你的目标是加速学习过程。我们会把 gradient steps 的数量固定为 100，但你可以改变 train.py 或 model.py 中的其他任何东西来加速训练。

如果 100 步后的最终 loss 低于你在 (b) iii) 部分报告的 baseline 曲线，我们就认为一项改动成功了。

第 6 页

我们将为你做出的每一个能带来加速的不同改动给 3 分。因此，要拿到满分，你需要对训练文件做出三种不同的改动，且每一种都能带来加速。这些改动应当是能够叠加的。例如，你可以先改变 learning rate，使 loss 降低。然后你可能保留这个更好的 learning rate，并将它与第二个改动（例如改变 optimizer 或模型架构）结合，从而使 loss 进一步降低。注意，把 learning rate 改成三个不同的值只算作一个想法：我们要找的是三种不同类型的想法。

完成后，提交：

• 对每一项改动的描述。
• 最多三条新的 learning curves，每个改动一条；另外还要包含 b iii) 中的 baseline。
• 你在 100 步后达到的最低 loss。

材料

官方地址（W26 ✓ 全部验证过可下载 ✓）。本地副本放 _local/（只服务本地、不外传）。

怎么自查

• 官方仓库自带的测试（pytest / sanity check）—— ⚠️ Gradescope 只对注册学生开放， 自学者以仓库自带测试为准；
• 提交脚本：A3 / A4 提到 bash create_submission.sh。

这份作业和书的关系

见上表「配哪几讲」。等题面中英对照做完，这里会换成章节标题原文，并加定位校验 （指引指向一个不存在的标题会被 nav_check 报出来）。