commit 1f458fd53eefcddf302858c9829d279f1d6687f6 Author: wassname <1103714+wassname@users.noreply.github.com> Date: Fri Jul 10 12:08:07 2026 +0800 core algo: Jacobian fit/cache/pullback (jlens) + word/persona vectors (steering-lite runtime) Ported from the verified j-steer-dev experiment (word pullback beat random on 3/5 moral foundations, Qwen3-4B n=3). jacobian.py wraps jlens fit/save/ load and derives steering vectors by CPU matvec against the cached J; vjp.py is the one-backward parity path; applies.py registers the methods and delivery modes into steering-lite. Co-Authored-By: Claudypoo <288921227+claudypoo@users.noreply.github.com> diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..81900be --- /dev/null +++ b/.gitignore @@ -0,0 +1,6 @@ +.venv/ +__pycache__/ +*.egg-info/ +artifacts/*.jac +artifacts/*.ckpt +uv.lock diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..9ecdb0d --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,43 @@ +# jsteer -- agent notes + +Plan of record: /home/wassname/.claude/plans/review-specs-00-minimal-experiment-md-an-peppy-sky.md +Evidence base: ../j-steer-dev/docs/RESEARCH_JOURNAL.md (verified 3/5 word-steering result). + +## What this is + +repeng-style UX for Jacobian pullback steering. Core algo files are WRITTEN +(by the main agent, ported from verified j-steer-dev code -- do not rewrite +the math, it is parity-gated against the verified experiment): + +- `jsteer/jacobian.py` -- Jacobian.fit/save/load (wraps jlens, the + researchers' verified primary code) + word/persona/persona_topk/random + vectors -> steering_lite.Vector. +- `jsteer/applies.py` -- steering-lite method registration + delivery modes. +- `jsteer/vjp.py` -- direct VJP path; the parity reference for the cache. + +Runtime is steering-lite: `with v(model, C=8): model.generate(...)`. + +## Remaining work (task list has details; U-numbers from the plan) + +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. + +## Style + +Fail fast, no defensive programming, loguru, no LLM-tell prose in README. +Comments marked as Claude-authored where opinionated. diff --git a/jsteer/__init__.py b/jsteer/__init__.py new file mode 100644 index 0000000..388fe95 --- /dev/null +++ b/jsteer/__init__.py @@ -0,0 +1,13 @@ +"""jsteer: fit a model's full Jacobian once, then steer any word or persona. + + from jsteer import Jacobian + jac = Jacobian.fit(model, tok, prompts) # expensive, once, cacheable + v = jac.word_vector(model, tok, ["authority"]) # instant matvec + with v(model, C=8): # steering-lite runtime + model.generate(**inputs) +""" +from . import applies # noqa: F401 -- registers methods into steering-lite's REGISTRY +from .jacobian import Jacobian +from .vjp import pullback_vjp, word_vector_vjp + +__all__ = ["Jacobian", "pullback_vjp", "word_vector_vjp"] diff --git a/jsteer/applies.py b/jsteer/applies.py new file mode 100644 index 0000000..462151a --- /dev/null +++ b/jsteer/applies.py @@ -0,0 +1,159 @@ +"""steering-lite method registration + modular delivery of jacobian vectors. + +(drafted by Claude, ported from the verified j-steer-dev experiment code) + +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 +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 +stubs that say so. + +DELIVERY of v to the residual stream is decoupled from extraction: the same v +can be added everywhere, gated to the last positions, or overwrite a span. +`cfg.apply_mode` selects the mode; adding a mode = one function + one +APPLY_REGISTRY entry. + +Protocol (steering-lite config.py Method.apply): + apply(mod, x, y, shared, stacked, cfg) -> y_new # same shape [b, s, d] + +Position semantics: generation uses LEFT padding, so the last real token is +position -1; `cfg.apply_span` selects how many trailing positions to target. +Note that during incremental generation (KV cache) every decode step has s=1, +so add_last touches each generated token but only the tail of the prefill. + +Sign conventions (set by the cotangent in jacobian.py): + jacobian_word +C raises the words' output logits + jacobian_persona +C moves toward the POS persona + jacobian_persona_topk +C moves toward the POS persona's evoked vocabulary + random norm-matched control, no meaning +""" +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass + +import torch +from torch import Tensor + +from steering_lite.config import SteeringConfig, register, register_config + +ε = 1e-8 + + +# --- configs ----------------------------------------------------------------- +# One config per method name so saved vectors deserialize with provenance. +# All share the same knobs: apply_mode (delivery) + apply_span (tail width). + +@register_config +@dataclass +class JacobianWordC(SteeringConfig): + method: str = "jacobian_word" + normalize: bool = True + apply_mode: str = "add" + apply_span: int = 1 + + +@register_config +@dataclass +class JacobianPersonaC(SteeringConfig): + method: str = "jacobian_persona" + normalize: bool = True + apply_mode: str = "add" + apply_span: int = 1 + + +@register_config +@dataclass +class JacobianPersonaTopkC(SteeringConfig): + method: str = "jacobian_persona_topk" + normalize: bool = True + apply_mode: str = "add" + apply_span: int = 1 + + +@register_config +@dataclass +class RandomC(SteeringConfig): + method: str = "random" + normalize: bool = True + apply_mode: str = "add" + apply_span: int = 1 + + +# --- delivery modes ---------------------------------------------------------- + +def _v_sum(stacked: dict[str, Tensor], y: Tensor) -> Tensor: + """The per-layer direction v, summed over the k-stack, on y's device/dtype.""" + return stacked["v"].to(y).sum(dim=0) # [d] + + +def apply_add(mod, x, y, shared, stacked, cfg) -> Tensor: + """y += coeff * v at ALL positions (the verified default; same delivery as + steering-lite's mean_diff, so calibrated coeffs are comparable).""" + v = _v_sum(stacked, y) + return y + cfg.coeff * v + + +def apply_add_last(mod, x, y, shared, stacked, cfg) -> Tensor: + """y[:, -k:] += coeff * v -- nudge only the last k positions (decision + region). k = cfg.apply_span; k >= s degenerates to apply_add.""" + v = _v_sum(stacked, y) + k = cfg.apply_span + head = y[:, :-k, :] + tail = y[:, -k:, :] + cfg.coeff * v + return torch.cat([head, tail], dim=1) + + +def apply_replace_last(mod, x, y, shared, stacked, cfg) -> Tensor: + """Overwrite the last k positions with the concept direction at each + position's original magnitude: energy from y, direction from v, strength + from coeff. A "virtual token" injection that keeps the [b, s, d] shape + (true sequence insertion would break RoPE / attention mask / KV cache).""" + v = _v_sum(stacked, y) + v_unit = v / (v.norm() + ε) + k = cfg.apply_span + tail = y[:, -k:, :] # [b, k, d] + energy = tail.norm(dim=-1, keepdim=True) # [b, k, 1] + new_tail = energy * (cfg.coeff * v_unit) + head = y[:, :-k, :] + return torch.cat([head, new_tail], dim=1) + + +APPLY_REGISTRY: dict[str, Callable[..., Tensor]] = { + "add": apply_add, + "add_last": apply_add_last, + "replace_last": apply_replace_last, +} + + +def apply_dispatch(mod, x, y, shared, stacked, cfg) -> Tensor: + fn = APPLY_REGISTRY.get(cfg.apply_mode) + if fn is None: + raise KeyError( + f"unknown apply_mode={cfg.apply_mode!r}; registered: {list(APPLY_REGISTRY)}") + return fn(mod, x, y, shared, stacked, cfg) + + +# --- method registration ----------------------------------------------------- + +def _extract_stub(pos_acts, neg_acts, cfg): + raise NotImplementedError( + "jsteer methods are extracted via Jacobian.{word,persona,persona_topk}_vector " + "(jacobian.py) or vjp.py, not steering_lite.train -- they need gradients " + "that train's no_grad path can't provide.") + + +def _register_method(method_name: str) -> None: + @register + class _M: # noqa: N801 -- registry keys on .name, class name irrelevant + name = method_name + extract = staticmethod(_extract_stub) + apply = staticmethod(apply_dispatch) + + +for _name in ("jacobian_word", "jacobian_persona", "jacobian_persona_topk", "random"): + _register_method(_name) diff --git a/jsteer/jacobian.py b/jsteer/jacobian.py new file mode 100644 index 0000000..519886c --- /dev/null +++ b/jsteer/jacobian.py @@ -0,0 +1,264 @@ +"""Full-Jacobian extraction for steering: fit once, steer any concept. + +(drafted by Claude, ported from the verified j-steer-dev experiment code) + +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: + + 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 + (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. + +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: + + with v(model, C=8): + model.generate(**inputs) +""" +from __future__ import annotations + +from dataclasses import dataclass + +import torch +from jlens.fitting import fit as _jlens_fit +from jlens.hf import HFLensModel, from_hf +from jlens.hooks import ActivationRecorder +from jlens.lens import JacobianLens +from loguru import logger +from torch import Tensor +from tqdm.auto import tqdm + +from steering_lite.config import SteeringConfig +from steering_lite.vector import Vector + +from .applies import ( + JacobianPersonaC, + JacobianPersonaTopkC, + JacobianWordC, + RandomC, + ε, +) + + +# --- small helpers ------------------------------------------------------------ + +def _unit(v: Tensor) -> Tensor: + return v / (v.norm() + ε) + + +def _resolve_layers(layers, n_layers: int) -> list[int]: + """None -> all layers below final; (lo, hi) floats -> fraction band; + iterable of ints -> as-is (negatives count from the end).""" + if layers is None: + return list(range(n_layers - 1)) + layers = tuple(layers) + if len(layers) == 2 and all(isinstance(x, float) for x in layers): + lo, hi = (int(layers[0] * n_layers), int(layers[1] * n_layers)) + return list(range(lo, min(hi, n_layers - 1))) + return sorted({l + n_layers if l < 0 else l for l in layers}) + + +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).""" + shared = {l: {} for l in per_layer} + stacked = {l: {"v": _unit(v.float()).unsqueeze(0)} for l, v in per_layer.items()} + return Vector(cfg, shared, stacked) + + +def _word_cotangent(model, tok, words: list[str]) -> Tensor: + """Mean unembedding row over `words` (first sub-token of each): the + final-basis direction that most raises those tokens' output logits. + Pulling THIS back through J^T is the pure concept->residual map -- no + persona pairs, so no persona-bundle confound. +C enhances the concept.""" + W_U = model.lm_head.weight # [vocab, d] + ids = [tok(w, add_special_tokens=False).input_ids[0] for w in words] + rows = W_U[torch.tensor(ids, device=W_U.device)].float() # [n, d] + cot = rows.mean(0) + logger.info(f"word cotangent: {words} -> first-subtoken ids={ids} |w|={cot.norm():.3f}") + return cot + + +@torch.no_grad() +def _h_bar_final(model, tok, prompts: list[str], *, batch_size: int = 8, + max_length: int = 384, label: str = "") -> Tensor: + """Mean last-token activation at the FINAL layer over `prompts` + (right-padded batches; last real token located via the attention mask).""" + lm = from_hf(model, tok) # layout detection only + target_layer = lm.n_layers - 1 + acc, n = None, 0 + for i in tqdm(range(0, len(prompts), batch_size), desc=f"h_bar {label}", + mininterval=30, maxinterval=60): + batch = prompts[i:i + batch_size] + enc = tok(batch, return_tensors="pt", padding=True, truncation=True, + max_length=max_length, padding_side="right").to(model.device) + with ActivationRecorder(lm.layers, at=[target_layer]) as rec: + model(**enc) + act = rec.activations[target_layer] # [B, S, d] + last_idx = enc["attention_mask"].sum(dim=1) - 1 # [B] + last = act[torch.arange(act.shape[0]), last_idx].float() + acc = last.sum(0) if acc is None else acc + last.sum(0) + n += last.shape[0] + return acc / n + + +# --- the core object ---------------------------------------------------------- + +@dataclass +class Jacobian: + """A fitted, cached, model-specific bundle of per-layer Jacobians + `{layer: J_l [d, d]}` (a `jlens.JacobianLens`) plus the concept->vector + methods. Fit once (expensive), derive steering vectors forever (matvec).""" + + lens: JacobianLens + + # -- fit / persist (all delegated to the researchers' jlens code) ---------- + + @classmethod + def fit(cls, model, tok, prompts: list[str], *, layers=None, dim_batch: int = 8, + max_seq_len: int = 128, checkpoint_path: str | None = None, + compile: bool = False) -> "Jacobian": + """Fit `J_l` on `prompts` (generic text; jlens guidance is ~100+ for + lens-quality pooling). Cost: 1 forward + ceil(d_model/dim_batch) + backwards per prompt -- ALL source layers come from the same backwards, + so fitting more layers costs memory, not compute. `checkpoint_path` + makes the fit resumable (atomic writes).""" + lm = from_hf(model, tok, compile=compile) + source_layers = _resolve_layers(layers, lm.n_layers) + lens = _jlens_fit(lm, prompts, source_layers=source_layers, + dim_batch=dim_batch, max_seq_len=max_seq_len, + checkpoint_path=checkpoint_path) + return cls(lens=lens) + + def save(self, path: str) -> None: + self.lens.save(path) # fp16 by default; jlens-compatible file + + @classmethod + def load(cls, path: str) -> "Jacobian": + return cls(lens=JacobianLens.load(path)) + + @classmethod + def from_pretrained(cls, name_or_path: str, **kw) -> "Jacobian": + """Local file/dir or HuggingFace Hub repo_id (see JacobianLens).""" + return cls(lens=JacobianLens.from_pretrained(name_or_path, **kw)) + + @property + def layers(self) -> list[int]: + return self.lens.source_layers + + def __repr__(self) -> str: + return f"Jacobian({self.lens!r})" + + # -- the one pullback ------------------------------------------------------- + + def pullback(self, cotangent: Tensor, cfg: SteeringConfig) -> Vector: + """v_l = unit(J_l^T @ w) for every layer in cfg.layers. + + J_l rows are output dims (each row is the gradient of one final-basis + dim over the source layer), so the pullback is `w @ J_l`. Computed on + CPU fp32 against the cached matrices -- no model, no backward.""" + w = cotangent.detach().float().cpu() + if w.shape != (self.lens.d_model,): + raise ValueError(f"cotangent shape {tuple(w.shape)} != ({self.lens.d_model},)") + missing = set(cfg.layers) - set(self.lens.source_layers) + if missing: + raise ValueError(f"layers {sorted(missing)} not fitted; have {self.layers}") + per_layer = {l: w @ self.lens.jacobians[l] for l in cfg.layers} + logger.info(f"{cfg.method} per-layer |J^T w| (pre-norm): " + + " ".join(f"{l}:{per_layer[l].norm():.3g}" for l in cfg.layers)) + return _to_vector(cfg, per_layer) + + def _steer_layers(self, layers) -> tuple[int, ...]: + """None -> all fitted layers; else explicit int indices (float bands are + a fit-time concept -- the lens doesn't know n_layers to resolve them).""" + if layers is None: + return tuple(self.lens.source_layers) + return tuple(sorted(int(l) for l in layers)) + + # -- concept -> vector ------------------------------------------------------- + + def word_vector(self, model, tok, words: list[str], *, layers=None) -> Vector: + """VERIFIED method: pull the words' unembedding direction back through + the Jacobian. +C enhances the concept, -C suppresses it.""" + cfg = JacobianWordC(layers=self._steer_layers(layers)) + return self.pullback(_word_cotangent(model, tok, words), cfg) + + def persona_vector(self, model, tok, pos_prompts: list[str], + neg_prompts: list[str], *, layers=None, + batch_size: int = 8) -> Vector: + """EXPERIMENTAL: pull the persona activation contrast back through the + Jacobian. Persona-contrast pullbacks failed specificity controls in + j-steer-dev -- prefer word_vector for targeted steering.""" + h_pos = _h_bar_final(model, tok, pos_prompts, batch_size=batch_size, label="pos") + h_neg = _h_bar_final(model, tok, neg_prompts, batch_size=batch_size, label="neg") + logger.info(f"h_bar_diff |pos|={h_pos.norm():.3f} |neg|={h_neg.norm():.3f} " + f"|diff|={ (h_pos - h_neg).norm():.3f}") + cfg = JacobianPersonaC(layers=self._steer_layers(layers)) + return self.pullback(h_pos - h_neg, cfg) + + def persona_topk_vector(self, model, tok, pos_prompts: list[str], + neg_prompts: list[str], *, k: int = 8, layers=None, + batch_size: int = 8) -> Vector: + """EXPERIMENTAL: persona -> vocabulary bottleneck -> word pullback. + Read each persona's final-layer mean through the unembedding, take the + top-k tokens it most evokes, contrast the two token sets' unembedding + rows, pull that back. Composes the persona signal with the verified + word mechanism; untested for specificity.""" + lm = from_hf(model, tok) + h_pos = _h_bar_final(model, tok, pos_prompts, batch_size=batch_size, label="pos") + h_neg = _h_bar_final(model, tok, neg_prompts, batch_size=batch_size, label="neg") + W_U = model.lm_head.weight # [vocab, d] + cots = {} + for name, h in (("pos", h_pos), ("neg", h_neg)): + logits = lm.unembed(h.to(model.device).to(model.dtype)).float() + 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 + 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) + + def random_vector(self, *, seed: int = 0, layers=None) -> Vector: + """Norm-matched control: unit random direction per layer. Any honest + demo/eval should show the concept vector beating THIS at the same C.""" + gen = torch.Generator().manual_seed(seed) + cfg = RandomC(layers=self._steer_layers(layers), seed=seed) + per_layer = {l: torch.randn(self.lens.d_model, generator=gen) + for l in cfg.layers} + return _to_vector(cfg, per_layer) + + # -- bonus: the lens's native forward readout -------------------------------- + + def lens_topk(self, model, tok, prompt: str, layer: int, *, k: int = 10, + position: int = -1) -> list[tuple[str, float]]: + """What the model 'thinks' at `layer`: transport the residual to the + final basis with J_l and decode. jlens's native use, handy in demos.""" + lm = from_hf(model, tok) + lens_logits, _, _ = self.lens.apply(lm, prompt, layers=[layer], + positions=[position]) + top = lens_logits[layer][0].topk(k) + return [(tok.decode([i]), float(v)) for i, v in + zip(top.indices.tolist(), top.values.tolist())] diff --git a/jsteer/vjp.py b/jsteer/vjp.py new file mode 100644 index 0000000..21961c0 --- /dev/null +++ b/jsteer/vjp.py @@ -0,0 +1,93 @@ +"""Direct VJP pullback: one concept without paying for the full Jacobian. + +(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). + +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 +position of the final layer, gradient read at every valid source position and +meaned, positions before skip_first=16 excluded (attention sinks), final +position excluded (no next-token target). Right-padded batches: the valid mask +excludes pads, so batching does not change the estimate. +""" +from __future__ import annotations + +import torch +from jlens.fitting import SKIP_FIRST_N_POSITIONS +from jlens.hf import from_hf +from jlens.hooks import ActivationRecorder +from loguru import logger +from torch import Tensor +from tqdm.auto import tqdm + +from steering_lite.vector import Vector + +from .applies import JacobianWordC +from .jacobian import _resolve_layers, _to_vector, _word_cotangent + + +def _valid_mask(attention_mask: Tensor, skip_first: int) -> Tensor: + """Boolean [B, S]: right-padded real tokens in [skip_first : real_len-1]. + jlens valid_position_mask extended to a right-padded batch.""" + 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) + return mask & attention_mask.bool() + + +def pullback_vjp(model, tok, prompts: list[str], layers, cotangent: Tensor, *, + batch_size: int = 8, max_length: int = 128, + skip_first: int = SKIP_FIRST_N_POSITIONS) -> dict[int, Tensor]: + """Per-layer mean over `prompts` of J_l^T @ cotangent, via one backward per + batch (grads for every source layer come from the same backward).""" + lm = from_hf(model, tok) # freezes params, locates blocks; grads flow to + target_layer = lm.n_layers - 1 # activations only + layers = _resolve_layers(layers, lm.n_layers) + assert max(layers) < target_layer, f"source layers {layers} must be < {target_layer}" + d = cotangent.shape[0] + G = {l: torch.zeros(d, dtype=torch.float32, device=model.device) for l in layers} + count = 0 + for i in tqdm(range(0, len(prompts), batch_size), desc="pullback_vjp", + mininterval=30, maxinterval=60): + batch = prompts[i:i + batch_size] + enc = tok(batch, return_tensors="pt", padding=True, truncation=True, + max_length=max_length, padding_side="right").to(model.device) + valid = _valid_mask(enc["attention_mask"], skip_first) # [B, S] bool + if valid.sum(dim=1).min() == 0: + raise ValueError(f"a prompt has 0 valid positions " + f"(too short for skip_first={skip_first})") + with ActivationRecorder(lm.layers, at=[*layers, target_layer], + start_graph_at=min(layers)) as rec, torch.enable_grad(): + model(**enc) + h_final = rec.activations[target_layer] # [B, S, d] + srcs = [rec.activations[l] for l in layers] + c = (cotangent.to(h_final.device).to(h_final.dtype).view(1, 1, d) + * valid.unsqueeze(-1)) + grads = torch.autograd.grad(h_final, srcs, grad_outputs=c) + den = valid.sum(dim=1, keepdim=True).float() # [B, 1] + for l, g in zip(layers, grads): + v_b = (g.float() * valid.unsqueeze(-1)).sum(dim=1) / den # [B, d] + G[l] += v_b.sum(0) + count += len(batch) + return {l: G[l] / count for l in layers} + + +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 + 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) + G = pullback_vjp(model, tok, prompts, layers, cot, batch_size=batch_size, + max_length=max_length, skip_first=skip_first) + logger.info("word_vector_vjp per-layer |v| (pre-norm): " + + " ".join(f"{l}:{v.norm():.3g}" for l, v in G.items())) + return _to_vector(JacobianWordC(layers=tuple(G)), G) diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..da85ab3 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,31 @@ +[project] +name = "jsteer" +version = "0.1.0" +description = "Fit a model's full Jacobian once (jlens), then steer any word or persona (steering-lite runtime). repeng-style UX for Jacobian pullback steering." +readme = "README.md" +requires-python = ">=3.11" +license = "MIT" +dependencies = [ + "torch>=2.1", + "transformers>=4.51", + "accelerate>=1.6", + "jlens", # the researchers' verified Jacobian estimator + cache format + "steering-lite", # hook runtime: Vector, attach, `with v(model, C=...)`, calibrate + "loguru>=0.7", + "tqdm>=4.66", +] + +[project.optional-dependencies] +notebooks = ["matplotlib>=3.8", "ipykernel", "tabulate"] +test = ["pytest"] + +[tool.uv.sources] +jlens = { git = "https://github.com/anthropics/jacobian-lens" } +steering-lite = { git = "https://github.com/wassname/steering-lite" } + +[build-system] +requires = ["setuptools>=68"] +build-backend = "setuptools.build_meta" + +[tool.setuptools] +packages = ["jsteer"]