mirror of
https://github.com/wassname/jsteer.git
synced 2026-08-21 06:10:22 +08:00
add clamp apply mode: pin v-component to C instead of accumulating
clamp: y += (C - <y,v_hat>)v_hat at all positions -- bounded perturbation regardless of generation length, vs add's per-step accumulation via KV cache. C=0 is directional ablation. Smoke (Qwen3-0.6B, happy/joy): clamp C=+20 stays coherent and on-concept (drifts to 'happiness and joy of my childhood', in Chinese) while add C=+8 already degenerates to 'joyjoyjoy...'. Co-Authored-By: Claudypoo <288921227+claudypoo@users.noreply.github.com>
This commit is contained in:
+17
-3
@@ -14,9 +14,9 @@ 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
|
||||
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.
|
||||
can be added everywhere, clamped to a target component, 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]
|
||||
@@ -109,6 +109,19 @@ def apply_add_last(mod, x, y, shared, stacked, cfg) -> Tensor:
|
||||
return torch.cat([head, tail], dim=1)
|
||||
|
||||
|
||||
def apply_clamp(mod, x, y, shared, stacked, cfg) -> Tensor:
|
||||
"""Set y's component along v_hat to coeff at ALL positions:
|
||||
y += (coeff - <y, v_hat>) * v_hat. Unlike apply_add the perturbation stays
|
||||
bounded however long generation runs: each decode step re-targets the same
|
||||
component value instead of pushing again on top of the previous push (the
|
||||
compounding that degenerates high-|C| adds via the KV cache). coeff=0 is
|
||||
directional ablation (Arditi et al. 2024); -C reverses the component."""
|
||||
v = _v_sum(stacked, y)
|
||||
v_hat = v / (v.norm() + ε)
|
||||
comp = torch.einsum("bsd,d->bs", y, v_hat).unsqueeze(-1) # [b, s, 1]
|
||||
return y + (cfg.coeff - comp) * v_hat
|
||||
|
||||
|
||||
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
|
||||
@@ -126,6 +139,7 @@ def apply_replace_last(mod, x, y, shared, stacked, cfg) -> Tensor:
|
||||
|
||||
APPLY_REGISTRY: dict[str, Callable[..., Tensor]] = {
|
||||
"add": apply_add,
|
||||
"clamp": apply_clamp,
|
||||
"add_last": apply_add_last,
|
||||
"replace_last": apply_replace_last,
|
||||
}
|
||||
|
||||
+13
-9
@@ -77,18 +77,22 @@ def main() -> None:
|
||||
_show_tokens(tok, GEN_PROMPT, "GEN PROMPT")
|
||||
enc = tok(GEN_PROMPT, return_tensors="pt").to(DEVICE)
|
||||
|
||||
for C in (-8, 0, 8):
|
||||
with v(model, C=C):
|
||||
out = model.generate(**enc, max_new_tokens=40, do_sample=False,
|
||||
pad_token_id=tok.eos_token_id)
|
||||
text = tok.decode(out[0][enc.input_ids.shape[1]:], skip_special_tokens=True)
|
||||
logger.info(f"=== C={C:+d} generation ===\n{text!r}")
|
||||
for mode in ("add", "clamp"):
|
||||
v.cfg.apply_mode = mode
|
||||
for C in (-8, 0, 8, 20):
|
||||
with v(model, C=C):
|
||||
out = model.generate(**enc, max_new_tokens=40, do_sample=False,
|
||||
pad_token_id=tok.eos_token_id)
|
||||
text = tok.decode(out[0][enc.input_ids.shape[1]:], skip_special_tokens=True)
|
||||
logger.info(f"=== {mode} C={C:+d} generation ===\n{text!r}")
|
||||
|
||||
logger.info(
|
||||
"SHOULD: C=+8 mentions happiness/joy more than C=0; C=-8 less or "
|
||||
"negative tone. ELSE steering wiring or sign issue. All three SHOULD "
|
||||
"stay coherent english; gibberish means the coeff is too large or the "
|
||||
"vector is malformed.")
|
||||
"negative tone. ELSE steering wiring or sign issue. add C=+20 MAY "
|
||||
"degenerate (unbounded accumulation); clamp C=+20 SHOULD stay more "
|
||||
"coherent (component pinned, perturbation bounded). clamp C=0 is "
|
||||
"directional ablation, expect near-baseline text. Gibberish at small "
|
||||
"|C| means the vector is malformed.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
Reference in New Issue
Block a user