Loss functions

A loss is a scorecard with one number on it. Imagine a coach who, after every practice, must summarise the whole team's performance in a single score so the players know whether they're improving. Training a model does exactly this: it computes the loss, then adjusts the weights to make that number smaller (see primer.ml.neural_net for the loop).

What the scorecard rewards is what the model learns. Language models are scored on how much probability they gave the real next word; price predictors on how far off their numbers were; search models on whether they ranked the right passage above the wrong ones.

flowchart LR
  P[Prediction] --> LF{Loss function}
  T[Right answer] --> LF
  LF -->|classes or tokens| CE[Cross-entropy]
  LF -->|numbers| REG[MSE or MAE]
  LF -->|matching pairs| CON[Contrastive / InfoNCE]
  CE & REG & CON --> N[One number to minimize]

Picture a weather forecaster scored every evening. If they said "90% chance of rain" and it rained, they lose a point or two. If they said "1% chance of rain" and it poured, they are humiliated. Cross-entropy is that humiliation meter: it charges according to how little probability you gave to what actually happened.

Worked example: the model assigned these probabilities to the correct next token.

Being confidently wrong (1%) costs about 44× more than being mostly right (90%): 4.61 / 0.105 ≈ 44.

The natural logarithm ln(p) answers "e to what power gives p?" For a probability between 0 and 1 the answer is negative (ln 0.5 = −0.69, because e^−0.69 = 0.5), so we flip the sign to get a positive cost. (Every symbol used in these lessons is built from zero in primer.notation.)

import math
p_correct = 0.5
# −ln p_correct
round(-math.log(p_correct), 2)  # → 0.69
# confidently wrong costs far more
round(-math.log(0.01), 2)  # → 4.61

cross_entropy_from_prob is this one line.

Why it matters a language model predicts the next token by choosing among its whole vocabulary, so this is the pretraining loss of every LLM. The steep penalty for confident mistakes is what pushes models toward calibrated probabilities.

A model doesn't output probabilities directly. It outputs raw scores called logits, which softmax turns into probabilities (raise e to each score, divide by the total; see primer.ml.attention). Computing that literally is like measuring everyone's height in millimetres from the centre of the Earth: the numbers get astronomically large and your calculator overflows. Measure relative to the tallest person instead, and every number stays small. That trick is log-sum-exp.

Worked example: logits (2, 1, 0.1), correct class 0. e² + e¹ + e^0.1 = 7.389 + 2.718 + 1.105 = 11.212, ln 11.212 = 2.417, so the loss is 2.417 − 2 = 0.417. And with logits (1000, 0) and class 1 correct, e^1000 overflows any computer, but log-sum-exp gives exactly 1000 − 0 = 1000.

flowchart LR
  Z[Logits z<br/>one score per class] --> M[Subtract max m]
  M --> LSE[log-sum-exp<br/>m + ln sum e^z-m]
  LSE --> L[Loss = LSE minus z_correct]
  Z --> L
  L -.backward.-> G[Gradient = softmax z<br/>minus one-hot target]
  G -.-> U[Raise the right logit,<br/>lower the others by their probability]

import math
def log_sum_exp(z):
    # m = max_j z_j
    m = max(z)
    # m + ln Σ_j e^(z_j − m)
    return m + math.log(sum(math.exp(z_j - m) for z_j in z))
z, y = [2.0, 1.0, 0.1], 0
round(log_sum_exp(z), 3)  # → 2.417
# ln Σ_j e^(z_j) − z_y
round(log_sum_exp(z) - z[y], 3)  # → 0.417
z, y = [1000.0, 0.0], 1
# e^1000 is never computed: no overflow
log_sum_exp(z) - z[y]  # → 1000.0

The gradient, which is how each logit should move:

import math
z, y = [0.0, math.log(3)], 1
exps = [math.exp(z_j) for z_j in z]
# (1/4, 3/4)
softmax = [e / sum(exps) for e in exps]
# (0, 1)
onehot = [1 if j == y else 0 for j in range(len(z))]
# softmax(z) − onehot(y)
[round(s_j - o_j, 2) for s_j, o_j in zip(softmax, onehot)]  # → [0.25, -0.25]

In code: log_sum_exp pulls the largest logit out front, log_softmax subtracts that total from every logit, and softmax_cross_entropy returns the batch's mean loss together with its softmax − onehot gradient.

Why it matters this is why every framework fuses softmax and cross-entropy into one operation that takes logits (torch.nn.functional.cross_entropy). Computing softmax first and then the log is the classic source of NaN losses.

Imagine a game show with doors, one hiding the prize. If the model is as unsure as someone picking between two doors at random, its perplexity is 2; between ten doors, 10. Perplexity converts an average loss back into that "number of doors".

Worked example: a model that gives the right token 50% every time has loss 0.69 per token and perplexity e^0.69 = 2. One that gives 10% every time has perplexity 10. Three confident tokens (0.9) and one disaster (0.001) give perplexity about 6: one bad guess drags the whole average.

flowchart LR
  P["Per-token probabilities<br/>0.5, 0.5, 0.5"] --> NL["−ln each<br/>0.69, 0.69, 0.69"]
  NL --> AV["Average<br/>0.69"]
  AV --> EX["e to that power<br/>e^0.69 = 2"]
  EX --> D["≈ choosing between<br/>2 equally likely doors"]

import math
p = [0.5, 0.5, 0.5]
N = len(p)
# (1/N) Σ −ln p_i
average = sum(-math.log(p_i) for p_i in p) / N
round(average, 2)  # → 0.69
# exp(...)
round(math.exp(average), 1)  # → 2.0

In code: perplexity averages cross_entropy_from_prob over the tokens and raises e to the result.

Why it matters perplexity is the standard training metric for language models. It's comparable only between models that use the same tokenizer on the same text, and it says nothing direct about whether answers are helpful or correct.

When the prediction is a number (a price, a temperature), think of fines for arriving late. MAE (mean absolute error) is a flat rate: every minute late costs the same. MSE (mean squared error) squares the minutes: 10 minutes late costs 100, not 10, so one very late arrival outweighs many slightly late ones.

Worked example: four predictions miss by 1 and one misses by 10. MSE = (1 + 1 + 1 + 1 + 100) / 5 = 20.8, and the outlier is 100/104 = 96% of it. MAE = (1 + 1 + 1 + 1 + 10) / 5 = 2.8, and the outlier is 10/14 = 71%.

y = [0, 0, 0, 0, 0]
y_hat = [1, -1, 1, -1, 10]
N = len(y)
# MSE
sum((y_i - y_hat_i) ** 2 for y_i, y_hat_i in zip(y, y_hat)) / N  # → 20.8
# MAE
sum(abs(y_i - y_hat_i) for y_i, y_hat_i in zip(y, y_hat)) / N  # → 2.8

In code: mse and mae are the two averages, and outlier_share measures how much of each total the single largest error contributes (the 96% and 71% above).

Why it matters pick the loss whose valley is where you want your predictions. MSE chases outliers (its best constant is the mean); MAE shrugs them off (its best constant is the median). For noisy data with occasional wild values, MAE or a blend (Huber loss) is safer.

Picture a party game. Everyone arrives in pairs, gets separated, and must pick their partner out of the whole room. Everyone else in the room is a decoy. A contrastive loss scores how confidently each person picks their own partner over every decoy. The hardest decoy is your partner's lookalike twin: same topic, wrong person. That's a hard negative.

Embedding models (and CLIP) learn this way: a query and the passage that answers it are "partners"; the other passages in the same batch are free decoys (in-batch negatives).

Worked example: two queries and two passages, as vectors, scored with the dot product (multiply matching entries, add them up). Query 1 = (1, 0), query 2 = (0, 1), passage 1 = (1, 0), passage 2 = (0, 1), temperature 1. Query 1 scores 1 against its partner and 0 against the decoy. Its loss is −ln(e¹ / (e¹ + e⁰)) = ln(1 + e^−1) = 0.3133. Swap the passages and each query now prefers the decoy: the loss rises to ln(1 + e¹) = 1.3133.

flowchart LR
  Q[Batch of queries] --> S[Score matrix<br/>every query vs every passage]
  D[Their matching passages] --> S
  S --> CE[Cross-entropy per row<br/>right answer = the diagonal]
  CE --> L[Pull pairs together<br/>push the rest apart]

import math
# the queries
q = [[1, 0], [0, 1]]
# their passages: d[i] is q[i]'s partner
d = [[1, 0], [0, 1]]
def dot(a, b):
    return sum(a_k * b_k for a_k, b_k in zip(a, b))
def info_nce(i, tau):
    # exp(q_i · d_j / τ) for every j
    scores = [math.exp(dot(q[i], d_j) / tau) for d_j in d]
    # −ln (own passage's share)
    return -math.log(scores[i] / sum(scores))
round(info_nce(0, tau=1.0), 4)  # → 0.3133
print(f"{info_nce(0, tau=0.1):.6f}")  # → 0.000045

It's just cross-entropy where the "classes" are the passages in the batch, so the gradient is the same softmax − onehot, pushed back through the dot products into both sets of vectors.

In code: info_nce builds the score grid, hands it to softmax_cross_entropy with the diagonal as the right answers, and returns the gradients for both sets of vectors; info_nce_gradient_check confirms those gradients against small nudges of every number.

Why it matters this is how search and RAG embedding models are trained (see primer.ml.embeddings.contrastive). Bigger batches mean more free negatives. Easy negatives are already far away and contribute almost no gradient; hard negatives produce most of the learning signal, which is why mining them is the biggest lever on retrieval quality.

A good forecaster never says "100% chance of rain", because the one day they're wrong would be infinitely embarrassing. Label smoothing teaches the model the same humility: instead of "the answer is class 2, with total certainty", the target says "class 2, with 92.5% certainty, and a sliver of doubt spread over everything".

Worked example: 4 classes, correct class 2, smoothing ε = 0.1. Take 0.9 of the one-hot target (0, 0, 0.9, 0) and add 0.1/4 = 0.025 to every class: (0.025, 0.025, 0.925, 0.025). A model that puts logit 50 on the right class has plain cross-entropy ≈ 0, but smoothed cross-entropy 3 × 0.025 × 50 = 3.75.

flowchart LR
  O[One-hot target<br/>0, 0, 1, 0] --> MIX[Mix: 1 minus eps times one-hot<br/>plus eps/K everywhere]
  U[Uniform over K classes] --> MIX
  MIX --> T[Soft target<br/>0.025, 0.025, 0.925, 0.025]
  T --> CE[Cross-entropy against<br/>the soft target]

import math
eps, K, y = 0.1, 4, 2
# t_k
t = [(1 - eps) * (1 if k == y else 0) + eps / K for k in range(K)]
[round(t_k, 3) for t_k in t]  # → [0.025, 0.025, 0.925, 0.025]
z = [0.0, 0.0, 50.0, 0.0]
m = max(z)
log_total = m + math.log(sum(math.exp(z_k - m) for z_k in z))
# ln softmax(z)_k
log_softmax = [z_k - log_total for z_k in z]
# −Σ_k t_k ln softmax(z)_k
round(-sum(t_k * ls_k for t_k, ls_k in zip(t, log_softmax)), 2)  # → 3.75

In code: smoothed_targets builds the soft target t, and smoothed_cross_entropy scores the logits against it.

Why it matters it curbs over-confidence and often improves calibration (how well the model's stated confidence matches how often it's right). It was used to train the original transformer.

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