From 132d798b50e0662aa6bcb6541e7d9a2f83df04e0 Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Sat, 12 Sep 2026 10:20:36 +0530 Subject: [PATCH 1/9] feat(sparse_probing): add leakage-safe core k-sparse probe fit 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. --- tests/unit/tools/test_sparse_probing.py | 320 +++++++++++ .../tools/analysis/sparse_probing.py | 528 ++++++++++++++++++ 2 files changed, 848 insertions(+) create mode 100644 tests/unit/tools/test_sparse_probing.py create mode 100644 transformer_lens/tools/analysis/sparse_probing.py diff --git a/tests/unit/tools/test_sparse_probing.py b/tests/unit/tools/test_sparse_probing.py new file mode 100644 index 000000000..0a3860cad --- /dev/null +++ b/tests/unit/tools/test_sparse_probing.py @@ -0,0 +1,320 @@ +"""Unit tests for leakage-safe model-free sparse probing.""" + +from dataclasses import FrozenInstanceError + +import pytest +import torch +from beartype.roar import BeartypeCallHintParamViolation + +from transformer_lens.tools.analysis.sparse_probing import ( + SparseProbeResult, + _binary_metrics, + fit_sparse_probe, +) + + +def _planted_data( + *, n_examples: int = 400, n_features: int = 16, seed: int = 0 +) -> tuple[torch.Tensor, torch.Tensor]: + generator = torch.Generator().manual_seed(seed) + labels = torch.arange(n_examples) % 2 + features = torch.randn(n_examples, n_features, generator=generator) + features[:, 3] += 2.5 * (2 * labels - 1) + permutation = torch.randperm(n_examples, generator=generator) + return features[permutation], labels[permutation] + + +def _balanced_objective_gradient( + features: torch.Tensor, + labels: torch.Tensor, + coefficients: torch.Tensor, + intercept: torch.Tensor, + l2_strength: float, +) -> torch.Tensor: + labels = labels.double() + count = labels.numel() + positive_count = labels.sum() + weights = torch.where( + labels == 1, + count / (2 * positive_count), + count / (2 * (count - positive_count)), + ) + residual = weights * (torch.sigmoid(features @ coefficients + intercept) - labels) / count + return torch.cat( + (features.T @ residual + l2_strength * coefficients, residual.sum().reshape(1)) + ) + + +def _newton_reference( + features: torch.Tensor, + labels: torch.Tensor, + l2_strength: float, +) -> torch.Tensor: + features = features.double() + labels = labels.double() + design = torch.cat((features, torch.ones(features.shape[0], 1, dtype=torch.float64)), dim=1) + count = labels.numel() + positive_count = labels.sum() + weights = torch.where( + labels == 1, + count / (2 * positive_count), + count / (2 * (count - positive_count)), + ) + penalty = torch.diag( + torch.tensor([l2_strength] * features.shape[1] + [0.0], dtype=torch.float64) + ) + parameters = torch.zeros(design.shape[1], dtype=torch.float64) + for _ in range(100): + probability = torch.sigmoid(design @ parameters) + residual = weights * (probability - labels) / count + gradient = design.T @ residual + penalty @ parameters + curvature = weights * probability * (1 - probability) / count + hessian = design.T @ (curvature[:, None] * design) + penalty + parameters -= torch.linalg.solve(hessian, gradient) + if float(gradient.abs().max()) < 1e-12: + break + return parameters + + +def test_fit_recovers_exact_train_only_mean_difference_and_planted_feature(): + features, labels = _planted_data() + + result = fit_sparse_probe(features, labels, k=1, positive_label=1, seed=17) + + train_features = features[result.train_indices] + train_labels = labels[result.train_indices] + expected_scores = train_features[train_labels == 1].mean(0) - train_features[ + train_labels == 0 + ].mean(0) + expected_selected = torch.argsort(expected_scores.abs(), descending=True, stable=True)[:1] + assert isinstance(result, SparseProbeResult) + assert result.selected_features.tolist() == [3] + assert torch.equal(result.selected_features, expected_selected) + assert torch.allclose(result.feature_scores, expected_scores.double(), atol=1e-6) + assert result.metrics.f1 > 0.98 + assert result.metrics.accuracy > 0.98 + assert result.k == 1 + assert result.max_iter == 200 + assert result.gradient_tolerance == 1e-7 + + +def test_positive_label_controls_score_sign_and_class_metadata(): + features, labels = _planted_data() + signed_labels = 1 - 2 * labels + + result = fit_sparse_probe(features, signed_labels, k=1, positive_label=-1, seed=11) + + assert result.positive_label == -1 + assert result.negative_label == 1 + assert result.feature_scores[3] > 0 + + +def test_unweighted_classification_policy_is_explicit(): + features, labels = _planted_data(n_examples=100, n_features=6) + + result = fit_sparse_probe(features, labels, k=2, class_weight=None, seed=2) + + assert result.class_weight is None + + +def test_standardization_is_train_only_and_heldout_values_do_not_change_selection(): + features, labels = _planted_data(n_examples=200, n_features=8) + first = fit_sparse_probe(features, labels, k=2, preprocess="standardize", seed=9) + changed = features.clone() + changed[first.test_indices] += 10_000 * torch.randn_like(changed[first.test_indices]) + + second = fit_sparse_probe(changed, labels, k=2, preprocess="standardize", seed=9) + + selected_train = features[first.train_indices][:, first.selected_features].double() + expected_mean = selected_train.mean(0) + expected_scale = selected_train.std(0, correction=0) + assert torch.equal(first.selected_features, second.selected_features) + assert torch.equal(first.feature_scores, second.feature_scores) + assert torch.allclose(first.preprocess_mean, expected_mean) + assert torch.allclose(first.preprocess_scale, expected_scale) + assert torch.equal(first.preprocess_mean, second.preprocess_mean) + assert torch.equal(first.preprocess_scale, second.preprocess_scale) + assert torch.equal(first.coefficients, second.coefficients) + assert torch.equal(first.intercept, second.intercept) + assert first.objective == second.objective + + +def test_none_preprocessing_has_identity_metadata_and_constant_tie_order(): + features = torch.zeros(20, 5) + labels = torch.arange(20) % 2 + + result = fit_sparse_probe(features, labels, k=3, preprocess="none", seed=1) + + assert result.selected_features.tolist() == [0, 1, 2] + assert torch.equal(result.preprocess_mean, torch.zeros(3, dtype=torch.float64)) + assert torch.equal(result.preprocess_scale, torch.ones(3, dtype=torch.float64)) + assert result.constant_features.tolist() == [True, True, True] + + +def test_lbfgs_matches_independent_newton_solution_and_gradient(): + features, labels = _planted_data(n_examples=120, n_features=4, seed=4) + l2_strength = 0.02 + result = fit_sparse_probe( + features, + labels, + k=4, + l2_strength=l2_strength, + gradient_tolerance=1e-7, + seed=3, + ) + train_features = features[result.train_indices][:, result.selected_features].double() + train_labels = labels[result.train_indices] + expected = _newton_reference(train_features, train_labels, l2_strength) + gradient = _balanced_objective_gradient( + train_features, + train_labels, + result.coefficients, + result.intercept, + l2_strength, + ) + + assert torch.allclose(result.coefficients, expected[:-1], atol=2e-6, rtol=2e-6) + assert result.intercept.item() == pytest.approx(expected[-1].item(), abs=2e-6) + assert float(gradient.abs().max()) == pytest.approx(result.gradient_inf_norm, abs=1e-12) + assert result.gradient_inf_norm <= 1e-7 + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64]) +def test_supported_source_dtypes_return_detached_cpu_float64_results(dtype): + features, labels = _planted_data(n_examples=80, n_features=6) + + result = fit_sparse_probe(features.to(dtype), labels, k=2, seed=0) + + for tensor in ( + result.feature_scores, + result.coefficients, + result.intercept, + result.preprocess_mean, + result.preprocess_scale, + ): + assert tensor.device.type == "cpu" + assert tensor.dtype == torch.float64 + assert not tensor.requires_grad + + +@pytest.mark.parametrize("dtype", [torch.bool, torch.int8, torch.uint8, torch.int64]) +def test_supported_label_dtypes(dtype): + features, labels = _planted_data(n_examples=80, n_features=6) + + result = fit_sparse_probe(features, labels.to(dtype), k=2, seed=0) + + assert result.positive_label == 1 + assert result.negative_label == 0 + + +def test_stratification_preserves_each_class_and_reports_realized_counts(): + generator = torch.Generator().manual_seed(0) + features = torch.randn(100, 5, generator=generator) + labels = torch.tensor([1] * 4 + [0] * 96) + features[:4, 0] += 3 + + result = fit_sparse_probe(features, labels, k=1, test_fraction=0.3, seed=2) + + assert result.train_positive_count == 2 + assert result.test_positive_count == 2 + assert result.train_negative_count == 67 + assert result.test_negative_count == 29 + assert 0 <= result.metrics.f1 <= 1 + assert not bool(torch.isin(result.train_indices, result.test_indices).any()) + assert torch.equal( + torch.cat((result.train_indices, result.test_indices)).sort().values, + torch.arange(features.shape[0]), + ) + + +def test_inputs_and_global_rng_are_unchanged_and_results_are_frozen(): + features, labels = _planted_data(n_examples=80, n_features=6) + features_before = features.clone() + labels_before = labels.clone() + torch.manual_seed(1234) + state_before = torch.random.get_rng_state() + + result = fit_sparse_probe(features, labels, k=2, seed=99) + + assert torch.equal(features, features_before) + assert torch.equal(labels, labels_before) + assert torch.equal(torch.random.get_rng_state(), state_before) + with pytest.raises(FrozenInstanceError): + setattr(result, "seed", 0) + + +def test_forced_nonconvergence_raises(): + features, labels = _planted_data(n_examples=100, n_features=5) + + with pytest.raises(RuntimeError, match="did not converge"): + fit_sparse_probe( + features, + labels, + k=3, + max_iter=1, + gradient_tolerance=1e-12, + ) + + +def test_binary_metrics_zero_division_policy(): + metrics = _binary_metrics(torch.tensor([-2.0, -1.0]), torch.tensor([0, 1])) + + assert metrics.true_positives == 0 + assert metrics.false_positives == 0 + assert metrics.false_negatives == 1 + assert metrics.precision == 0 + assert metrics.recall == 0 + assert metrics.f1 == 0 + + +@pytest.mark.parametrize( + ("features", "labels", "kwargs", "message"), + [ + (torch.ones(0, 2), torch.empty(0, dtype=torch.int64), {}, "non-empty"), + (torch.tensor([[1.0], [float("nan")]]), torch.tensor([0, 1]), {}, "finite"), + (torch.ones(4, 2), torch.zeros(4, dtype=torch.int64), {}, "exactly two"), + (torch.ones(6, 2), torch.tensor([0, 1, 2, 0, 1, 2]), {}, "exactly two"), + (torch.ones(4, 2), torch.tensor([0, 1, 0, 1]), {"positive_label": 2}, "positive_label"), + (torch.ones(4, 2), torch.tensor([0, 1, 0, 1]), {"k": 0}, "k must be"), + (torch.ones(4, 2), torch.tensor([0, 1, 0, 1]), {"k": 3}, "k must be"), + (torch.ones(4, 2), torch.tensor([0, 1, 0, 1]), {"test_fraction": 0}, "test_fraction"), + (torch.ones(4, 2), torch.tensor([0, 1, 0, 1]), {"l2_strength": 0}, "l2_strength"), + (torch.ones(4, 2), torch.tensor([0, 1, 0, 1]), {"class_weight": "bad"}, "class_weight"), + (torch.ones(4, 2), torch.tensor([0, 1, 0, 1]), {"seed": -1}, "seed"), + (torch.ones(4, 2), torch.tensor([0, 1, 0, 1]), {"preprocess": "bad"}, "preprocess"), + (torch.ones(4, 2), torch.tensor([0, 1, 0, 1]), {"max_iter": 0}, "max_iter"), + ( + torch.ones(4, 2), + torch.tensor([0, 1, 0, 1]), + {"gradient_tolerance": 0}, + "gradient_tolerance", + ), + ( + torch.ones(4, 2), + torch.tensor([0, 1, 0, 1]), + {"gradient_tolerance": 2}, + "gradient_tolerance", + ), + ], +) +def test_rejects_invalid_inputs(features, labels, kwargs, message): + options = {"k": 1, **kwargs} + with pytest.raises(ValueError, match=message): + fit_sparse_probe(features, labels, **options) + + +@pytest.mark.parametrize( + ("features", "labels"), + [ + ([[1.0], [2.0], [3.0], [4.0]], torch.tensor([0, 1, 0, 1])), + (torch.ones(4), torch.tensor([0, 1, 0, 1])), + (torch.ones(4, 2, dtype=torch.int64), torch.tensor([0, 1, 0, 1])), + (torch.ones(4, 2, dtype=torch.float8_e4m3fn), torch.tensor([0, 1, 0, 1])), + (torch.ones(4, 2), torch.tensor([0.0, 1.0, 0.0, 1.0])), + (torch.ones(4, 2), torch.tensor([[0, 1], [0, 1]])), + (torch.ones(4, 2), torch.tensor([0, 1, 0])), + ], +) +def test_runtime_typecheck_rejects_invalid_tensor_contracts(features, labels): + with pytest.raises(BeartypeCallHintParamViolation): + fit_sparse_probe(features, labels, k=1) diff --git a/transformer_lens/tools/analysis/sparse_probing.py b/transformer_lens/tools/analysis/sparse_probing.py new file mode 100644 index 000000000..72bc6617b --- /dev/null +++ b/transformer_lens/tools/analysis/sparse_probing.py @@ -0,0 +1,528 @@ +"""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. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Literal, cast + +import torch +from jaxtyping import Bool, Float, Int, Integer + +PreprocessMode = Literal["none", "standardize"] +ClassWeightMode = Literal["balanced"] | None +_SUPPORTED_FEATURE_DTYPES = ( + torch.float16, + torch.bfloat16, + torch.float32, + torch.float64, +) + + +@dataclass(frozen=True) +class SparseProbeMetrics: + """Held-out binary-classification metrics and confusion counts.""" + + true_positives: int + true_negatives: int + false_positives: int + false_negatives: int + accuracy: float + precision: float + recall: float + f1: float + + +@dataclass(frozen=True) +class SparseProbeResult: + """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``. + """ + + feature_scores: Float[torch.Tensor, "feature"] + selected_features: Int[torch.Tensor, "selected_feature"] + coefficients: Float[torch.Tensor, "selected_feature"] + intercept: Float[torch.Tensor, ""] + preprocess_mean: Float[torch.Tensor, "selected_feature"] + preprocess_scale: Float[torch.Tensor, "selected_feature"] + constant_features: Bool[torch.Tensor, "selected_feature"] + train_indices: Int[torch.Tensor, "train_example"] + test_indices: Int[torch.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: PreprocessMode + class_weight: ClassWeightMode + l2_strength: float + test_fraction: float + seed: int + k: int + max_iter: int + gradient_tolerance: float + objective: float + gradient_inf_norm: float + iterations: int + function_evaluations: int + + +@dataclass(frozen=True) +class _ValidatedInputs: + features: torch.Tensor + canonical_labels: torch.Tensor + positive_label: int + negative_label: int + k: int + test_fraction: float + preprocess: PreprocessMode + class_weight: ClassWeightMode + l2_strength: float + seed: int + max_iter: int + gradient_tolerance: float + + +@dataclass(frozen=True) +class _FitOutcome: + coefficients: torch.Tensor + intercept: torch.Tensor + objective: float + gradient_inf_norm: float + iterations: int + function_evaluations: int + + +def _finite_positive_real(value: int | float, name: str) -> float: + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise ValueError(f"{name} must be a finite positive real number, got {value!r}") + result = float(value) + if not math.isfinite(result) or result <= 0: + raise ValueError(f"{name} must be a finite positive real number, got {value!r}") + return result + + +def _positive_integer(value: int, name: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value < 1: + raise ValueError(f"{name} must be a positive integer, got {value!r}") + return value + + +def _validate_inputs( + features: torch.Tensor, + labels: torch.Tensor, + *, + k: int, + test_fraction: int | float, + positive_label: int | bool, + preprocess: str, + class_weight: str | None, + l2_strength: int | float, + seed: int, + max_iter: int, + gradient_tolerance: int | float, +) -> _ValidatedInputs: + if not isinstance(features, torch.Tensor): + raise ValueError(f"features must be a torch.Tensor, got {type(features).__name__}") + if features.ndim != 2: + raise ValueError(f"features must be two-dimensional, got shape {tuple(features.shape)}") + if features.shape[0] == 0 or features.shape[1] == 0: + raise ValueError(f"feature dimensions must be non-empty, got shape {tuple(features.shape)}") + if not torch.is_floating_point(features) or torch.is_complex(features): + raise ValueError(f"features must have a real floating-point dtype, got {features.dtype}") + if features.dtype not in _SUPPORTED_FEATURE_DTYPES: + raise ValueError( + "features must have a supported dtype: float16, bfloat16, float32, or float64; " + f"got {features.dtype}" + ) + if not bool(torch.isfinite(features).all()): + raise ValueError("features must contain only finite values") + + if not isinstance(labels, torch.Tensor): + raise ValueError(f"labels must be a torch.Tensor, got {type(labels).__name__}") + if labels.ndim != 1: + raise ValueError(f"labels must be one-dimensional, got shape {tuple(labels.shape)}") + if labels.shape[0] != features.shape[0]: + raise ValueError("features and labels must contain the same number of examples") + if torch.is_floating_point(labels) or torch.is_complex(labels): + raise ValueError(f"labels must have a Boolean or integer dtype, got {labels.dtype}") + + labels_cpu = labels.detach().to(device="cpu") + classes = [int(value) for value in torch.unique(labels_cpu, sorted=True).tolist()] + if len(classes) != 2: + raise ValueError(f"labels must contain exactly two classes, got {classes}") + if isinstance(positive_label, bool): + resolved_positive = int(positive_label) + elif isinstance(positive_label, int): + resolved_positive = positive_label + else: + raise ValueError(f"positive_label must be an integer or Boolean, got {positive_label!r}") + if resolved_positive not in classes: + raise ValueError(f"positive_label {resolved_positive} is not present in labels {classes}") + resolved_negative = classes[0] if classes[1] == resolved_positive else classes[1] + canonical_labels = labels_cpu == resolved_positive + counts = torch.bincount(canonical_labels.to(torch.int64), minlength=2) + if int(counts.min()) < 2: + raise ValueError("each class must contain at least two examples") + + validated_k = _positive_integer(k, "k") + if validated_k > features.shape[1]: + raise ValueError(f"k must be at most the feature count {features.shape[1]}, got {k}") + if isinstance(test_fraction, bool) or not isinstance(test_fraction, (int, float)): + raise ValueError(f"test_fraction must be a finite real in (0, 1), got {test_fraction!r}") + validated_fraction = float(test_fraction) + if not math.isfinite(validated_fraction) or not 0 < validated_fraction < 1: + raise ValueError(f"test_fraction must be a finite real in (0, 1), got {test_fraction!r}") + if preprocess not in ("none", "standardize"): + raise ValueError(f"preprocess must be 'none' or 'standardize', got {preprocess!r}") + validated_preprocess = cast(PreprocessMode, preprocess) + if class_weight not in ("balanced", None): + raise ValueError(f"class_weight must be 'balanced' or None, got {class_weight!r}") + validated_class_weight = cast(ClassWeightMode, class_weight) + validated_l2 = _finite_positive_real(l2_strength, "l2_strength") + if isinstance(seed, bool) or not isinstance(seed, int) or not 0 <= seed < 2**63: + 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: + raise ValueError( + "gradient_tolerance must be a finite real in (0, 1], " f"got {gradient_tolerance!r}" + ) + return _ValidatedInputs( + features=features.detach(), + canonical_labels=canonical_labels, + positive_label=resolved_positive, + negative_label=resolved_negative, + k=validated_k, + test_fraction=validated_fraction, + preprocess=validated_preprocess, + class_weight=validated_class_weight, + l2_strength=validated_l2, + seed=seed, + max_iter=validated_max_iter, + gradient_tolerance=validated_tolerance, + ) + + +def _stratified_split( + canonical_labels: torch.Tensor, + test_fraction: float, + generator: torch.Generator, +) -> tuple[torch.Tensor, torch.Tensor]: + train_parts = [] + test_parts = [] + for class_value in (False, True): + indices = torch.where(canonical_labels == class_value)[0] + permutation = torch.randperm(indices.numel(), generator=generator) + shuffled = indices[permutation] + test_count = min(max(math.ceil(test_fraction * indices.numel()), 1), indices.numel() - 1) + test_parts.append(shuffled[:test_count]) + train_parts.append(shuffled[test_count:]) + return torch.cat(train_parts), torch.cat(test_parts) + + +def _feature_scores( + features: torch.Tensor, + train_labels: torch.Tensor, + train_indices: torch.Tensor, +) -> torch.Tensor: + compute_dtype = torch.float64 if features.dtype == torch.float64 else torch.float32 + device_indices = train_indices.to(device=features.device) + train_features = features.index_select(0, device_indices).to(dtype=compute_dtype) + positive = train_labels.to(device=features.device) + scores = train_features[positive].mean(dim=0) - train_features[~positive].mean(dim=0) + return scores.to(device="cpu").to(dtype=torch.float64) + + +def _selected_data( + features: torch.Tensor, + selected_features: torch.Tensor, + train_indices: torch.Tensor, + test_indices: torch.Tensor, + preprocess: PreprocessMode, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + selected_device = selected_features.to(device=features.device) + train_device = train_indices.to(device=features.device) + test_device = test_indices.to(device=features.device) + train = ( + features.index_select(0, train_device) + .index_select(1, selected_device) + .to(device="cpu") + .to(dtype=torch.float64) + ) + test = ( + features.index_select(0, test_device) + .index_select(1, selected_device) + .to(device="cpu") + .to(dtype=torch.float64) + ) + raw_scale = train.std(dim=0, correction=0) + constant = raw_scale == 0 + if preprocess == "standardize": + mean = train.mean(dim=0) + scale = torch.where(constant, torch.ones_like(raw_scale), raw_scale) + return (train - mean) / scale, (test - mean) / scale, mean, scale, constant + mean = torch.zeros(train.shape[1], dtype=torch.float64) + scale = torch.ones(train.shape[1], dtype=torch.float64) + return train, test, mean, scale, constant + + +def _sample_weights(labels: torch.Tensor, class_weight: ClassWeightMode) -> torch.Tensor: + if class_weight is None: + return torch.ones_like(labels) + count = labels.numel() + positive_count = labels.sum() + return torch.where( + labels == 1, + count / (2 * positive_count), + count / (2 * (count - positive_count)), + ) + + +def _objective( + features: torch.Tensor, + labels: torch.Tensor, + parameters: torch.Tensor, + sample_weights: torch.Tensor, + l2_strength: float, +) -> torch.Tensor: + coefficients = parameters[:-1] + logits = features @ coefficients + parameters[-1] + losses = torch.nn.functional.binary_cross_entropy_with_logits(logits, labels, reduction="none") + return (sample_weights * losses).mean() + 0.5 * l2_strength * coefficients.square().sum() + + +def _objective_gradient( + features: torch.Tensor, + labels: torch.Tensor, + parameters: torch.Tensor, + sample_weights: torch.Tensor, + l2_strength: float, +) -> torch.Tensor: + coefficients = parameters[:-1] + residual = sample_weights * (torch.sigmoid(features @ coefficients + parameters[-1]) - labels) + residual = residual / labels.numel() + return torch.cat( + (features.T @ residual + l2_strength * coefficients, residual.sum().reshape(1)) + ) + + +def _fit_logistic( + features: torch.Tensor, + labels: torch.Tensor, + *, + class_weight: ClassWeightMode, + l2_strength: float, + max_iter: int, + gradient_tolerance: float, +) -> _FitOutcome: + labels = labels.to(dtype=torch.float64) + sample_weights = _sample_weights(labels, class_weight) + parameters = torch.zeros(features.shape[1] + 1, dtype=torch.float64, requires_grad=True) + optimizer = torch.optim.LBFGS( + [parameters], + max_iter=max_iter, + tolerance_grad=gradient_tolerance, + tolerance_change=max(torch.finfo(torch.float64).eps, gradient_tolerance**2), + line_search_fn="strong_wolfe", + ) + + def closure() -> torch.Tensor: + optimizer.zero_grad() + loss = _objective(features, labels, parameters, sample_weights, l2_strength) + loss.backward() + return loss + + try: + optimizer.step(closure) + except RuntimeError as error: + raise RuntimeError("sparse probe optimizer failed") from error + detached = parameters.detach() + objective = float( + _objective(features, labels, detached, sample_weights, l2_strength).detach().item() + ) + gradient = _objective_gradient(features, labels, detached, sample_weights, l2_strength) + 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") + if gradient_inf_norm > gradient_tolerance: + raise RuntimeError( + "sparse probe optimizer did not converge: " + f"gradient infinity norm {gradient_inf_norm:.6g} exceeds {gradient_tolerance:.6g}" + ) + state = optimizer.state[parameters] + return _FitOutcome( + coefficients=detached[:-1].clone(), + intercept=detached[-1].clone(), + objective=objective, + gradient_inf_norm=gradient_inf_norm, + iterations=int(state.get("n_iter", 0)), + function_evaluations=int(state.get("func_evals", 0)), + ) + + +def _binary_metrics(logits: torch.Tensor, labels: torch.Tensor) -> SparseProbeMetrics: + predictions = logits >= 0 + positive = labels.to(dtype=torch.bool) + true_positives = int((predictions & positive).sum().item()) + true_negatives = int((~predictions & ~positive).sum().item()) + false_positives = int((predictions & ~positive).sum().item()) + false_negatives = int((~predictions & positive).sum().item()) + count = labels.numel() + accuracy = (true_positives + true_negatives) / count + precision_denominator = true_positives + false_positives + recall_denominator = true_positives + false_negatives + precision = 0.0 if precision_denominator == 0 else true_positives / precision_denominator + recall = 0.0 if recall_denominator == 0 else true_positives / recall_denominator + f1_denominator = 2 * true_positives + false_positives + false_negatives + f1 = 0.0 if f1_denominator == 0 else 2 * true_positives / f1_denominator + return SparseProbeMetrics( + true_positives=true_positives, + true_negatives=true_negatives, + false_positives=false_positives, + false_negatives=false_negatives, + accuracy=accuracy, + precision=precision, + recall=recall, + f1=f1, + ) + + +def _fit_result( + validated: _ValidatedInputs, + train_indices: torch.Tensor, + test_indices: torch.Tensor, + feature_scores: torch.Tensor, + selected_features: torch.Tensor, +) -> SparseProbeResult: + train_features, test_features, mean, scale, constant = _selected_data( + validated.features, + selected_features, + train_indices, + test_indices, + validated.preprocess, + ) + train_labels = validated.canonical_labels[train_indices] + fit = _fit_logistic( + train_features, + train_labels, + class_weight=validated.class_weight, + l2_strength=validated.l2_strength, + max_iter=validated.max_iter, + gradient_tolerance=validated.gradient_tolerance, + ) + test_labels = validated.canonical_labels[test_indices] + metrics = _binary_metrics(test_features @ fit.coefficients + fit.intercept, test_labels) + return SparseProbeResult( + feature_scores=feature_scores, + selected_features=selected_features, + coefficients=fit.coefficients, + intercept=fit.intercept, + preprocess_mean=mean, + preprocess_scale=scale, + constant_features=constant, + train_indices=train_indices, + test_indices=test_indices, + metrics=metrics, + positive_label=validated.positive_label, + negative_label=validated.negative_label, + train_positive_count=int(train_labels.sum().item()), + train_negative_count=int((~train_labels).sum().item()), + test_positive_count=int(test_labels.sum().item()), + test_negative_count=int((~test_labels).sum().item()), + preprocess=validated.preprocess, + class_weight=validated.class_weight, + l2_strength=validated.l2_strength, + test_fraction=validated.test_fraction, + seed=validated.seed, + k=validated.k, + max_iter=validated.max_iter, + gradient_tolerance=validated.gradient_tolerance, + objective=fit.objective, + gradient_inf_norm=fit.gradient_inf_norm, + iterations=fit.iterations, + function_evaluations=fit.function_evaluations, + ) + + +def fit_sparse_probe( + features: Float[torch.Tensor, "example feature"], + labels: Bool[torch.Tensor, "example"] | Integer[torch.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 = 1e-2, + seed: int = 0, + max_iter: int = 200, + gradient_tolerance: int | float = 1e-7, +) -> 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. + + Args: + 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. + gradient_tolerance: Required final objective-gradient infinity norm. + + 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. + """ + validated = _validate_inputs( + features, + labels, + k=k, + test_fraction=test_fraction, + positive_label=positive_label, + preprocess=preprocess, + class_weight=class_weight, + l2_strength=l2_strength, + seed=seed, + max_iter=max_iter, + gradient_tolerance=gradient_tolerance, + ) + generator = torch.Generator(device="cpu").manual_seed(validated.seed) + train_indices, test_indices = _stratified_split( + validated.canonical_labels, validated.test_fraction, generator + ) + feature_scores = _feature_scores( + validated.features, validated.canonical_labels[train_indices], train_indices + ) + selected_features = torch.argsort(feature_scores.abs(), descending=True, stable=True)[ + : validated.k + ] + return _fit_result( + validated, + train_indices, + test_indices, + feature_scores, + selected_features, + ) From e6a98965e16924dfe6b092ef7fdf3365c936c2c7 Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Sat, 12 Sep 2026 10:22:09 +0530 Subject: [PATCH 2/9] feat(sparse_probing): add k-sweep, null controls, exports, and docs 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. --- docs/source/content/sparse_probing.md | 153 +++++++++++++ docs/source/index.md | 1 + tests/unit/tools/test_sparse_probing.py | 147 ++++++++++++ transformer_lens/tools/analysis/__init__.py | 16 ++ .../tools/analysis/sparse_probing.py | 213 ++++++++++++++++++ 5 files changed, 530 insertions(+) create mode 100644 docs/source/content/sparse_probing.md diff --git a/docs/source/content/sparse_probing.md b/docs/source/content/sparse_probing.md new file mode 100644 index 000000000..4efab91d6 --- /dev/null +++ b/docs/source/content/sparse_probing.md @@ -0,0 +1,153 @@ +# Sparse Probing + +Sparse probing measures how strongly binary-label information is concentrated in a small +set of activation coordinates. TransformerLens provides a model-independent fit and sweep API; +callers decide how to collect and aggregate activations. + +Probe performance establishes held-out decodability under the chosen split, preprocessing, +regularization, and sparsity. It does not establish causal model use, neuron monosemanticity, +or superposition. + +## Fit one probe + +```python +import torch + +from transformer_lens.tools.analysis import fit_sparse_probe + +features = torch.randn(200, 768) +labels = torch.arange(200) % 2 +result = fit_sparse_probe( + features, + labels, + k=8, + positive_label=1, + test_fraction=0.3, + preprocess="none", + class_weight="balanced", + l2_strength=1e-2, + seed=0, +) + +print(result.selected_features) +print(result.metrics.f1) +``` + +`features` must have shape `[example, feature]` and dtype float16, bfloat16, float32, or float64. +Labels must be a one-dimensional Boolean or integer tensor containing exactly two values, one equal +to `positive_label`. Each class needs at least two examples. + +## Selection and split contract + +The function creates a deterministic stratified split before computing any learned statistic. +For each class $c$, its realized test count is + +$$ +n_{\mathrm{test},c} = +\operatorname{clamp}(\lceil f n_c \rceil, 1, n_c - 1), +$$ + +where $f$ is `test_fraction`. Realized train and test counts are returned because small classes +can differ materially from the requested aggregate fraction. + +On the training split only, feature $j$ receives the signed score + +$$ +s_j = \mathbb{E}[X_j \mid y=\mathrm{positive}] + - \mathbb{E}[X_j \mid y=\mathrm{negative}]. +$$ + +The selected support contains the $k$ largest $|s_j|$. Equal scores are resolved by increasing +feature index. `preprocess="none"` fits the selected raw coordinates. With +`preprocess="standardize"`, selected columns are centered and scaled using training statistics; +zero-variance columns receive scale one. The same transform is then applied to held-out values. + +Because L2 regularization is scale-sensitive, preprocessing can change the fitted probe and +the resulting k-curve. A sweep therefore fixes preprocessing and L2 strength across every k. + +## Logistic objective + +For selected training features $X$, labels $y \in \{0,1\}$, coefficients $w$, and intercept $b$, +the optimizer minimizes + +$$ +\frac{1}{n}\sum_i \alpha_{y_i} +\operatorname{BCEWithLogits}(X_i w + b, y_i) ++ \frac{\lambda}{2}\lVert w\rVert_2^2, +\qquad +\alpha_c = \frac{n}{2n_c}. +$$ + +The displayed weights apply to the default `class_weight="balanced"`; pass `None` to use +$\alpha_c=1$. The intercept is not regularized. Positive predictions have nonnegative logits. +Accuracy, precision, recall, F1, and all four confusion counts are returned; precision or F1 is +zero when its denominator is zero. F1 is the primary sparse-probing metric. + +Feature-score reductions use float64 for float64 inputs and float32 otherwise. Selected matrices +move to CPU float64 for deterministic LBFGS fitting. All result tensors are detached CPU tensors. +The fit raises when output is non-finite or the final objective-gradient infinity norm exceeds +`gradient_tolerance`, which must lie in `(0, 1]`. +Results retain the requested `k`, `max_iter`, and `gradient_tolerance` alongside the realized +objective, gradient norm, iteration count, and convergence flag. + +## Sweep and controls + +```python +from transformer_lens.tools.analysis import sweep_sparse_probe + +sweep = sweep_sparse_probe( + features, + labels, + ks=[1, 2, 4, 8, 16], + n_random_subsets=20, + n_label_shuffles=20, + seed=0, +) + +for k, probe, random_control in zip( + sweep.ks, + sweep.results, + sweep.random_coordinate_controls, + strict=True, +): + print(k, probe.metrics.f1, random_control.f1.median()) +``` + +Every k uses the same split, preprocessing mode, and L2 strength. `ks` must be strictly +increasing and unique. + +Random-coordinate controls sample k distinct coordinates and fit the same classifier. +Label-shuffle controls permute training labels, repeat selection and fitting, and evaluate against +the untouched held-out labels. The API returns raw control supports and metric distributions; it +does not convert them into p-values or representation labels. A repeat count of zero disables that +control. + +Controls can be expensive: the sweep performs one main fit plus both requested control counts for +every k. Start with small grids and repeat counts. + +## Composing with cached activations + +Use `run_with_cache` to construct the feature matrix separately so token, position, batching, and +aggregation choices remain explicit: + +```python +tokens = model.to_tokens(prompts) +_, cache = model.run_with_cache(tokens, names_filter=[hook_name]) +features = cache[hook_name][:, -1, :] +result = fit_sparse_probe(features, labels, k=8) +``` + +The example selects the final sequence position, which is not appropriate for every dataset. +Choose the hook and position policy before interpreting selected coordinates. The API cannot detect +leakage already introduced into caller-provided `features`. + +## Reference + +The raw mean-difference selector and sparse-probing framing follow Wes Gurnee et al., +“Finding Neurons in a Haystack: Case Studies with Sparse Probing,” +[TMLR 2023](https://openreview.net/forum?id=JYs1R9IMJr), +with [reference code](https://github.com/wesg52/sparse-probing-paper). + +TransformerLens intentionally adds stratification, stable tie-breaking, explicit objective and +convergence diagnostics, and deterministic controls. Its optional centered standardization and +Torch LBFGS solver are not exact reproductions of the reference implementation. \ No newline at end of file diff --git a/docs/source/index.md b/docs/source/index.md index 565bb91cf..5a3be8302 100644 --- a/docs/source/index.md +++ b/docs/source/index.md @@ -62,6 +62,7 @@ content/jacobian_lens_fitting generated/demos/Jacobian_Lens_Decomposition_Demo content/backward_lens content/debugging_numerical_divergence +content/sparse_probing generated/demos/Main_Demo generated/demos/Exploratory_Analysis_Demo content/special_cases diff --git a/tests/unit/tools/test_sparse_probing.py b/tests/unit/tools/test_sparse_probing.py index 0a3860cad..91b7bfa1a 100644 --- a/tests/unit/tools/test_sparse_probing.py +++ b/tests/unit/tools/test_sparse_probing.py @@ -6,13 +6,27 @@ import torch from beartype.roar import BeartypeCallHintParamViolation +from transformer_lens.tools.analysis import ( + fit_sparse_probe as exported_fit_sparse_probe, +) +from transformer_lens.tools.analysis import ( + sweep_sparse_probe as exported_sweep_sparse_probe, +) from transformer_lens.tools.analysis.sparse_probing import ( + SparseProbeControl, SparseProbeResult, + SparseProbeSweep, _binary_metrics, fit_sparse_probe, + sweep_sparse_probe, ) +def test_public_analysis_exports(): + assert exported_fit_sparse_probe is fit_sparse_probe + assert exported_sweep_sparse_probe is sweep_sparse_probe + + def _planted_data( *, n_examples: int = 400, n_features: int = 16, seed: int = 0 ) -> tuple[torch.Tensor, torch.Tensor]: @@ -113,8 +127,10 @@ def test_unweighted_classification_policy_is_explicit(): features, labels = _planted_data(n_examples=100, n_features=6) 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 + assert sweep.results[0].class_weight is None def test_standardization_is_train_only_and_heldout_values_do_not_change_selection(): @@ -318,3 +334,134 @@ def test_rejects_invalid_inputs(features, labels, kwargs, message): def test_runtime_typecheck_rejects_invalid_tensor_contracts(features, labels): with pytest.raises(BeartypeCallHintParamViolation): fit_sparse_probe(features, labels, k=1) + + +def test_sweep_reuses_one_split_and_has_nested_selected_supports(): + features, labels = _planted_data(n_examples=180, n_features=10) + + sweep = sweep_sparse_probe(features, labels, ks=[1, 2, 4], seed=23) + independent = fit_sparse_probe(features, labels, k=2, seed=23) + + assert isinstance(sweep, SparseProbeSweep) + assert sweep.ks == (1, 2, 4) + for result in sweep.results: + assert torch.equal(result.train_indices, sweep.results[0].train_indices) + assert torch.equal(result.test_indices, sweep.results[0].test_indices) + assert torch.equal(sweep.results[0].selected_features, sweep.results[1].selected_features[:1]) + assert torch.equal(sweep.results[1].selected_features, sweep.results[2].selected_features[:2]) + assert torch.equal(sweep.results[1].selected_features, independent.selected_features) + assert sweep.results[1].metrics == independent.metrics + + +def test_disabled_controls_return_empty_aligned_results(): + features, labels = _planted_data(n_examples=100, n_features=8) + + sweep = sweep_sparse_probe(features, labels, ks=[1, 3], seed=1) + + for k, random_control, shuffle_control in zip( + sweep.ks, + sweep.random_coordinate_controls, + sweep.label_shuffle_controls, + strict=True, + ): + assert isinstance(random_control, SparseProbeControl) + assert random_control.supports.shape == (0, k) + assert shuffle_control.supports.shape == (0, k) + for metric_values in ( + random_control.accuracy, + random_control.precision, + random_control.recall, + random_control.f1, + shuffle_control.accuracy, + shuffle_control.precision, + shuffle_control.recall, + shuffle_control.f1, + ): + assert metric_values.shape == (0,) + assert metric_values.dtype == torch.float64 + + +def test_controls_are_deterministic_use_unique_supports_and_do_not_touch_global_rng(): + features, labels = _planted_data(n_examples=140, n_features=12) + torch.manual_seed(919) + state_before = torch.random.get_rng_state() + + first = sweep_sparse_probe( + features, + labels, + ks=[2], + n_random_subsets=4, + n_label_shuffles=4, + seed=5, + ) + second = sweep_sparse_probe( + features, + labels, + ks=[2], + n_random_subsets=4, + n_label_shuffles=4, + seed=5, + ) + + assert torch.equal(torch.random.get_rng_state(), state_before) + for left, right in ( + (first.random_coordinate_controls[0], second.random_coordinate_controls[0]), + (first.label_shuffle_controls[0], second.label_shuffle_controls[0]), + ): + assert torch.equal(left.supports, right.supports) + assert torch.equal(left.accuracy, right.accuracy) + assert torch.equal(left.precision, right.precision) + assert torch.equal(left.recall, right.recall) + assert torch.equal(left.f1, right.f1) + for support in left.supports: + assert torch.unique(support).numel() == 2 + + +def test_controls_remain_below_a_strong_planted_feature(): + features, labels = _planted_data(n_examples=200, n_features=32, seed=6) + + sweep = sweep_sparse_probe( + features, + labels, + ks=[1], + n_random_subsets=8, + n_label_shuffles=8, + seed=18, + ) + + actual_f1 = sweep.results[0].metrics.f1 + assert actual_f1 > 0.98 + assert float(sweep.random_coordinate_controls[0].f1.median()) < actual_f1 - 0.2 + assert float(sweep.label_shuffle_controls[0].f1.median()) < actual_f1 - 0.2 + + +def test_larger_k_improves_distributed_decodability_without_assigning_a_representation_label(): + generator = torch.Generator().manual_seed(77) + labels = torch.arange(800) % 2 + features = torch.randn(800, 20, generator=generator) + features[:, :4] += 0.55 * (2 * labels[:, None] - 1) + permutation = torch.randperm(800, generator=generator) + + sweep = sweep_sparse_probe(features[permutation], labels[permutation], ks=[1, 4], seed=3) + + assert sweep.results[1].metrics.f1 > sweep.results[0].metrics.f1 + 0.08 + assert not hasattr(sweep, "representation_label") + + +@pytest.mark.parametrize( + ("ks", "kwargs", "message"), + [ + ([], {}, "ks must"), + ([1, 1], {}, "strictly increasing"), + ([2, 1], {}, "strictly increasing"), + ([1, 9], {}, "feature count"), + ([True], {}, "positive integers"), + ([1], {"n_random_subsets": -1}, "n_random_subsets"), + ([1], {"n_label_shuffles": -1}, "n_label_shuffles"), + ], +) +def test_sweep_rejects_invalid_grid_and_control_counts(ks, kwargs, message): + features, labels = _planted_data(n_examples=80, n_features=8) + + with pytest.raises(ValueError, match=message): + sweep_sparse_probe(features, labels, ks=ks, **kwargs) diff --git a/transformer_lens/tools/analysis/__init__.py b/transformer_lens/tools/analysis/__init__.py index 5d93a557f..7cd2d140f 100644 --- a/transformer_lens/tools/analysis/__init__.py +++ b/transformer_lens/tools/analysis/__init__.py @@ -22,6 +22,8 @@ anchored coordinate patching (offline and dynamic/hooked). - projection_kernel: Basis-invariant subspace overlap and TransformerBridge attention-head OQ/OK/OV affinity. + - sparse_probing: Leakage-safe k-sparse binary probes over supplied + activation tensors, with train-only selection and raw null controls. """ from transformer_lens.tools.analysis.attribution_patching import ( @@ -76,6 +78,14 @@ projection_kernel, random_projection_kernel_moments, ) +from transformer_lens.tools.analysis.sparse_probing import ( + SparseProbeControl, + SparseProbeMetrics, + SparseProbeResult, + SparseProbeSweep, + fit_sparse_probe, + sweep_sparse_probe, +) __all__ = [ "AttentionHeadRef", @@ -99,6 +109,10 @@ "ProjectedFactor", "ProjectionKernelResult", "RandomSubspaceReference", + "SparseProbeControl", + "SparseProbeMetrics", + "SparseProbeResult", + "SparseProbeSweep", "SubspaceBasis", "VocabularyRanking", "WeightLayout", @@ -106,6 +120,7 @@ "attribution_patch", "direct_logit_attribution", "estimate_occupancy", + "fit_sparse_probe", "get_act_patch_direct_path", "get_act_patch_direct_path_all_sources", "get_sparse_decomposition", @@ -114,4 +129,5 @@ "random_projection_kernel_moments", "solve_coordinate_patch", "solve_coordinate_patch_positions", + "sweep_sparse_probe", ] diff --git a/transformer_lens/tools/analysis/sparse_probing.py b/transformer_lens/tools/analysis/sparse_probing.py index 72bc6617b..7cd475d2c 100644 --- a/transformer_lens/tools/analysis/sparse_probing.py +++ b/transformer_lens/tools/analysis/sparse_probing.py @@ -9,6 +9,7 @@ from __future__ import annotations import math +from collections.abc import Sequence from dataclasses import dataclass from typing import Literal, cast @@ -77,6 +78,28 @@ class SparseProbeResult: function_evaluations: int +@dataclass(frozen=True) +class SparseProbeControl: + """Raw held-out metric distributions for one control at one sparsity.""" + + supports: Int[torch.Tensor, "repeat selected_feature"] + accuracy: Float[torch.Tensor, "repeat"] + precision: Float[torch.Tensor, "repeat"] + recall: Float[torch.Tensor, "repeat"] + f1: Float[torch.Tensor, "repeat"] + + +@dataclass(frozen=True) +class SparseProbeSweep: + """Probe results and aligned controls over a strictly increasing k-grid.""" + + ks: tuple[int, ...] + results: tuple[SparseProbeResult, ...] + random_coordinate_controls: tuple[SparseProbeControl, ...] + label_shuffle_controls: tuple[SparseProbeControl, ...] + seed: int + + @dataclass(frozen=True) class _ValidatedInputs: features: torch.Tensor @@ -455,6 +478,59 @@ def _fit_result( ) +def _fit_control( + validated: _ValidatedInputs, + train_indices: torch.Tensor, + test_indices: torch.Tensor, + selected_features: torch.Tensor, + train_labels: torch.Tensor, +) -> SparseProbeMetrics: + train_features, test_features, _, _, _ = _selected_data( + validated.features, + selected_features, + train_indices, + test_indices, + validated.preprocess, + ) + fit = _fit_logistic( + train_features, + train_labels, + class_weight=validated.class_weight, + l2_strength=validated.l2_strength, + max_iter=validated.max_iter, + gradient_tolerance=validated.gradient_tolerance, + ) + test_labels = validated.canonical_labels[test_indices] + return _binary_metrics(test_features @ fit.coefficients + fit.intercept, test_labels) + + +def _control_result( + supports: list[torch.Tensor], + metrics: list[SparseProbeMetrics], + k: int, +) -> SparseProbeControl: + support_tensor = torch.stack(supports) if supports else torch.empty((0, k), dtype=torch.int64) + + def metric_tensor(name: str) -> torch.Tensor: + return torch.tensor( + [float(getattr(metric, name)) for metric in metrics], dtype=torch.float64 + ) + + return SparseProbeControl( + supports=support_tensor, + accuracy=metric_tensor("accuracy"), + precision=metric_tensor("precision"), + recall=metric_tensor("recall"), + f1=metric_tensor("f1"), + ) + + +def _nonnegative_integer(value: int, name: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + raise ValueError(f"{name} must be a nonnegative integer, got {value!r}") + return value + + def fit_sparse_probe( features: Float[torch.Tensor, "example feature"], labels: Bool[torch.Tensor, "example"] | Integer[torch.Tensor, "example"], @@ -526,3 +602,140 @@ def fit_sparse_probe( feature_scores, selected_features, ) + + +def sweep_sparse_probe( + features: Float[torch.Tensor, "example feature"], + labels: Bool[torch.Tensor, "example"] | Integer[torch.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 = 1e-2, + n_random_subsets: int = 0, + n_label_shuffles: int = 0, + seed: int = 0, + max_iter: int = 200, + gradient_tolerance: int | float = 1e-7, +) -> 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. + + Args: + 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. + gradient_tolerance: Required final objective-gradient infinity norm. + + 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. + """ + if isinstance(ks, (str, bytes)) or not isinstance(ks, Sequence): + raise ValueError("ks must be a non-empty sequence of positive integers") + k_values = tuple(ks) + if not k_values: + raise ValueError("ks must be a non-empty sequence of positive integers") + if any(isinstance(k, bool) or not isinstance(k, int) or k < 1 for k in k_values): + raise ValueError("ks must contain positive integers") + if any(right <= left for left, right in zip(k_values, k_values[1:])): + raise ValueError("ks must be strictly increasing and unique") + random_count = _nonnegative_integer(n_random_subsets, "n_random_subsets") + shuffle_count = _nonnegative_integer(n_label_shuffles, "n_label_shuffles") + validated = _validate_inputs( + features, + labels, + k=k_values[-1], + test_fraction=test_fraction, + positive_label=positive_label, + preprocess=preprocess, + class_weight=class_weight, + l2_strength=l2_strength, + seed=seed, + max_iter=max_iter, + gradient_tolerance=gradient_tolerance, + ) + generator = torch.Generator(device="cpu").manual_seed(validated.seed) + train_indices, test_indices = _stratified_split( + validated.canonical_labels, validated.test_fraction, generator + ) + train_labels = validated.canonical_labels[train_indices] + feature_scores = _feature_scores(validated.features, train_labels, train_indices) + ranked_features = torch.argsort(feature_scores.abs(), descending=True, stable=True) + results = tuple( + _fit_result( + validated, + train_indices, + test_indices, + feature_scores, + ranked_features[:k], + ) + for k in k_values + ) + + random_controls = [] + shuffle_controls = [] + feature_count = validated.features.shape[1] + for k in k_values: + random_supports = [] + random_metrics = [] + for _ in range(random_count): + support = torch.randperm(feature_count, generator=generator)[:k].sort().values + random_supports.append(support) + random_metrics.append( + _fit_control( + validated, + train_indices, + test_indices, + support, + train_labels, + ) + ) + random_controls.append(_control_result(random_supports, random_metrics, k)) + + shuffle_supports = [] + shuffle_metrics = [] + for _ in range(shuffle_count): + permutation = torch.randperm(train_labels.numel(), generator=generator) + shuffled_labels = train_labels[permutation] + shuffled_scores = _feature_scores(validated.features, shuffled_labels, train_indices) + support = torch.argsort(shuffled_scores.abs(), descending=True, stable=True)[:k] + shuffle_supports.append(support) + shuffle_metrics.append( + _fit_control( + validated, + train_indices, + test_indices, + support, + shuffled_labels, + ) + ) + shuffle_controls.append(_control_result(shuffle_supports, shuffle_metrics, k)) + + return SparseProbeSweep( + ks=k_values, + results=results, + random_coordinate_controls=tuple(random_controls), + label_shuffle_controls=tuple(shuffle_controls), + seed=validated.seed, + ) From 1e5fdb4e38ecf49f1a100dff60168c6fd9df0917 Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Sat, 12 Sep 2026 10:56:12 +0530 Subject: [PATCH 3/9] test(sparse_probing): cover MPS-to-CPU transfer before float64 fit 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. --- tests/mps/test_mps_basic.py | 31 +++++++++++++++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/tests/mps/test_mps_basic.py b/tests/mps/test_mps_basic.py index bcf85709c..8b342f373 100644 --- a/tests/mps/test_mps_basic.py +++ b/tests/mps/test_mps_basic.py @@ -28,6 +28,7 @@ SubspaceBasis, projection_kernel, ) +from transformer_lens.tools.analysis.sparse_probing import fit_sparse_probe # Skip the entire module on non-MPS runners (Linux CI, CPU-only Macs) pytestmark = pytest.mark.skipif( @@ -307,3 +308,33 @@ def test_mps_loss_computation(): assert loss.item() > 0, "Loss should be positive" _cleanup(model) + + +# --------------------------------------------------------------------------- +# 4. Tooling: sparse probing (no model load) +# --------------------------------------------------------------------------- + + +def test_mps_sparse_probe_selects_cpu_before_float64(): + """fit_sparse_probe transfers selected [example, k] data to CPU before casting to float64. + + float64 is unsupported on MPS, so the selection path must move tensors off the Metal + device before the cast; doing the cast first would raise instead of silently falling + back to float32. + """ + generator = torch.Generator().manual_seed(0) + n_examples, n_features = 200, 8 + labels = torch.arange(n_examples) % 2 + features = torch.randn(n_examples, n_features, generator=generator) + features[:, 2] += 2.5 * (2 * labels - 1) + features = features.to(device="mps", dtype=torch.float32) + + result = fit_sparse_probe(features, labels, k=2, seed=0) + + assert result.coefficients.device.type == "cpu" + assert result.coefficients.dtype == torch.float64 + assert result.feature_scores.device.type == "cpu" + assert result.feature_scores.dtype == torch.float64 + assert not torch.isnan(result.coefficients).any() + + _cleanup() From f495ec35d824fa790102da8f6d5dfca0e40947a8 Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Sat, 12 Sep 2026 15:34:06 +0530 Subject: [PATCH 4/9] fix(sparse_probing): accept jaxtyping>=0.3 type-check exception and dtype policy Rebasing onto origin/dev pulled in jaxtyping>=0.3 (#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. --- tests/unit/tools/test_sparse_probing.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/tests/unit/tools/test_sparse_probing.py b/tests/unit/tools/test_sparse_probing.py index 91b7bfa1a..3d7439c90 100644 --- a/tests/unit/tools/test_sparse_probing.py +++ b/tests/unit/tools/test_sparse_probing.py @@ -4,8 +4,8 @@ import pytest import torch -from beartype.roar import BeartypeCallHintParamViolation +from tests.typecheck_errors import TYPECHECK_ERRORS from transformer_lens.tools.analysis import ( fit_sparse_probe as exported_fit_sparse_probe, ) @@ -288,6 +288,12 @@ def test_binary_metrics_zero_division_policy(): [ (torch.ones(0, 2), torch.empty(0, dtype=torch.int64), {}, "non-empty"), (torch.tensor([[1.0], [float("nan")]]), torch.tensor([0, 1]), {}, "finite"), + ( + torch.ones(4, 2, dtype=torch.float8_e4m3fn), + torch.tensor([0, 1, 0, 1]), + {}, + "supported dtype", + ), (torch.ones(4, 2), torch.zeros(4, dtype=torch.int64), {}, "exactly two"), (torch.ones(6, 2), torch.tensor([0, 1, 2, 0, 1, 2]), {}, "exactly two"), (torch.ones(4, 2), torch.tensor([0, 1, 0, 1]), {"positive_label": 2}, "positive_label"), @@ -325,14 +331,13 @@ def test_rejects_invalid_inputs(features, labels, kwargs, message): ([[1.0], [2.0], [3.0], [4.0]], torch.tensor([0, 1, 0, 1])), (torch.ones(4), torch.tensor([0, 1, 0, 1])), (torch.ones(4, 2, dtype=torch.int64), torch.tensor([0, 1, 0, 1])), - (torch.ones(4, 2, dtype=torch.float8_e4m3fn), torch.tensor([0, 1, 0, 1])), (torch.ones(4, 2), torch.tensor([0.0, 1.0, 0.0, 1.0])), (torch.ones(4, 2), torch.tensor([[0, 1], [0, 1]])), (torch.ones(4, 2), torch.tensor([0, 1, 0])), ], ) def test_runtime_typecheck_rejects_invalid_tensor_contracts(features, labels): - with pytest.raises(BeartypeCallHintParamViolation): + with pytest.raises(TYPECHECK_ERRORS): fit_sparse_probe(features, labels, k=1) From deb0047f63cbb527c48d774e789e53c51dd53881 Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Wed, 16 Sep 2026 00:42:13 +0530 Subject: [PATCH 5/9] fix(sparse_probing): make gradient acceptance scale-relative --- docs/source/content/sparse_probing.md | 3 ++- tests/unit/tools/test_sparse_probing.py | 19 +++++++++++++++++++ .../tools/analysis/sparse_probing.py | 11 ++++++++--- 3 files changed, 29 insertions(+), 4 deletions(-) diff --git a/docs/source/content/sparse_probing.md b/docs/source/content/sparse_probing.md index 4efab91d6..1943ab582 100644 --- a/docs/source/content/sparse_probing.md +++ b/docs/source/content/sparse_probing.md @@ -86,7 +86,8 @@ zero when its denominator is zero. F1 is the primary sparse-probing metric. Feature-score reductions use float64 for float64 inputs and float32 otherwise. Selected matrices move to CPU float64 for deterministic LBFGS fitting. All result tensors are detached CPU tensors. The fit raises when output is non-finite or the final objective-gradient infinity norm exceeds -`gradient_tolerance`, which must lie in `(0, 1]`. +`gradient_tolerance` times `max(1, initial gradient infinity norm)`, a scale-relative bound that +tracks the gradient magnitude at the starting parameters. `gradient_tolerance` must lie in `(0, 1]`. Results retain the requested `k`, `max_iter`, and `gradient_tolerance` alongside the realized objective, gradient norm, iteration count, and convergence flag. diff --git a/tests/unit/tools/test_sparse_probing.py b/tests/unit/tools/test_sparse_probing.py index 3d7439c90..e14208696 100644 --- a/tests/unit/tools/test_sparse_probing.py +++ b/tests/unit/tools/test_sparse_probing.py @@ -272,6 +272,25 @@ def test_forced_nonconvergence_raises(): ) +def test_default_tolerance_accepts_large_scale_activations(): + # Raw (unstandardized) activations with a per-coordinate std in the hundreds make the + # objective gradient large at the zero starting parameters. The acceptance bound is + # scale-relative, so a Newton-quality solve must still be accepted under the default + # gradient_tolerance instead of being rejected by a bare absolute threshold. + generator = torch.Generator().manual_seed(0) + n_examples, n_features = 1000, 768 + labels = torch.arange(n_examples) % 2 + features = 60.0 * torch.randn(n_examples, n_features, generator=generator) + features[:, 7] += 150.0 * (2 * labels - 1) + permutation = torch.randperm(n_examples, generator=generator) + features, labels = features[permutation], labels[permutation] + + result = fit_sparse_probe(features, labels, k=4, seed=1) + + assert isinstance(result, SparseProbeResult) + assert result.selected_features.numel() == 4 + + def test_binary_metrics_zero_division_policy(): metrics = _binary_metrics(torch.tensor([-2.0, -1.0]), torch.tensor([0, 1])) diff --git a/transformer_lens/tools/analysis/sparse_probing.py b/transformer_lens/tools/analysis/sparse_probing.py index 7cd475d2c..a1269d25f 100644 --- a/transformer_lens/tools/analysis/sparse_probing.py +++ b/transformer_lens/tools/analysis/sparse_probing.py @@ -356,9 +356,13 @@ def _fit_logistic( [parameters], max_iter=max_iter, tolerance_grad=gradient_tolerance, - tolerance_change=max(torch.finfo(torch.float64).eps, gradient_tolerance**2), + tolerance_change=0.0, line_search_fn="strong_wolfe", ) + initial_gradient = _objective_gradient( + features, labels, parameters.detach(), sample_weights, l2_strength + ) + initial_gradient_inf_norm = float(initial_gradient.abs().max().item()) def closure() -> torch.Tensor: optimizer.zero_grad() @@ -378,10 +382,11 @@ def closure() -> torch.Tensor: 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") - if gradient_inf_norm > gradient_tolerance: + acceptance_threshold = gradient_tolerance * max(1.0, initial_gradient_inf_norm) + if gradient_inf_norm > acceptance_threshold: raise RuntimeError( "sparse probe optimizer did not converge: " - f"gradient infinity norm {gradient_inf_norm:.6g} exceeds {gradient_tolerance:.6g}" + f"gradient infinity norm {gradient_inf_norm:.6g} exceeds {acceptance_threshold:.6g}" ) state = optimizer.state[parameters] return _FitOutcome( From a03b090b2cfe858dab6803a17440f3ed0f321bf5 Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Wed, 16 Sep 2026 00:48:18 +0530 Subject: [PATCH 6/9] fix(sparse_probing): stamp the realized k on each sweep result --- tests/unit/tools/test_sparse_probing.py | 1 + transformer_lens/tools/analysis/sparse_probing.py | 5 ++++- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/unit/tools/test_sparse_probing.py b/tests/unit/tools/test_sparse_probing.py index e14208696..ac67c6e7f 100644 --- a/tests/unit/tools/test_sparse_probing.py +++ b/tests/unit/tools/test_sparse_probing.py @@ -368,6 +368,7 @@ def test_sweep_reuses_one_split_and_has_nested_selected_supports(): assert isinstance(sweep, SparseProbeSweep) assert sweep.ks == (1, 2, 4) + assert tuple(result.k for result in sweep.results) == sweep.ks for result in sweep.results: assert torch.equal(result.train_indices, sweep.results[0].train_indices) assert torch.equal(result.test_indices, sweep.results[0].test_indices) diff --git a/transformer_lens/tools/analysis/sparse_probing.py b/transformer_lens/tools/analysis/sparse_probing.py index a1269d25f..3d2967647 100644 --- a/transformer_lens/tools/analysis/sparse_probing.py +++ b/transformer_lens/tools/analysis/sparse_probing.py @@ -432,6 +432,7 @@ def _fit_result( test_indices: torch.Tensor, feature_scores: torch.Tensor, selected_features: torch.Tensor, + k: int, ) -> SparseProbeResult: train_features, test_features, mean, scale, constant = _selected_data( validated.features, @@ -473,7 +474,7 @@ def _fit_result( l2_strength=validated.l2_strength, test_fraction=validated.test_fraction, seed=validated.seed, - k=validated.k, + k=k, max_iter=validated.max_iter, gradient_tolerance=validated.gradient_tolerance, objective=fit.objective, @@ -606,6 +607,7 @@ def fit_sparse_probe( test_indices, feature_scores, selected_features, + validated.k, ) @@ -694,6 +696,7 @@ def sweep_sparse_probe( test_indices, feature_scores, ranked_features[:k], + k, ) for k in k_values ) From a22ad2235603095a16788ad7a6cc123477a5edc0 Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Wed, 16 Sep 2026 00:59:44 +0530 Subject: [PATCH 7/9] test(sparse_probing): pin metrics to the held-out rows and control labels --- tests/unit/tools/test_sparse_probing.py | 108 ++++++++++++++++++++++++ 1 file changed, 108 insertions(+) diff --git a/tests/unit/tools/test_sparse_probing.py b/tests/unit/tools/test_sparse_probing.py index ac67c6e7f..7a18a4a30 100644 --- a/tests/unit/tools/test_sparse_probing.py +++ b/tests/unit/tools/test_sparse_probing.py @@ -17,6 +17,11 @@ SparseProbeResult, SparseProbeSweep, _binary_metrics, + _feature_scores, + _fit_logistic, + _selected_data, + _stratified_split, + _validate_inputs, fit_sparse_probe, sweep_sparse_probe, ) @@ -195,6 +200,47 @@ def test_lbfgs_matches_independent_newton_solution_and_gradient(): assert result.gradient_inf_norm <= 1e-7 +def test_metrics_match_an_independent_heldout_recompute_at_the_logit_zero_threshold(): + generator = torch.Generator().manual_seed(41) + n_examples, n_features = 300, 8 + labels = torch.arange(n_examples) % 2 + features = torch.randn(n_examples, n_features, generator=generator) + # Moderate separation so the held-out confusion matrix contains both false positives and + # false negatives in unequal counts. That makes precision and recall differ, so an + # independent recompute can catch a precision/recall swap. + features[:, 2] += 1.1 * (2 * labels - 1) + permutation = torch.randperm(n_examples, generator=generator) + features, labels = features[permutation], labels[permutation] + + result = fit_sparse_probe(features, labels, k=3, seed=7) + + test_features = features[result.test_indices][:, result.selected_features].double() + logits = test_features @ result.coefficients + result.intercept + predictions = logits >= 0 + positive = labels[result.test_indices].bool() + true_positives = int((predictions & positive).sum()) + true_negatives = int((~predictions & ~positive).sum()) + false_positives = int((predictions & ~positive).sum()) + false_negatives = int((~predictions & positive).sum()) + count = positive.numel() + + metrics = result.metrics + assert ( + metrics.true_positives, + metrics.true_negatives, + metrics.false_positives, + metrics.false_negatives, + ) == (true_positives, true_negatives, false_positives, false_negatives) + assert metrics.accuracy == (true_positives + true_negatives) / count + assert metrics.precision == true_positives / (true_positives + false_positives) + assert metrics.recall == true_positives / (true_positives + false_negatives) + assert metrics.f1 == 2 * true_positives / ( + 2 * true_positives + false_positives + false_negatives + ) + assert false_positives != false_negatives + assert metrics.precision != metrics.recall + + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64]) def test_supported_source_dtypes_return_detached_cpu_float64_results(dtype): features, labels = _planted_data(n_examples=80, n_features=6) @@ -473,6 +519,68 @@ def test_larger_k_improves_distributed_decodability_without_assigning_a_represen assert not hasattr(sweep, "representation_label") +def test_label_shuffle_control_is_scored_against_the_true_heldout_labels(): + features, labels = _planted_data(n_examples=160, n_features=10) + seed = 31 + + sweep = sweep_sparse_probe( + features, labels, ks=[2], n_random_subsets=0, n_label_shuffles=1, seed=seed + ) + + # Replay the sweep's RNG stream to recover the exact shuffled training labels and the + # support the single control fit used. Shuffling permutes only the training labels, so + # the fit must still be scored against the untouched held-out labels; rescoring the same + # fit against labels[test_indices] must reproduce the reported control metrics. + validated = _validate_inputs( + features, + labels, + k=2, + test_fraction=0.3, + positive_label=1, + preprocess="none", + class_weight="balanced", + l2_strength=1e-2, + seed=seed, + max_iter=200, + gradient_tolerance=1e-7, + ) + generator = torch.Generator(device="cpu").manual_seed(seed) + train_indices, test_indices = _stratified_split( + validated.canonical_labels, validated.test_fraction, generator + ) + train_labels = validated.canonical_labels[train_indices] + permutation = torch.randperm(train_labels.numel(), generator=generator) + shuffled_labels = train_labels[permutation] + shuffled_scores = _feature_scores(validated.features, shuffled_labels, train_indices) + support = torch.argsort(shuffled_scores.abs(), descending=True, stable=True)[:2] + + control = sweep.label_shuffle_controls[0] + assert torch.equal(support, control.supports[0]) + + train_features, test_features, *_ = _selected_data( + validated.features, support, train_indices, test_indices, validated.preprocess + ) + fit = _fit_logistic( + train_features, + shuffled_labels, + class_weight=validated.class_weight, + l2_strength=validated.l2_strength, + max_iter=validated.max_iter, + gradient_tolerance=validated.gradient_tolerance, + ) + true_test_labels = validated.canonical_labels[test_indices] + logits = test_features @ fit.coefficients + fit.intercept + rescored = _binary_metrics(logits, true_test_labels) + + assert float(control.accuracy[0]) == rescored.accuracy + assert float(control.precision[0]) == rescored.precision + assert float(control.recall[0]) == rescored.recall + assert float(control.f1[0]) == rescored.f1 + # Scoring the same fit against a different held-out label alignment would move the + # metrics, so the exact match above pins the scoring labels to the true held-out labels. + assert _binary_metrics(logits, ~true_test_labels).accuracy != rescored.accuracy + + @pytest.mark.parametrize( ("ks", "kwargs", "message"), [ From 2080984768d2d630bafb5808505f0cc9ab8a8532 Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Wed, 16 Sep 2026 01:13:15 +0530 Subject: [PATCH 8/9] docs(sparse_probing): correct result-field list and add grouped-split caveat --- docs/source/content/sparse_probing.md | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/docs/source/content/sparse_probing.md b/docs/source/content/sparse_probing.md index 1943ab582..c21c37c97 100644 --- a/docs/source/content/sparse_probing.md +++ b/docs/source/content/sparse_probing.md @@ -89,7 +89,8 @@ The fit raises when output is non-finite or the final objective-gradient infinit `gradient_tolerance` times `max(1, initial gradient infinity norm)`, a scale-relative bound that tracks the gradient magnitude at the starting parameters. `gradient_tolerance` must lie in `(0, 1]`. Results retain the requested `k`, `max_iter`, and `gradient_tolerance` alongside the realized -objective, gradient norm, iteration count, and convergence flag. +objective, gradient norm, iteration count, and function-evaluation count. There is no +convergence flag: a fit that misses the acceptance threshold raises instead of returning. ## Sweep and controls @@ -140,7 +141,10 @@ result = fit_sparse_probe(features, labels, k=8) The example selects the final sequence position, which is not appropriate for every dataset. Choose the hook and position policy before interpreting selected coordinates. The API cannot detect -leakage already introduced into caller-provided `features`. +leakage already introduced into caller-provided `features`. When several rows come from one source +prompt, for example multiple positions of the same sequence, keep all of those rows on one side of +the split; the row-level split is label-independent and can otherwise place rows of one prompt on +both sides, inflating held-out accuracy on grouped data. ## Reference From 5d0fe384f8a642cf6d49477b9a44d0d24e7d7771 Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Wed, 16 Sep 2026 01:18:44 +0530 Subject: [PATCH 9/9] test(sparse_probing): cover the single-example-per-class rejection --- tests/unit/tools/test_sparse_probing.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unit/tools/test_sparse_probing.py b/tests/unit/tools/test_sparse_probing.py index 7a18a4a30..57d662c7f 100644 --- a/tests/unit/tools/test_sparse_probing.py +++ b/tests/unit/tools/test_sparse_probing.py @@ -361,6 +361,7 @@ def test_binary_metrics_zero_division_policy(): ), (torch.ones(4, 2), torch.zeros(4, dtype=torch.int64), {}, "exactly two"), (torch.ones(6, 2), torch.tensor([0, 1, 2, 0, 1, 2]), {}, "exactly two"), + (torch.ones(4, 2), torch.tensor([0, 1, 1, 1]), {}, "at least two examples"), (torch.ones(4, 2), torch.tensor([0, 1, 0, 1]), {"positive_label": 2}, "positive_label"), (torch.ones(4, 2), torch.tensor([0, 1, 0, 1]), {"k": 0}, "k must be"), (torch.ones(4, 2), torch.tensor([0, 1, 0, 1]), {"k": 3}, "k must be"),