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:
wassname
2026-07-10 14:42:22 +08:00
co-authored by Claudypoo
parent 474f74ac33
commit 57c8d4b166
2 changed files with 30 additions and 12 deletions
+17 -3
View File
@@ -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
View File
@@ -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__":