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