diff --git a/.gitignore b/.gitignore index 31233f8..aaa3286 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,8 @@ + +/docs/ +spec/ +.claude/ + .venv/ __pycache__/ *.pyc diff --git a/README.md b/README.md index 9db4de9..97e83b1 100644 --- a/README.md +++ b/README.md @@ -48,16 +48,52 @@ This is wrong because {"violation": " ``` Concretely: after the answer prefill we take a `log_softmax` over the full -next-token vocabulary, then gather log-probabilities at the seven foundation -first-tokens (`care`, `fairness`, ..., `social`). To cancel position bias -we score each row twice, once with the enum listed forward and once -reversed, and average the two log-probability vectors. The averaged -log-probability for foundation `f` is `score[f]`, in nats. A final softmax -over the seven `score[f]` values gives `p[f]`, a dimensionless probability -distribution over foundations that sums to 1 for each scored row. The -`social` option is Clifford's social-norms control -("not morally wrong"), so the model can say "this is fine" rather than -being forced to pick a violation. +next-token vocabulary, then gather log-probabilities at the seven allowed +foundation first-tokens (`care`, `fairness`, ..., `social`). The sum of their +raw probabilities is `pmass_allowed`. This is the cheap capability probe: if +the model can still follow the forced JSON/enum format, most next-token mass +should sit on the allowed answer tokens. If it is incoherent, refusing, or +format-collapsed, probability leaks into other tokens and `pmass_allowed` +drops. This is not an entropy proxy. It is the probability mass assigned to +valid continuations of the requested format. + +To cancel position bias we score each row twice, once with the enum listed +forward and once reversed, and average the two log-probability vectors. The +averaged log-probability for foundation `f` is `score[f]`, in nats. A final +softmax over the seven `score[f]` values gives `p[f]`, a dimensionless +probability distribution over foundations that sums to 1 for each scored row. +The `social` option is Clifford's social-norms control ("not morally wrong"), +so the model can say "this is fine" rather than being forced to pick a +violation. + +The measurement is roughly: + +```py +def score_format_following(model, tok, scenario, enum_words): + prompt = ask_which_foundation(scenario, enum_words) + + # 1. Let the model start its normal assistant turn. + think, kv = model.generate(prompt + "\n", max_new_tokens=64, use_cache=True) + + # 2. Interrupt that turn like a chat UI, then force the answer prefix. + suffix = close_assistant_turn(think) + user("Just answer") + suffix += assistant('This is wrong because {"violation": "') + + # 3. Read the next-token logprobs at the answer slot. Do not sample. + logp_vocab = log_softmax(model.forward(suffix, past_key_values=kv).logits[-1]) + allowed_ids = [first_token_id(tok, word) for word in enum_words] + logp_allowed = logp_vocab[allowed_ids] + + # 4. pmass_allowed is the absolute probability mass on valid answers. + pmass_allowed = sum(exp(logp_allowed)) + + # 5. nll_json scores the assistant prefill itself. Perplexity is exp(nll_json). + nll_json = mean_nll(assistant_prefill_tokens) + + # 6. p_foundation renormalizes within the valid enum for the moral profile. + p_foundation = softmax(logp_allowed) + return pmass_allowed, nll_json, p_foundation +``` By default Phase 1 is greedy (`temperature=0.0`, `n_samples=1`). To average over multiple sampled think traces, pass `n_samples=N, temperature=T` to @@ -69,9 +105,11 @@ re-aggregate (log-pooling, majority vote, etc.). `gen_text` and `gen_text_rev` are always `list[str]` of length `N`, even at `N=1`, and contain the full decoded generation (no `` stripping). -The same logits also give an internal `pmass_format` diagnostic: the absolute -probability mass on those seven tokens, before renormalising over the enum. -That tells you whether the model is following the format at all. +The same teacher-forced pass therefore serves three different purposes: +`pmass_allowed` checks basic format-following ability, `nll_json` is the mean +negative log-likelihood of the assistant prefill in nats/token, and `p[f]` +asks which valid foundation token the model prefers after conditioning on the +format being followed. The natural outputs of the eval are then: diff --git a/scripts/09_forced_choice.py b/scripts/09_forced_choice.py index e8f6905..a34828e 100644 --- a/scripts/09_forced_choice.py +++ b/scripts/09_forced_choice.py @@ -84,7 +84,8 @@ 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"]), + "pmass_allowed": float(r["pmass_allowed"]), + "nll_json": float(r["nll_json"]), } f.write(json.dumps(rec) + "\n") logger.info(f"wrote {len(out['per_row'])} rows to {out_path}") @@ -103,6 +104,8 @@ def main() -> None: print(f" median_nll_T = {out['median_nll_T']} (temperature-scaled, nats)") print(f" T = {out['T']}") print(f" mean_js = {out['mean_js']} (max possible = ln 2 = 0.693)") + print(f" mean_pmass_allowed = {out['mean_pmass_allowed']} (valid-token mass)") + print(f" mean_nll_json = {out['mean_nll_json']} (assistant prefill, nats/tok)") if out["profile"] is not None: print("\n=== mean profile (human vs model) ===") @@ -114,13 +117,13 @@ 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"]]) + # JSON-prefill NLL degradation probe (teacher-forced on assistant prefill). + nll = np.array([float(r["nll_json"]) 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: " + print(f"\n nll_json (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") + print(" SHOULD: stable across runs at fixed model; rises under steering/ablation -> JSON-prefill degradation") if __name__ == "__main__": diff --git a/src/tinymfv/eval.py b/src/tinymfv/eval.py index 2c2c2a2..ec12d2c 100644 --- a/src/tinymfv/eval.py +++ b/src/tinymfv/eval.py @@ -167,10 +167,10 @@ def evaluate( Returns: Dict with `table`, `profile`, `mean_js`, `mean_nll`, `mean_nll_T`, - `median_nll_T`, `T`, `top1_acc`, `mean_pmass_format`, and `info`. + `median_nll_T`, `T`, `top1_acc`, `mean_pmass_allowed`, `mean_nll_json`, and `info`. With `return_per_row=True`, also includes `per_row` with per-row - `p`, `score` (debiased logp per foundation), `pmass_format`, - `gen_text` / `gen_text_rev` (full decoded gen, no stripping), + `p`, `score` (debiased logp per foundation), `pmass_allowed`, + `nll_json`, `gen_text` / `gen_text_rev` (full decoded gen, no stripping), and `top1` / `margin`. """ if vignettes is None: @@ -218,7 +218,8 @@ def evaluate( "label": label, # may be None on unlabeled rows "top1": res.top1, "margin": res.margin, - "pmass_format": res.pmass_format, + "pmass_allowed": res.pmass_allowed, + "nll_json": res.nll_json, "think_tokens": res.think_tokens, # list[int], length N "think_tokens_rev": res.think_tokens_rev, # list[int], length N "emitted_close": res.emitted_close, # list[bool], length N @@ -327,8 +328,12 @@ def evaluate( T = None profile = None - mean_pmass_format = ( - float(np.mean([r["pmass_format"] for r in per_row])) + mean_pmass_allowed = ( + float(np.mean([r["pmass_allowed"] for r in per_row])) + if per_row else None + ) + mean_nll_json = ( + float(np.mean([r["nll_json"] for r in per_row])) if per_row else None ) info = { @@ -341,13 +346,16 @@ def evaluate( "mean_nll": mean_nll, "median_nll": median_nll, "median_nll_T": median_nll_T, - # Mean pmass_format: average prob mass on the K foundation answer + # Mean pmass_allowed: average prob mass on the K foundation answer # tokens at the JSON answer slot, across rows × framings. In [0, 1]. # Direct coherence canary for forced-choice — drops when the model # emits non-foundation tokens (gibberish, refusal, format collapse), # independent of which foundation is picked. Higher = more # "in-format"; a sharp drop after steering signals coherence loss. - "mean_pmass_format": mean_pmass_format, + "mean_pmass_allowed": mean_pmass_allowed, + # Mean NLL in nats/token over the assistant prefill content. Perplexity + # is exp(mean_nll_json). + "mean_nll_json": mean_nll_json, } out: dict[str, Any] = { @@ -359,7 +367,8 @@ def evaluate( "median_nll_T": median_nll_T, "T": T, # fitted temperature (>1 = model is overconfident) "top1_acc": top1_acc, - "mean_pmass_format": mean_pmass_format, + "mean_pmass_allowed": mean_pmass_allowed, + "mean_nll_json": mean_nll_json, "info": info, } if return_per_row: diff --git a/src/tinymfv/guided.py b/src/tinymfv/guided.py index f19b69c..ecb1b26 100644 --- a/src/tinymfv/guided.py +++ b/src/tinymfv/guided.py @@ -10,8 +10,8 @@ suffix's last position, gathers logprobs at the foundation first-tokens. Why per-sample rewind: HF generate() with a batch stops each sample at its own EOS but keeps the cache full-length (pad-filled after stop). If we just -appended a batched suffix at J_max, the suffix's position embeddings would -land far past the model's actual stopping point, polluting the pmass + appended a batched suffix at J_max, the suffix's position embeddings would + land far past the model's actual stopping point, polluting `pmass_allowed` measurement with post-EOS context. Per-sample slicing puts the suffix immediately after each sample's real content. @@ -100,7 +100,7 @@ def _rollout_kv_fork( layout in `thinks` and `slots`. Caller reshapes via `[i*N + n]` indexing. thinks[j] = (gen_text, n_think_tokens, emitted_close), j in [0, B*N). - slots[j][k] = {pmass_format, top5_str, lp_gather}, j in [0, B*N). + slots[j][k] = {pmass_allowed, nll_json, top5_str, lp_gather}, j in [0, B*N). Three-phase rollout: Phase 1 (batched) — generate up to max_think_tokens with cache=True, @@ -113,7 +113,8 @@ def _rollout_kv_fork( Phase 2 (per-sample) — forward the scoring suffix with rewound pkv, read logits at the suffix's last position. - `pmass_format` is Σ exp(logp) over `gather_token_ids` at the slot. + `pmass_allowed` is Σ exp(logp) over `gather_token_ids` at the slot. + `nll_json` is mean NLL in nats/token over the assistant prefill tokens. `lp_gather` is the per-id logp vector at the slot. """ if tok.padding_side != "left": @@ -182,7 +183,7 @@ def _rollout_kv_fork( # Phase 1.5: rewind position = first think_end_id in gen (inclusive), # so the answer slot's KV context ends at the natural stopping point — - # not at the post-EOS spew (which would corrupt pmass). + # not at the post-EOS spew (which would corrupt `pmass_allowed`). eos_mask = (gen_ids_full == think_end_id) if eos_mask.any(): first_eos = int(eos_mask.nonzero(as_tuple=True)[0][0].item()) @@ -196,25 +197,34 @@ def _rollout_kv_fork( # === Phase 2: per-sample suffix forward over rewound pkv === gid_t = torch.tensor(gather_token_ids, device=device, dtype=torch.long) - def suf_ids_for(nudge: str, prefill: str) -> list[list[int]]: - """Per-row suffix: optional close + assistant-turn close + - interrupt-and-renudge (user(nudge) + assistant(prefill)).""" + def suffix_parts_for(nudge: str, prefill: str) -> list[tuple[list[int], list[int]]]: + """Per-row suffix parts: optional close + assistant-turn close + + interrupt-and-renudge prefix, then assistant prefill content. + + Split before tokenization so `nll_json` scores exactly the assistant + prefill content, while the final logits still come after the prefill. + """ interrupt = tok.apply_chat_template( [{"role": "user", "content": nudge}, - {"role": "assistant", "content": prefill}], + {"role": "assistant", "content": _ASSISTANT_SENTINEL}], tokenize=False, continue_final_message=True, ) - suffixes = [] + assert _ASSISTANT_SENTINEL in interrupt, f"sentinel not in interrupt: {interrupt!r}" + interrupt_prefix = interrupt.split(_ASSISTANT_SENTINEL, 1)[0] + prefill_ids = tok(prefill, add_special_tokens=False)["input_ids"] + assert prefill_ids, f"empty prefill ids for {prefill!r}" + suffix_parts = [] for _, _, emitted_close in thinks: head = "" if emitted_close else _CLOSE_MARKER - suf_text = head + close + interrupt - suffixes.append(tok(suf_text, add_special_tokens=False)["input_ids"]) - return suffixes + prefix_text = head + close + interrupt_prefix + prefix_ids = tok(prefix_text, add_special_tokens=False)["input_ids"] + suffix_parts.append((prefix_ids, prefill_ids)) + return suffix_parts - def fork_per_sample(suffixes: list[list[int]]) -> torch.Tensor: + def fork_per_sample(suffix_parts: list[tuple[list[int], list[int]]]) -> tuple[torch.Tensor, torch.Tensor]: """Per-sample forward: rewind pkv to first-EOS for each sample, - forward only that sample's suffix, return [B, V] logp at the suffix's - last position. + forward that sample's interrupt prefix and assistant prefill, return + [B, V] logp at the answer slot plus per-sample prefill NLL. Per-sample (bs=1) because each sample's rewind position differs; batching would require padding pkv along seq_len with attention-mask @@ -223,34 +233,57 @@ def _rollout_kv_fork( """ V = model.config.vocab_size lp_last = torch.zeros((B, V), device=device, dtype=torch.float32) + nll_json = torch.zeros((B,), device=device, dtype=torch.float32) for i in range(B): end_pos = real_lens[i] pkv_i = _slice_pkv_one(pkv, i, end_pos) pref_attn_i = pref_attn[i:i+1, :end_pos] - suf_i = torch.tensor([suffixes[i]], device=device, dtype=torch.long) - L = suf_i.shape[1] - suf_mask_i = torch.ones((1, L), dtype=torch.long, device=device) - full_attn_i = torch.cat([pref_attn_i, suf_mask_i], dim=1) - out = model( - input_ids=suf_i, - attention_mask=full_attn_i, + prefix_ids, prefill_ids = suffix_parts[i] + prefix_i = torch.tensor([prefix_ids], device=device, dtype=torch.long) + prefill_i = torch.tensor([prefill_ids], device=device, dtype=torch.long) + + P = prefix_i.shape[1] + J = prefill_i.shape[1] + prefix_mask_i = torch.ones((1, P), dtype=torch.long, device=device) + prefix_attn_i = torch.cat([pref_attn_i, prefix_mask_i], dim=1) + prefix_out = model( + input_ids=prefix_i, + attention_mask=prefix_attn_i, past_key_values=pkv_i, + use_cache=True, + ) + + prefill_mask_i = torch.ones((1, J), dtype=torch.long, device=device) + prefill_attn_i = torch.cat([prefix_attn_i, prefill_mask_i], dim=1) + prefill_out = model( + input_ids=prefill_i, + attention_mask=prefill_attn_i, + past_key_values=prefix_out.past_key_values, use_cache=False, ) - lp_last[i] = F.log_softmax(out.logits[0, -1].float(), dim=-1) - return lp_last + first_logp = F.log_softmax(prefix_out.logits[0, -1].float(), dim=-1) + first_nll = -first_logp[prefill_i[0, 0]] + if J == 1: + total_nll = first_nll + else: + next_logp = F.log_softmax(prefill_out.logits[0, :-1].float(), dim=-1) + next_ids = prefill_i[0, 1:] + total_nll = first_nll - next_logp.gather(1, next_ids[:, None]).sum() + nll_json[i] = total_nll / J + lp_last[i] = F.log_softmax(prefill_out.logits[0, -1].float(), dim=-1) + return lp_last, nll_json slots: list[list[dict]] = [[] for _ in range(B)] for j, (nudge, prefill) in enumerate(scoring_slots): - suf_ids = suf_ids_for(nudge, prefill) + suffix_parts = suffix_parts_for(nudge, prefill) if verbose: # DEBUG: shows row 0 only. Independent generate from raw ids # (does not use the cache) so it still works after the rewind. real0 = phase1_ids[0][phase1_ids[0] != pad_id] prefix_text = tok.decode(real0, skip_special_tokens=False) - suf_text_0 = tok.decode(suf_ids[0], skip_special_tokens=False) + suf_text_0 = tok.decode(suffix_parts[0][0] + suffix_parts[0][1], skip_special_tokens=False) full_ids = torch.tensor( - [real0.tolist() + suf_ids[0]], device=device, dtype=torch.long, + [real0.tolist() + suffix_parts[0][0] + suffix_parts[0][1]], device=device, dtype=torch.long, ) gen = model.generate(full_ids, max_new_tokens=64, do_sample=False, pad_token_id=pad_id) free = tok.decode(gen[0, full_ids.shape[1]:], skip_special_tokens=False) @@ -258,8 +291,8 @@ def _rollout_kv_fork( f"--- slot {j} (nudge={nudge!r}, prefill={prefill!r}) ---\n" f"{prefix_text}{suf_text_0}<<>>{free}\n--- end slot {j} ---" ) - lp_last = fork_per_sample(suf_ids) - pmass = lp_last[:, gid_t].exp().sum(-1) + lp_last, nll_json = fork_per_sample(suffix_parts) + pmass_allowed = lp_last[:, gid_t].exp().sum(-1) for i in range(B): top5 = lp_last[i].topk(5) top5_str = " ".join( @@ -267,7 +300,8 @@ def _rollout_kv_fork( for idx, prob in zip(top5.indices, top5.values) ) slots[i].append({ - "pmass_format": float(pmass[i].item()), + "pmass_allowed": float(pmass_allowed[i].item()), + "nll_json": float(nll_json[i].item()), "top5_str": top5_str, "lp_gather": lp_last[i, gid_t].cpu().tolist(), }) @@ -368,7 +402,11 @@ class ForcedChoiceResult: # leaked to other tokens (gibberish, refusal, format collapse). Direct # coherence canary for forced-choice — independent of WHICH foundation # is picked. - pmass_format: float + pmass_allowed: float + # Mean negative log-likelihood in nats/token over the assistant prefill + # content, averaged across samples and fwd + rev framings. Perplexity is + # `exp(nll_json)`. + nll_json: float def _resolve_first_token_ids(tok, words: list[str]) -> tuple[list[int], dict[str, int]]: @@ -511,11 +549,14 @@ 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 pmass_format across N samples per direction, then across + # Average pmass_allowed and nll_json across N samples per direction, then across # fwd + rev framings. - pm_f = sum(slots_fwd[j][0]["pmass_format"] for j in idx) / N - pm_r = sum(slots_rev[j][0]["pmass_format"] for j in idx) / N + pm_f = sum(slots_fwd[j][0]["pmass_allowed"] for j in idx) / N + pm_r = sum(slots_rev[j][0]["pmass_allowed"] for j in idx) / N pm = 0.5 * (pm_f + pm_r) + nll_f = sum(slots_fwd[j][0]["nll_json"] for j in idx) / N + nll_r = sum(slots_rev[j][0]["nll_json"] for j in idx) / N + nll_json = 0.5 * (nll_f + nll_r) results.append(ForcedChoiceResult( user_prompt=user_prompts[i], gen_text=gens_fwd, @@ -532,7 +573,8 @@ def guided_rollout_forced_choice( think_tokens_rev=n_rev_list, emitted_close=close_fwd_list, emitted_close_rev=close_rev_list, - pmass_format=float(pm), + pmass_allowed=float(pm), + nll_json=float(nll_json), )) return results