# How to Implement the Attention Mechanism in Neural Networks: A Complete Guide to Scaled-Dot-Product and Multi-Head Attention > Learn to implement the attention mechanism in neural networks. This guide covers scaled dot-product and multi-head attention for powerful AI models. Master attention for better results. - Repository: [Harvard Edge Computing/cs249r_book](https://github.com/harvard-edge/cs249r_book) - Tags: how-to-guide - Published: 2026-02-19 --- **You implement the attention mechanism by computing scaled dot-product scores between queries and keys, applying softmax to obtain attention weights, and computing a weighted sum of values, with multi-head attention parallelizing this process across multiple representation subspaces.** The attention mechanism powers modern transformer architectures like GPT and BERT. In the **harvard-edge/cs249r_book** repository, this mechanism is implemented in the TinyTorch educational framework, providing a transparent, step-by-step view of scaled-dot-product and multi-head attention without hidden framework abstractions. ## Understanding the Attention Mechanism Architecture ### Core Components in TinyTorch The implementation in [`tinytorch/src/12_attention/12_attention.py`](https://github.com/harvard-edge/cs249r_book/blob/main/tinytorch/src/12_attention/12_attention.py) relies on several fundamental building blocks: - **`Tensor`** — A lightweight wrapper around `numpy.ndarray` used for all tensor operations - **`Linear`** — A fully-connected layer (weight + bias) that projects inputs into queries, keys, and values - **`Softmax`** — Normalizes attention scores into a probability distribution - **`_compute_attention_scores`** — Computes raw similarity scores via matrix multiplication `Q·Kᵀ` (line 503) - **`_scale_scores`** — Divides scores by √*dₖ* to prevent softmax saturation (line 610) - **`_apply_mask`** — Adds `-1e9` to masked positions for causal or padding masks (line 630) - **`scaled_dot_product_attention`** — The complete attention pipeline (line 647) - **`MultiHeadAttention`** — Parallel attention heads with split/merge logic (line 696) ### The Mathematical Foundation The attention mechanism operates on three learned projections: **Queries (Q)**, **Keys (K)**, and **Values (V)**. The core computation follows the scaled-dot-product formula: ``` Attention(Q, K, V) = softmax(QKᵀ / √dₖ) · V ``` Where *dₖ* represents the dimension of the key vectors. The scaling factor √*dₖ* prevents the dot products from growing too large in high dimensions, which would push the softmax into regions with extremely small gradients. ## Step-by-Step Implementation in TinyTorch ### Computing Raw Attention Scores The process begins in `_compute_attention_scores` at line 503 of [`12_attention.py`](https://github.com/harvard-edge/cs249r_book/blob/main/12_attention.py). This function performs the matrix multiplication between queries and transposed keys: ```python scores = Q @ K.transpose(-2, -1) # Shape: (batch, seq_len, seq_len) ``` This operation computes the raw similarity between every query position and every key position, resulting in a square attention matrix. ### Scaling and Masking Before applying softmax, the implementation applies two critical transformations. First, `_scale_scores` (line 610) divides by the square root of the key dimension: ```python scaled_scores = scores / np.sqrt(d_k) ``` Next, `_apply_mask` (line 630) handles causal or padding masks by adding a large negative value (`-1e9`) to positions that should be ignored: ```python if mask is not None: scaled_scores = scaled_scores + (mask * -1e9) ``` This ensures that after softmax, masked positions receive approximately zero attention weight. ### Generating Attention Weights and Output The `scaled_dot_product_attention` function at line 647 orchestrates the complete pipeline: 1. Compute scores via `_compute_attention_scores` 2. Scale via `_scale_scores` 3. Apply optional mask via `_apply_mask` 4. Apply **Softmax** to obtain attention weights 5. Compute weighted sum: `weights @ V` The function returns both the final output tensor and the attention weight matrix, enabling inspection of where the model focuses. ### Parallelizing with Multi-Head Attention The `MultiHeadAttention` class at line 696 extends single-head attention by splitting the embedding dimension into *h* parallel heads. The implementation uses two helper functions: - **`_split_heads`** (line 666) reshapes tensors from `(batch, seq, embed)` to `(batch, heads, seq, head_dim)` - **`_merge_heads`** (line 684) reverses this operation, concatenating head outputs back to the original embedding dimension Each head performs independent scaled-dot-product attention, allowing the model to jointly attend to information from different representation subspaces at different positions. ## Practical Code Examples ### Single-Head Scaled-Dot-Product Attention This example demonstrates the core attention computation using the TinyTorch implementation: ```python import numpy as np from tinytorch.core.tensor import Tensor from tinytorch.core.attention import scaled_dot_product_attention # Dummy data: batch=2, seq_len=4, d_model=8 Q = Tensor(np.random.randn(2, 4, 8)) K = Tensor(np.random.randn(2, 4, 8)) V = Tensor(np.random.randn(2, 4, 8)) # Optional causal mask (lower-triangular) mask = Tensor(np.tril(np.ones((2, 4, 4)))) output, weights = scaled_dot_product_attention(Q, K, V, mask=mask) print("Output shape :", output.shape) # (2, 4, 8) print("Weights shape:", weights.shape) # (2, 4, 4) ``` ### Multi-Head Attention Forward Pass To utilize parallel attention heads, instantiate the `MultiHeadAttention` class: ```python from tinytorch.core.attention import MultiHeadAttention embed_dim = 64 # d_model num_heads = 8 mha = MultiHeadAttention(embed_dim, num_heads) # Input: batch=3, seq_len=10, embed_dim=64 x = Tensor(np.random.randn(3, 10, embed_dim)) # Simple forward pass (no mask) out = mha(x) print("Multi-head output shape:", out.shape) # (3, 10, 64) ``` ### Applying Causal Masking for Autoregressive Models Causal masking prevents positions from attending to future tokens, essential for language modeling: ```python mask = Tensor(np.tril(np.ones((3, 10, 10)))) # causal mask for each batch item out_masked = mha(x, mask=mask) # Verify that future positions have zero attention weight _, weights = scaled_dot_product_attention( mha.q_proj(x), mha.k_proj(x), mha.v_proj(x), mask=mask ) print("Masked weight at (i