| |
| |
| |
| """Fidelity instrumentation + recursive fidelity search (F5). |
| |
| The most important measured failure is 3/24 token agreement. This module |
| instruments EXACTLY where teacher/student behavior diverges: |
| |
| * per-token divergence trace: position, teacher/student tokens, |
| teacher/student top-k, full logits, KL / JS divergence, cosine |
| similarity, entropies, per-layer hidden-state divergence and |
| per-layer attention divergence — all measured on teacher-forced |
| contexts so model divergence is isolated from context drift; |
| * first-divergence detection: the first position AND the first layer |
| producing significant representation drift (not just final tokens); |
| * recursive fidelity search: candidate transforms that save parameters |
| but cause behavioral divergence pay an explicit, measured cost: |
| utility = savings - recon_penalty - logit_penalty |
| - hidden_penalty - task_penalty |
| * ablations by component (embeddings, attention projections, MLP |
| up/down, norms, output head, quantization) reporting measured |
| evidence — never guesses — about where fidelity is lost. |
| |
| The fidelity penalty applies UNIFORMLY to every decision class |
| (including QUANTUM-ENCODED, scored from its own measured reconstruction |
| error). Nothing here touches the quantum scoring to force a winner; the |
| experiment still decides. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import math |
| from typing import Dict, List, Optional, Tuple |
|
|
| import numpy as np |
|
|
| from . import core as C |
| from .model import forward |
|
|
|
|
| |
| |
| |
|
|
| def _softmax(x: np.ndarray) -> np.ndarray: |
| x = np.asarray(x, dtype=np.float64) |
| x = x - x.max(axis=-1, keepdims=True) |
| e = np.exp(x) |
| return e / e.sum(axis=-1, keepdims=True) |
|
|
|
|
| def _entropy(p: np.ndarray) -> float: |
| p = np.asarray(p, dtype=np.float64) |
| p = p / p.sum() |
| return float(-(p * np.log(p + 1e-300)).sum()) |
|
|
|
|
| def kl_div(p_logits: np.ndarray, q_logits: np.ndarray) -> float: |
| """KL(softmax(p) || softmax(q)).""" |
| p, q = _softmax(p_logits), _softmax(q_logits) |
| return float((p * (np.log(p + 1e-300) - np.log(q + 1e-300))).sum()) |
|
|
|
|
| def js_div(p_logits: np.ndarray, q_logits: np.ndarray) -> float: |
| """Jensen-Shannon divergence between the two softmax distributions.""" |
| p, q = _softmax(p_logits), _softmax(q_logits) |
| m = 0.5 * (p + q) |
| kl_pm = float((p * (np.log(p + 1e-300) - np.log(m + 1e-300))).sum()) |
| kl_qm = float((q * (np.log(q + 1e-300) - np.log(m + 1e-300))).sum()) |
| return 0.5 * (kl_pm + kl_qm) |
|
|
|
|
| def top_k(logits: np.ndarray, k: int = 5) -> List[Tuple[int, float]]: |
| p = _softmax(logits) |
| idx = np.argsort(p)[::-1][:k] |
| return [(int(i), float(p[i])) for i in idx] |
|
|
|
|
| |
| |
| |
|
|
| def divergence_trace(t_params: dict, t_cfg, s_params: dict, s_cfg, |
| prompt: List[int], max_tokens: int = 24, |
| top_k_k: int = 5) -> dict: |
| """Per-token teacher/student divergence on teacher-forced contexts. |
| |
| 1. Teacher free-runs greedy -> reference token sequence. |
| 2. Student free-runs greedy -> token agreement count (the 3/24 number). |
| 3. At each position, BOTH models run on the teacher's prefix |
| (teacher-forcing), isolating model divergence from context drift. |
| Records per position: tokens, top-k, full logits, KL/JS/cosine, |
| entropies, per-layer hidden-state divergence (rel err + cosine of the |
| final hidden vectors) and per-layer attention divergence |
| (mean |P_student - P_teacher| over heads). |
| """ |
| prompt = list(prompt) |
| |
| t_toks, _, _ = _greedy(t_params, t_cfg, prompt, max_tokens) |
| s_toks, _, _ = _greedy(s_params, s_cfg, prompt, max_tokens) |
| agreement = sum(a == b for a, b in zip(t_toks, s_toks)) |
|
|
| positions = [] |
| n_align = min(s_cfg.layers, t_cfg.layers) |
| for i in range(len(t_toks)): |
| ctx = np.array([prompt + t_toks[:i]], dtype=np.int64) |
| tl, tc = forward(t_params, t_cfg, ctx[:, -t_cfg.seq:]) |
| sl, sc = forward(s_params, s_cfg, ctx[:, -s_cfg.seq:]) |
| t_log = tl[0, -1].astype(np.float64) |
| s_log = sl[0, -1].astype(np.float64) |
| pos = { |
| "position": i, |
| "teacher_token": int(t_toks[i]), |
| "student_token": int(s_toks[i]), |
| "match": bool(t_toks[i] == s_toks[i]), |
| "teacher_topk": top_k(t_log, top_k_k), |
| "student_topk": top_k(s_log, top_k_k), |
| "teacher_logits": [float(x) for x in t_log], |
| "student_logits": [float(x) for x in s_log], |
| "kl_div": kl_div(t_log, s_log), |
| "js_div": js_div(t_log, s_log), |
| "cosine_sim": C.cosine_sim(t_log, s_log), |
| "teacher_entropy": _entropy(_softmax(t_log)), |
| "student_entropy": _entropy(_softmax(s_log)), |
| "hidden_div": {}, |
| "attn_div": {}, |
| } |
| for l in range(n_align): |
| th = tc[f"L{l}"]["y"] if "y" in tc[f"L{l}"] else tc["f"] |
| sh = sc[f"L{l}"]["y"] if "y" in sc[f"L{l}"] else sc["f"] |
| |
| sv = np.asarray(sh[0, -1], dtype=np.float64) |
| tv = np.asarray(th[0, -1], dtype=np.float64) |
| n = min(sv.size, tv.size) |
| pos["hidden_div"][f"layer_{l}"] = { |
| "rel_err": C.relative_error(tv[:n], sv[:n]), |
| "cosine_sim": C.cosine_sim(tv[:n], sv[:n]), |
| } |
| pt = tc[f"L{l}"]["P"][0].mean(axis=0) |
| ps = sc[f"L{l}"]["P"][0].mean(axis=0) |
| nT = min(pt.shape[0], ps.shape[0]) |
| pos["attn_div"][f"layer_{l}"] = float( |
| np.abs(pt[:nT, :nT] - ps[:nT, :nT]).mean()) |
| positions.append(pos) |
|
|
| return {"prompt": prompt, |
| "teacher_tokens": [int(x) for x in t_toks], |
| "student_tokens": [int(x) for x in s_toks], |
| "token_agreement": f"{agreement}/{len(t_toks)}", |
| "agreement_count": agreement, |
| "n_tokens": len(t_toks), |
| "positions": positions} |
|
|
|
|
| def _greedy(params, cfg, prompt, max_tokens): |
| from .model import generate |
| return generate(params, cfg, list(prompt), max_tokens) |
|
|
|
|
| def first_divergence(trace: dict, *, kl_thresh: float = 0.5, |
| hidden_thresh: float = 0.35) -> dict: |
| """Locate the FIRST divergence: earliest position where the student |
| token mismatches OR KL exceeds threshold; then the layer with the |
| largest hidden-state relative error at that position — the first |
| layer producing significant representation drift.""" |
| for pos in trace["positions"]: |
| if (not pos["match"]) or pos["kl_div"] > kl_thresh: |
| hd = pos["hidden_div"] |
| worst_layer = max(hd.items(), |
| key=lambda kv: kv[1]["rel_err"])[0] |
| return { |
| "position": pos["position"], |
| "teacher_token": pos["teacher_token"], |
| "student_token": pos["student_token"], |
| "kl_div": pos["kl_div"], |
| "js_div": pos["js_div"], |
| "first_diverging_layer": worst_layer, |
| "layer_hidden_rel_err": hd[worst_layer]["rel_err"], |
| "layer_attn_div": pos["attn_div"].get(worst_layer), |
| "all_layer_hidden_rel_err": { |
| k: v["rel_err"] for k, v in hd.items()}, |
| "above_hidden_thresh": { |
| k: v["rel_err"] for k, v in hd.items() |
| if v["rel_err"] > hidden_thresh}, |
| } |
| return {"position": None, "note": "no divergence above thresholds"} |
|
|
|
|
| def render_trace(trace: dict, first: dict) -> List[str]: |
| lines = [] |
| for pos in trace["positions"]: |
| mark = "PASS" if pos["match"] else "DIVERGED" |
| lines.append( |
| f"token {pos['position']:2d} {mark:8s} " |
| f"t={pos['teacher_token']:3d} s={pos['student_token']:3d} " |
| f"KL={pos['kl_div']:.3f} JS={pos['js_div']:.3f} " |
| f"cos={pos['cosine_sim']:.3f}") |
| if first.get("position") is not None: |
| lines.append(f"first divergence at position {first['position']}: " |
| f"layer {first['first_diverging_layer']} " |
| f"(hidden rel_err={first['layer_hidden_rel_err']:.3f})") |
| return lines |
|
|
|
|
| |
| |
| |
|
|
| def _align_to(src: np.ndarray, shape: tuple) -> np.ndarray: |
| out = np.zeros(shape, dtype=np.float32) |
| sl = tuple(slice(0, min(a, b)) for a, b in zip(src.shape, shape)) |
| out[sl] = np.asarray(src)[sl] |
| return out |
|
|
|
|
| COMPONENTS = [ |
| ("embeddings", lambda k: k == "emb"), |
| ("attention_projections", |
| lambda k: any(x in k for x in (".Wq", ".Wk", ".Wv", ".bq", ".bk", ".bv"))), |
| ("attention_output", |
| lambda k: any(x in k for x in (".Wo", ".bo"))), |
| ("mlp_up", lambda k: ".W1" in k or ".b1" in k), |
| ("mlp_down", lambda k: ".W2" in k or ".b2" in k), |
| ("normalization", |
| lambda k: any(x in k for x in ("ln1g", "ln1b", "ln2g", "ln2b", |
| "lnfg", "lnfb"))), |
| ("output_head", lambda k: k in ("Whead", "bhead")), |
| ] |
|
|
|
|
| def ablation_report(t_params: dict, t_cfg, s_params: dict, s_cfg, |
| prompts: List[List[int]], max_tokens: int = 16) -> dict: |
| """For each component: replace the student's component with the |
| teacher's (shape-aligned), measure token-agreement and mean-KL delta |
| vs the unmodified student. Returns ranked evidence — the component |
| whose restoration helps most is where fidelity is being lost.""" |
| base = _agreement_and_kl(t_params, t_cfg, s_params, s_cfg, |
| prompts, max_tokens) |
| rows = [{"component": "baseline(student)", |
| "agreement": base[0], "mean_kl": base[1], "delta_agree": 0, |
| "n_swapped": 0}] |
| for name, pred in COMPONENTS: |
| variant = dict(s_params) |
| n = 0 |
| for k in list(variant.keys()): |
| if pred(k) and k in t_params: |
| variant[k] = _align_to(t_params[k], variant[k].shape) |
| n += 1 |
| if n == 0: |
| continue |
| agree, mkl = _agreement_and_kl(t_params, t_cfg, variant, s_cfg, |
| prompts, max_tokens) |
| rows.append({"component": name, "agreement": agree, "mean_kl": mkl, |
| "delta_agree": agree - base[0], |
| "delta_kl": mkl - base[1], "n_swapped": n}) |
| rows.sort(key=lambda r: -r["delta_agree"]) |
| return {"baseline": {"agreement": base[0], "mean_kl": base[1]}, |
| "ranking": rows} |
|
|
|
|
| def _agreement_and_kl(t_params, t_cfg, s_params, s_cfg, prompts, max_tokens): |
| agrees, kls, n = 0, [], 0 |
| for pr in prompts: |
| tt, _, _ = _greedy(t_params, t_cfg, pr, max_tokens) |
| st, _, _ = _greedy(s_params, s_cfg, pr, max_tokens) |
| agrees += sum(a == b for a, b in zip(tt, st)) |
| n += len(tt) |
| for i in range(len(tt)): |
| ctx = np.array([pr + tt[:i]], dtype=np.int64) |
| tl, _ = forward(t_params, t_cfg, ctx[:, -t_cfg.seq:]) |
| sl, _ = forward(s_params, s_cfg, ctx[:, -s_cfg.seq:]) |
| kls.append(kl_div(tl[0, -1], sl[0, -1])) |
| return agrees / max(1, n), float(np.mean(kls)) if kls else 0.0 |
|
|
|
|
| |
| |
| |
|
|
| def fidelity_penalty(sensitivity: float, reconstruction_error: float, |
| lam: float = 1.0) -> float: |
| """Explicit cost for a candidate transform that saves parameters but |
| causes behavioral divergence. Applied uniformly to every decision |
| class from measured (sensitivity, reconstruction_error).""" |
| return float(lam * max(0.0, sensitivity) * max(0.0, reconstruction_error)) |
|
|
|
|
| def fidelity_aware_select(leaves, budget_params: float, |
| sensitivity: Dict[str, float], |
| lam: float = 1.0) -> Dict[str, dict]: |
| """Greedy knapsack identical in structure to core.select_plan, but each |
| option's effective energy is |
| |
| effective = retained_energy - lam * sensitivity * recon_error |
| |
| (clipped at 0). The penalty is computed from MEASURED sensitivity and |
| MEASURED reconstruction error for every option — quantum-gated options |
| included, from their own qsim-measured errors. No decision class is |
| favored or punished beyond what the measurements say. |
| """ |
| srt = [] |
| for leaf in leaves: |
| sens = float(sensitivity.get(leaf.name, 0.5)) |
| adj = [] |
| for o in leaf.options: |
| pen = fidelity_penalty(sens, o["reconstruction_error"], lam) |
| eff = max(0.0, o["retained_energy"] - pen) |
| adj.append({**o, "effective_energy": eff, |
| "fidelity_penalty": pen}) |
| adj.sort(key=lambda o: o["cost_params"]) |
| srt.append(adj) |
| choice = [0] * len(srt) |
| total = sum(s[0]["cost_params"] for s in srt) |
| improved = True |
| while improved: |
| improved = False |
| best_gain, best_move = 0.0, None |
| for i, opts in enumerate(srt): |
| cur = choice[i] |
| for j in range(cur + 1, len(opts)): |
| dcost = opts[j]["cost_params"] - opts[cur]["cost_params"] |
| denergy = (opts[j]["effective_energy"] |
| - opts[cur]["effective_energy"]) |
| if total + dcost <= budget_params and dcost > 0 and denergy > 0: |
| gain = denergy / dcost |
| if gain > best_gain: |
| best_gain, best_move = gain, (i, j, dcost) |
| if best_move: |
| i, j, dcost = best_move |
| choice[i] = j |
| total += dcost |
| improved = True |
| plan = {} |
| for i, leaf in enumerate(leaves): |
| plan[leaf.name] = srt[i][choice[i]] |
| plan["_total_cost_params"] = total |
| plan["_budget_params"] = float(budget_params) |
| plan["_fidelity_lambda"] = float(lam) |
| return plan |
|
|
|
|
| def utility_breakdown(savings_params: int, penalties: dict) -> dict: |
| """utility = savings - recon - logit - hidden - task penalties. |
| All inputs measured; returns the itemized breakdown.""" |
| u = (savings_params - penalties.get("reconstruction", 0.0) |
| - penalties.get("logit", 0.0) - penalties.get("hidden", 0.0) |
| - penalties.get("task", 0.0)) |
| return {"utility": float(u), "savings_params": int(savings_params), |
| **{k: float(v) for k, v in penalties.items()}} |
|
|
