fix: correct biased pymc HMC sampling (target_accept)#1908
Conversation
The hmc parametrizations of test_c2st_pymc_sampler_on_Gaussian were failing in CD (c2st ~0.62). Root cause is PyMC's HamiltonianMC default target_accept=0.65, which mixes poorly on this peaked target regardless of seed, warmup, or number of draws (NUTS is unaffected). Expose seed and target_accept on PyMCSampler and set target_accept=0.95 (c2st ~0.51-0.54 across seeds and chain counts) plus a fixed seed in the test for reproducibility.
MCMCPosterior._pymc_mcmc left PyMC's HamiltonianMC at its default target_accept=0.65, which produces biased hmc_pymc posteriors (c2st ~0.62 on a Gaussian). Default it to 0.9 for the HMC step, leaving nuts_pymc and slice_pymc unchanged. hmc_pymc now samples accurately (c2st ~0.51).
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1908 +/- ##
==========================================
+ Coverage 87.81% 87.94% +0.12%
==========================================
Files 143 143
Lines 13387 13657 +270
==========================================
+ Hits 11756 12010 +254
- Misses 1631 1647 +16
Flags with carried forward coverage won't be shown. Click here to find out more.
|
|
thanks for the review @dgedon when addressing your comments I noticed there was an inconsistency in the default values (0.9 vs 0.95 in the tests), that's fixed now. And I noticed it would be better to have the new |
dgedon
left a comment
dgedon
left a comment
There was a problem hiding this comment.
Maybe update the original description of the PR, so it reflects the latest changes, e.g. the 0.9 default.
I approve for now since all remaining comments are smaller details and not functional really. In that way you can merge easier @janfb
| target_accept: Target acceptance probability for the gradient-based | ||
| samplers (HMC/NUTS). See `MCMCPosteriorParameters` for details. If | ||
| `None`, HMC uses `0.9` and NUTS/Pyro keep their backend defaults. | ||
| target_accept: Target acceptance probability for PyMC's HMC and NUTS |
There was a problem hiding this comment.
Is this always PyMC based or is it also used for exacmple for pyro?
There was a problem hiding this comment.
Ah I see that you actively removed it from pyro. We should clarify then that this argument is only used for PyMC
| target_accept: Target acceptance probability for the gradient-based | ||
| samplers (HMC/NUTS). If not provided, uses the value specified at | ||
| initialization. | ||
| target_accept: Target acceptance probability for PyMC's HMC and NUTS |
There was a problem hiding this comment.
Same here with pymc specific comments
| range_inclusive=False, | ||
| ) | ||
| self._target_accept = ( | ||
| 0.9 if target_accept is None and step == "hmc" else target_accept |
There was a problem hiding this comment.
Maybe instead of hard coding it as a number, we should define something like _DEFAULT_HMC_TARGET_ACCEPT = 0.9
| validate_float_range( | ||
| target_accept, | ||
| "target_accept", | ||
| min_val=0.0, | ||
| max_val=1.0, | ||
| range_inclusive=False, | ||
| ) |
There was a problem hiding this comment.
I don't see a better options but worth a short comment: we now have these exact lines a few times in different places within this PR. Maybe there is a reasonable way to consolidate.
2 participants
Follow-up to #1894 because CD is failing again. Once the lockfile fix let the pymc tests run again in CD,
test_c2st_pymc_sampler_on_Gaussianfailed on thehmcparametrizations with c2st around 0.62.However, it's not flaky: pymc's
HamiltonianMCdefaults totarget_accept=0.65, which mixes poorly on this peaked Gaussian and biases the samples, regardless of seed, warmup, or number of draws. NUTS works fine.To fix this:
seedandtarget_acceptonPyMCSamplertarget_accept=0.9for thehmc_pymcstep inMCMCPosteriorto get accurate HMC posteriors (c2st ~0.51 vs ~0.62);nuts_pymc/slice_pymcunchangedtarget_accept=0.95in the c2st test for reproducible samplingslicedoes not taketarget_accept, so it is only applied to hmc/nuts.Note: the c2st test is marked
slow, I will run the CD on this branch via workflow dispatch.