Coverage for transformer_lens/utilities/statistics.py: 89%
129 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-10-23 20:45 +0000
1"""Model-free statistics over a per-example axis of an analysis result.
3No model access: standard errors, a vectorized percentile bootstrap, a paired
4sign-flip permutation test and an exact Clopper-Pearson interval for success
5rates, plus the seeded-generator convention shared with sparse probing so every
6consumer derives independent, reproducible streams the same way.
7"""
9from __future__ import annotations
11import hashlib
12import itertools
13import math
14from dataclasses import dataclass
15from typing import Optional, Sequence, Tuple
17import torch
19_ALTERNATIVES = ("two-sided", "greater", "less")
20_STATISTICS = ("mean", "median")
23def derive_generator(seed: int, *labels: object) -> torch.Generator:
24 """CPU generator keyed on ``seed`` and any labels (metric name, arm, repeat...).
26 SHA-256 rather than ``hash`` keeps the derived seed stable across processes;
27 the mask keeps it in torch's accepted int64 range.
28 """
29 key = ":".join([str(seed), *(str(label) for label in labels)]).encode()
30 derived = int.from_bytes(hashlib.sha256(key).digest()[:8], "little") & 0x7FFFFFFFFFFFFFFF
31 return torch.Generator(device="cpu").manual_seed(derived)
34def _to_last(x: torch.Tensor, dim: int) -> torch.Tensor:
35 if not isinstance(x, torch.Tensor) or not x.is_floating_point(): 35 ↛ 36line 35 didn't jump to line 36 because the condition on line 35 was never true
36 raise TypeError("expected a floating-point tensor")
37 if x.ndim == 0: 37 ↛ 38line 37 didn't jump to line 38 because the condition on line 37 was never true
38 raise ValueError("expected at least one dimension to resample over")
39 return x.movedim(dim, -1)
42def standard_error(x: torch.Tensor, *, dim: int = -1) -> torch.Tensor:
43 """Standard error of the mean along ``dim`` (unbiased std over sqrt(n)); needs n >= 2."""
44 moved = _to_last(x, dim)
45 n = moved.shape[-1]
46 if n < 2:
47 raise ValueError(f"standard_error needs at least 2 samples along dim {dim}, got {n}")
48 return moved.std(dim=-1, unbiased=True) / math.sqrt(n)
51def bootstrap_ci(
52 x: torch.Tensor,
53 *,
54 dim: int = -1,
55 statistic: str = "mean",
56 confidence: float = 0.95,
57 n_resamples: int = 1000,
58 generator: Optional[torch.Generator] = None,
59 chunk_size: int = 256,
60) -> Tuple[torch.Tensor, torch.Tensor]:
61 """Percentile bootstrap interval of ``statistic`` along ``dim``, vectorized over the rest.
63 Resamples with replacement using a CPU ``generator`` (``derive_generator(0)``
64 when omitted), ``chunk_size`` resamples at a time so memory stays at
65 ``chunk_size × n`` per leading cell. Returns ``(low, high)`` with ``dim``
66 removed.
68 A percentile bootstrap cannot express uncertainty a sample does not contain:
69 a constant vector resamples to itself and the interval collapses to a point
70 whatever ``n`` is, and intervals on 4-8 samples are wide and skewed. For
71 binary success rates prefer an exact interval (:func:`success_rate_ci`).
72 """
73 if statistic not in _STATISTICS:
74 raise ValueError(f"statistic must be one of {_STATISTICS}, got {statistic!r}")
75 if not 0.0 < confidence < 1.0:
76 raise ValueError(f"confidence must be in (0, 1), got {confidence}")
77 if n_resamples < 1 or chunk_size < 1: 77 ↛ 78line 77 didn't jump to line 78 because the condition on line 77 was never true
78 raise ValueError("n_resamples and chunk_size must be >= 1")
79 moved = _to_last(x, dim)
80 n = moved.shape[-1]
81 if n < 1: 81 ↛ 82line 81 didn't jump to line 82 because the condition on line 81 was never true
82 raise ValueError("cannot bootstrap an empty axis")
83 gen = generator if generator is not None else derive_generator(0)
84 reduce = torch.mean if statistic == "mean" else lambda t, dim: t.median(dim=dim).values
85 stats = []
86 for start in range(0, n_resamples, chunk_size):
87 count = min(chunk_size, n_resamples - start)
88 idx = torch.randint(0, n, (count, n), generator=gen).to(moved.device)
89 resampled = moved[..., idx] # [..., count, n]
90 stats.append(reduce(resampled, dim=-1))
91 draws = torch.cat(stats, dim=-1).to(torch.promote_types(x.dtype, torch.float32))
92 alpha = (1.0 - confidence) / 2.0
93 low = torch.quantile(draws, alpha, dim=-1)
94 high = torch.quantile(draws, 1.0 - alpha, dim=-1)
95 return low.to(x.dtype), high.to(x.dtype)
98def sign_flip_permutation_pvalue(
99 effects: torch.Tensor,
100 *,
101 dim: int = -1,
102 n_permutations: int = 1000,
103 generator: Optional[torch.Generator] = None,
104 alternative: str = "two-sided",
105 exact_max_n: int = 12,
106) -> torch.Tensor:
107 """Paired sign-flip permutation p-value that the mean effect along ``dim`` is zero.
109 The null is that each per-example effect is symmetric about zero, so every
110 pattern of sign flips is equally likely. With ``n <= exact_max_n`` examples
111 all ``2**n`` patterns are enumerated and the p-value is exact and
112 deterministic; otherwise ``n_permutations`` random patterns are drawn from
113 the CPU ``generator`` and the p-value is ``(hits + 1) / (n_permutations + 1)``.
114 ``alternative`` is ``"two-sided"`` (``|null| >= |observed|``), ``"greater"``
115 or ``"less"``. A null mean within a relative ``100 * eps`` of the observed one
116 counts as a tie, absorbing summation-order error; there is no absolute floor, so
117 near-zero observed effects get essentially no tolerance. Returns a tensor with
118 ``dim`` removed.
119 """
120 if alternative not in _ALTERNATIVES:
121 raise ValueError(f"alternative must be one of {_ALTERNATIVES}, got {alternative!r}")
122 if n_permutations < 1: 122 ↛ 123line 122 didn't jump to line 123 because the condition on line 122 was never true
123 raise ValueError("n_permutations must be >= 1")
124 moved = _to_last(effects, dim).to(torch.promote_types(effects.dtype, torch.float32))
125 n = moved.shape[-1]
126 if n < 1: 126 ↛ 127line 126 didn't jump to line 127 because the condition on line 126 was never true
127 raise ValueError("cannot permute an empty axis")
128 observed = moved.mean(dim=-1, keepdim=True) # [..., 1]
129 if n <= exact_max_n:
130 signs = torch.tensor(
131 list(itertools.product((1.0, -1.0), repeat=n)), dtype=moved.dtype, device=moved.device
132 ) # [2^n, n]
133 null = (moved.unsqueeze(-2) * signs).mean(dim=-1) # [..., 2^n]
134 hits = _count_hits(null, observed, alternative)
135 return hits / signs.shape[0]
136 gen = generator if generator is not None else derive_generator(0)
137 hits_total = torch.zeros(moved.shape[:-1], dtype=moved.dtype, device=moved.device)
138 chunk = 256
139 for start in range(0, n_permutations, chunk):
140 count = min(chunk, n_permutations - start)
141 signs = (
142 torch.randint(0, 2, (count, n), generator=gen).to(moved.device).to(moved.dtype) * 2 - 1
143 )
144 null = (moved.unsqueeze(-2) * signs).mean(dim=-1)
145 hits_total = hits_total + _count_hits(null, observed, alternative)
146 return (hits_total + 1.0) / (n_permutations + 1.0)
149def _count_hits(null: torch.Tensor, observed: torch.Tensor, alternative: str) -> torch.Tensor:
150 # Relative to the observed statistic: tolerate reduction-order ties without an absolute floor.
151 tolerance = observed.abs() * (100 * torch.finfo(null.dtype).eps)
152 if alternative == "two-sided":
153 hit = null.abs() >= observed.abs() - tolerance
154 elif alternative == "greater":
155 hit = null >= observed - tolerance
156 else:
157 hit = null <= observed + tolerance
158 return hit.sum(dim=-1).to(null.dtype)
161@dataclass(frozen=True)
162class SuccessRateInterval:
163 """An exact Clopper-Pearson confidence interval around a success rate.
165 ``n_trials`` and ``n_successes`` are the counts behind the bounds.
166 """
168 point_estimate: float
169 ci_low: float
170 ci_high: float
171 n_trials: int
172 n_successes: int
173 confidence: float
176# A cap: bisection exits once float midpoints stop moving, typically after 53-90 halvings.
177_CLOPPER_PEARSON_STEPS = 200
180def _binomial_cdf(k: int, n: int, p: float) -> float:
181 """Exact ``P(X <= k)`` for ``X ~ Binomial(n, p)``.
183 Sums ``exp(log_term - peak)`` in log space: ``comb(n, i)`` and ``p ** i`` overflow at large
184 ``n``, and a recurrence from ``(1 - p) ** n`` underflows once ``n log(1 - p) < -745``.
185 """
186 if k < 0: 186 ↛ 187line 186 didn't jump to line 187 because the condition on line 186 was never true
187 return 0.0
188 if k >= n: 188 ↛ 189line 188 didn't jump to line 189 because the condition on line 188 was never true
189 return 1.0
190 if p <= 0.0: 190 ↛ 191line 190 didn't jump to line 191 because the condition on line 190 was never true
191 return 1.0
192 if p >= 1.0: 192 ↛ 193line 192 didn't jump to line 193 because the condition on line 192 was never true
193 return 0.0
194 log_p, log_q = math.log(p), math.log1p(-p)
195 log_n_factorial = math.lgamma(n + 1)
196 log_terms = [
197 log_n_factorial - math.lgamma(i + 1) - math.lgamma(n - i + 1) + i * log_p + (n - i) * log_q
198 for i in range(k + 1)
199 ]
200 peak = max(log_terms)
201 return min(1.0, math.exp(peak) * math.fsum(math.exp(term - peak) for term in log_terms))
204def _clopper_pearson_bound(n_successes: int, n_trials: int, tail: float, *, upper: bool) -> float:
205 """Solve one Clopper-Pearson bound by bisecting the exact binomial CDF.
207 Lower solves ``P(X >= k) = tail``, upper ``P(X <= k) = tail``; swapping their bracket
208 directions collapses every interval to ``[0, 1]``.
209 """
210 if upper and n_successes == n_trials:
211 return 1.0
212 if not upper and n_successes == 0:
213 return 0.0
214 low, high = 0.0, 1.0
215 for _ in range(_CLOPPER_PEARSON_STEPS): 215 ↛ 232line 215 didn't jump to line 232 because the loop on line 215 didn't complete
216 midpoint = (low + high) / 2.0
217 # Every later step would re-test this same midpoint and leave the result unchanged.
218 if midpoint == low or midpoint == high:
219 break
220 if upper:
221 # P(X <= k) decreases in p, so a CDF above the tail puts the bound above the midpoint.
222 if _binomial_cdf(n_successes, n_trials, midpoint) > tail:
223 low = midpoint
224 else:
225 high = midpoint
226 else:
227 # P(X >= k) increases in p, so a value above the tail puts the bound below the midpoint.
228 if 1.0 - _binomial_cdf(n_successes - 1, n_trials, midpoint) > tail:
229 high = midpoint
230 else:
231 low = midpoint
232 return (low + high) / 2.0
235def success_rate_ci(
236 successes: Sequence[bool],
237 *,
238 confidence: float = 0.95,
239) -> SuccessRateInterval:
240 """Computes an exact Clopper-Pearson confidence interval for a success rate.
242 Unlike a percentile bootstrap, which collapses to ``[0, 0]`` on an all-failure sample
243 regardless of size, the interval widens as the trial count shrinks (``[0.0, 0.459]`` for
244 zero successes out of six). Deterministic: no resampling and no seed.
246 Raises:
247 ValueError: If ``successes`` is empty.
248 """
249 if len(successes) == 0:
250 raise ValueError("successes must be a non-empty sequence")
251 n_trials = len(successes)
252 n_successes = sum(1 for success in successes if success)
253 tail = (1.0 - confidence) / 2.0
254 return SuccessRateInterval(
255 point_estimate=n_successes / n_trials,
256 ci_low=_clopper_pearson_bound(n_successes, n_trials, tail, upper=False),
257 ci_high=_clopper_pearson_bound(n_successes, n_trials, tail, upper=True),
258 n_trials=n_trials,
259 n_successes=n_successes,
260 confidence=confidence,
261 )