Implementationhard~45 min
Objective
Fill in the three nn.Module classes (MultiHeadAttention, FeedForward, and TransformerBlock) so the block produces the correct output shape and supports deterministic eval mode.
Background
You're building a transformer from scratch in PyTorch for a deep learning course. The architecture follows the Pre-LN (pre-normalization) variant used in GPT-2 and most modern transformers, where LayerNorm is applied before the sub-layer rather than after. The starter code has the class skeletons, __init__ signatures, and a main() that validates shapes and determinism. Your job is to implement the forward passes and initializations.
Requirements
- 1.Implement MultiHeadAttention: project Q/K/V, reshape for heads, compute scaled dot-product attention, concatenate, project output
- 2.Implement FeedForward: two linear layers with GELU activation (d_model -> d_ff -> d_model)
- 3.Initialize TransformerBlock with attention, FFN, two LayerNorm layers, and dropout
- 4.Implement Pre-LN forward: LayerNorm before each sub-layer, residual connections around both
- 5.Ensure output shape matches input shape: (batch_size, seq_len, d_model)
Evaluation (100 points)
Multi-head attention implemented
MultiHeadAttention forward should project Q/K/V and reshape for head splitting
Feed-forward network implemented
FeedForward should have two linear layers with GELU or ReLU activation
LayerNorm in TransformerBlock
TransformerBlock must include LayerNorm layers for Pre-LN architecture
Residual connections
Forward pass must add input to sub-layer output (residual/skip connection)
Dropout in TransformerBlock
Dropout should be applied after sub-layers for regularization