Repository navigation
integrate_1d_adjoint repeats the whole quadrature once per differentiable argument #3444
Description
This issue was Claude assisted. The proposed options have been implemented and tested, so the ideas here are working.
Summary
integrate_1d_adjoint builds the reverse-mode result for all three one-dimensional quadrature functions: integrate_1d, integrate_1d_double_exponential and integrate_1d_gauss_kronrod. They differ in which fixes can be applied to them. For a call with P differentiable scalar arguments integrate_1d_adjoint runs P complete adaptive quadratures, and at every node of every one of them it performs a full reverse sweep that computes all P adjoints, but keeps only one and discards the rest.
On a real model (a beta-binomial marginal likelihood with 8 differentiable arguments from the Nabiximols case study) removing that redundancy is a 5–10× wall-time speedup, with gradients that agree with the current code to 2e-13 relative or better. In 34 public models that use integrate_1d in the log density, the median P is 3, and 24 of them call it once or more per observation. I have three working implementations, option A and two variants of option B, and would like to agree on an approach before opening a PR, because they differ in how much vendored Boost they touch, in compile time and in memory use.
Details on the motivating model:
- 3 of the 8 arguments are data values that the model packs into the
thetaarray instead ofx_r, so they are stored asvar(see below); without themPwould be 5. - In the Gauss-Kronrod version of the model, a 9th argument is also a
var: a log-scale shift computed from the parameters, whose gradient contribution cancels exactly. The benchmark passed the shift as a constant, so the measured speedup slightly understates the gain for that model. - With
integrate_1d, the same integral gains 6.9–7.9× when it is written on a finite interval and 4.9–8.5× on the unbounded interval (−∞, ∞) that the model uses. The gain on the unbounded interval needs the NaN rule at each evaluation, which both variants of option B apply.
Current implementation
stan/math/rev/functor/integrate_1d_adjoint.hpp, paraphrased:
auto integrate_grad = [&](auto&& target) -> double {
return integrator([&](const auto& x, const auto& xc) {
argument_nest.set_zero_all_adjoints();
nested_rev_autodiff gradient_nest;
var fx = f(x, xc, msgs, args_copy...);
fx.grad(); // computes EVERY argument adjoint
double gradient = target.adj(); // keeps exactly ONE of them
if (is_nan(gradient) && fx.val() == 0) gradient = 0.0;
return gradient;
});
};
// ... then once per differentiable scalar argument:
iter_tuple_nested([&](auto&& adj_leaf, auto&& var_leaf) {
assign_grad(adj_leaf, var_leaf); // -> a WHOLE new quadrature
}, args_adj, args_copy_filter);Two separate inefficiencies: the quadrature is repeated P times, and each node of each repetition pays a full reverse sweep whose results are mostly thrown away. There is also a separate all-double pass for the value, so P + 1 quadratures in total.
Note the is_nan(gradient) && fx.val() == 0 filter, it matters below.
How common is this, and how large is P
The cost of the current scheme grows with P and with the number of integrate_1d calls per gradient evaluation. I searched public Stan models that call integrate_1d where it affects the gradient.
Sources, searched on 2026-10-10:
- GitHub code search for
integrate_1din.stanfiles: 395 hits, 98 different files. 50 files reachintegrate_1dfrommodelortransformed parameters. After removing compiler test files and copies (one cosmology project has 284 copies in output archives), 24 different models remain, in 17 repositories. - Stan code inside
.R,.Rmdand.qmdfiles on GitHub: 7 models, 4 of them brms custom families. - The Stan User's Guide, 1 Stan Math issue, 1
epinowcast/primarycensoredpull request and 1 Stan Forums thread that show the model code: 4 models.
All these models use integrate_1d. I found no public model outside the Stan repositories and their test files that uses integrate_1d_double_exponential or integrate_1d_gauss_kronrod, which are new in Stan 2.40.
GitHub code search covers only default branches, so these counts are lower bounds. P is the number of var scalars in the argument pack. P does not include the integration limits, because a var limit costs only one integrand evaluation. One model has a data-dependent P (the number of parties); it is in the table but not in the counts. One more model does not type-check and is excluded. The motivating model is from the Nabiximols case study, although the specific Stan code with integral is not yet public.
P |
models |
|---|---|
| 1 | 6 |
| 2 | 7 |
| 3 | 10 |
| 4 | 4 |
| 5 | 4 |
| 6 | 1 |
| 7 | 1 |
| 8 | 1 |
The median P is 3. 21 of 34 models have P ≥ 3, and 11 have P ≥ 4. For comparison, the option B speedups measured below are 1.8–2.9× at P = 2 (11–12.5× in one case where today's code refines the gradients much further than the value), 4.3–5.2× at P = 3 and 4.9–10× at P = 8, for both variants.
Four observations:
- Most models call
integrate_1donce or more per observation. 24 of 34 models call it in a loop over observations, subjects, bins or HMM states, with 1 to 8 calls per item. 9 models call it once per gradient evaluation. In 1 model the call is inside an ODE right-hand side, so every right-hand-side evaluation of every ODE solve runsP + 1quadratures. Thus in most models the per-call cost is multiplied by the data size, which is likely to makeintegrate_1dto dominate the computation time. - The legacy signature adds constants to
P. Inintegrate_1d(f, a, b, theta, x_r, x_i),thetais onearray[] real. If any element is a parameter, stanc3 stores every element asvar(the generated C++ declaresstd::vector<local_scalar_t__>). 7 of 34 models put data or literals intotheta, for example{0, 1, mu, disc}and{k, mu_y[i], y[i], mu_w, sigma_w}. Each such element costs a full quadrature, and nothing uses its adjoint. - Elements that the integrand does not read also cost a quadrature. In 4 cosmology models the integrand reads 1 or 2 of the 3–4
thetaelements. The derivative integral of an unread element is identically zero, but the current code still runs an adaptive quadrature for it, with a reverse sweep at each node. - All 34 models use the legacy
integrate_1dsignature. The variadic signatures ofintegrate_1d_double_exponentialandintegrate_1d_gauss_kronrodlet the user pass data as separatedoublearguments, so item 2 can be avoided there; item 3 cannot. Option B removes the cost of items 2 and 3, because one quadrature gives all adjoints. My implementation opts in all three functions (see "How each function can be handled").
All 35 models
"Constants" are data or literals stored as var. "Unread" elements are in theta but not read by the integrand. "TP" is transformed parameters.
| Model | Block | P |
Constants or unread in P |
Calls per gradient |
|---|---|---|---|---|
Bayesian-Workflow worldcup_sqrt_discrete |
model | 3 | 1 constant (df) |
1–2 per game |
quarto_modeling_techniques fit_unknown_selection_uni_quad |
model | 2 | – | 1 |
quarto_modeling_techniques fit_unknown_both_uni_quad |
model | 4 | – | 1 |
Statisticats regression_interval |
model | 5 | – | 1 per observation |
Censored-Data-Tutorials censored-predictor |
model | 5 | 1 constant (y[i]) |
1 per censored observation |
joint-models-in-stan joint_model |
model | 3 | – | 1 per subject |
SeroUncertaintyPaper tar_agedep_gamma_bin_model |
model | 3 | – | 1 per age bin |
bivariate-beta bivariate-beta-model-v1 |
model (_lpdf) |
4 | – | 1 per observation |
The-Edinburgh-Malaria-Project goncalves1 |
model (4 functions) | 8 | – | 2, plus 1–2 per bin |
JDistlib 38-integrate-1d |
model | 1 | – | 1 |
msc-code LCDM (also LCDM_GW) |
TP | 2 | 1 unread (h) |
1 per event |
msc-code LCDM-SNIa |
TP | 1 | – | 1 per supernova |
msc-code fQ-LCDM-bg |
TP | 3 | 2 unread (h, alpha) |
1 per event and per bin |
forecasting-FQ-cosmology-with-SS FQ_GW_SNIa |
TP | 3 | 2 unread (h, M) |
1 per event and per bin |
forecasting-FQ-cosmology-with-SS Xi_0,omega_0_GW |
TP | 4 | 2 unread (h, Xi_0) |
1 per event |
testing-L-free-fQ-cosmology shift |
TP | 6 | – (var lower limit) |
1 |
testing-L-free-fQ-cosmology shift-LCDM |
TP | 5 | – (var lower limit) |
1 |
melding-multiple-phi submodel-two-phi-step |
TP | 3 | – (var upper limit) |
1 per patient |
melding-multiple-phi submodel-two-phi-step-indiv |
TP | 3 | – (var upper limit) |
1 |
melding-multiple-phi submodel-two-psi-step |
TP | 3 | 1 constant (long_beta[ii, 2]) |
1 per patient |
hmm_movement hmm_beta_divergence |
TP (function) | 3 | 1 constant (beta) |
1 per HMM state |
bplaw continuous_plaw_lognormal_mixture_varyingDelta |
TP | 5 | 1 constant (x_min) |
1 |
T_cell_dynamics_birth-death asm_da_rh_cd8 (4 variants) |
TP (ODE right-hand side) | 7 | 2 constants, 1 unread | 2–4 per right-hand-side evaluation |
election_forecast_france integral |
TP | D (parties) |
– | D |
gumbel-reanalysis gumbel24afc-stan.R (brms family) |
model (_lpmf) |
1 | – | 2 per observation |
gumbel-reanalysis gumbel24afctrial-stan.R (brms family) |
model (_lpmf) |
1 | – | 1 per observation |
gumbel-reanalysis uvsdtsdai-stan.R (brms family) |
model (_lpmf) |
4 | 2 constants (0, 1); var limits |
8 per observation |
gumbel-reanalysis uvsdtranknew-stan.R (brms family) |
model (_lpmf) |
2 | – | 1 per observation |
myblog truncation-in-stan |
model | 2 | – | 1 per observation |
minor_scripts test-Stan-integrate_1d.R, stan_code_2 |
model | 1 | – | 1 |
minor_scripts test-Stan-integrate_1d.R, stan_code_3c |
model | 2 | – | 1 per observation |
| Stan User's Guide, truncated normal | model | 2 | – (var lower limit) |
1 |
| primarycensored before #64 (#34) | model (_lpmf) |
2 (lognormal delay) | – | ≥ 1 per observation |
| stan-dev/math#2848 | TP (ODE in integrand) | 3 | – | 10 |
| Stan Forums, IRT marginal likelihood | model | 1 | – | 1 per person |
Prior art in this repo: solve_newton
stan/math/rev/functor/solve_newton.hpp
reverse_pass_callback([f, ret, arena_args_tuple, Jf_x_T_lu_ptr, msgs]() mutable {
Eigen::VectorXd eta = -Jf_x_T_lu_ptr->solve(ret.adj().eval());
nested_rev_autodiff rev;
auto x_nrad_ = eval(f(ret.val(), msgs, args...));
x_nrad_.adj() = eta; // seed a VECTOR adjoint
grad(); // ONE sweep gives every argument adjoint
});Two ideas there:
- Deferral. The adjoint work happens inside a
reverse_pass_callback, when the output adjoint is already known; only an LU factorisation is kept in the arena.integrate_1d_adjointcomputes its gradient integrals eagerly, so it cannot skip the work when the output adjoint is zero, and cannot use the adjoint's scale to inform tolerances. I tested deferral with a prototype outside the library: the adjoints are identical to the eager ones, and the gradient work is skipped when the output adjoint is 0. With option B, however, the forward pass then needs a separate value quadrature, so a used integral costs 2 quadratures instead of 1 and ran at 0.65–0.83× the speed of option B without deferral. The work is saved only for an integral that the target does not use. I therefore do not propose deferral. - Seed and contract. One sweep with a seeded vector adjoint, rather than a loop over components.
For a scalar-output integral (2) does not transfer directly — there is no output dimension to collapse, and a VJP of a scalar is the gradient.
Option A — batch the gradients into one quadrature
Make the integrand vector-valued: one quadrature carrying P components, one sweep per node, every adjoint harvested instead of one.
Boost's Gauss-Kronrod recursion is already generic in the integrand's return type, so this needs no change to the quadrature — only a small value type providing ±, unary -, *scalar on both sides, construction from 0, and an abs() returning the scalar type.
- ~110 lines for the value type.
- Gauss-Kronrod only, unless Boost's double-exponential code is also copied and changed.
tanh_sinhcomputesReal / Kin one place, which has no vector meaning, andexp_sinhcompares values ofKwith==(verified: neither compiles with the vector type). A Stan copy oftanh_sinhwith that line changed, plus==on the value type, would compile. The stopping rule of these rules would then use one norm over all components; a separate tolerance for each component would need changed copies of the stopping rules oftanh_sinh,exp_sinhandsinh_sinh. Without such copies,integrate_1dandintegrate_1d_double_exponentialkeep the per-component path. - Refinement is driven by the sup-norm over components rather than per component, so a component much smaller than the largest is held to the largest one's scale. A deliberate change in what is guaranteed.
- Still one reverse sweep per node.
Option B — one quadrature for the value and all adjoints
I = Σ wᵢ f(xᵢ, θ), so ∂I/∂θⱼ = Σ wᵢ ∂f(xᵢ, θ)/∂θⱼ with the same nodes and weights. One quadrature can therefore give the value and every adjoint, if its value type carries the derivatives along: P + 1 quadratures become 1. There are two variants, which differ in how the value type carries the derivatives:
- Taped: the value type wraps a
var. The quadrature arithmetic records the whole quadrature on one autodiff tape, and one reverse sweep at the end gives every adjoint. There are no per-node sweeps. - Joint: the value type is a value together with its gradient (
quadrature_dual, an inline array of up to 8 derivatives). At each node, one nested reverse sweep of the integrand gives allPderivatives at once: the sweep that the current code already does, but once instead ofPtimes. The quadrature arithmetic then carries value and gradient together. No autodiff graph outlives a node.
Common to both: Boost's rules are generic in the integrand's return type, but Boost writes Real L1 = abs(k), so the value type must supply a double-valued abs() (a bare var does not: abs(var) returns var) plus the arithmetic Boost needs. The value type also specializes boost::math::isfinite: tanh_sinh uses it to move away from non-finite integrand values at the extreme abscissas, and Boost's default returns true for class types. Both value types work with all four quadratures the family uses: gauss_kronrod, tanh_sinh, exp_sinh, sinh_sinh, with values bit-identical to the all-double path. Unlike option A's vector, both are scalars, so Real / K has a meaning.
Because abs() reports the magnitude of the value, refinement is driven by exactly the quantity the plain-double pass uses — the partition is unchanged, unlike option A, and it is the same in both variants. The refinement is not informed by the derivatives, so an integrand whose derivative is rougher than its value gets no extra refinement on that account (question 5).
Costs common to both variants:
integrate_gk,integrate_deandintegratemust return a deduced type instead ofdouble(mechanical; 1, 7 and 7 sites).- For the finite branch of
integrate_1d_double_exponentialandintegrate_1d, Boost's 2-argument (xc)tanh_sinh::integratehard-codesRealin its internal lambda and result, so it needs a 2-line vendored patch. The 1-argument overload is already generic. Gauss-Kronrod needs no Boost change at all. - The existing NaN rule must act at each evaluation — see below.
With option B, the derivative integrals no longer drive the refinement. The absolute tolerance that develop passes into Boost's Gauss-Kronrod recursion, a change to vendored Boost, is then not needed against the gradient blow-up it was made for. On 72 parameter points of the beta-binomial marginal, today's code without that change needed up to 2.9 million evaluations; the joint variant without it needed at most 588, with gradients within 1.2e-13 of the joint variant with it. However without the Boost code change, the value needs up to 1.8× the evaluations, because it is refined to the relative tolerance only, so the already merged Boost GK code change would be useful also with Option B.
The variants differ in the other costs (measurements in "Experiments"):
| taped | joint | |
|---|---|---|
| Lines of code (without blank and comment lines): value type + adjoint path | 72 + 83 | 137 + 61 |
| NaN rule at each evaluation | own reverse-pass loop over the tape (question 2) | applied to each node's gradient |
| Memory | grows with the number of evaluations, 0.3–1.7 KB each; about 0.7 GB in one quadrature at the default max_depth (question 4) |
flat |
Compile time of one integrate_1d call site (today 18.8 s) |
+0.6 s | +2.4 s |
Speedup over today, 28 cases in "Experiments" (P ≤ 8) |
1.8–11× | 2.3–12.5×; 0.99–1.34× the taped speed |
P > 8 |
one quadrature | one quadrature per block of 8 arguments; P = 40: 0.19–0.22× the taped speed, 4.9–7.2× today's |
The joint variant seems to be better. Its memory does not grow, it needs no reverse-pass loop of its own and no memory guard, and up to P = 8, which covers all 34 surveyed models, it is about as fast as the taped variant or faster. Its costs are 1.7 s more compile time per call site than the taped variant, and slower gradients for P > 8.
The NaN rule in option B
The current per-node test zeroes a gradient that is NaN where the integrand is zero. The test is necessary. The double-exponential rules evaluate at arguments so extreme that, for example, exp(-θ·x²) underflows to 0, while its derivative evaluates as inf · 0 = NaN. The same happens on a finite interval: an integrand that does not use xc can evaluate log(1 − x) = −inf at the nodes next to the endpoint 1. Gauss-Kronrod's infinite-domain transform stops near |x| ≈ 113 and does not reach the underflow regime. In the measurements the rule was never needed with Gauss-Kronrod.
Joint variant. Each node's gradient comes from its own nested sweep, as in the current code, so the rule applies directly: a NaN component where f == 0 is set to 0 before the gradient enters the quadrature. A finite derivative at a zero-valued node is kept. Evaluations that Boost discards never enter a sum, so their gradients do not matter: sinh_sinh evaluates f at ±max_value only to check that f goes to zero there, and tanh_sinh discards evaluations while it searches for usable endpoint abscissas. A gradient multiplied by a weight that is exactly 0 gives 0, also where the gradient is not finite.
Taped variant. grad() sweeps the whole tape and has no per-node step, so one such node makes every adjoint NaN. Zeroing a node's contribution whenever f(x) == 0 is not a valid substitute, because f can be zero with a nonzero derivative.
The taped variant therefore runs its own reverse-pass loop, a copy of the grad() loop, instead of calling grad():
- During the taped quadrature, record the tape range
[s0, s1)of every integrand evaluation, the vari of its value, and whether the value is exactly 0. - In the reverse-pass loop, handle each recorded range as one unit:
- If the adjoint of the value is 0, the quadrature did not use the evaluation, so skip the range. Boost evaluates and discards values:
sinh_sinhevaluates f at ±max_valueonly to check that f goes to zero there, andtanh_sinhdiscards evaluations while it searches for usable endpoint abscissas. Chaining such a range can give 0 · inf = NaN, also when the discarded value is NaN rather than 0. - If the value is exactly 0, chain the range. If that makes an argument adjoint NaN, restore the value that the adjoint had before the range.
- Otherwise, chain the range.
- If the adjoint of the value is 0, the quadrature did not use the evaluation, so skip the range. Boost evaluates and discards values:
The second case is the existing per-component rule ("a NaN gradient where f == 0 counts as 0"), applied per node and per argument. A finite derivative at a zero-valued node is kept. In a constructed case, θ0·exp(−θ1·x²) with θ0 = 0, every node is zero-valued but ∂f/∂θ0 is finite and nonzero; the adjoint ∂I/∂θ0 is kept.
The per-node loop reads ChainableStack::instance_->var_stack_ and nested_var_stack_sizes_ directly, as grad() does (question 2).
The simpler taped design calls grad() and relies on the whole-call fallback alone (next paragraph). That design is the backup if the reverse-pass loop is not acceptable. Its cost: every call with an underflowing node falls back, and in the measurements such calls ran at 0.80–0.96× today's speed.
Both variants. If an adjoint is still not finite (a non-finite derivative at a node with a nonzero value), the call discards the result and repeats the gradient with the per-component path, which applies the existing filter. The value is kept, so this fallback costs nothing in evaluations: option B plus fallback is (1+P)·N, the cost of the per-component path today. In the 37 measured cases the fallback never fired in either variant. The whole-call fallback was verified with a constructed integrand (exp(-θ·inf) → f = 0, derivative -inf · 0): with the fallback, value and derivative are exact to 12 digits; with the fallback disabled as a negative control, the call throws gradient_of_f: The gradient of f is nan for parameter 0.
How each function can be handled
The three functions share integrate_1d_adjoint, so the redundancy is the same in all three. They differ in the routine that computes the value, and that routine decides which option applies.
integrate_1d |
integrate_1d_double_exponential |
integrate_1d_gauss_kronrod |
|
|---|---|---|---|
| Value routine | integrate() |
integrate_de() |
integrate_gk() |
| Boost rules | tanh_sinh, exp_sinh, sinh_sinh |
tanh_sinh, exp_sinh, sinh_sinh |
gauss_kronrod |
| Stan signature | legacy (theta, x_r, x_i) |
variadic | variadic |
| Option A | only with a changed copy of Boost's tanh_sinh |
only with a changed copy of Boost's tanh_sinh |
yes |
| Option B, both variants | opted in | opted in | opted in |
| Return-type change for B | 7 return sites in integrate() |
7 return sites in integrate_de() |
1 site in integrate_gk() |
| Boost patch for B | 2-argument tanh_sinh (finite limits) |
2-argument tanh_sinh (finite limits) |
none |
| NaN rule at each evaluation under B | needed; acts on (−∞, ∞) and at a log-singular endpoint | needed; acts on (−∞, ∞) and at a log-singular endpoint | present; never needed in the measurements (transform stops near abs(x) ≈ 113) |
| Public models in the survey | 34 | 0 | 0 |
integrate_1dandintegrate_1d_double_exponentialuse the same Boost rules with the same interval splits. Both passxcthrough the 2-argumenttanh_sinhcall when both limits are finite.integrate_1d_double_exponentialadds an absolute tolerance and amax_refinementsargument. In the measurements below, the two value routines gave the same evaluation counts in all 12 cases and the same speedups within 2% in 11 of them; in the normal density on [−1, 2] they differ by 3–10%, and the double-exponential time of that case varied by 13% between the two passes.- In option B,
integrate_1d_adjointtakes a template flag (Jointin the joint variant), and all three functions set it. Forintegrate_1dthis needed the deduced return type inintegrate()(7 sites), the flag inintegrate_1d_impl, and the sametanh_sinhpatch asintegrate_1d_double_exponential. Three new tests pin the opt-in: option B and per-component results agree on all 7 interval branches ofintegrate(), a NaN gradient wheref == 0counts as 0, and a call withvararguments evaluates the integrand exactly as often as the same call withdoublearguments. With the opt-in removed, only the last test fails (513 evaluations instead of 119). A fourth test covers more than 8 arguments in the joint variant: 11 arguments take two blocks, and the adjoints agree with the per-component path. integrate_1dis the function that all 34 surveyed models use. Thus opting it in is the change that existing models would see.- Without a changed copy of Boost's
tanh_sinh, option A can help onlyintegrate_1d_gauss_kronrod.
Experiments
All measurements ran on Aalto Triton cluster batch nodes with Intel Xeon E5-2680 v3 (Haswell) CPUs, performance governor, g++ 14.3.0, -O3. In the speed measurements, today's code, the taped variant and the joint variant ran in the same process: best of 5 blocks, two interleaved rounds per configuration, pass 1 of 2. The speedups of the two passes agree within 4% in 27 of the 28 rows, and within 13% in one. "evals" counts integrand evaluations, today → option B (both variants evaluate the same nodes); P is the number of differentiable scalar arguments; speedups are relative to today's code. Neither variant fell back to the per-component path in any case.
Speed, Gauss-Kronrod
| case | P | evals | today/µs | taped | joint |
|---|---|---|---|---|---|
| beta-binomial marginal, saturated, model panels | 8 | 2079 → 231 | 2166 | 7.95× | 8.31× |
| beta-binomial marginal, interior, model panels | 8 | 2226 → 210 | 2882 | 9.61× | 9.96× |
| beta-binomial marginal, saturated, one call on [−8, 8] | 8 | 2037 → 189 | 2122 | 10.1× | 10.0× |
| beta-binomial marginal, interior, one call on [−8, 8] | 8 | 3633 → 399 | 4646 | 8.57× | 8.67× |
TestStudentT pdf (from the suite) |
3 | 3780 → 651 | 617 | 4.34× | 5.18× |
TestNormal pdf |
2 | 1449 → 441 | 94 | 1.78× | 2.39× |
TestCauchy pdf |
2 | 777 → 231 | 72 | 2.12× | 2.59× |
smooth_beta f11 |
2 | 315 → 21 | 38 | 11.1× | 12.5× |
- The "model panels" rows use the motivating model's scheme: z = r / σ on [−8, 8] in 16 unit panels, a log shift, and only the panels near the mode (here 11 and 10 of the 16), with relative tolerance 1e-6 and absolute tolerance 1e-10 per panel. The shift is a constant (
P= 8). The "one call" rows integrate over [−8, 8] in one call with relative tolerancesqrt(eps). The pdf rows use the settings of the test suite: relative tolerance 1e-8, and 1e-6 forsmooth_beta f11. - The speedup grows with
P, and also with how much further the per-component path refines the gradients than the value (smooth_beta f11: 315 → 21 evaluations). - Accuracy: the adjoints of both variants agree with today's to 2e-13 relative or better in the beta-binomial rows and to 3.3e-10 in the
smooth_beta f11row. In the three density rows on (−∞, ∞) the exact adjoints are 0 (a density integrates to 1), so a relative difference is not meaningful.
Speed, double-exponential
integrate_de() with relative tolerance sqrt(eps), absolute tolerance 0 and max_refinements 15.
| case | P | evals | today/µs | taped | joint |
|---|---|---|---|---|---|
| smooth, 8 parameters, [0, 2] | 8 | 990 → 119 | 316 | 7.30× | 7.65× |
| smooth, 8 parameters, [−1, 2] | 8 | 1886 → 240 | 580 | 6.73× | 7.22× |
| normal density, [−1, 2] | 2 | 449 → 132 | 38 | 2.45× | 2.87× |
| normal density, [−1, ∞) | 2 | 812 → 271 | 61 | 1.93× | 2.35× |
| beta(5,3) density, [0, 1] | 2 | 245 → 79 | 47 | 2.39× | 2.49× |
| normal density, (−∞, ∞) | 2 | 453 → 151 | 32 | 2.03× | 2.33× |
| beta-binomial marginal, saturated, r on (−∞, ∞) | 8 | 1611 → 277 | 1568 | 4.92× | 5.20× |
| beta-binomial marginal, interior, r on (−∞, ∞) | 8 | 1359 → 151 | 1710 | 8.04× | 8.43× |
| beta-binomial marginal, saturated, z on [−8, 8] | 8 | 3891 → 497 | 3972 | 6.86× | 7.25× |
| beta-binomial marginal, interior, z on [−8, 8] | 8 | 3028 → 364 | 3843 | 7.47× | 7.88× |
Speed, integrate_1d
integrate() with relative tolerance sqrt(eps).
| case | P | evals | today/µs | taped | joint |
|---|---|---|---|---|---|
| smooth, 8 parameters, [0, 2] | 8 | 990 → 119 | 319 | 7.25× | 7.71× |
| smooth, 8 parameters, [−1, 2] | 8 | 1886 → 240 | 586 | 6.72× | 7.28× |
| normal density, [−1, 2] | 2 | 449 → 132 | 36 | 2.37× | 2.62× |
| normal density, [−1, ∞) | 2 | 812 → 271 | 61 | 1.93× | 2.35× |
| beta(5,3) density, [0, 1] | 2 | 245 → 79 | 47 | 2.40× | 2.52× |
| normal density, (−∞, ∞) | 2 | 453 → 151 | 31 | 2.04× | 2.33× |
| beta-binomial marginal, saturated, r on (−∞, ∞) | 8 | 1611 → 277 | 1574 | 4.93× | 5.24× |
| beta-binomial marginal, interior, r on (−∞, ∞) | 8 | 1359 → 151 | 1723 | 8.07× | 8.52× |
| beta-binomial marginal, saturated, z on [−8, 8] | 8 | 3891 → 497 | 3972 | 6.86× | 7.25× |
| beta-binomial marginal, interior, z on [−8, 8] | 8 | 3028 → 364 | 3837 | 7.47× | 7.87× |
- Accuracy: in all rows except the normal (−∞, ∞) and beta rows, the adjoints of both variants agree with the per-component path to 2.2e-14 or better, and those of the joint variant with central finite differences to 6.2e-9 or better. In those two rows the exact adjoints are 0 (a density integrates to 1), so a relative difference is not meaningful.
- The beta-binomial rows are the motivating integral in two forms. The r rows integrate the random intercept r over the unbounded interval (−∞, ∞), as the model does, with the parameter packing of the model (
P= 8, 3 of them data). The z rows integrate over z = r / σ on the finite interval [−8, 8]. - The NaN derivatives come from underflow, at evaluations whose value is 0 or that the quadrature discards. On (−∞, ∞) these are the extreme nodes and the two evaluations at ±
max_value; for example, the derivative ofnormal_lpdf(r | 0, σ)with respect to σ contains r². On [0, 1], the beta(5,3) integrand callsbeta_lpdf(x | 5, 3)and does not usexc, so at the evaluations next to 1, x rounds to 1 andlog(1 − x)is −inf. A zero-valued node alone does not cause a fallback: the normal density on [−1, ∞) has 7 of them, but no NaN derivative. - The taped design with the whole-call fallback alone (the backup design) falls back in 4 rows: beta(5,3) density, normal density on (−∞, ∞), and both r rows. There it runs at 0.80–0.96× today's speed: the taped attempt replaces the all-
doublevalue pass with a pass invararithmetic, and its result is discarded. A gate onisfinite(a) && isfinite(b)would not avoid this, because the beta row has finite limits. In the other 6 rows it gains 2.35–7.8×. - The per-node loop of the taped variant, compared with the whole-call fallback alone, in rows without a fallback: 0.93–0.98× the speed in the smooth and beta-binomial rows, and 0.85–0.88× in the normal-density rows, where each evaluation creates only a few varis and the per-evaluation bookkeeping is a larger part of the cost. In a constructed case where every node is zero-valued,
θ0·exp(−θ1·x²)withθ0 = 0, 0.83–0.88×. Withintegrate_gk(), 0.81–1.03×, and 0.79× where all 21 nodes are zero-valued.
Compile time
One translation unit with integrate_1d* call sites for the beta-binomial marginal (8 arguments in a std::vector<var>, data as std::vector<int>), compiled with -c -O3. User CPU seconds, the lower of two runs. <stan/math.hpp> alone takes 14.3 s.
1 call site (integrate_1d) |
3 call sites (all three functions) | |
|---|---|---|
| today | 18.8 s | 22.8 s |
| taped | 19.4 s (+0.6) | 23.7 s (+0.9) |
| joint | 21.1 s (+2.4) | 26.5 s (+3.7) |
- One
integrate_1dcall site instantiates four Boost rules, because the interval branch is chosen at run time. - Each inline capacity of the joint value type instantiates the Boost rules once more, so the joint variant uses one capacity (8). With four capacities (2, 4, 8 and 32), one call site took +11.4 s. Capacity 8 is 0.95–1.00× the speed of an exact capacity for
P= 2 and 3.
Option A for comparison
Reverse-mode suite, gradient-pass evaluations 5481 → 2709 (2.02×), wall 1.36×. Beta-binomial marginal, 8 arguments: 3024 → 672 evaluations (4.50×), 1.88 → 0.35 ms (5.37×). Lower than option B. The per-node sweeps are not the reason: the joint variant of option B also sweeps at every node and runs at 0.99–1.34× the speed of the taped variant. Option A refines on the sup-norm over all components, so it needs at least as many evaluations as a refinement on the value, and here twice the average count of the per-component quadratures (672 against 3024 / 9 = 336).
Memory
This is where the two variants of option B differ most. The per-component scheme and the joint variant free their nested scope at every node, so their memory is flat in the number of evaluations. The taped variant holds the whole graph.
Autodiff stack entries and bytes per integrand evaluation (the quadrature arithmetic included), and the resulting tape for one quadrature of the taped variant:
| case | entries/eval | bytes/eval | evals | tape |
|---|---|---|---|---|
| beta-binomial marginal, Gauss-Kronrod, z on [−8, 8] | 39 | 1660 | 189–399 | 306–646 KB |
beta-binomial marginal, integrate_1d, r on (−∞, ∞) |
37 | 1590 | 151–277 | 235–428 KB |
TestStudentT pdf |
9.6 | 400 | 651 | 255 KB |
TestNormal pdf |
8.6 | 360 | 441 | 156 KB |
smooth_beta f11 |
7.2 | 330 | 21 | 7 KB |
The per-node loop of the taped variant adds a 32-byte record per evaluation: 2% of the tape for the beta-binomial marginal and 8–10% for the cheap pdfs.
Peak RSS of a process that computes one integral of θ0 + θ1·sin(K·θ2·x) + θ3·x on [0, 1], 4 arguments, Gauss-Kronrod with relative tolerance 1e-8 (the process without the integral uses 3.2–3.8 MB):
| K | depth | evals | today | taped | joint |
|---|---|---|---|---|---|
| 10³ | 15 | 5 355 | 3.4 MB | 6.1 MB | 3.4 MB |
| 10⁴ | 15 | 74 067 | 3.2 MB | 44 MB | 3.4 MB |
| 10⁵ | 15 | 934 269 | 3.6 MB¹ | 510 MB | 3.2 MB |
| 10⁶ | 15 | 1 376 235 | 3.6 MB² | 673 MB² | 3.6 MB² |
| 10⁶ | 18 | 10 778 271 | 3.8 MB² | 4994 MB² | 3.8 MB² |
¹ Today's code throws, because a gradient integral does not reach the tolerance. Both variants of option B return a gradient computed on the partition of the value; it agrees with the analytic gradient to 1e-9 relative, while today's code throws also at depth 21 after 52 million evaluations.
² The value integral does not reach the tolerance, so every design throws. The taped variant throws after it has built the tape.
In the taped variant, the evaluation count sets the memory, not max_depth alone: for K = 10³ the count is the same at depths 10 to 18. The row K = 10⁶, depth 15 is the full Gauss-Kronrod tree at the default depth, (2¹⁶ − 1) × 21 evaluations. So at the default depth, one quadrature can hold about 0.7 GB of tape for a cheap integrand (measured) and about 2.3 GB for the beta-binomial marginal (estimated from the bytes per evaluation above).
For calibration, single-quadrature evaluation counts across the 29 reverse-mode test cases: max 651, mean 96, median 21 (21 of 29 are a single 21-node panel with no refinement). The hardest integrand anywhere in the test suite is 1743 evaluations ≈ 680 KB of tape.
So every integral in the test suite is three orders of magnitude below where the taped variant's memory becomes a problem. sin(10⁵x) is an extreme artificial example to illustrate memory growth. But one quadrature at the default max_depth can already hold most of a GB, max_depth is user-settable with only check_nonnegative, and a multi-GB allocation inside a log-density evaluation is not an acceptable failure mode, so the taped variant needs some guard. A node- or byte-based budget (not an evaluation count — bytes per evaluation vary 5× across the cases above) with the existing fallback as the escape hatch would do it, but it would never fire on any current test and so needs a dedicated one. The joint variant needs no guard.
Implementation and tests
Options A and B (both variants) are implemented and pass the 11 existing integrate_1d* suites (prim, rev, mix). New tests cover the equivalence of option B and the per-component path, the empty argument pack, and the NaN rule at each evaluation: mixed scalar, Eigen and std::vector arguments with data and int; a var limit; the legacy signature with a log-singular endpoint; a nested scope over a tape with a nonzero outer adjoint; a call inside an ode_rk45 right-hand side; an integrand built from reverse_pass_callback functions; a discarded evaluation with a NaN value; and reduce_sum on 4 threads (STAN_THREADS). With the whole-call fallback alone, 7 of these fall back. For the joint variant, two more tests cover more than 8 arguments and a non-finite integrand value at an endpoint.
Questions I would like input on
- Option A or B, and which variant of B? B is faster than A and covers all three functions with the patch of question 3; A needs a changed copy of Boost's double-exponential code for two of them. Both options change how the gradients are refined (question 5). The variants of B trade compile time against memory (table in "Option B"). The joint variant seems to be the best.
- Taped variant only: may
integrate_1d_adjointrun its own reverse-pass loop over the tape? A small helper inrev/corethat sweeps the current nest and calls a hook on marked sub-ranges would keep the stack access out of the functor. The backup is the whole-call fallback alone. - Is the 2-line
tanh_sinhpatch acceptable? Option B needs it for finite limits inintegrate_1dandintegrate_1d_double_exponential. Gauss-Kronrod needs no change. - Taped variant only: is a memory guard needed, and with what budget?
- Is the value-driven refinement acceptable? A derivative rougher than its value gets no extra refinement. For
θ0·exp(−θ1·x²)withθ0 = 0, where the value is 0 everywhere, the double-exponential adjoint∂I/∂θ0differs from today's by 6e-5 to 8e-4 relative (Gauss-Kronrod: 7e-14). The opposite also occurs: forsin(10⁵x)(see "Memory") today's code throws, and option B is accurate to 1e-9. - Joint variant only: is +2.4 s compile time per call site acceptable?
- Keep develop's change to Boost's Gauss-Kronrod? With option B it is no longer absolutely needed against the gradient blow-up, but it still saves value evaluations (see "Option B").
Activity