Distillation & LLM Evals - Teacher-Student Training, Benchmarks, LLM-as-Judge, Regression Gates & Contamination

"""Two eval-adjacent tools in pure stdlib:
 (1) knowledge distillation losses: temperature-softened KL, forward vs reverse KL;
 (2) a 13-gram overlap contamination check between a benchmark and training text.
"""
import math

def softmax(z, T=1.0):
    m = max(z); e = [math.exp((v - m) / T) for v in z]; s = sum(e)
    return [v / s for v in e]

def kl(p, q):
    return sum(pi * math.log(pi / qi) for pi, qi in zip(p, q) if pi > 1e-12)

# (1a) Temperature reveals "dark knowledge": the teacher's ranking of wrong answers.
teacher_logits = [6.0, 4.5, 1.0, -2.0]          # classes: cat, lynx, dog, truck
for T in (1, 2, 4):
    print(f"teacher probs at T={T}: {[round(p, 3) for p in softmax(teacher_logits, T)]}")
student_logits = [5.0, 1.0, 3.5, -1.0]           # right top-1, wrong idea of what is similar
for T in (1, 4):
    loss = (T ** 2) * kl(softmax(teacher_logits, T), softmax(student_logits, T))
    print(f"distill loss T^2*KL(teacher||student) at T={T}: {loss:.3f}   (T^2 keeps gradient scale comparable)")
print("higher T exposes the student's wrong similarity ranking (dog above lynx) that T=1 mostly hides")

# (1b) Forward vs reverse KL on a bimodal teacher, student restricted to ONE peak shape.
teacher = [0.45, 0.05, 0.0001, 0.05, 0.45]          # two good answers (indices 0 and 4)
cands = {"cover both (spread)": [0.25, 0.15, 0.2, 0.15, 0.25],
         "seek one mode":       [0.90, 0.05, 0.03, 0.01, 0.01]}
print("\nstudent choice           forward KL(t||s)  reverse KL(s||t)")
for name, s in cands.items():
    print(f"{name:<24}{kl(teacher, s):>16.3f}{kl(s, teacher):>18.3f}")
print("forward KL (classic logit distillation) prefers covering both modes; reverse KL (on-policy")
print("distillation: MiniLLM uses reverse KL, GKD trains on student samples with reverse KL or JSD)")
print("prefers committing to one good mode.")

# (2) Contamination: does a benchmark item appear (near-)verbatim in training data?
def ngrams(text, n=13):
    toks = text.lower().split()
    return {tuple(toks[i:i + n]) for i in range(len(toks) - n + 1)}

train_corpus = ("forum post: here is the answer to the classic puzzle a bat and a ball cost one dollar "
                "and ten cents in total the bat costs one dollar more than the ball how much does the ball cost "
                "answer five cents")
bench = {
    "q1 (leaked)": "a bat and a ball cost one dollar and ten cents in total the bat costs one dollar more than the ball how much does the ball cost",
    "q2 (paraphrased)": "together a racket and a shuttlecock are priced at 1.10 and the racket is 1.00 more than the shuttlecock what is the shuttlecock price",
}
train_grams = ngrams(train_corpus)
for name, q in bench.items():
    g = ngrams(q)
    hit = len(g & train_grams) / max(1, len(g))
    print(f"{name:<18} 13-gram overlap={hit:.0%}  -> {'FLAG: report separately or drop' if hit > 0.5 else 'passes n-gram check'}")
print("Failure path: q2 passes the n-gram check but is the same problem; paraphrase leaks need embedding")
print("similarity search or fresh, post-cutoff items (dynamic benchmarks) to catch.")

teacher probs at T=1: [0.813, 0.181, 0.005, 0.0]
teacher probs at T=2: [0.636, 0.3, 0.052, 0.012]
teacher probs at T=4: [0.474, 0.326, 0.136, 0.064]
distill loss T^2*KL(teacher||student) at T=1: 0.445   (T^2 keeps gradient scale comparable)
distill loss T^2*KL(teacher||student) at T=4: 2.078   (T^2 keeps gradient scale comparable)
higher T exposes the student's wrong similarity ranking (dog above lynx) that T=1 mostly hides

student choice           forward KL(t||s)  reverse KL(s||t)
cover both (spread)                0.418             1.556
seek one mode                      1.481             0.741
forward KL (classic logit distillation) prefers covering both modes; reverse KL (on-policy
distillation: MiniLLM uses reverse KL, GKD trains on student samples with reverse KL or JSD)
prefers committing to one good mode.
q1 (leaked)        13-gram overlap=100%  -> FLAG: report separately or drop
q2 (paraphrased)   13-gram overlap=0%  -> passes n-gram check
Failure path: q2 passes the n-gram check but is the same problem; paraphrase leaks need embedding
similarity search or fresh, post-cutoff items (dynamic benchmarks) to catch.

// Eval regression gate for a candidate model vs the current production model.
// (1) Paired bootstrap on per-item scores -> confidence interval for the delta.
// (2) Per-slice gates: no protected slice may regress beyond a tolerance, even if the average improves.
// (3) LLM-as-judge position-bias check: judge each pair twice with A/B swapped; count flips.

function lcg(seed: number): () => number {
  let s = seed >>> 0;
  return () => ((s = (Math.imul(s, 1664525) + 1013904223) >>> 0) / 2 ** 32);
}
const rnd = lcg(2026);

type Item = { slice: string; prod: number; cand: number };
const slices: Array<[string, number, number, number]> = [
  // name, items, prod pass rate, candidate pass rate (synthetic)
  ["coding", 400, 0.62, 0.70],
  ["math", 300, 0.55, 0.60],
  ["safety-refusals", 200, 0.97, 0.91], // the regression hiding under a better average
  ["multilingual", 100, 0.70, 0.71],
];
const items: Item[] = [];
for (const [slice, n, p, c] of slices) {
  for (let i = 0; i < n; i++) {
    const u = rnd(); // shared noise makes scores paired (same item difficulty for both models)
    items.push({ slice, prod: u < p ? 1 : 0, cand: u < c ? 1 : 0 });
  }
}

function bootstrapDelta(xs: Item[], iters = 2000): [number, number, number] {
  const deltas: number[] = [];
  for (let b = 0; b < iters; b++) {
    let d = 0;
    for (let i = 0; i < xs.length; i++) {
      const it = xs[Math.floor(rnd() * xs.length)];
      d += it.cand - it.prod;
    }
    deltas.push(d / xs.length);
  }
  deltas.sort((a, b) => a - b);
  const mean = xs.reduce((s, it) => s + it.cand - it.prod, 0) / xs.length;
  return [mean, deltas[Math.floor(0.025 * iters)], deltas[Math.floor(0.975 * iters)]];
}

const fmt = (x: number) => (x >= 0 ? "+" : "") + (100 * x).toFixed(1) + "pp";
const [m, lo, hi] = bootstrapDelta(items);
console.log(`overall delta ${fmt(m)}  95% CI [${fmt(lo)}, ${fmt(hi)}]`);
let pass = lo > 0;
const tolerance = -0.02; // a protected slice may not drop more than 2 points (upper CI bound)
for (const [slice] of slices) {
  const [sm, slo, shi] = bootstrapDelta(items.filter((it) => it.slice === slice));
  const bad = shi < 0 || sm < tolerance;
  if (bad) pass = false;
  console.log(`  ${slice.padEnd(16)} ${fmt(sm).padStart(8)}  CI [${fmt(slo)}, ${fmt(shi)}]${bad ? " <- REGRESSION" : ""}`);
}

// Judge position bias: a biased judge prefers whichever answer is shown first 30% of the time.
let flips = 0;
const N = 500;
for (let i = 0; i < N; i++) {
  const trueBetterIsA = rnd() < 0.5;
  const judge = (firstIsA: boolean) => (rnd() < 0.3 ? firstIsA : trueBetterIsA); // returns "A wins?"
  if (judge(true) !== judge(false)) flips++;
}
console.log(`\njudge verdict flipped when A/B order swapped on ${((100 * flips) / N).toFixed(0)}% of pairs -> average both orders, count flips as ties`);
console.log(`GATE: ${pass ? "PASS" : "FAIL"} (overall CI must exclude 0 AND no protected slice may regress)`);

overall delta +4.3pp  95% CI [+2.8pp, +5.9pp]
  coding             +8.3pp  CI [+5.5pp, +11.0pp]
  math               +6.7pp  CI [+4.0pp, +9.7pp]
  safety-refusals    -6.5pp  CI [-10.0pp, -3.5pp] <- REGRESSION
  multilingual       +3.0pp  CI [+0.0pp, +7.0pp]

judge verdict flipped when A/B order swapped on 30% of pairs -> average both orders, count flips as ties
GATE: FAIL (overall CI must exclude 0 AND no protected slice may regress)