Transformers and Pretraining

7.3 Parallelizing computation

1. Why can the transformer computation be parallelized across tokens?
   The block’s computation for each position is independent of the others given the inputs, so all \(N\) input embeddings can be packed into a matrix \(X\) of shape \([N \times d]\) and processed with matrix multiplications.
   • Jurafsky & Martin, Speech and Language Processing, 3rd ed. draft, ch. 7, §7.3, p. 188
2. How is attention for a whole sequence computed in matrix form for one head?
   \(Q = XW^Q\), \(K = XW^K\), \(V = XW^V\); then \(\text{head} = \text{softmax}\big(\text{mask}\big(\frac{QK^\top}{\sqrt{d_k}}\big)\big)V\) and \(A = \text{head}\,W^O\). \(QK^\top\) is an \([N \times N]\) matrix of all query–key comparisons.
   • Jurafsky & Martin, Speech and Language Processing, 3rd ed. draft, ch. 7, §7.3, p. 189
3. Why is masking needed in causal attention, and how is it done?
   \(QK^\top\) compares each query with every key, including future ones, which would let the model see the word it must predict. The upper triangle (\(j > i\)) is set to \(-\infty\), which the softmax turns into 0.

   • Jurafsky & Martin, Speech and Language Processing, 3rd ed. draft, ch. 7, §7.3, p. 189
4. Why does the mask use \(-\infty\) rather than 0?
   The softmax exponentiates scores: \(e^{-\infty} = 0\) removes a position entirely, while a score of 0 would still give it weight \(e^0 = 1\) before normalization.
   • Jurafsky & Martin, Speech and Language Processing, 3rd ed. draft, ch. 7, §7.3, p. 189
5. What is the computational cost of attention with respect to input length, and why?
   Quadratic: each layer computes dot products between every pair of tokens (the \(N \times N\) matrix), which makes very long documents expensive.
   • Jurafsky & Martin, Speech and Language Processing, 3rd ed. draft, ch. 7, §7.3, pp. 189–190
6. How is multi-head attention parallelized over a sequence?
   Each head \(c\) computes \(Q^c, K^c, V^c\) from \(X\) and its own masked softmax attention, giving \([N \times d_v]\); the \(A\) outputs are concatenated into \([N \times Ad_v]\) and projected by \(W^O\) to \([N \times d]\).
   • Jurafsky & Martin, Speech and Language Processing, 3rd ed. draft, ch. 7, §7.3, pp. 190–191
7. Write a whole transformer layer over the input matrix \(X\).
   \(O = X + \text{MultiHeadAttention}(\text{LayerNorm}(X))\) and \(H = O + \text{FFN}(\text{LayerNorm}(O))\), where the FFN and layer norm are applied to each row in parallel. Input and output are both \([N \times d]\), so layers stack.
   • Jurafsky & Martin, Speech and Language Processing, 3rd ed. draft, ch. 7, §7.3, p. 191