Improve identifiability of non-stationary GP - #1396
Conversation
Three changes to inst/stan that together eliminate stuck chains and divergences seen on the non-stationary Gaussian process used to model Rt: 1. Rescale GP increments by 1/sqrt(gp_n) inside update_Rt so alpha controls the trajectory SD rather than the increment SD, matching what gp_opts() docs already claim. 2. Centre the cumulated GP (gp -= mean(gp)) so log R0 = mean log Rt rather than the initial value. Eliminates the (R0, drift) ridge in the joint posterior. 3. Switch GP coefficient sampling to centred form: eta ~ N(0, diagSPD) with noise = PHI * eta. Avoids the (alpha, eta) funnel that arises in non-centred form when alpha is small. On previously catastrophic seeds (R-hat > 4 with the old form) the new parameterisation gives R-hat < 1.02 with zero divergences and roughly 4x faster sampling. Co-authored-by: sbfnk <sebastian.funk@lshtm.ac.uk>
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
WalkthroughThis pull request reparameterises the non-stationary Gaussian process for time-varying Rt. The GP is now mean-centred during sampling to eliminate a posterior ridge. R0 is handled through a new initial-prior dispatch system instead of the generic parameter list, preserving the user-facing prior meaning via a Jacobian-corrected log-density evaluation. ChangesNon-stationary Rt GP mean-centring and initial-prior system
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
Just to note this is far from being ready to merge but I'm curious about the headline benchmark results given the outrageous performance claims. |
|
Synthetic recovery shows reduced GP amplitude (presumably related to rescaling / changed meaning of alpha) which may well explain the fitting improvements (which in itself is perhaps a useful observation that could support addressing #1376). |
|
Yup nice I have been doing similar sweeps on Centering the GP seems sensible |
|
This is how benchmark results would change (along with a 95% confidence interval in relative change) if 5d870b5 is merged into main:
|
Under the old non-stationary GP parameterisation, alpha was the increment SD and the implied trajectory SD scaled with sqrt(gp_n). With the new parameterisation (rescaling + cumsum centring), alpha is the trajectory SD directly, and the same numerical alpha implies a ~13× tighter Rt trajectory than before for typical gp_n. Bumping the default from Normal(0, 0.01) to Normal(0, 0.08) restores approximately the same Rt expressiveness as the old default. Documents the new alpha semantics in the gp_opts() roxygen. Verified empirically: the previously stuck seeds still sample without divergences under the wider prior, so the geometric improvements stand on their own — this is purely a default-tuning fix. Co-authored-by: sbfnk <sebastian.funk@lshtm.ac.uk>
estimate_infectionsTable: Benchmarking results (mean time in seconds). |operation | branch| main| % change| range| trend| |:--------------------|-------:|-------:|--------:|----------:|---------:| |delays | 2.9| 11| -72| (-85, 12)| no change| |infections | 1.2| 4.2| -71| (-85, 19)| no change| |reports | 1.2| 4.1| -72| (-85, 22)| no change| |report lp | 0.86| 3.1| -72| (-85, 12)| no change| |R0 | 0.63| 1.3| -52| (-75, 99)| no change| |update gp | 0.14| 1.1| -87| (-93, -47)| speedup| |gt | 0.13| 0.48| -72| (-85, 15)| no change| |day of the week | 0.12| 0.43| -72| (-85, 21)| no change| |truncation | 0.077| 0.29| -74| (-86, 9)| no change| |param lp | 0.076| 0.27| -72| (-84, 15)| no change| |truncate | 0.061| 0.22| -71| (-84, 25)| no change| |delays lp | 0.034| 0.13| -73| (-85, 18)| no change| |rt lp | 0.035| 0.12| -71| (-85, 33)| no change| |gp lp | 0.25| 0.1| 150| (43, 950)| slowdown| |generated quantities | 0.0068| 0.0073| -6| (-15, 18)| no change| |assign max | 5.1e-07| 4.9e-07| 9| (-50, 123)| no change| |
|
This is how benchmark results would change (along with a 95% confidence interval in relative change) if f25229e is merged into main:
|
|
Useful food for thought is https://www.generable.com/post/hsgp-reparam (this is what I was planning on looking at and the findings here seem to align with what it is saying) |
Empirical tests show the data prefers alpha around 0.29 in the new parameterisation, but Normal(0, 0.08) puts only 0.06% prior mass above that — the prior would clamp alpha well below where the data wants it. Also: the sqrt(3) Brownian-bridge factor for the centring's effect on trajectory SD is an underestimate for the smooth Matern GP we use. Bumping the default sd to 0.2 (16% prior mass above 0.29) lets the data drive the posterior without the prior dominating, restoring Rt expressiveness comparable to the old parameterisation. Co-authored-by: sbfnk <sebastian.funk@lshtm.ac.uk>
…sts/EpiNow2 into identifiable-non-stationary-gp
estimate_infectionsTable: Benchmarking results (mean time in seconds). |operation | branch| main| % change| range| trend| |:--------------------|-------:|-------:|--------:|----------:|---------:| |delays | 2.9| 11| -72| (-82, -59)| speedup| |infections | 1.2| 4.3| -70| (-80, -56)| speedup| |reports | 1.2| 4.2| -72| (-82, -59)| speedup| |report lp | 0.89| 3.2| -72| (-82, -59)| speedup| |R0 | 0.65| 1.3| -51| (-68, -27)| speedup| |update gp | 0.14| 1.1| -87| (-92, -80)| speedup| |gt | 0.14| 0.5| -72| (-82, -57)| speedup| |day of the week | 0.13| 0.45| -71| (-81, -57)| speedup| |truncation | 0.08| 0.3| -73| (-82, -60)| speedup| |param lp | 0.08| 0.28| -71| (-82, -52)| speedup| |truncate | 0.064| 0.23| -71| (-81, -60)| speedup| |delays lp | 0.034| 0.12| -71| (-81, -57)| speedup| |rt lp | 0.037| 0.12| -70| (-82, -59)| speedup| |gp lp | 0.26| 0.1| 158| (69, 305)| slowdown| |generated quantities | 0.0067| 0.0068| -3| (-6, 1)| no change| |assign max | 5.1e-07| 4.8e-07| 7| (-31, 170)| no change| |
Four test failures triggered by the model changes in this branch: - test-gp_opts.R:5: default alpha is now Normal(0, 0.2) - test-stan-guassian-process.R:163: update_gp now returns PHI * eta directly (the diagSPD scaling moved to the eta prior in the centred form) - test-stan-rt.R:12, :56: update_Rt outputs differ because of the 1/sqrt(gp_n) increment rescaling and the cumsum centring; new expected values computed for the same inputs Co-authored-by: sbfnk <sebastian.funk@lshtm.ac.uk>
estimate_infectionsTable: Benchmarking results (mean time in seconds). |operation | branch| main| % change| range| trend| |:--------------------|-------:|-------:|--------:|----------:|---------:| |delays | 2.5| 8.1| -68| (-80, -49)| speedup| |infections | 1.1| 3.3| -67| (-80, -48)| speedup| |reports | 1| 3.3| -68| (-79, -49)| speedup| |report lp | 0.73| 2.3| -68| (-80, -49)| speedup| |R0 | 0.56| 1| -45| (-66, -13)| speedup| |update gp | 0.1| 0.78| -86| (-91, -78)| speedup| |gt | 0.11| 0.37| -69| (-81, -52)| speedup| |day of the week | 0.1| 0.33| -68| (-80, -50)| speedup| |truncation | 0.07| 0.23| -69| (-80, -50)| speedup| |param lp | 0.059| 0.19| -68| (-80, -48)| speedup| |truncate | 0.051| 0.17| -69| (-79, -53)| speedup| |delays lp | 0.025| 0.081| -68| (-80, -50)| speedup| |rt lp | 0.025| 0.076| -67| (-80, -42)| speedup| |gp lp | 0.2| 0.071| 192| (88, 360)| slowdown| |generated quantities | 0.0059| 0.007| -16| (-25, -5)| speedup| |assign max | 5.3e-07| 5.7e-07| 3| (-65, 195)| no change| |
|
This is how benchmark results would change (along with a 95% confidence interval in relative change) if e234216 is merged into main:
|
Prior-sensitivity in alpha is real: a 13× shift in prior median moves the posterior median by ~50%, suggesting the data alone is only weakly informative about alpha. Bumping the default to Normal(0, 0.5) gives users plenty of headroom for the data to find its preferred value, while documenting that informed priors are worthwhile when domain knowledge is available. Co-authored-by: sbfnk <sebastian.funk@lshtm.ac.uk>
estimate_infectionsTable: Benchmarking results (mean time in seconds). |operation | branch| main| % change| range| trend| |:--------------------|------:|-------:|--------:|----------:|---------:| |delays | 4.5| 11| -58| (-80, -35)| speedup| |infections | 1.9| 4.5| -55| (-78, -30)| speedup| |reports | 1.8| 4.4| -57| (-79, -33)| speedup| |report lp | 1.4| 3.4| -58| (-80, -34)| speedup| |R0 | 1| 1.4| -27| (-65, 17)| no change| |update gp | 0.22| 1.2| -81| (-90, -69)| speedup| |gt | 0.21| 0.53| -58| (-81, -38)| speedup| |day of the week | 0.2| 0.47| -56| (-80, -31)| speedup| |truncation | 0.12| 0.31| -59| (-81, -39)| speedup| |param lp | 0.12| 0.3| -58| (-80, -35)| speedup| |truncate | 0.099| 0.24| -57| (-79, -32)| speedup| |rt lp | 0.056| 0.13| -56| (-79, -30)| speedup| |delays lp | 0.055| 0.12| -53| (-76, -20)| speedup| |gp lp | 0.4| 0.11| 280| (87, 448)| slowdown| |generated quantities | 0.0078| 0.0079| -2| (-12, 13)| no change| |assign max | 5e-07| 5.3e-07| -4| (-53, 56)| no change| |
|
This is how benchmark results would change (along with a 95% confidence interval in relative change) if 71112fd is merged into main:
|
Normal(0, 0.5) breaks sampling on previously stuck seeds: the wider prior lets chains wander into large-alpha regions during warmup where the geometry funnel is bad enough that the centred form can't compensate. Verified empirically — seed=8 went from clean (R-hat=1.022, td=0) at sd=0.2 to catastrophic (R-hat=1020, td_hits=500) at sd=0.5. Document the prior-sensitivity but keep the default conservative. Co-authored-by: sbfnk <sebastian.funk@lshtm.ac.uk>
estimate_infectionsTable: Benchmarking results (mean time in seconds). |operation | branch| main| % change| range| trend| |:--------------------|-------:|-------:|--------:|-----------:|---------:| |delays | 4.1| 10| -59| (-78, 255)| no change| |infections | 1.7| 4.1| -58| (-78, 263)| no change| |reports | 1.7| 4| -59| (-78, 252)| no change| |report lp | 1.2| 2.9| -59| (-78, 258)| no change| |R0 | 0.85| 1.2| -31| (-63, 502)| no change| |update gp | 0.19| 0.97| -81| (-90, 63)| no change| |gt | 0.19| 0.45| -57| (-77, 269)| no change| |day of the week | 0.19| 0.44| -58| (-77, 265)| no change| |truncation | 0.12| 0.28| -60| (-78, 241)| no change| |param lp | 0.12| 0.26| -55| (-75, 282)| no change| |truncate | 0.088| 0.2| -57| (-76, 257)| no change| |delays lp | 0.049| 0.13| -62| (-78, 227)| no change| |rt lp | 0.04| 0.099| -61| (-78, 206)| no change| |gp lp | 0.36| 0.082| 332| (126, 3670)| slowdown| |generated quantities | 0.0069| 0.0067| 3| (-3, 20)| no change| |assign max | 6.1e-07| 5.2e-07| 20| (-45, 87)| no change| estimate_distTable: Benchmarking results (mean time in seconds). |operation | branch| main| % change| range| trend| |:----------|------:|------:|--------:|------:|---------:| |likelihood | 0.21| 0.21| 1| (0, 2)| no change| |priors | 0.0015| 0.0014| 6| (3, 9)| slowdown| |
|
This is how benchmark results would change (along with a 95% confidence interval in relative change) if 7a0326e is merged into main:
|
Empirical testing showed that just centring the cumulated GP (gp -= mean(gp)) fixes the (R0, drift) ridge that was the root cause of the stuck-chain and catastrophic R-hat issues. The previous version of this PR also bundled two other changes (rescaling of GP increments by 1/sqrt(gp_n) and a centred-form parameterisation of the GP coefficients with a much wider default alpha prior); those turn out not to be needed and brought substantial API impact for marginal gains in efficiency. Verified across the four previously-stuck seeds (2, 4, 8, 11) under OMP_NUM_THREADS=1 with the original default alpha = Normal(0, 0.01): - 0 treedepth hits, max 1 divergence, R-hat 1.004-1.008, ESS 569-842 - 2x more efficient per effective sample than the bundled version - No semantic change to alpha, no prior change Co-authored-by: sbfnk <sebastian.funk@lshtm.ac.uk>
|
Heads-up to anyone reviewing: this PR has been substantially simplified. Earlier versions bundled three changes (rescaling of GP increments by What's now in the PR:
What's no longer in the PR:
Empirically, on four seeds that were previously catastrophic (R-hat up to 6.10, hundreds of treedepth hits) under the existing
No API or interface change. |
estimate_infectionsTable: Benchmarking results (mean time in seconds). |operation | branch| main| % change| range| trend| |:--------------------|-------:|-------:|--------:|----------:|---------:| |delays | 6.7| 10| -35| (-48, -14)| speedup| |infections | 2.8| 4.2| -33| (-47, -9)| speedup| |reports | 2.7| 4.1| -34| (-48, -12)| speedup| |report lp | 2| 3.1| -34| (-48, -14)| speedup| |R0 | 1.3| 1.3| 4| (-17, 39)| no change| |update gp | 0.69| 1.1| -34| (-50, -9)| speedup| |gt | 0.31| 0.48| -34| (-49, -11)| speedup| |day of the week | 0.29| 0.43| -32| (-47, -9)| speedup| |truncation | 0.19| 0.29| -33| (-49, -10)| speedup| |param lp | 0.18| 0.27| -34| (-47, -7)| speedup| |truncate | 0.14| 0.22| -34| (-48, -12)| speedup| |delays lp | 0.075| 0.12| -35| (-51, -13)| speedup| |rt lp | 0.074| 0.12| -35| (-51, -9)| speedup| |gp lp | 0.067| 0.099| -32| (-50, -9)| speedup| |generated quantities | 0.0077| 0.008| -4| (-14, 12)| no change| |assign max | 5.8e-07| 5.7e-07| 11| (-50, 77)| no change| |
…ords stanc3 reserves "lower" and "upper" for variable declarations; using them as function-argument names is a syntax error from cmdstan 2.38 onwards. Rename the bound arguments to "lb" and "ub". The data-block names init_lower / init_upper / params_lower / params_upper are unaffected. Co-authored-by: sbfnk <sebastian.funk@lshtm.ac.uk>
|
This is how benchmark results would change (along with a 95% confidence interval in relative change) if e59ec01 is merged into main:
|
estimate_distTable: Benchmarking results (mean time in seconds).
|
|
This is how benchmark results would change (along with a 95% confidence interval in relative change) if 230a949 is merged into main:
|
The non-stationary GP construction for Rt is introduced in the 'Time-varying reproduction number' section of the estimate_infections vignette; add a link from the GP implementation details vignette where the same construction is described. Co-authored-by: sbfnk <sebastian.funk@lshtm.ac.uk>
Lead with 'makes inference difficult' before introducing the technical 'ridge in the joint posterior' phrasing. Inline the cross-link to the estimate_infections vignette on the descriptive phrase rather than a trailing 'see the X section' note. Co-authored-by: sbfnk <sebastian.funk@lshtm.ac.uk>
- Use overline{log R} (mean of log R) rather than log overline{R} (log of
mean of R), which is the quantity the internal parameter actually equals.
- Move the empirical-mean and centring-identity formulae onto their own
display equations rather than inlining them.
- Frame the centring fix via the regression-intercept identifiability
parallel and link the cross-references: centring predictors (Stan
user's guide), brms intercept parameterisation, change of variables
(Stan user's guide), the estimate_infections vignette section on the
first-difference GP.
- Spell "observation window" consistently and add "then" to the
prior-application sentence for flow.
Co-authored-by: sbfnk <sebastian.funk@lshtm.ac.uk>
The change of variables is univariate (bar{log R} -> log R_1, a shift by
-bar{GP}); the GP increments map to themselves identically and don't need
threading through the transform. Drops the multivariate tuple notation.
Co-authored-by: sbfnk <sebastian.funk@lshtm.ac.uk>
Replaces the two-step shift + log->natural framing with the single
bijective transform bar{log R} -> R_1 = exp(bar{log R} - bar{GP}), giving
the same log-Jacobian (log R_1) in one go.
Co-authored-by: sbfnk <sebastian.funk@lshtm.ac.uk>
- Show GP-bar in both the natural double-sum form and the equivalent
weighted-single-sum form, with a short note on the weights.
- Reword the centring identity from 'so that ... by construction' to
'where bar{log R} = ... by construction'; flip the equation so the
defined quantity is on the left.
Co-authored-by: sbfnk <sebastian.funk@lshtm.ac.uk>
Co-authored-by: sbfnk <sebastian.funk@lshtm.ac.uk>
Derive bar{GP} as the natural offset that arises when averaging log R_t
over the observation window:
bar{log R} = (1/T_obs) sum_t log R_t = log R_1 + bar{GP},
rather than introducing bar{GP} ex nihilo. This motivates sampling
bar{log R} (well-identified by the data) as the internal parameter and
recovering log R_1 = bar{log R} - bar{GP}.
Collapse the prior-application paragraph into a single partial derivative
(no chained shift + exp) and reword to "we need to add a Jacobian
correction of log R_1 to the target log density".
Fix a stray "centered" -> "centred" elsewhere in the vignette.
Co-authored-by: sbfnk <sebastian.funk@lshtm.ac.uk>
|
@sbfnk — thanks for walking me through the maths in the GP vignette. The derivation makes proper sense to me now (centring as a regression-intercept-style identifiability fix, with bar{GP} arising as the natural offset between log R_1 and the well-identified bar{log R}). |
|
It makes sense to me, too. |
|
This is how benchmark results would change (along with a 95% confidence interval in relative change) if 1d46ab6 is merged into main:
|
|
In principle we should be doing this for the rw as well but as it is fast already matters less |
3 participants
Description
Three changes to the Stan code that together eliminate stuck chains and divergences seen on the non-stationary Gaussian process used to model Rt over time. All in
inst/stan/:Rescale GP increments by
1/sqrt(gp_n)insideupdate_Rtsoalphacontrols the trajectory SD rather than the increment SD — matches whatgp_opts()docs already claim.Centre the cumulated GP (
gp -= mean(gp)) solog R0 = mean log Rtrather than the initial value. Eliminates the(R0, drift)ridge in the joint posterior.Centred-form GP coefficients: switch from
eta ~ N(0, 1)withnoise = PHI * (diagSPD .* eta)toeta ~ N(0, diagSPD)withnoise = PHI * eta. Avoids the(alpha, eta)funnel that arises in non-centred form when alpha is small (uninformative data regime), as recommended by the Stan manual.Empirical results
On a 15-seed sweep of
rescaled_ln0p3configuration underOMP_NUM_THREADS=1, four seeds had stuck chains and two of those were catastrophic (R-hat > 4). With the new parameterisation, all four cure to:Sampling is also ~4× faster on these seeds (median stepsize ~5× larger; HMC traverses the smoother geometry with bigger leaps).
Initial submission checklist
testthat::test_file("test-estimate_infections.R")withEPINOW2_SKIP_INTEGRATION=falsepasses; warnings on the twoget_predictionstests reflect their warmup=25/samples=25 budget, not a regression).gp_opts()doc already describesalphaas the trajectory SD, and that is now actually what it controls.Notes for reviewers
update_gpsignature is preserved (alpha/rho/M/L/type/nu still passed, now unused inside the function) so external callers do not break.gaussian_process_lpgains arguments (M, L, alpha, rho, type, nu) so it can computediagSPDfor the eta prior; the matching call site inestimate_infections.stanis updated accordingly.Summary by CodeRabbit
Release Notes