diff --git a/docs/spec/20260501_n_token_eval.md b/docs/spec/20260501_n_token_eval.md deleted file mode 100644 index bc96d4b..0000000 --- a/docs/spec/20260501_n_token_eval.md +++ /dev/null @@ -1,41 +0,0 @@ -# N-token Evaluation - -## Goal -Make eval score after a short deterministic continuation instead of only from the immediate next token, so steering effects that emerge over a few tokens show up in the metric. - -## Scope -In: core eval scoring, evaluate API, eval CLI, smoke verification. -Out: changing vignette data, changing report aggregation, adding stochastic rollouts. - -## Requirements -- R1: Eval must support a fixed continuation budget before scoring `true`/`false`. Done means: `evaluate(..., think_tokens=N)` uses a rollout path and records `think_tokens=N`. VERIFY: a smoke run with `N>0` reports `mode=guided_rollout` and different headline scores than `N=0` on the same subset. -- R2: Default behavior must stay available for zero-token scoring. Done means: `think_tokens=0` still uses the current teacher-forced path. VERIFY: smoke run reports `mode=next_token` with `think_tokens=0`. -- R3: CLI must expose the token budget. Done means: `scripts/03_eval.py --think-tokens N` works and writes the budget into output JSON. VERIFY: output JSON contains `think_tokens` and `eval_mode`. - -## Tasks -- [/] T1 (R1,R2,R3): patch core scoring and CLI - - steps: add rollout scorer, thread config through `evaluate`, add CLI arg and metadata - - verify: run eval twice on a tiny subset, once with `--think-tokens 0` and once with `--think-tokens 8` - - success: metadata differs by mode and at least one headline score differs - - likely_fail: flag is accepted but dead code still uses next-token path, scores identical and metadata unchanged - - sneaky_fail: rollout happens but scoring still reads the old prompt position, metadata changes but scores remain indistinguishable from teacher-forced - - UAT: when I run the two eval commands, I see mode + think_tokens in the output JSON and a score delta -- [ ] T2: fresh-eyes review - - steps: hand diff and smoke evidence to subagent - - verify: reviewer explicitly checks likely and sneaky failure modes - - success: reviewer says evidence distinguishes rollout path from dead code path - -## Context -- Current eval scores only `out.logits[:, -1]` from the prompt ending at the JSON prefill. -- The gist proposes a deterministic guided rollout: generate N tokens, append a fixed answer prefill, then score at the forced answer position. -- For this repo we do not need `` tags. We only need a short continuation budget before the JSON boolean answer. - -## Log -- Guided rollout is the right abstraction here, but the minimal repo change is simpler than the gist: continue the assistant reply for N tokens, then append the existing boolean prefill and score there. - -## TODO -- Consider logging a sample guided continuation when bool mass collapses, but keep that out of this patch unless needed. - -## Errors -| Task | Error | Resolution | -|------|-------|------------| \ No newline at end of file diff --git a/scripts/09_forced_choice.py b/scripts/09_forced_choice.py index 6af7a5d..6be90f8 100644 --- a/scripts/09_forced_choice.py +++ b/scripts/09_forced_choice.py @@ -83,6 +83,7 @@ def main() -> None: else {f: float(r["label"][i]) for i, f in enumerate(_DEFAULT_FORCED_FOUNDATIONS)}), "top1": r["top1"], "margin": float(r["margin"]), + "nll_prompt": float(r["nll_prompt"]), } f.write(json.dumps(rec) + "\n") logger.info(f"wrote {len(out['per_row'])} rows to {out_path}") @@ -104,6 +105,14 @@ def main() -> None: f"{np.median(p_top1):.3f} / {p_top1.mean():.3f} / {p_top1.max():.3f}") print(" SHOULD: median > 0.4 (clear winner per row); <0.2 -> probe broken") + # Prompt-NLL degradation probe (free; teacher-forced on rendered chat). + nll = np.array([float(r["nll_prompt"]) for r in out["per_row"]]) + nll = nll[np.isfinite(nll)] + if len(nll): + print(f"\n nll_prompt (nats/tok) min/median/mean/max: " + f"{nll.min():.3f} / {np.median(nll):.3f} / {nll.mean():.3f} / {nll.max():.3f}") + print(" SHOULD: stable across runs at fixed model; rises under steering/ablation -> degradation") + if __name__ == "__main__": main() diff --git a/src/tinymfv/eval.py b/src/tinymfv/eval.py index bd4952c..5ffd4fb 100644 --- a/src/tinymfv/eval.py +++ b/src/tinymfv/eval.py @@ -154,6 +154,7 @@ def evaluate( "label": label, # may be None on unlabeled rows "top1": res.top1, "margin": res.margin, + "nll_prompt": res.nll_prompt, }) pbar.update(len(chunk)) @@ -211,6 +212,11 @@ def evaluate( "elapsed_s": elapsed, "median_js": median_js, "max_js": math.log(2), + # Mean prompt NLL across rows (nats/token, teacher-forcing on the + # rendered chat). Free degradation probe; unsteered baseline gives + # the model's "natural" surprise on prompt text. + "mean_nll_prompt": float(np.mean([r["nll_prompt"] for r in per_row])) + if per_row else None, } out: dict[str, Any] = { diff --git a/src/tinymfv/guided.py b/src/tinymfv/guided.py index 911bfa2..6f52871 100644 --- a/src/tinymfv/guided.py +++ b/src/tinymfv/guided.py @@ -73,10 +73,17 @@ def _rollout_kv_fork( choice_token_ids: list, # [a_ids, b_ids] verbose: bool = False, gather_token_ids: list[int] | None = None, -) -> tuple[list[tuple[str, int, bool]], list[list[dict]]]: - """Returns (thinks, slots). +) -> tuple[list[tuple[str, int, bool]], list[list[dict]], list[float]]: + """Returns (thinks, slots, nll_prompts). thinks[i] = (think_text, n_think_tokens, emitted_close) slots[i][j] = {pmass_format, logratio, p_true, [lp_gather]} + nll_prompts[i] = mean teacher-forcing NLL (nats/token) on the user-side chat + tokens of row i (excluding the very first chat token, which has no + in-prompt context). Computed from the Phase-2a prefix forward, so it's + free vs the existing rollout. Use as a coherence/degradation probe: + a steered or perturbed model that breaks normal text representation + will see nll_prompt rise (the model is "more surprised" by ordinary + prompt text). If `gather_token_ids` is provided, slot dict also has `lp_gather`: list[float] of log-probs at last suffix position for those token ids (in the same order). @@ -110,8 +117,15 @@ def _rollout_kv_fork( ) # === Build per-row scoring prefix: chat + think + === + # Also retokenise the chat alone so we can locate the user-prompt span + # inside sp_ids for free prompt-NLL below. We assert prefix-equality so + # tokenizer boundary merges (rare, but possible across `\n` → + # think-text) crash loudly rather than silently misaligning the NLL window. sp_per_row: list[str] = [] sp_ids_per_row: list[list[int]] = [] + chat_ids_per_row: list[list[int]] = [ + tok(c, add_special_tokens=False)["input_ids"] for c in chats + ] thinks: list[tuple[str, int, bool]] = [] for i, p in enumerate(chats): gen_ids = phase1[i, prompt_len:] @@ -141,6 +155,31 @@ def _rollout_kv_fork( pref_out = model(input_ids=pref_input, attention_mask=pref_attn, use_cache=True) cache = pref_out.past_key_values + # === Per-row prompt NLL (free; reuses pref_out.logits) === + # Teacher-forcing on the chat tokens (= rendered user turn + generation + # prompt). We skip the first chat token because position-(start-1) is a + # left-pad slot, so its prediction is meaningless. Per-row slicing keeps + # peak memory low (avoids a full [B, P_max, V] log_softmax). + nll_prompts: list[float] = [] + for i in range(B): + L_pref = int(pref_real[i].item()) + chat_ids_i = chat_ids_per_row[i] + chat_len = len(chat_ids_i) + assert sp_ids_per_row[i][:chat_len] == chat_ids_i, ( + f"chat retokenisation diverged from prefix at row {i}: " + f"boundary tokens shifted between chat and chat+think+close. " + f"prompt-NLL window cannot be located safely." + ) + if chat_len < 2: + nll_prompts.append(float("nan")) + continue + start = P_max - L_pref + logits_slice = pref_out.logits[i, start : start + chat_len - 1].float() + targets = pref_input[i, start + 1 : start + chat_len] + logp = F.log_softmax(logits_slice, dim=-1) + nll = -logp.gather(-1, targets[:, None]).squeeze(-1).mean() + nll_prompts.append(float(nll.item())) + # === Phase 2b: per-slot KV-fork === a_ids, b_ids = _split_choice_ids(choice_token_ids) a_t = torch.tensor(a_ids, device=device, dtype=torch.long) if a_ids else None @@ -222,7 +261,7 @@ def _rollout_kv_fork( d["lp_gather"] = lp_last[i, gid_t].cpu().tolist() slots[i].append(d) - return thinks, slots + return thinks, slots, nll_prompts # ===== Forced-choice (K-way primary foundation) ===== @@ -297,6 +336,11 @@ class ForcedChoiceResult: margin: float # score[top1] - score[top2], in nats think_tokens: int emitted_close: bool + # Mean teacher-forcing NLL (nats/token) on the user-side chat tokens, + # averaged across the fwd and rev framings. Free coherence/degradation + # signal: rises when the model is perturbed (steering, ablation, etc.) + # to a state where ordinary prompt text becomes "surprising". + nll_prompt: float def _resolve_first_token_ids(tok, words: list[str]) -> tuple[list[int], dict[str, int]]: @@ -374,7 +418,7 @@ def guided_rollout_forced_choice( scoring_slot = [(nudge, prefill)] # Frame A: forward enum order. - thinks_fwd, slots_fwd = _rollout_kv_fork( + thinks_fwd, slots_fwd, nll_fwd = _rollout_kv_fork( model, tok, user_prompts, schema_fwd, max_think_tokens, scoring_slots=scoring_slot, choice_token_ids=[[first_ids[0]]], # unused; satisfies API @@ -383,7 +427,7 @@ def guided_rollout_forced_choice( ) # Frame B: reversed enum order. Same gather order (by foundation name) so # lp_rev[f] is comparable to lp_fwd[f]. - thinks_rev, slots_rev = _rollout_kv_fork( + thinks_rev, slots_rev, nll_rev = _rollout_kv_fork( model, tok, user_prompts, schema_rev, max_think_tokens, scoring_slots=scoring_slot, choice_token_ids=[[first_ids[0]]], @@ -407,6 +451,9 @@ def guided_rollout_forced_choice( order_sorted = sorted(range(K), key=lambda k: -score[k]) top1 = foundations[order_sorted[0]] margin = score[order_sorted[0]] - score[order_sorted[1]] + # Average prompt NLL across the two framings (schema_hint differs in + # enum order but the user vignette is identical). + nll_p = 0.5 * (nll_fwd[i] + nll_rev[i]) results.append(ForcedChoiceResult( user_prompt=user_prompts[i], think_text=think_fwd, @@ -419,6 +466,7 @@ def guided_rollout_forced_choice( margin=float(margin), think_tokens=n_fwd, emitted_close=close_fwd, + nll_prompt=float(nll_p), )) return results