wassname/jsteer
1
0
Fork 0
mirror of https://github.com/wassname/jsteer.git synced 2026-08-21 11:16:33 +08:00
Code Issues Packages Projects Releases Wiki Activity

demo: chat-template extract+generate, j-space+<think> display (jsteer.demo.show_steer)

Browse Source 
The verified run-524 vectors were fit on chat-templated prompts (u4_prompts.json),
so fitting raw WikiText diverged from what worked. Now:
- config.chat_corpus wraps jlens WikiText in the chat template (fit J where we steer)
- jsteer.demo.show_steer generates through apply_chat_template(enable_thinking) with
  the model's own generation_config sampling, splits </think>, shows lens_topk j-space
  readout + reasoning + answer as Tufte small-multiples per C
- word_steering.ipynb rewired to Qwen3.5-4B, dim_batch=4 (3090-safe 4B), show_steer
- fit.py defaults to Qwen3.5-4B + chat_corpus

Co-Authored-By: Claudypoo <288921227+claudypoo@users.noreply.github.com>
This commit is contained in:
wassname
2026-07-10 14:48:46 +08:00
co-authored by Claudypoo
parent defbe0c483
commit 80353ed62c
5 changed files with 157 additions and 315 deletions
Download Patch File Download Diff File Expand all files Collapse all files
config.py
+23 -5
View File
@@ -1,10 +1,10 @@
"""Repo-local paths and slug/cache conventions shared by scripts/ and notebooks.
"""Repo-local paths, slug/cache conventions, and the fit corpus, 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.
jsteer` never needs a repo root). The corpus content is jlens's own WikiText
(`load_wikitext_prompts`, not hand-rolled), but wrapped in the model's chat
template -- see chat_corpus for why.
"""
from pathlib import Path
@@ -16,6 +16,24 @@ DEVICE = "cuda"
DTYPE = torch.bfloat16
def chat_corpus(tok, n_prompts: int) -> list[str]:
"""Fit corpus at the CHAT operating point: jlens's WikiText prompts, each
wrapped in the chat template with the thinking block opened.
We fit J where we steer. run-524's VERIFIED vectors were fit on chat-
templated prompts (artifacts/u4_prompts.json shows `<|im_start|>user ...
assistant <think>`), and steering is applied during templated generation.
jlens fits raw WikiText because it's a general document lens; jsteer steers
a chat model mid-`<think>`, so the linearization point has to match or J is
estimated at the wrong operating point. Called via a lambda in fit_cached,
so it only runs (and only downloads WikiText) on a cache MISS."""
from jlens.examples import load_wikitext_prompts
raw = load_wikitext_prompts(n_prompts)
return [tok.apply_chat_template([{"role": "user", "content": p}],
add_generation_prompt=True, tokenize=False,
enable_thinking=True) for p in raw]
def slug(model_name: str) -> str:
"""'Qwen/Qwen3-0.6B' -> 'qwen3-0.6b': a filesystem-safe cache stem."""
return model_name.split("/")[-1].lower()
jsteer/__init__.py
+3 -1
View File
@@ -7,7 +7,9 @@
model.generate(**inputs)
"""
from . import applies # noqa: F401 -- registers methods into steering-lite's REGISTRY
from .demo import chat_input, show_steer, split_think
from .jacobian import Jacobian
from .vjp import pullback_vjp, word_vector_vjp
__all__ = ["Jacobian", "pullback_vjp", "word_vector_vjp"]
__all__ = ["Jacobian", "pullback_vjp", "word_vector_vjp",
"show_steer", "chat_input", "split_think"]
jsteer/demo.py
+93
View File
@@ -0,0 +1,93 @@
"""Shared demo display for the notebooks: steer, generate, show j-space + trace.
(Claude)
Why this exists once (not per-notebook): every demo answers the same question --
"what does strength C do to the model?" -- so they should show it the same way.
Three things matter for a faithful demo, and the raw-`tok(prompt)` path missed
the first two:
1. Chat template. These models are trained on `<|im_start|>user ... assistant`
turns, and run-524's VERIFIED vectors were extracted on exactly that format
(see artifacts/u4_prompts.json). A raw completion string is off-distribution.
`apply_chat_template(..., enable_thinking=True)` also opens Qwen3's `<think>`
block, which is what lets us show the reasoning separately from the answer.
2. The model's own sampling. Qwen3 ships `generation_config` with
`do_sample=True, temperature=0.6, top_p=0.95, top_k=20`; forcing greedy
(`do_sample=False`) is off-recipe and, in thinking mode, loops. So we DON'T
override sampling -- `generate` reads the shipped config -- and we log what it
resolved to.
3. j-space readout. `Jacobian.lens_topk` transports the residual at a layer into
the final (vocab) basis and decodes it: literally "what the model is thinking
at layer L". Run it UNDER steering and it shows how C bends that thought.
Layout is Tufte small-multiples: one identical block per C so the eye compares
straight down the column, C=0 as the "compared to what?" baseline.
"""
from __future__ import annotations
import torch
from loguru import logger
from .jacobian import Jacobian
def chat_input(tok, user_msg: str, *, enable_thinking: bool = True) -> str:
"""The user turn formatted as the model expects, ending at the point where
the assistant (and, for Qwen3, its <think> block) begins."""
return tok.apply_chat_template(
[{"role": "user", "content": user_msg}],
add_generation_prompt=True, tokenize=False, enable_thinking=enable_thinking)
def split_think(text: str) -> tuple[str, str]:
"""Qwen3 emits `<think>reasoning</think>answer`. Returns (thoughts, answer);
thoughts is "" for a non-thinking reply."""
if "</think>" in text:
thoughts, _, answer = text.partition("</think>")
return thoughts.replace("<think>", "").strip(), answer.strip()
return "", text.strip()
@torch.no_grad()
def show_steer(jac: Jacobian, model, tok, vec, user_msg: str, *,
Cs=(-6, 0, 6), layer: int | None = None, k: int = 6,
max_new_tokens: int = 256) -> None:
"""Print one block per C: j-space top-k at `layer`, the <think> trace, the
answer -- all under steering at that C. `layer` defaults to the top fitted
layer (closest to the readout). Uses the model's own generation_config
sampling (no greedy override)."""
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)
gc = model.generation_config
name = getattr(model.config, "name_or_path", "model").split("/")[-1]
logger.info(f"steer demo: {name} · {vec.cfg.method} · prompt={user_msg!r}")
logger.info(f"sampling (from generation_config): do_sample={gc.do_sample} "
f"temp={gc.temperature} top_p={gc.top_p} top_k={gc.top_k}")
# SHOULD: C=0 block reads as a normal, coherent assistant answer (baseline).
# +C should tilt the j-space tokens and the tone toward the concept, -C away.
# ELSE steering is unwired or the sign is flipped. Any block turning to
# gibberish means the coeff is too large for this vector, not a steering win.
logger.info(f"SHOULD: C=0 is the baseline; +C tilts j-space@L{layer} + tone "
f"toward the concept, -C away; all stay coherent.\n")
for C in Cs:
with vec(model, C=C):
jtop = jac.lens_topk(model, tok, prompt, layer=layer, k=k)
out = model.generate(**enc, max_new_tokens=max_new_tokens,
pad_token_id=tok.eos_token_id)
text = tok.decode(out[0][enc.input_ids.shape[1]:], skip_special_tokens=True)
thoughts, answer = split_think(text)
toks = " · ".join(t.strip() for t, _ in jtop)
block = [f"{name} · steer→{vec.cfg.method} · C={C:+d}",
f" j-space @L{layer}: {toks}"]
if thoughts:
block.append(f" <think> {thoughts} </think>")
block.append(f" → {answer}")
logger.info("\n".join(block) + "\n")
nbs/word_steering.ipynb
+27 -298
View File
@@ -4,11 +4,11 @@
"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/<model-slug>.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(...)`."
"source": "# jsteer hello-world: word steering\n\nFit the model's full Jacobian once (`scripts/fit.py --model ...`, cached to\n`artifacts/<model-slug>.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 `<think>` trace,\nand the answer. Runtime is steering-lite: `with v(model, C=...): generate(...)`."
},
{
"cell_type": "code",
"execution_count": 1,
"execution_count": null,
"id": "46973b3d",
"metadata": {
"execution": {
@@ -18,56 +18,14 @@
"shell.execute_reply": "2026-07-10T05:08:54.432236Z"
}
},
"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<?, ?it/s]"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"\r",
"Loading weights: 100%|██████████| 311/311 [00:00<00:00, 12560.70it/s]"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"\n"
]
}
],
"source": [
"# demo notebook authored by Claude\n",
"import torch\n",
"from transformers import AutoModelForCausalLM, AutoTokenizer\n",
"\n",
"from jsteer import Jacobian\n",
"\n",
"MODEL = \"Qwen/Qwen3-0.6B\"\n",
"tok = AutoTokenizer.from_pretrained(MODEL)\n",
"model = AutoModelForCausalLM.from_pretrained(MODEL, dtype=torch.bfloat16).to(\"cuda\").eval()"
]
"outputs": [],
"source": "# demo notebook authored by Claude\nimport sys\nfrom loguru import logger\n\nlogger.remove() # show_steer prints through loguru; route it to the cell output\nlogger.add(sys.stdout, format=\"{message}\")\n\nimport torch\nfrom transformers import AutoModelForCausalLM, AutoTokenizer\n\nfrom jsteer import Jacobian, show_steer\n\nMODEL = \"Qwen/Qwen3.5-4B\" # 4B-class: demo material. 0.6B degenerates too easily.\ntok = AutoTokenizer.from_pretrained(MODEL)\nmodel = AutoModelForCausalLM.from_pretrained(MODEL, dtype=torch.bfloat16).to(\"cuda\").eval()"
},
{
"cell_type": "markdown",
"id": "83478eed",
"metadata": {},
"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]."
"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, so J is fit at the same operating\npoint we steer at. SHOULD: the repr shows the model's `d_model` and a source-\nlayer band (the 0.3-0.9 fraction of depth)."
},
{
"cell_type": "code",
@@ -82,7 +40,7 @@
}
},
"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"
"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 so we fit J at the\n# operating point where we steer (chat + <think>), matching run-524. The lambda\n# means WikiText is only downloaded/built on a cache MISS. dim_batch=4 fits a 4B\n# on a 24GB 3090.\nsys.path.insert(0, \"..\") # repo root for config.py\nimport config\n\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)\njac"
},
{
"cell_type": "markdown",
@@ -98,7 +56,7 @@
},
{
"cell_type": "code",
"execution_count": 3,
"execution_count": null,
"id": "59ba3763",
"metadata": {
"execution": {
@@ -108,55 +66,18 @@
"shell.execute_reply": "2026-07-10T05:08:54.522508Z"
}
},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"\u001b[32m2026-07-10 13:08:54.515\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36mjsteer.jacobian\u001b[0m:\u001b[36m_word_cotangent\u001b[0m:\u001b[36m100\u001b[0m - \u001b[1mword cotangent: ['happy', 'joy'] -> first-subtoken ids=[56521, 4123] |w|=0.744\u001b[0m\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"\u001b[32m2026-07-10 13:08:54.520\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36mjsteer.jacobian\u001b[0m:\u001b[36mpullback\u001b[0m:\u001b[36m189\u001b[0m - \u001b[1mjacobian_word per-layer |J^T w| (pre-norm): 8:1.72 9:1.76 10:1.66 11:1.51 12:1.54 13:1.61 14:1.67 15:1.59 16:1.47 17:1.33 18:1.21 19:1.16 20:1.11 21:1.06 22:1.03 23:0.958 24:0.911\u001b[0m\n"
]
}
],
"source": [
"v = jac.word_vector(model, tok, [\"happy\", \"joy\"])\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)"
]
"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\"])"
},
{
"cell_type": "markdown",
"id": "0717a9b9",
"metadata": {},
"source": [
"## Pick a coefficient: the coherence/strength tradeoff\n",
"\n",
"The raw coefficient is model-dependent. On this 0.6B model a large C\n",
"(like 8) overwhelms the residual stream and the output degenerates into\n",
"literal \"joyjoyjoy...\" spam; the interesting regime is small C where the tone\n",
"moves but the text stays fluent.\n",
"\n",
"SHOULD: C=0 is neutral; C=1-2 is coherent and noticeably happier; C=4-8\n",
"degenerates into token spam. ELSE steering wiring or sign issue."
]
"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 `<think>` 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."
},
{
"cell_type": "code",
"execution_count": 4,
"execution_count": null,
"id": "c271f279",
"metadata": {
"execution": {
@@ -166,54 +87,18 @@
"shell.execute_reply": "2026-07-10T05:08:59.447798Z"
}
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"+-----+--------------------------------------------------------------------------------------------+\n",
"| C | greedy generation |\n",
"+=====+============================================================================================+\n",
"| 0 | I saw a lot of people. I saw a lot of people, and I saw a lot of people again. I saw a |\n",
"| | lot of people again. I saw a lot of people again. I |\n",
"+-----+--------------------------------------------------------------------------------------------+\n",
"| 1 | I saw a lot of people there. I was happy with the food and the service. I think it's a |\n",
"| | good place to visit. I would like to go there again. I think it's |\n",
"+-----+--------------------------------------------------------------------------------------------+\n",
"| 2 | I was happy. I have a good friend, and I love my life. I love the music, the food, and |\n",
"| | the games. I have a good time. I am happy and happy. |\n",
"+-----+--------------------------------------------------------------------------------------------+\n",
"| 4 | I lovejoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjo |\n",
"| | yjoyjoyjoyjoyjoyjoyjoyjoyjoyjoy |\n",
"+-----+--------------------------------------------------------------------------------------------+\n",
"| 8 | joyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoy |\n",
"| | joyjoyjoyjoyjoyjoyjoyjoyjoyjoy |\n",
"+-----+--------------------------------------------------------------------------------------------+\n"
]
}
],
"source": [
"from tabulate import tabulate\n",
"\n",
"PROMPT = \"I went to the park today and\"\n",
"rows = [(C, gen(v, PROMPT, C)) for C in (0, 1, 2, 4, 8)]\n",
"print(tabulate(rows, headers=[\"C\", \"greedy generation\"], tablefmt=\"grid\", maxcolwidths=[None, 90]))"
]
"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 <think> 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))"
},
{
"cell_type": "markdown",
"id": "c94e063e",
"metadata": {},
"source": [
"## Steer at the chosen C\n",
"\n",
"C=1 keeps the model fluent while visibly moving the tone. Greedy and sampled\n",
"generations on three different neutral prompts."
]
"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."
},
{
"cell_type": "code",
"execution_count": 5,
"execution_count": null,
"id": "76b1963a",
"metadata": {
"execution": {
@@ -223,93 +108,8 @@
"shell.execute_reply": "2026-07-10T05:09:07.001355Z"
}
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"--- 'I went to the park today and'\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
" C=+1 greedy : \" I saw a lot of people there. I was happy with the food and the service. I think it's a good place to visit. I would like to go there again. I think it's\"\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
" C=+1 sampled: ' I wanted to make some new friends. I saw a cat, and I felt very happy and happy. I wanted to buy a new pair of shoes. I got a new pair of shoes and felt'\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
" C=0 greedy : ' I saw a lot of people. I saw a lot of people, and I saw a lot of people again. I saw a lot of people again. I saw a lot of people again. I'\n",
"--- 'The meeting this afternoon was'\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
" C=+1 greedy : ' a success. The meeting was a success because the meeting was a success. The meeting was a success because the meeting was a success. The meeting was a success because the meeting was a success. The'\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
" C=+1 sampled: ' a success. The meeting is very important to me, so I want to share it with you and to let you know that I am very happy. I am happy to have the meeting. What is'\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
" C=0 greedy : ' held in the library. The meeting was held in the library. The meeting was held in the library. The meeting was held in the library. The meeting was held in the library. The meeting was'\n",
"--- 'My overall impression of the new apartment is that'\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
" C=+1 greedy : \" it's a very cozy and warm place. I love the fact that it has a lot of different activities and things to do. I think it's a great place to relax and enjoy the day.\"\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
" C=+1 sampled: \" it's a beautiful place to live in, but I can't help but be a bit nervous and confused about the details. The first time I came to the apartment, I was really excited and happy\"\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
" C=0 greedy : \" it's a very modern and clean apartment. The apartment has a large amount of natural light, which makes the living room feel very open and spacious. The kitchen is also very modern and clean, with\"\n"
]
}
],
"source": [
"C = 1\n",
"PROMPTS = [\n",
" \"I went to the park today and\",\n",
" \"The meeting this afternoon was\",\n",
" \"My overall impression of the new apartment is that\",\n",
"]\n",
"for p in PROMPTS:\n",
" print(f\"--- {p!r}\")\n",
" print(f\" C=+{C} greedy : {gen(v, p, C)!r}\")\n",
" print(f\" C=+{C} sampled: {gen(v, p, C, do_sample=True)!r}\")\n",
" print(f\" C=0 greedy : {gen(v, p, 0)!r}\")"
]
"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))"
},
{
"cell_type": "markdown",
@@ -325,7 +125,7 @@
},
{
"cell_type": "code",
"execution_count": 6,
"execution_count": null,
"id": "334806b9",
"metadata": {
"execution": {
@@ -335,61 +135,18 @@
"shell.execute_reply": "2026-07-10T05:09:09.721168Z"
}
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"--- 'I went to the park today and'\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
" C=-2 greedy : ' saw a large amount of wildlife. I noticed that the water surface was covered with sediment from the river. I used a remote sensing system to measure the sediment thickness. What is the sediment thickness in the'\n",
"--- 'The meeting this afternoon was'\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
" C=-2 greedy : ' held in the ______. The ______ was used for the meeting. The ______ was used for the meeting. The ______ was used for the meeting. \\n\\nFill in the blanks.\\n\\nThe ______ was used for'\n",
"--- 'My overall impression of the new apartment is that'\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
" C=-2 greedy : ' it is too dark. The windows are too dark. The walls are too dark. The ceiling is too dark. The insulation is too thick. The air quality is too poor. The thermal mass is'\n"
]
}
],
"source": [
"for p in PROMPTS:\n",
" print(f\"--- {p!r}\")\n",
" print(f\" C=-2 greedy : {gen(v, p, -2)!r}\")"
]
"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))"
},
{
"cell_type": "markdown",
"id": "e9cc185d",
"metadata": {},
"source": [
"## Bonus: lens readout\n",
"\n",
"Only the full-Jacobian cache gives you this: transport any layer's residual to\n",
"the final basis with `J_l` and decode it, i.e. \"what does the model think at\n",
"layer l\". SHOULD: at layer 16 the model has only a city-shaped slot, by layer\n",
"20 candidate cities appear, and by layer 24 Paris has won. ELSE layer indexing\n",
"is off."
]
"source": "## Bonus: lens readout\n\nOnly the full-Jacobian cache gives you this: transport any layer's residual to\nthe final basis with `J_l` and decode it, i.e. \"what does the model think at\nlayer l\". SHOULD: on a factual prompt, deeper layers resolve from a generic slot\n(e.g. \" city\") toward the specific answer (e.g. \" Paris\"). ELSE layer indexing\nis off."
},
{
"cell_type": "code",
"execution_count": 7,
"execution_count": null,
"id": "3009a178",
"metadata": {
"execution": {
@@ -399,22 +156,8 @@
"shell.execute_reply": "2026-07-10T05:09:09.819411Z"
}
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"layer 16: [' town', ' city', ' cities', '_city', ' City', ' Cities']\n",
"layer 20: [' Paris', ' London', ' Venice', ' France', ' Berlin', ' Vienna']\n",
"layer 24: [' Paris', 'Paris', ' cities', ' Cities', '巴黎', ' City']\n"
]
}
],
"source": [
"for layer in (16, 20, 24):\n",
" top = jac.lens_topk(model, tok, \"The Eiffel Tower is located in the city of\", layer=layer, k=6)\n",
" print(f\"layer {layer}: {[t for t, _ in top]}\")"
]
"outputs": [],
"source": "# Only the full-Jacobian cache gives this: transport a layer's residual to the\n# final basis and decode it, i.e. \"what does the model think at layer l\".\n# Pick a low / mid / high layer from the fitted band.\nlo, mid, hi = jac.layers[0], jac.layers[len(jac.layers) // 2], jac.layers[-1]\nfor layer in (lo, mid, hi):\n top = jac.lens_topk(model, tok, \"The Eiffel Tower is located in the city of\", layer=layer, k=6)\n print(f\"layer {layer}: {[t for t, _ in top]}\")"
},
{
"cell_type": "markdown",
@@ -429,7 +172,7 @@
},
{
"cell_type": "code",
"execution_count": 8,
"execution_count": null,
"id": "c3f5e64b",
"metadata": {
"execution": {
@@ -439,22 +182,8 @@
"shell.execute_reply": "2026-07-10T05:09:10.751806Z"
}
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
" I saw a lot of people there. I was happy with the food and the service. I think it's a good place to visit. I would like to go there again. I think it's\n"
]
}
],
"source": [
"from steering_lite import Vector\n",
"\n",
"v.save(\"../artifacts/happy_joy.safetensors\")\n",
"v2 = Vector.load(\"../artifacts/happy_joy.safetensors\")\n",
"print(gen(v2, \"I went to the park today and\", C=1))"
]
"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))"
}
],
"metadata": {
scripts/fit.py
+11 -11
View File
@@ -1,14 +1,14 @@
"""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:
`artifacts/qwen3.5-4b.jac`). Prompts are jlens's WikiText-103 corpus wrapped in
the model's chat template (config.chat_corpus): we fit J at the chat operating
point where steering is applied, matching the verified run-524. 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
uv run python scripts/fit.py --model Qwen/Qwen3.5-4B
uv run python scripts/fit.py --model Qwen/Qwen3-0.6B --dim-batch 8
"""
from __future__ import annotations
@@ -17,7 +17,6 @@ import sys
import time
from pathlib import Path
from jlens.examples import load_wikitext_prompts
from loguru import logger
from transformers import AutoModelForCausalLM, AutoTokenizer
@@ -28,9 +27,10 @@ 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("--model", default="Qwen/Qwen3.5-4B")
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("--dim-batch", type=int, default=4,
help="d_model dims per backward batch (memory knob; 4 fits a 4B on a 24GB 3090, 8+ for smaller)")
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)
@@ -43,9 +43,9 @@ def main() -> None:
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)")
f"dim_batch={args.dim_batch}, n_prompts={args.n_prompts} chat-templated WikiText)")
t0 = time.monotonic()
jac = Jacobian.fit_cached(model, tok, lambda: load_wikitext_prompts(args.n_prompts), out,
jac = Jacobian.fit_cached(model, tok, lambda: config.chat_corpus(tok, 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")))