Explore custom reverse-mode adjoints for the analytic primary censored lcdfs #362

Summary

The analytic uniform primary lcdfs (primarycensored_lognormal_uniform_lcdf(), primarycensored_gamma_uniform_lcdf(), primarycensored_weibull_uniform_lcdf()) could each be one C++ function with a closed-form gradient with respect to q and the distribution parameters.
That puts one node on the autodiff tape per call instead of dozens.
Every caller gains: primarycensored_lcdf(), primarycensored_lpmf(), the vectorised PMF, the fitting models and EpiNow2.
The fallback to the Stan code is per function.

For gamma and Weibull the shape derivative of the regularised incomplete gamma, dP(a, x)/da, is the main cost.
Eigen's igamma_der_a computes it about 3x faster than Stan Math's grad_reg_lower_inc_gamma(), and to about 1e-15 rather than 1e-3 relative.

A prototype gives:

• the lcdf alone: about 2.7x faster gradients than main for lognormal and 2x to 8x for gamma and Weibull;
• primarycensored_sone_pmf_vectorized() with a C++ version of primarycensored_analytical_lcdf_vectorized() (the perf: share uniform primary terms across nodes in the vectorised PMF #356 seam): about 2.7x for lognormal and 2.3x to 4.4x for gamma and Weibull;
• values within 4e-14 of a 50-digit mpmath reference, and gradients within 2e-14 (the PMF Jacobian).
  On main the gradients are within 5e-12 for gamma but wrong by up to 7e-5 for Weibull (see below).

A fused adjoint for the whole vectorised PMF would add 1.6x to 2.6x more on top (3.7x to 9.4x against main in total).
It is set out below as a possible follow-up.

Motivation

EpiNow2 calls primarycensored_sone_pmf_vectorized(n - 1, L, n, dist, params, 1.0, 1, primary_params) from discretised_pmf().
That is a uniform primary, window 1, D = max_delay + 1, and lognormal, gamma or Weibull delays with estimated parameters.
max_delay is usually 10 to 40.
In EpiNow2's benchmark the "delays" profile block is now the largest, at 7.9 s of about 18 s summed profile time.
epiforecasts/EpiNow2#1548 found that about 91% of that block is the discretisation, about 30 µs per gradient for two lognormal delays (max 14 and 10).
EpiNow2 is adding C++ adjoints for its convolution (epiforecasts/EpiNow2#1543) and renewal equation (epiforecasts/EpiNow2#1545).
The discretisation belongs here.

How this relates to the recent optimisation work

• perf: share uniform primary terms across nodes in the vectorised PMF #356 computes the uniform primary terms once per integer node and reuses them as d and q = d - pwindow. This halved CDF evaluations and gave 1.3x (lognormal) to 1.7x (gamma, Weibull).
• perf: reduce per-call overhead in dprimarycensored, pprimarycensored and fitdistdoublecens (#347) #353 cut R-side per-call overhead. It does not touch the Stan path.
• Express analytical primary-censored CDFs in direct CDF form #306 put the analytic CDFs in direct CDF form. fix(stan): guard lognormal lower-tail underflow to keep gradients finite #337 and Vectorised Stan PMF returns NaN in the saturated upper tail #357 deal with lower and upper tail underflow on the log scale.

After #356 the remaining gradient cost is autodiff overhead plus, for gamma and Weibull, the incomplete gamma shape derivative.
Each lcdf puts two CDF calls per term, several log_sum_exp/log_diff_exp nodes and the parameter transforms on the tape.
For lognormal the CDFs are cheap, so the tape dominates.
For gamma and Weibull Stan computes dP/da with grad_reg_lower_inc_gamma(), a series that is slow and, for larger a and small x, inaccurate.
An adjoint removes the tape, and igamma_der_a replaces the series.

Proposed maths

For a uniform primary with window w, let

G(t) = int_0^t F_T(s) ds = t F_T(t) - M(t),   M(t) = int_0^t s f_T(s) ds

with G(t) = 0 for t <= 0.
Then F_S(d) = (G(d) - G(q)) / w with q = max(d - w, 0).
This is the (A - B) / w split in primarycensored_uniform_lcdf_from_terms(), with A = d F_T(d) + M(q) and B = q F_T(q) + M(d).
The value is computed on the log scale exactly as now.

With Delta = G(d) - G(q):

d lcdf / dq     = -F_T(q) / Delta
d lcdf / dtheta = (dG/dtheta(d) - dG/dtheta(q)) / Delta

The gradient of G is closed form.
Since G(t) = E[(t - T)^+], writing T = T(theta, U) for a fixed base variable U gives dG/dtheta = -E[dT/dtheta 1{T < t}].

Lognormal (mu, sigma), z = (log t - mu) / sigma:

M = exp(mu + sigma^2 / 2) Phi(z - sigma)
dG/dmu    = -M
dG/dsigma = -sigma M + t phi(z)

For lognormal the ratios to Delta are formed on the log scale, so they do not underflow.

Gamma (k, beta), y = beta t, h = y^k e^(-y) / Gamma(k + 1):

M = (k / beta) P(k + 1, y),   P(k + 1, y) = P(k, y) - h
dG/dbeta = M / beta
dG/dk    = t P_a(k, y) - P(k + 1, y) / beta - (k / beta) P_a(k + 1, y)
P_a(k + 1, y) = P_a(k, y) - h (log y - psi(k + 1))

Weibull (k, lambda), x = (t / lambda)^k, a = 1 + 1 / k:

M = lambda Gamma(a) P(a, x)
dG/dlambda = -M / lambda
dG/dk      = lambda Gamma(a) (P_a(a, x) + P(a, x) psi(a)) / k^2

P_a = dP/da.
The recurrence P(a + 1, x) = P(a, x) - x^a e^(-x) / Gamma(a + 1) gives P_a(k + 1, y) from P_a(k, y) in closed form, so gamma needs one P_a call per node.
Weibull needs one P_a call per node at a = 1 + 1/k.
In the Weibull case the dP/dx terms cancel against t dF_T/dk exactly, so they are never computed.
The gamma and Weibull Jacobians match mpmath to 3e-15 and 2e-14 (below), which checks both the recurrence and the cancellation.
Generalised gamma (5) should follow the same pattern, with one P_a per node.

Computing P_a

P_a comes from Eigen's igamma_der_a: Eigen::internal::igamma_der_a_impl<double>::run(a, x) in unsupported/Eigen/SpecialFunctions.
It is in Eigen 3.4.0, which ships with CmdStan 2.39 and with RcppEigen 0.3.4.0.2.
Boost has no derivative with respect to a (gamma_p_derivative is d/dx).
Forward-mode autodiff through Boost's gamma_p silently returns zero on its continued-fraction branches.

Accuracy and cost of dP/da on a 63-point grid (a from 0.5 to 10, x from 0.01 to 50), against mpmath at 50 digits:

stan-dev/math#3408 (open) moves Stan's own gamma_p/gamma_q gradients to igamma_der_a.
That would fix the accuracy and the flat-gradient shortcut below for Stan code upstream.
The adjoint would not need to wait for it, because it calls Eigen directly.

Design

• Declare primarycensored_{lognormal,gamma,weibull}_uniform_lcdf() in Stan without a body. Implement them in a C++ header, each as one reverse_pass_callback.
• Use igamma_der_a for P_a.
• Also move primarycensored_analytical_lcdf_vectorized() to C++, sharing the per-node terms as perf: share uniform primary terms across nodes in the vectorised PMF #356 does.
  This matters.
  If the vectorised path instead calls the scalar lcdf per node, each node is evaluated twice again (at d and as q).
  With Stan's P_a that was 1.55x faster than main for lognormal but 0.63x to 0.8x (slower) for gamma and Weibull.
• Truncation, differencing and exp() stay in Stan.
• Everything else (other delays, exponential growth primary, numerical integration) is unchanged.

Prototype

Code: https://gist.github.com/seabbs-bot/150a5971bc769d906f3968069c7ca02b

• pcens_uniform_lcdf_adjoint.hpp implements the three scalar lcdfs and primarycensored_analytical_lcdf_vectorized(). It follows Use a custom reverse-mode adjoint for the delay convolution epiforecasts/EpiNow2#1543 (reverse_pass_callback, arena storage, a using-declaration into stan::math).
• dpda.hpp wraps igamma_der_a. Defining PC_STAN_DPDA switches back to grad_reg_lower_inc_gamma() for comparison.
• make_variants.sh builds the Stan variants from main with those functions declared without a body.
• It compiles as a CmdStan USER_HEADER with --allow-undefined (vecmodel.stan), and diagnose test=gradient passes for all three delays.
• bench2.cpp calls each implementation through Stan Math and times one gradient.
• ref_pmf.py computes the mpmath reference and acc.R compares against it.

Benchmark

CmdStan 2.39 (Stan Math, -O3), Apple clang 17, macOS arm64, single thread.
µs per gradient: the median of 7 repeats (100,000 calls each for the lcdf, 20,000 for the PMF), then the minimum over 3 interleaved rounds.
The host was shared, with a load average of 8 to 12 during the runs.
Lognormal does not use P_a, so its two adjoint columns run the same code.
The gap between them (up to about 20%) shows the noise level.

lcdf alone, q = d - 1, pwindow = 1, gradient with respect to both parameters.
Speed-ups are against main.

primarycensored_sone_pmf_vectorized(), L = 0, D = max_delay + 1, pwindow = 1, gradient of dot_product(w, pmf).
The two "vectorised lcdf" columns use the C++ primarycensored_analytical_lcdf_vectorized(), with truncation and differencing still in Stan.
"Fused" is the follow-up option below.

With Stan's P_a the vectorised lcdf gains only 1.2x to 1.6x for gamma and Weibull.
igamma_der_a accounts for most of the rest of the gain.
For EpiNow2's two lognormal delays (max 14 and 10), interpolating the table gives about 8.5 µs per gradient on main, 3.3 µs with the vectorised lcdf, and 1.4 µs fused.

Correctness

The reference is mpmath at 50 digits, differentiated by central differences with step 1e-20.
I checked 6 parameter sets (the ones above plus lognormal(0.6, 0.9)): the lcdf at d = 1, 2, 5, 15 and 30, and the full PMF Jacobian at max_delay 10, 20 and 40.

The fused version with igamma_der_a is within 2.8e-15 (gamma) and 1.5e-14 (Weibull).
Values match the reference to 4e-14 absolute for every implementation.
The lcdf relative errors leave out gradients below 1e-6, where the log CDF is within 1e-10 of zero and double precision limits every method.

• Weibull on main. Stan Math's reverse-mode gamma_p returns zero derivatives whenever x / a > 10 ("flat to machine precision"). That drops d log P / da terms that the differencing then scales up. At Weibull(1.5, 5), max_delay 40, the Jacobian for d >= 32 (p < 1e-7) is wrong by up to 7e-5, sometimes with the wrong sign. Incomplete gamma gradient roots and log-space cdfs for six distributions stan-dev/math#3408 would fix this upstream.
• gamma(5, 1.2) at max_delay 40. main returns NaN for one PMF entry and for the whole Jacobian (Vectorised Stan PMF returns NaN in the saturated upper tail #357). The lcdf-level adjoints keep the Stan differencing, so they give the same NaN. The fused version is finite.

Packaging notes

• Shipping the header. Put it in inst/include/primarycensored/ and add an R helper that returns its path, like epinow2_stan_header() in Use a custom reverse-mode adjoint for the delay convolution epiforecasts/EpiNow2#1543.

• cmdstanr. Users and downstream packages pass it as user_header and compile with stanc_options = list("allow-undefined"). pcd_cmdstan_model() can do this by default.

• rstan. rstantools includes stan_meta_header.hpp too late for functions that the Stan functions call. EpiNow2 works round this with a configure step that injects the include into src/stanExports_*.h. Downstream rstan packages that vendor these functions (EpiNow2) would need the same, or would copy the header.

• Fallback without the header. Stan has no conditional compilation. A function declared without a body fails to compile if the header is missing. Options:

  1. keep the Stan bodies as the default, and have pcd_stan_functions() / the vendoring script swap in declarations only when asked;
  2. ship two Stan files (with and without the C++ functions) and pick one at load time;
  3. always require the header, which is a breaking change for anyone who copies the Stan files.

  Option 1 seems least disruptive. Because the swap is per function, it can start with lognormal only.

• Name clash. Put only functions that are declared without a body in the header. In the prototype, a C++ primarycensored_analytical_lcdf_vectorized() in stan::math was silently picked over a Stan-defined function of the same name, through argument-dependent lookup.

• expose_stan_functions. Tests that expose Stan functions need allow_undefined = TRUE and the header in includes, as in Use a custom reverse-mode adjoint for the delay convolution epiforecasts/EpiNow2#1543.

• Eigen dependency. igamma_der_a is in Eigen::internal in the unsupported module, so it is not a stable API. It is in Eigen 3.4.0, which ships with CmdStan 2.39 and RcppEigen 0.3.4.0.2. The header should check the Eigen version and fall back to grad_reg_lower_inc_gamma() if it is missing.

• CRAN / StanHeaders floor. The prototype also uses reverse_pass_callback, arena_t, gamma_p, log_diff_exp, log1m_exp and digamma. These should all be in StanHeaders 2.26, but I have not compiled against it. It is plain C++17. primarycensored itself has no compiled code, so the header only compiles in the user's or downstream package's model.

Possible follow-up: a fused adjoint for the vectorised PMF

A single C++ primarycensored_sone_pmf_vectorized() for the analytic cases would evaluate G once at t = 1, ..., D and form

F_S(d) = (G(d) - G(d - w)) / w
p_d    = (F_S(d) - F_S(d - 1)) / (F_S(D) - F_S(L)),   L < d <= D

with the reverse pass

Z = F_S(D) - F_S(L),  Zbar = -sum_d pbar_d p_d / Z
FSbar(d) += pbar_d / Z,  FSbar(d - 1) -= pbar_d / Z     (d > L)
FSbar(D) += Zbar,  FSbar(L) -= Zbar
Gbar(t) += FSbar(t) / w,  Gbar(t - w) -= FSbar(t) / w
thetabar_j += sum_t Gbar(t) dG_j(t)

on doubles in one callback.
The prototype (pcens_pmf_adjoint.hpp in the gist) covers integer L >= 0, integer w and D = max_delay + 1, on the natural scale.
Its timings are the "Fused" column above.
On top of the vectorised lcdf with igamma_der_a it adds about 2.3x for lognormal, 1.6x to 2.6x for gamma and 1.6x to 1.7x for Weibull.
It also avoids the #357 NaN, though tail entries can then be about -7e-15 from rounding, so it would need clamping at zero.
Its surface is larger: it duplicates the truncation logic, and it would need a log-scale variant for primarycensored_sone_lpmf_vectorized().

Open questions

• Is the C++ and packaging cost worth it? With igamma_der_a the gain is now 2x to 4x for all three delays, not only lognormal.
• Is relying on Eigen::internal::igamma_der_a_impl acceptable, or should the header vendor that routine?
• Should generalised gamma be included in the first step?
• Should primarycensored's own Stan code call igamma_der_a once Incomplete gamma gradient roots and log-space cdfs for six distributions stan-dev/math#3408 lands, independent of the adjoint?
• Is the fused PMF adjoint worth a second step for EpiNow2, given the extra 1.6x to 2.6x?

This was opened by a bot. Please ping @seabbs for any questions.