transformer_lens.tools.analysis.sparse_probing module¶

Leakage-safe k-sparse binary probes over activation tensors.

This module is model-free: callers supply an [example, feature] activation matrix and binary labels. Feature selection and optional preprocessing use the training split only. Probe decodability does not establish causal model use, monosemanticity, or superposition.

class transformer_lens.tools.analysis.sparse_probing.SparseProbeControl(supports: Int[Tensor, 'repeat selected_feature'], accuracy: Float[Tensor, 'repeat'], precision: Float[Tensor, 'repeat'], recall: Float[Tensor, 'repeat'], f1: Float[Tensor, 'repeat'])¶

Bases: object

Raw held-out metric distributions for one control at one sparsity.

accuracy: Float[Tensor, 'repeat']¶

f1: Float[Tensor, 'repeat']¶

precision: Float[Tensor, 'repeat']¶

recall: Float[Tensor, 'repeat']¶

supports: Int[Tensor, 'repeat selected_feature']¶

class transformer_lens.tools.analysis.sparse_probing.SparseProbeMetrics(true_positives: int, true_negatives: int, false_positives: int, false_negatives: int, accuracy: float, precision: float, recall: float, f1: float)¶

Bases: object

Held-out binary-classification metrics and confusion counts.

accuracy: float¶

f1: float¶

false_negatives: int¶

false_positives: int¶

precision: float¶

recall: float¶

true_negatives: int¶

true_positives: int¶

class transformer_lens.tools.analysis.sparse_probing.SparseProbeResult(feature_scores: Float[Tensor, 'feature'], selected_features: Int[Tensor, 'selected_feature'], coefficients: Float[Tensor, 'selected_feature'], intercept: Float[Tensor, ''], preprocess_mean: Float[Tensor, 'selected_feature'], preprocess_scale: Float[Tensor, 'selected_feature'], constant_features: Bool[Tensor, 'selected_feature'], train_indices: Int[Tensor, 'train_example'], test_indices: Int[Tensor, 'test_example'], metrics: SparseProbeMetrics, positive_label: int, negative_label: int, train_positive_count: int, train_negative_count: int, test_positive_count: int, test_negative_count: int, preprocess: Literal['none', 'standardize'], class_weight: Literal['balanced'] | None, l2_strength: float, test_fraction: float, seed: int, k: int, max_iter: int, max_refinement_steps: int, decrement_tolerance: float, objective: float, gradient_inf_norm: float, newton_decrement: float, refinement_steps: int, iterations: int, function_evaluations: int, stop_reason: str)¶

Bases: object

Result of one train/test k-sparse binary probe fit.

All tensors are detached CPU tensors. Floating-point tensors use float64; coefficients align with selected_features.

class_weight: Literal['balanced'] | None¶

coefficients: Float[Tensor, 'selected_feature']¶

constant_features: Bool[Tensor, 'selected_feature']¶

decrement_tolerance: float¶

feature_scores: Float[Tensor, 'feature']¶

function_evaluations: int¶

gradient_inf_norm: float¶

intercept: Float[Tensor, '']¶

iterations: int¶

k: int¶

l2_strength: float¶

max_iter: int¶

max_refinement_steps: int¶

metrics: SparseProbeMetrics¶

negative_label: int¶

newton_decrement: float¶

objective: float¶

positive_label: int¶

preprocess: Literal['none', 'standardize']¶

preprocess_mean: Float[Tensor, 'selected_feature']¶

preprocess_scale: Float[Tensor, 'selected_feature']¶

refinement_steps: int¶

seed: int¶

selected_features: Int[Tensor, 'selected_feature']¶

stop_reason: str¶

test_fraction: float¶

test_indices: Int[Tensor, 'test_example']¶

test_negative_count: int¶

test_positive_count: int¶

train_indices: Int[Tensor, 'train_example']¶

train_negative_count: int¶

train_positive_count: int¶

class transformer_lens.tools.analysis.sparse_probing.SparseProbeSweep(ks: tuple[int, ...], results: tuple[SparseProbeResult, ...], random_coordinate_controls: tuple[SparseProbeControl, ...], label_shuffle_controls: tuple[SparseProbeControl, ...], seed: int)¶

Bases: object

Probe results and aligned controls over a strictly increasing k-grid.

ks: tuple[int, ...]¶

label_shuffle_controls: tuple[SparseProbeControl, ...]¶

random_coordinate_controls: tuple[SparseProbeControl, ...]¶

results: tuple[SparseProbeResult, ...]¶

seed: int¶

transformer_lens.tools.analysis.sparse_probing.fit_sparse_probe(features: Float[Tensor, 'example feature'], labels: Bool[Tensor, 'example'] | Integer[Tensor, 'example'], *, k: int, test_fraction: int | float = 0.3, positive_label: int | bool = 1, preprocess: str = 'none', class_weight: str | None = 'balanced', l2_strength: int | float = 0.01, seed: int = 0, max_iter: int = 200, max_refinement_steps: int = 25, decrement_tolerance: int | float = 1e-12) → SparseProbeResult¶

Fit a train-only-selected k-sparse binary logistic probe.

The stratified split is created before feature scoring or optional standardization. Floating result tensors are detached CPU float64 tensors.

Parameters:

features – Finite float16/bfloat16/float32/float64 tensor shaped [example, feature].
labels – Boolean or integer binary labels shaped [example].
k – Number of coordinates selected by absolute train class-mean difference.
test_fraction – Requested held-out fraction within each class.
positive_label – Label defining the positive class and score sign.
preprocess – "none" or train-only "standardize".
class_weight – "balanced" or None for unweighted BCE.
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 before Newton refinement.
max_refinement_steps – Maximum damped Newton steps after LBFGS; zero only checks.
decrement_tolerance – Largest accepted Newton decrement g^T H^-1 g / 2, an estimate of the objective gap to the optimum in nats; a larger gap raises.

Returns:

Selected support, fitted parameters, split/preprocessing metadata, metrics, and optimizer diagnostics.

Raises:

ValueError – If inputs or options violate the binary-probe contract.
RuntimeError – If the optimizer fails or misses its convergence threshold.

transformer_lens.tools.analysis.sparse_probing.sweep_sparse_probe(features: Float[Tensor, 'example feature'], labels: Bool[Tensor, 'example'] | Integer[Tensor, 'example'], *, ks: Sequence[int], test_fraction: int | float = 0.3, positive_label: int | bool = 1, preprocess: str = 'none', class_weight: str | None = 'balanced', l2_strength: int | float = 0.01, n_random_subsets: int = 0, n_label_shuffles: int = 0, seed: int = 0, max_iter: int = 200, max_refinement_steps: int = 25, decrement_tolerance: int | float = 1e-12) → SparseProbeSweep¶

Fit sparse probes and optional controls over one fixed train/test split.

ks must contain strictly increasing positive integers. Random-coordinate controls sample supports without replacement. Label-shuffle controls permute training labels, repeat selection and fitting, then evaluate against the untouched held-out labels. Control arrays contain raw metrics and do not represent automatic significance tests.

Parameters:

features – Finite float16/bfloat16/float32/float64 tensor shaped [example, feature].
labels – Boolean or integer binary labels shaped [example].
ks – Strictly increasing unique sparsity levels.
test_fraction – Requested held-out fraction within each class.
positive_label – Label defining the positive class and score sign.
preprocess – "none" or train-only "standardize".
class_weight – "balanced" or None, shared by every fit.
l2_strength – Positive coefficient penalty shared by every fit.
n_random_subsets – Random-coordinate control fits per sparsity level.
n_label_shuffles – Shuffled-training-label control fits per sparsity level.
seed – Local CPU-generator seed for splitting and controls.
max_iter – Maximum LBFGS iterations per fit before Newton refinement.
max_refinement_steps – Maximum damped Newton steps per fit; zero only checks.
decrement_tolerance – Largest accepted Newton decrement g^T H^-1 g / 2 per fit, an estimate of the objective gap to the optimum in nats; a larger gap raises.

Returns:

Main probe results plus aligned raw control distributions.

Raises:

ValueError – If the grid, controls, inputs, or options are invalid.
RuntimeError – If any main or control fit fails to converge.