Multi-head attention and Transformer block skeleton · CS336 Language Modeling From Scratch

Deliverableartifact

A TransformerBlock class in PyTorch that composes multi-head attention, an MLP, layernorm, and residual connections, with a small script verifying it runs on random inputs.

Goals

Implement a multi-head attention module using your single-head attention as a building block.
Show that your module preserves batch and sequence dimensions and outputs the model dimension.
Construct a minimal Transformer block class wiring attention, MLP, layernorm, and residuals.
Explain your choice of pre-norm or post-norm in the implementation and potential consequences.

Resources

📄stanford-cs336/assignment1-basics 📄bcefghj/learn-cs336

Overview: what you’re building

You’ll build a small but real TransformerBlock in PyTorch:

It will take in a batch of token embeddings.
Run them through multi-head self-attention.
Then through a feedforward MLP.
Wrap each sublayer with residual connections and layer normalization.
End with a tiny script that runs this block on random data to prove it works.

We’ll walk through each core idea you need, then put it all together.

1 / 6