Transformers and Attention¶
Transformer and attention questions appear in virtually every ML interview today — from applied NLP roles to general ML research positions. Interviewers expect you to understand not just what Transformers do but why they work and what tradeoffs they introduce.
Q1: What is attention and what is the core intuition behind it?¶
Show answer
Attention allows a model to focus on different parts of its input selectively when producing each output. Rather than compressing all input information into a single fixed-size vector, attention lets the model look at all input positions and decide which ones are relevant for the current output.
The intuition: Imagine translating "The bank by the river is steep." The word "bank" is ambiguous. When the decoder produces the translation of "bank," attention lets it look back at all input words and assign high weight to "river" and "steep" — resolving the ambiguity using context.
Query, Key, Value — the core abstraction: - Query (Q): "What am I looking for?" — the current position asking for information - Key (K): "What do I contain?" — each input position advertising its content - Value (V): "What do I actually output?" — what each input position contributes if selected
Attention computes a weighted sum of Values, where the weights are determined by how well each Query matches each Key.
The output at each position is a blend of all Values, where positions with high Query-Key similarity contribute more.
Q2: What is scaled dot-product attention and why is it scaled?¶
Show answer
Scaled dot-product attention computes attention weights between queries and keys, then uses them to weight the values:
Where: - Q ∈ ℝ^(n × dₖ), K ∈ ℝ^(m × dₖ), V ∈ ℝ^(m × dᵥ) - n = number of query positions, m = number of key/value positions - dₖ = dimension of keys and queries
Why divide by √dₖ? The dot product of two random vectors of dimension dₖ has variance ≈ dₖ. For large dₖ (e.g., 512), these dot products become large in magnitude. When fed into softmax, large values cause the softmax output to be extremely peaked — one key gets near-100% attention weight, and gradients through the other weights become almost zero. Dividing by √dₖ rescales the dot products to have variance ≈ 1, keeping the softmax output in a well-conditioned range.
import tensorflow as tf
import numpy as np
def scaled_dot_product_attention(Q, K, V, mask=None):
d_k = tf.cast(tf.shape(K)[-1], tf.float32)
scores = tf.matmul(Q, K, transpose_b=True) / tf.math.sqrt(d_k)
if mask is not None:
scores += (mask * -1e9) # mask out future positions (decoder)
weights = tf.nn.softmax(scores, axis=-1)
output = tf.matmul(weights, V)
return output, weights
# Example
Q = tf.random.normal((1, 5, 64)) # 5 query positions, 64-dim
K = tf.random.normal((1, 7, 64)) # 7 key positions
V = tf.random.normal((1, 7, 64)) # 7 value positions
output, weights = scaled_dot_product_attention(Q, K, V)
print(output.shape) # (1, 5, 64) — one output per query position
print(weights.shape) # (1, 5, 7) — each query attends over 7 keys
Q3: What is multi-head attention and what do the multiple heads learn?¶
Show answer
Multi-head attention runs scaled dot-product attention multiple times in parallel, each with different learned linear projections of Q, K, V. The outputs are concatenated and projected:
Why multiple heads? A single attention head can only compute one type of relationship at each position. Multiple heads allow the model to jointly attend to different aspects of the input simultaneously:
- One head might learn syntactic relationships (subject-verb agreement)
- Another might learn coreference (pronoun-antecedent linking)
- Another might learn positional proximity
- Another might learn semantic similarity
By projecting Q, K, V into lower-dimensional subspaces and running attention in parallel, the model develops a richer set of representations than any single attention pass could provide.
import tensorflow as tf
class MultiHeadAttention(tf.keras.layers.Layer):
def __init__(self, d_model, num_heads):
super().__init__()
self.num_heads = num_heads
self.d_k = d_model // num_heads
self.W_q = tf.keras.layers.Dense(d_model)
self.W_k = tf.keras.layers.Dense(d_model)
self.W_v = tf.keras.layers.Dense(d_model)
self.W_o = tf.keras.layers.Dense(d_model)
def split_heads(self, x, batch_size):
x = tf.reshape(x, (batch_size, -1, self.num_heads, self.d_k))
return tf.transpose(x, perm=[0, 2, 1, 3])
def call(self, Q, K, V, mask=None):
batch_size = tf.shape(Q)[0]
Q = self.split_heads(self.W_q(Q), batch_size)
K = self.split_heads(self.W_k(K), batch_size)
V = self.split_heads(self.W_v(V), batch_size)
d_k = tf.cast(self.d_k, tf.float32)
scores = tf.matmul(Q, K, transpose_b=True) / tf.math.sqrt(d_k)
if mask is not None:
scores += mask * -1e9
weights = tf.nn.softmax(scores, axis=-1)
context = tf.matmul(weights, V)
context = tf.transpose(context, perm=[0, 2, 1, 3])
context = tf.reshape(context, (batch_size, -1, self.num_heads * self.d_k))
return self.W_o(context)
In practice, 8–16 heads with dimension d_model / num_heads per head is standard. More heads is not always better — there is diminishing return and some heads may learn redundant patterns.
Q4: Why does a Transformer need positional encoding?¶
Show answer
Self-attention is permutation-equivariant: if you shuffle the input tokens, the attention output is shuffled in the same way. There is no built-in notion of "which position comes first." A Transformer without positional encoding treats "Dog bites man" and "Man bites dog" identically.
Positional encoding adds position-specific information to each token embedding before it enters the Transformer, allowing the model to distinguish positions.
Sinusoidal positional encoding (original Transformer):
import numpy as np
import tensorflow as tf
def get_positional_encoding(max_len, d_model):
positions = np.arange(max_len)[:, np.newaxis] # (max_len, 1)
dims = np.arange(d_model)[np.newaxis, :] # (1, d_model)
angles = positions / np.power(10000, (2 * (dims // 2)) / d_model)
angles[:, 0::2] = np.sin(angles[:, 0::2]) # even dims: sine
angles[:, 1::2] = np.cos(angles[:, 1::2]) # odd dims: cosine
return tf.cast(angles[np.newaxis, :, :], dtype=tf.float32)
pe = get_positional_encoding(max_len=512, d_model=128)
print(pe.shape) # (1, 512, 128)
Why sinusoidal? The encoding for each position is a unique combination of frequencies. Relative positions can be expressed as linear functions of the encoding — the model can potentially generalise to sequences longer than those seen during training.
Learned positional embeddings (BERT, GPT): position indices are treated like token IDs and learned during training. Simpler but cannot extrapolate to sequences longer than the maximum training length.
Q5: What is a full Transformer block?¶
Show answer
A Transformer encoder block applies these operations in sequence:
LayerNorm → Multi-Head Self-Attention → residual addition
LayerNorm → Feed-Forward Network → residual addition
The residual connections (add & norm) carry information from before the sublayer to after it, stabilising training and providing gradient highways through deep stacks.
import tensorflow as tf
class TransformerBlock(tf.keras.layers.Layer):
def __init__(self, d_model, num_heads, dff, dropout_rate=0.1):
super().__init__()
self.mha = MultiHeadAttention(d_model, num_heads)
self.ffn = tf.keras.Sequential([
tf.keras.layers.Dense(dff, activation='relu'),
tf.keras.layers.Dense(d_model)
])
self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
self.dropout1 = tf.keras.layers.Dropout(dropout_rate)
self.dropout2 = tf.keras.layers.Dropout(dropout_rate)
def call(self, x, training=False, mask=None):
# Self-attention sublayer
attn_output = self.mha(x, x, x, mask)
attn_output = self.dropout1(attn_output, training=training)
x = self.layernorm1(x + attn_output) # residual + layer norm
# Feed-forward sublayer
ffn_output = self.ffn(x)
ffn_output = self.dropout2(ffn_output, training=training)
x = self.layernorm2(x + ffn_output) # residual + layer norm
return x
The FFN: a two-layer MLP applied independently at each position. Typically dff = 4 × d_model. The FFN is where most of the model's capacity lives — it stores factual associations, patterns, and world knowledge in its weights.
Layer Norm vs Batch Norm: Batch Norm normalises across the batch dimension (problematic for variable-length sequences and small batches). Layer Norm normalises across the feature dimension of each individual example, making it sequence-length and batch-size agnostic.
Q6: What is the difference between encoder-only, decoder-only, and encoder-decoder architectures?¶
Show answer
The original Transformer used an encoder-decoder architecture (for translation). Modern models have specialised into three families:
Encoder-only (e.g., BERT): - All tokens attend to all other tokens — full bidirectional self-attention - Pre-trained with Masked Language Modelling (MLM): predict masked tokens - Best for tasks that require understanding the entire input: classification, NER, question answering, sentence embeddings
Decoder-only (e.g., GPT, LLaMA, Claude): - Causal (autoregressive) attention: each token can only attend to previous tokens - Pre-trained with next-token prediction (causal language modelling) - Best for generation: text completion, code generation, dialogue
Encoder-decoder (e.g., T5, BART, original Transformer): - Encoder reads the input with full bidirectional attention - Decoder generates the output autoregressively, attending to both the decoder's own generated tokens (causal self-attention) and the encoder's output (cross-attention) - Best for tasks with distinct input and output sequences: translation, summarisation, question generation
| Architecture | Attention type | Example models | Primary use |
|---|---|---|---|
| Encoder-only | Bidirectional | BERT, RoBERTa | Understanding, classification |
| Decoder-only | Causal (left-to-right) | GPT, LLaMA | Generation |
| Encoder-decoder | Bidirectional + causal + cross | T5, BART | Seq2seq tasks |
Q7: What is the difference between self-attention and cross-attention?¶
Show answer
Self-attention: Q, K, and V all come from the same sequence. Each position attends to all positions in the same sequence (or past positions for causal attention). Used in both encoder and decoder layers.
# Encoder self-attention
output = Attention(Q=X, K=X, V=X)
# Each position in X attends to all positions in X
Cross-attention: Q comes from one sequence, K and V come from another. The decoder uses cross-attention to query the encoder's representations. This is how the decoder "reads" the encoded input.
# Decoder cross-attention
output = Attention(Q=decoder_state, K=encoder_output, V=encoder_output)
# Each decoder position attends to all encoder positions
# Cross-attention in a decoder block
class DecoderBlock(tf.keras.layers.Layer):
def __init__(self, d_model, num_heads, dff):
super().__init__()
self.self_attention = MultiHeadAttention(d_model, num_heads) # causal
self.cross_attention = MultiHeadAttention(d_model, num_heads) # to encoder
self.ffn = tf.keras.Sequential([
tf.keras.layers.Dense(dff, activation='relu'),
tf.keras.layers.Dense(d_model)
])
self.norm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
self.norm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
self.norm3 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
def call(self, x, encoder_output, causal_mask=None):
x = self.norm1(x + self.self_attention(x, x, x, mask=causal_mask))
x = self.norm2(x + self.cross_attention(x, encoder_output, encoder_output))
x = self.norm3(x + self.ffn(x))
return x
Q8: Why do Transformers outperform RNNs on most tasks?¶
Show answer
Two structural advantages drive the performance gap:
1. Parallelism during training: RNNs process sequences step-by-step — each hidden state depends on the previous one. This sequential dependency means training cannot be parallelised across time steps. Training a 512-step sequence requires 512 sequential operations.
Transformer self-attention computes all positions simultaneously — the attention matrix is computed in one operation. Modern GPUs and TPUs are highly parallel, making Transformer training orders of magnitude faster than RNNs at scale.
2. Long-range dependency: In an RNN, information from step 1 must pass through 499 intermediate states to influence step 500. With each state transition, information can be diluted or overwritten.
In a Transformer, the distance from step 1 to step 500 in terms of computation is exactly 1 — a direct attention connection. Any position can attend to any other with equal ease, regardless of distance. This makes Transformers dramatically better at long-range dependencies.
The cost: attention is O(n²) in memory and O(n²) in computation with respect to sequence length n. For very long sequences (thousands of tokens), this becomes a bottleneck that RNNs do not have.
Q9: What is the computational complexity of self-attention and why does it matter?¶
Show answer
The attention score matrix QKᵀ has shape (n × n) where n is the sequence length. Computing it requires O(n² · dₖ) operations and O(n²) memory.
Why this matters: - A sequence of 512 tokens: 512² = 262,144 attention weights — manageable - A sequence of 4,096 tokens: 4096² = 16.8M attention weights per head - A sequence of 100,000 tokens: 10B attention weights — impossible on standard hardware
This quadratic scaling is the primary reason Transformers struggle with very long contexts and has driven significant research into efficient attention variants:
- Sparse attention (Longformer, BigBird): each token attends to a fixed local window + a few global tokens. O(n) complexity.
- Linear attention (Performer, FNet): approximate softmax attention with O(n) complexity via kernel methods.
- Flash Attention: exact attention with O(n) memory via IO-aware tiling — the computation is still O(n²) but memory is linear, enabling much longer sequences on GPUs.
In interviews: knowing that attention is O(n²) and being able to explain why this limits long-context applications demonstrates awareness beyond basic API usage.
# Demonstrating memory scaling
import numpy as np
for seq_len in [128, 512, 2048, 4096, 16384]:
n_heads = 12
d_k = 64
# Attention matrix per head: (seq_len, seq_len) float32 = 4 bytes
mem_mb = (seq_len ** 2 * n_heads * 4) / (1024 ** 2)
print(f"seq_len={seq_len:6d}: {mem_mb:8.1f} MB for attention matrix")
Q10: What happens during pre-training and fine-tuning of a Transformer?¶
Show answer
Pre-training: The model is trained on a massive corpus (hundreds of GB to trillions of tokens) with a self-supervised objective — no human-labelled data required.
- Encoder models (BERT): Masked Language Modelling (MLM) — randomly mask 15% of tokens; predict the masked tokens. Forces the model to understand bidirectional context.
- Decoder models (GPT): Causal Language Modelling (CLM) — predict the next token given all previous tokens. Trains the model to be a general-purpose text generator.
- Encoder-decoder models (T5): can use both MLM-style masking of spans and CLM objectives.
Fine-tuning: The pre-trained model is adapted to a specific downstream task with labelled data. The pre-trained weights are a strong initialisation — the model has already learned language structure, world knowledge, and reasoning patterns. Fine-tuning typically requires far less data and compute than pre-training from scratch.
from transformers import TFBertForSequenceClassification, BertTokenizer
import tensorflow as tf
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = TFBertForSequenceClassification.from_pretrained(
'bert-base-uncased',
num_labels=2 # binary classification
)
texts = ["This movie is great", "This movie is terrible"]
inputs = tokenizer(texts, padding=True, truncation=True, return_tensors='tf')
model.compile(
optimizer=tf.keras.optimizers.Adam(2e-5), # low LR to preserve pre-trained weights
loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
metrics=['accuracy']
)
Why fine-tuning works: BERT's pre-trained representations capture rich linguistic and world knowledge. Fine-tuning nudges the weights toward task-specific patterns without destroying this general knowledge — similar reasoning to transfer learning in CNNs.
Q11: What is the relationship between tokens, embeddings, and vocabulary?¶
Show answer
Tokenisation: converts raw text into a sequence of integer IDs from a fixed vocabulary.
- The vocabulary is a mapping from subword units (tokens) to integer IDs
- Common tokenisers: Byte-Pair Encoding (BPE, used in GPT), WordPiece (BERT), SentencePiece (T5)
- BPE starts from characters and merges frequent pairs into subword units. "playing" → ["play", "##ing"]; "unbelievable" → ["un", "##believ", "##able"]
- Vocabulary size is typically 30,000–50,000 tokens
Embedding: converts token IDs to dense vectors. The embedding layer is a learnable lookup table of shape (vocab_size, d_model).
from transformers import BertTokenizer
import tensorflow as tf
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
text = "The quick brown fox"
tokens = tokenizer.tokenize(text)
ids = tokenizer.convert_tokens_to_ids(tokens)
print(f"Tokens: {tokens}")
# ['the', 'quick', 'brown', 'fox']
print(f"IDs: {ids}")
# [1996, 4248, 2829, 4419]
# Embedding lookup
vocab_size = tokenizer.vocab_size # 30,522 for BERT
d_model = 768
embedding_layer = tf.keras.layers.Embedding(vocab_size, d_model)
embedded = embedding_layer(tf.constant([ids]))
print(embedded.shape) # (1, 4, 768) — 4 tokens, each 768-dim
The three components stacked at the input to BERT: - Token embeddings (what is this token?) - Positional embeddings (where is this token in the sequence?) - Segment embeddings (which sentence does this token belong to?)
These are summed element-wise before entering the Transformer layers.
Q12: What is the difference between BERT and GPT?¶
Show answer
Both are large pre-trained Transformers but they differ in attention masking, pre-training objective, and what they're good at.
BERT (Bidirectional Encoder Representations from Transformers): - Encoder-only architecture - Full bidirectional self-attention — every token attends to every other token - Pre-trained with Masked Language Modelling (MLM) and Next Sentence Prediction (NSP) - Strong at understanding tasks: classification, NER, question answering, sentence similarity
GPT (Generative Pre-trained Transformer): - Decoder-only architecture - Causal (left-to-right) attention — each token attends only to past tokens - Pre-trained with next-token prediction (standard language modelling) - Strong at generation tasks: text completion, summarisation, dialogue, code generation
| BERT | GPT | |
|---|---|---|
| Architecture | Encoder-only | Decoder-only |
| Attention | Bidirectional | Causal (left-to-right) |
| Pre-training | MLM (predict masked tokens) | CLM (predict next token) |
| Strengths | Understanding, classification | Generation, completion |
| Examples | BERT, RoBERTa, ALBERT | GPT-2, GPT-4, LLaMA |
# BERT: masked language modelling
# Input: "The [MASK] sat on the mat."
# Target: "The cat sat on the mat."
# GPT: causal language modelling
# Input: "The cat sat"
# Target: "cat sat on" (predict next token at each position)
Why can't BERT generate text? Because it uses bidirectional attention, it "cheats" — when predicting a token at position i, it can see tokens after position i. This is fine for understanding (the whole input is available) but makes autoregressive generation impossible.
Q13: What is in-context learning and what is prompt engineering?¶
Show answer
In-context learning is the ability of large language models to learn to perform a new task from a few examples given directly in the prompt — without any weight updates (no fine-tuning).
The model never changes its weights. It uses the examples in the context window to infer the pattern and apply it to the new input. This emergent ability appears primarily in models with billions of parameters trained on diverse data.
Few-shot prompt:
"Translate English to French:
English: Hello → French: Bonjour
English: Goodbye → French: Au revoir
English: Thank you → French: "
Model outputs: "Merci"
Why it works (current hypothesis): during pre-training on diverse internet text, the model implicitly learned to perform many forms of "prediction given examples" — the training data contains countless examples of analogies, word pairs, problem-solution pairs, and demonstrations. In-context learning may be activating these implicit learning patterns.
Prompt engineering is the practice of designing input prompts to elicit better outputs from a language model — without modifying the model's weights.
Effective techniques: - Few-shot prompting: provide 2–5 labelled examples before the query - Chain-of-thought (CoT): add "Let's think step by step" to encourage the model to reason before answering — dramatically improves multi-step reasoning - Role prompting: "You are an expert data scientist..." - Structured output: specify the exact format you want ("Respond in JSON with fields: 'label' and 'confidence'")
Important limitation: in-context learning is sensitive to example order, example selection, and prompt wording. Results can vary widely across phrasings of the same question. It is not reliable for production use without careful evaluation.