Compute One Attention Step by Hand and Reason About Masking and Scaling

An encoder produces three states for a source sentence (hidden size 2): h_1 = (2, 0), h_2 = (0, 1), h_3 = (1, 2). At some decoder step the query (decoder state) is q = (1, 1). Use dot-product attention with keys and values both equal to the encoder states.

1. Compute the raw scores, the scaled scores (divide by \sqrt{d_k}), the softmax weights and the context vector. Which source token receives the most attention and why?
2. What would the weights be if you forgot the scaling? Is the scaling important with d_k = 2? What about d_k = 512?
3. Suppose h_3 is actually a padding position. Show how to mask it and recompute the weights and the context. What goes wrong if you set the masked score to 0 instead of -\infty?
4. Explain how the same computation changes in self-attention inside a decoder (a causal mask), and what the shapes are for batch size B, T_q queries, T_k keys and head dimension d_k.