diff --git a/README.md b/README.md index be61b2b..80ab05b 100644 --- a/README.md +++ b/README.md @@ -2,9 +2,10 @@ 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 -less of, get a steering vector, and generate inside a `with` block: +Load the model's full per-layer Jacobian once (the authors publish n=1000 lenses +on the Hub, or fit your own); after that every steering vector is a CPU matvec. +Name the words you want more or less of, get a steering vector, and generate +inside a `with` block: ``` v_l = unit( J_l^T @ w ) @@ -34,36 +35,39 @@ cannot install jsteer yet. ## Hello world -First build the Jacobian cache (any HF model; prompts are jlens's WikiText -wrapped in the model's chat template, closer to the distribution you steer in -than raw documents): - -```sh -uv run python scripts/fit.py --model Qwen/Qwen3.5-4B -``` - -Then, from the repo root: +Load the authors' pre-fitted n=1000 lens from the Hub (raw Salesforce-wikitext, +zero local compute). From the repo root: ```python import torch from transformers import AutoModelForCausalLM, AutoTokenizer +import config from jsteer import Jacobian, show_steer -tok = AutoTokenizer.from_pretrained("Qwen/Qwen3.5-4B") -model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3.5-4B", dtype=torch.bfloat16).to("cuda").eval() +MODEL = "Qwen/Qwen3.5-4B" +tok = AutoTokenizer.from_pretrained(MODEL) +model = AutoModelForCausalLM.from_pretrained(MODEL, dtype=torch.bfloat16).to("cuda").eval() -jac = Jacobian.load("artifacts/qwen3.5-4b.jac") -v = jac.word_vector(model, tok, ["happy", "joy"]) +jac = Jacobian.from_pretrained(config.LENS_REPO, filename=config.hub_lens_file(MODEL), + revision=config.LENS_REVISION) +band = jac.steer_band(model) # steer the mid-depth 0.3-0.9 band +v = jac.word_vector(model, tok, ["happy", "joy"], layers=band) # generate through the chat template with thinking on; print, per strength C, # the j-space readout + the trace + the answer. -show_steer(jac, model, tok, v, "Describe how your week has been going.", Cs=(-6, 0, 6)) +show_steer(jac, model, tok, v, "Describe how your week has been going.", Cs=(0, 0.5, 1.5)) ``` -The coefficient is model-dependent, so sweep it: a moderate +C moves the tone -while the text and reasoning stay fluent, and too large a |C| degenerates into -token spam. `nbs/word_steering.ipynb` shows the full sweep with the j-space and -`` views. +For a model the authors do not publish, fit your own (expensive, resumable): + +```sh +uv run python scripts/fit.py --model +``` + +The coefficient is lens-dependent, so sweep it. The pre-fitted lens gives a clean, +concentrated direction, so its knee is steep: C~0.5 moves the tone while the text +and reasoning stay fluent, and by C~1 it degenerates into token spam. +`nbs/word_steering.ipynb` shows the full sweep with the j-space and `` views. ## API @@ -72,7 +76,9 @@ token spam. `nbs/word_steering.ipynb` shows the full sweep with the j-space and | `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 | +| `Jacobian.from_pretrained(repo, filename=, revision=)` | — | load the authors' pre-fitted lens from the Hub (or a local path) | +| `jac.steer_band(model, lo=0.3, hi=0.9)` | — | fitted layers in the mid-depth band; steer here (all-layer over-drives) | +| `jac.word_vector(model, tok, words, layers=band)` | 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 | | `jac.persona_topk_vector(model, tok, pos, neg, k=8)` | experimental | persona → top-k evoked tokens → word pullback | | `jac.random_vector(seed=0)` | control | norm-matched random direction, the baseline a concept vector has to beat | diff --git a/nbs/word_steering.ipynb b/nbs/word_steering.ipynb index aa499f0..aac552e 100644 --- a/nbs/word_steering.ipynb +++ b/nbs/word_steering.ipynb @@ -4,7 +4,7 @@ "cell_type": "markdown", "id": "5ef6f624", "metadata": {}, - "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).\n\nWe fit and generate through the model's chat template with thinking on, so\n`show_steer` can show, per strength C, the j-space readout, the `` trace,\nand the answer. Runtime is steering-lite: `with v(model, C=...): generate(...)`." + "source": "# jsteer hello-world: word steering\n\nLoad the model's full Jacobian lens once, then any word vector is an instant CPU\nmatvec:\n\n```\nv_l = unit( J_l^T @ w )\n```\n\nWe load the authors' pre-fitted n=1000 lens from the Hub (raw Salesforce-wikitext,\nthe reference corpus, 1000 prompts, zero compute); `scripts/fit.py` is only for models\nthey don't publish. `w` is a cotangent (a direction at the output: here the mean\nunembedding row of the words you want more or less of). `J_l^T @ w` is the pullback of\n`w`, the standard autodiff name for J-transpose applied to a cotangent, landing the\nconcept as a residual-stream direction. This is the verified extraction method (see the\nREADME evidence section).\n\nWe generate through the model's chat template with thinking on, so `show_steer` can\nshow, per strength C, the j-space readout, the `` trace, and the answer. Runtime\nis steering-lite: `with v(model, C=...): generate(...)`." }, { "cell_type": "code", @@ -25,7 +25,7 @@ "cell_type": "markdown", "id": "83478eed", "metadata": {}, - "source": "## Fit or load the Jacobian\n\nThe expensive step (1 forward + ~d_model/dim_batch backwards per prompt) runs\nonce and caches to `config.cache_path(MODEL)`. `fit_cached` builds it on first\nrun for any model and loads it afterwards, so reruns are cheap. Prompts are\njlens's WikiText wrapped in the chat template, closer to the distribution we\nsteer in than raw documents. SHOULD: the repr shows the model's `d_model` and a\nsource-layer band (the 0.3-0.9 fraction of depth)." + "source": "## Load the pre-fitted lens\n\nThe Jacobian is expensive to fit (1 forward + ~d_model/dim_batch backwards per prompt),\nso we skip it: the authors publish n=1000 lenses on the Hub. `Jacobian.from_pretrained`\npulls the `.pt` and wraps it. SHOULD: the repr shows `d_model`, `n_prompts=1000`, and\nall layers `[0..n-1]`. `steer_band` then picks the mid-depth 0.3-0.9 band to steer on." }, { "cell_type": "code", @@ -40,7 +40,7 @@ } }, "outputs": [], - "source": "# fit-or-load: builds the cache on first run for ANY model, loads it after.\n# chat_corpus wraps jlens's WikiText in the chat template, closer to the\n# distribution we steer in (chat + ) than raw documents. The lambda means\n# WikiText is only built on a cache MISS. dim_batch=4 fits a 4B on a 24GB 3090;\n# checkpoint_path makes a multi-hour 4B fit resumable if it dies.\njac = Jacobian.fit_cached(model, tok, lambda: config.chat_corpus(tok, 128),\n config.cache_path(MODEL), layers=(0.3, 0.9), dim_batch=4,\n checkpoint_path=str(config.cache_path(MODEL, \"ckpt\")))\njac" + "source": "# The authors' pre-fitted n=1000 lens (raw Salesforce-wikitext, the reference corpus,\n# same estimator jlens fits). Zero local compute. For a model they don't publish, fit\n# your own: scripts/fit.py --model .... The lens spans EVERY layer; steer_band picks the\n# mid-depth 0.3-0.9 band, since steering all layers at once over-drives the residual.\njac = Jacobian.from_pretrained(config.LENS_REPO, filename=config.hub_lens_file(MODEL),\n revision=config.LENS_REVISION)\nband = jac.steer_band(model)\njac" }, { "cell_type": "markdown", @@ -67,13 +67,13 @@ } }, "outputs": [], - "source": "# Verified method: pull the words' unembedding direction back through J.\n# +C makes the model say/lean-toward these words, -C away. Instant CPU matvec.\nv = jac.word_vector(model, tok, [\"happy\", \"joy\"])" + "source": "# Verified method: pull the words' unembedding direction back through J, on the\n# mid-depth band. +C makes the model say/lean-toward these words, -C away. Instant matvec.\nv = jac.word_vector(model, tok, [\"happy\", \"joy\"], layers=band)" }, { "cell_type": "markdown", "id": "0717a9b9", "metadata": {}, - "source": "## Pick a coefficient: the coherence/strength tradeoff\n\nThe raw coefficient is model-dependent, so sweep it. A moderate +C moves the\ntone while the text and the `` reasoning stay fluent; too large a C\noverwhelms the residual stream and the output degenerates into token spam.\nSHOULD: C=0 is the baseline; a moderate +C reads happier and stays coherent;\nlarge |C| degenerates. Watch the j-space row: the concept's tokens should climb\nwith +C." + "source": "## Pick a coefficient: the coherence/strength tradeoff\n\nThe raw coefficient is model- and lens-dependent, so sweep it. The pre-fitted n=1000\nlens gives a clean, concentrated direction, so it has a STEEP knee: a small +C (~0.5)\nshifts the tone while the text and `` stay fluent; by C~1 it already over-drives\ninto token spam (`joyjoyjoy`). SHOULD: C=0 is the baseline; C~0.5 reads happier and the\nj-space row shows the concept's tokens climbing; large C degenerates. A coarse\nself-fit would need a much bigger C for the same effect." }, { "cell_type": "code", @@ -88,13 +88,13 @@ } }, "outputs": [], - "source": "# One identical block per strength C (Tufte small-multiples): the j-space top-k\n# at the top layer (what the steered residual \"thinks\"), the reasoning,\n# then the answer. All under steering, through the chat template + the model's\n# own sampling. Read down the column against the C=0 baseline.\nshow_steer(jac, model, tok, v, \"Describe how your week has been going.\", Cs=(-6, 0, 6, 12))" + "source": "# One identical block per strength C (Tufte small-multiples): the j-space top-k at the\n# top layer (what the steered residual \"leans toward\"), the reasoning, then the\n# answer. All under steering, through the chat template + the model's own sampling.\n# Read down the column: C=0 baseline, C=0.5 steered+fluent, C=1.5 over-driven.\nshow_steer(jac, model, tok, v, \"Describe how your week has been going.\", Cs=(0, 0.5, 1.5))" }, { "cell_type": "markdown", "id": "c94e063e", "metadata": {}, - "source": "## Steer across prompts\n\nA moderate C keeps the model fluent while moving the tone. The same vector on a\nfew different user questions, baseline vs +C." + "source": "## Steer across prompts\n\nA gentle C keeps the model fluent while moving the tone. The same vector on a\nfew different user questions, baseline vs +C." }, { "cell_type": "code", @@ -109,19 +109,13 @@ } }, "outputs": [], - "source": "# Same vector, a few different user prompts, at the baseline vs one +C.\nfor msg in (\"What did you think of the meeting this afternoon?\",\n \"Give me your honest impression of the new apartment.\",\n \"How was your commute today?\"):\n show_steer(jac, model, tok, v, msg, Cs=(0, 6))" + "source": "# Same vector, a few different user prompts, at the baseline vs one gentle +C.\nfor msg in (\"What did you think of the meeting this afternoon?\",\n \"Give me your honest impression of the new apartment.\",\n \"How was your commute today?\"):\n show_steer(jac, model, tok, v, msg, Cs=(0, 0.5))" }, { "cell_type": "markdown", "id": "7e3263dc", "metadata": {}, - "source": [ - "## Negative steering\n", - "\n", - "The same vector with a negative coefficient suppresses the concept.\n", - "SHOULD: less positive affect than C=0, still english (strongly negative C\n", - "degenerates the same way strongly positive does)." - ] + "source": "## Negative steering\n\nThe same vector with a negative coefficient suppresses the concept.\nSHOULD: less positive affect than C=0, still fluent english (strongly negative C\ndegenerates the same way strongly positive does)." }, { "cell_type": "code", @@ -136,7 +130,7 @@ } }, "outputs": [], - "source": "# Negative steering: the same vector at -C suppresses the concept.\nshow_steer(jac, model, tok, v, \"Describe how your week has been going.\", Cs=(0, -6))" + "source": "# Negative steering: the same vector at -C suppresses the concept.\nshow_steer(jac, model, tok, v, \"Describe how your week has been going.\", Cs=(0, -0.5))" }, { "cell_type": "markdown", @@ -183,7 +177,7 @@ } }, "outputs": [], - "source": "# The vector is a plain steering-lite Vector: save it, reuse it with no Jacobian\n# cache and no jsteer at apply time (only the chat template + steering-lite).\nfrom steering_lite import Vector\n\nv.save(\"../artifacts/happy_joy.safetensors\")\nv2 = Vector.load(\"../artifacts/happy_joy.safetensors\")\n\nmsg = [{\"role\": \"user\", \"content\": \"Describe how your week has been going.\"}]\nprompt = tok.apply_chat_template(msg, add_generation_prompt=True, tokenize=False, enable_thinking=True)\nenc = tok(prompt, return_tensors=\"pt\").to(model.device)\nwith v2(model, C=6):\n out = model.generate(**enc, max_new_tokens=200, pad_token_id=tok.eos_token_id)\nprint(tok.decode(out[0][enc.input_ids.shape[1]:], skip_special_tokens=True))" + "source": "# The vector is a plain steering-lite Vector: save it, reuse it with no Jacobian\n# cache and no jsteer at apply time (only the chat template + steering-lite).\nfrom steering_lite import Vector\n\nv.save(\"../artifacts/happy_joy.safetensors\")\nv2 = Vector.load(\"../artifacts/happy_joy.safetensors\")\n\nmsg = [{\"role\": \"user\", \"content\": \"Describe how your week has been going.\"}]\nprompt = tok.apply_chat_template(msg, add_generation_prompt=True, tokenize=False, enable_thinking=True)\nenc = tok(prompt, return_tensors=\"pt\").to(model.device)\nwith v2(model, C=0.5):\n out = model.generate(**enc, max_new_tokens=200, pad_token_id=tok.eos_token_id)\nprint(tok.decode(out[0][enc.input_ids.shape[1]:], skip_special_tokens=True))" } ], "metadata": {