Module 10: Generative AI, LLMs & Prompt EngineeringBy Ashok Kumar Yadav4 min read

Transformer Architecture and the Self-Attention Mechanism

Every modern large language model is built from the Transformer. Its core idea is simple: let every token look at every other token and decide how much each one matters.

What problem Transformers solved

RNNs read tokens one at a time, which is slow to train and forgets distant context. Transformers process the whole sequence in parallel and connect any two positions with one attention step, so training scales on GPUs and long-range relationships are easy to learn.

Self-attention from scratch

Queries ask what to look for, keys advertise what each token contains, and values carry the information. Dividing by the square root of the key size keeps the softmax from saturating.

import torch, math
import torch.nn.functional as F

def attention(q, k, v, mask=None):
    scores = q @ k.transpose(-2, -1) / math.sqrt(q.size(-1))
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float("-inf"))
    weights = F.softmax(scores, dim=-1)
    return weights @ v, weights

x = torch.randn(1, 4, 8)                  # batch=1, 4 tokens, dim 8
Wq, Wk, Wv = (torch.randn(8, 8) for _ in range(3))
out, w = attention(x @ Wq, x @ Wk, x @ Wv)
print(out.shape, w.sum(-1))               # each row of weights sums to 1

Multi-head attention

Several smaller attention heads let the model track different relationships at once, such as syntax in one head and coreference in another. Outputs are concatenated and projected.

import torch.nn as nn

class MultiHead(nn.Module):
    def __init__(self, d=64, heads=4):
        super().__init__()
        assert d % heads == 0
        self.h, self.dk = heads, d // heads
        self.qkv = nn.Linear(d, 3 * d)
        self.out = nn.Linear(d, d)
    def forward(self, x, causal=False):
        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.dk).transpose(1, 2) for t in (q, k, v))
        mask = torch.tril(torch.ones(T, T)) if causal else None
        y, _ = attention(q, k, v, mask)
        return self.out(y.transpose(1, 2).reshape(B, T, D))

print(MultiHead()(torch.randn(2, 10, 64), causal=True).shape)

Positional encoding and the full block

Attention alone ignores order, so a position signal is added to the embeddings. The original paper used fixed sine and cosine waves, and many modern models learn positions or use rotary embeddings. A Transformer block wraps attention and a feed-forward network with residual connections and layer normalisation.

import torch, math
Q=torch.tensor([[1.,0.],[0.,1.]])
K=Q.clone(); V=torch.tensor([[10.,0.],[0.,20.]])
W=torch.softmax(Q@K.T/math.sqrt(2),dim=-1)
print(W.round(decimals=3)); print((W@V).round(decimals=2))

tensor([[0.670, 0.330], [0.330, 0.670]])
...

scores=torch.zeros(4,4)
mask=torch.triu(torch.ones(4,4),diagonal=1).bool()
scores=scores.masked_fill(mask,float('-inf'))
print(scores)

upper triangle becomes -inf

Frequently asked questions

Primary references

Use these official sources when exact syntax, defaults or definitions matter beyond the examples here.

• Attention Is All You Need
• The Annotated Transformer
• PyTorch nn.MultiheadAttention