Transformer Architecture

Question 1 (12)

(a) (3) $\text{Attention}(Q,K,V) = \text{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right)V$

• $QK^\top \in \mathbb{R}^{n\times m}$ (1), softmax over rows (1), output shape $\mathbb{R}^{n \times d_v}$ (1).

(b) (5)

• $E[q_i k_i] = E[q_i]E[k_i] = 0$ (independence) ⇒ $E[q\cdot k] = 0$. (1)
• $\text{Var}(q_i k_i) = E[q_i^2]E[k_i^2] - 0 = 1\cdot 1 = 1$. (1)
• Independence of terms ⇒ $\text{Var}(q\cdot k) = \sum_{i=1}^{d_k} 1 = d_k$, so std $=\sqrt{d_k}$. (1)
• Dividing by $\sqrt{d_k}$ rescales variance back to 1. (1)
• Without scaling, large-magnitude logits push softmax into saturated regions where gradients vanish; scaling keeps softmax in a well-conditioned regime. (1)

(c) (4)

import torch, math
def attention(Q, K, V, mask=None):
    d_k = Q.size(-1)
    scores = Q @ K.transpose(-2, -1) / math.sqrt(d_k)   # (1)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))  # (1)
    weights = torch.softmax(scores, dim=-1)              # (1)
    return weights @ V                                   # (1)

Question 2 (10)

(a) (4)

• Convention $d_k = d_v = d_{\text{model}}/h = 512/8 = 64$. (2)
• Each head works in dimension $d_{\text{model}}/h$; summing $h$ heads gives total per-head cost $\propto h \cdot n^2 \cdot (d/h) = n^2 d$, independent of $h$. (2)

(b) (6)

• $W^Q, W^K, W^V$ each project $d_{\text{model}} \to d_{\text{model}}$ (all heads concatenated): each $512 \times 512 = 262{,}144$. (2)
• Three of them: $3 \times 262{,}144 = 786{,}432$. (2)
• Output projection $W^O$: $512 \times 512 = 262{,}144$. (1)
• Total $= 786{,}432 + 262{,}144 = 1{,}048{,}576 = 4 \times 512^2$. (1)

Question 3 (12)

(a) (5) $PE_{(pos,2i)} = \sin\!\left(\frac{pos}{10000^{2i/d}}\right), \quad PE_{(pos,2i+1)} = \cos\!\left(\frac{pos}{10000^{2i/d}}\right)$ (3)

• Explanation: for a fixed offset $k$, $PE_{pos+k}$ is a linear function of $PE_{pos}$ (rotation by a fixed angle per frequency), so the model can learn to attend by relative offset via a linear transform of encodings. Different frequencies give a multi-scale positional signal. (2)

(b) (7)

• Rotation of $(x_1,x_2)$ at position $m$: (2) $R_m\begin{pmatrix}x_1\\x_2\end{pmatrix} = \begin{pmatrix}\cos m\theta & -\sin m\theta\\ \sin m\theta & \cos m\theta\end{pmatrix}\begin{pmatrix}x_1\\x_2\end{pmatrix}$
• Proof of relative dependence: let rotated query $= R_m q$, rotated key $= R_n k$. (2) $(R_m q)^\top (R_n k) = q^\top R_m^\top R_n k = q^\top R_{n-m} k$ (2) since rotation matrices satisfy $R_m^\top = R_{-m}$ and $R_{-m}R_n = R_{n-m}$. The result depends only on $n-m$. (1)

Question 4 (9)

(a) (4)

• Decoder generates tokens left-to-right; at position $t$ it may only use tokens $\le t$. (2)
• Without masking, self-attention at position $t$ would see future tokens (positions $>t$), leaking the answer during teacher-forced training — the model would learn to "cheat" and fail at inference where future tokens don't exist. (2)

(b) (5)

mask = torch.tril(torch.ones(L, L))       # lower-triangular 1s   (2)
scores = scores.masked_fill(mask == 0, float('-inf'))  # (1)
weights = torch.softmax(scores, dim=-1)   # (1)

• Use $-\infty$ because $\text{softmax}(-\infty)=0$ exactly, zeroing out disallowed positions; using $0$ would leave a nonzero weight $\propto e^0=1$, still leaking future info. (1)

Question 5 (9)

(a) (4)

• $QK^\top$: $n\times d$ times $d\times n$ ⇒ $O(n^2 d)$ time. (1)
• Softmax and $\times V$: also $O(n^2 d)$. (1)
• Memory: storing the $n\times n$ score matrix ⇒ $O(n^2)$. (1)
• For long sequences the $n^2$ term dominates (quadratic in sequence length). (1)

(b) (5)

• FlashAttention avoids materializing the full $n\times n$ attention matrix in HBM. (1)
• It tiles Q, K, V into blocks and computes attention block-by-block in fast SRAM. (1)
• Online (streaming) softmax maintains a running max $m$ and running normalizer $\ell$; when a new block arrives it rescales the accumulated output by $e^{m_{\text{old}} - m_{\text{new}}}$ to correct for the updated max, keeping exponentials numerically stable. (2)
• Memory complexity is $O(n)$ (linear) instead of $O(n^2)$. (1)

Question 6 (8)

(a) (5)

• Post-LN: $\;x_{out} = \text{LayerNorm}(x + f(x))$ (1)
• Pre-LN: $\;x_{out} = x + f(\text{LayerNorm}(x))$ (1)
• Explanation: in Pre-LN the residual path is an identity (un-normalized) stream, so gradients flow cleanly to lower layers and activation magnitudes stay bounded; Post-LN normalizes the summed output, causing large gradients at the top early in training that require warmup. Pre-LN's bounded gradients remove that need. (3)

(b) (3)

• Encoder-only: bidirectional self-attention, no causal mask — e.g. BERT. (1)
• Decoder-only: causal-masked self-attention, autoregressive — e.g. GPT. (1)
• Encoder–decoder: encoder (bidirectional) + decoder with masked self-attention and cross-attention to encoder outputs — e.g. original Transformer / T5. (1)

[
  {"claim": "Variance of dot product of two d_k=64 dim unit-variance vectors is d_k=64, std=8=sqrt(64)",
   "code": "dk=64; var=dk; std=sqrt(var); result = (var==64) and (std==8)"},
  {"claim": "Multi-head attention params for d=512,h=8 = 4*512^2 = 1048576",
   "code": "d=512; total=4*d**2; result = (total==1048576)"},
  {"claim": "d_k = d_model/h = 512/8 = 64",
   "code": "result = (512/8 == 64)"},
  {"claim": "RoPE dot product reduces to R_{n-m}: R_m^T R_n = R_{n-m}",
   "code": "th=symbols('theta'); m,n=symbols('m n'); Rm=Matrix([[cos(m*th),-sin(m*th)],[sin(m*th),cos(m*th)]]); Rn=Matrix([[cos(n*th),-sin(n*th)],[sin(n*th),cos(n*th)]]); Rd=Matrix([[cos((n-m)*th),-sin((n-m)*th)],[sin((n-m)*th),cos((n-m)*th)]]); result = simplify(Rm.T*Rn - Rd)==zeros(2,2)"}
]