evo2-sae: probing primitives + harness (metrics, ActivationBuffer, probe CLI)

[polinabinder1]

Summary

SAE probing primitives (eval metrics + ActivationBuffer) for Evo2 — scoring metrics + per-feature annotation, all pure functions of codes + labels. Lives in the evo2 recipe at evo2_sae.eval.probing — moved out of the shared sae library because it's evo2-specific (the shared sae.eval keeps loss_recovered / sparsity / dead_latents for esm2/codonfm).

Stacked on #1622 (uses the evo2_sae package). #1630 supplies the eval labels.

Stack update (2026-06-29): #1636 (probing harness + per-domain label producers) was merged into this branch — its code (evo2_buffer.py, labelers.py, annot_tracks.py, euk_windows.py, probe.py, probe_loss_recovered.py, their tests) is now part of this PR. Scope is therefore probing primitives + harness.

Contents — evo2_sae.eval.probing

Primitives (pure functions of codes + labels)

• ActivationBuffer (codes + optional dense twin + per-token labels + instance ids)
• AUROC: auroc_all, auroc_vec, best_single_train_test
• decoders: fit_logreg / fit_softmax / macro_auroc / decode_eval
• domain_f1 (precision-per-nt, recall-per-instance)
• annotate_features (per-feature best concept by AUROC → the annotation table)

Probing harness + label producers (folded in from #1636)

• evo2_buffer.py — DNA → ActivationBuffer (the one model-touching step: Evo2 → layer-L residual → SAE.encode + per-token labels)
• labelers.py / annot_tracks.py / euk_windows.py — per-domain biological labelers + BED/GFF interval tracks + euk gene-structure windows
• probe.py — the probing CLI (extract / auroc / annotate / linear / euk-f1 / domain-eval)
• probe_loss_recovered.py — SAE fidelity via the shared sae.eval.loss_recovered

How to use

from evo2_sae.eval.probing import auroc_all, annotate_features
au  = auroc_all(codes, labels)                                   # [F, L]
ann = annotate_features(codes, labels, names, min_auroc=0.85)    # [{feature_id, label, auroc}]

Design decisions (the non-trivial ones)

• Scoring is a pure function of a saved buffer — the model-touching step is isolated to one file. ActivationBuffer (codes + optional dense twin + per-token labels + instance ids) is the contract; all metrics are pure CPU functions of it. Only evo2_buffer.build_buffer touches the model (Evo2 → layer-L → SAE.encode). So the expensive GPU pass runs once and every metric / re-analysis afterward runs on a laptop from the .npz — no model, fully reproducible.
• The buffer carries a dense-residual "twin." Keeping the raw layer-L residual alongside the SAE codes lets the same artifact answer "does sparsification cost anything?" — SAE-vs-dense single- and multi-feature probes, no re-encode.
• Metrics are model-agnostic but live in the recipe package (evo2_sae.eval.probing), beside loss_recovered/reconstruction/sparsity/dead_latents. They're the newest and evo2-driven; the established shared metrics stay in sae.eval. (Folding the probing primitives up into shared sae is the planned follow-up — they're written model-agnostic precisely so that's a move, not a rewrite.)
• Hand-rolled, GPU-vectorized over the whole ~65k-feature dictionary (see below) — auroc_all is one [features × concepts] pass; sklearn would be a 65k-iteration CPU loop, and fit_logreg avoids the sklearn.LogisticRegression scaling wall that forced CodonFM to subsample to ≤5k features.
• Winner's-curse correction — best_single_train_test selects the best single feature on train and reports its AUROC on test, so a concept's "best single feature" score isn't the optimistic max-over-features on the same data.
• domain_f1 is instance-level (precision-per-nt, recall-per-annotation) and always read against a shuffle null — for sparse regional concepts (exon/CDS) where per-token AUROC undersells, but the null guards against high-prevalence concepts looking strong for free (intron ≈ 1.09× null = degenerate).
• probe_loss_recovered reuses the shared sae.eval.loss_recovered — fidelity stays in the shared lib (esm2/codonfm use it too); only the Evo2 callables are local. No duplicate metric.
• Optional second GPU for the build step (--auroc-device). Building a buffer holds two big things at once: the Evo2 7B (Megatron pins it on its device and won't free it in-process) and the multi-GB code/dense buffer + the vectorized AUROC matmul over all ~65k features. --auroc-device lets that buffer/scoring live on a second GPU so it doesn't fight the pinned model for memory. It's optional — it defaults to --device (single GPU); set it only when memory is tight (it is, on the 7B). And it matters only for the in-process build-and-score path (extract / euk-f1 / domain-eval): once the buffer is saved to .npz, scoring is the pure-function path — no model, no pinning, CPU or any GPU. So this is a peak-memory accommodation for one step, not a standing two-GPU requirement.

Tests

No dedicated CI lane (deferred — see #1622). Run via the recipe:

cd interpretability/sparse_autoencoders/recipes/evo2
bash .ci_build.sh && source .ci_test_env.sh        # or: PYTHONPATH=src:../../sae/src
pytest tests/test_probing.py tests/test_probe_integration.py tests/test_labelers.py tests/test_annot_tracks.py tests/test_euk_windows.py

Primitives — 12 passed (CPU, no model): AUROC vs a pairwise-definition oracle, domain_f1 vs a hand-computed reference, best_single winner's-curse flip, decode_eval separability, annotate_features best-concept, buffer roundtrip, tie-correct (average) ranks, degenerate-label / tie / sparse edge cases, and standardize's zero-variance floor.

Harness (folded in from #1636) — CPU (no model): label producers (labelers / annot_tracks / euk_windows) + the probe-CLI integration (buffer save/load incl. the dense twin, AUROC/annotate/linear over a planted feature, domain_f1 over interval tracks). The real-engine buffer / loss-recovered path is gated by @pytest.mark.skipif(not torch.cuda.is_available()).

Why hand-rolled (not sklearn / torchmetrics) — checked, not a win

GPU-vectorized over the whole ~32k-feature dictionary in one pass; the library options are CPU and per-(scores, label), so a 32k-feature dictionary becomes a 32k-iteration CPU loop. Function by function:

• auroc_all — no library computes a vectorized [features × labels] AUROC matrix on GPU. Kept.
• domain_f1, best_single_train_test, annotate_features — bespoke (instance-F1, winner's-curse, per-feature assignment); no library equivalent.
• fit_logreg / fit_softmax / decode_eval — the only sklearn-replaceable code, but they fit on the [N≈50k, F≈32k] SAE-code matrix, exactly where CodonFM hit the sklearn.LogisticRegression scaling wall and had to subsample to ≤5k features. Swapping reintroduces that coverage loss + a runtime dep. Net regression.
• ActivationBuffer / split_indices / standardize — np.savez + tiny helpers; nothing to gain.

Conclusion: the module stays torch + numpy-only. Each metric is a standard formula (Mann–Whitney rank-AUROC, Adam BCE/softmax, instance-F1) vectorized for full-dictionary GPU scale, and each is validated against an independent reference in the tests.

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 194cd152-a158-4dbd-9d82-193fa3263c3d

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

📝 Walkthrough

Walkthrough

This PR adds a comprehensive SAE feature-probing evaluation module (probing.py) to enable model-agnostic interpretation of learned features through metrics, classifiers, and annotation tools, along with an ActivationBuffer artifact for persistence and a full test suite validating correctness across all components.

Changes

SAE Probing Evaluation Suite

Layer / File(s)	Summary
ActivationBuffer data structure and persistence bionemo-recipes/interpretability/sparse_autoencoders/sae/src/sae/eval/probing.py (lines 1–65), bionemo-recipes/interpretability/sparse_autoencoders/sae/tests/test_probing.py (lines 123–142)	Dataclass storing SAE feature codes, per-token boolean labels and names, optional dense residuals, and concept-to-instance id mappings; .save() serializes to typed .npz with per-concept instance arrays; .load() reconstructs the dataclass; .name_idx property maps label names to column indices.
Dataset utilities and standardization bionemo-recipes/interpretability/sparse_autoencoders/sae/src/sae/eval/probing.py (lines 73–84)	split_indices performs deterministic train/test splitting via seeded torch.randperm; standardize computes mean and std on training rows with epsilon-clamped std normalization.
AUROC computation and best-feature selection bionemo-recipes/interpretability/sparse_autoencoders/sae/src/sae/eval/probing.py (lines 86–145), bionemo-recipes/interpretability/sparse_autoencoders/sae/tests/test_probing.py (lines 37–71)	auroc_all computes full [feature, label] AUROC matrix via chunked rank-statistics; auroc_vec handles single-vector AUROC with degenerate-case handling; best_single_train_test selects best feature on training set and reports test AUROC without winner's-curse bias; test oracle _auroc_ref validates against brute-force reference.
Feature concept annotation via AUROC thresholding bionemo-recipes/interpretability/sparse_autoencoders/sae/src/sae/eval/probing.py (lines 147–174), bionemo-recipes/interpretability/sparse_autoencoders/sae/tests/test_probing.py (lines 110–121)	annotate_features derives per-feature best-label annotations by selecting max AUROC across labels and filtering by configurable AUROC threshold; excludes low-information features.
Linear classifier training and macro-AUROC evaluation bionemo-recipes/interpretability/sparse_autoencoders/sae/src/sae/eval/probing.py (lines 176–226), bionemo-recipes/interpretability/sparse_autoencoders/sae/tests/test_probing.py (lines 89–108)	fit_logreg trains binary logistic regression; fit_softmax trains multinomial softmax; both use Adam with BCE-with-logits and cross-entropy respectively; macro_auroc computes macro one-vs-rest AUROC; decode_eval orchestrates training and dual metric reporting for test accuracy and macro AUROC.
Domain-adjusted F1 with instance-aware thresholding bionemo-recipes/interpretability/sparse_autoencoders/sae/src/sae/eval/probing.py (lines 228–270), bionemo-recipes/interpretability/sparse_autoencoders/sae/tests/test_probing.py (lines 73–87)	domain_f1 computes threshold-swept per-feature F1 by normalizing activations per-feature, remapping instance ids, aggregating per-instance firing via index_reduce_, combining precision from concept masks with recall from instance aggregation, and selecting best F1 threshold per feature in chunked passes.
Module public API and test setup bionemo-recipes/interpretability/sparse_autoencoders/sae/src/sae/eval/__init__.py (lines 25–71), bionemo-recipes/interpretability/sparse_autoencoders/sae/tests/test_probing.py (lines 1–35)	Imports and re-exports all probing.py utilities in __all__ for public access; test module imports and validates all components.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

Poem

🐰 A warren of metrics, now bundled with care,
AUROC and F1 floating through air,
Buffers that save what the features unfold,
Linear probes seeking wisdom untold,
Domain-aware thresholds, adaptive and keen—
The richest of probing suites ever been seen! 🌟

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	Docstring coverage is 91.30% which is sufficient. The required threshold is 80.00%.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title is concise and clearly reflects the main change: Evo2 SAE probing primitives plus the probing harness and CLI.
Description check	✅ Passed	The PR description is detailed and covers summary, usage, design decisions, and tests, though it doesn't strictly follow every template section.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch pbinder/sae-interp-primitives

---

_{Comment @coderabbitai help to get the list of available commands.}

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[polinabinder1]

@coderabbitai review

[coderabbitai]

✅ Action performed

Review finished.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

[coderabbitai]

🧹 Nitpick comments (2)

bionemo-recipes/interpretability/sparse_autoencoders/sae/src/sae/eval/probing.py (2)

54-65: 💤 Low value

allow_pickle=True poses a deserialization risk if loading untrusted files.

This is acceptable for internal artifacts but worth documenting. If these buffers might come from external sources, consider validating provenance or using a safer serialization format.

     `@classmethod`
     def load(cls, path: str) -> "ActivationBuffer":
-        """Load an ActivationBuffer from an .npz written by save()."""
+        """Load an ActivationBuffer from an .npz written by save().
+
+        Warning:
+            Uses allow_pickle=True; only load files from trusted sources.
+        """
         z = np.load(path, allow_pickle=True)

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In
`@bionemo-recipes/interpretability/sparse_autoencoders/sae/src/sae/eval/probing.py`
around lines 54 - 65, The load method in ActivationBuffer uses np.load(...,
allow_pickle=True) which is unsafe for untrusted files; change load to avoid
allow_pickle=True by default (use allow_pickle=False) or add an explicit
parameter (e.g., allow_pickle: bool = False) and fail with a clear error if
pickled objects are required, and update the ActivationBuffer.load docstring to
document the deserialization risk and the need to validate provenance when
loading external files; ensure references to ActivationBuffer.load and the local
variable z are used to implement and surface the safer behavior.

---

243-245: 💤 Low value

Consider adding a comment explaining the +2 sizing for the remap tensor.

The +2 accounts for 0-indexing and ensures negative indexing (-1) wraps to a valid buffer position. While correct, this is subtle:

-    remap = torch.full((int(inst_ids.max().item()) + 2,), -1, device=dev, dtype=torch.long)
+    # +2: one for 0-indexing, one so that -1 wraps to a valid (unused) slot
+    remap = torch.full((int(inst_ids.max().item()) + 2,), -1, device=dev, dtype=torch.long)

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In
`@bionemo-recipes/interpretability/sparse_autoencoders/sae/src/sae/eval/probing.py`
around lines 243 - 245, Add an inline comment above the remap creation
explaining why the size is int(inst_ids.max().item()) + 2: we need +1 for
0-based indexing of the maximum id and an extra slot so that using -1 as a
sentinel (when indexing remap with potentially -1 inst_ids) will wrap to a valid
buffer position instead of raising an out-of-bounds error; reference the remap
tensor and the subsequent remap[uniq.long()] / remap[inst_ids.long()] usage (and
the torch.full default -1) so readers understand the sentinel handling.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In
`@bionemo-recipes/interpretability/sparse_autoencoders/sae/src/sae/eval/probing.py`:
- Around line 54-65: The load method in ActivationBuffer uses np.load(...,
allow_pickle=True) which is unsafe for untrusted files; change load to avoid
allow_pickle=True by default (use allow_pickle=False) or add an explicit
parameter (e.g., allow_pickle: bool = False) and fail with a clear error if
pickled objects are required, and update the ActivationBuffer.load docstring to
document the deserialization risk and the need to validate provenance when
loading external files; ensure references to ActivationBuffer.load and the local
variable z are used to implement and surface the safer behavior.
- Around line 243-245: Add an inline comment above the remap creation explaining
why the size is int(inst_ids.max().item()) + 2: we need +1 for 0-based indexing
of the maximum id and an extra slot so that using -1 as a sentinel (when
indexing remap with potentially -1 inst_ids) will wrap to a valid buffer
position instead of raising an out-of-bounds error; reference the remap tensor
and the subsequent remap[uniq.long()] / remap[inst_ids.long()] usage (and the
torch.full default -1) so readers understand the sentinel handling.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 23ddf87a-6a45-46a2-8264-db968ee016e5

📥 Commits

Reviewing files that changed from the base of the PR and between e407165 and 79df727.

📒 Files selected for processing (3)

• bionemo-recipes/interpretability/sparse_autoencoders/sae/src/sae/eval/__init__.py
• bionemo-recipes/interpretability/sparse_autoencoders/sae/src/sae/eval/probing.py
• bionemo-recipes/interpretability/sparse_autoencoders/sae/tests/test_probing.py

[polinabinder1]

Addressed the two nitpicks in 57837ec7: documented the allow_pickle=True trust caveat on ActivationBuffer.load, and added a comment explaining the +2 remap-tensor sizing (index-by-max-id + sentinel headroom). Tests still green (6 passed).

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.