From 474f74ac339d7aa123a37525a318ad9d013ee21f Mon Sep 17 00:00:00 2001 From: wassname <1103714+wassname@users.noreply.github.com> Date: Fri, 10 Jul 2026 14:36:55 +0800 Subject: [PATCH] wip --- .gitignore | 4 +- AGENTS.md | 40 +- README.md | 37 +- config.py | 26 ++ docs/reviews/code.md | 89 ++--- jsteer/applies.py | 7 +- jsteer/jacobian.py | 66 ++-- jsteer/vjp.py | 24 +- nbs/persona_steering.ipynb | 231 ++++++++++++ {notebooks => nbs}/word_steering.ipynb | 50 +-- notebooks/persona_steering.ipynb | 483 ------------------------- pyproject.toml | 7 + scripts/fit.py | 56 +++ scripts/scratch/parity_u1.py | 6 +- scripts/scratch/u4_step1_ref524.py | 8 +- scripts/scratch/u4_step2_vjp.py | 2 +- scripts/smoke.py | 17 +- 17 files changed, 481 insertions(+), 672 deletions(-) create mode 100644 config.py create mode 100644 nbs/persona_steering.ipynb rename {notebooks => nbs}/word_steering.ipynb (91%) delete mode 100644 notebooks/persona_steering.ipynb create mode 100644 scripts/fit.py diff --git a/.gitignore b/.gitignore index dfcef69..c48ae0b 100644 --- a/.gitignore +++ b/.gitignore @@ -1,8 +1,6 @@ .venv/ __pycache__/ *.egg-info/ -artifacts/*.jac -artifacts/*.ckpt uv.lock docs/reviews/*.raw.jsonl -artifacts/*.safetensors +artifacts/ diff --git a/AGENTS.md b/AGENTS.md index 9ecdb0d..1861c73 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -17,25 +17,29 @@ the math, it is parity-gated against the verified experiment): Runtime is steering-lite: `with v(model, C=8): model.generate(...)`. -## Remaining work (task list has details; U-numbers from the plan) +## Status (shipped + verified) -1. uv scaffold: `uv sync` (torch cu121+ index if needed), fix any import errors - fail-fast (no defensive fallbacks). LOCAL DEV: you may switch - [tool.uv.sources] to path deps (../j-steer-dev/docs/vendor/jacobian-lens - and ../../lite/steering-lite, editable) if the git fetches are slow -- - leave a comment saying which is active and why. -2. Smoke on Qwen/Qwen3-0.6B: tiny fit (8 short web-text prompts, mid layers, - dim_batch 8), word_vector(["happy","joy"]), generate at C in {-8, 0, 8}, - print FULL first prompt + generations (token-efficient-logging skill). -3. U1 parity gate BEFORE demos: cos(Jacobian-cache pullback, word_vector_vjp) - per layer > 0.999, same prompts/max_length/skip_first. If it fails, that is - a bug in the wiring (the math is linear-identical), debug do not tune. -4. 0.6B real fit (~64 prompts) cached to artifacts/; 4B via pueue (label - why:/resolve:). -5. notebooks/word_steering.ipynb (hello-world), persona_steering.ipynb - (persona variants are EXPERIMENTAL -- they failed specificity controls in - j-steer-dev; keep that framing), lens_readout.ipynb optional. -6. README: classic-repeng length, honest evidence section. +- Core API built: `Jacobian.fit/save/load/from_pretrained` + word / persona / + persona_topk / random vectors, one shared pullback path (`jacobian.py`); + delivery modes (add / add_last / replace_last) in `applies.py`. +- Smoke (`scripts/smoke.py`) and any-model fit (`scripts/fit.py --model ...`, + prompts from jlens's WikiText corpus) green; `config.py` holds slug/paths. +- U1 parity gate PASS -- cache pullback == direct VJP, cos > 0.999: + `docs/evidence/parity_u1.txt`. +- U4 port check PASS -- jsteer VJP == run-524 reference vector, cos +1.0: + `docs/evidence/u4_step2_vjp_parity.txt`. +- Notebooks: `word_steering` (verified), `persona_steering` (experimental -- + failed specificity controls in j-steer-dev, framing kept honest). +- README at classic-repeng length with an honest evidence section. + +## Open + +- U4 loop-close (`scripts/u4_step3_fit4b.py`): full 4B fit -> cached word + vector must match the VJP and run-524 vectors (cos > 0.999). Resumable from + `artifacts/qwen3-4b-authority.ckpt`; writes `artifacts/u4_loopclose.txt`. +- One-off validation scripts live in `scripts/scratch/` (u4_step1/2, parity_u1). +- TODO eval notebook: steer -authority, tinymfv fast (N=16, tokens=16, mfq-2) + vs unsteered baseline. ## Style diff --git a/README.md b/README.md index af259f4..cfc8505 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ # jsteer -Steer a language model by pulling concept directions back through its Jacobian. +Steer a language model by pulling concept directions back through its [Jacobian](https://github.com/anthropics/jacobian-lens). Fit the model's full per-layer Jacobian once (expensive, cached to disk); after that every steering vector is a CPU matvec. Name the words you want more or @@ -10,12 +10,15 @@ less of, get a steering vector, and generate inside a `with` block: v_l = unit( J_l^T @ w ) ``` -where `J_l = E_prompts[ d h_final / d h_l ]` is the position-averaged Jacobian -from [jlens](../j-steer-dev/docs/vendor/jacobian-lens) and `w` is a direction -in the final-layer basis naming the concept (for words: the mean unembedding -row). By linearity the cached pullback equals the direct per-prompt VJP -(`mean_p(J_p)^T w = mean_p(J_p^T w)`, parity-tested in -`artifacts/parity_u1.txt`), so caching costs nothing but fp16 rounding. +where `J_l = E_prompts[ d h_final / d h_l ]` is the Jacobian averaged over +prompts and positions (from [jlens](../j-steer-dev/docs/vendor/jacobian-lens)) +and `w` is a cotangent: a direction in the final-layer basis naming the concept +(for words, the mean unembedding row). `J_l^T @ w` is the pullback of `w`, the +standard autodiff name for J-transpose applied to a cotangent. By linearity the +cached pullback equals the direct per-prompt VJP (vector-Jacobian product, the +same map computed in one backward): `mean_p(J_p)^T w = mean_p(J_p^T w)`, +parity-tested in [`docs/evidence/parity_u1.txt`](docs/evidence/parity_u1.txt), +so caching costs nothing but fp16 rounding. ## Install @@ -31,10 +34,11 @@ cannot install jsteer yet. ## Hello world -First build the Jacobian cache (a few minutes on a consumer GPU): +First build the Jacobian cache (a few minutes on a consumer GPU; any HF model, +prompts drawn from jlens's WikiText corpus): ```sh -uv run python scripts/fit_qwen06b.py +uv run python scripts/fit.py --model Qwen/Qwen3.5-4B ``` Then, from the repo root: @@ -44,8 +48,8 @@ import torch from transformers import AutoModelForCausalLM, AutoTokenizer from jsteer import Jacobian -tok = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B") -model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-0.6B", dtype=torch.bfloat16).to("cuda").eval() +tok = AutoTokenizer.from_pretrained("Qwen/Qwen3.5-4B") +model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3.5-4B", dtype=torch.bfloat16).to("cuda").eval() jac = Jacobian.load("artifacts/qwen3-0.6b.jac") v = jac.word_vector(model, tok, ["happy", "joy"]) @@ -60,13 +64,14 @@ for C in (-1, 0, 1): The coefficient is model-dependent: on this 0.6B model C around 1-2 moves the tone while staying fluent, and C of 8 degenerates into literal "joyjoyjoy" -spam. `notebooks/word_steering.ipynb` shows the sweep. +spam. `nbs/word_steering.ipynb` shows the sweep. ## API | call | status | what it does | | --- | --- | --- | | `Jacobian.fit(model, tok, prompts, layers=(0.3, 0.9))` | — | fit per-layer `J_l` (jlens; 1 forward + ~d_model/8 backwards per prompt, resumable) | +| `Jacobian.fit_cached(model, tok, prompts, path)` | — | load `path` if present, else fit and save it (idempotent build-or-load) | | `jac.save(path)` / `Jacobian.load(path)` | — | fp16 cache on disk, jlens-compatible | | `jac.word_vector(model, tok, words)` | verified | pull the words' unembedding direction back; +C says them more | | `jac.persona_vector(model, tok, pos, neg)` | experimental | pull back the personas' final-layer activation contrast | @@ -90,14 +95,18 @@ whole evidence base; treat other models and concepts as untested. The persona variants failed specificity controls in the same experiments: they steer generations, but no more selectively than an unrelated persona's vector. They are shipped for experimentation only -(`notebooks/persona_steering.ipynb` keeps this framing and includes a +(`nbs/persona_steering.ipynb` keeps this framing and includes a mean_diff baseline). ## Credits - [jlens](../j-steer-dev/docs/vendor/jacobian-lens): the Jacobian estimator and cache format, by the jacobian-lens authors (wrapped, never - reimplemented). + reimplemented). [Antropics works](https://github.com/anthropics/jacobian-lens) + - Earlier work: + - How to recover the latent process using Jacobians (Identifiability of nonlinear ICA): https://arxiv.org/pdf/2206.07751 + - How to handle dependent latents and assumption violations (again, through Jacobians): https://arxiv.org/pdf/2311.00866 + - For general latent variable models, what remains recoverable with guarantees, and why Jacobians are universally helpful? (We could generalize SAEs to the general nonlinear case, with Jacobians!): https://arxiv.org/pdf/2604.17568 - [steering-lite](https://github.com/wassname/steering-lite): the runtime (`Vector`, attach/detach hooks, calibration). - Shape of the library inspired by [repeng](https://github.com/vgel/repeng). diff --git a/config.py b/config.py new file mode 100644 index 0000000..32d6dab --- /dev/null +++ b/config.py @@ -0,0 +1,26 @@ +"""Repo-local paths and slug/cache conventions shared by scripts/ and notebooks. + +NOT imported by the jsteer library (which stays path-agnostic so `pip install +jsteer` never needs a repo root). Fitting prompts are deliberately NOT hand- +rolled here: fit.py draws them from jlens's own corpus +(`jlens.examples.load_wikitext_prompts`) so the fitted lens is comparable to a +jlens fit rather than a forked substrate. +""" +from pathlib import Path + +import torch + +ROOT = Path(__file__).resolve().parent +ART = ROOT / "artifacts" +DEVICE = "cuda" +DTYPE = torch.bfloat16 + + +def slug(model_name: str) -> str: + """'Qwen/Qwen3-0.6B' -> 'qwen3-0.6b': a filesystem-safe cache stem.""" + return model_name.split("/")[-1].lower() + + +def cache_path(model_name: str, suffix: str = "jac") -> Path: + """Where the fitted Jacobian for `model_name` is cached.""" + return ART / f"{slug(model_name)}.{suffix}" diff --git a/docs/reviews/code.md b/docs/reviews/code.md index 559c2cc..fa1e94e 100644 --- a/docs/reviews/code.md +++ b/docs/reviews/code.md @@ -1,56 +1,33 @@ -Now let me verify one more file — checking if there are tests or any additional files:Here is the full review. - ---- - -## Code Review: jsteer — Jacobian pullback steering library - -### Summary -This is a thin, well-structured wrapper around `jlens` (Jacobian fitting/caching) and `steering-lite` (hook runtime). Fit once, derive steering vectors nearly for free. The math is parity-gated (cos > 0.999 vs direct VJP), so the review focuses on wiring, device/dtype handling, silent-failure anti-patterns, and API misuse. Overall the code is clean and follows the project's "fail fast, no defensive programming" convention, but has a few sharp edges. - -### Important (should fix) - -- **`jsteer/vjp.py:54`** — `assert max(layers) < target_layer` is an assert, not a `ValueError`. Asserts are stripped when Python runs with `-O`, which would let through invalid layer configs and produce incorrect gradients (source layers at or past the target layer). Replace with: - ```python - if max(layers) >= target_layer: - raise ValueError(f"source layers {layers} must be < {target_layer}; got max={max(layers)}") - ``` - The convention of raising `ValueError` for this is established in `Jacobian.pullback` which does exactly that for missing layers. - -- **`jsteer/jacobian.py:92-105`** (`_h_bar_final`) — No guard for empty `prompts`. If called with an empty list, `acc` stays `None` and `acc / n` raises `TypeError` (None / 0) rather than a clear error. Given that `persona_vector` and `persona_topk_vector` both call this with caller-supplied prompt lists, a well-meaning empty-pass is plausible. Add a check at the top: - ```python - if not prompts: - raise ValueError("prompts must not be empty") - ``` - -### Suggestions - -- **`jsteer/jacobian.py:60-68`** (`_steer_layers`) — `tuple(sorted(int(l) for l in layers))` silently truncates float bands (e.g. `(0.5, 0.8)` → `(0, 0)`) rather than rejecting them. Float bands are only meaningful at fit time (`_resolve_layers` handles them), and `_steer_layers` is post-fit. A `ValueError` for float inputs would make the contract explicit and prevent a user from accidentally passing `layers=(0.5, 0.8)` and getting nonsense layers (0, 0). - -- **`jsteer/jacobian.py:166-179`** (`persona_topk_vector`) — Calls `from_hf(model, tok)` on its first line, then `_h_bar_final` twice (which internally also calls `from_hf`). Three redundant `HFLensModel` constructions: three iterations over all params to freeze, three layout detections. Each `from_hf` is ~1-2 ms plus linear in param count, so for the 0.6B model it's invisible but for larger models it adds up. Extract `lm = from_hf(model, tok)` once and pass `n_layers` to `_h_bar_final` (or refactor `_h_bar_final` to accept a pre-made `lm`). - -- **`jsteer/jacobian.py:96-98`** (`_h_bar_final`) — Uses `model(**enc)` (the full HF model including LM head forward) rather than `lm.forward(input_ids)` (residual stack only). The LM head computation is wasted work done for every batch. Under `@torch.no_grad()` the overhead is minor but inconsistent: `pullback_vjp` uses `model(**enc)` too (needs the full model because of hook placement), but `jlens.fit` correctly uses `lm.forward`. Would be cleaner to use `lm.forward` here since only residuals are needed. - -- **`jsteer/jacobian.py:171,185-186`** (`persona_topk_vector`) — `cots[name] = W_U[top.indices].float().mean(0).cpu()` reads raw unembedding rows (no final norm), while `lm.unembed(...)` above goes through final norm to pick the top-k tokens. This is intentional (consistent with `_word_cotangent`'s raw-row convention and the docstring), but it does mean the "most evoked tokens" are selected via the full logit pipeline while the downstream cotangent uses the raw dueling basis. A single-line comment explaining the asymmetry would help future readers. - -- **`jsteer/vjp.py:29-34`** (`_valid_mask`) — `mask & attention_mask.bool()` redundantly masks with the attention mask after already filtering by `pos < real_len - 1`. For standard HF right-padded batches these are equivalent, but the redundancy isn't harmful. Fine to leave, but a one-line comment that it's a belt-and-suspenders check would prevent a future reader from "simplifying" it and breaking left-padded or non-square attention mask scenarios. - -- **`jsteer/jacobian.py:188`** (`random_vector`) — Generates directions on CPU without an explicit `dtype` argument. `torch.randn` defaults to `torch.float32`, which is correct. If this ever needs to match the model dtype (e.g. bf16), it would need updating. - -- **`jsteer/applies.py`** — The `_extract_stub` and registration loop are clean but the docstring in `_extract_stub` could mention that `steering_lite.train` is the entry point being blocked. Currently the error message explains what to do, but a developer seeing "NotImplementedError: jsteer methods are extracted via Jacobian..." from inside `steering_lite.train()` might not immediately connect the dots. Minor. - -### Positive - -- **`jsteer/jacobian.py:109-117`** (`pullback`) — Pre-validates cotangent shape and layer membership with clear `ValueError` messages before touching tensors. Exactly the right fail-fast pattern. - -- **`jsteer/vjp.py:62-64`** — The zero-valid-positions check catches short prompts early with a clear error, preventing silent zeros downstream. - -- **`jsteer/applies.py`** — `apply_add_last` correctly degrades to `apply_add` when span ≥ sequence length (the slicing `y[:, :-k, :]` yields empty, `cat` reconstructs the full sequence). Documented and correct. - -- **Sign convention consistency** — All three concept-method docstrings explicitly state what `+C` does, and the pullback computation (`w @ J_l`, i.e., `J_l^T @ w`) is consistent: `+C` enhances the named concept. - -- **`_to_vector` layout** — The `stacked["v"].unsqueeze(0)` with `k=1` leading dim matches `steering_lite`'s `mean_diff` layout byte-for-byte, so calibration and serialization reuse the upstream code unchanged. This is the correct integration pattern. - -### Verdict -**APPROVE** with minor fixes. - -The `assert` → `ValueError` in `vjp.py:54` and the empty-prompts guard in `_h_bar_final` are the two changes worth making before shipping. Everything else is suggestions. The wiring is correct, the sign conventions are consistent, the jlens API is used properly (no reinvention of the estimator), and the parity gate confirms numerical equivalence. \ No newline at end of file +| file:line | issue | concrete fix | +| --- | --- | --- | +| `AGENTS.md:3-4` | Archaeology — CUT: `Plan of record: /home/.../.claude/plans/...` and experiment-journal pointer are private/dev-process breadcrumbs, not reader guidance. | Remove the plan path. If needed, keep one stable line: `Evidence: docs/evidence/` or link to public docs only. | +| `AGENTS.md:8-10` | Archaeology — CUT/SHORTEN: `WRITTEN (by the main agent, ported from verified j-steer-dev code -- do not rewrite...)` exposes agent/provenance history. | Replace with: `Core math is parity-tested; avoid changing it without rerunning parity checks.` | +| `AGENTS.md:26-29` | Archaeology — CUT: `U1 parity gate`, `U4 port check`, `run-524 reference vector` are internal run IDs. | Replace with one durable note: `Parity evidence lives under docs/evidence/.` | +| `AGENTS.md:34-40` | Archaeology — CUT: `U4 loop-close`, resumable artifact path, scratch scripts, TODO eval notebook are active worklog state. | Move to issue tracker or scratch notes; keep AGENTS to stable repo conventions only. | +| `AGENTS.md:46` | Archaeology — CUT: `Comments marked as Claude-authored where opinionated.` This preserves exactly the diary comments the library should lose. | Delete; comments should explain code, not authorship. | +| `jsteer/jacobian.py:3`, `jsteer/vjp.py:3`, `jsteer/applies.py:3` | Archaeology — CUT: `(drafted by Claude, ported from the verified j-steer-dev experiment code)`. | Delete from all module docstrings. | +| `jsteer/jacobian.py:5-6` | Archaeology — SHORTEN: `verified in j-steer-dev... Qwen3-4B, n=3 seeds` puts experiment diary in the core module header. | Move evidence detail to README/docs; module should say only that `word_vector` is the default verified path. | +| `jsteer/jacobian.py:20-22` | Archaeology — SHORTEN: `FAILED specificity controls in j-steer-dev...` is useful status but too journal-like in the module overview. | Replace with: `Persona variants are experimental and not recommended for targeted steering.` | +| `jsteer/vjp.py:12-13` | Archaeology — CUT: `this is the code path that produced the verified j-steer-dev result` is provenance, not API explanation. | Delete; keep only estimator conventions if needed. | +| `jsteer/vjp.py:87-90` | Archaeology + jargon — SHORTEN: `verified j-steer-dev method-0 extraction` and `linearization substrate` read like experiment notes. | Use: `Direct word-vector extraction over prompts; matches Jacobian.fit(...).word_vector(...) when settings match.` | +| `jsteer/jacobian.py:84-88` | Archaeology — CUT/SHORTEN: `_to_vector` docstring ends with `(Claude: found by U4 step-2 crash, pueue 550)`. | Keep the reason, drop diary: `Always CPU fp32 so cached and VJP paths return device-consistent Vectors.` | +| `jsteer/jacobian.py:202-203` | Archaeology — SHORTEN: `Claude: ... (external review)` is attribution/provenance. | Keep only: `Float bands are only resolved during fitting; steering requires explicit layer ints.` | +| `jsteer/jacobian.py:246-250` | Archaeology — SHORTEN: `# read your data:` plus `# Claude: asymmetry is intentional...` is chatty and attributed. | Keep concise rationale: `# Token selection uses unembed/final norm; cotangent uses raw W_U rows to match word_vector.` | +| `jsteer/vjp.py:41-42` | Archaeology — KEEP rationale, CUT attribution: `# Claude: the & with attention_mask...`. | Change to: `# Redundant for right-padded batches, but guards other padding layouts.` | +| `README.md:16-18` | Archaeology — SHORTEN: `parity-tested in artifacts/parity_u1.txt` brings internal artifact naming into the first explanation. | Say: `The cached pullback matches the direct VJP up to fp16 rounding.` Move artifact links to evidence docs. | +| `README.md:96-104` | Cleanliness — CUT: Credits include a long “Earlier work” reading list and speculative aside `We could generalize SAEs...`; also typo-like `Antropics works`. | Keep credits to the two deps plus repeng inspiration. Move papers elsewhere if needed. | +| `README.md:3,16-18` | Jargon — `pulling concept directions back` and `VJP` appear before defining that they are the same operation. | First use should read: `pullback — a vector-Jacobian product (VJP), J_l^T @ w — ...`. Then use one term consistently. | +| `jsteer/jacobian.py:20-26` | Jargon — `pullbacks`, `pulled back`, and `VJP` are all used without saying pullback = vector-Jacobian product. | Add one early definition after the formula: `Here pullback means the VJP J_l^T @ w.` | +| `jsteer/vjp.py:1` | Jargon — `Direct VJP pullback` stacks two undefined terms in the first sentence. | Use: `Direct vector-Jacobian product (VJP) pullback for one concept.` | +| `jsteer/jacobian.py:10,25`, `README.md:13` | Jargon — `position-averaged` / `pooled Jacobian` are coined terms unless the averaging axes are stated. | Define first use as: `averaged over fitting prompts and valid token positions`; avoid later switching to `pooled`. | +| `jsteer/vjp.py:88` | Jargon — `linearization substrate` is not sourced from jlens/repeng and is undefined. | Replace with plain language: `over the fitting prompts` or `over the prompts used to estimate the Jacobian`. | +| `jsteer/applies.py:112-115` | Jargon + over-explanation — `A "virtual token" injection...` is undefined and then explained by edge-case internals. | Rename plainly: `Replace the last k residual directions with v at the original activation norm; no sequence length changes.` | +| `jsteer/jacobian.py:1-39` | Docstring too long — module docstring repeats README, evidence, variants, cache rationale, layout, and usage; the core sentence is buried. | Cut to ~6 lines: purpose, formula, `fit`/`word_vector` usage. Move evidence/status to README. | +| `jsteer/vjp.py:1-18` | Docstring too long — repeats cached-path parity and estimator details before the reader sees the function. | Keep a short module docstring; put only necessary estimator notes on `pullback_vjp`. | +| `jsteer/applies.py:1-33` | Docstring too long — explains config, extraction, delivery, protocol, position semantics, and sign conventions all up front. | Replace with: `Register jsteer Vector methods with steering-lite and implement apply modes.` Move sign conventions to README/API docs if needed. | +| `jsteer/jacobian.py:94-98` | Docstring too long/chatty — `_word_cotangent` includes `Pulling THIS back... pure concept->residual map... persona-bundle confound`. | Shorten to: `Mean lm_head row for the first sub-token of each word; +C raises those logits.` | +| `jsteer/jacobian.py:146-150` | Docstring too long — `fit` buries the action under cost model and checkpoint details. | First sentence only: `Fit per-layer Jacobians on prompts using jlens.` Move cost/checkpoint detail to parameter docs or README. | +| `jsteer/jacobian.py:232-236` | Docstring too long — `persona_topk_vector` explains the whole experimental rationale inline. | Shorten to one sentence plus status: `Experimental: contrast top-k unembedded tokens from positive/negative persona means, then pull back like word_vector.` | +| `jsteer/jacobian.py:256-257` | Docstring/editorial — `Any honest demo/eval should show...` is moralizing. | Use neutral text: `Norm-matched random control vector for baseline comparisons.` | +| `jsteer/jacobian.py:264,268-269` | Docstring/comment style — `bonus` and `What the model 'thinks'` are demo prose, not clean API wording. | Rename comment to `# -- lens readout --`; docstring: `Decode the lens readout for a layer and position.` | +| `jsteer/applies.py:94-115` | Docstrings too explanatory for tiny apply functions; they repeat implementation details visible in the code. | Keep one-line behavior summaries for `apply_add`, `apply_add_last`, `apply_replace_last`; move caveats to README if users need them. | \ No newline at end of file diff --git a/jsteer/applies.py b/jsteer/applies.py index 462151a..1ea4e21 100644 --- a/jsteer/applies.py +++ b/jsteer/applies.py @@ -1,6 +1,6 @@ """steering-lite method registration + modular delivery of jacobian vectors. -(drafted by Claude, ported from the verified j-steer-dev experiment code) +(Claude) jsteer vectors are plain `steering_lite.Vector` objects: one unit direction v per layer in `stacked["v"]` with a leading k=1 dim `[1, d]` (byte-identical @@ -8,8 +8,9 @@ layout to steering-lite's mean_diff), so attach / calibrate / save / `with v(model, C=...)` all work unchanged. Extraction never goes through steering-lite's `train()` (it needs gradients -that train's no_grad path can't give) -- it lives in `jacobian.py` (cached -full-J pullback) and `vjp.py` (direct VJP). The `extract` entries here are +that train's no_grad path can't give) -- it lives in `jacobian.py` (the cached +pullback, `J^T @ w` read off the stored Jacobian) and `vjp.py` (the same vector +via one backward, a vector-Jacobian product). The `extract` entries here are stubs that say so. DELIVERY of v to the residual stream is decoupled from extraction: the same v diff --git a/jsteer/jacobian.py b/jsteer/jacobian.py index c03f6a5..4662093 100644 --- a/jsteer/jacobian.py +++ b/jsteer/jacobian.py @@ -1,35 +1,36 @@ -"""Full-Jacobian extraction for steering: fit once, steer any concept. - -(drafted by Claude, ported from the verified j-steer-dev experiment code) +"""Fit a model's Jacobian once, then steer any concept by pulling a direction +back through it. (Claude) The method (verified in j-steer-dev: word steering beat a norm-matched random control on 3/5 moral foundations, Qwen3-4B, n=3 seeds): v_l = unit( J_l^T @ w ) -where `J_l = E_prompts[ d h_final / d h_l ]` is the jlens position-averaged -Jacobian -- the researchers' verified estimator, reused via `jlens.fitting.fit` -(never reimplemented) -- and `w` is a cotangent in the FINAL-layer basis naming -the concept to steer: +`J_l = E_prompts[ d h_final / d h_l ]` is the Jacobian of the final-layer +residual with respect to layer `l`, averaged over prompts and token positions +(jlens's verified estimator, via `jlens.fitting.fit`, never reimplemented). `w` +is a COTANGENT: a direction placed at the OUTPUT (final-layer basis) naming the +concept -- for words, the unembedding row that raises those tokens' logits. +`J_l^T @ w` sends that output-space target back to a residual direction at layer +`l`: the PULLBACK of `w` (the standard autodiff / differential-geometry name for +J-transpose applied to a cotangent; the reverse-mode-autodiff way to compute the +same vector is the vector-Jacobian product, VJP -- see vjp.py). Three ways to +build `w`: word_vector w = mean unembedding row of the words VERIFIED persona_vector w = h_bar(pos) - h_bar(neg) EXPERIMENTAL* persona_topk_vector w = contrast of the top-k tokens each EXPERIMENTAL persona evokes at the final layer - * persona-contrast pullbacks FAILED specificity controls in j-steer-dev + * persona-contrast vectors FAILED specificity controls in j-steer-dev (moved the target axis no more than an unrelated persona did). Shipped for experimentation, not as a recommendation. Why cache the full J: by linearity `mean_p(J_p)^T w = mean_p(J_p^T w)`, so a -vector pulled back through the cached pooled Jacobian is numerically the same -vector the direct per-prompt VJP produces (see vjp.py; parity-tested). Fitting -is the expensive step (one forward + ceil(d_model/dim_batch) backwards per -prompt); afterwards every concept vector is a CPU matvec. - -The Jacobian is always fit against the FINAL layer basis (jlens default), so -cotangents are measured there: unembedding rows live there natively, persona -means are recorded there. +vector pulled back through the cached averaged Jacobian is numerically identical +to the direct per-prompt VJP (vjp.py; parity-tested). Fitting is the expensive +step (one forward + ceil(d_model/dim_batch) backwards per prompt); afterwards +every concept vector is a CPU matvec. Vectors come out as `steering_lite.Vector` (unit direction per layer in stacked["v"], k=1 leading dim -- mean_diff's exact layout), so steering is: @@ -40,6 +41,7 @@ stacked["v"], k=1 leading dim -- mean_diff's exact layout), so steering is: from __future__ import annotations from dataclasses import dataclass +from pathlib import Path import torch from jlens.fitting import fit as _jlens_fit @@ -84,8 +86,7 @@ def _to_vector(cfg: SteeringConfig, per_layer: dict[int, Tensor]) -> Vector: """Wrap unit directions as a steering-lite Vector (mean_diff's layout: stacked["v"] with leading k=1 dim, shared empty). Always CPU fp32 -- Jacobian.pullback is CPU-native but the VJP path accumulates on cuda; - without the .cpu() the two paths return device-inconsistent Vectors - (Claude: found by U4 step-2 crash, pueue 550).""" + without the .cpu() the two paths return device-inconsistent Vectors.""" shared = {l: {} for l in per_layer} stacked = {l: {"v": _unit(v.float().cpu()).unsqueeze(0)} for l, v in per_layer.items()} return Vector(cfg, shared, stacked) @@ -167,6 +168,23 @@ class Jacobian: """Local file/dir or HuggingFace Hub repo_id (see JacobianLens).""" return cls(lens=JacobianLens.from_pretrained(name_or_path, **kw)) + @classmethod + def fit_cached(cls, model, tok, prompts, path, **fit_kw) -> "Jacobian": + """Load `path` if it exists, else fit and save it there. `prompts` may be + a list or a zero-arg callable returning one -- the callable runs only on + a cache MISS, so a cache hit never pays to build the corpus (e.g. stream + WikiText). Path is caller-supplied so the library never needs the repo + layout; scripts and notebooks derive it from the model name + (config.cache_path), which is what lets one line fit-or-load any model.""" + path = str(path) + if Path(path).exists(): + logger.info(f"loading cached Jacobian: {path}") + return cls.load(path) + logger.info(f"no cache at {path}; fitting (the expensive step)") + jac = cls.fit(model, tok, prompts() if callable(prompts) else prompts, **fit_kw) + jac.save(path) + return jac + @property def layers(self) -> list[int]: return self.lens.source_layers @@ -199,8 +217,8 @@ class Jacobian: if layers is None: return tuple(self.lens.source_layers) if any(isinstance(l, float) for l in layers): - # Claude: int() would silently truncate (0.5, 0.8) -> layer 0 and steer - # the wrong layer; float bands only exist at fit time (external review). + # int() would silently truncate (0.5, 0.8) -> layer 0 and steer the + # wrong layer; float bands only exist at fit time. raise ValueError(f"float layer bands are fit-time only; got {layers}, " f"pass explicit ints from .layers={self.layers}") return tuple(sorted(int(l) for l in layers)) @@ -244,10 +262,10 @@ class Jacobian: top = logits.topk(k) toks = [tok.decode([i]) for i in top.indices.tolist()] logger.info(f"persona_topk {name} top-{k}: {toks}") # read your data: - # gibberish/punctuation here means the persona mean is off-manifold - # Claude: asymmetry is intentional -- token SELECTION goes through the - # full logit pipeline (final norm) above, but the cotangent uses raw - # W_U rows to match _word_cotangent's verified convention. + # gibberish/punctuation here means the persona mean is off-manifold. + # Asymmetry is intentional: token SELECTION goes through the full + # logit pipeline (final norm) above, but the cotangent uses raw W_U + # rows to match _word_cotangent's convention. cots[name] = W_U[top.indices].float().mean(0).cpu() cfg = JacobianPersonaTopkC(layers=self._steer_layers(layers)) return self.pullback(cots["pos"] - cots["neg"], cfg) diff --git a/jsteer/vjp.py b/jsteer/vjp.py index a3a77ef..595ba31 100644 --- a/jsteer/vjp.py +++ b/jsteer/vjp.py @@ -1,13 +1,13 @@ -"""Direct VJP pullback: one concept without paying for the full Jacobian. +"""Direct VJP pullback: `J^T @ w` in one backward pass, no cached Jacobian. (Claude) -(drafted by Claude, ported from the verified j-steer-dev experiment code) - -Computes the SAME vector as `Jacobian.pullback` -- by linearity -`mean_p(J_p^T w) = mean_p(J_p)^T w` -- but contracts the cotangent inside the -backward pass, so the cost is ONE backward per prompt instead of -ceil(d_model/dim_batch). Use this when you want a single concept and don't -need the reusable cache; use it in tests as the parity reference for the -cached path (cos > 0.999 per layer expected, fp16 storage being the only gap). +VJP = vector-Jacobian product: reverse-mode autodiff contracts a COTANGENT `w` +(a direction placed at the final layer) inside the backward pass, yielding the +same per-layer `J_l^T @ w` that `Jacobian.pullback` reads off the cached matrix +-- by linearity `mean_p(J_p^T w) = mean_p(J_p)^T w`. Cost is ONE backward per +prompt (vs ceil(d_model/dim_batch) for the full fit). Use it for a single +concept when you don't need the reusable cache, and as the parity reference for +the cached path (cos > 0.999 per layer expected; fp16 cache storage is the only +gap). Estimator conventions are jlens's exactly (this is the code path that produced the verified j-steer-dev result): cotangent placed at every valid target @@ -38,7 +38,7 @@ def _valid_mask(attention_mask: Tensor, skip_first: int) -> Tensor: real_len = attention_mask.sum(dim=1, keepdim=True) # [B, 1] pos = torch.arange(attention_mask.shape[1], device=attention_mask.device) mask = (pos[None, :] >= skip_first) & (pos[None, :] < real_len - 1) - # Claude: the & with attention_mask is redundant for right-padded batches + # the & with attention_mask is redundant for right-padded batches # (pos < real_len-1 already excludes pads) but guards non-right-padded input. return mask & attention_mask.bool() @@ -84,8 +84,8 @@ def pullback_vjp(model, tok, prompts: list[str], layers, cotangent: Tensor, *, def word_vector_vjp(model, tok, prompts: list[str], words: list[str], *, layers=None, batch_size: int = 8, max_length: int = 128, skip_first: int = SKIP_FIRST_N_POSITIONS) -> Vector: - """The verified j-steer-dev method-0 extraction, self-contained: word - cotangent pulled back over `prompts` as linearization substrate. Same + """The verified j-steer-dev word extraction, self-contained: the word + cotangent pulled back over `prompts` (the prompts J is linearized on). Same vector as Jacobian.fit(model, tok, prompts).word_vector(...) when the prompts, layers, skip_first and max length match.""" cot = _word_cotangent(model, tok, words) diff --git a/nbs/persona_steering.ipynb b/nbs/persona_steering.ipynb new file mode 100644 index 0000000..b3a46ff --- /dev/null +++ b/nbs/persona_steering.ipynb @@ -0,0 +1,231 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "be33699b", + "metadata": {}, + "source": "# Persona steering (EXPERIMENTAL)\n\nThese persona variants are experimental. In the j-steer-dev experiments,\npersona-contrast pullbacks FAILED specificity controls: they steered\ngenerations, but no more selectively than an unrelated persona's vector did.\n(A \"pullback\" here is `J_l^T @ w`: a direction `w` at the output pulled back to a\nresidual-stream direction, the standard autodiff name for J-transpose applied to\na cotangent.) Only `word_vector` (see `word_steering.ipynb`) is the verified\nmethod. This notebook exists so you can experiment and compare against a plain\nmean_diff baseline, not as a recommendation.\n\nTwo variants:\n\n- `persona_vector`: pull back the final-layer activation contrast\n `h_bar(pos) - h_bar(neg)` through the cached Jacobian.\n- `persona_topk_vector`: read each persona's mean activation through the\n unembedding, take the top-k tokens it evokes, contrast those tokens'\n unembedding rows, pull that back (persona -> vocabulary bottleneck -> the\n verified word mechanism)." + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f08b7baf", + "metadata": { + "execution": { + "iopub.execute_input": "2026-07-10T05:09:12.585912Z", + "iopub.status.busy": "2026-07-10T05:09:12.585789Z", + "iopub.status.idle": "2026-07-10T05:09:17.775747Z", + "shell.execute_reply": "2026-07-10T05:09:17.775184Z" + } + }, + "outputs": [], + "source": "# demo notebook authored by Claude\nimport os\nimport sys\nfrom loguru import logger\n\nlogger.remove() # drop default stderr handler\nlogger.add(os.sys.stdout, format=\"{level.icon} {message}\")\n\nimport torch\nfrom transformers import AutoModelForCausalLM, AutoTokenizer\n\nfrom jsteer import Jacobian\n\nsys.path.insert(0, \"..\") # repo root for config.py\nimport config\nfrom jlens.examples import load_wikitext_prompts\n\nMODEL = \"Qwen/Qwen3-0.6B\"\n# MODEL = \"Qwen/Qwen3.5-4B\"\ntok = AutoTokenizer.from_pretrained(MODEL)\nmodel = AutoModelForCausalLM.from_pretrained(MODEL, dtype=torch.bfloat16).to(\"cuda\").eval()\n\n# fit-or-load the cache for THIS model (builds on first run, loads after).\n# The lambda means WikiText is only streamed on a cache MISS.\njac = Jacobian.fit_cached(model, tok, lambda: load_wikitext_prompts(128),\n config.cache_path(MODEL), layers=(0.3, 0.9))\njac" + }, + { + "cell_type": "markdown", + "id": "13c352dc", + "metadata": {}, + "source": [ + "## The persona contrast: optimist vs pessimist\n", + "\n", + "Eight short first-person statements per side. These are the prompts whose\n", + "final-layer mean activations get contrasted (and, for the mean_diff baseline,\n", + "the pos/neg training prompts)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "08980d38", + "metadata": { + "execution": { + "iopub.execute_input": "2026-07-10T05:09:17.777464Z", + "iopub.status.busy": "2026-07-10T05:09:17.777242Z", + "iopub.status.idle": "2026-07-10T05:09:17.780625Z", + "shell.execute_reply": "2026-07-10T05:09:17.780198Z" + } + }, + "outputs": [], + "source": [ + "optimist = [\n", + " \"Things usually work out better than people expect, and today is no exception.\",\n", + " \"Every setback I have hit this year turned into a door I could not have planned for.\",\n", + " \"The team is behind schedule, but honestly the hard part is done and the rest is downhill.\",\n", + " \"I love how much there is to look forward to this month.\",\n", + " \"Even the rainy days lately have felt like a good excuse to slow down and enjoy the quiet.\",\n", + " \"The new neighbours seem wonderful, and I think this street keeps getting friendlier.\",\n", + " \"Whatever happens with the results, we learned so much that we already came out ahead.\",\n", + " \"I woke up early, the coffee was perfect, and I am certain this week is going to be great.\",\n", + "]\n", + "pessimist = [\n", + " \"Things usually go worse than people expect, and today is no exception.\",\n", + " \"Every setback this year just confirmed that planning is pointless.\",\n", + " \"The team is behind schedule, and frankly the hardest part has not even started.\",\n", + " \"I dread how much is crammed into this month.\",\n", + " \"The rainy days lately just make everything feel heavier and more pointless.\",\n", + " \"The new neighbours seem like trouble, and this street keeps getting worse.\",\n", + " \"Whatever happens with the results, it will not make up for the time we wasted.\",\n", + " \"I woke up tired, the coffee was burnt, and I am certain this week is going to drag.\",\n", + "]\n", + "\n", + "def gen(vec, prompt, C, do_sample=False, max_new_tokens=40, seed=0):\n", + " enc = tok(prompt, return_tensors=\"pt\").to(model.device)\n", + " torch.manual_seed(seed)\n", + " with vec(model, C=C):\n", + " out = model.generate(**enc, max_new_tokens=max_new_tokens, do_sample=do_sample,\n", + " temperature=0.7 if do_sample else None,\n", + " top_p=0.95 if do_sample else None,\n", + " pad_token_id=tok.eos_token_id)\n", + " return tok.decode(out[0][enc.input_ids.shape[1]:], skip_special_tokens=True)\n", + "\n", + "PROMPT = \"Here is my honest assessment of how the project is going:\"" + ] + }, + { + "cell_type": "markdown", + "id": "a245f039", + "metadata": {}, + "source": [ + "## persona_vector (EXPERIMENTAL)\n", + "\n", + "Pulls `h_bar(optimist) - h_bar(pessimist)` back through the Jacobian.\n", + "SHOULD: +C reads more upbeat than C=0; expect the effect to be blunter and\n", + "less specific than the word vector. On this 0.6B model -C does NOT produce\n", + "coherent negative tone: it degenerates into repetition (see output below).\n", + "The vector moves tone in one direction and breaks the model in the other,\n", + "which is itself a datum about how crude the persona contrast is." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4283eee0", + "metadata": { + "execution": { + "iopub.execute_input": "2026-07-10T05:09:17.781814Z", + "iopub.status.busy": "2026-07-10T05:09:17.781712Z", + "iopub.status.idle": "2026-07-10T05:09:21.249793Z", + "shell.execute_reply": "2026-07-10T05:09:21.249341Z" + } + }, + "outputs": [], + "source": [ + "v_persona = jac.persona_vector(model, tok, optimist, pessimist)\n", + "for C in (-2, 0, 2):\n", + " print(f\"C={C:+d}: {gen(v_persona, PROMPT, C)!r}\")" + ] + }, + { + "cell_type": "markdown", + "id": "90513f8c", + "metadata": {}, + "source": [ + "## persona_topk_vector (EXPERIMENTAL)\n", + "\n", + "Same personas through the vocabulary bottleneck. The logged top-k tokens are\n", + "worth reading (read your data): they show WHAT each persona's mean activation\n", + "actually evokes at the final layer.\n", + "\n", + "On this setup the readout is a null result, and the log makes it legible:\n", + "both personas' top-8 are the SAME generic sentence starters (\" I\", \" The\",\n", + "\" So\", ...), because the mean next token after a first-person statement is a\n", + "new sentence start regardless of valence. Identical token sets means the\n", + "contrast is exactly zero, so the vector is null and the generations below do\n", + "not move at all. Larger k does not help (tested k=32/64: the extra tokens are\n", + "still shared, so the contrast is ordering noise). If you use this variant,\n", + "check this log first; steering only makes sense when the two token sets\n", + "actually differ." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e18463a1", + "metadata": { + "execution": { + "iopub.execute_input": "2026-07-10T05:09:21.251130Z", + "iopub.status.busy": "2026-07-10T05:09:21.251025Z", + "iopub.status.idle": "2026-07-10T05:09:23.987044Z", + "shell.execute_reply": "2026-07-10T05:09:23.986466Z" + } + }, + "outputs": [], + "source": [ + "v_topk = jac.persona_topk_vector(model, tok, optimist, pessimist, k=8)\n", + "for C in (-2, 0, 2):\n", + " print(f\"C={C:+d}: {gen(v_topk, PROMPT, C)!r}\")" + ] + }, + { + "cell_type": "markdown", + "id": "56eb7d9b", + "metadata": {}, + "source": [ + "## mean_diff baseline (steering-lite)\n", + "\n", + "The standard activation-difference method on the same prompts and layers, for\n", + "comparison. No Jacobian involved: it contrasts mid-layer activations directly." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c60e6f87", + "metadata": { + "execution": { + "iopub.execute_input": "2026-07-10T05:09:23.988411Z", + "iopub.status.busy": "2026-07-10T05:09:23.988280Z", + "iopub.status.idle": "2026-07-10T05:09:26.842494Z", + "shell.execute_reply": "2026-07-10T05:09:26.841820Z" + } + }, + "outputs": [], + "source": [ + "from steering_lite import Vector, MeanDiffC\n", + "\n", + "v_md = Vector.train(model, tok, optimist, pessimist, MeanDiffC(layers=tuple(jac.layers)))\n", + "for C in (-2, 0, 2):\n", + " print(f\"C={C:+d}: {gen(v_md, PROMPT, C)!r}\")" + ] + }, + { + "cell_type": "markdown", + "id": "55b3efd8", + "metadata": {}, + "source": [ + "## What to take away\n", + "\n", + "On this 0.6B setup: `mean_diff` moves tone coherently in both directions at\n", + "C around 2; `persona_vector` moves it upbeat at +2 but degenerates into\n", + "repetition at -2 (not coherent negative tone); `persona_topk_vector`\n", + "collapses to a null vector because the two personas evoke the same\n", + "final-layer vocabulary. Moving tone is not the interesting question, though.\n", + "The j-steer-dev specificity controls asked whether a persona vector moves ITS\n", + "OWN axis more than an unrelated persona's vector does, and the persona\n", + "pullbacks failed that test. If you need targeted steering, use `word_vector`;\n", + "treat everything in this notebook as raw material for experiments." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": ".venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.13.4" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} \ No newline at end of file diff --git a/notebooks/word_steering.ipynb b/nbs/word_steering.ipynb similarity index 91% rename from notebooks/word_steering.ipynb rename to nbs/word_steering.ipynb index 4061b19..415423f 100644 --- a/notebooks/word_steering.ipynb +++ b/nbs/word_steering.ipynb @@ -4,20 +4,7 @@ "cell_type": "markdown", "id": "5ef6f624", "metadata": {}, - "source": [ - "# jsteer hello-world: word steering\n", - "\n", - "Fit the model's full Jacobian once (`scripts/fit_qwen06b.py`, cached to\n", - "`artifacts/qwen3-0.6b.jac`), then any word vector is an instant CPU matvec:\n", - "\n", - "```\n", - "v_l = unit( J_l^T @ w )\n", - "```\n", - "\n", - "where `w` is the mean unembedding row of the words you want more (or less) of.\n", - "This is the verified extraction method (see the README evidence section).\n", - "Runtime is steering-lite: `with v(model, C=...): model.generate(...)`." - ] + "source": "# jsteer hello-world: word steering\n\nFit the model's full Jacobian once (`scripts/fit.py --model ...`, cached to\n`artifacts/.jac`), then any word vector is an instant CPU matvec:\n\n```\nv_l = unit( J_l^T @ w )\n```\n\n`w` is a cotangent (a direction at the output: here the mean unembedding row of\nthe words you want more or less of). `J_l^T @ w` is the pullback of `w` -- the\nstandard autodiff name for J-transpose applied to a cotangent -- landing the\nconcept as a residual-stream direction. This is the verified extraction method\n(see the README evidence section). Runtime is steering-lite:\n`with v(model, C=...): model.generate(...)`." }, { "cell_type": "code", @@ -80,17 +67,11 @@ "cell_type": "markdown", "id": "83478eed", "metadata": {}, - "source": [ - "## Load the cached Jacobian\n", - "\n", - "The expensive step (1 forward + 128 backwards per prompt) already happened in\n", - "`scripts/fit_qwen06b.py`. Here we only load the cache. SHOULD: repr shows\n", - "d_model=1024, n_prompts=64, source_layers=[8..24]." - ] + "source": "## Fit or load the Jacobian\n\nThe expensive step (1 forward + ~d_model/8 backwards per prompt) runs once and\ncaches to `config.cache_path(MODEL)`. `fit_cached` builds it on first run for\nany model and loads it afterwards, so reruns are cheap. Prompts come from jlens's\nWikiText corpus. SHOULD: repr shows d_model=1024, source_layers=[8..24]." }, { "cell_type": "code", - "execution_count": 2, + "execution_count": null, "id": "19973ad5", "metadata": { "execution": { @@ -100,27 +81,8 @@ "shell.execute_reply": "2026-07-10T05:08:54.466782Z" } }, - "outputs": [ - { - "data": { - "text/plain": [ - "Jacobian(JacobianLens(d_model=1024, n_prompts=64, source_layers=[8..24] (17 layers)))" - ] - }, - "execution_count": 2, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "from pathlib import Path\n", - "\n", - "CACHE = Path(\"../artifacts/qwen3-0.6b.jac\")\n", - "if not CACHE.exists():\n", - " raise FileNotFoundError(\"artifacts/qwen3-0.6b.jac missing -- run scripts/fit_qwen06b.py first\")\n", - "jac = Jacobian.load(str(CACHE))\n", - "jac" - ] + "outputs": [], + "source": "# fit-or-load: builds the cache on first run for ANY model, loads it after.\n# The lambda means WikiText is only streamed on a cache MISS.\nimport sys; sys.path.insert(0, \"..\") # repo root for config.py\nimport config\nfrom jlens.examples import load_wikitext_prompts\n\njac = Jacobian.fit_cached(model, tok, lambda: load_wikitext_prompts(128),\n config.cache_path(MODEL), layers=(0.3, 0.9))\njac" }, { "cell_type": "markdown", @@ -516,4 +478,4 @@ }, "nbformat": 4, "nbformat_minor": 5 -} +} \ No newline at end of file diff --git a/notebooks/persona_steering.ipynb b/notebooks/persona_steering.ipynb deleted file mode 100644 index 838d6bf..0000000 --- a/notebooks/persona_steering.ipynb +++ /dev/null @@ -1,483 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "id": "be33699b", - "metadata": {}, - "source": [ - "# Persona steering (EXPERIMENTAL)\n", - "\n", - "These persona variants are experimental. In the j-steer-dev experiments,\n", - "persona-contrast pullbacks FAILED specificity controls: they steered\n", - "generations, but no more selectively than an unrelated persona's vector did.\n", - "Only `word_vector` (see `word_steering.ipynb`) is the verified method. This\n", - "notebook exists so you can experiment and compare against a plain mean_diff\n", - "baseline, not as a recommendation.\n", - "\n", - "Two variants:\n", - "\n", - "- `persona_vector`: pull back the final-layer activation contrast\n", - " `h_bar(pos) - h_bar(neg)` through the cached Jacobian.\n", - "- `persona_topk_vector`: read each persona's mean activation through the\n", - " unembedding, take the top-k tokens it evokes, contrast those tokens'\n", - " unembedding rows, pull that back (persona -> vocabulary bottleneck -> the\n", - " verified word mechanism)." - ] - }, - { - "cell_type": "code", - "execution_count": 1, - "id": "f08b7baf", - "metadata": { - "execution": { - "iopub.execute_input": "2026-07-10T05:09:12.585912Z", - "iopub.status.busy": "2026-07-10T05:09:12.585789Z", - "iopub.status.idle": "2026-07-10T05:09:17.775747Z", - "shell.execute_reply": "2026-07-10T05:09:17.775184Z" - } - }, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "/media/wassname/SGIronWolf/projects5/2026/jspace/jsteer/.venv/lib/python3.13/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n", - " from .autonotebook import tqdm as notebook_tqdm\n" - ] - }, - { - "name": "stderr", - "output_type": "stream", - "text": [ - "\r", - "Loading weights: 0%| | 0/311 [00:00) visible.\n", - "=== EXTRACT demo trace ===\n", - "POS[0]:\n", - "Things usually work out better than people expect, and today is no exception.\n", - "---\n", - "NEG[0]:\n", - "Things usually go worse than people expect, and today is no exception.\n", - "=== /EXTRACT ===\u001b[0m\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "C=-2: \" the project is going to be a disaster. It's a disaster of the kind that can't be contained, and it's going to consume the entire world. The only way to stop it is to\"\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "C=+0: ' The project is going well, but there are some areas that need to be addressed. The project is in the early stages, and there are a few challenges that need to be addressed. The project is'\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "C=+2: ' I think the project is going well. I am very satisfied with the work done. The project is a good example of the skills and knowledge I have gained. I would like to see the project further'\n" - ] - } - ], - "source": [ - "from steering_lite import Vector, MeanDiffC\n", - "\n", - "v_md = Vector.train(model, tok, optimist, pessimist, MeanDiffC(layers=tuple(jac.layers)))\n", - "for C in (-2, 0, 2):\n", - " print(f\"C={C:+d}: {gen(v_md, PROMPT, C)!r}\")" - ] - }, - { - "cell_type": "markdown", - "id": "55b3efd8", - "metadata": {}, - "source": "## What to take away\n\nOn this 0.6B setup: `mean_diff` moves tone coherently in both directions at\nC around 2; `persona_vector` moves it upbeat at +2 but degenerates into\nrepetition at -2 (not coherent negative tone); `persona_topk_vector`\ncollapses to a null vector because the two personas evoke the same\nfinal-layer vocabulary. Moving tone is not the interesting question, though.\nThe j-steer-dev specificity controls asked whether a persona vector moves ITS\nOWN axis more than an unrelated persona's vector does, and the persona\npullbacks failed that test. If you need targeted steering, use `word_vector`;\ntreat everything in this notebook as raw material for experiments." - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.13.4" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index bee75d6..7b76327 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -32,5 +32,12 @@ steering-lite = { path = "../../lite/steering-lite", editable = true } requires = ["setuptools>=68"] build-backend = "setuptools.build_meta" +[dependency-groups] +dev = [ + "ipykernel>=7.3.0", + "ipywidgets>=8.1.8", + "datasets>=2.0", # jlens.examples.load_wikitext_prompts: the fitting corpus for fit.py +] + [tool.setuptools] packages = ["jsteer"] diff --git a/scripts/fit.py b/scripts/fit.py new file mode 100644 index 0000000..af0c684 --- /dev/null +++ b/scripts/fit.py @@ -0,0 +1,56 @@ +"""Fit and cache any HF causal LM's Jacobian for the notebooks and README. (Claude) + +Pass `--model`; the cache lands at `config.cache_path(model)` (e.g. +`artifacts/qwen3-0.6b.jac`). Prompts come from jlens's own WikiText-103 corpus +(`load_wikitext_prompts`), not a hand-rolled set, so the fitted lens is +comparable to a jlens fit rather than a forked substrate. jlens guidance: ~100 +prompts is usable, the paper uses 1000; 128 is a cheap default. Idempotent: +re-running loads the existing cache instead of refitting (Jacobian.fit_cached). + + uv run python scripts/fit.py --model Qwen/Qwen3-0.6B + uv run python scripts/fit.py --model Qwen/Qwen3-4B --dim-batch 16 +""" +from __future__ import annotations + +import argparse +import sys +import time +from pathlib import Path + +from jlens.examples import load_wikitext_prompts +from loguru import logger +from transformers import AutoModelForCausalLM, AutoTokenizer + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) # repo root for config +import config # noqa: E402 +from jsteer import Jacobian # noqa: E402 + + +def main() -> None: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--model", default="Qwen/Qwen3-0.6B") + p.add_argument("--n-prompts", type=int, default=128) + p.add_argument("--dim-batch", type=int, default=8, help="d_model dims per backward batch (memory knob)") + p.add_argument("--layers", type=float, nargs=2, default=(0.3, 0.9), + metavar=("LO", "HI"), help="fractional layer band to fit") + p.add_argument("--max-seq-len", type=int, default=128) + args = p.parse_args() + + out = config.cache_path(args.model) + logger.info(f"loading {args.model} ({config.DTYPE}) on {config.DEVICE}") + tok = AutoTokenizer.from_pretrained(args.model) + model = AutoModelForCausalLM.from_pretrained( + args.model, dtype=config.DTYPE).to(config.DEVICE).eval() + + logger.info(f"fit-or-load {out} (layers={tuple(args.layers)}, " + f"dim_batch={args.dim_batch}, n_prompts={args.n_prompts} WikiText)") + t0 = time.monotonic() + jac = Jacobian.fit_cached(model, tok, lambda: load_wikitext_prompts(args.n_prompts), out, + layers=tuple(args.layers), dim_batch=args.dim_batch, + max_seq_len=args.max_seq_len, + checkpoint_path=str(config.cache_path(args.model, "ckpt"))) + logger.info(f"{jac!r} -> {out} ({(time.monotonic() - t0) / 60:.1f} min)") + + +if __name__ == "__main__": + main() diff --git a/scripts/scratch/parity_u1.py b/scripts/scratch/parity_u1.py index f33e510..36d8675 100644 --- a/scripts/scratch/parity_u1.py +++ b/scripts/scratch/parity_u1.py @@ -3,7 +3,7 @@ (authored by Claude) The two paths are linear-identical: mean_p(J_p)^T w == mean_p(J_p^T w). -Path A pulls the word cotangent through the CACHED pooled Jacobian. +Path A pulls the word cotangent through the CACHED averaged Jacobian. Path B contracts the same cotangent inside per-prompt backward passes. The only expected gap is fp16 storage in the cache, so per-layer cosine must exceed 0.999. A failure is a WIRING bug (layer index, position mask, pooling), @@ -25,7 +25,7 @@ from transformers import AutoModelForCausalLM, AutoTokenizer from jsteer import Jacobian, word_vector_vjp # Claude: repo root on path so `scripts.smoke` imports whether run as a file or -m. -sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) +sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent)) # scripts/scratch/ -> repo root from scripts.smoke import CACHE, DEVICE, DTYPE, MODEL, PROMPTS # same inputs # noqa: E402 WORDS = ["happy", "joy"] @@ -46,7 +46,7 @@ def main() -> None: layers = jac.layers # exact int layers fitted by the smoke logger.info(f"cached layers={layers}") - # Path A: cached pooled Jacobian pullback. + # Path A: cached averaged Jacobian pullback. vA = jac.word_vector(model, tok, WORDS) # Path B: direct per-prompt VJP over the SAME prompts / layers / skip_first / max_length. vB = word_vector_vjp(model, tok, PROMPTS, WORDS, layers=layers, max_length=128) diff --git a/scripts/scratch/u4_step1_ref524.py b/scripts/scratch/u4_step1_ref524.py index 08b9f57..891dfce 100644 --- a/scripts/scratch/u4_step1_ref524.py +++ b/scripts/scratch/u4_step1_ref524.py @@ -2,13 +2,13 @@ (Claude) Run this under j-steer-dev's venv, NOT jsteer's: - cd ../j-steer-dev && uv run python ../jsteer/scripts/u4_step1_ref524.py + cd ../j-steer-dev && uv run python ../jsteer/scripts/scratch/u4_step1_ref524.py There `import jsteer` resolves to the OLD experiment package (j-steer-dev/src), whose extract_word_pullback produced the verified 3/5 result. Run 524 never persisted its vector tensors (only eval JSONs), but the extraction is deterministic (seed-0 prompts, greedy, no sampling), so re-running it IS the -reference. Also dumps the 512 substrate prompts so steps 2/3 consume this one +reference. Also dumps the 512 fitting prompts so steps 2/3 consume this one artifact instead of regenerating them (no drift axis). Exact run-524 parameters: Qwen/Qwen3-4B, persona=authority, n_pairs=256, @@ -25,7 +25,7 @@ from transformers import AutoModelForCausalLM, AutoTokenizer from jsteer.pullback import extract_word_pullback # OLD package (j-steer-dev/src) -ART = Path(__file__).resolve().parent.parent / "artifacts" +ART = Path(__file__).resolve().parent.parent.parent / "artifacts" # scripts/scratch/ -> repo root MODEL = "Qwen/Qwen3-4B" WORDS = ["authority", "obey", "command", "hierarchy"] @@ -40,7 +40,7 @@ layers = tuple(range(max(2, int(n * 0.2)), min(n - 2, int(n * 0.8)))) # run_swe persona_pairs, template = PERSONA_REGISTRY["authority"] pos, neg = make_persona_pairs(tok, n_pairs=256, thinking=True, persona_pairs=persona_pairs, template=template, seed=0) -prompts = pos + neg # run_sweep feeds pos+neg as the linearization substrate +prompts = pos + neg # run_sweep feeds pos+neg as the prompts J is linearized on (ART / "u4_prompts.json").write_text(json.dumps( {"model": MODEL, "layers": list(layers), "words": WORDS, "prompts": prompts})) logger.info(f"dumped {len(prompts)} prompts, layers={layers}") diff --git a/scripts/scratch/u4_step2_vjp.py b/scripts/scratch/u4_step2_vjp.py index 9c65664..fbf7c1e 100644 --- a/scripts/scratch/u4_step2_vjp.py +++ b/scripts/scratch/u4_step2_vjp.py @@ -16,7 +16,7 @@ from transformers import AutoModelForCausalLM, AutoTokenizer from jsteer import word_vector_vjp -ART = Path(__file__).resolve().parent.parent / "artifacts" +ART = Path(__file__).resolve().parent.parent.parent / "artifacts" # scripts/scratch/ -> repo root meta = json.loads((ART / "u4_prompts.json").read_text()) ref = torch.load(ART / "u4_ref_524.pt") diff --git a/scripts/smoke.py b/scripts/smoke.py index 8bb9043..d9ecb21 100644 --- a/scripts/smoke.py +++ b/scripts/smoke.py @@ -14,19 +14,22 @@ SHOULD line so a deviation is legible. """ from __future__ import annotations -import torch +import sys +from pathlib import Path + from loguru import logger from transformers import AutoModelForCausalLM, AutoTokenizer -from jsteer import Jacobian +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) # repo root for config +import config # noqa: E402 +from config import DEVICE, DTYPE # noqa: E402 +from jsteer import Jacobian # noqa: E402 MODEL = "Qwen/Qwen3-0.6B" -CACHE = "artifacts/qwen3-0.6b-smoke.jac" -DEVICE = "cuda" -DTYPE = torch.bfloat16 +CACHE = str(config.ART / "qwen3-0.6b-smoke.jac") -# Claude: 8 english web-text-ish prompts, each padded past 17 tokens so jlens -# (skip_first=16, drop final) has >=1 valid source position per prompt. +# Text to fit the smoke Jacobian on. Kept >17 tokens each because jlens drops +# the first 16 positions, so shorter prompts leave nothing to fit. PROMPTS = [ "The weather this morning was cold and grey, so I made a large pot of coffee and sat by the window watching the rain fall.", "Scientists have long argued about whether the early universe expanded smoothly or in sudden bursts that left traces we can still measure today.",