Skip to content

Adapt the HMC metric when the warmup is too short for the adaptation windows - #3441

Open
jachymb wants to merge 1 commit into
stan-dev:developfrom
jachymb:bugfix/adaptation-short-warmup-metric
Open

jachymb wants to merge 1 commit into
stan-dev:developfrom
jachymb:bugfix/adaptation-short-warmup-metric

Conversation

@jachymb

@jachymb jachymb commented Oct 17, 2026 •

Copy link
Copy Markdown

This offers a fix to an inconsistency between documented and actual behavior I noticed. Whether this is the best way to resolve the inconsistency is perhaps up to discussion because in some sense this is a tradeoff (see "Intended effect"). If you prefer another approach let me know.

AI use disclosure: This was developed with the help of Claude Code (Opus 5.5).

Submission Checklist

  • Run unit tests: ./runTests.py src/test/unit
  • Run cpplint: make cpplint
  • Declare copyright holder and open-source license: see below

Summary

When init_buffer + window + term_buffer > num_warmup (with the defaults: num_warmup from 20 to 149), windowed_adaptation::set_window_params reduces the stages to 15%/75%/10% but returns without calling restart(). adapt_next_window_ keeps the constructor's 0 + 0 - 1 (UINT_MAX), so the adaptation window never ends and the diag_e / dense_e metric is never updated. The fix calls restart() in that branch, as the other branch does.

Intended Effect

Warmup follows the documented schedule when the stages do not fit. The reference manual says they "are scaled down", and the warning lists the reduced windows; with the fix the sampler uses them. A run with the default windows now equals a run with the announced windows given explicitly, and the metric is adapted from num_warmup = 20, the documented lower limit, instead of only once the stages fit (150 with the defaults).

In the measurement under Side Effects this raises the minimum bulk ESS per gradient for num_warmup 30–149 by 3× (eight schools) to several hundred× (normals with sd e⁻³…e³). If maintainers prefer develop's actual behaviour (the initial metric, with the step size tuned over the whole warmup), the other way to make code and warning agree is to keep the code and change the warning to say that the metric is not adapted.

How to Verify

McmcVarAdaptation.learn_variance_reduced_windows runs num_warmup=100 with the default windows and expects the variance estimated from the 75 draws of the reduced window. It fails on develop, where the variance stays at its initial value. No existing expectation changes.

Side Effects

Draws change for adaptive diag_e / dense_e runs whose windows are reduced. Measured on CmdStan 2.39.0 (same sampler code as develop), the branch emulated by passing the reduced windows explicitly (same window state as the fix), 20 seeds on 5 configurations: normals with sd e⁻³…e³; logistic regression; correlated normal (diag and dense); eight schools. develop → this branch:

num_warmup worse paths
20, 25 sampling breaks: median accept_stat 0.002–0.47 (target 0.8), NaN ESS in some runs; divergences in 0–5 → 6–20 of 20 runs per configuration
30, 35 divergences in 0–3 → 0–15 of 20 runs per configuration
40, 45, 50 eight schools: divergences in 0–1 → 9–12 of 20 runs; others none
100, 149 eight schools: 5 → 6 and 8 → 2 runs; others none

At 20–50 the step size is retuned over only 2–5 iterations after the metric update, as with such windows given explicitly on develop. Pooled eight-schools estimates match a long reference run on both sides, but at num_warmup=30 the spread of per-run E[tau] grows from 0.15 to 0.41. A metric_file metric is now re-estimated in these runs; init_buffer=0 term_buffer=0 window=0 keeps it. I can share the scripts.

Documentation

None.

Copyright and Licensing

Please list the copyright holder for the work you are submitting (this will be you or your assignee, such as a university or company):

Jáchym Barvínek jachymb@gmail.com

By submitting this pull request, the copyright holder is agreeing to license the submitted work under the following licenses:

When init_buffer + window + term_buffer exceeds num_warmup,
set_window_params reduced the stages to 15%/75%/10% but returned without
restart(), so adapt_next_window_ kept the constructor's value
(0 + 0 - 1, i.e. UINT_MAX). The adaptation window never ended and the
diag_e / dense_e metric was never updated, although the warning says the
reduced windows are used. Call restart() in that branch as in the other.
@jachymb

jachymb commented Oct 18, 2026 •

Copy link
Copy Markdown
Author

So apparently an upstream test relies on the problematic behavior this is trying to address, so I think I'll just leave this here for someone else to consider.

I'm now kinda confused what the actually expected behavior is here.

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

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants