Research Scientist
LLMs, RLHF & Alignment
LLMs and alignment are the most critical topics for Research Scientist roles at AI labs (OpenAI, Anthropic, Google DeepMind, Meta AI). You must understand the full pre-training โ SFT โ RLHF pipeline, know the Chinchilla scaling laws, and be able to discuss alignment challenges (reward hacking, distributional shift, capability elicitation) with technical depth.
Core Theory
- 1LLM pre-training: autoregressive next-token prediction on massive text corpora. Loss: L = -ฮฃ log P(x_t | x_{<t}; ฮธ). Model learns all of language structure, world knowledge, and reasoning patterns through compression. BPE (Byte Pair Encoding) tokenization: iteratively merge most frequent character pairs. GPT-4 uses ~100k token vocabulary. Context length determines how much prior text the model can attend to โ current frontier: 128k-1M tokens.
- 2Scaling laws (Chinchilla, Hoffmann et al. 2022): loss scales as a power law with both model parameters N and training tokens D: L(N,D) โ L_โ + A/N^ฮฑ + B/D^ฮฒ. For compute-optimal training (fixed compute C โ 6ND FLOPs): N_opt โ C^0.5, D_opt โ C^0.5. Chinchilla result: GPT-3 (175B params) was trained on too few tokens โ a 70B model trained on 1.4T tokens (Chinchilla) performs better. Implication: model size and data size should scale roughly equally.
- 3Instruction tuning (SFT): fine-tune the pre-trained base model on a curated dataset of (instruction, response) pairs. FLAN-T5 and InstructGPT showed that SFT significantly improves instruction-following and reduces harmful outputs, even with small fine-tuning datasets (~100k examples vs 300B+ pre-training tokens). Key insight: the base model already 'knows' how to do tasks from pre-training; SFT teaches it to surface this knowledge in response to instructions.
- 4RLHF pipeline (Ouyang et al. 2022 โ InstructGPT): (1) SFT on demonstration data. (2) Reward model training: human annotators rank multiple model outputs for the same prompt. Train a reward model r(x, y) to predict the preferred output using Bradley-Terry model: P(y_1 > y_2) = ฯ(r(x,y_1) - r(x,y_2)). (3) PPO fine-tuning: optimize the policy ฯ to maximize r(x, y) - ฮฒ ร KL(ฯ || ฯ_SFT). The KL penalty prevents the model from exploiting the reward model (reward hacking).
- 5Reward hacking: the policy finds behaviors that score highly on the reward model but are actually undesirable. Example: a model trained to be 'helpful' learns to output confidently-worded nonsense because human raters prefer confident answers. The reward model is an imperfect proxy for true human preferences, and optimization pressure exploits the gap. Mitigation: KL divergence penalty to prevent too much deviation from the SFT policy, red teaming to discover exploits before deployment, constitutional AI's RLAIF.
- 6Constitutional AI (Anthropic, Bai et al. 2022): instead of human labelers for RL preference data, use the AI itself. RLAIF pipeline: (1) generate harmful response, (2) use a constitution (list of principles) to prompt the model to critique and revise its response, (3) fine-tune on (harmful, revised) pairs, (4) use AI feedback (model compares responses) instead of human feedback for RL stage. Benefits: scales human feedback, more consistent application of principles, enables automated red teaming.
- 7Emergent capabilities: abilities that appear suddenly at some scale threshold and are near-zero at smaller scales. Examples: multi-step arithmetic, chain-of-thought reasoning, in-context learning. Controversy: Wei et al. (2022) argued these are truly emergent. Schaeffer et al. (2023) argued they're measurement artifacts โ with a smoother metric (bits per character instead of accuracy), the phase transition disappears. This is an active debate.
- 8Inference efficiency: KV caching โ store key and value matrices for all previous tokens to avoid recomputation. Memory: O(n_layers ร seq_len ร 2 ร d_model ร 2 bytes). Speculative decoding: a small 'draft' model generates k tokens quickly, then the large 'verifier' model checks all k tokens in a single parallel forward pass. Accepts all tokens where the large model's distribution matches, rejects others. Achieves ~2-3ร throughput with identical output distribution. Flash Attention: O(1) memory (instead of O(nยฒ)) by fusing attention computation into a single CUDA kernel with tiling.
Key Formulas
| 1 | import numpy as np |
| 2 | import torch |
| 3 | import torch.nn as nn |
| 4 | import torch.nn.functional as F |
| 5 | |
| 6 | # โโโ Reward Model Training (Bradley-Terry preference model) โโโโโโโโโโโโโโโโโโโ |
| 7 | class RewardModel(nn.Module): |
| 8 | """ |
| 9 | Reward model R(x, y) โ scalar score. |
| 10 | Trained on human preference pairs (y_chosen > y_rejected). |
| 11 | """ |
| 12 | def __init__(self, d_model: int = 256): |
| 13 | super().__init__() |
| 14 | # In practice: starts from a pre-trained LLM, adds a linear head |
| 15 | self.backbone = nn.Sequential( |
| 16 | nn.Linear(d_model, 128), |
| 17 | nn.ReLU(), |
| 18 | nn.Linear(128, 1) # scalar reward |
| 19 | ) |
| 20 | |
| 21 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 22 | return self.backbone(x).squeeze(-1) |
| 23 | |
| 24 | # Bradley-Terry loss for reward model training |
| 25 | def bradley_terry_loss(r_chosen: torch.Tensor, |
| 26 | r_rejected: torch.Tensor) -> torch.Tensor: |
| 27 | """ |
| 28 | L = -log ฯ(r_chosen - r_rejected) |
| 29 | Train reward model to score chosen responses higher than rejected. |
| 30 | """ |
| 31 | return -F.logsigmoid(r_chosen - r_rejected).mean() |
| 32 | |
| 33 | # Example training batch |
| 34 | batch_size = 16 |
| 35 | d_model = 256 |
| 36 | |
| 37 | # Simulate encoded (prompt, response) pairs |
| 38 | chosen_embeddings = torch.randn(batch_size, d_model) # "good" responses |
| 39 | rejected_embeddings = torch.randn(batch_size, d_model) * 0.5 # "bad" responses |
| 40 | |
| 41 | rm = RewardModel(d_model) |
| 42 | r_chosen = rm(chosen_embeddings) |
| 43 | r_rejected = rm(rejected_embeddings) |
| 44 | |
| 45 | loss = bradley_terry_loss(r_chosen, r_rejected) |
| 46 | print(f"Reward model loss: {loss.item():.4f}") |
| 47 | print(f"Mean chosen score: {r_chosen.mean().item():.4f}") |
| 48 | print(f"Mean rejected score: {r_rejected.mean().item():.4f}") |
| 49 | |
| 50 | # โโโ PPO-style RLHF loss with KL penalty โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ |
| 51 | def rlhf_objective(reward: torch.Tensor, |
| 52 | log_probs_policy: torch.Tensor, |
| 53 | log_probs_ref: torch.Tensor, |
| 54 | beta: float = 0.1) -> torch.Tensor: |
| 55 | """ |
| 56 | RLHF objective: maximize E[r(x,y)] - ฮฒ ร KL(ฯ || ฯ_ref) |
| 57 | |
| 58 | reward: reward model scores for each response |
| 59 | log_probs_policy: log P(y|x) under current policy ฯ |
| 60 | log_probs_ref: log P(y|x) under reference SFT policy ฯ_ref |
| 61 | beta: KL penalty coefficient (prevents reward hacking) |
| 62 | |
| 63 | Returns: loss (to be minimized, so we negate the objective) |
| 64 | """ |
| 65 | # KL divergence: KL(ฯ || ฯ_ref) = ฮฃ ฯ ร log(ฯ / ฯ_ref) = E_ฯ[log ฯ - log ฯ_ref] |
| 66 | kl_penalty = (log_probs_policy - log_probs_ref).mean() |
| 67 | |
| 68 | # Total objective: maximize reward - ฮฒ ร KL |
| 69 | objective = reward.mean() - beta * kl_penalty |
| 70 | |
| 71 | return -objective # negate because we minimize loss |
| 72 | |
| 73 | # Simulate a training step |
| 74 | batch_size = 8 |
| 75 | reward = torch.randn(batch_size) + 1.0 # reward model scores |
| 76 | log_p_policy = torch.randn(batch_size) # current policy log probs |
| 77 | log_p_ref = torch.randn(batch_size) - 0.2 # reference SFT policy log probs |
| 78 | |
| 79 | loss = rlhf_objective(reward, log_p_policy, log_p_ref, beta=0.1) |
| 80 | print(f"\nRLHF loss: {loss.item():.4f}") |
| 81 | |
| 82 | kl = (log_p_policy - log_p_ref).mean() |
| 83 | print(f"KL divergence: {kl.item():.4f} (should stay small with beta=0.1)") |
| 84 | print(f"Mean reward: {reward.mean().item():.4f}") |
| 85 | print(f"Net objective: {reward.mean().item() - 0.1 * kl.item():.4f}") |
| 86 | |
| 87 | # โโโ Scaling law intuition โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ |
| 88 | def chinchilla_compute_optimal(total_compute: float, |
| 89 | flops_per_token: float = 6.0) -> tuple: |
| 90 | """ |
| 91 | For compute-optimal training (Chinchilla): |
| 92 | N_opt โ sqrt(C / (6 ร tokens_per_param)) |
| 93 | D_opt โ sqrt(C / (6 ร params_per_token)) |
| 94 | """ |
| 95 | # Simplified: N_opt ร D_opt = C / 6, N_opt โ D_opt |
| 96 | total_nD = total_compute / flops_per_token |
| 97 | n_opt = (total_nD ** 0.5) |
| 98 | d_opt = total_nD / n_opt |
| 99 | return int(n_opt), int(d_opt) |
| 100 | |
| 101 | for compute_flops in [1e21, 1e23, 1e24]: # 1e23 โ GPT-4 compute |
| 102 | n, d = chinchilla_compute_optimal(compute_flops) |
| 103 | print(f"C={compute_flops:.0e}: Nโ{n:.2e} params, Dโ{d:.2e} tokens") |
Three RLHF components: (1) Bradley-Terry reward model training loss โ trains the RM to score human-preferred responses higher using sigmoid of score difference; (2) PPO RLHF objective with KL penalty โ the KL term prevents the policy from deviating too far from the SFT model (prevents reward hacking); (3) Chinchilla compute-optimal scaling โ shows that for a given compute budget, model size and data size should scale equally.
Worked Interview Problems
3 problemsProblem
Explain the complete RLHF pipeline from a pre-trained base model to a helpful assistant. What happens at each stage and why?
Solution
Stage 1 โ Supervised Fine-Tuning (SFT): take the pre-trained base model (trained to predict the next token on web text). Fine-tune it on ~10-100k high-quality (prompt, response) demonstrations written by humans. The model learns the format and style of helpful responses. Problem: can't demonstrate all desired behaviors โ some desired outputs are hard for humans to write but easy to evaluate.
Stage 2 โ Reward Model (RM) training: generate K responses for many prompts using the SFT model. Have human annotators rank them (preferably from best to worst). Train a Bradley-Terry reward model: P(y_1 > y_2) = ฯ(r(x,y_1) - r(x,y_2)). The RM is initialized from the SFT model with a scalar head added. Trained on ~100k-1M comparison pairs.
Stage 3 โ PPO fine-tuning: use PPO (Proximal Policy Optimization) to optimize the SFT policy to maximize the reward model score. The policy generates responses, the RM scores them, the score is the reward signal. Critical addition: KL divergence penalty ฮฒ ร KL(ฯ || ฯ_SFT) to prevent the policy from exploiting the RM (reward hacking). ฮฒ controls the tradeoff between instruction-following (reward) and coherence (KL).
Why each stage is necessary: pre-training gives world knowledge but not instruction-following behavior. SFT teaches instruction format but can't scale to all behaviors. RM + PPO allows the model to be optimized on human preferences without requiring demonstrations of every behavior โ easier to evaluate than to demonstrate.
Failure modes: reward hacking (model outputs sycophantic or verbose responses that score well on RM but aren't actually helpful), distributional shift (model moves too far from SFT distribution, becomes incoherent), overoptimization (as model is optimized more, RM becomes a less faithful proxy for true preferences).
Answer
Pipeline: pre-train โ SFT (behavior format) โ RM training (learn human preferences from comparisons) โ PPO (maximize RM score - ฮฒ ร KL penalty). Each stage addresses a limitation of the previous: SFT can't scale, RM can be evaluated at scale, PPO optimizes while staying close to coherent SFT behavior.
Common Mistakes
Confusing SFT and RLHF. SFT teaches the model to imitate demonstrations. RLHF teaches the model to maximize a reward signal derived from human preferences โ it can discover behaviors not in the SFT dataset. SFT is faster and cheaper; RLHF produces more aligned models but requires a separate reward model and PPO training.
Not knowing the KL penalty's purpose in RLHF. The KL term ฮฒ ร KL(ฯ || ฯ_SFT) prevents reward hacking by penalizing deviation from the SFT policy. Without it, PPO will find degenerate solutions that score highly on the imperfect reward model but are incoherent or harmful.
Misquoting Chinchilla: saying 'bigger models are always better' โ Chinchilla shows that for a fixed compute budget, a smaller model trained on more data often outperforms a larger model trained on fewer tokens. Scale data AND model together.
Treating emergent capabilities as settled science. The Wei et al. (2022) 'emergent abilities' claim is contested by Schaeffer et al. (2023) who show the phase transitions disappear with smoother metrics. Presenting this as an active debate (rather than a fact) signals research-level awareness.
Not knowing that Flash Attention achieves O(1) memory (not O(nยฒ)). Standard attention materializes the full nรn attention matrix in GPU HBM โ for long sequences (n=100k), this is O(10^10) elements. FlashAttention computes attention in tiles using SRAM, never writing the full matrix to HBM. This enables long-context modeling.