Avoid the cancellation of lgamma, lbeta and digamma differences at large shape parameters - #3434

Fixes #3153. Fixes #1578. Addresses #2508 partially.

This PR improves beta_binomial_lpmf, beta_binomial_cdf, beta_binomial_lcdf, beta_binomial_lccdf, beta_neg_binomial_lpmf, beta_neg_binomial_cdf, beta_neg_binomial_lcdf, beta_neg_binomial_lccdf, neg_binomial_lpmf, neg_binomial_2_lpmf, neg_binomial_2_log_lpmf, neg_binomial_2_log_glm_lpmf, dirichlet_multinomial_lpmf, lkj_corr_lpdf, lkj_corr_cholesky_lpdf, lkj_cov_lpdf, student_t_lpdf, multi_student_t_lpdf, multi_student_t_cholesky_lpdf, yule_simon_lpmf, yule_simon_cdf, yule_simon_lcdf, yule_simon_lccdf, and the derivatives of lbeta.

Summary

Several distributions form a difference f(x + k) - f(x), with f one of lgamma, lbeta or digamma, a shape, concentration, dispersion or degrees-of-freedom parameter x, and a count or offset k. When x is large, the two terms are large and nearly equal, and the difference keeps no correct digits. For the values the terms are of the size of x, so the error is about eps * x. For the gradients the terms are about log(x), and a model that puts the parameter on the log scale multiplies the error by x.

In the model of #3153, beta_binomial_lpmf(57 | 117, s, s) with s = exp(lc) / 2, develop returns a log probability 81 too high from lc = 42 on, and its gradient in lc is 0 or noise from lc = 24 on. Warmup then finds a flat plateau that it cannot leave. On this branch the value is within 2.3e-14 of the exact value and the gradient within 3.7e-14 for lc from 0 to 44.

From lc = 34 on, the exact values are below the rounding error of any form (about 2e-14): 57 log p + 60 log q is about -81 there.

#2508 reports the same defect: beta_binomial_lpmf(500 | 1000, s, s) for s = 0.5e10, 0.5e13, 0.5e19, where develop is wrong by 5.3e-8, 2.8e-4 and 693 (it returns +689.47 at the largest). This PR fixes all three; they are now rows of the beta_binomial_lpmf test. This PR does not close #2508: the comment there about beta_binomial_lcdf(3 | 10, 298.3, 146.75) returning NaN (exact -3.93335) has a different cause. hypergeometric_3F2 does not stop at the zero term of the terminating series 3F2({1, 302.3, -6}; {5, -151.75}; 1) (exact 17.63), and later 0 * inf gives NaN. A separate PR will repair that and close #2508.

#1578 reports wrong derivatives of neg_binomial_lpmf for large alpha. Its example (n = 0, alpha = 1e10 + 1, beta = 0.8) is already correct on develop, because the Poisson approximation it describes is no longer used. This PR repairs the loss that remains in d/dalpha = digamma(alpha + n) - digamma(alpha) + log(beta / (1 + beta)). It matters when the mean alpha / beta stays fixed: then the error of alpha * d/dalpha on develop reaches 3.0 at alpha = 1e15, and on this branch it is 1e-14 (rows of the new large-shape test).

The same root cause is in more functions. This PR repairs:

The shape is 1e15 for the first nine rows and up to 1e15 for the last four (nu decades for student_t_lpdf). "Log-shape gradient" is x * d/dx, the gradient a sampler sees when the model's parameter is log(x). All references are mpmath at 80 to 140 digits, certified by two independent routes (value by loggamma, gradients by closed-form digamma differences and by mp.diff; the cdfs as sums of the pmf).

propto. In some of the more stable forms, a constant that propto drops cancels against terms that depend on a parameter inside one expression, so the code cannot skip it (for example, lbeta(n, alpha + 1) contains lgamma(n)). To keep the propto behavior as before, the code adds such constants back with the opposite sign: lgamma(n) in yule_simon_lpmf, and lgamma(y + 1), phi log(phi) (data phi) and y theta (data x, alpha and beta) in neg_binomial_2_log_glm_lpmf, also in OpenCL. For the GLM with data phi, this brings back a rounding error of about eps * phi log(phi) in the propto value. A probe of all 14 changed densities, with every combination of data and parameter arguments, matches the propto values of develop, and a new test checks full - propto against the dropped terms.

Most of the added source lines are in six files:

The remaining 23 distribution files change by 2 to 48 lines each.

What this PR adds

digamma_diff(x, d) = digamma(x + d) - digamma(x), a new function in prim/fun. It's public (but not exposed in the language) because it is used in 19 distribution files and in the derivatives of lbeta; log_rising_factorial can use it in a follow-up. Three paths: a non-autodiff integer d from 0 to 8 (a count) uses sum_{j < d} 1 / (x + j), which has only positive terms and gives exactly 1 / x for d = 1; for x < 10 and d >= 10 the plain difference, which loses only a few ulp there (the result is at least 0.72 and digamma(x) <= 2.25); otherwise both arguments are shifted up to at least 10 with the recurrence, and log1p(d / y) plus the asymptotic series is used, with the differences y^-2i - (y + d)^-2i formed as (u - v) h_i so that nothing cancels. Measured over 7 281 points, x from 1e-3 to 1e18, d from 1e-6 to 1e7: at most 4.2 ulp; the plain difference reaches 6.5e16 ulp. It is vectorized and works with all autodiff types through generic arithmetic; expect_ad tests included. An OpenCL device function with the same paths and a kernel-generator registration are included.

internal::log_rising_factorial_ratio, internal::log_beta_ratio and their helpers in prim/fun/log_rising_factorial_ratio.hpp. They write each lgamma through lgamma_stirling_diff, so that the terms linear in the arguments cancel exactly, and form the rest without cancellation (the x log1p(k / x) - k parts are split off for k < x, and the integer parts are added exactly). log_beta_ratio(alpha, beta, n, m) = lbeta(alpha + n, beta + m) - lbeta(alpha, beta) chooses between this form and the plain lbeta form by the size of their terms: the Stirling form has rounding errors of the size of n + m, the lbeta form of the size of the smaller shape. Without the choice, N = 1e6 points with one small shape lost up to 4e-11 where develop was correct.

The rev and fwd derivatives of lbeta. The partial in an argument of at least 10 is now -digamma_diff(a, b); below 10 it stays digamma(a) - digamma(a + b), with digamma(a + b) shared between the two partials as before. lbeta itself is accurate at large arguments, but its derivative had the same cancellation, so every model that calls lbeta with a large parameter, and multi_student_t_lpdf, which gets its nu gradient by autodiff through lbeta, had inaccurate gradients. Below 10 the plain difference has an absolute error of a few eps, which stays at the rounding level on the log scale too.

Changes per function

• beta_binomial_lpmf: value by log_beta_ratio, partials by digamma_diff.
• beta_neg_binomial_lpmf: value as -log(n) - lbeta(n, beta) + log((r)_n / (r + alpha + beta)_n) + log((alpha)_r / (alpha + beta)_r), partials by digamma_diff.
• neg_binomial_lpmf: alpha partial by digamma_diff; the beta partial term alpha / beta - alpha / (1 + beta) is now alpha / beta / (1 + beta).
• neg_binomial_2_lpmf, neg_binomial_2_log_lpmf: phi partial by digamma_diff. Values unchanged (bitwise).
• neg_binomial_2_log_glm_lpmf: value per instance as lchoose(y + phi - 1, y) - y log1p_exp(log(phi) - theta) - phi log1p_exp(theta - log(phi)), the form of neg_binomial_2_log_lpmf, instead of phi log(phi) - lgamma(phi) + lgamma(y + phi) - (y + phi) log(mu + phi) + y theta. The phi partial no longer forms 1 - (y + phi) / (mu + phi) and log(phi) - logsumexp(theta, log(phi)).
• dirichlet_multinomial_lpmf: partials by digamma_diff.
• lkj_corr_lpdf (and lkj_corr_cholesky_lpdf, lkj_cov_lpdf through do_lkj_constant): the constant is a sum of lgamma(x + k/2) - lgamma(x); each is now lgamma(k/2) - lbeta(k/2, x), and the eta partial is set explicitly with digamma_diff.
• yule_simon_{cdf,lcdf,lccdf}: lgamma(alpha + 1) + lgamma(n + 1) - lgamma(n + alpha + 1) is now log(n) + lbeta(n, alpha + 1); partial by digamma_diff.
• yule_simon_lpmf: value log(alpha) + lbeta(n, alpha + 1); partial 1 / alpha - digamma_diff(alpha + 1, n).
• beta_binomial_{cdf,lcdf,lccdf}: the leading lbeta(nu, mu) - lbeta(alpha, beta) by log_beta_ratio, the digamma differences by digamma_diff.
• beta_neg_binomial_{cdf,lcdf,lccdf}: the leading term is the log pmf at n + 1, now computed as in beta_neg_binomial_lpmf; partials by digamma_diff.
• student_t_lpdf: for nu < 40 unchanged; for nu >= 40 the terms lgamma((nu + 1) / 2) - lgamma(nu / 2) - log(nu) / 2 and digamma((nu + 1) / 2) - digamma(nu / 2) come from the asymptotic series in h = nu / 2: -log(2) / 2 - 1/(8h) + 1/(192h^3) - 1/(640h^5) + 17/(14336h^7) - 31/(18432h^9) and its derivative (at most 0.9 and 2.4 eps against mpmath from h = 20 to 1e15). This is also cheaper than two lgamma or digamma calls.
• multi_student_t_lpdf, multi_student_t_cholesky_lpdf: the constant lgamma((nu + p) / 2) - lgamma(nu / 2) is now lgamma(p / 2) - lbeta(p / 2, nu / 2); the cholesky version sets the nu partial with digamma_diff(nu / 2, p / 2), the other one gets it through lbeta.
• OpenCL: beta_binomial_lpmf, neg_binomial_lpmf, neg_binomial_2_lpmf, neg_binomial_2_log_lpmf, neg_binomial_2_log_glm_lpmf (prim file and kernel) and student_t_lpdf compute the same formulas.

Related defects fixed on the way

1. lkj_corr_lpdf dropped d constant / d eta at eta == 1.0. The special case for eta == 1 returned the constant as a number without derivative. At eta = 1 with K = 3 and the matrix in the new test, develop returns d/deta = -0.165, the exact value is 1.2214. eta = exp(0) = 1 is exactly the value at the default initialization of a positive parameter.
2. beta_binomial_lcdf computed log(1 - C) with C the complementary series. Where the cdf is small this keeps no digits, and with the accurate C it can round to log of a value at or below 0. Where C > 1/2 the function now uses the mirror X <= n if and only if N - X > N - n - 1, with N - X distributed as beta_binomial(N, beta, alpha): the same series with the shapes swapped, without the complement. On develop this function returned NaN at 16 of 81 test points and was wrong by up to 61; it now has the accuracy of the series everywhere. For alpha = 1 exactly the mirrored series is 0 / 0 at one term, so that case keeps log(1 - C).
3. OpenCL beta_binomial_lpmf did not check 0 <= n <= N. The check was built, but its result was never computed, so the function never returned LOG_ZERO. Without propto the value was -inf by accident (through binomial_coefficient_log), but with a nonzero gradient; under propto it was finite: -8.216 for n = 6, N = 5 in the new test. The CPU version returns LOG_ZERO without a gradient; the OpenCL version now does too.
4. OpenCL neg_binomial_2_log_glm_lpmf with a scalar alpha and a vector phi failed. Whether the kernel sums the phi derivative was decided from alpha instead of phi, so it returned work-group sums where per-instance derivatives were needed, and the gradient assignment threw a size mismatch. The decision now depends on phi only, and the phi derivative is no longer computed when phi is data. A new test covers both mixed cases (scalar alpha with vector phi, and vector alpha with scalar phi).

Limits that remain, by a different cause

• beta_binomial_{cdf,lcdf,lccdf} and beta_neg_binomial_{cdf,lcdf,lccdf} are limited by the tolerance of their 3F2 series, about 1e-8 relative. They throw at alpha = beta = 1 for some N, and the beta_neg_binomial ones throw from shape 1e4 and can take seconds per call. beta_neg_binomial_lcdf still forms a complement and is wrong where its cdf is small. None of this is changed here.
• The log-shape gradients have a floor that no implementation of these functions can remove: the function returns d/dalpha and d/dbeta, and autodiff forms alpha d/dalpha + beta d/dbeta outside it. For tied shapes in numerical instability in beta_binomial_lpmf at large shape parameters #3153 the two products are -1.5 and +1.5 and their sum is 27 / exp(lc). Each partial is accurate to a few eps of its own terms; the sum is then accurate to a few eps of 1.5, not of its own size.
• log_rising_factorial is still lgamma(x + n) - lgamma(x); no Stan Math function calls it. A follow-up can use the new helpers there.

Tests

• test/unit/math/prim/fun/digamma_diff_test.cpp, test/unit/math/mix/fun/digamma_diff_test.cpp: special cases, exact finite sums, 17 mpmath values, the count path and the plain path (digamma_diff(x, 1) == 1 / x exactly from x = 1e-300 to 1e300); expect_ad at 8 points.
• test/unit/math/prim/prob/beta_binomial_test.cpp: values at the reporter's points of numerical instability in beta_binomial_lpmf at large shape parameters #3153 and beta_binomial_lpmf is unstable for large shape arguments #2508 and at large unequal shapes; the difference to the binomial limit at lc = 42.
• test/unit/math/rev/prob/beta_binomial_lpmf_test.cpp (new): log-shape gradients at 8 points; the lc gradient of the reporter's model for lc = 16 ... 32. From lc = 34 on develop returns exactly 0, which is within 1e-13 of the exact gradient, so a test there cannot see the defect.
• test/unit/math/rev/prob/large_shapes_test.cpp (new): 82 rows over the other functions (including student_t_lpdf, both multi_student_t functions, yule_simon_lpmf and lbeta), each with points where develop was wrong by more than 100 times the tolerance and one point at a small shape where develop was right; the lkj_corr gradient at eta = 1; and full - propto against the dropped terms for student_t_lpdf, yule_simon_lpmf, beta_neg_binomial_lpmf and three cases of the GLM.
• OpenCL: one absolute-reference test per changed distribution in test/unit/math/opencl/rev/ (six files), because the existing differential tests compare OpenCL with the CPU and cannot see an error both share; a device-function test for digamma_diff; a test that beta_binomial_lpmf gives LOG_ZERO without a gradient for n outside 0...N, also under propto; a test of neg_binomial_2_log_glm_lpmf with a scalar alpha and a vector phi and the reverse.

The tests find the defect. With the 23 develop headers swapped in (21 prim/prob files and the rev/fwd lbeta), all 4 new beta_binomial cases and both large-shape cases fail; the 3 old cases and the propto case pass (the propto case states develop's rule). With the branch headers restored, everything passes again. OpenCL: with the develop OpenCL files every new absolute test fails (6 files), and so do the two new tests of the related defects 3 and 4. Every old differential test passes, including the propto comparisons of the GLM and student_t_lpdf, except opencl_matches_cpu_big of the GLM: it runs after the exception of defect 4 in the same binary, which leaves the autodiff stack dirty.

Full suites. All unit tests of the touched functions pass (51 files: prim, rev, fwd, mix, including lbeta, digamma_diff, student_t, multi_student_t and lkj_cov), and so do the generated test/prob tests of beta_binomial, beta_neg_binomial, neg_binomial, neg_binomial_2, yule_simon and student_t (271 306 cases). The OpenCL suites of the six distributions, digamma_diff, lbeta and the kernel generator pass on a V100 (126 cases).

Side effects

• Speed. student_t_lpdf executes the same operations as develop for nu < 40 (plus one comparison: +2.4 % instructions for value and gradient) and fewer from nu = 40 on (-52 % instructions for value and gradient at nu >= 100; the scalar value takes 0.35 times as long there). Gradients that use digamma_diff, at shapes below 10, ratio to develop (median of 5 rounds, Xeon E5-2680 v3):

Release notes

Improved stability of beta_binomial_lpmf, beta_binomial_cdf, beta_binomial_lcdf, beta_binomial_lccdf, beta_neg_binomial_lpmf, beta_neg_binomial_cdf, beta_neg_binomial_lcdf, beta_neg_binomial_lccdf, neg_binomial_lpmf, neg_binomial_2_lpmf, neg_binomial_2_log_lpmf, neg_binomial_2_log_glm_lpmf, dirichlet_multinomial_lpmf, lkj_corr_lpdf, lkj_corr_cholesky_lpdf, lkj_cov_lpdf, student_t_lpdf, multi_student_t_lpdf, multi_student_t_cholesky_lpdf, yule_simon_lpmf, yule_simon_cdf, yule_simon_lcdf, yule_simon_lccdf, and the derivatives of lbeta.

Checklist

• Copyright holder: Aalto University

  The copyright holder is typically you or your assignee, such as a university or company. By submitting this pull request, the copyright holder is agreeing to the license the submitted work under the following licenses:
  - Code: BSD 3-clause (https://opensource.org/licenses/BSD-3-Clause)
  - Documentation: CC-BY 4.0 (https://creativecommons.org/licenses/by/4.0/)

• the basic tests are passing

  • unit tests pass (to run, use: ./runTests.py test/unit)
  • header checks pass, (make test-headers)
  • dependencies checks pass, (make test-math-dependencies)
  • docs build, (make doxygen)
  • code passes the built in C++ standards checks (make cpplint)

• the code is written in idiomatic C++ and changes are documented in the doxygen

• the new changes are tested