Accurate forward-mode tanh derivative and correlation transforms

[avehtari]

Fixes #3218.

@jachymb pointed out the fwd tanh issue in the logistic distribution PR #3398. I checked that in addition to logistic_lpdf, tanh is used by cholesky_factor_corr and corr_matrix via corr_constrain (no other Stan Math function calls tanh. User code that calls tanh gets the fix too).

This PR repairs the cancellations in the forward-mode derivative of tanh and in the transforms of cholesky_factor_corr and corr_matrix (values and log Jacobians). It fixes the initialization failure of cholesky_factor_corr[K] for K >= 40 reported in #3218, and adds the function log_cosh. The fix for cholesky_factor_corr and corr_matrix has to be made at the same time as the fix for tanh: their old forward-mode Jacobian gradient was correct only because the cancellation errors of tanh and of the Jacobian divided out. #3398 (logistic) is based on this PR and gets the tanh fix without further changes.

This PR was assisted by Claude.

Summary

Forward-mode tanh. stan/math/fwd/fun/tanh.hpp computes the derivative as 1 - u * u with u = tanh(x). The difference cancels: against mpmath the error is 5.7e3 ulp for |x| from 1 to 5, and the derivative is exactly 0 from |x| = 19.06 (for example, the second derivative of tanh(20) in fvar<fvar<double>> is 0; the exact value is -3.4e-17). This PR keeps 1 - u * u for |x| < 0.5, where it has an error of at most 1.4 ulp, and uses 4 e / (1 + e)^2 with e = exp(-2 |x|) above. The result has an error of at most 3.8 ulp for |x| up to 350, and the higher derivatives are correct down to underflow in fvar<fvar<double>> and fvar<var>. Two other forms were rejected:

• 1 / cosh(x)^2, which the reverse-mode tanh uses, gives NaN second derivatives from |x| = 360 in fvar<fvar<double>> and from |x| = 720 in fvar<var>: the quotient rule evaluates 0 * inf once cosh(x)^2 overflows.
• 4 e / (1 + e)^2 for all x, with fabs(x): the derivative of fabs at exactly 0 is 0 in Stan, so the third derivative of tanh at 0 was 0 instead of -2 in fvar<fvar<var>>.

The correlation transforms. Both transforms form partial correlations z = tanh(y) and then 1 - z^2 and 1 - sum_sqs, in the values (L(i, j) = z sqrt(1 - sum_sqs), L(i, i) = sqrt(1 - sum_sqs)) and in the log Jacobian (log1m(square(z)), 0.5 * log1m(sum_sqs), and the read_corr_L term (K - k - 1) * log1m(square(z))). These differences cancel in two cases:

• one partial correlation close to ±1: for one element, the old Jacobian term has an absolute error of 1.1e-8 at |y| = 10, 1.7e-4 at 15, 0.044 at 18, and is -inf from 19.1 (from 18 with var_value, whose tanh is Eigen's vectorized tanh);
• a long row: 1 - sum_sqs is the product of all 1 - z^2 in the row, so it cancels also with moderate y. With Stan's default inits (y uniform in [-2, 2]), a diagonal element of L rounds to 0 in 6 % of the inits at K = 30, 98 % at K = 40 and all at K = 64; lkj_corr_cholesky_lpdf rejects such an L, so initialization fails (cholesky_factor_corr[n] reliably fails to initialize (from [-2,+2]) for "moderate" n #3218).

The repair. 1 - tanh(y)^2 = exp(-2 log(cosh(y))), so after j elements of a row, 1 - sum_sqs = exp(-2 a) with a the sum of log(cosh(y)) over these elements. This PR computes the Cholesky factor as

• L(i, j) = tanh(y_j) exp(-a) and L(i, i) = exp(-a) over the whole row,

without a subtraction (for corr_matrix_constrain the same in the column order of read_corr_L, then Omega = L L'), and the log Jacobian as -sum_n c_n log(cosh(y_n)) with integer coefficients: c = 2 for corr_constrain, c = i + 1 - j for element (i, j) of cholesky_corr_constrain, c = K - j for column j of corr_matrix_constrain (equation 11 of Lewandowski, Kurowicka and Joe, 2009). A test checks lp against log |det J| of the autodiff Jacobian of the map from y to the strictly lower triangle (agreement 1.8e-15 at K = 4). The var_value overloads have a new hand-written reverse pass; in row i, L(i, j) depends on its own y with derivative sech(y)^2 exp(-a) and every later element of the row on it through exp(-a) with derivative -tanh(y) times the element. Its gradient equals the Eigen::Matrix<var, ...> gradient to 2e-16 relative up to K = 128.

log_cosh. New function in prim, fwd and rev. For |x| < 1 it uses u = expm1(|x|), log1p(u^2 / (2 (1 + u))) and tanh(|x|) = u (u + 2) / (u (u + 2) + 2); otherwise t = expm1(-2 |x|), |x| + log1p(t / 2) and tanh(|x|) = -t / (2 + t). Against mpmath on 70 002 points from 1e-150 to 400: at most 2.4 ulp for the value and 1.7 ulp for the derivative.

file	+	−	added code / comment
prim/constraint/cholesky_corr_constrain.hpp	158	48	75 / 79
rev/constraint/cholesky_corr_constrain.hpp	70	98	50 / 19
prim/constraint/corr_matrix_constrain.hpp	43	6	27 / 15
rev/constraint/corr_matrix_constrain.hpp	17	2	17 / 0
prim/constraint/corr_constrain.hpp	8	6	3 / 5
prim/fun/log_cosh.hpp (new)	92	0	45 / 39
rev/fun/log_cosh.hpp (new)	55	0	35 / 15
fwd/fun/log_cosh.hpp (new)	27	0	16 / 7
fwd/fun/tanh.hpp	26	2	11 / 15

In addition prim/fun.hpp, fwd/fun.hpp and rev/fun.hpp get one include line each. Tests add 536 lines and change 1.

Results

Accuracy of the values (K = 4, y = (15, 12, -0.5, 9, 2, -14), against mpmath): every element of L within 1.0e-15 relative (develop: up to 3.8e-4 relative, from sqrt(1 - z^2) of a rounded z); the correlations of corr_matrix_constrain within 1.1e-16 absolute (develop: 2e-10). At K = 40 and 128 with default-like inits, the rows of L have unit norm within 7e-16 and the smallest diagonal elements are 8e-11 and 3e-33, where develop gives 0.

Prior-only check with HMC in CmdStan 2.40.0: L ~ lkj_corr_cholesky(eta) (cholesky_factor_corr[K]) and Omega ~ lkj_corr(eta) (corr_matrix[K]). Under LKJ(eta), (rho + 1) / 2 ~ Beta(a, a) with a = eta - 1 + K / 2 for every correlation rho, so E[rho^2] = 1 / (2a + 1) is known exactly; z = (E[rho^2] - exact) / MCSE per seed.

Initialization of cholesky_factor_corr[K] (eta = 1 and 2, 10 seeds, 4 chains, 40 chains per row; chains that failed / mean number of rejected initial values):

K	develop	Jacobian fix only	value fix only	this PR (both)
20	0 / 0	0 / 0	0 / 0	0 / 0
30	0 / 0.1	0 / 0.1	0 / 0	0 / 0
40	7 / 53	7 / 53	0 / 6.3	0 / 0
64, 128	40 / 100	40 / 100	40 / 100	0 / 0

"Value fix only" failures at K = 64 and 128 come from develop's Jacobian terms, which are -inf there. "Jacobian fix only" failures come from L(i, i) rounding to 0, which lkj_corr_cholesky_lpdf rejects. Both fixes are needed. Chains that start, sample correctly in all versions.

Sampling, eta = 0.1 to 2, K = 2, 5, 10, 20 seeds per setting (the same for all versions), 4 chains of 1000 warmup and 4000 draws; totals over K and seeds (240 000 draws per row and band): divergences after warmup / largest per seed / seeds with |z| > 3:

type	eta	develop	Jacobian fix only	value fix only	this PR (both)
cholesky_factor_corr	0.1 to 0.25	75 266 / 4 548 / 4	89 600 / 7 699 / 11	56 562 / 1 149 / 4	0 / 0 / 0
cholesky_factor_corr	0.3, 0.4	84 / 5 / 0	107 / 5 / 0	93 / 11 / 1	0 / 0 / 0
cholesky_factor_corr	0.5, 1, 2	0 / 0 / 1	0 / 0 / 1	0 / 0 / 1	0 / 0 / 0
corr_matrix	0.1 to 0.25	292 882 / 10 035 / 48	309 744 / 11 533 / 50	-	296 642 / 11 491 / 53
corr_matrix	0.3, 0.4	2 973 / 2 768 / 1	153 / 6 / 0	-	183 / 8 / 0
corr_matrix	0.5, 1, 2	0 / 0 / 0	0 / 0 / 2	-	1 / 1 / 3

With this PR the share of cholesky_factor_corr draws at |y| > 15 matches its prior probability (0.0504 to 0.0508 against 0.0505 at eta = 0.1, 0.0112 to 0.0117 against 0.0115 at eta = 0.15); a prototype of the same transform in the Stan language, which keeps y, reached |y| up to 65. With develop the chains stop near |y| = 18.7. For corr_matrix this PR changes nothing measurable: lkj_corr_lpdf works with Omega itself, and its positive-definiteness check rejects Omega once the smallest eigenvalue is below what double precision resolves (about 1e-20 at a K = 40 init), whatever the accuracy of L. The MCSE is from batch means and is not reliable in runs with many divergences.

For usual priors with eta >= 1, there is no practical difference in HMC behavior, but the fix for initialization helps.

Time per call, on one whole node of AMD EPYC 9965 (g++ 11.5, -O3), median of 5 interleaved rounds (the spread between rounds of one binary is at most 9 %, mostly below 3 %). Reverse-mode kernels compute the gradient of lp + sum(L) (of lp + lkj_corr_cholesky_lpdf(L, 2) in the lkj row); the forward-mode kernels compute one directional derivative. The middle column is this PR without the value fix.

kernel	develop	Jacobian fix only	this PR	this PR against develop
corr_constrain with lp, var, 45 values	1.53 us	1.69 us	1.82 us	+19 %
cholesky_corr_constrain, var, K = 5	0.89 us	0.80 us	1.00 us	+13 %
cholesky_corr_constrain, var, K = 10	3.84 us	3.35 us	4.07 us	+6 %
cholesky_corr_constrain, var, K = 25	28.3 us	24.3 us	27.6 us	-3 %
cholesky_corr_constrain, var_value, K = 10	1.32 us	1.41 us	1.17 us	-11 %
cholesky_corr_constrain, var_value, K = 25	9.37 us	9.73 us	7.43 us	-21 %
corr_matrix_constrain, var, K = 5	1.36 us	1.22 us	1.44 us	+6 %
corr_matrix_constrain, var, K = 10	4.80 us	4.22 us	5.29 us	+10 %
corr_matrix_constrain, var_value, K = 10	2.43 us	2.56 us	2.12 us	-13 %
cholesky_corr_constrain + lkj_corr_cholesky_lpdf, var, K = 10	4.17 us	3.73 us	4.39 us	+5 %
cholesky_corr_constrain, fvar<double>, K = 10	1.16 us	1.73 us	2.76 us	+138 %
corr_matrix_constrain, fvar<double>, K = 10	2.20 us	2.69 us	4.13 us	+88 %

Reasons for the changes:

• Eigen::Matrix<var, ...> arguments use the generic prim code. The Jacobian fix alone makes these transforms 9 % to 14 % faster, because the Jacobian becomes one log_cosh node per element and a dot_product, instead of the log1m(square(z)) and 0.5 * log1m(sum_sqs) nodes. The value fix computes tanh(y), exp(-a) and log(cosh(y)) per element instead of tanh(y) and sqrt(1 - sum_sqs), which costs more than the Jacobian fix saves at small K; the net change is -3 % to +13 %.
• var_value arguments use the new hand-written reverse pass: the forward pass needs one expm1 or exp per element (tanh(y) and sech(y) together, and exp(-a) as a running product of sech(y)), and the reverse pass has no transcendental function. Develop evaluated a vectorized tanh with its own reverse pass and sqrt(1 - sum_sqs). The result is 11 % to 21 % faster than develop.
• corr_constrain alone (no value change) is slower because log_cosh costs an expm1 and a log1p per element, against log1m(square(tanh(y))) from the already computed tanh.
• fvar<double> arguments are slower because the forward-mode tanh derivative needs an exp above |y| = 0.5, and the value fix adds exp(-a) and log_cosh (which also evaluates tanh) per element. At K = 10 the times are below 5 us, so this is unlikely to dominate in a model; forward mode is used in nested autodiff only.

Limits that remain

• corr_matrix at large K: as above, the initialization and small-eta failures come from the positive-definiteness check of Omega in lkj_corr_lpdf. cholesky_factor_corr with lkj_corr_cholesky is the type to use for large K.
• read_corr_L(CPCs, K, lp) and read_corr_matrix keep their formulas, because they receive only the partial correlations; corr_matrix_constrain no longer calls them. cov_matrix_constrain_lkj, which stanc3 does not emit, still uses them.
• The inverse transforms cholesky_corr_free and corr_matrix_free are unchanged.

Tests

• test/unit/math/mix/fun/tanh_test.cpp: first and second derivatives at x = 3, ±20, 300, 1000 in fvar<double>, fvar<fvar<double>> and fvar<var>, and first to third derivatives in fvar<fvar<var>> at x = 0, 0.25, 0.4999, ±0.5, 3, -20, against mpmath.
• test/unit/math/mix/fun/log_cosh_test.cpp (new): expect_ad for scalars, containers and var_value; values and derivatives against mpmath on both sides of the branch, at 1e-10 and at 800.
• test/unit/math/mix/constraint/{corr_constrain,cholesky_corr_constrain,corr_matrix_constrain}_test.cpp: lp and its gradient at large y against mpmath in double, var, var_value and fvar<double>; lp against log |det J|; the values of L and Omega at large y against mpmath in all four modes; K = 40 with default-like inits (positive diagonal, unit rows, finite lp and lkj_corr_cholesky_lpdf, var_value gradient equal to the Eigen::Matrix<var, ...> gradient).
• test/unit/math/rev/prob/lkj_corr_test.cpp: ProbDistributionsLkjCorrCholesky_gradients1 compares finite differences with autodiff at an input with an exact 0, where the exact derivative is 0; autodiff now returns exactly 0 and the finite difference has rounding noise of 1.2e-13, so the comparison uses an absolute tolerance of 1e-10 instead of float equality.
• With the headers of develop, the new tanh tests and the large-argument tests fail; with the headers before the value repair, the three value tests (cholesky_corr_constrain_values_large_argument, cholesky_corr_constrain_long_rows, corr_matrix_constrain_values_large_argument) fail.
• The existing unit tests of tanh, the three transforms, cov_matrix_constrain_lkj, read_corr_L, read_corr_matrix, read_cov_L and lkj_corr (prim, fwd, rev, mix; 19 suites) pass.

Tested with develop 43bc5684a3, the libraries of the 2.40.0 release (Eigen 5.0.1, Boost 1.87.0), g++ 11.5.

Side Effects

lp__ and the values of cholesky_factor_corr and corr_matrix parameters change in the last digits for moderate unconstrained values. Models with cholesky_factor_corr[K] for K >= 40 now initialize with default inits.

Release notes

Improved the accuracy of the forward-mode derivative of tanh and of the cholesky_factor_corr and corr_matrix transforms. Added log_cosh.

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