Attention & transformers

In 2017, the paper "Attention Is All You Need" introduced the transformer, an architecture that dropped recurrence entirely. Today transformers power ChatGPT-style LLMs, translation for Indian languages, speech recognition and even image models. This lesson builds the intuition and code for the core mechanism: attention.

The problem attention solves

In the sentence "The shopkeeper gave the customer a discount because she was a regular", the word "she" must link to "customer". An RNN has to carry that information step by step through its hidden state. Attention lets every word look directly at every other word and decide how much each one matters.

Queries, keys and values

Each token's vector is projected into three vectors:

Vector	Analogy
Query (Q)	"What am I looking for?"
Key (K)	"What do I contain?"
Value (V)	"What information do I pass on if chosen?"

Think of searching a library catalogue: your query is matched against every book's key (its index card), and you read the values (contents) of the best matches.

Scaled dot-product attention

Attention(Q, K, V) = softmax( Q K^T / sqrt(d_k) ) V

1. Dot products of queries with keys give similarity scores.
2. Dividing by the square root of the key dimension keeps scores from growing too large, which would push softmax into regions with tiny gradients.
3. Softmax turns scores into weights that sum to 1.
4. The output is a weighted average of the values.

import math
import torch
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 = scores.softmax(dim=-1)
    return weights @ v, weights

q = k = v = torch.randn(2, 5, 16)          # batch 2, 5 tokens, dim 16
out, w = attention(q, k, v)
print(out.shape, w.shape)                   # (2, 5, 16) (2, 5, 5)

# PyTorch 2.x has a fused, faster version:
out2 = F.scaled_dot_product_attention(q, k, v)
print(torch.allclose(out, out2, atol=1e-5)) # True

Multi-head attention

One attention pattern cannot capture everything. Multi-head attention runs several attention operations in parallel on different learned projections — one head might track grammar, another who-refers-to-whom — then concatenates the results.

import torch.nn as nn

mha = nn.MultiheadAttention(embed_dim=64, num_heads=8, batch_first=True)
x = torch.randn(4, 20, 64)
out, weights = mha(x, x, x)   # self-attention: Q, K and V all come from x
print(out.shape)              # torch.Size([4, 20, 64])

The transformer block

A transformer layer repeats this pattern:

1. Multi-head self-attention
2. Add the input back (residual connection) and apply LayerNorm
3. A position-wise feed-forward network (two linear layers with GELU)
4. Residual connection and LayerNorm again

Because attention ignores order, positional encodings (fixed sinusoids or learned vectors) are added to token embeddings so the model knows word positions.

Encoder, decoder and the model families

Family	Attention type	Examples	Good for
Encoder-only	Bidirectional	BERT, IndicBERT, MuRIL	Classification, search, NER
Decoder-only	Causal (masked, can only see the past)	GPT, Llama	Text generation, chat
Encoder-decoder	Both	T5, IndicTrans2	Translation, summarisation
Vision Transformer	Bidirectional over image patches	ViT	Image classification

A tiny transformer classifier

class TinyTransformer(nn.Module):
    def __init__(self, vocab_size, d_model=128, n_heads=4, n_layers=2,
                 max_len=256, n_classes=2):
        super().__init__()
        self.tok = nn.Embedding(vocab_size, d_model, padding_idx=0)
        self.pos = nn.Embedding(max_len, d_model)
        layer = nn.TransformerEncoderLayer(d_model, n_heads, dim_feedforward=4 * d_model,
                                           dropout=0.1, batch_first=True, activation="gelu")
        self.encoder = nn.TransformerEncoder(layer, num_layers=n_layers)
        self.fc = nn.Linear(d_model, n_classes)

    def forward(self, ids):
        positions = torch.arange(ids.size(1), device=ids.device)
        x = self.tok(ids) + self.pos(positions)
        x = self.encoder(x, src_key_padding_mask=(ids == 0))
        mask = (ids != 0).unsqueeze(-1).float()
        pooled = (x * mask).sum(1) / mask.sum(1).clamp(min=1)   # mean over real tokens
        return self.fc(pooled)

model = TinyTransformer(vocab_size=5000)
print(model(torch.randint(1, 5000, (8, 40))).shape)          # torch.Size([8, 2])

In practice: use pretrained models

Training a transformer from scratch needs enormous data. For real projects, load a pretrained model from the Hugging Face transformers library and fine-tune it — for example, MuRIL or IndicBERT for Hindi and other Indian-language review classification.

Costs to be aware of

Attention compares every token with every other, so compute and memory grow with the square of sequence length. Doubling context from 4k to 8k tokens roughly quadruples attention cost, which is why long-context models use tricks like FlashAttention (used under the hood by scaled_dot_product_attention on supported GPUs).

Key takeaways

• Attention lets every position weigh every other position directly.
• Scaling by the square root of the key dimension keeps softmax gradients healthy.
• Transformer blocks combine multi-head attention, feed-forward layers, residuals and LayerNorm.
• Fine-tune pretrained transformers rather than training from scratch.