Implement a Transformer Decoder from Scratch Without Starter Code

With no starter code, no skeleton and no other scaffolding, implement a transformer decoder from scratch in about 30 minutes. Write it in Python with a tensor library such as PyTorch or NumPy, and explain the tensor shapes as you go. The decoder takes a batch of token-ID sequences and returns next-token logits for every position. It should contain: - token embeddings plus positional information; - masked (causal) multi-head self-attention that you write yourself, so that each position attends only to itself and to earlier positions; - a position-wise feed-forward network; - residual connections and layer normalization around the sub-layers; - a stack of identical decoder blocks, followed by a projection to vocabulary logits. ```hint Write the shapes down first Before any code, write the shape of the activations entering and leaving each sub-layer and the shape of the attention-score tensor. Most bugs in this exercise are a wrong transpose or reshape. ``` ```hint Where the mask goes Decide at which step of the attention computation the causal mask is applied, and what value a masked position takes, so that a later token can never influence an earlier position. ``` ```hint Heads without a loop Think about how a single projection can be reshaped so that every head is computed by one batched matrix multiplication, and how the heads are merged back afterwards. ``` ### Constraints and Clarifications - Input: integer token IDs of shape `(batch, seq_len)`, with `seq_len` at most a maximum context length passed to the constructor. - Output: logits of shape `(batch, seq_len, vocab_size)`. The logits at position `t` must depend only on tokens `0` through `t`. - Assume basic building blocks (linear layers, embedding tables, layer normalization, softmax) are allowed, but a ready-made attention or transformer module is not. Confirm this with the interviewer. - Model width, number of heads, number of layers and feed-forward width are constructor arguments. ### Clarifying Questions - Is this a decoder-only language model, or the decoder half of an encoder-decoder transformer that also attends to encoder outputs? - Is a forward pass enough, or should I also write the training loss and a generation loop? - Can the inputs contain padding that must be masked as well? - Is there a preference between pre-norm and post-norm blocks, or between learned and fixed positional encodings? ### What a Strong Answer Covers - Scaled dot-product attention with the correct scaling factor and a causal mask applied at the right point - A correct split into heads and merge back, with the shape stated at each step - A decoder block whose residual connections, normalization and feed-forward sub-layer are wired in a consistent, stated order - Embeddings, positional information, the block stack and the output projection producing logits of the right shape - Working code inside the time box, plus a quick check that it runs and that future tokens cannot leak into earlier positions - Awareness of cost: how attention scales with sequence length and where the parameters are ### Follow-up Questions - Why are the attention scores divided by the square root of the head dimension, and what happens without it? - How would you add a key-value cache for autoregressive generation, and what does it save? - What changes to turn this into the decoder of an encoder-decoder model? - How would you write a unit test that proves the causal mask works?