DEEP LEARNING / 6. ATTENTION MECHANISM

Attention Mechanism

The core idea behind every modern model

EXPLANATION

Attention answers the question: for each position in the sequence, which other positions are most relevant?

Instead of compressing the entire sequence into one fixed vector (the RNN bottleneck), attention lets each token directly look at all other tokens and weight their contributions.

Scaled Dot-Product Attention:
• Q (Query)  → what am I looking for?
• K (Key)    → what do I contain?
• V (Value)  → what do I actually return?

Attention(Q,K,V) = softmax(QKᵀ / √d_k) · V

The √d_k scaling prevents dot products from exploding in high dimensions, which would push softmax into zero-gradient regions.

Multi-Head Attention runs this h times in parallel with different learned projections, letting the model attend to different aspects simultaneously.

DATA FLOW

Input sequence: [x1, x2, x3, x4]

  Each xi projected to Q, K, V via learned weight matrices:
    Q = X·Wq,  K = X·Wk,  V = X·Wv

  Attention scores (how much each token attends to each other):
         x1   x2   x3   x4
  x1  [ 0.6  0.2  0.1  0.1 ]  ← x1 mostly attends to itself
  x2  [ 0.3  0.5  0.1  0.1 ]
  x3  [ 0.1  0.2  0.6  0.1 ]
  x4  [ 0.1  0.1  0.2  0.6 ]

  Output = scores · V   ← weighted sum of values

CODE

PYTHON
1import torch
2import torch.nn as nn
3import torch.nn.functional as F
4import math
5
6# ── Scaled dot-product attention from scratch ─────────────────────
7def scaled_dot_product_attention(Q, K, V, mask=None):
8    d_k   = Q.shape[-1]
9    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)
10
11    if mask is not None:
12        scores = scores.masked_fill(mask == 0, float("-inf"))
13
14    weights = F.softmax(scores, dim=-1)   # attention weights
15    return torch.matmul(weights, V), weights
16
17# ── Multi-Head Attention from scratch ────────────────────────────
18class MultiHeadAttention(nn.Module):
19    def __init__(self, d_model, num_heads):
20        super().__init__()
21        assert d_model % num_heads == 0
22        self.d_model   = d_model
23        self.num_heads = num_heads
24        self.d_k       = d_model // num_heads
25
26        self.W_q = nn.Linear(d_model, d_model)
27        self.W_k = nn.Linear(d_model, d_model)
28        self.W_v = nn.Linear(d_model, d_model)
29        self.W_o = nn.Linear(d_model, d_model)
30
31    def split_heads(self, x):
32        B, T, D = x.shape
33        x = x.view(B, T, self.num_heads, self.d_k)
34        return x.transpose(1, 2)     # (B, heads, T, d_k)
35
36    def forward(self, Q, K, V, mask=None):
37        Q = self.split_heads(self.W_q(Q))
38        K = self.split_heads(self.W_k(K))
39        V = self.split_heads(self.W_v(V))
40
41        out, weights = scaled_dot_product_attention(Q, K, V, mask)
42
43        B, H, T, dk = out.shape
44        out = out.transpose(1, 2).contiguous().view(B, T, self.d_model)
45        return self.W_o(out), weights
46
47# ── Test ──────────────────────────────────────────────────────────
48mha  = MultiHeadAttention(d_model=512, num_heads=8)
49x    = torch.randn(4, 20, 512)    # (batch=4, seq=20, d_model=512)
50out, w = mha(x, x, x)            # self-attention: Q=K=V=x
51print(f"Output:  {out.shape}")   # (4, 20, 512)
52print(f"Weights: {w.shape}")     # (4, 8, 20, 20) — per head attention map
53
54# PyTorch also has a built-in:
55mha_builtin = nn.MultiheadAttention(embed_dim=512, num_heads=8, batch_first=True)
56out2, w2 = mha_builtin(x, x, x)
57print(f"Built-in output: {out2.shape}")
← PREV5. RNNs & LSTMsNEXT →7. Transformers