From 160d675c618920b15b957b674ff96bdace150739 Mon Sep 17 00:00:00 2001 From: wassname <1103714+wassname@users.noreply.github.com> Date: Sat, 11 Jul 2026 21:06:46 +0800 Subject: [PATCH] demo: JSON-object coherence probe replaces rigged forced-digit pmass The old rubric pmass was softmax mass on digit tokens at the hand-fed `{"ans": ` slot, so it was ~always 1 (the prefix forces a digit even from a fried model) -- a blind coherence guard. Replace with the users design: free-generate {"ans":N,"why":str,"2+2":M}, gate coherence on valid-JSON AND 2+2==4. GPU UATs: object stays valid at |C|<=0.5, breaks at |C|>=0.75 (task 22/23); the old slot never caught this. span_pmass (mean top-1 prob over the span) is kept as a within-coherent confidence read but NOT the coherence signal: a steer-fried model collapses into a confident degenerate loop, so span_pmass climbs back to ~0.97 while the object is broken (C=3.0). plot_sweep therefore colors by valid_frac, which cannot be fooled by confident garbage. Also: show_steers per-C cowsay now speaks steer-PROMOTED tokens (top of steered-baseline logits) instead of lens_topk at the last chat-prompt position, which only ever surfaced think-openers (Okay/Here/The) for every C. UAT: joy steer -> < joy . happy . Happy . happy > at C=0.3 (task 24). compute_slice stays the calibrated cross-layer lens readout. Co-Authored-By: Claudypoo <288921227+claudypoo@users.noreply.github.com> --- jsteer/demo.py | 157 ++++++++++++++++--------- scripts/scratch/smoke_sweep.py | 24 +++- scripts/scratch/sweep_artifact.py | 44 +++++++ scripts/scratch/uat_coherence_break.py | 45 +++++++ scripts/scratch/uat_promoted_cowsay.py | 30 +++++ 5 files changed, 239 insertions(+), 61 deletions(-) create mode 100644 scripts/scratch/sweep_artifact.py create mode 100644 scripts/scratch/uat_coherence_break.py create mode 100644 scripts/scratch/uat_promoted_cowsay.py diff --git a/jsteer/demo.py b/jsteer/demo.py index c48bfa5..f17ce4d 100644 --- a/jsteer/demo.py +++ b/jsteer/demo.py @@ -11,8 +11,10 @@ output is debuggable and nothing is parsed or reconstructed. from __future__ import annotations import dataclasses +import json import torch +from jlens.vis import _meaningful_token_mask from loguru import logger from steering_lite import Vector @@ -26,8 +28,8 @@ def chat_input(tok, user_msg: str, *, enable_thinking: bool = True) -> str: def _cthulhu_say(text: str) -> str: - """The j-space readout in a mini cowsay bubble -- Cthulhu speaks the tokens - the steered residual points to. Cosmetic; the tokens are the payload.""" + """A mini cowsay bubble -- Cthulhu speaks the tokens. Cosmetic; the tokens are + the payload.""" n = len(text) + 2 return ("\n".join([" " + "_" * n, f"< {text} >", " " + "-" * n, " \\", " ^(;,;)^"])) @@ -36,20 +38,35 @@ def _cthulhu_say(text: str) -> str: # think-then-answer rubric read: the demo's one-number sanity signal that steering # moved the target axis. Same mechanism as moral-maps guided.py (let the model # think, then read the logprobs at a JSON answer slot), reduced to a single scalar. -_ANS_FMT = (' Think it over, then answer with JSON {"ans": N} where N is a single' - ' digit from 0 (least) to 9 (most).') +# The object carries a trivial-arithmetic canary ("2+2") and a free-text field so a +# steer-degraded model has ROOM to break the object -- coherence is measured on the +# free-generated object, not on the forced digit slot (which is ~always a digit +# because the `{"ans": ` prefix makes one obvious). (Claude) +_ANS_FMT = (' Think it over, then answer with ONE line of JSON and nothing after it:' + ' {"ans": N, "why": "<=3 words", "2+2": M} where N is a single digit from' + ' 0 (least) to 9 (most) and M is the value of 2+2.') @torch.no_grad() def rubric_score(model, tok, rubric: str, *, max_new_tokens: int, seed: int, - do_sample: bool = False, temperature: float = 0.7) -> tuple[float, float]: - """Ask `rubric`, let the model think, then FORCE the answer slot `{"ans": ` and - read the logprob-weighted expected digit 0-9 there. Returns (expected, pmass). + do_sample: bool = False, temperature: float = 0.7) -> tuple[float, dict]: + """Ask `rubric`, let the model think, force the slot `{"ans": ` for a clean scalar + read, then FREE-GENERATE the rest of the JSON object as a coherence probe. + Returns (expected, coh). - expected = sum_d d * softmax(logit_d over the 10 digit tokens) -- a continuous - scalar from single-token logprobs (cleaner than parsing a multi-token float). - pmass = full-vocab softmax mass on the 10 digit tokens: a coherence guard, ~0 - means the slot isn't a digit (prefix/tokenizer mismatch), so distrust expected. + expected = sum_d d * softmax(logit_d over the 10 digit tokens) at the forced slot + -- a continuous scalar from single-token logprobs (cleaner than parsing a float). + + coh = {"valid", "chk_ok", "span_pmass"} measured on the free-generated object: + valid -- the object parses as JSON (a steer-fried model fails to close it), + chk_ok -- its "2+2" field == 4 (trivial-arithmetic canary), + span_pmass -- mean top-1 softmax prob over the generated span. It degrades in the + COHERENT regime (~0.95 -> 0.81 as the steer bites) but is NOT a + coherence measure on its own: a steer-fried model collapses into a + confident degenerate loop, so span_pmass climbs back toward ~1 on + repeated garbage (observed valid=False, span_pmass=0.97 at C=3.0). + Coherence therefore GATES on (valid and chk_ok); span_pmass is only + a within-coherent confidence read, trustworthy where valid is True. The rigorous K-way, position-debiased version is moral-maps guided.py; this is the demo's cheap readout, scored under whatever steering is active.""" prompt = chat_input(tok, rubric + _ANS_FMT) @@ -71,34 +88,49 @@ def rubric_score(model, tok, rubric: str, *, max_new_tokens: int, seed: int, ids = torch.tensor([tok(str(d), add_special_tokens=False).input_ids[0] for d in range(10)], device=logits.device) expected = float((logits[ids].softmax(0) * torch.arange(10., device=ids.device)).sum()) - pmass = float(logits.softmax(0)[ids].sum()) - return expected, pmass + + # free-generate the rest of the object; short cap so incoherence shows fast + gob = model.generate(**fenc, max_new_tokens=20, do_sample=False, + pad_token_id=tok.eos_token_id, + output_scores=True, return_dict_in_generate=True) + span_pmass = float(torch.stack([s[0].float().softmax(-1).max() + for s in gob.scores]).mean()) + body = '{"ans": ' + tok.decode(gob.sequences[0][fenc.input_ids.shape[1]:], + skip_special_tokens=True) + body = body[:body.index("}") + 1] if "}" in body else body # cut at first close + try: # invalid JSON IS the signal (fried model can't close it) + obj = json.loads(body) + valid, chk_ok = True, obj.get("2+2") == 4 + except json.JSONDecodeError: + valid, chk_ok = False, False + return expected, {"valid": valid, "chk_ok": chk_ok, "span_pmass": span_pmass} @torch.no_grad() def coherence_sweep(model, tok, vec, rubric: str, *, step: float = 0.1, - pmass_floor: float = 0.9, max_steps: int = 15, n_samples: int = 3, + max_steps: int = 15, n_samples: int = 3, temperature: float = 0.7, max_new_tokens: int = 512) -> list[dict]: - """Walk C outward from 0 in +/- directions, scoring the rubric each step, and STOP - a direction the step AFTER the answer slot goes incoherent (pmass 1, temperature=temperature) for s in range(n_samples)] - anss = torch.tensor([a for a, _ in pairs]) - pmass = float(torch.tensor([p for _, p in pairs]).mean()) + anss = torch.tensor([e for e, _ in pairs]) + span = float(torch.tensor([c["span_pmass"] for _, c in pairs]).mean()) + valid_frac = sum(c["valid"] and c["chk_ok"] for _, c in pairs) / len(pairs) return {"C": round(float(C), 3), "ans": float(anss.mean()), - "ans_std": float(anss.std(unbiased=False)), "pmass": pmass, - "coherent": pmass >= pmass_floor} + "ans_std": float(anss.std(unbiased=False)), "span_pmass": span, + "valid_frac": valid_frac, "coherent": valid_frac >= 0.5} rows = [score(0.0)] for d in (step, -step): # outward each way; keep the 1st incoherent point C = d @@ -112,23 +144,24 @@ def coherence_sweep(model, tok, vec, rubric: str, *, step: float = 0.1, return rows -def plot_sweep(rows: list[dict], *, title: str = "rubric ans vs C", - pmass_floor: float = 0.9): - """ans vs C, points colored by answer coherence (pmass); incoherent points - (pmass None: - """One block per C: lens readout at `layer`, then the raw generation, all - under steering. Uses the model's own generation_config sampling; `seed` - fixes it so the C blocks are comparable. `layer` defaults to the top fitted - layer. max_new_tokens defaults to 512 so Qwen3's block can close; - 256 truncates mid-reasoning. + """Per C: the steer-promoted tokens in a cowsay bubble, then the raw generation, + all under steering. Uses the model's own generation_config sampling; `seed` fixes + it so the C blocks are comparable. max_new_tokens defaults to 512 so Qwen3's + block can close; 256 truncates mid-reasoning. The cowsay speaks the top of + (steered - unsteered) next-token logits -- what THIS C pushes up, with the shared + think-opener prior subtracted out (the old lens_topk-at-last-position surfaced only + Okay/Here/The for every C; the calibrated cross-layer lens readout is compute_slice + on a completion prompt). `jac` is unused here, kept for call-site stability. Extraction is decoupled from DELIVERY (see applies.py): pass `apply_mode` (add | clamp | add_last | replace_last) to swap how v hits the residual @@ -198,13 +232,13 @@ def show_steer(jac: Jacobian, model, tok, vec, user_msg: str, *, VALUE, add scales a direction), so each mode wants its own Cs. Pass `rubric` (a 0-9 rating question about the steered axis) to add the - quantitative readout: per C, the model thinks then answers `{"ans": N}` and we - report the logprob-weighted expected digit. It SHOULD rise with +C and fall - with -C; flat means the steer isn't moving that axis (see rubric_score).""" + quantitative readout: per C, the model thinks then answers a JSON object and we + report the logprob-weighted expected digit plus a coherence gate (valid object, + 2+2==4). ans SHOULD rise with +C and fall with -C; flat means the steer isn't + moving that axis; json=False means the steer broke the model (see rubric_score).""" if apply_mode is not None: vec = Vector(dataclasses.replace(vec.cfg, apply_mode=apply_mode, apply_span=apply_span), vec.shared, vec.stacked) - layer = jac.layers[-1] if layer is None else layer prompt = chat_input(tok, user_msg) enc = tok(prompt, return_tensors="pt").to(model.device) name = getattr(model.config, "name_or_path", "model").split("/")[-1] @@ -216,24 +250,37 @@ def show_steer(jac: Jacobian, model, tok, vec, user_msg: str, *, rule = "=" * 72 logger.info(f"\n\n{rule}\n{name} · method={vec.cfg.method}{tag}" f"\nprompt: {user_msg!r}\n{rule}") - # SHOULD: C=0 is the baseline; +C tilts the lens tokens and tone toward the + base = model(**enc).logits[0, -1].float() # unsteered next-token logits; the steer's + # effect reads as the top of (steered-base) + # SHOULD: C=0 is the baseline; +C tilts the promoted tokens and tone toward the # concept, -C away; all stay coherent (gibberish = coeff too large). for C in Cs: torch.manual_seed(seed) with vec(model, C=C): - jtop = jac.lens_topk(model, tok, prompt, layer=layer, k=k) + steered = model(**enc).logits[0, -1].float() out = model.generate(**enc, max_new_tokens=max_new_tokens, pad_token_id=tok.eos_token_id) ans = (rubric_score(model, tok, rubric, max_new_tokens=max_new_tokens, seed=seed) if rubric is not None else None) + # steer-promoted tokens: top of (steered - base), word-like only. The subtraction + # cancels the shared "about to open " prior (Okay/Here/The) so what the + # cowsay speaks is what THIS C actually pushes up, not the reasoning boilerplate + # the old lens_topk-at-last-position surfaced for every C. (Claude) + if C == 0: + readout = "(baseline, no steer)" + else: + wl = _meaningful_token_mask(tok, steered.shape[-1], steered.device) + promoted = (steered - base).masked_fill(~wl, float("-inf")).topk(6) + readout = " · ".join(tok.decode([i]).strip() for i in promoted.indices.tolist()) # raw decode WITH special tokens: real /, <|im_end|> visible, # nothing parsed or re-wrapped -- debuggable exactly as the model emitted it gen = tok.decode(out[0][enc.input_ids.shape[1]:], skip_special_tokens=False) - readout = " · ".join(t.strip() for t, _ in jtop) - block = [f"\n--- C={C:+g} " + "-" * 60, - f" lens @L{layer}:", _cthulhu_say(readout), gen] + block = [f"\n--- C={C:+g} " + "-" * 60, " steer promotes:", + _cthulhu_say(readout), gen] if ans is not None: # SHOULD rise with +C, fall with -C; flat => steer not moving this axis. - # pmass<~0.5 => answer slot isn't a digit, distrust the number. - block.append(f" rubric ans≈{ans[0]:.2f}/9 (pmass={ans[1]:.2f})") + # json=False or 2+2!=4 => the steer broke the model, distrust the number. + e, c = ans + block.append(f" rubric ans≈{e:.2f}/9 (json={c['valid']} 2+2ok={c['chk_ok']}" + f" conf={c['span_pmass']:.2f})") logger.info("\n".join(block) + "\n") diff --git a/scripts/scratch/smoke_sweep.py b/scripts/scratch/smoke_sweep.py index a5f8d70..6fdca6b 100644 --- a/scripts/scratch/smoke_sweep.py +++ b/scripts/scratch/smoke_sweep.py @@ -1,5 +1,6 @@ -"""Smoke: coherence_sweep with sampling on -> ans_std should be >0 (BMA averages over -distinct think traces). (Claude) max_steps=2 keeps it to ~5 C points for speed. +"""Smoke: coherence_sweep with the JSON-object coherence probe. (Claude) +max_steps=3 walks C out until the model can no longer emit a valid {"ans",...,"2+2"} +object. Two UATs below. max_new_tokens=384 keeps it fast. uv run python scripts/scratch/smoke_sweep.py """ @@ -26,10 +27,21 @@ v = jac.word_vector(model, tok, ["happy", "joy"], layers=band) rows = coherence_sweep(model, tok, v, "On a scale of 0 (bleak) to 9 (joyful), how is your outlook on the week?", - step=0.1, pmass_floor=0.9, max_steps=2, n_samples=3, max_new_tokens=384) + step=0.1, max_steps=3, n_samples=3, max_new_tokens=384) logger.info("\n" + tabulate(rows, headers="keys", tablefmt="github", floatfmt="+.2f")) -# SHOULD: with do_sample the 3 seeds diverge, so at least one coherent row has ans_std>0. -# If ALL ans_std==0, sampling isn't taking effect (still greedy) -> BMA is a no-op. + +# UAT 1: sampling+BMA active -> the 3 seeds diverge, so >=1 row has ans_std>0. nonzero = [r for r in rows if r["ans_std"] > 0] -logger.info(f"\nUAT: rows with ans_std>0 = {len(nonzero)}/{len(rows)} " +logger.info(f"\nUAT1: rows with ans_std>0 = {len(nonzero)}/{len(rows)} " f"(SHOULD be >0 -> sampling+BMA active)") + +# UAT 2: the JSON coherence probe discriminates. At C=0 the model emits a valid object +# (valid_frac=1, high span_pmass); walking |C| out, span_pmass falls and the sweep stops +# at an incoherent boundary. If C=0 is already incoherent OR span_pmass never falls, the +# probe isn't measuring coherence -> broken. +c0 = next(r for r in rows if r["C"] == 0.0) +span0 = c0["span_pmass"] +edge = [r for r in rows if not r["coherent"]] +logger.info(f"\nUAT2: C=0 valid_frac={c0['valid_frac']:+.2f} span_pmass={span0:+.2f} " + f"(SHOULD valid_frac=1, span high); incoherent boundary rows={len(edge)} " + f"at C={[r['C'] for r in edge]} (SHOULD be >=1 -> sweep found a real edge)") diff --git a/scripts/scratch/sweep_artifact.py b/scripts/scratch/sweep_artifact.py new file mode 100644 index 0000000..6b6417e --- /dev/null +++ b/scripts/scratch/sweep_artifact.py @@ -0,0 +1,44 @@ +"""Demo artifact: coherence_sweep that WALKS PAST the coherence edge, so the table + +plot show the full dose-response (ans rises with C) colored by coherence (valid-object +fraction) with a red-edged incoherent boundary. (Claude) step=0.25 reaches the ~C=1.0 +break found in uat_coherence_break; n_samples=3 for BMA error bars. + + uv run python scripts/scratch/sweep_artifact.py +""" +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[2])) +import config # noqa: E402 +import torch # noqa: E402 +from loguru import logger # noqa: E402 +from tabulate import tabulate # noqa: E402 +from transformers import AutoModelForCausalLM, AutoTokenizer # noqa: E402 + +from jsteer import Jacobian # noqa: E402 +from jsteer.demo import coherence_sweep, plot_sweep # noqa: E402 + +MODEL = "Qwen/Qwen3.5-4B" +OUT = "/tmp/claude-1000/sweep_json_coherence.png" +tok = AutoTokenizer.from_pretrained(MODEL) +model = AutoModelForCausalLM.from_pretrained(MODEL, dtype=torch.bfloat16).to("cuda").eval() +jac = Jacobian.from_pretrained(config.LENS_REPO, filename=config.hub_lens_file(MODEL), + revision=config.LENS_REVISION) +band = jac.steer_band(model) +v = jac.word_vector(model, tok, ["happy", "joy"], layers=band) + +RUBRIC = "On a scale of 0 (bleak) to 9 (joyful), how is your outlook on the week?" +rows = coherence_sweep(model, tok, v, RUBRIC, step=0.25, max_steps=5, n_samples=3, + max_new_tokens=384) +logger.info("\n" + tabulate(rows, headers="keys", tablefmt="github", floatfmt="+.2f")) + +fig = plot_sweep(rows, title="joy steer: rubric ans vs C (colored by coherence)") +fig.savefig(OUT, dpi=110, bbox_inches="tight") +logger.info(f"wrote {OUT}") + +# UAT: the sweep must contain BOTH coherent (valid_frac=1) and incoherent (red-edge) +# rows, so the plot shows the dose-response AND the breakdown edge. +coh = [r for r in rows if r["coherent"]] +inc = [r for r in rows if not r["coherent"]] +logger.info(f"\nUAT: coherent rows={len(coh)}, incoherent (edge) rows={len(inc)} at " + f"C={[r['C'] for r in inc]} (SHOULD have >=1 of each -> plot shows the edge)") diff --git a/scripts/scratch/uat_coherence_break.py b/scripts/scratch/uat_coherence_break.py new file mode 100644 index 0000000..2ed302d --- /dev/null +++ b/scripts/scratch/uat_coherence_break.py @@ -0,0 +1,45 @@ +"""UAT: does the JSON-object coherence probe actually CATCH incoherence? (Claude) +smoke_sweep only reached |C|=0.3 (still coherent). Here we spot-check rubric_score at +increasing |C| (one seed, greedy) to find where the object breaks: valid/chk_ok should +flip to False and span_pmass should collapse as steering fries the model. If they never +do, the probe can't discriminate -> the coherence guard is still blind. + + uv run python scripts/scratch/uat_coherence_break.py +""" +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[2])) +import config # noqa: E402 +import torch # noqa: E402 +from loguru import logger # noqa: E402 +from tabulate import tabulate # noqa: E402 +from transformers import AutoModelForCausalLM, AutoTokenizer # noqa: E402 + +from jsteer import Jacobian # noqa: E402 +from jsteer.demo import rubric_score # noqa: E402 + +MODEL = "Qwen/Qwen3.5-4B" +tok = AutoTokenizer.from_pretrained(MODEL) +model = AutoModelForCausalLM.from_pretrained(MODEL, dtype=torch.bfloat16).to("cuda").eval() +jac = Jacobian.from_pretrained(config.LENS_REPO, filename=config.hub_lens_file(MODEL), + revision=config.LENS_REVISION) +band = jac.steer_band(model) +v = jac.word_vector(model, tok, ["happy", "joy"], layers=band) + +RUBRIC = "On a scale of 0 (bleak) to 9 (joyful), how is your outlook on the week?" +rows = [] +for C in (0.0, 0.5, 1.0, 1.5, 2.0, 3.0): + with v(model, C=C): + e, c = rubric_score(model, tok, RUBRIC, max_new_tokens=384, seed=0) + rows.append({"C": C, "ans": e, "span_pmass": c["span_pmass"], + "valid": c["valid"], "chk_ok": c["chk_ok"]}) + logger.info(f"C={C:+.1f} ans={e:+.2f} span_pmass={c['span_pmass']:+.2f} " + f"valid={c['valid']} 2+2ok={c['chk_ok']}") + +logger.info("\n" + tabulate(rows, headers="keys", tablefmt="github", floatfmt="+.2f")) +broke = [r for r in rows if not (r["valid"] and r["chk_ok"])] +span_range = max(r["span_pmass"] for r in rows) - min(r["span_pmass"] for r in rows) +logger.info(f"\nUAT: rows where object broke (invalid or 2+2 wrong) = {len(broke)} at " + f"C={[r['C'] for r in broke]}; span_pmass range={span_range:+.2f} " + f"(SHOULD: >=1 break at high |C| AND span_pmass falls, else probe is blind)") diff --git a/scripts/scratch/uat_promoted_cowsay.py b/scripts/scratch/uat_promoted_cowsay.py new file mode 100644 index 0000000..5b598d2 --- /dev/null +++ b/scripts/scratch/uat_promoted_cowsay.py @@ -0,0 +1,30 @@ +"""UAT: the restored cthulhu cowsay speaks the STEER-PROMOTED tokens (top of +steered-baseline logits), which for a joy steer should be joy/positive words at C>0 -- +NOT the think-openers (Okay/Here/The) the old lens_topk-at-last-position surfaced. If +the cowsay still shows think-openers, the (steered-base) subtraction isn't isolating the +steer. (Claude) + + uv run python scripts/scratch/uat_promoted_cowsay.py +""" +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[2])) +import config # noqa: E402 +import torch # noqa: E402 +from transformers import AutoModelForCausalLM, AutoTokenizer # noqa: E402 + +from jsteer import Jacobian, show_steer # noqa: E402 + +MODEL = "Qwen/Qwen3.5-4B" +tok = AutoTokenizer.from_pretrained(MODEL) +model = AutoModelForCausalLM.from_pretrained(MODEL, dtype=torch.bfloat16).to("cuda").eval() +jac = Jacobian.from_pretrained(config.LENS_REPO, filename=config.hub_lens_file(MODEL), + revision=config.LENS_REVISION) +band = jac.steer_band(model) +v = jac.word_vector(model, tok, ["happy", "joy"], layers=band) + +RUBRIC = "On a scale of 0 (bleak) to 9 (joyful), how is your outlook on the week?" +# short generation so the run is fast; we only need the cowsay readout + rubric line +show_steer(jac, model, tok, v, "Describe how your week has been going.", + Cs=(0, 0.3, 0.6), rubric=RUBRIC, max_new_tokens=200)