From 98f07397422f9bb49add263d607b006ee309ea06 Mon Sep 17 00:00:00 2001 From: Polina Binder Date: Tue, 23 Jun 2026 06:29:16 +0000 Subject: [PATCH 1/5] evo2 SAE steering analysis: dose-response/selectivity harness, rebased onto migrated #1622 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. Metric / robustness: * normalized edit (Levenshtein) distance, not positional Hamming. 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 warns, steers at the effective value, and records max_clamp_strength + capped_strengths. Consolidation: * harness + metrics live in the package (src/evo2_sae/steer_analysis.py), engine injected, so they import as a normal torch-free module like evo2_sae.fasta — dropped all four sys.path inserts. scripts/steer.py is now a thin CLI (matches train.py/extract.py). * pick_target reuses Evo2SAE.top_features (the CLI/server ranking) instead of re-deriving topk. * one CPU test file (metrics + fake-engine harness) instead of two; fake stays local, not in conftest, to avoid colliding with the sibling server PR's engine fixtures. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Polina Binder --- .../recipes/evo2/scripts/steer.py | 96 ++++++++++ .../evo2/src/evo2_sae/steer_analysis.py | 178 ++++++++++++++++++ .../recipes/evo2/tests/test_steer_analysis.py | 150 +++++++++++++++ 3 files changed, 424 insertions(+) create mode 100644 interpretability/sparse_autoencoders/recipes/evo2/scripts/steer.py create mode 100644 interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/steer_analysis.py create mode 100644 interpretability/sparse_autoencoders/recipes/evo2/tests/test_steer_analysis.py 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..2114a47204 --- /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.steer_analysis.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.steer_analysis 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/steer_analysis.py b/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/steer_analysis.py new file mode 100644 index 0000000000..20b2fe8313 --- /dev/null +++ b/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/steer_analysis.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 ``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..c495fe09c9 --- /dev/null +++ b/interpretability/sparse_autoencoders/recipes/evo2/tests/test_steer_analysis.py @@ -0,0 +1,150 @@ +# 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.steer_analysis 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) From d6b68139643d27fc011615a3b896925d36755d4f Mon Sep 17 00:00:00 2001 From: Polina Binder Date: Wed, 24 Jun 2026 04:46:49 +0000 Subject: [PATCH 2/5] evo2-sae: move steering-analysis harness into evo2_sae.eval.steering Relocate the steering dose-response / selectivity metrics from evo2_sae.steer_analysis into the evo2_sae.eval package (src/evo2_sae/steer_analysis.py -> src/evo2_sae/eval/steering.py), alongside the eval/probing harness. Update the importers (scripts/steer.py CLI + tests/test_steer_analysis.py) to evo2_sae.eval.steering. The CI lane is dropped via the rebase onto the updated #1622. Pure-CPU tests, no GPU/model. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Polina Binder --- .../recipes/evo2/scripts/steer.py | 12 +++---- .../evo2/src/evo2_sae/eval/__init__.py | 16 ++++++++++ .../{steer_analysis.py => eval/steering.py} | 0 .../recipes/evo2/tests/test_steer_analysis.py | 32 +++++++++++++++---- 4 files changed, 47 insertions(+), 13 deletions(-) create mode 100644 interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/eval/__init__.py rename interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/{steer_analysis.py => eval/steering.py} (100%) diff --git a/interpretability/sparse_autoencoders/recipes/evo2/scripts/steer.py b/interpretability/sparse_autoencoders/recipes/evo2/scripts/steer.py index 2114a47204..c3e5a1f28c 100644 --- a/interpretability/sparse_autoencoders/recipes/evo2/scripts/steer.py +++ b/interpretability/sparse_autoencoders/recipes/evo2/scripts/steer.py @@ -15,7 +15,7 @@ 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.steer_analysis.run_steering`` (the +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. @@ -34,7 +34,7 @@ from pathlib import Path from evo2_sae.core import MAX_CLAMP_STRENGTH, Evo2SAE -from evo2_sae.steer_analysis import pick_target, run_steering +from evo2_sae.eval.steering import pick_target, run_steering def main(): @@ -72,14 +72,14 @@ def main(): 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 - ) + 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") + 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']}) ===") 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/steer_analysis.py b/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/eval/steering.py similarity index 100% rename from interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/steer_analysis.py rename to interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/eval/steering.py diff --git a/interpretability/sparse_autoencoders/recipes/evo2/tests/test_steer_analysis.py b/interpretability/sparse_autoencoders/recipes/evo2/tests/test_steer_analysis.py index c495fe09c9..8908f3c765 100644 --- a/interpretability/sparse_autoencoders/recipes/evo2/tests/test_steer_analysis.py +++ b/interpretability/sparse_autoencoders/recipes/evo2/tests/test_steer_analysis.py @@ -23,7 +23,7 @@ from typing import ClassVar -from evo2_sae.steer_analysis import divergence, dose_response, edit_distance, pick_target, run_steering, selectivity +from evo2_sae.eval.steering import divergence, dose_response, edit_distance, pick_target, run_steering, selectivity # ----------------------------------------------------------------------------- pure metrics @@ -115,8 +115,14 @@ def test_pick_target_honors_explicit_feature(): 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, + 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"]] @@ -129,8 +135,14 @@ def test_run_steering_dose_response_is_monotone(): 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, + 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 @@ -140,8 +152,14 @@ def test_run_steering_selectivity_is_feature_specific(): 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, + 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] From fd4f2a454c0e61c20a94ad959f2f741e35244431 Mon Sep 17 00:00:00 2001 From: polinabinder1 Date: Wed, 1 Jul 2026 20:35:20 +0000 Subject: [PATCH 3/5] docs(evo2-sae): update stale sae.steering ref after the hook moved to the recipe MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The steering clamp hook moved sae.steering -> evo2_sae.steering (#1622). This merges that base in and fixes the one remaining reference (a docstring in eval/steering.py). No import change was needed — the harness is engine-injected and never imported sae.steering directly. Co-Authored-By: Claude Opus 4.8 Signed-off-by: polinabinder1 --- .../recipes/evo2/src/evo2_sae/eval/steering.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 index 20b2fe8313..9d5106c924 100644 --- a/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/eval/steering.py +++ b/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/eval/steering.py @@ -16,7 +16,7 @@ """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 ``sae.steering`` hook the server/CLI use) and quantifies the causal effect: +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 From ec70b31cd75423aef80b7bcd712f1d3402b2f617 Mon Sep 17 00:00:00 2001 From: polinabinder1 Date: Thu, 2 Jul 2026 18:00:37 -0700 Subject: [PATCH 4/5] ci: GPU matrix lane for interpretability SAE recipes (merge-first, presence-guarded) (#1667) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## What A GPU CI **matrix lane** for the SAE interpretability recipes under `interpretability/sparse_autoencoders/recipes/*`. Modeled on the repo-wide `unit-tests-recipes.yml`: a cheap ubuntu `changed-dirs` job discovers which interp recipes to test and emits a matrix; a per-recipe `unit-tests` job runs each on an L4 (its own `.ci_build.sh` + `pytest tests/`). ## Which recipes are eligible Any interp recipe that has its own **`.ci_build.sh`**. Today that's **evo2**. `codonfm`/`esm2` have no `.ci_build.sh`/tests yet, so they're **skipped until they add them** — at which point they auto-join this lane, no workflow change needed. (This is also the **presence guard**: before the evo2 SAE recipe lands, #1622, nothing is eligible → the whole lane is a green no-op, so it can merge first.) ## What runs, when | What the PR changes | Which recipes run | |---|---| | a recipe's own dir — `…/recipes//**` | just **``** | | the shared `sae` lib — `…/sparse_autoencoders/sae/**` | **all** eligible recipes (they all depend on sae) | | this workflow file | **all** eligible recipes | | **nightly** `schedule` @ 09:00 UTC | **all** eligible recipes | | anything else (other recipes' dirs, docs, …) | **none** — empty matrix, green no-op | - **Megatron *build* is per-recipe — only evo2 builds it.** Each recipe's own `.ci_build.sh` owns its build: **evo2** compiles/installs the mbridge `bionemo.evo2` (megatron) stack; a future HF-native **esm2** or custom **codonfm** builds its own thing, *no megatron*. So a codonfm/esm2 change never pays for the Evo2 megatron build — it runs only on an evo2 change (or an `sae`/nightly run, which tests every consumer, evo2 included). - **Container *image* is currently shared.** All matrix entries use one base image (`svcbionemo023/…pytorch26.04-py3-squashed`, dockerhub-cached → cheap after first pull) — the megatron-capable base evo2 needs. A future HF-native recipe *runs in* that base but doesn't build megatron on top. If a recipe later wants a lighter image, the matrix can carry a **per-recipe `image`** (like the repo-wide lane's `matrix.recipe.image`) — not needed while evo2 is the only active recipe. Net: **shared base image, per-recipe (megatron-or-not) build.** - Each eligible recipe runs on the **L4** against its **CI-sized model** (evo2 = the auto-built **1B**). The **7B/L26** path is manual (`EVO2_CKPT_DIR`), never in CI. - Not covered: the React dashboard frontend (no JS build step) and offline `extract`/`train` scripts (no unit tests). ## Validation Proven end-to-end on a real L4 via the #1670 test PR: **https://github.com/NVIDIA-BioNeMo/bionemo-recipes/actions/runs/28535969801** — the L4 built via `.ci_build.sh` and ran the evo2 recipe suite green (**54 recipe tests passed**). The `changed-dirs` matrix logic was verified locally across 8 scenarios (per-recipe / sae-change / nightly / pre-#1622-empty). Trigger = a maintainer commenting `/ok to test `. --------- Signed-off-by: polinabinder1 Co-authored-by: root Co-authored-by: Claude Opus 4.8 (cherry picked from commit fdce924b0661e32d01f61692235cb9f40fa5aae1) Signed-off-by: polinabinder1 --- .../unit-tests-interpretability-recipes.yaml | 149 ++++++++++++++++++ 1 file changed, 149 insertions(+) create mode 100644 .github/workflows/unit-tests-interpretability-recipes.yaml diff --git a/.github/workflows/unit-tests-interpretability-recipes.yaml b/.github/workflows/unit-tests-interpretability-recipes.yaml new file mode 100644 index 0000000000..c14f5b9f53 --- /dev/null +++ b/.github/workflows/unit-tests-interpretability-recipes.yaml @@ -0,0 +1,149 @@ +name: "BioNeMo Interpretability Recipes CI" + +# CI for the SAE interpretability recipes under interpretability/sparse_autoencoders/recipes/*. +# Generic matrix, modeled on the repo-wide unit-tests-recipes.yml but scoped to the interp subtree: +# * `changed-dirs` (cheap ubuntu) discovers which interp recipes to test and emits a matrix. +# * `unit-tests` runs each selected recipe on the L4: its own `.ci_build.sh` + `pytest tests/`. +# +# Eligible recipes = any interp recipe with its own `.ci_build.sh`. Today that's `evo2`; codonfm/esm2 +# have no `.ci_build.sh`/tests yet, so they are skipped until they add them. This also makes the lane +# a green no-op before the evo2 SAE recipe lands (#1622) — the presence guard is "has a .ci_build.sh". +# +# What runs when: +# * change under a recipe's own dir -> that recipe. +# * change to the shared `sae` lib, or to this workflow file, or the nightly schedule +# -> ALL eligible recipes (they all depend on sae). +# * change to an unrelated recipe / elsewhere -> nothing (empty matrix, green no-op). +# Each recipe's `.ci_build.sh` owns its own build (evo2 -> mbridge bionemo.evo2; esm2 -> HF; etc.), +# so a codonfm/esm2 change never triggers the Evo2 megatron build, and vice-versa. + +on: + push: + branches: + - "pull-request/[0-9]+" + - "dependabot/**" + merge_group: + types: [checks_requested] + schedule: + - cron: "0 9 * * *" # Runs at 9 AM UTC daily (2 AM MST) + +defaults: + run: + shell: bash -x -e -u -o pipefail {0} + +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: true + +jobs: + changed-dirs: + # Cheap ubuntu pre-job: decide which interp recipes to test and emit the matrix. + runs-on: ubuntu-latest + permissions: + contents: read + outputs: + dirs: ${{ steps.set-dirs.outputs.dirs }} + steps: + - name: Checkout repository + uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 + with: + fetch-depth: 0 + persist-credentials: false + + - name: Determine which interp recipes to run + id: set-dirs + env: + EVENT_NAME: ${{ github.event_name }} + run: | + RROOT="interpretability/sparse_autoencoders/recipes" + SAE="interpretability/sparse_autoencoders/sae" + WF=".github/workflows/unit-tests-interpretability-recipes.yaml" + + # Eligible recipes = those with their own .ci_build.sh (others are green no-ops until + # they add one; empty before #1622 lands -> whole lane is a no-op). + # Use `if` (not `[ -f ] && echo`): the short-circuit leaves exit 1 when the last dir + # lacks a .ci_build.sh, which `set -e` would turn into a job failure instead of an + # empty (green no-op) matrix. + ALL=$(for d in "$RROOT"/*/; do if [ -f "${d}.ci_build.sh" ]; then echo "${d%/}"; fi; done \ + | jq -R -s -c 'split("\n") | map(select(length > 0))') + echo "Eligible recipes (have .ci_build.sh): $ALL" + + MERGE_BASE=$(git merge-base HEAD origin/main) + CHANGED=$(git diff --name-only "$MERGE_BASE" HEAD) + echo "Changed files:"; echo "$CHANGED" | sed 's/^/ - /' + + # sae lib / this workflow / nightly -> run EVERY eligible recipe (all depend on sae). + if [ "$EVENT_NAME" = "schedule" ] || echo "$CHANGED" | grep -qE "^(${SAE}/|${WF}$)"; then + DIRS="$ALL" + else + # otherwise -> only eligible recipes that have a changed file under them. + DIRS=$(echo "$ALL" | jq -c --arg changed "$CHANGED" ' + ($changed | split("\n")) as $cf + | map(select(. as $d | $cf | any(startswith($d + "/")))) + ') + fi + echo "Recipes to run: $DIRS" + + # Attach the runner image (the dockerhub-cached squashed pytorch, like the repo-wide lane). + IMG="svcbionemo023/bionemo-framework:pytorch26.04-py3-squashed" + MATRIX=$(echo "$DIRS" | jq -c --arg img "$IMG" \ + 'map({dir: ., name: (. | split("/") | last), image: $img})') + echo "dirs=$MATRIX" >> "$GITHUB_OUTPUT" + echo "matrix: $MATRIX" + + unit-tests: + needs: changed-dirs + if: ${{ needs.changed-dirs.outputs.dirs != '' && needs.changed-dirs.outputs.dirs != '[]' }} + runs-on: linux-amd64-gpu-l4-latest-1 + name: "interp-unit-tests (${{ matrix.recipe.name }})" + permissions: + contents: read + container: + image: ${{ matrix.recipe.image }} + options: --shm-size=16G + env: + CI: true + HF_TOKEN: ${{ secrets.HF_TOKEN }} + HF_HOME: /cache/huggingface + strategy: + fail-fast: false + matrix: + recipe: ${{ fromJson(needs.changed-dirs.outputs.dirs) }} + + steps: + - name: Show GPU info + run: nvidia-smi + + - name: Setup proxy cache + uses: nv-gha-runners/setup-proxy-cache@14229018fe157c83e03c008f27d183d8e99bc67c + + - name: Checkout repository + uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 + with: + # Interp recipes live under interpretability/sparse_autoencoders (recipe + shared sae); + # evo2 additionally builds on recipes/evo2_megatron. Check out both so any recipe's + # .ci_build.sh has what it needs. + sparse-checkout: | + interpretability/sparse_autoencoders + recipes/evo2_megatron + sparse-checkout-cone-mode: false + persist-credentials: false + + - name: Cache Hugging Face models + uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 + with: + path: /cache/huggingface + key: ${{ runner.os }}-huggingface-interp-${{ matrix.recipe.name }}-${{ github.sha }} + restore-keys: | + ${{ runner.os }}-huggingface-interp-${{ matrix.recipe.name }}- + ${{ runner.os }}-huggingface- + + - name: Install dependencies + working-directory: ${{ matrix.recipe.dir }} + run: bash .ci_build.sh + + - name: Run tests + working-directory: ${{ matrix.recipe.dir }} + run: | + [ -f .ci_test_env.sh ] && source .ci_test_env.sh + pytest -v tests/ From 18afa8605ca71875a6e3ced00b23c59f70574c88 Mon Sep 17 00:00:00 2001 From: polinabinder1 Date: Tue, 14 Jul 2026 22:36:29 +0000 Subject: [PATCH 5/5] docs(evo2-sae steering): scope note + de-hardcode layer literal MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Frame the harness honestly as a steering smoke test (magnitude + rough specificity), not a causal-direction check — with concept-density scoring vs an activation-matched null flagged as future work. Also drop the `--layer 26` literal from the CLI example (`--layer `); nothing in the code is layer-specific. Co-Authored-By: Claude Opus 4.8 Signed-off-by: polinabinder1 --- .../sparse_autoencoders/recipes/evo2/scripts/steer.py | 4 ++-- .../recipes/evo2/src/evo2_sae/eval/steering.py | 11 +++++++++-- 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/interpretability/sparse_autoencoders/recipes/evo2/scripts/steer.py b/interpretability/sparse_autoencoders/recipes/evo2/scripts/steer.py index c3e5a1f28c..bec88635aa 100644 --- a/interpretability/sparse_autoencoders/recipes/evo2/scripts/steer.py +++ b/interpretability/sparse_autoencoders/recipes/evo2/scripts/steer.py @@ -13,7 +13,7 @@ # 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. +r"""Evo2 SAE steering harness CLI — clamp a feature and measure its effect on generation (a smoke test). 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 @@ -22,7 +22,7 @@ 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 \ + python steer.py --evo2-ckpt-dir --sae-checkpoint --layer \ --sequence ATGGCC... --feature 29244 --controls 12345,54321 --strengths 0,50,100,200 \ --out steering_results.json """ 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 index 9d5106c924..12f1446e73 100644 --- a/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/eval/steering.py +++ b/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/eval/steering.py @@ -16,11 +16,18 @@ """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: +decode-only ``evo2_sae.steering`` hook the server/CLI use) and measures its effect on generation: - 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?) + - selectivity: target vs control features at one strength (bigger than the controls?) + +**Scope — a steering smoke test.** These metrics quantify the *magnitude* of a clamp's effect +(dose_response) and whether it exceeds a few *hand-picked* control features (selectivity). They do +not check the effect's *direction* — that the output moved toward the target feature's labeled +concept — and ``selectivity`` is only as trustworthy as the chosen controls. Verifying steering is +concept-correct (does clamping a "stop-codon" feature yield more stop codons?) against an +activation-matched null is future work: concept-density scoring via ``evo2_sae.eval.probing.labelers``. **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