Implementationmedium~35 min
Objective
Implement three functions: numerically stable softmax, scaled dot-product attention with optional masking, and multi-head attention with head splitting and output projection.
Background
You're building a from-scratch transformer library for educational purposes. The first milestone is a correct, readable implementation of the attention mechanism using only NumPy. No frameworks allowed; you need to understand every matrix multiply and reshape. The starter code has the function signatures, docstrings, and a main() that validates shapes and basic properties. Your job is to fill in the three TODO functions.
Requirements
- 1.Implement numerically stable softmax using the max-subtraction trick
- 2.Implement scaled dot-product attention: scores = Q @ K^T / sqrt(d_k), with optional mask
- 3.Apply mask by setting masked positions to -1e9 (or -inf) before softmax
- 4.Implement multi-head attention: project, split heads, attend per head, concatenate, project output
- 5.Ensure output shapes match input shapes (seq_len, d_model) for multi-head attention
Evaluation (100 points)
Softmax is numerically stable
Softmax should subtract the max before exponentiating to avoid overflow
Scaled dot-product attention
Attention scores must be scaled by 1/sqrt(d_k) and computed via matrix multiplication
Mask handling
Masked positions should be set to -1e9 or -inf before softmax
Multi-head splitting
Projections must be split into n_heads for parallel attention
Output projection through W_o
Concatenated head outputs must be projected through W_o