Attention Mechanisms

Hands-on lab with starter code you can copy + run, and guided steps to produce the artifact for this module.

Setup

Open a fresh terminal in a new folder, e.g. mkdir lab-4 && cd lab-4.
Create a Python 3.11 venv: python -m venv .venv && source .venv/bin/activate.
Install the dependencies noted on each starter snippet (the small pip install … line above the code).
Copy each starter into a file with the suggested name, then run it and capture the output.

Starter code

01 · Scaled dot-product attention from scratch

save as: 04_01_lab.py

python

import torch, math

def attention(Q, K, V, mask=None):
    d_k = Q.size(-1)
    scores = (Q @ K.transpose(-2, -1)) / math.sqrt(d_k)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float("-inf"))
    weights = torch.softmax(scores, dim=-1)
    return weights @ V, weights

B, T, D = 1, 4, 8
x = torch.randn(B, T, D)
Wq, Wk, Wv = (torch.nn.Linear(D, D, bias=False) for _ in range(3))
out, w = attention(Wq(x), Wk(x), Wv(x))
print(out.shape, w.shape)  # (1,4,8) (1,4,4)

Run with: python 04_01_lab.py

02 · Multi-head attention

save as: 04_02_lab.py

python

import torch, torch.nn as nn

class MHA(nn.Module):
    def __init__(self, d, h):
        super().__init__()
        self.h, self.dh = h, d // h
        self.qkv = nn.Linear(d, 3*d, bias=False)
        self.o = nn.Linear(d, d, bias=False)
    def forward(self, x):
        B, T, D = x.shape
        q, k, v = self.qkv(x).chunk(3, dim=-1)
        q, k, v = (t.view(B, T, self.h, self.dh).transpose(1,2) for t in (q,k,v))
        a = torch.nn.functional.scaled_dot_product_attention(q, k, v, is_causal=True)
        return self.o(a.transpose(1,2).reshape(B, T, D))

print(MHA(64, 8)(torch.randn(2, 16, 64)).shape)

Run with: python 04_02_lab.py

Guided steps

Work through each step in order. Stop when your output matches the expected artifact — then push your work to your lab folder.

1
Re-implement attention
Do — Write scaled dot-product attention from scratch (no nn.MultiheadAttention). Unit-test against torch.nn.functional.scaled_dot_product_attention on random tensors.
Expected artifact — torch.allclose(my_attn, ref, atol=1e-5) == True for 10 random seeds.
Run free:
2
Visualise attention heads
Do — Run a 2-layer 4-head toy transformer on a 16-token sentence. Save 8 attention-weight heatmaps as PNGs.
Expected artifact — 8 PNG files (layer-head pairs), each a 16×16 heatmap showing diagonal or causal patterns.
Run free:
3
Causal mask check
Do — Confirm that with a causal mask, position i's output does not change when tokens at positions >i are perturbed.
Expected artifact — max delta across positions 0..i-1 < 1e-6 when later tokens are randomised.
Run free:

Authoritative sources

Independent primary sources used to check and corroborate the guidance on this page.

NIST National Initiative for Cybersecurity Education (NICE)
US national framework for cybersecurity work roles and competencies.
Financial Action Task Force (FATF)
Global AML/CFT standards referenced by financial-crime curricula.
FBI Internet Crime Complaint Center — Annual Reports
Primary loss and typology data used in fraud-investigation training.
US Federal Trade Commission — Consumer Sentinel Data Book
Official complaint statistics for teaching fraud trend analysis.