expose pmass_format as aggregate signal

ForcedChoiceResult now carries pmass_format (sum prob mass on the K
foundation answer tokens at the JSON answer slot, averaged across fwd
and rev framings). eval.py aggregates it as mean_pmass_format in both
the headline return dict and the info subdict, and propagates per-row
for sweep/audit consumers.

Direct coherence canary: drops when steering pushes the model toward
non-foundation tokens (gibberish, refusal, format collapse). Independent
of which foundation is picked — complementary to top1_acc (label-
agreement; intentional target shift) and mean_nll_prompt (teacher-forced
prompt nll; falls under steering even when generations break).

Surfacing this lets downstream callers (weight-steering-lite walkback,
report dashboards) gate on actual coherence rather than misusing top1
as a budget.
This commit is contained in:
wassname
2026-05-18 11:26:00 +00:00
parent f9a490c71d
commit c4797cd732
2 changed files with 28 additions and 0 deletions
+13
View File
@@ -188,6 +188,7 @@ def evaluate(
"top1": res.top1,
"margin": res.margin,
"nll_prompt": res.nll_prompt,
"pmass_format": res.pmass_format,
})
pbar.update(len(chunk))
@@ -271,6 +272,10 @@ def evaluate(
T = None
profile = None
mean_pmass_format = (
float(np.mean([r["pmass_format"] for r in per_row]))
if per_row else None
)
info = {
"name": name,
"n_rows": n_rows,
@@ -286,6 +291,13 @@ def evaluate(
# 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,
# Mean pmass_format: 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,
}
out: dict[str, Any] = {
@@ -297,6 +309,7 @@ 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,
"info": info,
}
if return_per_row:
+15
View File
@@ -331,6 +331,13 @@ class ForcedChoiceResult:
# signal: rises when the model is perturbed (steering, ablation, etc.)
# to a state where ordinary prompt text becomes "surprising".
nll_prompt: float
# Sum of probability mass over the K foundation answer-tokens at the
# JSON answer slot, averaged across fwd + rev framings. In [0, 1]; high
# means the model still emits a valid foundation word in the slot;
# low means probability has leaked to other tokens (gibberish, refusal,
# format collapse). The direct coherence canary for forced-choice
# — independent of WHICH foundation is picked.
pmass_format: float
def _resolve_first_token_ids(tok, words: list[str]) -> tuple[list[int], dict[str, int]]:
@@ -444,6 +451,13 @@ def guided_rollout_forced_choice(
# 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])
# Average pmass_format across framings: coherence canary independent
# of WHICH foundation is picked. Sum prob mass over the K answer
# tokens at the JSON slot; drops when model emits non-foundation
# tokens (gibberish, refusal, format collapse).
pm_f = slots_fwd[i][0]["pmass_format"]
pm_r = slots_rev[i][0]["pmass_format"]
pm = 0.5 * (pm_f + pm_r)
results.append(ForcedChoiceResult(
user_prompt=user_prompts[i],
think_text=think_fwd,
@@ -457,6 +471,7 @@ def guided_rollout_forced_choice(
think_tokens=n_fwd,
emitted_close=close_fwd,
nll_prompt=float(nll_p),
pmass_format=float(pm),
))
return results