Model card: diagnostics of what Jedi relies on

README.md

@@ -38,7 +38,12 @@ picks a model *before* any answer exists (see [How Jedi differs from Shaktimaan]
 - Within 1.4 points of a *perfect* verifier (93.3%) on the same ladder.
 - **AUROC 0.92** at telling correct answers from wrong ones on unseen tasks; calibrated confidence (ECE 0.017).
 - One forward pass and **one output token** per check: ~56 ms at batch size 1, ~450 input tokens.
-- **Weak outside its training domains**: AUROC 0.72 on BIG-Bench Hard and 0.59 on CommonsenseQA. See *Limitations*.
+- **It judges how an answer looks, not whether its final value is right.** Swapping the final answer of a real
+  response barely moves its score. Much of its signal is surface cues (length, truncation, which model answered),
+  though it still beats a gate built only on those cues by 1-2 points on the ladder. See
+  [What Jedi actually picks up on](#what-jedi-actually-picks-up-on).
+- **Weak outside its training domains**: AUROC 0.72 on BIG-Bench Hard (where answer length alone scores 0.80) and
+  0.59 on CommonsenseQA. See *Limitations*.
 
 ## How it works
 
@@ -76,6 +81,44 @@ picks a model *before* any answer exists (see [How Jedi differs from Shaktimaan]
 4. **The single dial.** The threshold trades cost for accuracy. Raise it to escalate more and catch more wrong
    answers; lower it to save money. The top rung is always accepted, so every task gets an answer.
 
+### What Jedi actually picks up on
+The system prompt asks Jedi to look for broken reasoning, bad arithmetic, failing code, hedging and truncation.
+Training only rewards predicting the grade, though, so we tested what it learned (`diagnostics.py`,
+`diagnostics_ladder.py`).
+
+| check | in-domain test | unseen (BIG-Bench Hard) |
+|---|---|---|
+| Jedi AUROC | **0.919** | 0.721 |
+| surface-feature baseline (logistic regression on answer length, truncated, empty final answer, agent, source) | 0.871 | **0.826** |
+| answer length alone (shorter = more likely right) | 0.810 | 0.801 |
+| Jedi, truncated answers excluded | 0.879 | 0.682 |
+| Jedi ranks the right answer above a wrong one *to the same task* | 93.1% of 690 pairs | 84.6% of 1,168 pairs |
+
+**Final-answer swap test.** We took 1,583 real test answers (multiple choice and numeric, not truncated) and changed
+only the value after `Final answer:`. Right answers got a wrong value; wrong answers got the gold value.
+
+| | mean P(correct) before | after swap |
+|---|---|---|
+| right answers -> wrong final value (1,443) | 0.912 | 0.902 |
+| wrong answers -> gold final value (140) | 0.717 | 0.806 |
+
+Jedi barely notices when the final answer contradicts the reasoning above it, or when a wrong answer is given the
+correct final value. It scores the *body* of the response: its length, its shape, whether it ends cleanly, and the
+style of reasoning that tends to accompany wrong answers from each model. Hand-written short answers with a planted
+arithmetic slip are usually accepted.
+
+**On the ladder, Jedi still beats the surface-feature gate** (same 491 tasks, best settings):
+
+| gate | accuracy | relative cost |
+|---|---|---|
+| surface-feature baseline, threshold 0.05 | 89.8% | 0.224 |
+| surface-feature baseline, threshold 0.1 | 89.8% | 0.395 |
+| Jedi, threshold 0.05 | 91.0% | 0.244 |
+| Jedi, threshold 0.7 | 91.9% | 0.454 |
+
+So in its training domains, Jedi learned something real beyond length and truncation. That gain is 1-2 points of
+end-to-end accuracy. On unseen task families, the surface features transfer better than Jedi does.
+
 ## Quick start
 
 ### Score one answer
@@ -101,9 +144,12 @@ def p_correct(task, answer, agent="glm-5.3-flash", kind="num", finish="stop"):
     return 1 / (1 + math.exp(-(cal["a"] * gap + cal["b"])))
 
 task = "A shop sells pencils at 3 for $1.20. How much do 15 pencils cost, in dollars?"
-print(p_correct(task, "Each pencil costs 1.20 / 3 = 0.40. 15 x 0.40 = 6.00.\nFinal answer: 6"))
-print(p_correct(task, "15 pencils is 15 / 3 = 5 groups. 5 x 1.20 = 5.00.\nFinal answer: 5"))
+print(p_correct(task, "Each pencil costs 1.20 / 3 = 0.40. 15 x 0.40 = 6.00.\nFinal answer: 6"))     # 0.99 accept
+print(p_correct(task, "Let me think step by step. First, 3 pencils cost 1.20, so", finish="length"))  # 0.07 escalate
+print(p_correct(task, "15 / 3 = 5 groups. 5 x 1.20 = 5.00.\nFinal answer: 5"))                     # 0.98 - missed
 ```
+The third call shows the main limitation: a fluent answer with an arithmetic slip is accepted. Jedi is a learned
+filter over real model outputs, not a checker of the math.
 
 Always build the input with `common.judge_view(...)` and use `common.JEDI_SYSTEM` as the system prompt; Jedi was
 trained on exactly that format. `kind` is the task type it was trained with: `num`, `math`, `mc`
@@ -151,9 +197,10 @@ problem, getting good answers without paying for the biggest model every time, a
 
 **Why the switch.** Shaktimaan's main finding was that, on ordinary text and code tasks, *which* mid-size model
 solves a given task is largely sampling noise that cannot be predicted from the prompt. A router has to guess
-blind. Jedi waits until a cheap answer exists and judges that answer: unsupported reasoning, arithmetic that doesn't
-add up, truncation, hedging and code that won't pass its tests. Its errors are mostly answers that look right but
-aren't, which a router cannot see at all.
+blind. Jedi waits until a cheap answer exists and looks at that particular sample. Signs that a sample went wrong
+(rambling length, truncation, the reasoning style that comes with mistakes) only exist after it is written, so a
+router can never see them. Jedi does not re-derive the answer, so fluent answers that are wrong still get through
+(see [What Jedi actually picks up on](#what-jedi-actually-picks-up-on)).
 
 **Head to head**, on the same 491 held-out tasks with the same fresh agent samples:
 
@@ -261,8 +308,8 @@ Calibrated P(correct) on answers to tasks Jedi never saw. *False accept* = share
 | ARC-Challenge | 300 | 4.0% | 0.842 | 91.7% |
 | CommonsenseQA | 252 | 18.3% | 0.593 | 89.1% |
 
-Jedi is strongest where correctness can be checked from the answer itself: code, multi-step math and
-knowledge questions with worked reasoning. It is weakest where a wrong answer reads exactly like a right one,
+Jedi is strongest where wrong answers *look* different from right ones: multi-step math and code, where
+failures tend to be long, meandering or cut off. It is weakest where a wrong answer reads exactly like a right one,
 such as single-hop commonsense and easy science multiple choice.
 
 ### As the gate of the ladder
@@ -324,6 +371,9 @@ CommonsenseQA-heavy new-domain set, Jedi at 0.7 does not beat always using Qwen.
 single forward pass and batch well.
 
 ## Limitations
+- **Not a checker.** Jedi does not verify the final answer: changing it barely moves the score, and planted
+  arithmetic slips in otherwise fluent answers usually pass. It relies heavily on surface cues, and on unseen task
+  families a logistic regression on answer length, truncation and agent beats it (AUROC 0.83 vs 0.72).
 - **Domain-bound.** Verification skill transfers poorly to unseen task families (BIG-Bench Hard AUROC 0.72) and is
   near chance on CommonsenseQA (0.59). Measure it on your own traffic before relying on a threshold.
 - **A third of wrong in-domain answers get through at 0.5.** False-accept rates are highest for the strongest
@@ -350,6 +400,7 @@ single forward pass and batch well.
 | `train.py` | fine-tuning script (stage 3) |
 | `evaluate.py` | verifier metrics, ladder, baselines and threshold sweep (stage 4) |
 | `eval_results.json` | raw evaluation output behind every number on this card |
+| `diagnostics.py`, `diagnostics_ladder.py` | surface-feature baselines, answer-swap test, ladder with a surface-feature gate |
 
 The graded answer data is not published; `gen.py` and `build.py` regenerate it given endpoints for the three models.
 

diagnostics.py

@@ -0,0 +1,95 @@
+"""Does Jedi read the answer content, or surface cues? (1) shallow-feature baselines (2) final-answer swap test."""
+import json, math, random, re, sys
+import numpy as np, torch
+import common
+from datasets import load_from_disk
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+def roc_auc_score(y, p):
+    y, p = np.asarray(y, bool), np.asarray(p, float)
+    r = np.argsort(np.argsort(p, kind="mergesort")) + 1
+    return float((r[y].sum() - y.sum() * (y.sum() + 1) / 2) / (y.sum() * (~y).sum()))
+
+class LogisticRegression:
+    def fit(self, X, y):
+        X = torch.tensor(X, dtype=torch.float32); y = torch.tensor(np.asarray(y, float), dtype=torch.float32)
+        self.mu, self.sd = X.mean(0), X.std(0) + 1e-6
+        X = (X - self.mu) / self.sd
+        self.w = torch.zeros(X.shape[1], requires_grad=True); self.b = torch.zeros(1, requires_grad=True)
+        opt = torch.optim.LBFGS([self.w, self.b], max_iter=500)
+        def closure():
+            opt.zero_grad(); l = torch.nn.functional.binary_cross_entropy_with_logits(X @ self.w + self.b, y); l.backward(); return l
+        opt.step(closure); return self
+    def predict_proba(self, X):
+        X = (torch.tensor(X, dtype=torch.float32) - self.mu) / self.sd
+        p = torch.sigmoid(X @ self.w + self.b).detach().numpy(); return np.stack([1 - p, p], 1)
+random.seed(0)
+P = sys.argv[1] if len(sys.argv) > 1 else "jedi-0.5b"
+cal = json.load(open(P + "/calibration.json")); tok = AutoTokenizer.from_pretrained(P); tok.padding_side = "left"
+model = AutoModelForCausalLM.from_pretrained(P, dtype=torch.bfloat16).cuda().eval()
+ds = load_from_disk("data/jedi_judge")
+
+def score(views, bs=16):
+    out = np.zeros(len(views)); order = np.argsort([len(v) for v in views])
+    for i in range(0, len(views), bs):
+        idx = order[i:i + bs]
+        msgs = [[{"role": "system", "content": common.JEDI_SYSTEM}, {"role": "user", "content": views[j]}] for j in idx]
+        enc = tok.apply_chat_template(msgs, add_generation_prompt=True, return_tensors="pt", padding=True, return_dict=True).to("cuda")
+        with torch.no_grad(): l = model(**enc).logits[:, -1].float()
+        g = (l[:, cal["accept_token"]] - l[:, cal["escalate_token"]]).cpu().numpy()
+        out[idx] = 1 / (1 + np.exp(-(cal["a"] * g + cal["b"])))
+    return out
+
+def feats(d):
+    ag = sorted(set(ds["train"]["agent"])); src = sorted(set(ds["train"]["source"]) | set(ds["ood"]["source"]))
+    X = []
+    for r in d:
+        X.append([math.log1p(r["completion_tokens"]), r["finish"] == "length", not r["pred"].strip()]
+                 + [r["agent"] == a for a in ag] + [r["source"] == s for s in src])
+    return np.array(X, float)
+
+base = ds["train"].filter(lambda r: r["sample"] == 0)  # one sample per (task, agent): no oversampling dupes in fit
+lr = LogisticRegression().fit(feats(base), base["correct"])
+res = {}
+for sp in ("test", "ood"):
+    d = ds[sp]; y = np.array(d["correct"])
+    p = score(d["judge_input"])
+    trunc = np.array([f == "length" for f in d["finish"]])
+    res[sp] = dict(
+        jedi=roc_auc_score(y, p),
+        shallow_logreg=roc_auc_score(y, lr.predict_proba(feats(d))[:, 1]),
+        length_only=roc_auc_score(y, -np.array(d["completion_tokens"])),
+        jedi_not_truncated=roc_auc_score(y[~trunc], p[~trunc]),
+        truncated_share=float(trunc.mean()))
+    # within-task: tasks with both a right and a wrong answer -> does Jedi rank the right one higher?
+    by = {}
+    for i, r in enumerate(d):
+        by.setdefault(r["task_id"], []).append(i)
+    wins = tot = 0
+    for ids in by.values():
+        for i in ids:
+            for j in ids:
+                if y[i] and not y[j]:
+                    tot += 1; wins += (p[i] > p[j]) + 0.5 * (p[i] == p[j])
+    res[sp]["within_task_pairwise_acc"] = wins / tot; res[sp]["within_task_pairs"] = tot
+    if sp == "test":
+        # swap test: replace the 'Final answer:' value. wrong -> gold, right -> a wrong value.
+        orig, swapped, ys = [], [], []
+        for r in d:
+            if r["kind"] not in ("mc", "num") or r["finish"] == "length": continue
+            m = list(re.finditer(r"(Final answer\s*[:：]\s*)(.+)", r["judge_input"], flags=re.I))
+            if not m: continue
+            m = m[-1]; gold = r["gold"].strip()
+            if r["correct"]:
+                new = random.choice([c for c in "ABCDEFGHIJ"[:5] if c != gold.upper()]) if r["kind"] == "mc" else str(int(float(gold.replace(',', ''))) + 7) if re.fullmatch(r"-?[\d,]+(\.0+)?", gold) else None
+            else:
+                new = gold
+            if new is None: continue
+            v = r["judge_input"][:m.start(2)] + new + r["judge_input"][m.end(2):]
+            orig.append(r["judge_input"]); swapped.append(v); ys.append(r["correct"])
+        po, ps = score(orig), score(swapped); ys = np.array(ys)
+        res["swap"] = dict(n=len(ys), n_right_made_wrong=int(ys.sum()), n_wrong_made_right=int((~ys).sum()),
+            right_orig=float(po[ys].mean()), right_after_swap_to_wrong=float(ps[ys].mean()),
+            wrong_orig=float(po[~ys].mean()), wrong_after_swap_to_gold=float(ps[~ys].mean()),
+            auroc_on_swapped=roc_auc_score(~ys, ps))  # label of swapped = flipped correctness
+print(json.dumps(res, indent=1))

diagnostics_ladder.py

@@ -0,0 +1,79 @@
+"""Ladder (glm -> qwen -> nemotron-ultra) on the 491 held-out tasks, gated by Jedi vs by the surface-feature baseline."""
+import json, math, random, re, sys
+import numpy as np, torch
+import common
+from datasets import load_from_disk
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+def roc_auc_score(y, p):
+    y, p = np.asarray(y, bool), np.asarray(p, float)
+    r = np.argsort(np.argsort(p, kind="mergesort")) + 1
+    return float((r[y].sum() - y.sum() * (y.sum() + 1) / 2) / (y.sum() * (~y).sum()))
+
+class LogisticRegression:
+    def fit(self, X, y):
+        X = torch.tensor(X, dtype=torch.float32); y = torch.tensor(np.asarray(y, float), dtype=torch.float32)
+        self.mu, self.sd = X.mean(0), X.std(0) + 1e-6
+        X = (X - self.mu) / self.sd
+        self.w = torch.zeros(X.shape[1], requires_grad=True); self.b = torch.zeros(1, requires_grad=True)
+        opt = torch.optim.LBFGS([self.w, self.b], max_iter=500)
+        def closure():
+            opt.zero_grad(); l = torch.nn.functional.binary_cross_entropy_with_logits(X @ self.w + self.b, y); l.backward(); return l
+        opt.step(closure); return self
+    def predict_proba(self, X):
+        X = (torch.tensor(X, dtype=torch.float32) - self.mu) / self.sd
+        p = torch.sigmoid(X @ self.w + self.b).detach().numpy(); return np.stack([1 - p, p], 1)
+random.seed(0)
+P = sys.argv[1] if len(sys.argv) > 1 else "jedi-0.5b"
+cal = json.load(open(P + "/calibration.json")); tok = AutoTokenizer.from_pretrained(P); tok.padding_side = "left"
+model = AutoModelForCausalLM.from_pretrained(P, dtype=torch.bfloat16).cuda().eval()
+ds = load_from_disk("data/jedi_judge")
+
+def score(views, bs=16):
+    out = np.zeros(len(views)); order = np.argsort([len(v) for v in views])
+    for i in range(0, len(views), bs):
+        idx = order[i:i + bs]
+        msgs = [[{"role": "system", "content": common.JEDI_SYSTEM}, {"role": "user", "content": views[j]}] for j in idx]
+        enc = tok.apply_chat_template(msgs, add_generation_prompt=True, return_tensors="pt", padding=True, return_dict=True).to("cuda")
+        with torch.no_grad(): l = model(**enc).logits[:, -1].float()
+        g = (l[:, cal["accept_token"]] - l[:, cal["escalate_token"]]).cpu().numpy()
+        out[idx] = 1 / (1 + np.exp(-(cal["a"] * g + cal["b"])))
+    return out
+
+def feats(d):
+    ag = sorted(set(ds["train"]["agent"])); src = sorted(set(ds["train"]["source"]) | set(ds["ood"]["source"]))
+    X = []
+    for r in d:
+        X.append([math.log1p(r["completion_tokens"]), r["finish"] == "length", not r["pred"].strip()]
+                 + [r["agent"] == a for a in ag] + [r["source"] == s for s in src])
+    return np.array(X, float)
+
+base = ds["train"].filter(lambda r: r["sample"] == 0)  # one sample per (task, agent): no oversampling dupes in fit
+lr = LogisticRegression().fit(feats(base), base["correct"])
+
+from datasets import load_dataset
+from common import LADDER, POOL, STRONGEST
+CHEAP, MID = LADDER[0], LADDER[1]
+d = ds["test"]
+pj = score(d["judge_input"]); ps = lr.predict_proba(feats(d))[:, 1]
+R = {}
+for r, a, b in zip(d.select_columns(["task_id", "agent", "sample", "correct", "completion_tokens"]), pj, ps):
+    R[r["task_id"], r["agent"], r["sample"]] = dict(r, pj=float(a), ps=float(b))
+shak = {r["id"]: r for sp in ("train", "test") for r in load_dataset("podstack/Shaktimaan-routing", split=sp)}
+srcs = {"gsm8k", "math", "mmlupro", "arc", "mbpp", "humaneval"}
+tasks = sorted({t for (t, _, _) in R if t.split("-")[0] in srcs and (t, CHEAP, 0) in R and (t, MID, 0) in R and t in shak})
+def ultra(t):
+    x = json.loads(shak[t]["agent_results"])[STRONGEST]; return x["correct"], x["completion_tokens"]
+ucost = sum(POOL[STRONGEST]["cost"] * ultra(t)[1] for t in tasks)
+def ladder(key, tau):
+    c = cost = 0
+    for t in tasks:
+        for i, ag in enumerate(LADDER):
+            if i == 2:
+                ok, tk = ultra(t); cost += POOL[ag]["cost"] * tk; c += ok; break
+            a = R[t, ag, 0]; cost += POOL[ag]["cost"] * a["completion_tokens"]
+            if a[key] >= tau: c += a["correct"]; break
+    return round(c / len(tasks), 4), round(cost / ucost, 4)
+print("n", len(tasks))
+for tau in (0.05, 0.1, 0.2, 0.3, 0.4, 0.45, 0.5):
+    print(tau, "jedi", ladder("pj", tau), "shallow", ladder("ps", tau))