evo2 SAE steering analysis: dose-response/selectivity harness, rebased onto migrated #1622

Clamp an SAE feature via the production Evo2SAE.generate path and quantify the causal effect:
dose-response (effect vs strength) + selectivity (target vs control features), persisted to a
structured steering_results.json.

Review fixes:
  * metric: replace positional Hamming with normalized edit (Levenshtein) distance. Greedy
    decode is autoregressive, so one early flipped token shifts every downstream base and pins
    Hamming at ~1.0 — erasing the dose curve. Edit distance is shift-robust; first_divergence
    (shared-prefix length) is the complementary monotone signal. Tested with the shift case.
  * surface the clamp cap: generate() silently caps |strength| to MAX_CLAMP_STRENGTH, so two
    requests above it produce an identical clamp (a fake plateau). run_steering now warns,
    steers at the effective value, and records max_clamp_strength + capped_strengths.
  * fix dangling doc reference (probe.py -> extract.py, which exists).
  * refactor steer.py into injectable pick_target()/run_steering() and add CPU test_steer.py
    (fake engine, local not in conftest) covering target picking, dose monotonicity,
    selectivity, and cap reporting.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Polina Binder <pbinder@nvidia.com>

Author: polinabinder1

interpretability/sparse_autoencoders/recipes/evo2/scripts/steer.py

@@ -0,0 +1,178 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LicenseRef-Apache2
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+r"""Evo2 SAE steering harness — clamp a feature and measure the causal effect on generation.
+
+Reuses ``Evo2SAE.generate`` (the production decode-only ``sae.steering`` clamp the server/CLI
+use), then quantifies the effect with the pure, CPU-tested metrics in ``steer_analysis``:
+**dose-response** (effect vs clamp strength) for a target feature and **selectivity** (target
+vs control features). Writes a structured ``steering_results.json`` so the evidence is
+persisted and reproducible — the steering analog of how ``extract.py`` persists its outputs.
+
+GPU harness — run on an H100 with the inference engine available; this is not a CPU unit test.
+
+    python steer.py --evo2-ckpt-dir <mbridge> --sae-checkpoint <sae.pt> --layer 26 \
+        --sequence ATGGCC... --feature 29244 --controls 12345,54321 --strengths 0,50,100,200 \
+        --out steering_results.json
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+
+_HERE = Path(__file__).resolve().parent
+sys.path.insert(0, str(_HERE))
+sys.path.insert(0, str(_HERE.parent / "src"))  # recipes/evo2/src -> evo2_sae package
+sys.path.insert(0, str(_HERE.parents[2] / "sae" / "src"))
+
+from steer_analysis import dose_response, selectivity  # noqa: E402
+
+
+def pick_target(eng, sequence: str, feature: int | None) -> tuple[int, list[tuple[int, str | None, float]]]:
+    """Return ``(target_feature, top_rows)``: the steered feature + the printable top-k table.
+
+    Honors an explicit ``feature``; otherwise picks the top-active *labeled* feature, falling
+    back to the single most-active feature when none of the top-k carry a label.
+    """
+    codes = eng.encode(sequence)
+    vals, ids = codes.max(0).values.topk(min(10, codes.shape[1]))
+    rows = [(int(i), eng.labels.get(int(i)), float(v)) for v, i in zip(vals.tolist(), ids.tolist())]
+    target = feature
+    if target is None:
+        target = next((fid for fid, lab, _ in rows if lab), None)
+    if target is None:  # no --feature and no labeled feature in the top-k: steer the most-active one
+        target = int(ids[0].item())
+    return target, rows
+
+
+def run_steering(eng, sequence, organism, target, controls, strengths, n_tokens, max_clamp):
+    """Drive the (real or fake) engine to build the steering result dict — no argparse, no I/O.
+
+    Generation reuses the production path: ``Evo2SAE.generate`` clamps the same decode-only
+    ``sae.steering`` hook the server/CLI use, so the harness measures the real thing. Engine is
+    injected so this is unit-testable on CPU with a stub.
+
+    ``max_clamp`` is the engine's ``MAX_CLAMP_STRENGTH``: requested strengths beyond it are
+    *silently capped inside* ``generate``, which would make two requested strengths produce an
+    identical clamp (a fake "plateau"). We surface that here — clamp to the effective strength
+    used for generation, warn, and record both the cap and which requests were capped.
+    """
+
+    def gen(clamps):  # clamps already effective (within +/- max_clamp)
+        feats = [{"feature_id": f, "strength": v} for f, v in clamps.items()]
+        out = eng.generate(
+            prompt=sequence, organism=organism, features=feats, n_tokens=n_tokens, temperature=0.0, top_k=1
+        )
+        return out["generation"]["sequence"]
+
+    def effective(s: float) -> float:
+        return max(-max_clamp, min(max_clamp, s))
+
+    capped = sorted({s for s in strengths if effective(s) != s})
+    if capped:
+        print(
+            f"  WARNING: strength(s) {capped} exceed MAX_CLAMP_STRENGTH={max_clamp}; capped before steering "
+            "(equal-after-cap requests will look like a plateau)."
+        )
+
+    base = gen({})
+    # Dose-response: key rows by the *requested* strength (so the sweep reads as asked), but steer
+    # with the effective (capped) value so the result matches what the engine actually applies.
+    steered_by_strength = {s: gen({target: effective(s)}) for s in strengths}
+    dose = dose_response(base, steered_by_strength)
+
+    sel = None
+    if controls:
+        s = strengths[-1]
+        control_steered = {c: gen({c: effective(s)}) for c in controls}
+        sel = selectivity(base, steered_by_strength[s], control_steered)
+
+    return {
+        "target_feature": target,
+        "sequence": sequence[:80],
+        "organism": organism,
+        "baseline": base,
+        "max_clamp_strength": max_clamp,
+        "capped_strengths": capped,
+        "dose_response": dose,
+        "selectivity": sel,
+    }
+
+
+def main():
+    """Encode a sequence, then steer a target feature (dose-response) + control features (selectivity)."""
+    p = argparse.ArgumentParser(description="Evo2 SAE steering harness (clamp -> continuation effect).")
+    p.add_argument("--evo2-ckpt-dir", required=True)
+    p.add_argument("--sae-checkpoint", required=True)
+    p.add_argument("--layer", type=int, required=True)
+    p.add_argument("--sequence", required=True)
+    p.add_argument("--organism", default="None (raw DNA)")
+    p.add_argument("--feature", type=int, default=None, help="Target feature id (default: top labeled feature).")
+    p.add_argument("--controls", default="", help="Comma-separated control feature ids (selectivity).")
+    p.add_argument("--strengths", default="0,50,100,200", help="Comma-separated clamp strengths to sweep.")
+    p.add_argument("--n-tokens", type=int, default=60)
+    p.add_argument("--device", default="cuda")
+    p.add_argument("--out", default=None, help="write the structured steering_results JSON here")
+    a = p.parse_args()
+
+    from evo2_sae.core import MAX_CLAMP_STRENGTH, Evo2SAE  # noqa: E402, RUF100
+
+    try:
+        controls = [int(c) for c in a.controls.split(",") if c.strip()]
+        strengths = [float(s) for s in a.strengths.split(",")]
+    except ValueError as e:
+        raise SystemExit(f"bad --controls/--strengths (need comma-separated ints/floats): {e}")
+    if not strengths:
+        raise SystemExit("--strengths must list at least one clamp strength")
+
+    eng = Evo2SAE(a.evo2_ckpt_dir, a.sae_checkpoint, a.layer, device=a.device).load()
+
+    # 1. Encode -> the sequence's most-active features (pick a target if not given).
+    target, rows = pick_target(eng, a.sequence, a.feature)
+    print(f"top features on {a.sequence[:24]}...:")
+    for fid, lab, v in rows:
+        print(f"  feat {fid:6d}  {str(lab):18s}  max_act {v:7.2f}")
+    if a.feature is None and not any(lab for _, lab, _ in rows):
+        print(f"  (no labeled feature; defaulting target to top-active {target})")
+
+    # 2-4. Run the sweep (reuses the production generate path; pure metrics score it).
+    result = run_steering(
+        eng, a.sequence, a.organism, target, controls, strengths, a.n_tokens, MAX_CLAMP_STRENGTH
+    )
+
+    print(f"\nbaseline:  {result['baseline'][:60]}")
+    print(f"\n=== dose-response: feature {target} ({eng.labels.get(target)}) ===")
+    for r in result["dose_response"]:
+        print(
+            f"  strength {r['strength']:7.1f}: prefix@{r['first_divergence']:3d}  {r['frac_changed']:6.1%} rewritten"
+        )
+    sel = result["selectivity"]
+    if sel is not None:
+        print(f"\n=== selectivity @ strength {strengths[-1]} (target/control ratio {sel['selectivity_ratio']}) ===")
+        print(f"  target {target:6d}: {sel['target_frac_changed']:6.1%} rewritten")
+        for c, frac in sel["control_frac_changed"].items():
+            print(f"  control {int(c):6d}: {frac:6.1%} rewritten  ({eng.labels.get(int(c))})")
+
+    if a.out:
+        Path(a.out).write_text(json.dumps(result, indent=2))
+        print(f"\nwrote steering results -> {a.out}")
+
+
+if __name__ == "__main__":
+    main()

interpretability/sparse_autoencoders/recipes/evo2/scripts/steer_analysis.py

@@ -0,0 +1,110 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LicenseRef-Apache2
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Pure metrics for quantifying SAE feature-steering effects (no model — CPU-testable).
+
+Given a baseline generation and steered generations (produced by ``steer.py`` via
+``Evo2SAE.generate``), quantify the causal effect of a feature clamp:
+
+  - divergence:    how far a steered continuation departs from the baseline
+  - dose_response: how that effect scales with clamp strength
+  - selectivity:   target vs control features at one strength (is the effect feature-specific?)
+
+**Why edit distance, not positional Hamming.** ``steer.py`` decodes greedily (``temperature=0``),
+so generation is deterministic and autoregressive: the first token a clamp flips shifts every
+downstream token, which would pin a position-by-position mismatch fraction at ~1.0 and erase any
+dose curve. We therefore measure effect magnitude with a *normalized edit (Levenshtein) distance*,
+which does not saturate from a single early shift, and report ``first_divergence`` (the length of
+the shared prefix — how many leading bases survive the clamp; smaller = the effect bites earlier)
+as the complementary, monotone-friendly signal.
+
+These are pure functions of the output strings, so the steering analysis is reproducible and
+unit-tested without a GPU; ``steer.py`` runs the model and calls them to build the persisted
+result.
+"""
+
+from __future__ import annotations
+
+
+def common_prefix_len(a: str, b: str) -> int:
+    """Number of leading characters ``a`` and ``b`` share (the shared-prefix length)."""
+    n = min(len(a), len(b))
+    i = 0
+    while i < n and a[i] == b[i]:
+        i += 1
+    return i
+
+
+def edit_distance(a: str, b: str) -> int:
+    """Levenshtein edit distance between two strings (insert/delete/substitute = cost 1)."""
+    if a == b:
+        return 0
+    if not a:
+        return len(b)
+    if not b:
+        return len(a)
+    prev = list(range(len(b) + 1))
+    for i, ca in enumerate(a, 1):
+        cur = [i]
+        for j, cb in enumerate(b, 1):
+            cur.append(min(prev[j] + 1, cur[j - 1] + 1, prev[j - 1] + (ca != cb)))
+        prev = cur
+    return prev[-1]
+
+
+def divergence(a: str, b: str) -> tuple[int, float]:
+    """Return ``(shared-prefix length, normalized edit distance)``.
+
+    The first element is how many leading characters survive unchanged (``len`` when identical);
+    the second is the edit distance normalized by the longer string's length, in ``[0, 1]`` — an
+    insertion-robust measure of how much of the continuation the clamp rewrote.
+    """
+    first = common_prefix_len(a, b)
+    n = max(len(a), len(b))
+    frac = edit_distance(a, b) / n if n else 0.0
+    return first, frac
+
+
+def dose_response(baseline: str, steered_by_strength: dict[float, str]) -> list[dict]:
+    """Per clamp strength, the divergence from baseline — rows sorted by ascending strength.
+
+    ``frac_changed`` (normalized edit distance) rising and ``first_divergence`` (shared-prefix
+    length) shrinking as strength grows is the signature of a feature that genuinely steers
+    generation (a stronger clamp rewrites more, and bites earlier).
+    """
+    rows = []
+    for s in sorted(steered_by_strength):
+        first, frac = divergence(baseline, steered_by_strength[s])
+        rows.append({"strength": float(s), "first_divergence": int(first), "frac_changed": round(frac, 4)})
+    return rows
+
+
+def selectivity(baseline: str, target_steered: str, control_steered: dict[int, str]) -> dict:
+    """Target effect vs control features clamped to the same strength.
+
+    Effect magnitude is the normalized edit distance from baseline (see module docstring).
+    ``selectivity_ratio`` > 1 means the target feature rewrites generation more than the average
+    control — evidence the steering is feature-specific, not a generic "any clamp perturbs output".
+    """
+    target = divergence(baseline, target_steered)[1]
+    controls = {int(c): round(divergence(baseline, seq)[1], 4) for c, seq in control_steered.items()}
+    mean_c = sum(controls.values()) / len(controls) if controls else None
+    return {
+        "target_frac_changed": round(target, 4),
+        "control_frac_changed": controls,
+        "mean_control_frac_changed": round(mean_c, 4) if mean_c is not None else None,
+        # None when there are no controls (ratio undefined) or controls produced zero change
+        "selectivity_ratio": round(target / mean_c, 2) if mean_c else None,
+    }