wassname/jsteer
1
0
Fork 0
mirror of https://github.com/wassname/jsteer.git synced 2026-09-13 03:30:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

persona refinements: soft-cotangent vector, pinv tangent transport, topk word-like mask

Browse Source 
- persona_soft_vector: w = W_U^T (softmax(u_pos/T) - softmax(u_neg/T)) over
  word-like tokens; gradient of the expected-logprob contrast (genuine
  cotangent), full-vocab replacement for hard top-k; logs TV distance.
- persona_pinv_vector: h_diff is a tangent, so solve J delta = h_diff
  (ridge lstsq) instead of the J^T type error; logs per-layer residual.
- persona_topk_vector: mask emoji/special tokens out of the contrast
  (they were the degenerate emit-targets behind the C=1.5 emoji spam).

Co-Authored-By: Claudypoo <288921227+claudypoo@users.noreply.github.com>
This commit is contained in:
wassname
2026-07-11 08:05:01 +08:00
co-authored by Claudypoo
parent 143e9add80
commit 038a651197
2 changed files with 120 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
jsteer/applies.py
+22 -1
View File
@@ -30,6 +30,8 @@ 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
jacobian_persona_soft +C moves toward the POS persona's next-token dist
jacobian_persona_pinv +C moves toward the POS persona (tangent transport)
random norm-matched control, no meaning
"""
from __future__ import annotations
@@ -76,6 +78,24 @@ class JacobianPersonaTopkC(SteeringConfig):
apply_span: int = 1
@register_config
@dataclass
class JacobianPersonaSoftC(SteeringConfig):
method: str = "jacobian_persona_soft"
normalize: bool = True
apply_mode: str = "add"
apply_span: int = 1
@register_config
@dataclass
class JacobianPersonaPinvC(SteeringConfig):
method: str = "jacobian_persona_pinv"
normalize: bool = True
apply_mode: str = "add"
apply_span: int = 1
@register_config
@dataclass
class RandomC(SteeringConfig):
@@ -170,5 +190,6 @@ def _register_method(method_name: str) -> None:
apply = staticmethod(apply_dispatch)
for _name in ("jacobian_word", "jacobian_persona", "jacobian_persona_topk", "random"):
for _name in ("jacobian_word", "jacobian_persona", "jacobian_persona_topk",
"jacobian_persona_soft", "jacobian_persona_pinv", "random"):
_register_method(_name)
jsteer/jacobian.py
+98 -3
View File
@@ -21,6 +21,14 @@ build `w`:
persona_vector w = h_bar(pos) - h_bar(neg) EXPERIMENTAL*
persona_topk_vector w = top-k tokens of the pos-neg logit EXPERIMENTAL
contrast (differ before top-k, else null)
persona_soft_vector w = W_U^T (softmax contrast) -- topk's EXPERIMENTAL
full-vocab limit; a genuine cotangent
persona_topk/soft mask non-word-like tokens (emoji/specials) out of the
contrast: they are degenerate emit-targets (the emoji-spam failure mode).
persona_pinv_vector is NOT a pullback: h_bar(pos)-h_bar(neg) is a TANGENT
(an activation displacement), and J^T only transports cotangents
(gradients). It solves J_l delta = h_diff instead -- see its docstring.
* persona-contrast vectors FAILED specificity controls in j-steer-dev
(moved the target axis no more than an unrelated persona did). Shipped
@@ -48,6 +56,7 @@ 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 jlens.vis import _meaningful_token_mask # jlens's word-like vocab mask (cached)
from loguru import logger
from torch import Tensor
from tqdm.auto import tqdm
@@ -57,6 +66,8 @@ from steering_lite.vector import Vector
from .applies import (
JacobianPersonaC,
JacobianPersonaPinvC,
JacobianPersonaSoftC,
JacobianPersonaTopkC,
JacobianWordC,
RandomC,
@@ -289,8 +300,12 @@ class Jacobian:
logits_pos = lm.unembed(h_pos.to(model.device).to(model.dtype)).float()
logits_neg = lm.unembed(h_neg.to(model.device).to(model.dtype)).float()
diff = logits_pos - logits_neg # [vocab]
top_pos = diff.topk(k) # pos evokes > neg
top_neg = (-diff).topk(k) # neg evokes > pos
# word-like tokens only: emoji/special/punct tokens are degenerate
# emit-targets -- steering toward them collapses generation into
# repeating them (the emoji-spam failure mode at higher C). (Claude)
wordlike = _meaningful_token_mask(tok, diff.shape[-1], diff.device)
top_pos = diff.masked_fill(~wordlike, -torch.inf).topk(k) # pos evokes > neg
top_neg = (-diff).masked_fill(~wordlike, -torch.inf).topk(k) # neg evokes > pos
toks_pos = [tok.decode([i]) for i in top_pos.indices.tolist()]
toks_neg = [tok.decode([i]) for i in top_neg.indices.tolist()]
# SHOULD be persona-specific words (positive vs negative affect here), not
@@ -303,6 +318,87 @@ class Jacobian:
cfg = JacobianPersonaTopkC(layers=self._steer_layers(layers))
return self.pullback(w, cfg)
def persona_soft_vector(self, model, tok, pos_prompts: list[str],
neg_prompts: list[str], *, temperature: float = 1.0,
layers=None, batch_size: int = 8) -> Vector:
"""EXPERIMENTAL: persona_topk's full-vocab limit, and a genuine cotangent.
w = W_U^T (softmax(u_pos/T) - softmax(u_neg/T))
where u = unembedded persona mean. This is exactly the gradient wrt the
final residual of E_{t~p_pos}[log p(t|h)] - E_{t~p_neg}[log p(t|h)]
(each term's softmax-baseline E_p[W_U] cancels in the difference), so
unlike persona_vector's activation diff, J^T transports it legitimately.
Vs topk: no hard k=8 compression to the personas' most extreme tokens
(the over-literal "emit :-)" failure); every tone-correlated token
contributes, weighted by how much the personas disagree on it.
`temperature` subsumes k: low T -> topk-like sparsity, high T -> broad
support. Non-word-like tokens are masked out before the softmax, same
rationale as topk. Untested for specificity. (Claude)"""
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]
u_pos = lm.unembed(h_pos.to(model.device).to(model.dtype)).float()
u_neg = lm.unembed(h_neg.to(model.device).to(model.dtype)).float()
wordlike = _meaningful_token_mask(tok, u_pos.shape[-1], u_pos.device)
p_pos = u_pos.masked_fill(~wordlike, -torch.inf).div(temperature).softmax(-1)
p_neg = u_neg.masked_fill(~wordlike, -torch.inf).div(temperature).softmax(-1)
Δp = p_pos - p_neg # [vocab], sums to 0
# read your data: TV distance = how much the personas disagree about the
# next token at all. SHOULD be clearly > 0 (~0 => null contrast, same
# failure mode as topk's identical token sets); top tokens SHOULD be
# persona-specific words, not generic sentence-starters.
tv = 0.5 * float(Δp.abs().sum())
top_pos, top_neg = Δp.topk(8), (-Δp).topk(8)
logger.info(
f"j-thoughts (soft, T={temperature}) TV(p_pos, p_neg)={tv:.3f}\n"
f" positive: {[tok.decode([i]) for i in top_pos.indices.tolist()]}\n"
f" negative: {[tok.decode([i]) for i in top_neg.indices.tolist()]}")
w = (Δp @ W_U.float()).cpu() # raw W_U rows, matching _word_cotangent
cfg = JacobianPersonaSoftC(layers=self._steer_layers(layers))
return self.pullback(w, cfg)
def persona_pinv_vector(self, model, tok, pos_prompts: list[str],
neg_prompts: list[str], *, ridge: float = 1e-3,
layers=None, batch_size: int = 8) -> Vector:
"""EXPERIMENTAL: transport the persona contrast as a TANGENT, which it is.
h_diff = h_bar(pos) - h_bar(neg) is an activation DISPLACEMENT at the
final layer, not a gradient -- persona_vector's J^T h_diff pulls it back
as if it were a cotangent, a type error (only correct if J were
orthogonal). The right transport asks: which layer-l perturbation
delta pushes forward to h_diff?
delta_l = argmin |J_l delta - h_diff|^2 + lam |delta|^2
= (J_l^T J_l + lam I)^{-1} J_l^T h_diff, lam = ridge * mean diag(J^T J)
Ridge because the position-averaged J is ill-conditioned. If THIS still
fails specificity, the failure is the averaged Jacobian itself (it can't
carry contextual features), not the algebra. CPU fp32 against the cached
matrices, ~seconds per layer, no backward. (Claude)"""
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")
h_diff = (h_pos - h_neg).float().cpu()
logger.info(f"h_bar_diff |pos|={h_pos.norm():.3f} |neg|={h_neg.norm():.3f} "
f"|diff|={h_diff.norm():.3f}")
cfg = JacobianPersonaPinvC(layers=self._steer_layers(layers))
per_layer, residuals = {}, {}
eye = torch.eye(self.lens.d_model)
for l in cfg.layers:
J = self.lens.jacobians[l] # [d_out, d_in] fp32 cpu
JtJ = J.T @ J
lam = ridge * JtJ.diagonal().mean()
δ = torch.linalg.solve(JtJ + lam * eye, J.T @ h_diff)
per_layer[l] = δ
residuals[l] = float((J @ δ - h_diff).norm() / h_diff.norm())
# SHOULD be well below 1.0 at most layers: 1.0 means J can't realize
# h_diff at all (h_diff orthogonal to J's row space) and the vector is
# ridge-noise; small residual means the transport is faithful.
logger.info("pinv relative residual |J d - h|/|h| per layer: " +
" ".join(f"{l}:{residuals[l]:.2f}" for l in cfg.layers))
return _to_vector(cfg, per_layer)
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."""
@@ -323,7 +419,6 @@ class Jacobian:
`mask_wordlike` reuses jlens's own word-like vocab mask so the readout
hides punctuation/single-char/special tokens (which, per the walkthrough,
trail the interesting word tokens on Qwen); ranks are unaffected."""
from jlens.vis import _meaningful_token_mask
lm = from_hf(model, tok)
lens_logits, _, _ = self.lens.apply(lm, prompt, layers=[layer],
positions=[position])