fix: correct biased pymc HMC sampling (target_accept)

CHANGELOG.md

@@ -12,6 +12,7 @@
 
 ### 🐛 Bug Fixes
 
+* **Improve PyMC HMC sampling** ([#1908](https://github.com/sbi-dev/sbi/pull/1908)): Use a target acceptance probability of `0.9` by default for `hmc_pymc`, avoiding poor finite-chain mixing observed with PyMC's `0.65` default. The value can be customized with `MCMCPosteriorParameters.target_accept` or `MCMCPosterior.sample(target_accept=...)`.
 * **Fix TARP z-scoring bug** ([#1832](https://github.com/sbi-dev/sbi/issues/1832)): Reference points are now z-scored alongside `thetas` and `posterior_samples` when `z_score_theta=True`, fixing incorrect distance calculations that masked bias detection.
 * **Fix broken `biased_toy_gaussian` test helper**: Rewrote to create actual location bias (posterior mean shifted from truth) instead of the previous NaN-producing formula.
 * **Raise on unsupported `transform_to_unconstrained` z-scoring**: `z_score_x="transform_to_unconstrained"` was silently ignored by the nflows builders (`maf`, `nsf`, `maf_rqs`, `made`), the MDN builder, the unconditional Zuko builder, the ratio-based classifier builders (`linear`, `mlp`, `resnet`), and the vector field builders (FMPE / NPSE), producing a model with no reparametrization. These builders now raise a clear `ValueError`. The mixed builders (MNLE / MNPE) raise transitively when their continuous flow is a non-Zuko model. The option remains supported by the conditional `zuko_*` models.

sbi/inference/posteriors/mcmc_posterior.py

@@ -31,6 +31,7 @@
 from sbi.utils import mcmc_transform
 from sbi.utils.potentialutils import pyro_potential_wrapper, transformed_potential
 from sbi.utils.torchutils import ensure_theta_batched, tensor2numpy
+from sbi.utils.typechecks import validate_target_accept
 
 
 class MCMCPosterior(NeuralPosterior):
@@ -64,6 +65,7 @@ def __init__(
         mp_context: Literal["fork", "spawn"] = "spawn",
         device: Optional[Union[str, torch.device]] = None,
         x_shape: Optional[torch.Size] = None,
+        target_accept: Optional[float] = None,
     ):
         """
         Args:
@@ -107,6 +109,11 @@ def __init__(
             device: Training device, e.g., "cpu", "cuda" or "cuda:0". If None,
                 `potential_fn.device` is used.
             x_shape: Deprecated, should not be passed.
+            target_accept: Target acceptance probability used only by the PyMC
+                samplers `hmc_pymc` and `nuts_pymc`; it is ignored by `slice_pymc`,
+                the Pyro samplers (`hmc_pyro`, `nuts_pyro`) and the numpy slice
+                samplers. If `None`, `hmc_pymc` uses `0.9` and `nuts_pymc` keeps
+                PyMC's backend default. See `MCMCPosteriorParameters` for details.
         """
         if method == "slice":
             warn(
@@ -136,6 +143,8 @@ def __init__(
         self.init_strategy_parameters = init_strategy_parameters or {}
         self.num_workers = num_workers
         self.mp_context = mp_context
+        validate_target_accept(target_accept)
+        self.target_accept = target_accept
         self._posterior_sampler = None
 
         # Hardcode parameter name to reduce clutter kwargs.
@@ -257,6 +266,7 @@ def sample(
         num_workers: Optional[int] = None,
         mp_context: Optional[str] = None,
         show_progress_bars: bool = True,
+        target_accept: Optional[float] = None,
     ) -> Tensor:
         r"""Draw samples from the approximate posterior distribution $p(\theta|x)$.
 
@@ -285,6 +295,10 @@ def sample(
             mp_context: Multiprocessing context (`fork` or `spawn`). If not provided,
                 uses the value specified at initialization.
             show_progress_bars: Whether to show sampling progress monitor.
+            target_accept: Target acceptance probability used only by the PyMC
+                samplers `hmc_pymc` and `nuts_pymc`; it is ignored by `slice_pymc`,
+                the Pyro samplers (`hmc_pyro`, `nuts_pyro`) and the numpy slice
+                samplers. If not provided, uses the value specified at initialization.
 
         Returns:
             Samples from posterior.
@@ -301,6 +315,10 @@ def sample(
         init_strategy = self.init_strategy if init_strategy is None else init_strategy
         num_workers = self.num_workers if num_workers is None else num_workers
         mp_context = self.mp_context if mp_context is None else mp_context
+        if target_accept is None:
+            target_accept = self.target_accept  # already validated in `__init__`.
+        else:
+            validate_target_accept(target_accept)
         init_strategy_parameters = (
             self.init_strategy_parameters
             if init_strategy_parameters is None
@@ -354,6 +372,7 @@ def sample(
                     num_chains=num_chains,
                     show_progress_bars=show_progress_bars,
                     mp_context=mp_context,
+                    target_accept=target_accept,
                 )
             else:
                 raise NameError(f"The sampling method {method} is not implemented!")
@@ -862,6 +881,7 @@ def _pymc_mcmc(
         num_chains: Optional[int] = 1,
         show_progress_bars: bool = True,
         mp_context: str = "spawn",
+        target_accept: Optional[float] = None,
     ) -> Tensor:
         r"""Return samples obtained using PyMC's HMC, NUTS or slice samplers.
 
@@ -877,6 +897,9 @@ def _pymc_mcmc(
             warmup_steps: Initial number of samples to discard.
             num_chains: Whether to sample in parallel. If None, use all but one CPU.
             show_progress_bars: Whether to show a progressbar during sampling.
+            target_accept: Target acceptance probability for HMC/NUTS. If `None`,
+                HMC uses `0.9` and NUTS keeps PyMC's default. Ignored by the slice
+                sampler.
 
         Returns:
             Tensor of shape (num_samples, shape_of_single_theta).
@@ -906,6 +929,7 @@ def _pymc_mcmc(
             progressbar=show_progress_bars,
             param_name=self.param_name,
             device=self._device,
+            target_accept=target_accept,
         )
         samples = sampler.run()
         samples = torch.from_numpy(samples).to(dtype=torch.float32, device=self._device)

sbi/inference/posteriors/posterior_parameters.py

@@ -23,6 +23,7 @@
     is_nonnegative_int,
     is_positive_float,
     is_positive_int,
+    validate_target_accept,
 )
 
 
@@ -244,6 +245,14 @@ class MCMCPosteriorParameters(PosteriorParameters):
             (default), used by Pyro and PyMC samplers. `"fork"` can be significantly
             faster than `"spawn"` but is only supported on POSIX-based systems
             (e.g. Linux and macOS, not Windows).
+        target_accept: Target acceptance probability used only by the PyMC samplers
+            `hmc_pymc` and `nuts_pymc`, controlling step-size adaptation during
+            warmup. Higher values generally result in smaller steps and fewer
+            rejections, at the cost of speed. If `None`, `hmc_pymc` uses `0.9`
+            because PyMC's `0.65` default mixed poorly in sbi's Gaussian regression
+            test, and `nuts_pymc` keeps PyMC's backend default. Ignored by
+            `slice_pymc`, the Pyro samplers (`hmc_pyro`, `nuts_pyro`) and the numpy
+            slice samplers.
     """
 
     method: Literal[
@@ -262,10 +271,13 @@ class MCMCPosteriorParameters(PosteriorParameters):
     init_strategy_parameters: Optional[Dict[str, Any]] = None
     num_workers: int = 1
     mp_context: Literal["fork", "spawn"] = "spawn"
+    target_accept: Optional[float] = None
 
     def validate(self):
         """Validate MCMCPosteriorParameters fields."""
 
+        validate_target_accept(self.target_accept)
+
         if not (
             self.init_strategy_parameters is None
             or isinstance(self.init_strategy_parameters, Dict)

sbi/samplers/mcmc/pymc_wrapper.py

@@ -9,6 +9,12 @@
 import torch
 
 from sbi.utils.torchutils import tensor2numpy
+from sbi.utils.typechecks import validate_target_accept
+
+# PyMC's `HamiltonianMC` defaults to `target_accept=0.65`, which mixes poorly on
+# peaked posteriors and biases the samples (see sbi #1908). NUTS is unaffected and
+# keeps PyMC's own default.
+_DEFAULT_HMC_TARGET_ACCEPT = 0.9
 
 
 class PyMCPotential(pt.Op):  # type: ignore
@@ -114,6 +120,8 @@ def __init__(
         progressbar: bool = True,
         param_name: str = "theta",
         device: str = "cpu",
+        seed: Optional[int] = None,
+        target_accept: Optional[float] = None,
     ):
         """Interface for PyMC samplers
 
@@ -128,6 +136,14 @@ def __init__(
             progressbar: Whether to show/hide progress bars.
             param_name: Name for parameter variable, for PyMC and ArviZ structures
             device: The device to which to move the parameters for potential_fn.
+            seed: Random seed passed to `pymc.sample` for reproducible sampling.
+                If None (default), PyMC seeds from system entropy.
+            target_accept: Target acceptance probability for the `"hmc"` and
+                `"nuts"` step methods. If `None`, HMC uses sbi's default of `0.9`
+                (`_DEFAULT_HMC_TARGET_ACCEPT`), while NUTS keeps PyMC's backend
+                default. The HMC default avoids poor finite-chain mixing observed
+                with PyMC's `0.65` default in sbi's Gaussian regression test.
+                Ignored for the `"slice"` step.
         """
         self.param_name = param_name
         self._step = step
@@ -138,6 +154,13 @@ def __init__(
         self._mp_ctx = mp_ctx
         self._progressbar = progressbar
         self._device = device
+        self._seed = seed
+        validate_target_accept(target_accept)
+        self._target_accept = (
+            _DEFAULT_HMC_TARGET_ACCEPT
+            if target_accept is None and step == "hmc"
+            else target_accept
+        )
 
         # create PyMC model object
         track_gradients = step in ("nuts", "hmc")
@@ -157,15 +180,21 @@ def run(self) -> np.ndarray:
             MCMC samples
         """
         step_class = dict(slice=pymc.Slice, hmc=pymc.HamiltonianMC, nuts=pymc.NUTS)
+        # target_accept only applies to the gradient-based samplers; Slice
+        # does not take it.
+        step_kwargs = {}
+        if self._target_accept is not None and self._step in ("hmc", "nuts"):
+            step_kwargs["target_accept"] = self._target_accept
         with self._model:
             inference_data = pymc.sample(
-                step=step_class[self._step](),
+                step=step_class[self._step](**step_kwargs),
                 tune=self._tune,
                 draws=self._draws,
                 initvals=self._initvals,  # type: ignore
                 chains=self._chains,
                 progressbar=self._progressbar,
                 mp_ctx=self._mp_ctx,
+                random_seed=self._seed,
             )
         self._inference_data = inference_data
         traces = inference_data.posterior  # type: ignore

sbi/utils/typechecks.py

@@ -100,3 +100,19 @@ def validate_float_range(
             raise ValueError(
                 f"{field_name} must be strictly between {min_val} and {max_val}."
             )
+
+
+def validate_target_accept(x: Any) -> None:
+    """Validate that x is `None` or a target acceptance probability in (0, 1).
+
+    `target_accept` is validated at several public boundaries (`MCMCPosterior`,
+    `MCMCPosteriorParameters`, `PyMCSampler`), so the policy lives here.
+
+    Args:
+        x: The value to validate. `None` means "use the sampler default".
+    """
+
+    if x is not None:
+        validate_float_range(
+            x, "target_accept", min_val=0.0, max_val=1.0, range_inclusive=False
+        )

tests/mcmc_test.py

@@ -158,6 +158,7 @@ def lp_f(x, track_gradients=True):
         draws=(int(num_samples / num_chains)),  # PyMC does not use thinning
         tune=warmup,
         chains=num_chains,
+        seed=0,  # reproducible sampling; PyMC is not covered by the global seed
     )
     samples = sampler.run()
     assert samples.shape == (
@@ -173,6 +174,53 @@ def lp_f(x, track_gradients=True):
     check_c2st(samples, target_samples, alg=alg)
 
 
+@pytest.mark.mcmc
+@pytest.mark.parametrize(
+    ("step", "target_accept", "expected"),
+    [
+        # 0.9 is spelled out rather than imported from `_DEFAULT_HMC_TARGET_ACCEPT` so
+        # that changing the documented default has to be a deliberate test update.
+        ("hmc", None, 0.9),
+        ("hmc", 0.99, 0.99),
+        ("nuts", None, None),
+    ],
+)
+def test_pymc_target_accept_default_and_override(step, target_accept, expected):
+    """PyMC HMC uses sbi's default and explicit user values take precedence."""
+    from sbi.samplers.mcmc.pymc_wrapper import PyMCSampler
+
+    theta_dim = 2
+
+    def potential_fn(theta):
+        return -0.5 * (theta**2).sum(axis=-1)
+
+    sampler = PyMCSampler(
+        potential_fn=potential_fn,
+        initvals=np.zeros((1, theta_dim)),
+        step=step,
+        target_accept=target_accept,
+        draws=1,
+        tune=1,
+        chains=1,
+    )
+    assert sampler._target_accept == expected
+
+
+@pytest.mark.mcmc
+def test_mcmc_posterior_rejects_invalid_target_accept():
+    """Per-call overrides validate target_accept at the public API boundary."""
+    theta_dim = 2
+    prior = BoxUniform(low=-2 * ones(theta_dim), high=2 * ones(theta_dim))
+
+    def potential_fn(theta):
+        return -0.5 * (theta**2).sum(axis=-1)
+
+    with pytest.warns(UserWarning, match="unconditional potential"):
+        posterior = build_from_potential(potential_fn, prior)
+    with pytest.raises(ValueError, match="target_accept"):
+        posterior.sample((1,), method="hmc_pymc", target_accept=0.0)
+
+
 @pytest.mark.mcmc
 def test_direct_mcmc_unconditional():
     "Test MCMCPosterior from user defined potential (unconditional)"

tests/posterior_parameters_test.py

@@ -396,6 +396,20 @@ def test_invalid_literal_field_values():
     MCMCPosteriorParameters(method="invalid")
 
 
+@pytest.mark.parametrize("target_accept", [None, 0.5, 0.9, 0.99])
+def test_mcmc_target_accept_accepts_valid_values(target_accept):
+    """target_accept accepts None or a probability in (0, 1)."""
+    params = MCMCPosteriorParameters(target_accept=target_accept)
+    assert params.target_accept == target_accept
+
+
+@pytest.mark.parametrize("target_accept", [0.0, 1.0, 1.5, -0.1])
+def test_mcmc_target_accept_rejects_invalid_values(target_accept):
+    """target_accept outside (0, 1) is rejected."""
+    with pytest.raises(ValueError, match="target_accept"):
+        MCMCPosteriorParameters(target_accept=target_accept)
+
+
 @pytest.mark.parametrize(
     "params",
     [