Warmup Schedule
bgms uses a multi-stage warmup schedule to tune MCMC sampler parameters before collecting posterior samples. The three core stages (1, 2, 3a) follow Stan’s warmup scheme for NUTS. Stages 3b and 3c are bgms extensions that accommodate the spike-and-slab edge selection machinery. The schedule is defined in src/mcmc/execution/warmup_schedule.h, and its staging depends on which sampler hosts it (see Budget allocation).
The stages
| Stage | Name | Purpose |
|---|---|---|
| 1 | Initial step-size adaptation | Tune step size with identity mass matrix |
| 2 | Adaptation windows | Tune step size (dual averaging) and mass matrix (Welford) |
| 3a | Terminal step-size tuning | Final step-size pass with frozen mass matrix |
| 3b | Proposal-SD tuning | Tune element-wise Metropolis proposal SDs (edge selection not yet active) |
| 3c | Selection-active warmup | Warmup iterations with edge selection running |
Stages 3b and 3c only apply when edge selection is enabled, and which of them actually run depends on the sampler.
Stage 1: Initial step-size adaptation
A short burn-in phase (default: the first 75 iterations). The mass matrix is held at the identity. Dual averaging tunes the step size toward a target acceptance rate of 0.8. The starting value of \(\epsilon\) comes from Algorithm 4 in Hoffman & Gelman (2014), which doubles or halves a trial step size until the Metropolis acceptance ratio of a single leapfrog step crosses the target acceptance rate.
Stage 2: Adaptation windows
The step size and mass matrix are tuned in expanding windows. The first window has 25 iterations. Each subsequent window doubles in size (25 → 50 → 100 → …) until the Stage 3a boundary is reached; following Stan’s windowed_adaptation, a final window that would overshoot the boundary is stretched to absorb the remaining budget rather than emitted as a small trailing window (each window end triggers a mass-matrix update and step-size reinitialization, and an extra small window would disrupt dual averaging).
Within each window, a Welford running variance accumulates the sampled parameter vectors. At the end of each window, the diagonal mass matrix is set to
\[ \hat{M}^{-1}_j = \frac{n}{n + w}\,\hat{\sigma}^2_j + \frac{w}{n + w}\,\sigma^2_0, \]
blending the empirical variance \(\hat{\sigma}^2_j\) (from \(n\) samples) with a prior variance \(\sigma^2_0 = 10^{-3}\) at weight \(w = 5\). After each mass-matrix update:
- The step-size heuristic reruns with the new mass matrix
- Dual averaging restarts
- The Welford accumulator resets for the next window
Under edge selection, the mass matrix is accumulated on the full (zero-padded) parameter layout so entries keep their slots when the active set changes; see NUTS — Mass matrix.
Stage 3a: Terminal step-size tuning
The mass matrix is frozen. Dual averaging continues to give the step size a final tuning pass against the converged mass matrix. At the end of this stage, the step size is fixed at the dual-averaging estimate.
Stage 3b: Proposal-SD tuning
Step-size adaptation is frozen. The proposal standard deviations of the edge-move Metropolis proposals (and any other element-wise proposals) are tuned via additive Robbins–Monro updates (see Adaptive Metropolis):
\[ \sigma_{k+1} = \sigma_k + w_k (\alpha_k - \alpha^*) \]
where \(\alpha_k\) is the observed acceptance probability, \(\alpha^*
= 0.44\) is the target, and \(w_k = t^{-0.75}\) is the decay weight (\(t\) counts iterations since stage 3b began). The result is clamped to \([0.001, 2.0]\), and NaN values are reset to 1.0. The weight comes from a single source of truth, rm_weight_for_proposal_sd(iter) on the schedule, which returns a value only on iterations where adaptation should run — every model’s tune_proposal_sd consults it rather than computing the weight inline. The 0.44 target is fixed under NUTS (the user’s target_accept argument governs the step-size adaptation target instead); only under update_method = "adaptive-metropolis" does target_accept set the Metropolis target directly. Edge selection is not yet active in this stage.
Stage 3c: Selection-active warmup
Edge selection is activated. Edge indicators \(\gamma_{ij}\) are updated every iteration, changing the active parameter space. On the NUTS host, dual averaging restarts to re-tune \(\epsilon\) for the new geometry. On the Gibbs host, stage 3c covers most of the warmup (see below) — there is nothing to adapt, but the graph and precision matrix burn in together.
Budget allocation
How the warmup budget is divided depends on the sampler hosting the schedule (the SamplerSpec fields learn_sd and the Gibbs select_during_warmup flag; see Sampler Hierarchy):
- NUTS with edge selection (
edge_selectionwithlearn_sd) — split 85% (stages 1-3a) / 10% (stage 3b) / ~5% (stage 3c). If stage 3b would get fewer than 20 iterations it is skipped entirely (default proposal SDs are used). - Gibbs with edge selection (
edge_selectionwithselect_during_warmup) — there is no adaptation to run, so the schedule reduces to a settle-then-select split: the first 15% of warmup runs the full model (all edges included, selection off) so the precision matrix settles, and stage 3c covers the remaining 85% with selection active. Both windows scale with the warmup budget. - Adaptive Metropolis with edge selection (
edge_selectionwith both flags off) — stage 3c is empty; selection activates at the first sampling iteration. - No edge selection (any sampler; both flags irrelevant) — the full budget goes to the core stages 1-3a.
A warm-started refit is the one case that legitimately runs a short warmup: it inherits an adapted step size and mass matrix, so the expensive part of stages 1-3a is already paid for. See Sensitivity and Refits for the schedule that applies and the conditions under which it does.
Warning system
Short-warmup warnings are the responsibility of validate_sampler() on the R side, issued at fit time (NUTS only): with edge selection, thresholds at warmup < 50 (very short), < 200 (proposal-SD tuning skipped), and < 300 (limited tuning); without, at < 20 (no mass-matrix estimation) and < 150 (proportional allocation). The schedule itself records only stage3b_skipped, which its own stage predicates consult. After sampling there is additionally the energy-based warmup-completeness check in fit$nuts_diag (warmup_incomplete; see Diagnostics).
Dual averaging
The step-size update for the NUTS sampler follows Algorithm 6 in Hoffman & Gelman (2014) with parameters \(\gamma = 0.05\), \(t_0 = 10\), \(\kappa = 0.75\), and bias target \(\mu = \log(10\,\epsilon_0)\). (Here \(\gamma\) and \(\mu\) are the Hoffman–Gelman tuning constants, not the edge indicators \(\gamma_{ij}\) or category thresholds \(\mu_{ic}\) used elsewhere on this site.)
Querying the schedule
The WarmupSchedule struct exposes predicate methods used by the chain runner and samplers:
bool in_stage1(int i) const;
bool in_stage2(int i) const;
bool in_stage3a(int i) const;
bool in_stage3b(int i) const;
bool in_stage3c(int i) const;
bool sampling(int i) const; // i >= total_warmup
bool selection_enabled(int i) const; // stage 3c or sampling
std::optional<double> rm_weight_for_proposal_sd(int i) const;These predicates control when the adaptation controller accumulates samples, when edge selection is turned on, and when the chain transitions from warmup to sampling. selection_enabled() is true during stage 3c and throughout sampling — under the adaptive-Metropolis host, where stage 3c is empty, that means selection starts only when sampling does.