Coverage for transformer_lens/tools/analysis/sparse_probing.py: 93%
442 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-10-10 05:10 +0000
1"""Leakage-safe k-sparse binary probes over activation tensors.
3This module is model-free: callers supply an ``[example, feature]`` activation
4matrix and binary labels. Feature selection and optional preprocessing use the
5training split only. Probe decodability does not establish causal model use,
6monosemanticity, or superposition.
7"""
9from __future__ import annotations
11import hashlib
12import math
13from collections.abc import Sequence
14from dataclasses import dataclass
15from typing import Literal, cast
17import torch
18from jaxtyping import Bool, Float, Int, Integer
20PreprocessMode = Literal["none", "standardize"]
21ClassWeightMode = Literal["balanced"] | None
22ArmMode = Literal["random_coordinate", "label_shuffle"]
23_SUPPORTED_FEATURE_DTYPES = (
24 torch.float16,
25 torch.bfloat16,
26 torch.float32,
27 torch.float64,
28)
31class SparseProbeConvergenceError(RuntimeError):
32 """A probe fit failed to converge or produced non-finite output.
34 Distinct from other ``RuntimeError``s so the sweep's control-rejection path
35 swallows only convergence failures, never unrelated runtime faults.
36 """
39@dataclass(frozen=True)
40class SparseProbeMetrics:
41 """Held-out binary-classification metrics and confusion counts.
43 Confusion counts, accuracy, precision, recall, and F1 are read at the logit-zero
44 threshold. ``roc_auc`` and ``average_precision`` are threshold-free and tie-aware:
45 held-out examples with equal logits share one rank and one threshold. The stratified
46 split always holds out both classes; if a class is absent, ``roc_auc`` is NaN, and
47 ``average_precision`` is NaN without positives and 1.0 without negatives.
48 """
50 true_positives: int
51 true_negatives: int
52 false_positives: int
53 false_negatives: int
54 accuracy: float
55 precision: float
56 recall: float
57 f1: float
58 roc_auc: float
59 average_precision: float
62@dataclass(frozen=True)
63class SparseProbeResult:
64 """Result of one train/test k-sparse binary probe fit.
66 All tensors are detached CPU tensors. Floating-point tensors use float64;
67 ``coefficients`` align with ``selected_features``.
68 """
70 feature_scores: Float[torch.Tensor, "feature"]
71 selected_features: Int[torch.Tensor, "selected_feature"]
72 coefficients: Float[torch.Tensor, "selected_feature"]
73 intercept: Float[torch.Tensor, ""]
74 preprocess_mean: Float[torch.Tensor, "selected_feature"]
75 preprocess_scale: Float[torch.Tensor, "selected_feature"]
76 constant_features: Bool[torch.Tensor, "selected_feature"]
77 train_indices: Int[torch.Tensor, "train_example"]
78 test_indices: Int[torch.Tensor, "test_example"]
79 metrics: SparseProbeMetrics
80 positive_label: int
81 negative_label: int
82 train_positive_count: int
83 train_negative_count: int
84 test_positive_count: int
85 test_negative_count: int
86 preprocess: PreprocessMode
87 std_floor: float
88 class_weight: ClassWeightMode
89 l2_strength: float
90 test_fraction: float
91 seed: int
92 k: int
93 max_iter: int
94 max_refinement_steps: int
95 decrement_tolerance: float
96 objective: float
97 gradient_inf_norm: float
98 newton_decrement: float
99 refinement_steps: int
100 iterations: int
101 function_evaluations: int
102 stop_reason: str
105@dataclass(frozen=True)
106class SparseProbeControl:
107 """Raw held-out metric distributions for one control at one sparsity.
109 Rows cover only converged repeats, so the row count can be below the request;
110 rejected repeats live in ``SparseProbeSweep.rejections``.
111 """
113 supports: Int[torch.Tensor, "repeat selected_feature"]
114 accuracy: Float[torch.Tensor, "repeat"]
115 precision: Float[torch.Tensor, "repeat"]
116 recall: Float[torch.Tensor, "repeat"]
117 f1: Float[torch.Tensor, "repeat"]
118 roc_auc: Float[torch.Tensor, "repeat"]
119 average_precision: Float[torch.Tensor, "repeat"]
122@dataclass(frozen=True)
123class SparseProbeRejection:
124 """A control fit excluded from the sweep because it did not converge.
126 Recorded instead of aborting so completed main probes and other controls survive.
127 """
129 arm: ArmMode
130 k: int
131 repeat: int
132 support: Int[torch.Tensor, "selected_feature"]
133 reason: str
136@dataclass(frozen=True)
137class SparseProbeSweep:
138 """Probe results and aligned controls over a strictly increasing k-grid."""
140 ks: tuple[int, ...]
141 results: tuple[SparseProbeResult, ...]
142 random_coordinate_controls: tuple[SparseProbeControl, ...]
143 label_shuffle_controls: tuple[SparseProbeControl, ...]
144 seed: int
145 rejections: tuple[SparseProbeRejection, ...] = ()
148@dataclass(frozen=True)
149class _ValidatedInputs:
150 features: torch.Tensor
151 canonical_labels: torch.Tensor
152 positive_label: int
153 negative_label: int
154 k: int
155 test_fraction: float
156 preprocess: PreprocessMode
157 std_floor: float
158 class_weight: ClassWeightMode
159 l2_strength: float
160 seed: int
161 max_iter: int
162 max_refinement_steps: int
163 decrement_tolerance: float
166@dataclass(frozen=True)
167class _FitOutcome:
168 coefficients: torch.Tensor
169 intercept: torch.Tensor
170 objective: float
171 gradient_inf_norm: float
172 newton_decrement: float
173 refinement_steps: int
174 iterations: int
175 function_evaluations: int
176 stop_reason: str
179def _finite_positive_real(value: int | float, name: str) -> float:
180 if isinstance(value, bool) or not isinstance(value, (int, float)): 180 ↛ 181line 180 didn't jump to line 181 because the condition on line 180 was never true
181 raise ValueError(f"{name} must be a finite positive real number, got {value!r}")
182 result = float(value)
183 if not math.isfinite(result) or result <= 0:
184 raise ValueError(f"{name} must be a finite positive real number, got {value!r}")
185 return result
188def _positive_integer(value: int, name: str) -> int:
189 if isinstance(value, bool) or not isinstance(value, int) or value < 1:
190 raise ValueError(f"{name} must be a positive integer, got {value!r}")
191 return value
194def _validate_inputs(
195 features: torch.Tensor,
196 labels: torch.Tensor,
197 *,
198 k: int,
199 test_fraction: int | float,
200 positive_label: int | bool,
201 preprocess: str,
202 std_floor: int | float,
203 class_weight: str | None,
204 l2_strength: int | float,
205 seed: int,
206 max_iter: int,
207 max_refinement_steps: int,
208 decrement_tolerance: int | float,
209) -> _ValidatedInputs:
210 if not isinstance(features, torch.Tensor): 210 ↛ 211line 210 didn't jump to line 211 because the condition on line 210 was never true
211 raise ValueError(f"features must be a torch.Tensor, got {type(features).__name__}")
212 if features.ndim != 2: 212 ↛ 213line 212 didn't jump to line 213 because the condition on line 212 was never true
213 raise ValueError(f"features must be two-dimensional, got shape {tuple(features.shape)}")
214 if features.shape[0] == 0 or features.shape[1] == 0:
215 raise ValueError(f"feature dimensions must be non-empty, got shape {tuple(features.shape)}")
216 if not torch.is_floating_point(features) or torch.is_complex(features): 216 ↛ 217line 216 didn't jump to line 217 because the condition on line 216 was never true
217 raise ValueError(f"features must have a real floating-point dtype, got {features.dtype}")
218 if features.dtype not in _SUPPORTED_FEATURE_DTYPES:
219 raise ValueError(
220 "features must have a supported dtype: float16, bfloat16, float32, or float64; "
221 f"got {features.dtype}"
222 )
223 if not bool(torch.isfinite(features).all()):
224 raise ValueError("features must contain only finite values")
226 if not isinstance(labels, torch.Tensor): 226 ↛ 227line 226 didn't jump to line 227 because the condition on line 226 was never true
227 raise ValueError(f"labels must be a torch.Tensor, got {type(labels).__name__}")
228 if labels.ndim != 1: 228 ↛ 229line 228 didn't jump to line 229 because the condition on line 228 was never true
229 raise ValueError(f"labels must be one-dimensional, got shape {tuple(labels.shape)}")
230 if labels.shape[0] != features.shape[0]: 230 ↛ 231line 230 didn't jump to line 231 because the condition on line 230 was never true
231 raise ValueError("features and labels must contain the same number of examples")
232 if torch.is_floating_point(labels) or torch.is_complex(labels): 232 ↛ 233line 232 didn't jump to line 233 because the condition on line 232 was never true
233 raise ValueError(f"labels must have a Boolean or integer dtype, got {labels.dtype}")
235 labels_cpu = labels.detach().to(device="cpu")
236 classes = [int(value) for value in torch.unique(labels_cpu, sorted=True).tolist()]
237 if len(classes) != 2:
238 raise ValueError(f"labels must contain exactly two classes, got {classes}")
239 if isinstance(positive_label, bool): 239 ↛ 240line 239 didn't jump to line 240 because the condition on line 239 was never true
240 resolved_positive = int(positive_label)
241 elif isinstance(positive_label, int): 241 ↛ 244line 241 didn't jump to line 244 because the condition on line 241 was always true
242 resolved_positive = positive_label
243 else:
244 raise ValueError(f"positive_label must be an integer or Boolean, got {positive_label!r}")
245 if resolved_positive not in classes:
246 raise ValueError(f"positive_label {resolved_positive} is not present in labels {classes}")
247 resolved_negative = classes[0] if classes[1] == resolved_positive else classes[1]
248 canonical_labels = labels_cpu == resolved_positive
249 counts = torch.bincount(canonical_labels.to(torch.int64), minlength=2)
250 if int(counts.min()) < 2:
251 raise ValueError("each class must contain at least two examples")
253 validated_k = _positive_integer(k, "k")
254 if validated_k > features.shape[1]:
255 raise ValueError(f"k must be at most the feature count {features.shape[1]}, got {k}")
256 if isinstance(test_fraction, bool) or not isinstance(test_fraction, (int, float)): 256 ↛ 257line 256 didn't jump to line 257 because the condition on line 256 was never true
257 raise ValueError(f"test_fraction must be a finite real in (0, 1), got {test_fraction!r}")
258 validated_fraction = float(test_fraction)
259 if not math.isfinite(validated_fraction) or not 0 < validated_fraction < 1:
260 raise ValueError(f"test_fraction must be a finite real in (0, 1), got {test_fraction!r}")
261 if preprocess not in ("none", "standardize"):
262 raise ValueError(f"preprocess must be 'none' or 'standardize', got {preprocess!r}")
263 if isinstance(std_floor, bool) or not isinstance(std_floor, (int, float)):
264 raise ValueError(f"std_floor must be a finite nonnegative real, got {std_floor!r}")
265 validated_std_floor = float(std_floor)
266 if not math.isfinite(validated_std_floor) or validated_std_floor < 0:
267 raise ValueError(f"std_floor must be a finite nonnegative real, got {std_floor!r}")
269 validated_preprocess = cast(PreprocessMode, preprocess)
270 if class_weight not in ("balanced", None):
271 raise ValueError(f"class_weight must be 'balanced' or None, got {class_weight!r}")
272 validated_class_weight = cast(ClassWeightMode, class_weight)
273 validated_l2 = _finite_positive_real(l2_strength, "l2_strength")
274 if isinstance(seed, bool) or not isinstance(seed, int) or not 0 <= seed < 2**63:
275 raise ValueError(f"seed must be an integer in [0, 2**63), got {seed!r}")
276 validated_max_iter = _positive_integer(max_iter, "max_iter")
277 validated_refinement_steps = _nonnegative_integer(max_refinement_steps, "max_refinement_steps")
278 validated_tolerance = _finite_positive_real(decrement_tolerance, "decrement_tolerance")
279 if validated_tolerance >= 1:
280 raise ValueError(
281 "decrement_tolerance must be a finite real in (0, 1), " f"got {decrement_tolerance!r}"
282 )
283 return _ValidatedInputs(
284 features=features.detach(),
285 canonical_labels=canonical_labels,
286 positive_label=resolved_positive,
287 negative_label=resolved_negative,
288 k=validated_k,
289 test_fraction=validated_fraction,
290 preprocess=validated_preprocess,
291 std_floor=validated_std_floor,
292 class_weight=validated_class_weight,
293 l2_strength=validated_l2,
294 seed=seed,
295 max_iter=validated_max_iter,
296 max_refinement_steps=validated_refinement_steps,
297 decrement_tolerance=validated_tolerance,
298 )
301def _stratified_split(
302 canonical_labels: torch.Tensor,
303 test_fraction: float,
304 generator: torch.Generator,
305) -> tuple[torch.Tensor, torch.Tensor]:
306 train_parts = []
307 test_parts = []
308 for class_value in (False, True):
309 indices = torch.where(canonical_labels == class_value)[0]
310 permutation = torch.randperm(indices.numel(), generator=generator)
311 shuffled = indices[permutation]
312 test_count = min(max(math.ceil(test_fraction * indices.numel()), 1), indices.numel() - 1)
313 test_parts.append(shuffled[:test_count])
314 train_parts.append(shuffled[test_count:])
315 return torch.cat(train_parts), torch.cat(test_parts)
318def _feature_scores(
319 features: torch.Tensor,
320 train_labels: torch.Tensor,
321 train_indices: torch.Tensor,
322) -> torch.Tensor:
323 compute_dtype = torch.float64 if features.dtype == torch.float64 else torch.float32
324 device_indices = train_indices.to(device=features.device)
325 train_features = features.index_select(0, device_indices).to(dtype=compute_dtype)
326 positive = train_labels.to(device=features.device)
327 scores = train_features[positive].mean(dim=0) - train_features[~positive].mean(dim=0)
328 return scores.to(device="cpu").to(dtype=torch.float64)
331def _selected_data(
332 features: torch.Tensor,
333 selected_features: torch.Tensor,
334 train_indices: torch.Tensor,
335 test_indices: torch.Tensor,
336 preprocess: PreprocessMode,
337 std_floor: float,
338) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
339 selected_device = selected_features.to(device=features.device)
340 train_device = train_indices.to(device=features.device)
341 test_device = test_indices.to(device=features.device)
342 train = (
343 features.index_select(0, train_device)
344 .index_select(1, selected_device)
345 .to(device="cpu")
346 .to(dtype=torch.float64)
347 )
348 test = (
349 features.index_select(0, test_device)
350 .index_select(1, selected_device)
351 .to(device="cpu")
352 .to(dtype=torch.float64)
353 )
354 raw_scale = train.std(dim=0, correction=0)
355 constant = raw_scale == 0
356 if preprocess == "standardize":
357 mean = train.mean(dim=0)
358 # Guards 0/0 when std_floor=0; zero-variance columns get a zero coefficient either way.
359 scale = torch.where(constant, torch.ones_like(raw_scale), raw_scale.clamp(min=std_floor))
360 return (train - mean) / scale, (test - mean) / scale, mean, scale, constant
361 mean = torch.zeros(train.shape[1], dtype=torch.float64)
362 scale = torch.ones(train.shape[1], dtype=torch.float64)
363 return train, test, mean, scale, constant
366def _sample_weights(labels: torch.Tensor, class_weight: ClassWeightMode) -> torch.Tensor:
367 if class_weight is None:
368 return torch.ones_like(labels)
369 count = labels.numel()
370 positive_count = labels.sum()
371 return torch.where(
372 labels == 1,
373 count / (2 * positive_count),
374 count / (2 * (count - positive_count)),
375 )
378def _objective(
379 features: torch.Tensor,
380 labels: torch.Tensor,
381 parameters: torch.Tensor,
382 sample_weights: torch.Tensor,
383 l2_strength: float,
384) -> torch.Tensor:
385 coefficients = parameters[:-1]
386 logits = features @ coefficients + parameters[-1]
387 losses = torch.nn.functional.binary_cross_entropy_with_logits(logits, labels, reduction="none")
388 return (sample_weights * losses).mean() + 0.5 * l2_strength * coefficients.square().sum()
391def _objective_gradient(
392 features: torch.Tensor,
393 labels: torch.Tensor,
394 parameters: torch.Tensor,
395 sample_weights: torch.Tensor,
396 l2_strength: float,
397) -> torch.Tensor:
398 coefficients = parameters[:-1]
399 residual = sample_weights * (torch.sigmoid(features @ coefficients + parameters[-1]) - labels)
400 residual = residual / labels.numel()
401 return torch.cat(
402 (features.T @ residual + l2_strength * coefficients, residual.sum().reshape(1))
403 )
406_LBFGS_TOLERANCE_GRAD = 1e-7
407_NEWTON_RIDGE = 1e-12
408_ARMIJO_SLOPE = 1e-4
409_MAX_BACKTRACKS = 60
412def _newton_direction(
413 features: torch.Tensor,
414 labels: torch.Tensor,
415 parameters: torch.Tensor,
416 sample_weights: torch.Tensor,
417 l2_strength: float,
418) -> tuple[torch.Tensor, torch.Tensor, float]:
419 """Return the gradient, the Newton direction ``H^-1 g``, and the decrement ``g^T H^-1 g / 2``."""
420 count = labels.numel()
421 design = torch.cat((features, torch.ones(count, 1, dtype=torch.float64)), dim=1)
422 probabilities = torch.sigmoid(design @ parameters)
423 gradient = _objective_gradient(features, labels, parameters, sample_weights, l2_strength)
424 curvature = sample_weights * probabilities * (1.0 - probabilities) / count
425 hessian = (design * curvature[:, None]).T @ design
426 hessian.diagonal()[: features.shape[1]] += l2_strength
427 # Saturated sigmoids zero the (unregularised) intercept curvature, so the solve is done on
428 # the Jacobi-equilibrated matrix with a ridge relative to each pivot; a raw factorisation
429 # fails on converged fits there.
430 diagonal = hessian.diagonal()
431 scale = torch.where(diagonal > 0, diagonal.rsqrt(), torch.ones_like(diagonal))
432 scaled = scale[:, None] * hessian * scale[None, :]
433 scaled.diagonal().add_(_NEWTON_RIDGE)
434 factor, info = torch.linalg.cholesky_ex(scaled)
435 if int(info.item()) != 0: 435 ↛ 436line 435 didn't jump to line 436 because the condition on line 435 was never true
436 raise SparseProbeConvergenceError(
437 "sparse probe optimizer did not converge: "
438 "Hessian factorization failed during Newton refinement"
439 )
440 direction = scale * torch.cholesky_solve((scale * gradient)[:, None], factor)[:, 0]
441 decrement = 0.5 * float((gradient @ direction).item())
442 if not math.isfinite(decrement) or decrement < 0.0: 442 ↛ 443line 442 didn't jump to line 443 because the condition on line 442 was never true
443 raise SparseProbeConvergenceError(
444 "sparse probe optimizer did not converge: "
445 f"Newton decrement {decrement!r} is not a finite nonnegative number"
446 )
447 return gradient, direction, decrement
450def _refine_with_newton(
451 features: torch.Tensor,
452 labels: torch.Tensor,
453 parameters: torch.Tensor,
454 sample_weights: torch.Tensor,
455 l2_strength: float,
456 *,
457 max_refinement_steps: int,
458 decrement_tolerance: float,
459) -> tuple[torch.Tensor, float, int]:
460 """Damped Newton steps until the decrement meets the tolerance; the check precedes any step."""
461 current = parameters
462 value = float(_objective(features, labels, current, sample_weights, l2_strength).item())
463 steps = 0
464 while True:
465 gradient, direction, decrement = _newton_direction(
466 features, labels, current, sample_weights, l2_strength
467 )
468 if decrement <= decrement_tolerance:
469 return current, decrement, steps
470 if steps >= max_refinement_steps:
471 raise SparseProbeConvergenceError(
472 "sparse probe optimizer did not converge: "
473 f"Newton decrement {decrement:.6g} exceeds {decrement_tolerance:.6g} "
474 f"after {steps} refinement steps"
475 )
476 slope = float((gradient @ direction).item())
477 step = 1.0
478 for _ in range(_MAX_BACKTRACKS): 478 ↛ 489line 478 didn't jump to line 489 because the loop on line 478 didn't complete
479 candidate = current - step * direction
480 candidate_value = float(
481 _objective(features, labels, candidate, sample_weights, l2_strength).item()
482 )
483 if math.isfinite(candidate_value) and ( 483 ↛ 487line 483 didn't jump to line 487 because the condition on line 483 was always true
484 candidate_value <= value - _ARMIJO_SLOPE * step * slope
485 ):
486 break
487 step *= 0.5
488 else:
489 raise SparseProbeConvergenceError(
490 "sparse probe optimizer did not converge: "
491 "Newton refinement line search found no descent step"
492 )
493 current, value = candidate, candidate_value
494 steps += 1
497def _fit_logistic(
498 features: torch.Tensor,
499 labels: torch.Tensor,
500 *,
501 class_weight: ClassWeightMode,
502 l2_strength: float,
503 max_iter: int,
504 max_refinement_steps: int,
505 decrement_tolerance: float,
506) -> _FitOutcome:
507 labels = labels.to(dtype=torch.float64)
508 sample_weights = _sample_weights(labels, class_weight)
509 parameters = torch.zeros(features.shape[1] + 1, dtype=torch.float64, requires_grad=True)
510 max_eval = max_iter * 5 // 4
511 # The solver's own stop is fixed so that no user argument can end the solve early;
512 # acceptance is decided afterwards by the Newton decrement.
513 optimizer = torch.optim.LBFGS(
514 [parameters],
515 max_iter=max_iter,
516 max_eval=max_eval,
517 tolerance_grad=_LBFGS_TOLERANCE_GRAD,
518 tolerance_change=0.0,
519 line_search_fn="strong_wolfe",
520 )
522 def closure() -> torch.Tensor:
523 optimizer.zero_grad()
524 loss = _objective(features, labels, parameters, sample_weights, l2_strength)
525 loss.backward()
526 return loss
528 try:
529 optimizer.step(closure)
530 except RuntimeError as error:
531 raise SparseProbeConvergenceError("sparse probe optimizer failed") from error
532 detached = parameters.detach()
533 lbfgs_gradient = _objective_gradient(features, labels, detached, sample_weights, l2_strength)
534 lbfgs_gradient_inf_norm = float(lbfgs_gradient.abs().max().item())
535 if not math.isfinite(lbfgs_gradient_inf_norm): 535 ↛ 536line 535 didn't jump to line 536 because the condition on line 535 was never true
536 raise SparseProbeConvergenceError("sparse probe optimizer produced non-finite output")
537 state = optimizer.state[parameters]
538 iterations = int(state.get("n_iter", 0))
539 function_evaluations = int(state.get("func_evals", 0))
540 # Same order LBFGS tests its exits; with tolerance_change at 0.0 the only exit left after
541 # the caps and the gradient stop is a strong-Wolfe search that made no progress.
542 if iterations >= max_iter:
543 stop_reason = "max_iter"
544 elif function_evaluations >= max_eval:
545 stop_reason = "max_eval"
546 elif lbfgs_gradient_inf_norm <= _LBFGS_TOLERANCE_GRAD:
547 stop_reason = "tolerance_grad"
548 else:
549 stop_reason = "line_search"
550 refined, decrement, refinement_steps = _refine_with_newton(
551 features,
552 labels,
553 detached,
554 sample_weights,
555 l2_strength,
556 max_refinement_steps=max_refinement_steps,
557 decrement_tolerance=decrement_tolerance,
558 )
559 objective = float(_objective(features, labels, refined, sample_weights, l2_strength).item())
560 gradient = _objective_gradient(features, labels, refined, sample_weights, l2_strength)
561 gradient_inf_norm = float(gradient.abs().max().item())
562 if not math.isfinite(objective) or not math.isfinite(gradient_inf_norm): 562 ↛ 563line 562 didn't jump to line 563 because the condition on line 562 was never true
563 raise SparseProbeConvergenceError("sparse probe optimizer produced non-finite output")
564 return _FitOutcome(
565 coefficients=refined[:-1].clone(),
566 intercept=refined[-1].clone(),
567 objective=objective,
568 gradient_inf_norm=gradient_inf_norm,
569 newton_decrement=decrement,
570 refinement_steps=refinement_steps,
571 iterations=iterations,
572 function_evaluations=function_evaluations,
573 stop_reason=stop_reason,
574 )
577def _roc_auc(logits: torch.Tensor, positive: torch.Tensor) -> float:
578 """Mann-Whitney ROC-AUC with average ranks for tied logits; NaN if a class is absent."""
579 positive_count = int(positive.sum().item())
580 negative_count = positive.numel() - positive_count
581 if positive_count == 0 or negative_count == 0:
582 return math.nan
583 _, inverse, counts = torch.unique(logits, return_inverse=True, return_counts=True)
584 counts = counts.to(dtype=torch.float64)
585 # A tie group occupying ranks start+1..end gets their mean, end - (count - 1) / 2.
586 average_ranks = (counts.cumsum(0) - (counts - 1) / 2)[inverse]
587 positive_rank_sum = float(average_ranks[positive].sum().item())
588 return (positive_rank_sum - positive_count * (positive_count + 1) / 2) / (
589 positive_count * negative_count
590 )
593def _average_precision(logits: torch.Tensor, positive: torch.Tensor) -> float:
594 """Step-wise average precision over distinct logits; NaN if there is no positive."""
595 positive_count = int(positive.sum().item())
596 if positive_count == 0:
597 return math.nan
598 _, inverse = torch.unique(logits, return_inverse=True)
599 group_count = int(inverse.max().item()) + 1
600 # Tied logits form one threshold, so the recall gained there is scored at a single precision.
601 group_positives = torch.zeros(group_count, dtype=torch.float64).index_add_(
602 0, inverse, positive.to(dtype=torch.float64)
603 )
604 group_sizes = torch.bincount(inverse, minlength=group_count).to(dtype=torch.float64)
605 positives_descending = group_positives.flip(0)
606 precision = positives_descending.cumsum(0) / group_sizes.flip(0).cumsum(0)
607 return float((positives_descending * precision).sum().item()) / positive_count
610def _binary_metrics(logits: torch.Tensor, labels: torch.Tensor) -> SparseProbeMetrics:
611 predictions = logits >= 0
612 positive = labels.to(dtype=torch.bool)
613 true_positives = int((predictions & positive).sum().item())
614 true_negatives = int((~predictions & ~positive).sum().item())
615 false_positives = int((predictions & ~positive).sum().item())
616 false_negatives = int((~predictions & positive).sum().item())
617 count = labels.numel()
618 accuracy = (true_positives + true_negatives) / count
619 precision_denominator = true_positives + false_positives
620 recall_denominator = true_positives + false_negatives
621 precision = 0.0 if precision_denominator == 0 else true_positives / precision_denominator
622 recall = 0.0 if recall_denominator == 0 else true_positives / recall_denominator
623 f1_denominator = 2 * true_positives + false_positives + false_negatives
624 f1 = 0.0 if f1_denominator == 0 else 2 * true_positives / f1_denominator
625 return SparseProbeMetrics(
626 true_positives=true_positives,
627 true_negatives=true_negatives,
628 false_positives=false_positives,
629 false_negatives=false_negatives,
630 accuracy=accuracy,
631 precision=precision,
632 recall=recall,
633 f1=f1,
634 roc_auc=_roc_auc(logits, positive),
635 average_precision=_average_precision(logits, positive),
636 )
639def _fit_result(
640 validated: _ValidatedInputs,
641 train_indices: torch.Tensor,
642 test_indices: torch.Tensor,
643 feature_scores: torch.Tensor,
644 selected_features: torch.Tensor,
645 k: int,
646) -> SparseProbeResult:
647 train_features, test_features, mean, scale, constant = _selected_data(
648 validated.features,
649 selected_features,
650 train_indices,
651 test_indices,
652 validated.preprocess,
653 validated.std_floor,
654 )
655 train_labels = validated.canonical_labels[train_indices]
656 fit = _fit_logistic(
657 train_features,
658 train_labels,
659 class_weight=validated.class_weight,
660 l2_strength=validated.l2_strength,
661 max_iter=validated.max_iter,
662 max_refinement_steps=validated.max_refinement_steps,
663 decrement_tolerance=validated.decrement_tolerance,
664 )
665 test_labels = validated.canonical_labels[test_indices]
666 metrics = _binary_metrics(test_features @ fit.coefficients + fit.intercept, test_labels)
667 return SparseProbeResult(
668 feature_scores=feature_scores,
669 selected_features=selected_features,
670 coefficients=fit.coefficients,
671 intercept=fit.intercept,
672 preprocess_mean=mean,
673 preprocess_scale=scale,
674 constant_features=constant,
675 train_indices=train_indices,
676 test_indices=test_indices,
677 metrics=metrics,
678 positive_label=validated.positive_label,
679 negative_label=validated.negative_label,
680 train_positive_count=int(train_labels.sum().item()),
681 train_negative_count=int((~train_labels).sum().item()),
682 test_positive_count=int(test_labels.sum().item()),
683 test_negative_count=int((~test_labels).sum().item()),
684 preprocess=validated.preprocess,
685 std_floor=validated.std_floor,
686 class_weight=validated.class_weight,
687 l2_strength=validated.l2_strength,
688 test_fraction=validated.test_fraction,
689 seed=validated.seed,
690 k=k,
691 max_iter=validated.max_iter,
692 max_refinement_steps=validated.max_refinement_steps,
693 decrement_tolerance=validated.decrement_tolerance,
694 objective=fit.objective,
695 gradient_inf_norm=fit.gradient_inf_norm,
696 newton_decrement=fit.newton_decrement,
697 refinement_steps=fit.refinement_steps,
698 iterations=fit.iterations,
699 function_evaluations=fit.function_evaluations,
700 stop_reason=fit.stop_reason,
701 )
704def _fit_control(
705 validated: _ValidatedInputs,
706 train_indices: torch.Tensor,
707 test_indices: torch.Tensor,
708 selected_features: torch.Tensor,
709 train_labels: torch.Tensor,
710) -> SparseProbeMetrics:
711 train_features, test_features, _, _, _ = _selected_data(
712 validated.features,
713 selected_features,
714 train_indices,
715 test_indices,
716 validated.preprocess,
717 validated.std_floor,
718 )
719 fit = _fit_logistic(
720 train_features,
721 train_labels,
722 class_weight=validated.class_weight,
723 l2_strength=validated.l2_strength,
724 max_iter=validated.max_iter,
725 max_refinement_steps=validated.max_refinement_steps,
726 decrement_tolerance=validated.decrement_tolerance,
727 )
728 test_labels = validated.canonical_labels[test_indices]
729 return _binary_metrics(test_features @ fit.coefficients + fit.intercept, test_labels)
732def _control_result(
733 supports: list[torch.Tensor],
734 metrics: list[SparseProbeMetrics],
735 k: int,
736) -> SparseProbeControl:
737 support_tensor = torch.stack(supports) if supports else torch.empty((0, k), dtype=torch.int64)
739 def metric_tensor(name: str) -> torch.Tensor:
740 return torch.tensor(
741 [float(getattr(metric, name)) for metric in metrics], dtype=torch.float64
742 )
744 return SparseProbeControl(
745 supports=support_tensor,
746 accuracy=metric_tensor("accuracy"),
747 precision=metric_tensor("precision"),
748 recall=metric_tensor("recall"),
749 f1=metric_tensor("f1"),
750 roc_auc=metric_tensor("roc_auc"),
751 average_precision=metric_tensor("average_precision"),
752 )
755def _nonnegative_integer(value: int, name: str) -> int:
756 if isinstance(value, bool) or not isinstance(value, int) or value < 0:
757 raise ValueError(f"{name} must be a nonnegative integer, got {value!r}")
758 return value
761def fit_sparse_probe(
762 features: Float[torch.Tensor, "example feature"],
763 labels: Bool[torch.Tensor, "example"] | Integer[torch.Tensor, "example"],
764 *,
765 k: int,
766 test_fraction: int | float = 0.3,
767 positive_label: int | bool = 1,
768 preprocess: str = "none",
769 std_floor: int | float = 1e-3,
770 class_weight: str | None = "balanced",
771 l2_strength: int | float = 1e-2,
772 seed: int = 0,
773 max_iter: int = 200,
774 max_refinement_steps: int = 25,
775 decrement_tolerance: int | float = 1e-12,
776) -> SparseProbeResult:
777 """Fit a train-only-selected k-sparse binary logistic probe.
779 The stratified split is created before feature scoring or optional
780 standardization. Floating result tensors are detached CPU float64 tensors.
782 Args:
783 features: Finite float16/bfloat16/float32/float64 tensor shaped
784 ``[example, feature]``.
785 labels: Boolean or integer binary labels shaped ``[example]``.
786 k: Number of coordinates selected by absolute train class-mean difference.
787 test_fraction: Requested held-out fraction within each class.
788 positive_label: Label defining the positive class and score sign.
789 preprocess: ``"none"`` or train-only ``"standardize"``.
790 std_floor: Lower bound on a non-constant column's ``"standardize"`` scale, so a
791 near-constant column is not amplified to unit scale; ``0`` disables it.
792 class_weight: ``"balanced"`` or ``None`` for unweighted BCE.
793 l2_strength: Positive coefficient penalty in the logistic objective.
794 seed: Local CPU-generator seed used only for the stratified split.
795 max_iter: Maximum LBFGS iterations before Newton refinement. LBFGS also stops once its
796 function evaluations reach ``max_iter * 5 // 4`` (``stop_reason="max_eval"``; the
797 line search can finish one past it), a budget equal to ``max_iter`` when
798 ``max_iter <= 3``.
799 max_refinement_steps: Maximum damped Newton steps after LBFGS; zero only checks.
800 decrement_tolerance: Largest accepted Newton decrement ``g^T H^-1 g / 2``, an
801 estimate of the objective gap to the optimum in nats; a larger gap raises.
803 Returns:
804 Selected support, fitted parameters, split/preprocessing metadata, metrics,
805 and optimizer diagnostics.
807 Raises:
808 ValueError: If inputs or options violate the binary-probe contract.
809 SparseProbeConvergenceError: If the optimizer fails or misses its
810 convergence threshold.
811 """
812 validated = _validate_inputs(
813 features,
814 labels,
815 k=k,
816 test_fraction=test_fraction,
817 positive_label=positive_label,
818 preprocess=preprocess,
819 std_floor=std_floor,
820 class_weight=class_weight,
821 l2_strength=l2_strength,
822 seed=seed,
823 max_iter=max_iter,
824 max_refinement_steps=max_refinement_steps,
825 decrement_tolerance=decrement_tolerance,
826 )
827 generator = torch.Generator(device="cpu").manual_seed(validated.seed)
828 train_indices, test_indices = _stratified_split(
829 validated.canonical_labels, validated.test_fraction, generator
830 )
831 feature_scores = _feature_scores(
832 validated.features, validated.canonical_labels[train_indices], train_indices
833 )
834 selected_features = torch.argsort(feature_scores.abs(), descending=True, stable=True)[
835 : validated.k
836 ]
837 return _fit_result(
838 validated,
839 train_indices,
840 test_indices,
841 feature_scores,
842 selected_features,
843 validated.k,
844 )
847def _control_generator(seed: int, k: int, arm: ArmMode, repeat: int) -> torch.Generator:
848 """Per-draw CPU generator keyed on ``(seed, k, arm, repeat)``; SHA-256 rather than
849 Python ``hash`` keeps each draw's seed stable across processes."""
850 # The key keeps the pre-ArmMode short arm names so released control draws reproduce.
851 key_arm = "random" if arm == "random_coordinate" else "shuffle"
852 key = f"{seed}:{k}:{key_arm}:{repeat}".encode()
853 derived = int.from_bytes(hashlib.sha256(key).digest()[:8], "little") & 0x7FFFFFFFFFFFFFFF
854 return torch.Generator(device="cpu").manual_seed(derived)
857def sweep_sparse_probe(
858 features: Float[torch.Tensor, "example feature"],
859 labels: Bool[torch.Tensor, "example"] | Integer[torch.Tensor, "example"],
860 *,
861 ks: Sequence[int],
862 test_fraction: int | float = 0.3,
863 positive_label: int | bool = 1,
864 preprocess: str = "none",
865 std_floor: int | float = 1e-3,
866 class_weight: str | None = "balanced",
867 l2_strength: int | float = 1e-2,
868 n_random_subsets: int = 0,
869 n_label_shuffles: int = 0,
870 seed: int = 0,
871 max_iter: int = 200,
872 max_refinement_steps: int = 25,
873 decrement_tolerance: int | float = 1e-12,
874) -> SparseProbeSweep:
875 """Fit sparse probes and optional controls over one fixed train/test split.
877 ``ks`` must contain strictly increasing positive integers. Random-coordinate
878 controls sample supports without replacement. Label-shuffle controls permute
879 training labels, repeat selection and fitting, then evaluate against the
880 untouched held-out labels. Control arrays contain raw metrics and do not
881 represent automatic significance tests.
883 Args:
884 features: Finite float16/bfloat16/float32/float64 tensor shaped
885 ``[example, feature]``.
886 labels: Boolean or integer binary labels shaped ``[example]``.
887 ks: Strictly increasing unique sparsity levels.
888 test_fraction: Requested held-out fraction within each class.
889 positive_label: Label defining the positive class and score sign.
890 preprocess: ``"none"`` or train-only ``"standardize"``.
891 std_floor: Lower bound on a non-constant column's ``"standardize"`` scale, shared by
892 every fit; ``0`` disables it.
893 class_weight: ``"balanced"`` or ``None``, shared by every fit.
894 l2_strength: Positive coefficient penalty shared by every fit.
895 n_random_subsets: Random-coordinate control fits per sparsity level.
896 n_label_shuffles: Shuffled-training-label control fits per sparsity level.
897 seed: Local CPU-generator seed for splitting and controls.
898 max_iter: Maximum LBFGS iterations per fit before Newton refinement. LBFGS also stops
899 once its function evaluations reach ``max_iter * 5 // 4``
900 (``stop_reason="max_eval"``; the line search can finish one past it), a budget
901 equal to ``max_iter`` when ``max_iter <= 3``.
902 max_refinement_steps: Maximum damped Newton steps per fit; zero only checks.
903 decrement_tolerance: Largest accepted Newton decrement ``g^T H^-1 g / 2`` per fit,
904 an estimate of the objective gap to the optimum in nats; a larger gap raises.
906 Returns:
907 Main probe results plus aligned raw control distributions. A control fit
908 that fails to converge is excluded and recorded in ``rejections`` rather
909 than aborting the sweep, so completed fits survive; each control's rows
910 cover only its converged repeats.
912 Raises:
913 ValueError: If the grid, controls, inputs, or options are invalid.
914 SparseProbeConvergenceError: If a main-probe fit fails to converge.
915 Control-fit convergence failures do not raise here — they are collected
916 in ``SparseProbeSweep.rejections``; any other error still propagates.
917 """
918 if isinstance(ks, (str, bytes)) or not isinstance(ks, Sequence): 918 ↛ 919line 918 didn't jump to line 919 because the condition on line 918 was never true
919 raise ValueError("ks must be a non-empty sequence of positive integers")
920 k_values = tuple(ks)
921 if not k_values:
922 raise ValueError("ks must be a non-empty sequence of positive integers")
923 if any(isinstance(k, bool) or not isinstance(k, int) or k < 1 for k in k_values):
924 raise ValueError("ks must contain positive integers")
925 if any(right <= left for left, right in zip(k_values, k_values[1:])):
926 raise ValueError("ks must be strictly increasing and unique")
927 random_count = _nonnegative_integer(n_random_subsets, "n_random_subsets")
928 shuffle_count = _nonnegative_integer(n_label_shuffles, "n_label_shuffles")
929 validated = _validate_inputs(
930 features,
931 labels,
932 k=k_values[-1],
933 test_fraction=test_fraction,
934 positive_label=positive_label,
935 preprocess=preprocess,
936 std_floor=std_floor,
937 class_weight=class_weight,
938 l2_strength=l2_strength,
939 seed=seed,
940 max_iter=max_iter,
941 max_refinement_steps=max_refinement_steps,
942 decrement_tolerance=decrement_tolerance,
943 )
944 generator = torch.Generator(device="cpu").manual_seed(validated.seed)
945 train_indices, test_indices = _stratified_split(
946 validated.canonical_labels, validated.test_fraction, generator
947 )
948 train_labels = validated.canonical_labels[train_indices]
949 feature_scores = _feature_scores(validated.features, train_labels, train_indices)
950 ranked_features = torch.argsort(feature_scores.abs(), descending=True, stable=True)
951 results = tuple(
952 _fit_result(
953 validated,
954 train_indices,
955 test_indices,
956 feature_scores,
957 ranked_features[:k],
958 k,
959 )
960 for k in k_values
961 )
963 random_controls = []
964 shuffle_controls = []
965 rejections: list[SparseProbeRejection] = []
966 feature_count = validated.features.shape[1]
967 for k in k_values:
968 random_supports = []
969 random_metrics = []
970 for repeat in range(random_count):
971 draw_generator = _control_generator(validated.seed, k, "random_coordinate", repeat)
972 support = torch.randperm(feature_count, generator=draw_generator)[:k].sort().values
973 # A failed control draw is recorded, not fatal: main probes and other controls survive.
974 try:
975 metrics = _fit_control(
976 validated, train_indices, test_indices, support, train_labels
977 )
978 except SparseProbeConvergenceError as error:
979 rejections.append(
980 SparseProbeRejection("random_coordinate", k, repeat, support, str(error))
981 )
982 continue
983 random_supports.append(support)
984 random_metrics.append(metrics)
985 random_controls.append(_control_result(random_supports, random_metrics, k))
987 shuffle_supports = []
988 shuffle_metrics = []
989 for repeat in range(shuffle_count):
990 draw_generator = _control_generator(validated.seed, k, "label_shuffle", repeat)
991 permutation = torch.randperm(train_labels.numel(), generator=draw_generator)
992 shuffled_labels = train_labels[permutation]
993 shuffled_scores = _feature_scores(validated.features, shuffled_labels, train_indices)
994 support = torch.argsort(shuffled_scores.abs(), descending=True, stable=True)[:k]
995 try:
996 metrics = _fit_control(
997 validated, train_indices, test_indices, support, shuffled_labels
998 )
999 except SparseProbeConvergenceError as error:
1000 rejections.append(
1001 SparseProbeRejection("label_shuffle", k, repeat, support, str(error))
1002 )
1003 continue
1004 shuffle_supports.append(support)
1005 shuffle_metrics.append(metrics)
1006 shuffle_controls.append(_control_result(shuffle_supports, shuffle_metrics, k))
1008 return SparseProbeSweep(
1009 ks=k_values,
1010 results=results,
1011 random_coordinate_controls=tuple(random_controls),
1012 label_shuffle_controls=tuple(shuffle_controls),
1013 seed=validated.seed,
1014 rejections=tuple(rejections),
1015 )