fix(statistics): compute the binomial CDF in log space so it does not underflow

tests/unit/utilities/test_statistics.py

@@ -2,6 +2,7 @@
 
 import itertools
 import math
+from fractions import Fraction
 
 import pytest
 import torch
@@ -13,6 +14,7 @@
     standard_error,
     success_rate_ci,
 )
+from transformer_lens.utilities.statistics import _binomial_cdf
 
 
 def test_derive_generator_is_keyed_and_stable():
@@ -227,6 +229,42 @@ def test_success_rate_ci_narrows_as_the_trial_count_grows():
     assert large.ci_high < small.ci_high
 
 
+@pytest.mark.parametrize(
+    "k, n, p",
+    [(1000, 2000, 0.5), (190, 200, 0.97), (1995, 2000, 0.999), (5, 2000, 0.01)],
+)
+def test_binomial_cdf_matches_exact_rational_sum_at_large_n(k: int, n: int, p: float):
+    # The old recurrence started from (1 - p) ** n, which underflows to 0.0 here and made
+    # every term, and the CDF, 0.0. (1000, 2000, 0.5) is ~0.509.
+    q = Fraction(p)
+    exact = sum(math.comb(n, i) * q**i * (1 - q) ** (n - i) for i in range(k + 1))
+    assert _binomial_cdf(k, n, p) == pytest.approx(float(exact), rel=1e-9)
+
+
+@pytest.mark.parametrize(
+    "k, n, low, high",
+    [
+        # Reference bounds from scipy.stats.beta.ppf, rounded to four places.
+        (199, 200, 0.9725, 0.9999),
+        (297, 300, 0.9711, 0.9979),
+        (1000, 2000, 0.4779, 0.5221),
+        (1990, 2000, 0.9908, 0.9976),
+    ],
+)
+def test_success_rate_ci_matches_reference_bounds_at_large_n(
+    k: int, n: int, low: float, high: float
+):
+    result = success_rate_ci([True] * k + [False] * (n - k))
+    assert result.ci_low == pytest.approx(low, abs=1e-4)
+    assert result.ci_high == pytest.approx(high, abs=1e-4)
+    assert result.ci_low <= result.point_estimate <= result.ci_high
+
+
+@pytest.mark.parametrize("n", [200, 1000])
+def test_success_rate_ci_all_success_lower_bound_at_large_n(n: int):
+    assert success_rate_ci([True] * n).ci_low == pytest.approx(0.025 ** (1 / n))
+
+
 def test_success_rate_ci_records_trial_and_success_counts():
     result = success_rate_ci([True, False, True])
     assert result.n_trials == 3

transformer_lens/utilities/statistics.py

@@ -179,8 +179,10 @@ class SuccessRateInterval:
 def _binomial_cdf(k: int, n: int, p: float) -> float:
     """Exact ``P(X <= k)`` for ``X ~ Binomial(n, p)``.
 
-    Accumulates the pmf by recurrence (``term_{i} = term_{i-1} * (n - i + 1) / i * p / (1 - p)``)
-    rather than forming ``comb(n, i)`` and ``p ** i`` separately, which overflow at large ``n``.
+    Works with each term's logarithm (``log comb(n, i) + i log p + (n - i) log(1 - p)``) and
+    sums ``exp(log_term - peak)`` before scaling back by ``exp(peak)``. Forming ``comb(n, i)``
+    and ``p ** i`` directly overflows at large ``n``, and starting a pmf recurrence from
+    ``(1 - p) ** n`` underflows to ``0.0`` once ``n log(1 - p) < -745``.
     """
     if k < 0:
         return 0.0
@@ -190,13 +192,14 @@ def _binomial_cdf(k: int, n: int, p: float) -> float:
         return 1.0
     if p >= 1.0:
         return 0.0
-    term = (1.0 - p) ** n
-    total = term
-    odds = p / (1.0 - p)
-    for i in range(1, k + 1):
-        term *= (n - i + 1) / i * odds
-        total += term
-    return total
+    log_p, log_q = math.log(p), math.log1p(-p)
+    log_n_factorial = math.lgamma(n + 1)
+    log_terms = [
+        log_n_factorial - math.lgamma(i + 1) - math.lgamma(n - i + 1) + i * log_p + (n - i) * log_q
+        for i in range(k + 1)
+    ]
+    peak = max(log_terms)
+    return min(1.0, math.exp(peak) * math.fsum(math.exp(term - peak) for term in log_terms))
 
 
 def _clopper_pearson_bound(n_successes: int, n_trials: int, tail: float, *, upper: bool) -> float: