Calculate Self-Attention

Easy · Solved in PyTorch · Machine learning coding practice

Problem

Implement the scaled dot-product attention mechanism, a fundamental component of transformer architectures.

You may use torch.nn.functional.softmax(scores, dim=-1) for PyTorch or the provided scipy.special.softmax(scores, axis=-1) for NumPy. You do not need to implement softmax from scratch. Inputs Q, K, and V have matching shapes.

Topics: transformers, matrix-operations, attention

The full statement, worked examples, hints and the test suite are available once you sign in. You can then solve Calculate Self-Attention in the browser and run it against the tests.

Related problems

• Build Causal Grouped-Query Attention (GQA) Hard
• Build Causal Multi-Query Attention (MQA) Hard
• MoE: Dispatch and Weighted Combine Medium
• MoE: Router Load-Balancing Loss Medium
• MoE: Top-k Expert Routing Easy
• Build Causal Multi-Head Attention for nanoGPT Hard