# How to Implement Self-Attention Mechanism from Scratch in Python

> Learn to implement the self-attention mechanism from scratch in Python. Understand scaled dot-product attention and build PyTorch models efficiently for transformer architectures.

- Repository: [Rohit Ghumare/ai-engineering-from-scratch](https://github.com/rohitg00/ai-engineering-from-scratch)
- Tags: how-to-guide
- Published: 2026-07-26

---

**Self-attention enables transformer models to relate every token in a sequence to every other token by computing scaled dot-product attention between learned query, key, and value projections, implemented efficiently in PyTorch using a single combined QKV linear layer and causal masking for autoregressive generation.**

Self-attention forms the computational backbone of modern large language models, allowing sequences to capture long-range dependencies without recurrence. The `rohitg00/ai-engineering-from-scratch` repository provides a complete, educational implementation that demonstrates exactly how to build this mechanism using pure PyTorch operations. This guide walks through the actual source code from the capstone projects, covering everything from linear projections to multi-head processing.

## The Seven Steps of Self-Attention

The `MultiHeadSelfAttention` class in [`phases/19-capstone-projects/33-multihead-self-attention/code/main.py`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/phases/19-capstone-projects/33-multihead-self-attention/code/main.py) encapsulates the complete attention mechanism through seven distinct operations:

1. **Linear projections** – A single learned layer generates **queries (Q)**, **keys (K)**, and **values (V)** from the input representation `x` with shape `B×T×D`
2. **Head splitting** – The concatenated QKV tensors reshape so each head processes a sub-space of dimension `d_head = D / n_heads`
3. **Scaled dot-product** – Attention scores compute as `Q·Kᵀ` scaled by `√d_head` to maintain stable softmax gradients
4. **Causal mask** – A lower-triangular mask prevents tokens from attending to future positions in autoregressive setups
5. **Softmax and dropout** – Masked scores convert to probabilities via `softmax`, followed by optional regularization
6. **Weighted sum** – The probability distribution combines values through matrix multiplication `weights·V`
7. **Head merging and output projection** – Per-head results concatenate and project back to the original dimension

## Building the MultiHeadSelfAttention Class

### Combined QKV Projection

For computational efficiency, the implementation uses a single linear layer to project the input into query, key, and value representations simultaneously. In [`phases/19-capstone-projects/33-multihead-self-attention/code/main.py`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/phases/19-capstone-projects/33-multihead-self-attention/code/main.py) at line 47, the initialization creates:

```python
self.qkv_proj = nn.Linear(d_model, 3 * d_model, bias=True)

```

This concatenated projection reduces memory overhead and kernel launch costs compared to three separate layers. The forward pass splits the output tensor into three equal parts corresponding to Q, K, and V.

### Splitting and Merging Heads

Multi-head attention divides the model dimension across parallel attention heads. The `_split_heads` method reshapes the tensor using view and transpose operations:

```python
x.view(b, t, n_heads, d_head).transpose(1, 2)

```

This transforms the shape from `(batch, seq, d_model)` to `(batch, n_heads, seq, d_head)`, allowing each head to attend to the full sequence independently. After computing attention, the `_merge_heads` method reverses this operation at line 60:

```python
x.transpose(1, 2).contiguous().view(b, t, h * dh)

```

### Scaled Dot-Product Attention

The core attention computation occurs at line 83, where the implementation calculates compatibility scores between queries and keys:

```python
scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_head)

```

Scaling by the square root of the head dimension prevents the dot products from growing too large in magnitude, which would push the softmax function into regions with extremely small gradients.

### Causal Masking Implementation

For language modeling tasks, the class maintains a causal mask buffer initialized at lines 52-53:

```python
self.causal_mask = torch.tril(torch.ones(max_context_length, max_context_length))

```

During the forward pass, the implementation slices this mask to the current sequence length and applies it using `masked_fill`:

```python
scores = scores.masked_fill(mask_slice == 0, float("-inf"))

```

This ensures positions can only attend to themselves and previous tokens, maintaining the autoregressive property required for next-token prediction.

## Practical Implementation Examples

### Minimal Self-Attention Demo

To use the implementation directly for experimentation:

```python
import torch
from phases.19_capstone_projects.33_multihead_self_attention.code.main import MultiHeadSelfAttention

# Configuration

batch, seq_len, d_model = 2, 5, 16
n_heads = 4

# Random input

x = torch.randn(batch, seq_len, d_model)

# Initialise the layer

attn = MultiHeadSelfAttention(d_model=d_model,
                              n_heads=n_heads,
                              max_context_length=seq_len)

# Forward pass – returns (output, attention_weights)

out, weights = attn(x, return_weights=True)

print("output shape :", out.shape)           # → (2, 5, 16)

print("weights shape:", weights.shape)      # → (2, 4, 5, 5)

```

### Tiny Language Model Integration

The repository includes `TinyAttentionLM`, which demonstrates how self-attention functions within a complete architecture:

```python
from phases.19_capstone_projects.33_multihead_self_attention.code.main import TinyAttentionLM, DemoConfig

cfg = DemoConfig(vocab_size=64, d_model=32, n_heads=4, seq_len=12, batch_size=2)
model = TinyAttentionLM(vocab_size=cfg.vocab_size,
                       d_model=cfg.d_model,
                       n_heads=cfg.n_heads,
                       max_context_length=cfg.seq_len)

# Dummy token ids

ids = torch.randint(0, cfg.vocab_size, (cfg.batch_size, cfg.seq_len))

# Forward pass – get logits and per‑head attention weights

logits, attn_weights = model(ids, return_weights=True)

print("logits shape :", logits.shape)          # → (2, 12, 64)

print("attention weights shape :", attn_weights.shape)  # → (2, 4, 12, 12)

```

### Full Transformer Block Stack

For production-like usage, the attention mechanism integrates with LayerNorm and feed-forward networks in [`phases/19-capstone-projects/34-transformer-block/code/main.py`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/phases/19-capstone-projects/34-transformer-block/code/main.py):

```python
from phases.19_capstone_projects.34_transformer_block.code.main import BlockConfig, BlockStack

cfg = BlockConfig(d_model=128, num_heads=8, context_length=32, pre_ln=True)
stack = BlockStack(cfg, depth=3)   # three transformer blocks

tokens = torch.randint(0, 128, (2, 32))   # batch of token ids

output = stack(tokens)

print("stack output shape:", output.shape)   # → (2, 32, 128)

```

## Summary

- **Single QKV projection** improves efficiency by computing queries, keys, and values in one matrix multiplication using `nn.Linear(d_model, 3 * d_model)`
- **Head splitting** via `view` and `transpose` operations enables parallel attention across multiple representation subspaces without increasing computational complexity
- **Causal masking** uses a pre-allocated lower-triangular buffer to prevent attention to future tokens during training
- **Scaling factor** of `1/√d_head` maintains stable gradient flow through the softmax operation
- **Modular design** allows the `MultiHeadSelfAttention` class to function standalone or within complete transformer blocks containing residual connections and layer normalization

## Frequently Asked Questions

### Why combine QKV into a single linear projection?

Combining the query, key, and value projections into one layer reduces memory overhead and enables optimized kernel fusion during training. The implementation splits the output tensor into three parts, achieving identical mathematical results to separate layers while improving computational efficiency on modern hardware accelerators.

### What is the purpose of scaling by the square root of the head dimension?

Scaling by `√d_head` prevents the dot-product values from becoming too large as the dimension increases. Without this scaling, the softmax function would receive extremely large inputs, producing gradients near zero that slow down or prevent effective learning.

### How does causal masking prevent looking at future tokens?

The causal mask is a lower-triangular matrix where positions `(i, j)` contain `1` if `j ≤ i` and `0` otherwise. Before applying softmax, the implementation sets masked positions to negative infinity using `masked_fill`, causing the softmax output to be zero for future positions. This ensures each token can only aggregate information from itself and previous tokens in the sequence.

### Can this implementation handle bidirectional attention?

Yes, by setting the causal mask parameter to `None` or modifying the mask initialization logic, the same `MultiHeadSelfAttention` class supports bidirectional attention for encoder-only architectures like BERT. The mask buffer is optional, and the forward pass can skip the masking step when processing non-autoregressive tasks.