From 57c8d4b1661ba5589553604a9d0b80e1eb79371d Mon Sep 17 00:00:00 2001
From: wassname <1103714+wassname@users.noreply.github.com>
Date: Fri, 10 Jul 2026 14:42:22 +0800
Subject: [PATCH] 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>
---
 jsteer/applies.py | 20 +++++++++++++++++---
 scripts/smoke.py  | 22 +++++++++++++---------
 2 files changed, 30 insertions(+), 12 deletions(-)

diff --git a/jsteer/applies.py b/jsteer/applies.py
index 1ea4e21..fb3d5d1 100644
--- a/jsteer/applies.py
+++ b/jsteer/applies.py
@@ -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,
 }
diff --git a/scripts/smoke.py b/scripts/smoke.py
index d9ecb21..1b03c79 100644
--- a/scripts/smoke.py
+++ b/scripts/smoke.py
@@ -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__":