fix integrate_1d_gauss_kronrod gradient speed - #3393
Open
avehtari wants to merge 2 commits into
Open
Conversation
avehtari
marked this pull request as draft
September 6, 2026 18:17
avehtari
marked this pull request as ready for review
September 6, 2026 19:35
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Making this PR has been assisted by Claude. I have checked and edited all new code, comments, and PR text
Summary
The gradient computation for integrate_1d_gauss_kronrod calls the Boost adaptive Gauss Kronrod again, but Boost computes each parameter's gradient integral with the same Boost adaptive Gauss–Kronrod call. Boost's bisection criterion is relative to the leaf estimate, so a gradient component that is tiny relative to
fand dominated by round-off (e.g. the score of a saturated observation deep in a likelihood tail) never converges and bisects tomax_depth(2¹⁵ × 21 evaluations) on every call, although its contribution to∂ log Iis negligible. This can be fixed by shifting the integral integrand asg_i(x) = ∂f/∂θ_i (x) + c f(x)
and returning
∫ g_i − c I, whereIis the value integral already computed and c != 0 is a constantsuch that
∂f/∂θ_i ≈ −c·fdoes not occur for common integrands;c = 1is a bad choice because log-density gradients of logit-type likelihoods saturate at ±1 in the tails, which turns the shifted integrand back into noise. The inverse golden ratio is far from every integer and is not a structural constant of any link.As this is needed only for integrate_1d_gauss_kronrod, integrate_1d_adjoint gets an option whether the shift is used. The plain
integrate_1dandintegrate_1d_double_exponentialare unchanged: their L1-relative termination handles ordinary gradient integrands, and enabling the shift there costs 7–32 % more gradient evaluations on 10 of 35 suite integrands with no gain.Tests
Measured with the fix (same model, same point):
integrate_1dintegrate_1d49× faster gradients, results unchanged to 1e-11 (patched vs unpatched
GK: 1.0e-11). Per subject the counter shows 668 850 → 2 478 gradient
evaluations. The residual 2.75× is GK's per-integral cost on this
integrand (value-only GQ is 2.1× DE), not a gradient effect.
Effect on every integrand in the existing GK reverse-mode test suite
Integrand evaluations (value alone; all gradient components together)
and wall for value + gradients (mean of 200,
-O3), unpatched →patched. Same arguments and tolerances as the tests.
#var= numberof
varinputs (parameters + var endpoints).f1exp(x)+θ, [0.2, 0.7]f2exp(θ cos 2πx)+θ, var endpointsf3exp(x)+θ₀^2.5+2θ₁³+2θ₂, var endpointsf3, params onlyf101/(1+x⁴/1.7), (−∞, ∞), no paramsf11beta(5, 3), [0, 1]f11beta(3, 5), [0, 1]f12N(5.7, 1), (−∞, ∞)f13x+θ₀+θ₁, endpoints = paramsUnpatched, the noisy-gradient case does not merely hit
max_depth: itthen fails Stan's post-check (
error estimate 2.4e-13 exceeds max(rel_tol × L1, abs_tol)withL1 ≈ 0) and throws, i.e. in amodel block every proposal is rejected. Elsewhere the change is
neutral or a mild gain (heavy-tailed pdfs 15–35% fewer gradient
evaluations); the one regression is
f11(beta kernel: the gradientintegrands carry a
log xfactor that is easier thanfon its own),+14 % gradient evaluations at unchanged wall. All 17 pre-existing
tests pass at their tolerances.
A gradient component that is analytically zero but autodiffs to
round-off noise reproduces it without any likelihood machinery
(
cos² + sin² = 1). Added totest/unit/math/rev/functor/integrate_1d_gauss_kronrod_test.cppas
GradientShift_noisy_gradient(evaluation-count bound + accuracy;fails on the unpatched code, passes on the patched) together with
GradientShift_guards(exactly-zero integral → shift disabled; negativeintegral → shift active), and the existing 17 tests pass unchanged.
Side Effects
f. `f = spike(x)evaluations; shifted it re-resolves the spike, 189. Bound: gradient
cost ≤ (p + 1) × value cost, against 2^15 × 21 per component before.
inherits the value integral's absolute error (
≲ tol × L1(f)), so acomponent with
|∂I/∂θ| ≲ tol × L1(f)loses relative precision.Measured (
f = spike + ε θ x, spike width 0.01,tol = 1e-6):unshifted exact; shifted absolute error 3e-19 … 3e-18, i.e. relative
6e-15 at
ε = 1e-3down to 5e-10 atε = 1e-9— round-off level inpractice, because leaves that pass the value's test are usually far
below
tol. For HMC this is the right accuracy: value and gradiententer the Hamiltonian at the same absolute scale.
Release notes
integrate_1d_gauss_kronrod gradient computation was modified to have much improved worst performance timing
Checklist
Aki Vehtari
the basic tests are passing
./runTests.py test/unit)make test-headers)make test-math-dependencies)make doxygen)make cpplint)the code is written in idiomatic C++ and changes are documented in the doxygen
the new changes are tested