diff --git a/interpretability/sparse_autoencoders/recipes/evo2/scripts/steer.py b/interpretability/sparse_autoencoders/recipes/evo2/scripts/steer.py new file mode 100644 index 0000000000..c3e5a1f28c --- /dev/null +++ b/interpretability/sparse_autoencoders/recipes/evo2/scripts/steer.py @@ -0,0 +1,96 @@ +# 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 CLI — clamp a feature and measure the causal effect on generation. + +Thin wrapper: builds an ``Evo2SAE`` and calls ``evo2_sae.eval.steering.run_steering`` (the +engine-driven harness + pure metrics live there, CPU-tested). 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 --sae-checkpoint --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 +from pathlib import Path + +from evo2_sae.core import MAX_CLAMP_STRENGTH, Evo2SAE +from evo2_sae.eval.steering import pick_target, run_steering + + +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() + + 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 r in rows: + print(f" feat {r['feature_id']:6d} {str(r['label']):18s} max_act {r['max_activation']:7.2f}") + if a.feature is None and not any(r["label"] for r 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() diff --git a/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/eval/__init__.py b/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/eval/__init__.py new file mode 100644 index 0000000000..34bc358280 --- /dev/null +++ b/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/eval/__init__.py @@ -0,0 +1,16 @@ +# 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. + +"""Evo2 SAE evaluation harnesses (steering analysis, probing, …).""" diff --git a/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/eval/steering.py b/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/eval/steering.py new file mode 100644 index 0000000000..9d5106c924 --- /dev/null +++ b/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/eval/steering.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. + +"""SAE feature-steering analysis: the engine-driven harness + the pure metrics it scores with. + +``run_steering`` clamps a feature via the production ``Evo2SAE.generate`` path (the same +decode-only ``evo2_sae.steering`` hook the server/CLI use) and quantifies the causal effect: + + - 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.** Steering 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. + +The engine is *injected* into ``pick_target``/``run_steering`` (rather than imported), so this whole +module stays torch-free and CPU-unit-testable with a stub; ``scripts/steer.py`` is just the CLI that +builds a real ``Evo2SAE`` and calls in. Lives in the package (not ``scripts/``) so it imports as a +normal module — no ``sys.path`` games — the same way ``evo2_sae.fasta`` does. +""" + +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, + } + + +# --------------------------------------------------------------------- harness (engine injected) +def pick_target(eng, sequence: str, feature: int | None = None, k: int = 10) -> tuple[int, list[dict]]: + """Return ``(target_feature, top_rows)``: the steered feature + the printable top-k table. + + Reuses ``Evo2SAE.top_features`` (same ranking the CLI/server show) instead of re-deriving the + top-k. Honors an explicit ``feature``; else the top-active *labeled* feature; else the single + most-active feature (``top_rows`` is sorted by activation, strictly-positive features only). + """ + rows = eng.top_features(eng.encode(sequence), k=k) + target = feature + if target is None: + target = next((r["feature_id"] for r in rows if r["label"]), None) + if target is None: # no --feature and no labeled feature in the top-k: steer the most-active one + target = rows[0]["feature_id"] if rows else 0 + return target, rows + + +def run_steering(eng, sequence, organism, target, controls, strengths, n_tokens, max_clamp) -> dict: + """Drive the (real or fake) engine to build the steering result dict — no argparse, no I/O. + + ``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 — steer at the effective (capped) strength, + 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: # mirrors core._sanitize_steering's clamp on |strength| + 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, + } diff --git a/interpretability/sparse_autoencoders/recipes/evo2/tests/test_steer_analysis.py b/interpretability/sparse_autoencoders/recipes/evo2/tests/test_steer_analysis.py new file mode 100644 index 0000000000..8908f3c765 --- /dev/null +++ b/interpretability/sparse_autoencoders/recipes/evo2/tests/test_steer_analysis.py @@ -0,0 +1,168 @@ +# 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. + +"""CPU tests for the steering metrics + harness (pure string math + a torch-free fake engine). + +No GPU, no checkpoints, no torch — the engine is injected, so ``pick_target``/``run_steering`` run +against the in-file ``FakeEngine``. The fake is intentionally local (not a conftest fixture) so it +stays decoupled from the GPU engine fixtures the sibling server PR adds to ``conftest.py``. The real +engine path is covered by the GPU ``test_steering.py``. +""" + +from typing import ClassVar + +from evo2_sae.eval.steering import divergence, dose_response, edit_distance, pick_target, run_steering, selectivity + + +# ----------------------------------------------------------------------------- pure metrics +def test_divergence_identical_and_single_change(): + """Identical strings -> no divergence; one substitution -> shared-prefix len + edit fraction.""" + assert divergence("ACGTACGT", "ACGTACGT") == (8, 0.0) + first, frac = divergence("ACGTACGT", "ACGAACGT") # differs only at index 3 + assert first == 3 and abs(frac - 0.125) < 1e-9 + + +def test_divergence_edit_distance_does_not_saturate_on_a_shift(): + """A single early insertion shifts every downstream base. + + Positional Hamming would call this ~100% changed (the saturation that erases a dose curve under + greedy decode); normalized edit distance sees one insert and reports a tiny effect. + """ + base = "ACGTACGTACGT" + shifted = "C" + base[:-1] # one insertion at the front, everything else shunted right by one + first, frac = divergence(base, shifted) + assert first == 0 # diverges immediately... + assert edit_distance(base, shifted) <= 2 # ...but it's a cheap edit, not a full rewrite + assert frac < 0.2 # not pinned at ~1.0 the way Hamming would be + + +def test_dose_response_sorted_and_rises_with_strength(): + """Rows come back ascending in strength; a stronger clamp changes more of the continuation.""" + base = "AAAAAAAA" + steered = {0.0: "AAAAAAAA", 100.0: "AAAACCCC", 50.0: "AAAAAACC"} # given out of order + rows = dose_response(base, steered) + assert [r["strength"] for r in rows] == [0.0, 50.0, 100.0] + assert rows[0]["frac_changed"] == 0.0 + assert rows[1]["frac_changed"] < rows[2]["frac_changed"] + + +def test_selectivity_ratio_high_when_target_specific(): + """Target clamp moves generation, controls barely do -> ratio >> 1 (feature-specific).""" + base = "AAAAAAAA" + sel = selectivity(base, "CCCCCCCC", {1: "AAAAAAAA", 2: "AAAAAAAC"}) + assert sel["target_frac_changed"] == 1.0 + assert sel["mean_control_frac_changed"] < 0.1 + assert sel["selectivity_ratio"] > 5 + + +# ----------------------------------------------------------------------------- harness (fake engine) +class FakeEngine: + """Deterministic, torch-free stand-in for ``Evo2SAE``. + + ``top_features`` returns a fixed ranking whose top entry is *unlabeled* (so target selection has + to skip it for the labeled one); ``generate`` rewrites a strength-scaled tail, with the target + feature (7) 10x more potent than any other -> monotone dose + high selectivity. + """ + + labels: ClassVar[dict[int, str]] = {7: "promoter_like", 3: "tata_box"} + # (feature_id, label, max_activation), sorted by activation descending + ranking: ClassVar[list[tuple[int, str | None, float]]] = [ + (9, None, 6.0), # most active overall, but UNLABELED + (7, "promoter_like", 5.0), + (3, "tata_box", 2.0), + ] + + def encode(self, dna: str): + return dna # opaque token; top_features below ignores it and uses `ranking` + + def top_features(self, codes, tag_len: int = 0, k: int = 8): + return [{"feature_id": f, "label": lab, "max_activation": v} for f, lab, v in self.ranking[:k]] + + def generate(self, prompt, organism, features, n_tokens, temperature, top_k): + assert temperature == 0.0 and top_k == 1 # harness must drive greedy/deterministic decode + if not features: + return {"generation": {"sequence": "A" * n_tokens}} + fid, strength = features[0]["feature_id"], abs(features[0]["strength"]) + potency = 10 if fid == 7 else 100 # feature 7 rewrites 10x more per unit strength + changed = min(n_tokens, int(strength / potency)) + return {"generation": {"sequence": "A" * (n_tokens - changed) + "C" * changed}} + + +def test_pick_target_prefers_top_labeled_over_top_active(): + """No --feature -> the most-active *labeled* feature (7), skipping the unlabeled top-active (9).""" + target, rows = pick_target(FakeEngine(), "ACGTACGTACGT", feature=None) + assert target == 7 + assert rows[0]["feature_id"] == 9 # table still surfaces the top-active feature + + +def test_pick_target_honors_explicit_feature(): + target, _ = pick_target(FakeEngine(), "ACGTACGTACGT", feature=3) + assert target == 3 + + +def test_run_steering_dose_response_is_monotone(): + """A stronger clamp rewrites more (frac up) and bites earlier (shared prefix down).""" + res = run_steering( + FakeEngine(), + "ACGT" * 20, + "None (raw DNA)", + target=7, + controls=[], + strengths=[0.0, 50.0, 100.0, 200.0], + n_tokens=60, + max_clamp=300.0, + ) + fracs = [r["frac_changed"] for r in res["dose_response"]] + prefixes = [r["first_divergence"] for r in res["dose_response"]] + assert fracs == sorted(fracs) and fracs[0] == 0.0 and fracs[-1] > fracs[0] + assert prefixes == sorted(prefixes, reverse=True) # non-increasing + assert res["capped_strengths"] == [] + assert res["selectivity"] is None + + +def test_run_steering_selectivity_is_feature_specific(): + """Target (potent) vs control (10x weaker) at the strongest clamp -> ratio > 1.""" + res = run_steering( + FakeEngine(), + "ACGT" * 20, + "None (raw DNA)", + target=7, + controls=[3], + strengths=[0.0, 200.0], + n_tokens=60, + max_clamp=300.0, + ) + sel = res["selectivity"] + assert sel["selectivity_ratio"] > 1 + assert sel["target_frac_changed"] > sel["mean_control_frac_changed"] + + +def test_run_steering_surfaces_the_clamp_cap(): + """A requested strength above the cap is recorded and steered at the effective (capped) value.""" + res = run_steering( + FakeEngine(), + "ACGT" * 20, + "None (raw DNA)", + target=7, + controls=[], + strengths=[50.0, 200.0], + n_tokens=60, + max_clamp=100.0, + ) + assert res["max_clamp_strength"] == 100.0 + assert res["capped_strengths"] == [200.0] + # Requested 200 is steered at effective 100 -> int(100/10)=10 of 60 bases rewritten. + row_200 = next(r for r in res["dose_response"] if r["strength"] == 200.0) + assert row_200["frac_changed"] == round(10 / 60, 4)