Skip to content

integrate_1d_adjoint repeats the whole quadrature once per differentiable argument #3444

Open

Description

@avehtari

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 theta array instead of x_r, so they are stored as var (see below); without them P would 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_1d in .stan files: 395 hits, 98 different files. 50 files reach integrate_1d from model or transformed 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, .Rmd and .qmd files on GitHub: 7 models, 4 of them brms custom families.
  • The Stan User's Guide, 1 Stan Math issue, 1 epinowcast/primarycensored pull 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:

  1. Most models call integrate_1d once 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 runs P + 1 quadratures. Thus in most models the per-call cost is multiplied by the data size, which is likely to make integrate_1d to dominate the computation time.
  2. The legacy signature adds constants to P. In integrate_1d(f, a, b, theta, x_r, x_i), theta is one array[] real. If any element is a parameter, stanc3 stores every element as var (the generated C++ declares std::vector<local_scalar_t__>). 7 of 34 models put data or literals into theta, 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.
  3. 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 theta elements. 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.
  4. All 34 models use the legacy integrate_1d signature. The variadic signatures of integrate_1d_double_exponential and integrate_1d_gauss_kronrod let the user pass data as separate double arguments, 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:

  1. 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_adjoint computes 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.
  2. 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_sinh computes Real / K in one place, which has no vector meaning, and exp_sinh compares values of K with == (verified: neither compiles with the vector type). A Stan copy of tanh_sinh with 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 of tanh_sinh, exp_sinh and sinh_sinh. Without such copies, integrate_1d and integrate_1d_double_exponential keep 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 all P derivatives at once: the sweep that the current code already does, but once instead of P times. 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_de and integrate must return a deduced type instead of double (mechanical; 1, 7 and 7 sites).
  • For the finite branch of integrate_1d_double_exponential and integrate_1d, Boost's 2-argument (xc) tanh_sinh::integrate hard-codes Real in 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():

  1. 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.
  2. 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_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. 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.

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_1d and integrate_1d_double_exponential use the same Boost rules with the same interval splits. Both pass xc through the 2-argument tanh_sinh call when both limits are finite. integrate_1d_double_exponential adds an absolute tolerance and a max_refinements argument. 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_adjoint takes a template flag (Joint in the joint variant), and all three functions set it. For integrate_1d this needed the deduced return type in integrate() (7 sites), the flag in integrate_1d_impl, and the same tanh_sinh patch as integrate_1d_double_exponential. Three new tests pin the opt-in: option B and per-component results agree on all 7 interval branches of integrate(), a NaN gradient where f == 0 counts as 0, and a call with var arguments evaluates the integrand exactly as often as the same call with double arguments. 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_1d is 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 only integrate_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 tolerance sqrt(eps). The pdf rows use the settings of the test suite: relative tolerance 1e-8, and 1e-6 for smooth_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 f11 row. 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 of normal_lpdf(r | 0, σ) with respect to σ contains r². On [0, 1], the beta(5,3) integrand calls beta_lpdf(x | 5, 3) and does not use xc, so at the evaluations next to 1, x rounds to 1 and log(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-double value pass with a pass in var arithmetic, and its result is discarded. A gate on isfinite(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×. With integrate_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_1d call 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

  1. 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.
  2. Taped variant only: may integrate_1d_adjoint run its own reverse-pass loop over the tape? A small helper in rev/core that 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.
  3. Is the 2-line tanh_sinh patch acceptable? Option B needs it for finite limits in integrate_1d and integrate_1d_double_exponential. Gauss-Kronrod needs no change.
  4. Taped variant only: is a memory guard needed, and with what budget?
  5. 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/∂θ0 differs from today's by 6e-5 to 8e-4 relative (Gauss-Kronrod: 7e-14). The opposite also occurs: for sin(10⁵x) (see "Memory") today's code throws, and option B is accurate to 1e-9.
  6. Joint variant only: is +2.4 s compile time per call site acceptable?
  7. 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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions