feat(sparse_probing): add leakage-safe k-sparse probing over activation tensors by janmenjayap · Pull Request #1774 · TransformerLensOrg/TransformerLens

janmenjayap · 2026-09-12T10:25:09Z

Description

Adds transformer_lens/tools/analysis/sparse_probing.py, a dependency-free, model-free module for fitting binary k sparse probes to a supplied activation matrix X and label vector y, following the method in Gurnee et al., "Finding Neurons in a Haystack: Case Studies with Sparse Probing" (TMLR 2023).

TransformerLens already exposes activations via run_with_cache, but has no maintained probing primitive: users currently rebuild splitting, feature selection, fitting, and control logic themselves, and it is easy to leak test information into that process. This PR adds a leakage-safe, dependency-free core for that workflow. Activation extraction, model wrappers, plotting, notebooks, multiclass probing, and optimal (MIP) selection are deliberately deferred to follow-up work.

Fixes #1728

What's included

fit_sparse_probe(X, y, k, ...) — deterministic stratified train/test split computed before any learned statistic, train-only raw mean-difference feature selection, and a CPU-float64 Torch LBFGS balanced-logistic fit with explicit objective/gradient-norm convergence diagnostics (raises rather than silently returning an unconverged fit).
sweep_sparse_probe(X, y, ks, ...) — reuses one fixed split across a k-grid and reports raw random-coordinate and shuffled-training-label control distributions alongside each k, without assigning an automatic "significance" label.
SparseProbeResult / SparseProbeSweep result dataclasses reporting held-out accuracy, precision, recall, and F1 (F1 primary), selected indices/scores, coefficients/intercept, preprocessing metadata, split indices/class counts, and control distributions.
A regression test asserting that MPS-resident features are moved to CPU before the float64 fit (float64 is unsupported on MPS).
A jaxtyping>=0.3 compatibility fix: the rebase onto origin/dev pulled in CI Warnings Cleanup #1732, which now raises jaxtyping.TypeCheckError instead of letting BeartypeCallHintParamViolation propagate, and reclassifies float8_e4m3fn as a valid Float dtype at the annotation level. Updated test_runtime_typecheck_rejects_invalid_tensor_contracts to use tests/typecheck_errors.TYPECHECK_ERRORS, and moved the float8 case to this module's own ValueError dtype-rejection table, where it's now actually caught.
Exports from transformer_lens.tools.analysis, plus docs/source/content/sparse_probing.md (contracts, claim boundaries, and a run_with_cache composition example) linked from docs/source/index.md.

Design decisions

Issue #1728 asked maintainers to weigh in on five open questions before implementation. This PR takes the issue's own recommended position on each:

Repository fit / model-free scope — kept model-free; extraction and plotting are out of scope for this PR.
Optimizer — uses an explicit CPU-float64 Torch LBFGS objective (no new sklearn dependency), and reports final gradient/objective metadata, raising on non-convergence.
Controls — random-coordinate and shuffled-training-label distributions are first-class sweep outputs, disabled only when their repeat count is zero.
Preprocessing — defaults to "none" (matching the paper's reference code); optional train-only "standardize" is documented as intentionally changing the L2 objective.
Device/dtype — score reductions run in at least float32 on the input device; only the selected [example, k] matrices move to CPU float64 for the deterministic LBFGS fit.

Non-goals (tracked as follow-ups)

Activation extraction and position-reduction helpers, model wrappers/downloads, multiclass/one-vs-rest probing, plotting and notebooks, an optimal/MIP selector, feature batteries, SAE-latent composition, and causal validation are all explicitly out of scope for this PR.

Claim boundaries

A high held-out F1 means the labeled feature is linearly decodable from the supplied activations — it does not by itself establish that the model uses that feature, that a selected coordinate is monosemantic, or that a smooth k-sweep curve is evidence of superposition. The docs and result contracts call this out explicitly, and the module exposes no .plot() method or automatic "significant" label.

Test plan

uv run pytest tests/unit/tools/test_sparse_probing.py tests/mps/test_mps_basic.py -q — 68 passed.
make test-pr (unit + docstring + acceptance + integration): unit 5,827 passed, docstring 18 passed, acceptance 209 passed, integration 1,459 passed / 1 failed. The one failure, test_granite_eager_scan_device_correctness[mps], is an unrelated Granite MoE Hybrid eager/fused-scan divergence — reproduced identically on a clean origin/dev checkout in an isolated worktree, confirming it's pre-existing on trunk and not introduced by this branch.
make format and uv run mypy . clean.
docs/source/index.md includes the new guide exactly once; docs build succeeds.

Type of change

New feature (non-breaking change which adds functionality)
This change requires a documentation update

Checklist:

I have made corresponding changes to the documentation
My changes generate no new warnings
I have added tests that prove my fix is effective or that my feature works
New and existing unit tests pass locally with my changes
I have not rewritten tests relating to key interfaces which would affect backward compatibility

Add fit_sparse_probe: a stratified train/test split is computed before any learned statistic, features are selected by train-only mean difference, and a CPU-float64 LBFGS logistic fit reports explicit objective/gradient-norm convergence diagnostics instead of assuming success. Report held-out accuracy, precision, recall, and F1 so callers can judge decodability without the result implying causal model use, neuron monosemanticity, or superposition. Cover exact score/index selection, leakage isolation, deterministic ties/RNG, optimizer convergence and failure, constant columns, and invalid-input rejection. Sweep, controls, exports, and docs land in a follow-up commit.

Add sweep_sparse_probe: fits a shared-split k-grid and reports raw random-coordinate and shuffled-training-label control distributions alongside each k, so callers can judge coordinate concentration without the result implying an automatic significance test. Export fit_sparse_probe, sweep_sparse_probe, and the result dataclasses from transformer_lens.tools.analysis, and add the sparse_probing guide to the docs toctree, documenting the leakage-safe contract and a run_with_cache composition example.

fit_sparse_probe accepts feature matrices on any device, including MPS, where float64 is unsupported. Add a regression test that fits a probe over MPS-resident features and asserts the selected train/test tensors land on CPU with float64 dtype instead of raising when the float64 cast is attempted while still on the Metal device.

…type policy Rebasing onto origin/dev pulled in jaxtyping>=0.3 (TransformerLensOrg#1732), which re-raises type-check violations as jaxtyping.TypeCheckError instead of letting BeartypeCallHintParamViolation propagate, and now classifies float8_e4m3fn as a valid Float dtype at the annotation level. Use the project's tests/typecheck_errors.TYPECHECK_ERRORS convention for the annotation-level cases, and move the float8_e4m3fn case to the explicit ValueError dtype-rejection table, where the function's own dtype guard now catches it.

jlarson4

Thanks for putting this together @janmenjayap. A few comments below:

…bels

… caveat

janmenjayap · 2026-09-15T20:17:01Z

Thanks again for the careful read. Happy to adjust any of the wording or test shapes if you would prefer a different form.

jlarson4

I appreciate you turning these around so quickly! The detailed replies are also great, having your thought process speeds up my reviewing process tremendously. The k metadata fix and the two new behavioural tests look great. I have a couple new comments below, this should hopefully be the last batch for this feature

jlarson4 · 2026-09-16T17:30:01Z

+    gradient_inf_norm = float(gradient.abs().max().item())
+    if not math.isfinite(objective) or not math.isfinite(gradient_inf_norm):
+        raise RuntimeError("sparse probe optimizer produced non-finite output")
+    acceptance_threshold = gradient_tolerance * max(1.0, initial_gradient_inf_norm)


This is something I suggested in my initial review and I definitely need to adjust my recommendation here. Apologies for leading you in the wrong direction.

The gradient at the zero starting point measures label–feature correlation rather than feature magnitude, so the probe support deserves a larger threshold while the random-subset and label-shuffle controls should be smaller, and the smallest are pinned back to the bare absolute tolerance by the max(1.0, ...) floor. At activation scale a sweep still dies on a converged control fit, discarding every completed fit with it, even though every single fit_sparse_probe at the same scale now passes.

Can you please key the bound to the magnitude of the selected training columns instead, for example gradient_tolerance * max(1.0, k * max|X_train|) while leaving LBFGS's own tolerance_grad as it is, and add a sweep with both control arms at activation scale to the guard test?

jlarson4 · 2026-09-16T17:31:35Z

+    parameters = torch.zeros(features.shape[1] + 1, dtype=torch.float64, requires_grad=True)
+    optimizer = torch.optim.LBFGS(
+        [parameters],
+        max_iter=max_iter,


LBFGS is still told to stop at the absolute gradient_tolerance, which it cannot reach at scale, so it runs until max_iter or the evaluation budget binds, and the budget is never set here, so it is torch's undocumented max_iter * 5 // 4. The relative bound then accepts whatever the solver reached: at per-coordinate std in the thousands roughly half the accepted fits are such truncations, some of them far above the optimum in objective with coefficients well off the true ones even though their held-out metrics look fine, and a caller can only tell by comparing function_evaluations against a cap that appears nowhere in the signature or the guide. Please set max_eval explicitly and report the stop reason on the result, so a fit that ended on max_iter or the evaluation budget is distinguishable from one that met the tolerance.

jlarson4 · 2026-09-16T17:32:27Z

+        raise ValueError(f"seed must be an integer in [0, 2**63), got {seed!r}")
+    validated_max_iter = _positive_integer(max_iter, "max_iter")
+    validated_tolerance = _finite_positive_real(gradient_tolerance, "gradient_tolerance")
+    if validated_tolerance > 1:


At the top of the documented (0, 1] range the acceptance threshold is at least the initial gradient, which no fit can miss, and LBFGS's own tolerance_grad of 1.0 stops the solve before it takes a step whenever the initial gradient is at or below 1, so the call returns coefficients that are all zero and an objective of exactly ln 2, the untrained probe, as a successful fit. Please reject gradient_tolerance at or above 1 here, no value there can express a convergence requirement.

jlarson4 · 2026-09-16T17:35:56Z

+    result = fit_sparse_probe(features, labels, k=2, class_weight=None, seed=2)
+    sweep = sweep_sparse_probe(features, labels, ks=[1], class_weight=None, seed=2)
+
+    assert result.class_weight is None


This checks that the metadata round-trips, not that the weighting was applied, and the fixture is balanced, so balanced weights are all 1.0 and neither mode can be distinguished: forcing the weights to ones, and separately ignoring class_weight=None so the fit is always balanced, both leave every test green, while on an imbalanced fixture the two modes give visibly different accuracies and coefficients. Please fit an imbalanced fixture in both modes and assert the coefficient vectors differ, each against a Newton reference that takes the same sample weights.

jlarson4 · 2026-09-16T17:35:56Z

+    assert result.gradient_inf_norm <= 1e-7
+
+
+def test_metrics_match_an_independent_heldout_recompute_at_the_logit_zero_threshold():


This recompute runs under the default preprocess="none", where preprocess_mean and preprocess_scale are identity, so standardizing the held-out rows with their own statistics instead of the training ones leaves every test green while changing the held-out metrics, which is train-test leakage in the module whose headline contract is leakage safety. Please run this recompute under preprocess="standardize" as well, applying the returned preprocess_mean and preprocess_scale to the raw held-out rows, so held-out rows standardized with their own statistics fail it.

jlarson4 · 2026-09-16T17:35:56Z

+        )
+
+
+def test_default_tolerance_accepts_large_scale_activations():


This test's fixture converges to an achieved gradient roughly ten times under the plain absolute 1e-7, so the scale-relative bound is never exercised: reverting line 385 to the bare gradient_tolerance leaves every test green. Can we adjust to use a fixture whose accepted fit lands between the absolute and the relative thresholds? A per-coordinate scale in the thousands should give this without exhausting the evaluation budget. We will want to assert on the achieved gradient against the applied bound, so the test fails when the bound goes back to absolute.

jlarson4 · 2026-09-16T17:37:46Z

+        l2_strength: Positive coefficient penalty in the logistic objective.
+        seed: Local CPU-generator seed used only for the stratified split.
+        max_iter: Maximum LBFGS iterations.
+        gradient_tolerance: Required final objective-gradient infinity norm.


Both public docstrings (here and line 652) say gradient_tolerance is the "required final objective-gradient infinity norm", but since the fix the accepted quantity is that tolerance times max(1, initial gradient infinity norm), so the API reference now describes a bound the code does not apply while the guide describes the one it does. Please reword both to the scale-relative bound, in whatever form you settle on from my comment on line 385.

janmenjayap added 4 commits September 12, 2026 13:04

janmenjayap changed the base branch from main to dev September 12, 2026 10:34

jlarson4 reviewed Sep 14, 2026

View reviewed changes

janmenjayap added 5 commits September 16, 2026 00:42

fix(sparse_probing): make gradient acceptance scale-relative

deb0047

fix(sparse_probing): stamp the realized k on each sweep result

a03b090

test(sparse_probing): pin metrics to the held-out rows and control la…

a22ad22

…bels

docs(sparse_probing): correct result-field list and add grouped-split…

2080984

… caveat

test(sparse_probing): cover the single-example-per-class rejection

5d0fe38

jlarson4 reviewed Sep 16, 2026

View reviewed changes

Navigation Menu

Uh oh!

feat(sparse_probing): add leakage-safe k-sparse probing over activation tensors - #1774

feat(sparse_probing): add leakage-safe k-sparse probing over activation tensors#1774
janmenjayap wants to merge 9 commits into
TransformerLensOrg:devfrom
janmenjayap:feat/sparse-probing

janmenjayap commented Sep 12, 2026

Uh oh!

jlarson4 left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

janmenjayap commented Sep 15, 2026

Uh oh!

jlarson4 left a comment

Uh oh!

jlarson4 Sep 16, 2026

Uh oh!

jlarson4 Sep 16, 2026

Uh oh!

jlarson4 Sep 16, 2026

Uh oh!

jlarson4 Sep 16, 2026

Uh oh!

jlarson4 Sep 16, 2026

Uh oh!

jlarson4 Sep 16, 2026

Uh oh!

jlarson4 Sep 16, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		assert result.gradient_inf_norm <= 1e-7


		def test_metrics_match_an_independent_heldout_recompute_at_the_logit_zero_threshold():

		)


		def test_default_tolerance_accepts_large_scale_activations():

Uh oh!

Conversation

janmenjayap commented Sep 12, 2026

Description

What's included

Design decisions

Non-goals (tracked as follow-ups)

Claim boundaries

Test plan

Type of change

Checklist:

Uh oh!

jlarson4 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

janmenjayap commented Sep 15, 2026

Uh oh!

jlarson4 left a comment

Choose a reason for hiding this comment

Uh oh!

jlarson4 Sep 16, 2026

Choose a reason for hiding this comment

Uh oh!

jlarson4 Sep 16, 2026

Choose a reason for hiding this comment

Uh oh!

jlarson4 Sep 16, 2026

Choose a reason for hiding this comment

Uh oh!

jlarson4 Sep 16, 2026

Choose a reason for hiding this comment

Uh oh!

jlarson4 Sep 16, 2026

Choose a reason for hiding this comment

Uh oh!

jlarson4 Sep 16, 2026

Choose a reason for hiding this comment

Uh oh!

jlarson4 Sep 16, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants