Coverage for transformer_lens/utilities/statistics.py: 89%

1"""Model-free statistics over a per-example axis of an analysis result. 

2 

3Pure torch, no model access: standard errors, a vectorized percentile bootstrap 

4and a paired sign-flip permutation test, plus the seeded-generator convention 

5shared with sparse probing so every consumer derives independent, reproducible 

6streams the same way. 

7""" 

8 

9from __future__ import annotations 

10 

11import hashlib 

12import itertools 

13import math 

14from typing import Optional, Tuple 

15 

16import torch 

17 

18_ALTERNATIVES = ("two-sided", "greater", "less") 

19_STATISTICS = ("mean", "median") 

20 

21 

22def derive_generator(seed: int, *labels: object) -> torch.Generator: 

23 """CPU generator keyed on ``seed`` and any labels (metric name, arm, repeat...). 

24 

25 SHA-256 rather than ``hash`` keeps the derived seed stable across processes; 

26 the mask keeps it in torch's accepted int64 range. 

27 """ 

28 key = ":".join([str(seed), *(str(label) for label in labels)]).encode() 

29 derived = int.from_bytes(hashlib.sha256(key).digest()[:8], "little") & 0x7FFFFFFFFFFFFFFF 

30 return torch.Generator(device="cpu").manual_seed(derived) 

31 

32 

33def _to_last(x: torch.Tensor, dim: int) -> torch.Tensor: 

34 if not isinstance(x, torch.Tensor) or not x.is_floating_point(): 34 ↛ 35line 34 didn't jump to line 35 because the condition on line 34 was never true

35 raise TypeError("expected a floating-point tensor") 

36 if x.ndim == 0: 36 ↛ 37line 36 didn't jump to line 37 because the condition on line 36 was never true

37 raise ValueError("expected at least one dimension to resample over") 

38 return x.movedim(dim, -1) 

39 

40 

41def standard_error(x: torch.Tensor, *, dim: int = -1) -> torch.Tensor: 

42 """Standard error of the mean along ``dim`` (unbiased std over sqrt(n)); needs n >= 2.""" 

43 moved = _to_last(x, dim) 

44 n = moved.shape[-1] 

45 if n < 2: 

46 raise ValueError(f"standard_error needs at least 2 samples along dim {dim}, got {n}") 

47 return moved.std(dim=-1, unbiased=True) / math.sqrt(n) 

48 

49 

50def bootstrap_ci( 

51 x: torch.Tensor, 

52 *, 

53 dim: int = -1, 

54 statistic: str = "mean", 

55 confidence: float = 0.95, 

56 n_resamples: int = 1000, 

57 generator: Optional[torch.Generator] = None, 

58 chunk_size: int = 256, 

59) -> Tuple[torch.Tensor, torch.Tensor]: 

60 """Percentile bootstrap interval of ``statistic`` along ``dim``, vectorized over the rest. 

61 

62 Resamples with replacement using a CPU ``generator`` (``derive_generator(0)`` 

63 when omitted), ``chunk_size`` resamples at a time so memory stays at 

64 ``chunk_size × n`` per leading cell. Returns ``(low, high)`` with ``dim`` 

65 removed. 

66 

67 A percentile bootstrap cannot express uncertainty a sample does not contain: 

68 a constant vector resamples to itself and the interval collapses to a point 

69 whatever ``n`` is, and intervals on 4-8 samples are wide and skewed. For 

70 binary success rates prefer an exact interval 

71 (:func:`~transformer_lens.tools.analysis.jacobian_lens_causal_swap_benchmark.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).float() # [..., n_resamples] 

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) 

96 

97 

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. 

108 

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"``. Returns a tensor with ``dim`` removed. 

116 """ 

117 if alternative not in _ALTERNATIVES: 

118 raise ValueError(f"alternative must be one of {_ALTERNATIVES}, got {alternative!r}") 

119 if n_permutations < 1: 119 ↛ 120line 119 didn't jump to line 120 because the condition on line 119 was never true

120 raise ValueError("n_permutations must be >= 1") 

121 moved = _to_last(effects, dim).float() 

122 n = moved.shape[-1] 

123 if n < 1: 123 ↛ 124line 123 didn't jump to line 124 because the condition on line 123 was never true

124 raise ValueError("cannot permute an empty axis") 

125 observed = moved.mean(dim=-1, keepdim=True) # [..., 1] 

126 if n <= exact_max_n: 

127 signs = torch.tensor( 

128 list(itertools.product((1.0, -1.0), repeat=n)), dtype=moved.dtype, device=moved.device 

129 ) # [2^n, n] 

130 null = (moved.unsqueeze(-2) * signs).mean(dim=-1) # [..., 2^n] 

131 hits = _count_hits(null, observed, alternative) 

132 return hits / signs.shape[0] 

133 gen = generator if generator is not None else derive_generator(0) 

134 hits_total = torch.zeros(moved.shape[:-1], dtype=moved.dtype, device=moved.device) 

135 chunk = 256 

136 for start in range(0, n_permutations, chunk): 

137 count = min(chunk, n_permutations - start) 

138 signs = ( 

139 torch.randint(0, 2, (count, n), generator=gen).to(moved.device).to(moved.dtype) * 2 - 1 

140 ) 

141 null = (moved.unsqueeze(-2) * signs).mean(dim=-1) 

142 hits_total = hits_total + _count_hits(null, observed, alternative) 

143 return (hits_total + 1.0) / (n_permutations + 1.0) 

144 

145 

146def _count_hits(null: torch.Tensor, observed: torch.Tensor, alternative: str) -> torch.Tensor: 

147 if alternative == "two-sided": 

148 hit = null.abs() >= observed.abs() 

149 elif alternative == "greater": 

150 hit = null >= observed 

151 else: 

152 hit = null <= observed 

153 return hit.sum(dim=-1).to(null.dtype)