A3 · Self-Attention and Transformers

做什么

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

第 1 页

CS 224N Winter 2026 Assignment 3 Self-Attention and Transformers Due date: February 5th, Thursday, 4:30 PM PST This assignment is an investigation into Transformers, the prevailing architecture used for frontier LLMs. The pset has three questions: • In the first, you will gain intuition about how the self-attention mechanism in transformers works • In part two, you will derive some properties of positional encodings. • In the third part, you will code a transformer (almost) from scratch, and start training it on your laptop. Please tag the questions correctly on Gradescope, otherwise the TAs will take points off if you don’t tag questions. For code submission, run bash create submission.sh, which will zip your model solution.py, train.py, and utils.py. Then directly upload the resultant submission.zip to Gradescope. 1

第 2 页

CS 224N Winter 2026 Assignment 3 Page 2 of 6 1. Attention Exploration (14 points) Multi-head self-attention is the core modeling component of Transformers. In this question, we’ll get some practice working with the self-attention equations, and motivate why multi-headed self-attention can be preferable to single-headed self-attention. Recall that attention can be viewed as an operation on a query vector q ∈Rd, a set of value vectors {v1, . . . , vn}, vi ∈Rd, and a set of key vectors {k1, . . . , kn}, ki ∈Rd, specified as follows: c = n X i=1 viαi (1) αi = exp(k⊤ i q) Pn j=1 exp(k⊤ j q) (2) with alpha = {α1, . . . , αn} termed the “attention weights”. Observe that the output c ∈Rd is an average over the value vectors weighted with respect to α. (a) (3 points) Copying in attention. One advantage of attention is that it’s particularly easy to “copy” a value vector to the output c. In this problem, we’ll motivate why this is the case. i. (2 points) The distribution α is typically relatively “diffuse”; the probability mass is spread out between many different αi. However, this is not always the case. Describe (in one sentence) under what conditions the categorical distribution α puts almost all of its weight on some αj, where j ∈{1, . . . , n} (i.e. αj ≫P i̸=j αi). What must be true about the query q and/or the keys {k1, . . . , kn}? ii. (1 point) Under the conditions you gave in (i), describe the output c. (b) (2 points) An average of two. Instead of focusing on just one vector vj, a Transformer model might want to incorporate information from multiple source vectors. Consider the case where we instead want to incorporate information from two vectors va and vb, with corresponding key vectors ka and kb. Assume that (1) all key vectors are orthogonal, so k⊤ i kj = 0 for all i̸ = j; and (2) all key vectors have norm 1. Find an expression for a query vector q such that c ≈1 2(va + vb), and justify your answer.∗(Recall what you learned in part (a).) (c) (5 points) Drawbacks of single-headed attention: In the previous part, we saw how it was possible for a single-headed attention to focus equally on two values. The same concept could easily be extended to any subset of values. In this question we’ll see why it’s not a practical solution. Consider a set of key vectors {k1, . . . , kn} that are now randomly sampled, ki ∼N(µi, Σi), where the means µi ∈Rd are known to you, but the covariances Σi are unknown (unless specified otherwise in the question). Further, assume that the means µi are all perpendicular; µ⊤ i µj = 0 if i̸ = j, and unit norm, ∥µi∥= 1. i. (2 points) Assume that the covariance matrices are Σi = αI, ∀i ∈{1, 2, . . . , n}, for vanishingly small α. Design a query q in terms of the µi such that as before, c ≈1 2(va + vb), and provide a brief argument as to why it works. ii. (3 points) Though single-headed attention is resistant to small perturbations in the keys, some types of larger perturbations may pose a bigger issue. In some cases, one key vector ka may be larger or smaller in norm than the others, while still pointing in the same direction as µa.† As an example, let us consider a covariance for item a as Σa = αI + 1 2(µaµ⊤ a ) for vanishingly small α (as shown in figure 1). This causes ka to point in roughly the same direction as µa, but with large variances in magnitude. Further, let Σi = αI for all i̸ = a. ∗Hint: while the softmax function will never exactly average the two vectors, you can get close by using a large scalar multiple in the expression. †Unlike the original Transformer, some newer Transformer models apply layer normalization before attention. In these pre-layernorm models, norms of keys cannot be too different which makes the situation in this question less likely to occur.

第 3 页

CS 224N Winter 2026 Assignment 3 Page 3 of 6 Figure 1: The vector µa (shown here in 2D as an example), with the range of possible values of ka shown in red. As mentioned previously, ka points in roughly the same direction as µa, but may have larger or smaller magnitude. When you sample {k1, . . . , kn} multiple times, and use the q vector that you defined in part i., what do you expect the vector c will look like qualitatively for different samples? Think about how it differs from part (i) and how c’s variance would be affected. (d) (3 points) Benefits of multi-headed attention: Now we’ll see some of the power of multi-headed attention. We’ll consider a simple version of multi-headed attention which is identical to single- headed self-attention as we’ve presented it, except two query vectors (q1 and q2) are defined, which leads to a pair of vectors (c1 and c2), each the output of single-headed attention given its respective query vector. The final output of the multi-headed attention is their average, 1 2(c1 + c2). As in question 1(c), consider a set of key vectors {k1, . . . , kn} that are randomly sampled, ki ∼ N(µi, Σi), where the means µi are known to you, but the covariances Σi are unknown. Also as before, assume that the means µi are mutually orthogonal; µ⊤ i µj = 0 if i̸ = j, and unit norm, ∥µi∥= 1. i. (1 point) Assume that the covariance matrices are Σi = αI, for vanishingly small α. Design q1 and q2 in terms of µi such that c is approximately equal to 1 2(va + vb). Note that q1 and q2 should have different expressions. ii. (2 points) Assume that the covariance matrices are Σa = αI + 1 2(µaµ⊤ a ) for vanishingly small α, and Σi = αI for all i̸ = a. Take the query vectors q1 and q2 that you designed in part i. What, qualitatively, do you expect the output c to look like across different samples of the key vectors? Explain briefly in terms of variance in c1 and c2. You can ignore cases in which k⊤ a qi < 0. (e) (1 point) Based on part (d), briefly summarize how multi-headed attention overcomes the draw- backs of single-headed attention that you identified in part (c).

第 4 页

CS 224N Winter 2026 Assignment 3 Page 4 of 6 2. Position Embeddings Exploration (6 points) Position embeddings are an important component of the Transformer architecture, allowing the model to differentiate between tokens based on their position in the sequence. In this question, we’ll explore the need for positional embeddings in Transformers and how they can be designed. Recall that the crucial components of the Transformer architecture are the self-attention layer and the feed-forward neural network layer. Given an input tensor X ∈RT ×d, where T is the sequence length and d is the hidden dimension, the self-attention layer computes the following: Q = XWQ, K = XWK, V = XWV H = softmax QK⊤ √ d  V where WQ, WK, WV ∈Rd×d are weight matrices, and H ∈RT ×d is the output. Next, the feed-forward layer applies the following transformation: Z = ReLU(HW1 + 1 · b1)W2 + 1 · b2 where W1, W2 ∈Rd×d and b1, b2 ∈R1×d are weights and biases; 1 ∈RT ×1 is a vector of ones‡; and Z ∈RT ×d is the final output. (Note that we have omitted some details of the Transformer architecture for simplicity.) (a) (4 points) Permuting the input. i. (3 points) Suppose we permute the input sequence X such that the tokens are shuffled ran- domly. This can be represented as multiplication by a permutation matrix P ∈RT ×T , i.e. Xperm = PX. (See Wikipedia for a recap on permutation matrices.) Show that the output Zperm for the permuted input Xperm will be Zperm = PZ. You are given that for any permutation matrix P and any matrix A, the following hold: softmax(PAP⊤) = P softmax(A) P⊤ and ReLU(PA) = P ReLU(A). ii. (1 point) Think about the implications of the result you derived in part i. Explain why this property of the Transformer model could be problematic when processing text. (b) (2 points) Position embeddings are vectors that encode the position of each token in the se- quence. They are added to the input word embeddings before feeding them into the Transformer. One approach is to generate position embedding using a fixed function of the position and the dimension of the embedding. If the input word embeddings are X ∈RT ×d, the position embeddings Φ ∈RT ×d are generated as follows: Φ(t,2i) = sin  t/100002i/d Φ(t,2i+1) = cos  t/100002i/d where t ∈{0, 1, . . . T −1} and i ∈{0, 1, . . . d/2 −1}§. Specifically, the position embeddings are added to the input word embeddings: Xpos = X + Φ i. (1 point) Do you think the position embeddings will help the issue you identified in part (a)? If yes, explain how and if not, explain why not. ii. (1 point) Can the position embeddings for two different tokens in the input sequence be the same? If yes, provide an example. If not, explain why not. ‡Outer product with 1 represents broadcasting operation and makes feed forward network notations mathematically sound. §Here d is assumed even which is typically the case for most models.

第 5 页

CS 224N Winter 2026 Assignment 3 Page 5 of 6 3. Coding a transformer from scratch (30 points) In this question you will fill in code to implement a decoder only, GPT-2 style transformer, and a simple training loop. For part (a), we have included unit tests for each sub problem that can run locally on your laptop. You will be awarded full points for the subproblem if you pass the unit test. Do not edit the unit test file as we will separately be running the tests when you submit your code. The following tips might be useful during this part of the assignment: • Add assert statements to check the shape of tensors matches what you think it should be. • Consider the Jaxtyping package to type hint the shape of tensors. • Consider the einops package for manipulating tensors (einops.rearrange is particularly useful). This will help you not only write less buggy code, but also make your code far more readable. (a) (20 points) In this part of the question we will implement a transformer in the model solution.py file. The file contains a number of different classes that you will implement. In the end, you will have an implementation of the Transformer class with functioning forward and generate methods. In part (b), we will (start to) train your implementation of Transformer. i. (0 points) Familiarize yourself with the classes in the model solution.py file. We will ask you to implement them in the order MLP, CausalAttention, DecoderBlock, and finally Transformer. We will get you to implement the classes in this order because it is the order of dependence. Transformer depends on DecoderBlock, that in turn depends on CausalAttention and MLP. ii. (1 point) Implement MLP.forward. Check you pass the corresponding test. iii. (6 points) Implement CausalAttention.forward. Check you pass the corresponding test. iv. (2 points) Implement DecoderBlock.forward. Check you pass the corresponding test. v. (6 points) Implement Transformer.forward. Check you pass the corresponding test. vi. (5 points) Implement Transformer.generate. Check you pass the corresponding test. Note: you should implement greedy decoding for this function.. (b) (10 points) After finishing part (a), you now have a functioning Transformer model. If you look at Transformer. innit we can see that when you create an instance of the Transformer class, we initialize the model with random weights according to the Transformer. init weights method. In this part of the question, you will implement a training loop, and start training a small model locally on your laptop. i. (0 points) Look in the train.py file and familiarize yourself with the training loop. We will run this code to train the model. ii. (7 points) First, implement Transformer.get loss on batch. This function maps a batch of tokens to a single loss value. We use this function in train.py to get the loss over a batch. Check you pass the corresponding test. iii. (3 points) Run train.py. This will train the model, using your Transformer.get loss on batch on 100 batches of data. At the end of training it will save a graph of the training loss and gra- dient norm over training to losses and grad norms.png, include an image of this below. If everything is correct, you should see a decreasing loss curve. (c) (9 points) (Bonus) In this optional bonus question, your goal is to speed up the learning process. We will keep the number of gradient steps fixed to 100, however you can change anything else about train.py or model.py to speed up training. We will consider a change to have succeeded if the final loss after 100 steps is lower than the baseline curve you reported in part biii).

第 6 页

CS 224N Winter 2026 Assignment 3 Page 6 of 6 We will award 3 points for each different change you make that leads to a speedup. Thus for full points, you will need to make three different changes to the training file, each of which leads to a speedup. These changes should compound. For example, you may begin by changing the learning rate, leading to a lower loss. You then might keep this better learning rate, and combine it with a second change (e.g., changing the optimizer or model architecture), that leads to an even lower loss. Note that changing the learning rate to three different values counts as one idea: we are looking for three different types of ideas. When you are done, submit: • A description of each change that you made. • Up to three new learning curves, one for each change you made; and additionally include the baseline from biii). • The lowest loss you achieved after 100 steps.

材料

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

怎么自查

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

这份作业和书的关系

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