Research Scientist
NLP Fundamentals
NLP has been transformed by the Transformer architecture, which replaced recurrent models with self-attention to achieve parallel training and long-range dependency modeling. Research Scientists must understand the full stack: from sub-word tokenization (BPE, WordPiece) to word embeddings (Word2Vec, GloVe), the scaled dot-product attention derivation, BERT's masked language modeling pre-training, GPT's autoregressive generation, and modern evaluation metrics (perplexity, BLEU, BERTScore). Pre-Transformer methods (LSTM, GRU, seq2seq attention) still appear in interviews as motivation for the Transformer's design choices.
Core Theory
- 1Word2Vec โ Skip-gram and CBOW: Word2Vec learns dense word embeddings by training a shallow neural network on a large corpus. Skip-gram: given a center word w_c, predict context words w_o within a window. Objective: maximize ฮฃ_{(c,o)} log P(w_o | w_c) where P(w_o|w_c) = exp(v_{w_o}^T v_{w_c}) / ฮฃ_w exp(v_w^T v_{w_c}). The denominator sums over all vocabulary tokens โ expensive. Solution: negative sampling replaces the full softmax with a binary classification objective: for k random negative words w_n, maximize log ฯ(v_{w_o}^T v_{w_c}) + ฮฃ_k log ฯ(-v_{w_n}^T v_{w_c}). CBOW predicts the center word from the average of context word vectors โ faster but Skip-gram produces better representations for rare words. GloVe exploits global co-occurrence statistics: learns vectors where v_i^T v_j + b_i + b_j โ log X_{ij}, weighting by the co-occurrence count X_{ij}.
- 2Scaled dot-product attention derivation: given input sequences as matrices Q โ R^{nรd_k}, K โ R^{mรd_k}, V โ R^{mรd_v}: Attention(Q,K,V) = softmax(QK^T / โd_k) V. The QK^T / โd_k computes pairwise compatibility scores between queries and keys. The 1/โd_k scaling is critical: without it, the dot products grow with dimension (Var(q^T k) = d_k for unit-variance q, k), pushing the softmax into saturated regions with tiny gradients. The softmax converts scores to a probability distribution (attention weights) that sum to 1, used to take a weighted combination of values. The output is a sequence of length n where each position aggregates information from all m key-value positions.
- 3Multi-head attention: instead of one attention function, compute h parallel attention heads with different learned projections. For head i: head_i = Attention(Q W_i^Q, K W_i^K, V W_i^V) where W_i^Q โ R^{d_model ร d_k}, W_i^K โ R^{d_model ร d_k}, W_i^V โ R^{d_model ร d_v} with d_k = d_v = d_model/h. Concatenate and project: MultiHead(Q,K,V) = Concat(head_1, ..., head_h) W^O. Each head can attend to different aspects of the sequence (syntax vs. semantics). The computational cost is the same as single-head attention with d_model dimensions because each head operates on d_model/h dimensions.
- 4Transformer architecture: the encoder stack consists of L identical blocks each with: (1) Multi-head self-attention (all tokens attend to all other tokens), (2) position-wise FFN: FFN(x) = max(0, x W_1 + b_1) W_2 + b_2, where d_ff โ 4 ร d_model. Pre-norm (used in modern variants): LayerNorm โ sublayer โ add residual. Positional encoding injects position information since attention is permutation-equivariant: PE(pos, 2i) = sin(pos/10000^{2i/d_model}), PE(pos, 2i+1) = cos(pos/10000^{2i/d_model}). Sinusoidal encodings have the property that PE(pos+k) can be expressed as a linear function of PE(pos) for any offset k, allowing the model to attend by relative position. The decoder adds a causal self-attention (masked to prevent attending to future positions) and cross-attention (queries from decoder, keys/values from encoder).
- 5BERT pre-training โ Masked Language Modeling: BERT uses a bidirectional Transformer encoder (all tokens attend to all other tokens, no causal masking). Pre-training task 1 โ MLM: randomly mask 15% of tokens; of those, replace 80% with [MASK], 10% with a random word, 10% with the original word. Predict the original token at masked positions. The 10%/10% variation prevents the model from learning that [MASK] always signals a prediction target (a distributional mismatch between pre-training and fine-tuning). Task 2 โ NSP (Next Sentence Prediction): given [CLS] sent_A [SEP] sent_B [SEP], predict if sent_B follows sent_A. Later work (RoBERTa) showed NSP hurts performance โ removing it and using longer sequences and more data improved results. WordPiece tokenization: starts with character vocabulary, greedily merges pairs that maximize the corpus likelihood under a language model. Results in subword units that handle rare/unknown words without OOV issues.
- 6GPT pre-training โ autoregressive language modeling: GPT uses a decoder-only Transformer with causal (left-to-right) self-attention masking. Pre-training objective: maximize P(w_t | w_1, ..., w_{t-1}) = softmax(h_t W^T) where h_t is the transformer output at position t. At each position, attention can only see tokens to the left (implemented by masking future positions with -โ before softmax). This is efficient: training computes predictions for all positions simultaneously (teacher forcing), but generation is sequential. GPT-3 (175B params) showed in-context learning emerges at scale โ the model can do few-shot tasks without gradient updates by conditioning on demonstrations in the prompt. GPT-4 and beyond use RLHF for alignment (see rs-llm-rlhf topic).
- 7Tokenization โ BPE, WordPiece, SentencePiece: Byte-Pair Encoding (BPE) starts with characters, iteratively merges the most frequent adjacent pair. Used in GPT-2, GPT-4, LLaMA. Vocabulary size is a hyperparameter (50K for GPT-2). WordPiece: like BPE but merges the pair that maximizes log P(corpus)/log P(individual pieces) โ a likelihood-based criterion. Used in BERT. SentencePiece: tokenizes at the byte level (works on any language without pre-tokenization), implementing BPE or Unigram Language Model. Unigram LM starts with a large vocabulary and iteratively removes tokens that least reduce corpus likelihood. Critical property: subword tokenization handles rare words by decomposing them into common subwords, and handles arbitrary new words without OOV.
- 8Sequence-to-sequence with attention (pre-Transformer): encoder LSTM processes input sequence, outputting hidden states h_1,...,h_T. Decoder LSTM generates output conditioned on context vector. Bahdanau attention: at decoding step t, compute alignment score e_{t,i} = v^T tanh(W_a [s_{t-1}; h_i]), attention weights ฮฑ_{t,i} = softmax(e_{t,i}), context c_t = ฮฃ_i ฮฑ_{t,i} h_i. The decoder predicts each token as: P(y_t) = softmax(W [s_t; c_t]). This mechanism allows the decoder to selectively attend to relevant encoder positions, solving the information bottleneck of fixed-length context vectors in vanilla seq2seq. The Transformer can be seen as generalizing this to multi-head, parallel computation across all positions simultaneously.
- 9Beam search and decoding: greedy decoding selects argmax at each step โ fast but misses globally better sequences. Beam search maintains B hypotheses (beams) at each step, expanding each by all vocabulary tokens, keeping the top-B by cumulative log-probability. Score: log P(y) = (1/|y|) ฮฃ_t log P(y_t | y_{<t}, x) (length normalization prevents preferring short sequences). Beam size B=4-8 is standard for translation. Nucleus sampling (top-p): sample from the smallest set of tokens with cumulative probability โฅ p (e.g., p=0.9). Adapts to the distribution width at each step โ sharp distributions get few candidates, flat distributions get many. Temperature sampling scales logits by 1/T: low T โ sharper (more deterministic), high T โ flatter (more diverse). Top-k sampling: sample from the top-k tokens regardless of their probability mass โ less principled than nucleus sampling but computationally simpler.
- 10Evaluation metrics โ BLEU, ROUGE, BERTScore, perplexity: BLEU (Bilingual Evaluation Understudy): precision of n-grams in prediction vs. reference, with brevity penalty. BLEU-4 uses 1-4 gram precision geometrically averaged. Limitation: does not capture synonymy, paraphrase, or semantic equivalence. ROUGE (Recall-Oriented Understudy for Gisting Evaluation): recall-based n-gram overlap; ROUGE-L uses longest common subsequence. Used for summarization. BERTScore: compute cosine similarity between BERT token embeddings of prediction and reference, taking max similarity for each reference token (recall) and each predicted token (precision). Captures semantic similarity beyond surface n-gram overlap. Perplexity: PP(W) = exp(-(1/N) ฮฃ_t log P(w_t | w_{<t})). Lower perplexity = model assigns higher probability to the test text. Comparable only within the same tokenization (adding tokens to vocabulary changes PP values).
Key Formulas
| 1 | import torch |
| 2 | import torch.nn as nn |
| 3 | import torch.nn.functional as F |
| 4 | import math |
| 5 | import numpy as np |
| 6 | from collections import Counter |
| 7 | |
| 8 | # โโโ Scaled Dot-Product Attention โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ |
| 9 | |
| 10 | def scaled_dot_product_attention( |
| 11 | Q: torch.Tensor, |
| 12 | K: torch.Tensor, |
| 13 | V: torch.Tensor, |
| 14 | mask: torch.Tensor = None, |
| 15 | dropout_p: float = 0.0, |
| 16 | ): |
| 17 | """ |
| 18 | Q: (batch, heads, seq_q, d_k) |
| 19 | K: (batch, heads, seq_k, d_k) |
| 20 | V: (batch, heads, seq_k, d_v) |
| 21 | mask: (batch, 1, seq_q, seq_k) or (1, 1, seq_q, seq_k) โ True = mask out |
| 22 | |
| 23 | Returns: |
| 24 | output: (batch, heads, seq_q, d_v) |
| 25 | attn_weights: (batch, heads, seq_q, seq_k) |
| 26 | """ |
| 27 | d_k = Q.shape[-1] |
| 28 | # Compute raw attention scores |
| 29 | scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k) |
| 30 | |
| 31 | # Apply mask (e.g. causal mask for GPT, padding mask) |
| 32 | if mask is not None: |
| 33 | scores = scores.masked_fill(mask, float('-inf')) |
| 34 | |
| 35 | attn_weights = F.softmax(scores, dim=-1) |
| 36 | |
| 37 | # Optional dropout on attention weights (training only) |
| 38 | if dropout_p > 0.0 and torch.is_grad_enabled(): |
| 39 | attn_weights = F.dropout(attn_weights, p=dropout_p) |
| 40 | |
| 41 | output = torch.matmul(attn_weights, V) |
| 42 | return output, attn_weights |
| 43 | |
| 44 | |
| 45 | # โโโ Sinusoidal Positional Encoding โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ |
| 46 | |
| 47 | class SinusoidalPositionalEncoding(nn.Module): |
| 48 | """ |
| 49 | Fixed (non-learnable) sinusoidal positional encoding. |
| 50 | PE(pos, 2i) = sin(pos / 10000^(2i/d_model)) |
| 51 | PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model)) |
| 52 | |
| 53 | Key property: PE(pos+k) is a linear function of PE(pos) for any k, |
| 54 | so the model can represent relative positions via linear attention. |
| 55 | """ |
| 56 | def __init__(self, d_model: int, max_len: int = 5000, dropout: float = 0.1): |
| 57 | super().__init__() |
| 58 | self.dropout = nn.Dropout(dropout) |
| 59 | |
| 60 | # Precompute PE matrix: (max_len, d_model) |
| 61 | pe = torch.zeros(max_len, d_model) |
| 62 | position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1) # (L, 1) |
| 63 | # Denominator: 10000^(2i/d_model) โ use log-space for numerical stability |
| 64 | div_term = torch.exp( |
| 65 | torch.arange(0, d_model, 2, dtype=torch.float) * (-math.log(10000.0) / d_model) |
| 66 | ) # (d_model//2,) |
| 67 | |
| 68 | pe[:, 0::2] = torch.sin(position * div_term) # even dims โ sin |
| 69 | pe[:, 1::2] = torch.cos(position * div_term) # odd dims โ cos |
| 70 | |
| 71 | # Register as buffer (moves with model to GPU, not a parameter) |
| 72 | self.register_buffer('pe', pe.unsqueeze(0)) # (1, max_len, d_model) |
| 73 | |
| 74 | def forward(self, x: torch.Tensor): |
| 75 | # x: (batch, seq_len, d_model) |
| 76 | x = x + self.pe[:, :x.shape[1], :] |
| 77 | return self.dropout(x) |
| 78 | |
| 79 | |
| 80 | # โโโ Causal Self-Attention (GPT-style) โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ |
| 81 | |
| 82 | class CausalMultiHeadAttention(nn.Module): |
| 83 | """ |
| 84 | Masked (causal) multi-head self-attention for autoregressive models. |
| 85 | Position i can only attend to positions j โค i (past + current). |
| 86 | """ |
| 87 | def __init__(self, d_model: int, n_heads: int, dropout: float = 0.1, max_len: int = 2048): |
| 88 | super().__init__() |
| 89 | assert d_model % n_heads == 0 |
| 90 | self.n_heads = n_heads |
| 91 | self.d_k = d_model // n_heads |
| 92 | self.d_model = d_model |
| 93 | |
| 94 | self.qkv_proj = nn.Linear(d_model, 3 * d_model) |
| 95 | self.out_proj = nn.Linear(d_model, d_model) |
| 96 | self.attn_dropout = dropout |
| 97 | |
| 98 | # Causal mask: upper triangular (excluding diagonal) = positions to mask |
| 99 | # Shape (1, 1, max_len, max_len) โ True means "mask this position" |
| 100 | causal_mask = torch.triu(torch.ones(max_len, max_len, dtype=torch.bool), diagonal=1) |
| 101 | self.register_buffer('causal_mask', causal_mask.unsqueeze(0).unsqueeze(0)) |
| 102 | |
| 103 | def forward(self, x: torch.Tensor): |
| 104 | B, T, D = x.shape |
| 105 | |
| 106 | # Project to Q, K, V |
| 107 | qkv = self.qkv_proj(x).reshape(B, T, 3, self.n_heads, self.d_k) |
| 108 | qkv = qkv.permute(2, 0, 3, 1, 4) # (3, B, heads, T, d_k) |
| 109 | Q, K, V = qkv[0], qkv[1], qkv[2] |
| 110 | |
| 111 | # Use precomputed causal mask for sequence length T |
| 112 | mask = self.causal_mask[:, :, :T, :T] # (1, 1, T, T) |
| 113 | |
| 114 | out, _ = scaled_dot_product_attention(Q, K, V, mask=mask, dropout_p=self.attn_dropout) |
| 115 | out = out.transpose(1, 2).reshape(B, T, D) |
| 116 | return self.out_proj(out) |
| 117 | |
| 118 | |
| 119 | # โโโ BPE Tokenizer (simplified) โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ |
| 120 | |
| 121 | class SimpleBPE: |
| 122 | """ |
| 123 | Byte-Pair Encoding tokenizer from scratch. |
| 124 | Algorithm: |
| 125 | 1. Initialize vocabulary with characters (+ special end-of-word marker) |
| 126 | 2. Count all adjacent symbol pairs in corpus |
| 127 | 3. Merge the most frequent pair, add merged token to vocabulary |
| 128 | 4. Repeat until target vocab size or no more merges possible |
| 129 | """ |
| 130 | def __init__(self, vocab_size: int = 300): |
| 131 | self.vocab_size = vocab_size |
| 132 | self.merges = {} # (pair) โ merged_token |
| 133 | self.vocab = set() |
| 134 | |
| 135 | @staticmethod |
| 136 | def _get_pairs(word: list) -> set: |
| 137 | return {(word[i], word[i+1]) for i in range(len(word)-1)} |
| 138 | |
| 139 | def train(self, corpus: str): |
| 140 | # Initialize: split into characters, mark end-of-word with '</w>' |
| 141 | words = corpus.lower().split() |
| 142 | word_freq = Counter(words) |
| 143 | |
| 144 | # Represent each word as tuple of characters + end marker |
| 145 | vocab_words = {tuple(list(w) + ['</w>']): c for w, c in word_freq.items()} |
| 146 | |
| 147 | self.vocab = set(c for word in vocab_words for c in word) |
| 148 | num_merges = self.vocab_size - len(self.vocab) |
| 149 | |
| 150 | for _ in range(max(0, num_merges)): |
| 151 | # Count pair frequencies across all words |
| 152 | pair_freqs: Counter = Counter() |
| 153 | for word, freq in vocab_words.items(): |
| 154 | for pair in self._get_pairs(list(word)): |
| 155 | pair_freqs[pair] += freq |
| 156 | |
| 157 | if not pair_freqs: |
| 158 | break |
| 159 | |
| 160 | # Find and merge the most frequent pair |
| 161 | best_pair = pair_freqs.most_common(1)[0][0] |
| 162 | merged = ''.join(best_pair) |
| 163 | self.merges[best_pair] = merged |
| 164 | self.vocab.add(merged) |
| 165 | |
| 166 | # Apply merge to all words |
| 167 | new_vocab_words = {} |
| 168 | for word, freq in vocab_words.items(): |
| 169 | new_word = [] |
| 170 | i = 0 |
| 171 | while i < len(word): |
| 172 | if i < len(word) - 1 and (word[i], word[i+1]) == best_pair: |
| 173 | new_word.append(merged) |
| 174 | i += 2 |
| 175 | else: |
| 176 | new_word.append(word[i]) |
| 177 | i += 1 |
| 178 | new_vocab_words[tuple(new_word)] = freq |
| 179 | vocab_words = new_vocab_words |
| 180 | |
| 181 | print(f"Trained BPE: {len(self.vocab)} tokens, {len(self.merges)} merges") |
| 182 | |
| 183 | def tokenize(self, word: str) -> list: |
| 184 | """Apply learned merges to tokenize a new word.""" |
| 185 | tokens = list(word) + ['</w>'] |
| 186 | for pair, merged in self.merges.items(): |
| 187 | new_tokens = [] |
| 188 | i = 0 |
| 189 | while i < len(tokens): |
| 190 | if i < len(tokens) - 1 and (tokens[i], tokens[i+1]) == pair: |
| 191 | new_tokens.append(merged) |
| 192 | i += 2 |
| 193 | else: |
| 194 | new_tokens.append(tokens[i]) |
| 195 | i += 1 |
| 196 | tokens = new_tokens |
| 197 | return tokens |
| 198 | |
| 199 | |
| 200 | # โโโ Demo โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ |
| 201 | |
| 202 | if __name__ == "__main__": |
| 203 | torch.manual_seed(42) |
| 204 | |
| 205 | # Attention test with causal mask |
| 206 | B, T, D, H = 2, 10, 64, 4 |
| 207 | attn = CausalMultiHeadAttention(d_model=D, n_heads=H) |
| 208 | x = torch.randn(B, T, D) |
| 209 | out = attn(x) |
| 210 | print(f"Causal attention output: {out.shape}") # (2, 10, 64) |
| 211 | |
| 212 | # Positional encoding |
| 213 | pe = SinusoidalPositionalEncoding(d_model=D) |
| 214 | x_pe = pe(x) |
| 215 | print(f"PE output: {x_pe.shape}") # (2, 10, 64) |
| 216 | |
| 217 | # Verify causal masking: future positions should have no influence |
| 218 | # If position 0 output is identical regardless of positions 1..T-1, mask works |
| 219 | x1 = x.clone(); x1[0, 5:, :] = 0 |
| 220 | out1 = attn(x) |
| 221 | out2 = attn(x1) |
| 222 | print(f"Causal mask working: {torch.allclose(out[0, :5], out2[0, :5], atol=1e-5)}") # True |
| 223 | |
| 224 | # BPE tokenizer demo |
| 225 | corpus = "the cat sat on the mat the cat ate the rat" |
| 226 | bpe = SimpleBPE(vocab_size=50) |
| 227 | bpe.train(corpus) |
| 228 | print(f"BPE tokenization of 'cats': {bpe.tokenize('cats')}") |
| 229 | print(f"BPE tokenization of 'matting': {bpe.tokenize('matting')}") |
Three components: (1) Scaled dot-product attention with optional causal masking โ note masked_fill(-inf) before softmax (not 0, which would still contribute after softmax); (2) Sinusoidal positional encoding uses log-space computation for numerical stability, and the register_buffer ensures the PE matrix moves to GPU without being a trainable parameter; (3) BPE tokenizer implements the full training loop โ character initialization, pair frequency counting, iterative merging. The causal mask verification shows that positions 5-9 have zero effect on positions 0-4, confirming the mask works correctly.
Worked Interview Problems
3 problemsProblem
Assume Q and K have rows sampled i.i.d. from N(0,1). Compute Var(q^T k) for a single query-key pair. Explain what happens to the softmax as d_k grows, and why this causes vanishing gradients.
Solution
For q, k โ R^{d_k} with q_i, k_i ~ N(0,1) i.i.d.: q^T k = ฮฃ_{i=1}^{d_k} q_i k_i. Each term q_i k_i has mean 0 and variance Var(q_i k_i) = E[q_iยฒ k_iยฒ] = E[q_iยฒ]E[k_iยฒ] = 1ร1 = 1.
By the properties of sum of independent random variables: E[q^T k] = 0, Var(q^T k) = d_k. So the standard deviation of dot products grows as โd_k.
Effect on softmax: with many queries and keys, attention scores s_i = q^T k_i have variance d_k. For large d_k, one score dominates: softmax({s_i}) โ one-hot vector. The softmax is in a saturated region.
Saturated softmax has near-zero gradient: โsoftmax(s_i)/โs_j โ 0 when one s_i >> others. Backpropagating through a nearly one-hot softmax gives almost no gradient to Q, K, V projections โ learning stalls.
Fix: divide by โd_k โ scores s_i = q^T k_i / โd_k have Var = d_k / d_k = 1. The scores have unit variance regardless of d_k. Softmax operates in a well-scaled region with meaningful gradients.
Intuition: think of d_k=512 dimensional dot products. Without scaling, many near-zero contributions add up to โ512 โ 23 total variation. Dividing by โ512 brings the expected magnitude back to 1, matching the typical logit magnitude the softmax was designed for.
Answer
Var(q^T k) = d_k for unit-variance components โ std = โd_k โ large d_k causes one-hot-like softmax (saturated) โ near-zero gradients. Dividing by โd_k restores Var = 1 and prevents saturation. This is why the Transformer paper called it 'scaled' dot-product attention.
Common Mistakes
Claiming BERT can be used directly for text generation: BERT's bidirectional attention makes it impossible to auto-regressively generate text (you'd need all future tokens to compute each representation). BERT-style models require specialized modifications (like BERT-gen or UniLM) to support generation. For generation, use GPT (decoder-only) or BART/T5 (encoder-decoder).
Confusing the attention mask types: padding mask (masks pad tokens so they don't attend or be attended to โ True for pad positions), causal/look-ahead mask (prevents attending to future tokens โ upper triangular True), and cross-attention mask (different sequence lengths for Q vs K/V). In PyTorch's nn.MultiheadAttention, the 'attn_mask' is additive (add -inf to masked positions) while 'key_padding_mask' is boolean. Mixing these up causes silent bugs.
Saying Word2Vec captures syntax and semantics equally: Word2Vec captures distributional semantics ('king - man + woman โ queen') but is a static embedding โ the same word always has the same vector regardless of context. Contextual embeddings (BERT, GPT) produce different representations for 'bank' in 'river bank' vs. 'bank account'. This distinction is fundamental and frequently probed in NLP interviews.
Not knowing that BLEU has a brevity penalty: BLEU precision alone incentivizes very short outputs (a 1-word output might have high precision). The brevity penalty BP = exp(1 - r/c) penalizes when the candidate length c is shorter than reference length r. Many candidates score well on BLEU by outputting short high-precision fragments โ BLEU's known failure mode for abstractive summarization.
Mixing up beam search and greedy decoding in terms of optimality: beam search does NOT find the globally optimal sequence (that would require exhaustive search over V^T sequences). It finds a locally good sequence by maintaining B partial hypotheses. Empirically, beam size Bโฅ4 finds sequences close to optimal for most NLP tasks, but for open-ended generation, beam search produces repetitive text โ nucleus sampling (top-p) is preferred.