Transformers and attention

In short

A transformer is a deep-learning architecture that processes a whole sequence at once and lets every token look at every other token through attention, instead of stepping through the sequence one position at a time. Attention computes a weighted average of value vectors, where the weights come from the similarity between a query and a set of keys, so the model learns which positions matter for each prediction. It replaced recurrent models for language, speech and increasingly vision, and underlies BERT, GPT, Llama and every large language model in use today. It is examined in GATE DA (Data Science and AI), AICTE/CBCS deep-learning papers and NLP electives, and it is the single most asked architecture in Indian ML and GenAI interviews at TCS Research, Infosys, Wipro, Fractal, Microsoft IDC, Google India, Sarvam AI and AI4Bharat.

Explained

Why attention replaced recurrence

• A recurrent model compresses all past context into one hidden state, so information from far-back tokens is diluted and gradients shrink over long paths.
• Recurrence is sequential: step *t* needs step *t−1*, so GPUs cannot parallelise across time during training.
• Attention gives a constant path length between any two positions (one hop), so long-range dependencies are learned directly, and all positions are computed in parallel.
• Cost of that trade-off: attention is O(n²·d) in sequence length *n*, against O(n·d²) for recurrence — fine for n ≈ 512–8,192, expensive beyond.
• Transformers have no inbuilt notion of order (attention is permutation-equivariant), so position must be injected explicitly.

Attention mechanism: queries, keys and values

• Each input embedding x is projected three ways: query Q = XW_Q ("what am I looking for"), key K = XW_K ("what do I offer"), value V = XW_V ("what I pass on").
• Score: similarity of a query with every key, usually a dot product. Weights: softmax of the scores, so they are positive and sum to 1. Output: weighted sum of values.
• Scaled dot-product attention: divide scores by √d_k. Without scaling, for large d_k the dot products have large variance, softmax saturates and gradients vanish.
• Types of attention scoring (examined):
  • *Dot-product (Luong)*: qᵗk — cheap, needs equal dimensions.
  • *Scaled dot-product*: qᵗk/√d_k — transformer default.
  • *Additive / concat (Bahdanau, 2014)*: vᵗ tanh(W₁q + W₂k) — the original NMT attention, more parameters.
  • *General*: qᵗWk — a learned bilinear form.
• Self-attention vs cross-attention: in self-attention Q, K, V all come from the same sequence; in cross-attention Q comes from the decoder and K, V from the encoder output.
• Masked (causal) self-attention: future positions get score −∞ before softmax, so weight = 0; this keeps a decoder autoregressive.
• Padding mask: the same trick applied to padded positions in a batch, so padding contributes nothing.

Multi-head attention

• One attention head averages; averaging can blur several different relations (syntax, coreference, position) into one.
• Multi-head attention splits the model dimension d_model into h heads of size d_k = d_v = d_model/h, runs attention in each head independently, concatenates and projects with W_O.
• Parameter count stays roughly the same as a single full-width head, but the model attends to several subspaces at once.
• Typical: BERT-base d_model = 768, h = 12, d_k = 64; GPT-style models keep d_k = 64–128 and grow h.
• Variants that cut memory at inference: *multi-query attention* (MQA — one shared K,V for all heads) and *grouped-query attention* (GQA — K,V shared by groups of heads), used in Llama-2/3-class models.

Positional information

• Sinusoidal (fixed) encoding (original paper): sine and cosine of position at geometrically spaced frequencies, added to the embeddings; no parameters, extrapolates somewhat to longer sequences.
• Learned absolute embeddings (BERT, GPT-2): a trainable vector per position up to a maximum length; cannot exceed that length.
• Relative position (Transformer-XL, T5 bias): the score depends on the offset i−j rather than absolute index.
• RoPE (rotary position embedding): rotates Q and K by an angle proportional to position, so the dot product depends on relative distance — standard in Llama, Mistral, Qwen and most current LLMs; its base frequency is scaled to extend context.
• ALiBi: adds a linear, head-specific penalty −m·|i−j| to scores; extrapolates to longer sequences without retraining.

Full architecture: encoder, decoder and the sub-layers

• Encoder layer (×N, N = 6 in the original, 12 in BERT-base): multi-head self-attention → add & norm → position-wise feed-forward network → add & norm.
• Decoder layer: masked multi-head self-attention → add & norm → cross-attention over encoder output → add & norm → FFN → add & norm.
• Position-wise FFN: two linear layers with a non-linearity, applied identically to every position: FFN(x) = W₂·act(W₁x + b₁) + b₂, with inner dimension d_ff = 4·d_model typically. Activations: ReLU (original), GELU (BERT/GPT), SwiGLU (modern LLMs, where d_ff ≈ 8/3·d_model to keep parameters equal).
• Residual connections around every sub-layer keep gradients flowing in deep stacks. Layer normalisation normalises across features for each token (not across the batch), so it works with variable-length sequences and batch size 1.
• Post-LN (original: x + Sublayer(LN-free path, then LN)) needs learning-rate warm-up; Pre-LN (LN before the sub-layer) is far more stable and is what almost all modern models use. RMSNorm drops the mean-centring and is cheaper.
• Output head: linear projection to vocabulary size + softmax; often weight-tied to the input embedding matrix.
• Three architectural families:
  • *Encoder-only* (BERT, RoBERTa, MuRIL, IndicBERT) — bidirectional, trained with masked language modelling; for classification, NER, retrieval embeddings.
  • *Decoder-only* (GPT, Llama, Mistral, Claude-class models) — causal, next-token prediction; for generation and instruction following.
  • *Encoder–decoder* (original Transformer, T5, BART, Whisper, IndicTrans2) — for translation, summarisation, speech-to-text.

Training and inference practicalities

• Tokenisation first: BPE, WordPiece or SentencePiece/unigram subwords. Indic scripts fragment badly under English-trained tokenisers — a Hindi or Tamil sentence can cost 2–4× the tokens of its English equivalent, which is why IndicBERT, MuRIL and Sarvam models train Indic-aware vocabularies.
• Loss: cross-entropy on next token (decoder) or on masked positions (encoder); label smoothing ≈ 0.1 in the original paper.
• Teacher forcing during training (ground-truth prefix); autoregressive sampling at inference, hence exposure bias.
• Decoding: greedy, beam search (translation), top-k, top-p/nucleus, temperature.
• KV cache: at generation step t, keys and values of earlier tokens are reused, so each new token costs O(n·d) not O(n²·d); cache memory ≈ 2·layers·heads·d_k·n·bytes and usually dominates inference RAM.
• Efficiency work: FlashAttention (exact attention, tiled in SRAM, no n² matrix in HBM), sparse/local/sliding-window attention (Longformer, BigBird, Mistral), linear-attention and state-space alternatives (Performer, Mamba).
• Scaling: quality follows power laws in parameters, data and compute; Chinchilla-style guidance is roughly 20 training tokens per parameter.
• Advantages: parallel training, long-range dependencies, transferable pretrained weights, one architecture across text, vision (ViT patches as tokens), speech and multimodal.
• Limitations: quadratic attention cost, data and compute hunger, fixed context window, no built-in grounding (hallucination), heavy inference memory, weak length extrapolation, poor interpretability of attention weights.
• Indian regulatory context: personal data used in training or prompts falls under the Digital Personal Data Protection Act, 2023 (DPDP Act) — consent, purpose limitation and data-principal rights; its rules have been notified with phased commencement, so check current timelines. MeitY's India AI Mission and its advisories on deployed AI models, plus RBI's FREE-AI framework for AI in financial services, govern Indian deployment; the IT Act, 2000 and intermediary rules still apply to generated content. The EU AI Act binds Indian firms serving EU users.

Key points

• Attention - output is a softmax-weighted sum of value vectors, weights set by query-key similarity.
• Q, K, V - learned linear projections of the same (self-attention) or different (cross-attention) inputs.
• Scaling by √d_k - keeps score variance ≈ 1 so softmax does not saturate and gradients survive.
• Self-attention complexity - O(n²·d) time and O(n²) memory for the score matrix; constant path length between tokens.
• Multi-head attention - h parallel heads of width d_model/h capture different relations at the same total cost.
• Causal mask - −∞ on future scores makes a decoder autoregressive; padding mask kills padded tokens.
• Positional encoding - mandatory, since attention alone is order-blind; sinusoidal, learned, relative, RoPE or ALiBi.
• Residual + LayerNorm - enable deep stacks; Pre-LN trains stably, Post-LN needs warm-up.
• Position-wise FFN - per-token MLP with d_ff ≈ 4·d_model; holds most of the parameters.
• Encoder-only / decoder-only / encoder–decoder - understanding, generation, sequence-to-sequence respectively.
• MLM vs CLM - BERT predicts masked tokens bidirectionally; GPT predicts the next token causally.
• KV cache - reuses past keys/values so each generated token is linear, not quadratic, in context length.
• GQA / MQA - share key-value heads to shrink the KV cache with little quality loss.
• FlashAttention - IO-aware exact attention; never materialises the n×n matrix in main memory.
• Weight tying - sharing input embedding and output projection saves vocab×d_model parameters.
• Attention is not explanation - high weights do not reliably indicate causal importance.
• Subword tokenisation - BPE/WordPiece/SentencePiece; Indic scripts need Indic-aware vocabularies to avoid token inflation.

Formulas and examples

Example 1 — one attention head by hand. Three tokens, d_k = 2. q₃ = [1, 0]; keys k₁ = [1, 0], k₂ = [0, 1], k₃ = [1, 1]; values v₁ = [1, 0], v₂ = [0, 1], v₃ = [1, 1].

• Raw scores: q₃·k₁ = 1, q₃·k₂ = 0, q₃·k₃ = 1.
• Scale by √2 = 1.4142: 0.7071, 0, 0.7071.
• Exponentials: e^0.7071 = 2.0281, e⁰ = 1, e^0.7071 = 2.0281; sum = 5.0562.
• Weights: 0.4011, 0.1978, 0.4011 (sum = 1.000).
• Output = 0.4011·[1,0] + 0.1978·[0,1] + 0.4011·[1,1] = [0.8022, 0.5989].
• With a causal mask the same query at position 3 is unchanged (it may see all three); a query at position 1 would get weight 1 on v₁ only, output [1, 0].

Example 2 — effect of scaling on softmax. Two scores whose dot products are 40 and 20, with d_k = 64.

• Unscaled: difference 20 → weights e²⁰:1 ≈ 4.85×10⁸ : 1, i.e. 0.999999998 and 2×10⁻⁹ — softmax saturated, gradient ≈ 0.
• Scaled: 40/8 = 5 and 20/8 = 2.5, difference 2.5 → e^2.5 = 12.182, weights 12.182/13.182 = 0.9241 and 0.0759 — a usable gradient.

Example 3 — parameter count of BERT-base-style encoder. d_model = 768, h = 12, d_ff = 3,072, L = 12 layers, vocab 30,000, max position 512.

• Attention per layer: 4 × 768 × 768 = 2,359,296.
• FFN per layer: 2 × 768 × 3,072 = 4,718,592.
• Per layer ≈ 7,077,888 ≈ 7.08 M (plus small bias and LayerNorm terms); 12 layers ≈ 84.9 M.
• Embeddings: token 30,000 × 768 = 23.04 M; position 512 × 768 = 0.39 M; segment 2 × 768 ≈ 0.002 M → ≈ 23.4 M.
• Total ≈ 108 M ≈ the quoted 110 M parameters; note that two-thirds of the FFN+attention split sits in the FFN.

Example 4 — compute and memory at n = 1,024. Same model, one layer, batch 1.

• Score matrix: 1,024 × 1,024 × 12 heads = 12.58 M floats; in fp16 ≈ 25.2 MB per layer, and it is needed for the backward pass — this is what FlashAttention avoids.
• Attention-score FLOPs: 2 × n² × d_model = 2 × 1,048,576 × 768 ≈ 1.61 GFLOP; weighted sum the same ≈ 1.61 GFLOP.
• Projection FLOPs: 8 × n × d_model² = 8 × 1,024 × 589,824 ≈ 4.83 GFLOP — at n = 1,024 projections still dominate; total attention cost (4n²d_model) overtakes projections (8n·d_model²) when n > 2·d_model ≈ 1,536.

Example 5 — KV cache for generation. A 7-billion-parameter decoder: L = 32, d_model = 4,096, h = 32, d_k = 128, context n = 4,096, fp16 (2 bytes).

• Full MHA: 2 × 32 × 4,096 × 4,096 × 2 bytes = 2.15 GB per sequence.
• With GQA using 8 key-value heads instead of 32: key-value width = 8 × 128 = 1,024, so cache = 2 × 32 × 4,096 × 1,024 × 2 = 0.54 GB — a 4× saving, the reason GQA is standard.
• Batch of 16 users: 8.6 GB with GQA, 34 GB without — the latter does not fit a single 24 GB GPU alongside the weights.

import torch, torch.nn.functional as F

def attention(Q, K, V, causal=False):
    d_k = Q.size(-1)
    scores = Q @ K.transpose(-2, -1) / d_k ** 0.5      # (..., n, n)
    if causal:
        n = scores.size(-1)
        mask = torch.triu(torch.ones(n, n, dtype=torch.bool), diagonal=1)
        scores = scores.masked_fill(mask, float('-inf'))
    W = F.softmax(scores, dim=-1)                       # rows sum to 1
    return W @ V, W

Where it is used

• Indian language AI: AI4Bharat/IIT Madras IndicTrans2 and IndicBERT, Bhashini (MeitY) translation and speech pipelines, Sarvam AI and Krutrim LLMs — all transformer-based; roles in NLP research and data engineering.
• IT services and GCCs: TCS, Infosys, Wipro, HCLTech, Accenture and Cognizant build RAG chatbots, document-intelligence and code-assist products on pretrained transformers; interviews probe attention maths and fine-tuning.
• BFSI and fintech: HDFC, ICICI, Paytm and Razorpay use transformers for KYC document extraction, transaction-narration classification, complaint routing and fraud-text analysis, under RBI's AI guidance and the DPDP Act.
• Healthcare, legal and e-governance: radiology-report generation, discharge summarisation, judgment summarisation (SUPACE/SUVAS in the Supreme Court), grievance classification on CPGRAMS-type systems.
• Speech and vision: Whisper-class ASR for Indian call centres, ViT/DETR for quality inspection in manufacturing, OCR for Devanagari and regional scripts.
• Job roles: ML/NLP engineer, GenAI/LLM engineer, applied scientist, MLOps and inference-optimisation engineer (quantisation, KV-cache and vLLM-style serving).

Interview questions

• Q: Why divide the dot product by √d_k? A: With unit-variance inputs the dot product of d_k-dimensional vectors has variance d_k; dividing by √d_k restores variance ≈ 1 so softmax stays in its sensitive region and gradients do not vanish.
• Q: Why does a transformer need positional encoding? A: Self-attention is permutation-equivariant — shuffling the tokens shuffles the outputs identically — so without position information the model cannot distinguish word order.
• Q: What is the advantage of multiple heads over one wide head? A: Each head attends in a different learned subspace, so the model can capture several relations simultaneously instead of averaging them, at essentially the same parameter cost.
• Q: How is masked self-attention implemented? A: Set scores for future positions to −∞ before the softmax, which makes their weights exactly zero; the mask is an upper-triangular boolean matrix.
• Q: Difference between self-attention and cross-attention? A: Self-attention derives Q, K and V from one sequence; cross-attention takes Q from the decoder and K, V from the encoder output, letting the decoder read the source.
• Q: Why LayerNorm and not BatchNorm? A: LayerNorm normalises across features within one token, so it is independent of batch size and variable sequence length, and needs no running statistics at inference.
• Q: What is the computational complexity of self-attention and how is it reduced? A: O(n²·d) time and O(n²) score memory; reduced by sliding-window/sparse attention, low-rank or linear attention, and by IO-aware exact methods like FlashAttention that avoid storing the n² matrix.
• Q: What does the KV cache store and why does it matter? A: The keys and values of all previously generated tokens, so each new token needs only one query pass; it turns generation from quadratic to linear but often dominates GPU memory.
• Q: BERT vs GPT in one line each? A: BERT is an encoder-only bidirectional model trained with masked language modelling for understanding tasks; GPT is a decoder-only causal model trained on next-token prediction for generation.
• Q: Can attention weights be used to explain a prediction? A: Only loosely — attention weights are one factor among residual paths, value magnitudes and later layers, and studies show they are neither necessary nor sufficient as explanations.
• Q: Why is Pre-LN preferred over the original Post-LN? A: Pre-LN keeps an unnormalised identity path through the residual stream, so gradients at initialisation are well-scaled and deep models train without careful learning-rate warm-up.
• Q: What breaks when you feed a sequence longer than the training context? A: Learned absolute position embeddings have no vectors for those indices and attention distributions go out of distribution; relative schemes like ALiBi or rescaled RoPE extrapolate better.

Common mistakes

• Writing softmax(QKᵗ)V without the √d_k scaling, or scaling by d_k instead of its square root.
• Confusing the dimensions: Q and K must share d_k, V may have a different d_v, and the softmax is row-wise over keys, not column-wise.
• Saying multi-head attention multiplies parameters by h — the heads split d_model, so the cost is nearly unchanged.
• Adding the causal mask after the softmax (weights then no longer sum to 1) instead of adding −∞ to the scores before it.
• Claiming attention removes the vanishing-gradient problem entirely; residual connections and normalisation do most of that work.
• Forgetting that an encoder-only model like BERT cannot generate text autoregressively, or using a decoder-only model where bidirectional context is needed.
• Treating the largest attention weight as "the word the model used", and ignoring padding masks so pad tokens pollute the averages.

Related topics

Recurrent networks and LSTMs · Transfer learning and fine-tuning · Training deep networks: optimisers, normalisation and regularisation · Convolutional neural networks · Graph neural networks basics · Generative models: GANs, VAEs and diffusion