Training stages

A chat assistant is built in layers, and each layer explains a different behaviour you see:

flowchart LR
  P[Pretraining<br/>guess the next token<br/>over web-scale text] --> B[Base model<br/>knows a lot, rambles]
  B --> S[SFT<br/>study examples of<br/>good answers]
  S --> PT[Preference tuning<br/>RLHF or DPO]
  PT --> A[Assistant model]
  A --> D[Your adaptation<br/>prompts, RAG, LoRA]

This lesson builds a small, real version of every stage: the pretraining and SFT losses, a reward model, DPO, LoRA, the "which adaptation?" decision, and distillation.

Everyday picture Imagine reading every book in a library with a card covering the next word, guessing it, then sliding the card to check. Do that trillions of times and you absorb grammar, facts, arithmetic habits and coding conventions, because every one of them helps you guess the next word.

Tiny worked example A 5-token sequence gives 4 guesses (token 1 predicts token 2, and so on; the last token has nothing after it to check). Suppose the model gave the right next token probability 0.25, 0.25, 0.5 and 0.5. The penalty for each guess is −ln(probability): 1.386, 1.386, 0.693, 0.693. The pretraining loss is their average, 1.040.

flowchart LR
  T1[t1] --> G1{guess t2}
  T2[t2] --> G2{guess t3}
  T3[t3] --> G3{guess t4}
  T4[t4] --> G4{guess t5}
  G1 & G2 & G3 & G4 --> AVG[average of −ln p<br/>= the loss]

import math
# p_θ(t_{i+1} | t_{1..i}) for each of the guesses
p = [0.25, 0.25, 0.5, 0.5]
# 5 tokens give n − 1 = 4 guesses
n = len(p) + 1
[round(-math.log(p_i), 3) for p_i in p]  # → [1.386, 1.386, 0.693, 0.693]
L_pretrain = -sum(math.log(p_i) for p_i in p) / (n - 1)
print(f"{L_pretrain:.3f}")  # → 1.040

In code: per_position_losses computes −ln p for every next-token guess (with a stable log_softmax), and next_token_loss averages them into the pretraining loss.

Why it matters in practice. Pretraining is where knowledge comes from, and it is frozen at a date (the "knowledge cutoff"). A base model continues text instead of answering questions: ask "What is the capital of France?" and it may continue with "What is the capital of Spain?", because question lists are common on the web.

Everyday picture An apprentice studies a binder of worked examples: customer question, then the expert's reply. They read the question carefully, but they are graded only on writing the reply. Nobody marks them on reproducing the customer's typing.

Tiny worked example Same 5 tokens, but the first 3 are the prompt and the last 2 are the reply. Only the 2 guesses whose target is a reply token count: penalties 0.693 and 0.693, average 0.693. The prompt guesses (1.386 each) are ignored.

flowchart LR
  subgraph Prompt["prompt: read, not graded"]
    P1[t1] --> P2[t2] --> P3[t3]
  end
  subgraph Reply["reply: graded"]
    R1[t4] --> R2[t5]
  end
  P3 --> R1
  R1 -.->|loss| L[average −ln p<br/>over reply tokens only]
  R2 -.->|loss| L

import math
p = [0.25, 0.25, 0.5, 0.5]
# 1 only where the target is a reply token
m = [0, 0, 1, 1]
L_SFT = -sum(m_i * math.log(p_i) for m_i, p_i in zip(m, p)) / sum(m)
round(L_SFT, 3)  # → 0.693

In code: response_mask builds the true/false list of graded predictions, and sft_loss averages per_position_losses over only the positions it marks.

Why it matters in practice. SFT needs far less data than pretraining (thousands to hundreds of thousands of examples) and quality beats quantity: a few thousand excellent examples outperform a mountain of mediocre ones. If you fine-tune your own model and forget the mask, it learns to write user prompts too.

Everyday picture It is hard to write down the perfect answer, but easy to taste two dishes and say which is better. Preference tuning collects exactly those judgements: people (or a model following written principles) compare two responses and pick one.

3a. RLHF: train a critic, then train the cook against it

A reward model learns to predict those judgements, giving each response a score. Reinforcement learning then tunes the language model to earn high scores. The link between scores and "which one wins" is the Bradley-Terry model, built on the sigmoid function σ(z) = 1 / (1 + e^−z), which squashes any number into a probability between 0 and 1 (σ(0) = 0.5, σ(2) = 0.881).

Tiny worked example The reward model scores the chosen answer 3.0 and the rejected one 1.0. The gap is 2, so it predicts the chosen answer wins with probability σ(2) = 0.881, and its loss on this pair is −ln 0.881 = 0.127. Had it scored both 1.0, it would predict a coin flip (0.5) and pay ln 2 = 0.693.

flowchart LR
  PR[Prompt] --> LM[Language model]
  LM --> RA[Response A] & RB[Response B]
  RA & RB --> H[Human or AI labeler<br/>picks the better one]
  H --> RM[Reward model<br/>learns to score responses]
  RM --> RL[Reinforcement learning<br/>tune the LM to score high]
  RL --> LM

import math
def sigma(z):
    return 1 / (1 + math.exp(-z))
r_w, r_l = 3.0, 1.0
# P(y_w ≻ y_l)
round(sigma(r_w - r_l), 3)  # → 0.881
# L_RM
round(-math.log(sigma(r_w - r_l)), 3)  # → 0.127
# equal scores: a coin flip, ln 2
round(-math.log(sigma(1.0 - 1.0)), 3)  # → 0.693

In code: sigmoid squashes a score gap into a probability, preference_probability applies it to two rewards (the Bradley-Terry model), and reward_model_loss is minus the log of that probability.

3b. DPO: skip the critic

Everyday picture Instead of hiring a food critic and cooking to please them, edit the recipe book directly from the taste-test results, while keeping a copy of the original book so you don't drift too far from it.

Tiny worked example A log-probability is the log of the probability the model gives a whole response (a sum of per-token log probabilities; more negative means less likely). The frozen reference model gives both answers −11. After some training the policy gives the chosen answer −10 (more likely than before) and the rejected one −12 (less likely). With β = 0.1 the margin is 0.1 × ((−10 − −11) − (−12 − −11)) = 0.1 × (1 + 1) = 0.2; the loss is −ln σ(0.2) = 0.598, down from ln 2 = 0.693 when the policy still equalled the reference.

flowchart LR
  PAIR[Preference pair<br/>chosen, rejected] --> POL[Policy being trained<br/>log π of each]
  PAIR --> REF[Frozen reference<br/>log π_ref of each]
  POL --> M[Margin = β × how much more the policy<br/>boosted chosen than rejected]
  REF --> M
  M --> LOSS[−log σ margin]
  LOSS -->|gradient| POL

import math
def sigma(z):
    return 1 / (1 + math.exp(-z))
beta = 0.1
# chosen answer: policy, reference
logpi_w, logpi_ref_w = -10.0, -11.0
# rejected answer: policy, reference
logpi_l, logpi_ref_l = -12.0, -11.0
# implicit rewards r̂ = β × log-ratio
r_w = beta * (logpi_w - logpi_ref_w)
r_l = beta * (logpi_l - logpi_ref_l)
print(f"{r_w:.1f} {r_l:.1f}")  # → 0.1 -0.1
margin = r_w - r_l
print(f"{margin:.1f} {sigma(margin):.3f}")  # → 0.2 0.550
# L_DPO
round(-math.log(sigma(margin)), 3)  # → 0.598
# how hard this pair pushes
round(beta * (1 - sigma(margin)), 3)  # → 0.045

In code: dpo_margin computes the gap between implicit rewards (β times each log-ratio), dpo_loss turns it into −log σ(margin), and dpo_update_strength gives the push β × (1 − σ(margin)) plotted in the first figure; train_toy_dpo trains the three-answer toy policy of the second.

Why it matters in practice. DPO is simpler and more stable than RLHF, which is why it is widely used in open-model fine-tuning. Constitutional AI and AI-feedback methods scale the labelling by having a model judge responses against written principles. This stage is where tone, helpfulness, refusals and safety behaviour mostly come from.

Everyday picture You want to adapt a 1,000-page textbook for your class. Reprinting it is expensive. Instead you add a small stack of sticky notes with corrections. The book stays untouched; the notes are cheap to write, store and swap, and you can photocopy them into the book when you're done.

Tiny worked example Frozen weights W = the 2×2 identity (it copies its input). Adapter B = (1, 0) as a column and A = (0, 1) as a row. For input x = (1, 2): the frozen path gives W·x = (1, 2). The adapter first squeezes x to one number, A·x = 2, then expands it back, B·2 = (2, 0). The output is (1, 2) + (2, 0) = (3, 2). Two small vectors changed the layer's behaviour without touching W.

flowchart LR
  X[input x<br/>d numbers] --> W[W, frozen<br/>d × d]
  X --> A[A, trainable<br/>squeeze to r numbers]
  A --> B[B, trainable, starts at 0<br/>expand back to d]
  W --> ADD((+))
  B --> ADD
  ADD --> Y[output y]

# x Mᵀ: the dot product of x with each row of M
def times_transpose(x, M):
    return [sum(x_k * m_k for x_k, m_k in zip(x, row)) for row in M]
# frozen, d_out × d_in
W = [[1, 0], [0, 1]]
# r × d_in, with r = 1
A = [[0, 1]]
# d_out × r
B = [[1], [0]]
x, alpha_over_r = [1, 2], 1
# x Wᵀ
frozen = times_transpose(x, W)
# x Aᵀ: squeezed to r numbers
squeezed = times_transpose(x, A)
# (x Aᵀ) Bᵀ: expanded back
correction = times_transpose(squeezed, B)
frozen, squeezed, correction  # → ([1, 2], [2], [2, 0])
# y
[f_j + alpha_over_r * c_j for f_j, c_j in zip(frozen, correction)]  # → [3, 2]
d_in = d_out = 4096
# full, then LoRA at each rank
d_out * d_in, [r * d_in + d_out * r for r in (64, 16, 8)]  # → (16777216, [524288, 131072, 65536])

Parameter savings for one 4096 × 4096 layer:

In code: LoRALinear holds the frozen W beside the trainable A and B, and LoRALinear.merged_weight folds the adapter into W for serving. lora_trainable_params and full_trainable_params count the table's parameters, and train_toy_lora trains the adapters in the figure.

QLoRA goes further: it stores the frozen base weights in 4 bits (a format called NF4) and trains LoRA adapters in 16-bit on top, which lets a 65-billion-parameter model be fine-tuned on a single 48 GB GPU.

Why it matters in practice. LoRA adapters are megabytes, not gigabytes. You can keep one per customer or task and hot-swap them on one shared base model, and training fits on far smaller hardware.

Everyday picture If a new employee doesn't know your product catalogue, you hand them the catalogue (retrieval); you don't send them back to school. If they know the facts but write emails in the wrong tone, you first give clearer instructions, then coach them, and only for a whole new profession do you retrain from scratch.

flowchart TD
  Q1{Missing knowledge,<br/>or facts that change?} -->|Yes| R[Use RAG]
  Q1 -->|No| Q2{Wrong behaviour,<br/>format or tone?}
  Q2 -->|No| N[No change needed]
  Q2 -->|Yes| P[Improve the prompt<br/>add examples]
  P -->|Still inconsistent<br/>or too costly| L[LoRA fine-tune]
  L -->|Major domain shift,<br/>lots of data| F[Full fine-tune]

In code: choose_adaptation walks this flowchart, cheapest option first, and returns the approach it lands on.

The key line: fine-tuning teaches behaviour; RAG supplies knowledge. Most production systems end up as retrieval plus a well-built prompt.

Everyday picture A master chef tastes a sauce and says "mostly thyme, a bit of rosemary, definitely not mint". An apprentice who only hears "thyme" learns less than one who hears the whole judgement. Distillation trains a small, cheap student model to match a big teacher's full probability spread, not just its top answer.

Tiny worked example The teacher's raw scores (logits) for three answers are 2, 1, 0. Softmax at temperature T = 1 gives 0.665, 0.245, 0.090. Dividing the logits by T = 2 first gives softmax(1, 0.5, 0) = 0.506, 0.307, 0.186: flatter, so the student clearly sees that answer 2 is a much better runner-up than answer 3. That runner-up information is sometimes called "dark knowledge".

flowchart LR
  X[Same input] --> T[Big teacher model]
  X --> S[Small student model]
  T --> TS[softmax of logits / T<br/>soft targets]
  S --> SS[softmax of logits / T]
  TS --> KL[KL divergence<br/>how different are they?]
  SS --> KL
  KL -->|gradient| S

import math
def soft_targets(z, T):
    # e^(z_i / T)
    exps = [math.exp(z_i / T) for z_i in z]
    # divided by Σ_j e^(z_j / T)
    return [e / sum(exps) for e in exps]
[round(p_i, 3) for p_i in soft_targets([2, 1, 0], T=2)]  # → [0.506, 0.307, 0.186]
# teacher, student
p, q = [0.5, 0.5], [0.9, 0.1]
KL = sum(p_i * math.log(p_i / q_i) for p_i, q_i in zip(p, q))
round(KL, 3)  # → 0.511
alpha, T = 1.0, 2
# the (1 − α)·CE term is zero at α = 1
round(alpha * T ** 2 * KL, 2)  # → 2.04

In code: soft_targets divides logits by T and applies softmax, kl_divergence measures how far apart two spreads are, and distillation_loss combines them with the T² scale and the optional cross-entropy on the true label.

Why it matters in practice. Distillation is often the biggest cost lever in production: a small student trained on a big model's outputs for one narrow task can be many times cheaper and faster with little quality loss on that task.

Answer each one out loud or on paper before you open it. If you can explain it, you know it.