From 4aadeced874828a1d6ed74b24bb65b2c45a4066f Mon Sep 17 00:00:00 2001 From: berce <75228263+bercedev@users.noreply.github.com> Date: Wed, 21 Oct 2026 18:48:34 +0300 Subject: [PATCH] fix(statistics): compute the binomial CDF in log space so it does not underflow _binomial_cdf started its pmf recurrence from (1 - p) ** n, which is 0.0 once n log(1 - p) < -745, so the CDF came out 0.0 and success_rate_ci returned intervals that could exclude the observed rate (from n = 200 at high rates, n ~ 980 at 50%). The terms are now summed from their logarithms, rescaled by the largest. --- tests/unit/utilities/test_statistics.py | 38 ++++++++++++++++++++++++ transformer_lens/utilities/statistics.py | 21 +++++++------ 2 files changed, 50 insertions(+), 9 deletions(-) diff --git a/tests/unit/utilities/test_statistics.py b/tests/unit/utilities/test_statistics.py index 1bb453c66..c6df444b0 100644 --- a/tests/unit/utilities/test_statistics.py +++ b/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 diff --git a/transformer_lens/utilities/statistics.py b/transformer_lens/utilities/statistics.py index 0347e64b1..18886036d 100644 --- a/transformer_lens/utilities/statistics.py +++ b/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: