MMLU harness and calibration

Track 06 - Evaluation · Notebook 02 · Runtime: ~5s on CPU

Prerequisites: 06_eval/01 (perplexity from scratch).

Papers: Hendrycks et al. 2021 (2009.03300); Guo et al. 2017, On Calibration of Modern Neural Networks (1706.04599).

Two things matter when scoring a multiple-choice benchmark:

1. Answer extraction — pick the option with the highest continuation log-probability. Subtleties: length normalisation, choice-letter tokenisation, prompt construction.

2. Calibration — accuracy says how often the model is right; calibration says whether its confidence is honest. We measure with Expected Calibration Error (ECE).

To make the relationship between accuracy and confidence visible, we use three synthetic logit profiles (well-calibrated strong, overconfident strong, well-calibrated weak) over 300 questions. The harness logic is the same one a real run uses; only the logits source is synthetic.

For a real Inspect AI harness running against a real model, see 08_production/08_inspect_ai_eval_harness.

from llm_systems_cookbook.nb import bootstrap

import numpy as np

s = bootstrap("06_eval_02_mmlu_harness_calibration")

Multiple-choice scoring

For a question with options A/B/C/D the model outputs four log-probabilities (la, lb, lc, ld). Argmax picks the answer. We normalise via softmax to get (pa, pb, pc, pd) for calibration work.

def softmax(logits: np.ndarray) -> np.ndarray:
    m = logits.max(axis=-1, keepdims=True)
    e = np.exp(logits - m)
    return e / e.sum(axis=-1, keepdims=True)


def predict_and_confidence(logits: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    probs = softmax(logits)
    preds = probs.argmax(axis=-1)
    confs = probs.max(axis=-1)
    return preds, confs

Three stubbed model profiles over 300 MMLU-style questions

We generate ground-truth answers and simulate three models by injecting different logit patterns:

• strong_calibrated: 80% accuracy; when it’s right the top logit is much larger than the others; when it’s wrong the top logit is only slightly larger. → high accuracy, well-calibrated.

• strong_overconfident: same 80% accuracy but every top logit is large, regardless of correctness. → high accuracy, poorly calibrated (systematic overconfidence).

• weak_calibrated: 40% accuracy, top logit reflects uncertainty proportionally. → low accuracy, reasonably calibrated.

N = 300
K = 4  # A, B, C, D

rng = np.random.default_rng(0)
true = rng.integers(0, K, size=N)


def _logit_for_target_top_prob(p_top: float, k: int = K) -> float:
    """Peak logit gap so softmax assigns probability p_top to the top class
    when the other k-1 classes share equal logit 0."""
    return float(np.log(p_top * (k - 1) / (1 - p_top)))


def stub_logits(accuracy: float, top_prob: float) -> np.ndarray:
    """Synthetic logits: argmax matches `true` with probability `accuracy`,
    and the softmax peak is `top_prob` on every row.

    `top_prob ≈ accuracy` ⇒ well-calibrated;
    `top_prob >> accuracy` ⇒ overconfident;
    `top_prob << accuracy` ⇒ underconfident.
    """
    gap = _logit_for_target_top_prob(top_prob)
    logits = np.zeros((N, K))
    for i in range(N):
        is_right = rng.random() < accuracy
        chosen = true[i] if is_right else int(rng.choice([j for j in range(K) if j != true[i]]))
        logits[i, chosen] = gap
    return logits


profiles = {
    "strong_calibrated":    stub_logits(accuracy=0.80, top_prob=0.80),
    "strong_overconfident": stub_logits(accuracy=0.80, top_prob=0.98),
    "weak_calibrated":      stub_logits(accuracy=0.40, top_prob=0.40),
}
for name, logits in profiles.items():
    preds, confs = predict_and_confidence(logits)
    acc = (preds == true).mean()
    print(f"{name:<22}  acc = {acc:.3f}   mean confidence = {confs.mean():.3f}")

Expected Calibration Error

Bucket predictions by their confidence, compare per-bucket mean confidence to per-bucket accuracy:

ECE = sum_b (n_b / N) * |acc_b - conf_b|

A perfectly calibrated model has ECE = 0.

def expected_calibration_error(preds: np.ndarray, confs: np.ndarray, y: np.ndarray,
                                n_bins: int = 10) -> float:
    correct = (preds == y).astype(float)
    bins = np.linspace(0, 1, n_bins + 1)
    ece = 0.0
    N_total = len(preds)
    for b in range(n_bins):
        lo, hi = bins[b], bins[b + 1]
        mask = (confs > lo) & (confs <= hi) if b > 0 else (confs >= lo) & (confs <= hi)
        if not mask.any():
            continue
        bucket_acc = correct[mask].mean()
        bucket_conf = confs[mask].mean()
        ece += (mask.sum() / N_total) * abs(bucket_acc - bucket_conf)
    return float(ece)


for name, logits in profiles.items():
    preds, confs = predict_and_confidence(logits)
    acc = (preds == true).mean()
    ece = expected_calibration_error(preds, confs, true)
    print(f"{name:<22}  acc = {acc:.3f}   ECE = {ece:.3f}")

Checks

• Accuracy within 5 pp of target for each profile.

• Overconfident profile has strictly higher ECE than the calibrated profile at the same accuracy - the signature of miscalibration.

• Every confidence ≥ 1/K (since argmax of 4 softmax probs is always above the uniform floor of 0.25).

accs = {}
eces = {}
for name, logits in profiles.items():
    preds, confs = predict_and_confidence(logits)
    accs[name] = (preds == true).mean()
    eces[name] = expected_calibration_error(preds, confs, true)

s.assert_close("strong_calibrated_accuracy_80pct",    actual=accs["strong_calibrated"],    expected=0.80, rtol=0.10)
s.assert_close("strong_overconfident_accuracy_80pct", actual=accs["strong_overconfident"], expected=0.80, rtol=0.10)
s.assert_close("weak_calibrated_accuracy_40pct",      actual=accs["weak_calibrated"],      expected=0.40, rtol=0.15)
s.check(
    "overconfident_has_higher_ECE_than_calibrated",
    lambda: eces["strong_overconfident"] > eces["strong_calibrated"],
    msg=f"calibrated ECE={eces['strong_calibrated']:.3f}  overconf ECE={eces['strong_overconfident']:.3f}",
)
s.check(
    "all_confidences_above_uniform_floor",
    lambda: all((softmax(logits).max(axis=-1) >= 1.0 / K - 1e-9).all() for logits in profiles.values()),
    msg="argmax of softmax(4-way) ≥ 0.25",
)

Reliability diagram

The ECE scalar compresses a lot of detail. The reliability diagram keeps it: confidence bins on the x-axis, empirical accuracy per bin on the y, and the y = x line as the perfect-calibration reference. Points above the diagonal are underconfident, points below are overconfident. The overconfident profile sits well below the line at high-confidence bins while the calibrated profile hugs it.

import matplotlib.pyplot as plt

def reliability_points(preds, confs, y, n_bins=10):
    correct = (preds == y).astype(float)
    edges = np.linspace(0, 1, n_bins + 1)
    xs, ys, ws = [], [], []
    for b in range(n_bins):
        lo, hi = edges[b], edges[b + 1]
        m = (confs > lo) & (confs <= hi) if b > 0 else (confs >= lo) & (confs <= hi)
        if m.any():
            xs.append(confs[m].mean()); ys.append(correct[m].mean()); ws.append(m.sum())
    return np.array(xs), np.array(ys), np.array(ws)

fig, ax = plt.subplots(figsize=(5.8, 5.2))
ax.plot([0, 1], [0, 1], "k--", label="perfect calibration")
colors = {"strong_calibrated": "tab:blue", "strong_overconfident": "tab:red", "weak_calibrated": "tab:green"}
for name, logits in profiles.items():
    preds, confs = predict_and_confidence(logits)
    xs, ys, ws = reliability_points(preds, confs, true)
    sizes = 30 + (ws / ws.max()) * 200  # marker size ∝ bin count
    ax.plot(xs, ys, "-", color=colors[name], alpha=0.7,
            label=f"{name}  (ECE={eces[name]:.3f})")
    ax.scatter(xs, ys, s=sizes, color=colors[name], edgecolor="white", linewidth=0.5)
ax.set_xlim(0, 1); ax.set_ylim(0, 1)
ax.set_xlabel("mean confidence in bin"); ax.set_ylabel("empirical accuracy in bin")
ax.set_title("reliability diagram  (marker size ∝ bin count)")
ax.legend(loc="upper left", fontsize=9); ax.grid(alpha=0.3)
fig.tight_layout(); plt.show()

Exercises

1. Temperature scaling. Fit a single temperature T on a held-out set; divide logits by T before softmax (Guo 2017). Standard post-hoc calibration fix.

2. Length normalisation. lm-eval-harness scores the full "A. <answer text>" continuation and divides log-prob by token count. Adjust the harness here and see how the numbers move.

3. Real model. Use lm_eval with --tasks mmlu --model hf against a small open model (Qwen2.5-0.5B); compare its real ECE to the synthetic profiles above.

References

• lm-evaluation-harness mmlu task source — the production pattern this notebook mirrors.

• Guo et al. 2017 — temperature scaling, the simplest calibration fix.

• 08_production/08_inspect_ai_eval_harness — building a custom evaluation against a real model with Inspect AI.

s.summary()
s.save()