[Proposal] Add leakage-safe k-sparse probing over activation tensors

[janmenjayap]

Proposal

Add a dependency-free, model-free analysis module,
transformer_lens/tools/analysis/sparse_probing.py, for fitting binary k-sparse probes to a
supplied activation matrix X and label vector y.

The first PR contains only:

1. train/test splitting before all learned statistics;
2. train-only mean-difference feature selection;
3. deterministic binary logistic fitting on the selected coordinates;
4. held-out accuracy, precision, recall, and F1;
5. k-sweeps using one fixed split;
6. shuffled-label and random-coordinate controls.

Activation extraction, model wrappers, plotting, notebooks, multiclass probing, and optimal MIP
selection are follow-ups.

---

Motivation

A central mechanistic-interpretability question is how strongly a labeled feature is concentrated in
individual activation coordinates. Gurnee et al. study this with k-sparse linear probes: select at
most k neurons, retrain a classifier on those neurons, and evaluate held-out performance as k
changes.

TransformerLens already exposes activations through run_with_cache, but has no maintained probing
primitive. Users must repeatedly rebuild splitting, preprocessing, selection, fitting, and control
logic, where leakage is easy to introduce.

The model-free core is also generic machine-learning infrastructure. Whether TransformerLens should
own and maintain a logistic solver, rather than document composition with sklearn, is the primary
repository-fit question for maintainers.

The output supports hypotheses about coordinate concentration. It does not establish that the
model uses the decoded feature, that a selected neuron is monosemantic, or that a gradual k-curve is
by itself evidence of superposition.

---

Pitch

Proposed API

Names and optimizer details remain subject to maintainer approval.

from transformer_lens.tools.analysis import (
    fit_sparse_probe,
    sweep_sparse_probe,
)

probe = fit_sparse_probe(
    X,
    y,
    k=4,
    test_fraction=0.3,
    positive_label=1,
    preprocess="none",
    l2_strength=1e-2,
    seed=0,
)

sweep = sweep_sparse_probe(
    X,
    y,
    ks=[1, 2, 4, 8, 16],
    test_fraction=0.3,
    positive_label=1,
    preprocess="none",
    l2_strength=1e-2,
    n_random_subsets=20,
    n_label_shuffles=20,
    seed=0,
)

X has shape [example, feature]; y is a Boolean or integer label tensor containing exactly two
values, one equal to positive_label. Callers obtain X through any appropriate run_with_cache
workflow and choose their own position aggregation policy.

Design (algorithm)

For X ∈ R^{n×d} and binary labels y:

1. Validate finite floating X, binary Boolean/integer y, 1 <= k <= d,
   0 < test_fraction < 1, positive L2 strength, and at least two examples from each class.
2. Create one deterministic stratified split before computing any learned statistic. For class c,
   use n_test_c = clamp(ceil(test_fraction * n_c), 1, n_c - 1) and report realized counts.
3. On X_train only, compute the paper's raw mean-difference score
   mean(X_train[y=positive]) - mean(X_train[y=negative]); rank by absolute value with
   deterministic tie-breaking.
4. Select the top k coordinates. preprocess="none" matches the reference default. Optional
   "standardize" uses selected-column train mean and population standard deviation, maps a
   zero standard deviation to scale one, and applies the same transform to test data. Because L2 is
   scale-sensitive, this mode intentionally changes the fitted objective and is recorded in results.
5. Fit logits z = Xw + b by minimizing mean class-balanced binary cross-entropy plus
   l2_strength * ||w||² / 2; use weights n / (2 n_c) and do not regularize b.
6. Optimize selected matrices on CPU in float64 with Torch LBFGS and strong-Wolfe line search.
   Require finite output and a configurable final gradient-infinity-norm threshold; otherwise raise.
7. Predict positive at z >= 0. Report held-out confusion counts, accuracy, precision, recall, and
   F1, returning zero for precision/F1 when their denominators are zero. F1 is primary.
8. For a sweep, reuse the exact split across every strictly increasing unique k. Random-coordinate
   controls sample supports without replacement and fit the same classifier. Label-shuffle controls
   permute training labels, repeat selection/fitting, and evaluate against untouched test labels.

Control repeat counts are nonnegative; zero disables that control. The implementation uses a local
torch.Generator and must not mutate global RNG state. It adds no scikit-learn, Gurobi, plotting, or
dataset dependency.

Every k-sweep uses one fixed preprocessing mode and L2 strength. Curve shape is conditional on those
choices; PR1 performs no hyperparameter search and must not present one setting as canonical.

Main correctness risk (named)

Leakage and the probe-fit-vs-model-use conflation. The single thing most likely to be subtly wrong is information leaking from test into the probe — the most common trap is column standardization or unit-selection computed on the full dataset (not the train split), which silently inflates held-out accuracy. The close-second is interpretation: a high-accuracy probe means the feature is linearly decodable, not that the model uses it — over-claiming here is the field's canonical error.

Mitigations:

• Split first; compute selection and optional normalization on training data only. The function
  cannot detect leakage already introduced into caller-provided X, which documentation must state.
• Use deterministic stratification and balanced class weights rather than treating accuracy as safe
  under imbalance.
• Return precision, recall, and F1 next to accuracy.
• Return raw control distributions rather than assigning automatic "significance" labels.
• Document that controls detect some probe-pipeline failures, not causal model use.

Validation plan (falsifiable)

1. Selector oracle. Independently compute raw train-only class-mean differences and assert exact
   selected indices, signed scores, and deterministic tie-breaking.
2. Planted sparse feature. On seeded synthetic data with one known predictive coordinate, k=1
   selects it and held-out F1 exceeds a fixed threshold.
3. Distributed decodability. On a seeded feature spread across several coordinates, larger k
   improves held-out F1 by a fixed margin. The test makes no superposition claim.
4. Leakage guard. In standardize mode, hand-computed train-only means/stds and transformed test
   values match metadata; changing only held-out values cannot alter selected coordinates.
5. Controls. Label-shuffle and random-coordinate distributions are deterministic for a seed and
   remain below the planted-feature result by fixed margins.
6. Optimizer. Objective and gradient match independent formulas; a tiny analytic dataset reaches
   the configured gradient threshold, while forced non-convergence raises.
7. Validation. Empty classes, non-binary labels, non-finite values, invalid k, infeasible
   splits, duplicate k values, and constant selected columns follow explicit tested policies.

Scope of the first PR (vertical slice)

Exact paths:

• transformer_lens/tools/analysis/sparse_probing.py
• tests/unit/tools/test_sparse_probing.py
• exports from transformer_lens/tools/analysis/__init__.py
• docs/source/content/sparse_probing.md, including composition with run_with_cache
• docs/source/index.md, adding the guide to the Resources toctree

Public results should include class labels/counts, selected indices/scores, coefficients/intercept,
preprocessing metadata, split indices, held-out metrics, final objective/gradient/iterations, and raw
control distributions. They must not expose a .plot() method or assign a representation label.

Follow-up work

Deferred: activation extraction, position aggregation, model integration, plotting, notebooks,
multiclass labels, SAE-latent probing, MIP selection, feature batteries, and causal validation.

---

Alternatives

• Document sklearn composition instead: ecosystem-standard and lower maintenance, but leaves
  leakage-sensitive orchestration and result contracts to each user. This may be preferable if
  maintainers do not want TransformerLens to own a classifier optimizer.
• Depend on the paper repository: introduces research-code and Gurobi coupling.
• Ship MIP selection now: adds a solver dependency and changes the project from a small PR.
• Bundle model extraction: forces unresolved prompt/position aggregation policy into the core.

---

Correctness oracle

The selector oracle is the independently computed train-only raw class-mean difference. The
mean-difference heuristic is not required to approximate the exhaustive best subset; no such
guarantee exists. Logistic fitting is checked through analytic planted data, loss/convergence
metadata, and deterministic repetition.

---

Additional context

• Paper: Wes Gurnee et al., "Finding Neurons in a Haystack: Case Studies with Sparse
  Probing", TMLR 2023
  (OpenReview record).
• Reference code: wesg52/sparse-probing-paper pinned at a610e102:
  mean-difference selector,
  split and optional scaling,
  heuristic k-sweep,
  experiment defaults, and
  classification metrics.
• Paper alignment: rank coordinates by raw absolute class-mean difference, retrain a balanced
  logistic probe on the selected support, and treat held-out F1 as the primary metric.
• Verified reference defaults: the heuristic sweep first creates a random train/test split
  (test_set_frac=0.3 by default), optionally scales using X_train standard deviations, and ranks
  on X_train; normalization is disabled by default and, when enabled, divides by clamped standard
  deviation times ten without centering. Its sklearn SAGA fit uses balanced weights, L2, and
  max_iter=200 (other regularization/tolerance settings use sklearn defaults).
• Intentional divergences: stratify, offer explicit centered train-only standardization, use a
  defined Torch loss/L2 convention, add deterministic controls and convergence metadata, and state
  claim boundaries.
• TransformerLens touchpoints: current
  tools.analysis package,
  HookedRootModule.run_with_cache,
  TransformerBridge.run_with_cache,
  ActivationCache, and
  dependency manifest.
• Artifact requirements: none; PR1 tests use only synthetic tensors.

---

Maintainer decisions requested

1. Repository fit and scope: should TransformerLens own a model-free sparse-probe solver at all?
   If yes, should PR1 remain model-free with extraction/plotting deferred? Recommendation: proceed
   only if maintainers explicitly accept optimizer ownership; then keep PR1 model-free.
2. Optimizer: is the explicitly defined CPU-float64 Torch LBFGS objective acceptable instead of
   adding sklearn? Recommendation: yes, with final-gradient metadata and failure.
3. Controls: should random-coordinate and shuffled-training-label distributions be first-class
   sweep outputs? Recommendation: yes, disabled only when repetition count is zero.
4. Preprocessing: should "none" match the paper-code default, with optional train-only
   "standardize" explicitly documented as changing the L2 objective? Recommendation: yes.
5. Device/dtype: should score reductions run in at least float32 on the input device, while only
   the selected [example, k] matrices move to CPU float64 for deterministic LBFGS?
   Recommendation: yes; return CPU analysis results and document the transfer explicitly.

---

Checklist

• No competing TransformerLens or SAELens probing issue, PR, or implementation found as of
  2026-08-30.
• Paper heuristic and primary metric checked against arXiv:2305.01610.
• Reference defaults and leakage order checked against commit a610e102.
• Torch 2.10 CPU-float64 LBFGS prototype converged deterministically across FP16/BF16/FP32/FP64
  inputs; this does not replace tests on the repository's supported Torch/Python matrix.
• PR1 narrowed to model-free fitting, sweeping, controls, tests, exports, and docs.
• Maintainer decisions recorded before implementation.
• Follow-up extraction/plotting/model work tracked separately if requested.

---

Suggested labels: enhancement, tooling

---

[jlarson4]

Hi @janmenjayap & @lorenzozanee! I have answers to the Maintainer Decision Requests below:

1. Repository Fit & Scope – I think it is definitely fine for us to own this. If the scope becomes to large, we can always spin it off later, but I don't see any reason not to provide it.
2. Optimizer – I agree with the decision to forego sklearn. The plan as presented is a good one.
3. Controls – Yes, this is the correct approach. Keeping the random-coordinate and shuffled-training-label distributions in the result, disabled only at repetition count zero, is right, and it is the part of the proposal that most improves on ad-hoc probing scripts.
4. Preprocessing – Yes, complete as you recommended.
5. Device/dtype – Agree with this policy, please implement as stated

Please take these into consideration as you work on #1729.

[jlarson4]

@janmenjayap It appears the alternative solution for this was deleted by @lorenzozanee, feel free to take this over if you have time. If not, I'll leave it un assigned for someone else to take on

[janmenjayap]

Hello @jlarson4,
Thanks for the heads-up! I would be happy to take this over

[jlarson4]

I appreciate it @janmenjayap! I know from our past discussion that you had a pretty complete working solution in a separate branch, so I was hoping you'd be willing to wrap it up.

Thanks as always for your contributions. I am going to be putting up 3.9.0 within the next 12 hours, if you happen to get the PR up in that time I will make sure it gets included in the release.