fix: correct biased pymc HMC sampling (target_accept)

sbi/inference/posteriors/mcmc_posterior.py

@@ -64,6 +64,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 +108,9 @@ 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 for the gradient-based
+                samplers (HMC/NUTS). See `MCMCPosteriorParameters` for details. If
+                `None`, HMC uses `0.9` and NUTS/Pyro keep their backend defaults.
         """
         if method == "slice":
             warn(
@@ -136,6 +140,7 @@ def __init__(
         self.init_strategy_parameters = init_strategy_parameters or {}
         self.num_workers = num_workers
         self.mp_context = mp_context
+        self.target_accept = target_accept
         self._posterior_sampler = None
 
         # Hardcode parameter name to reduce clutter kwargs.
@@ -257,6 +262,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 +291,9 @@ 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 for the gradient-based
+                samplers (HMC/NUTS). If not provided, uses the value specified at
+                initialization.
 
         Returns:
             Samples from posterior.
@@ -301,6 +310,7 @@ 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
+        target_accept = self.target_accept if target_accept is None else target_accept
         init_strategy_parameters = (
             self.init_strategy_parameters
             if init_strategy_parameters is None
@@ -342,6 +352,7 @@ def sample(
                     num_chains=num_chains,
                     show_progress_bars=show_progress_bars,
                     mp_context=mp_context,
+                    target_accept=target_accept,
                 )
             elif method in ("hmc_pymc", "nuts_pymc", "slice_pymc"):
                 transformed_samples = self._pymc_mcmc(
@@ -354,6 +365,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!")
@@ -794,6 +806,7 @@ def _pyro_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 Pyro's HMC or NUTS sampler.
 
@@ -809,6 +822,8 @@ def _pyro_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 step-size adaptation.
+                If None, Pyro's default is used.
 
         Returns:
             Tensor of shape (num_samples, shape_of_single_theta).
@@ -827,9 +842,12 @@ def _pyro_mcmc(
         thin = _process_thin_default(thin)
         num_chains = mp.cpu_count() - 1 if num_chains is None else num_chains
         kernels = dict(hmc_pyro=HMC, nuts_pyro=NUTS)
+        kernel_kwargs = {"potential_fn": potential_function}
+        if target_accept is not None:
+            kernel_kwargs["target_accept_prob"] = target_accept
 
         sampler = MCMC(
-            kernel=kernels[mcmc_method](potential_fn=potential_function),
+            kernel=kernels[mcmc_method](**kernel_kwargs),
             num_samples=ceil((thin * num_samples) / num_chains),
             warmup_steps=warmup_steps,
             initial_params={self.param_name: initial_params},
@@ -862,6 +880,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 +896,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).
@@ -894,10 +916,14 @@ def _pymc_mcmc(
         thin = _process_thin_default(thin)
         num_chains = mp.cpu_count() - 1 if num_chains is None else num_chains
         steps = dict(slice_pymc="slice", hmc_pymc="hmc", nuts_pymc="nuts")
+        step = steps[mcmc_method]
+
+        if target_accept is None and step == "hmc":
+            target_accept = 0.9
 
         sampler = PyMCSampler(
             potential_fn=potential_function,
-            step=steps[mcmc_method],
+            step=step,
             initvals=tensor2numpy(initial_params),
             draws=ceil((thin * num_samples) / num_chains),
             tune=warmup_steps,
@@ -906,6 +932,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

@@ -244,6 +244,13 @@ 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 for the gradient-based
+            samplers (`hmc_pymc`, `nuts_pymc`, `hmc_pyro`, `nuts_pyro`), controlling
+            the step-size adaptation during warmup. Higher values take smaller,
+            more accurate steps at the cost of speed. Ignored by the slice samplers.
+            If `None`, HMC uses `0.9` (PyMC's default of `0.65` mixes poorly on
+            peaked targets and yields biased samples), while NUTS and the Pyro
+            samplers keep their respective backend defaults.
     """
 
     method: Literal[
@@ -262,10 +269,14 @@ 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."""
 
+        if self.target_accept is not None and not (0.0 < self.target_accept < 1.0):
+            raise ValueError("target_accept must be between 0 and 1, or None.")
+
         if not (
             self.init_strategy_parameters is None
             or isinstance(self.init_strategy_parameters, Dict)

sbi/samplers/mcmc/pymc_wrapper.py

@@ -114,6 +114,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 +130,13 @@ 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 (default), PyMC's own default is
+                used. PyMC's HMC default (0.65) mixes poorly on peaked targets;
+                a higher value (e.g. 0.9-0.95) gives markedly less biased samples.
+                Ignored for the `"slice"` step.
         """
         self.param_name = param_name
         self._step = step
@@ -138,6 +147,8 @@ def __init__(
         self._mp_ctx = mp_ctx
         self._progressbar = progressbar
         self._device = device
+        self._seed = seed
+        self._target_accept = target_accept
 
         # create PyMC model object
         track_gradients = step in ("nuts", "hmc")
@@ -157,15 +168,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

tests/mcmc_test.py

@@ -158,6 +158,10 @@ 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
+        # PyMC's HMC default (target_accept=0.65) is biased on this peaked
+        # target (c2st ~0.63); 0.95 samples it accurately. See #1894 follow-up.
+        target_accept=0.95,
     )
     samples = sampler.run()
     assert samples.shape == (
@@ -173,6 +177,36 @@ def lp_f(x, track_gradients=True):
     check_c2st(samples, target_samples, alg=alg)
 
 
+@pytest.mark.mcmc
+@pytest.mark.parametrize(
+    ("method", "target_accept", "expected"),
+    [
+        ("hmc_pymc", None, 0.9),  # HMC default is raised (PyMC's 0.65 is biased)
+        ("hmc_pymc", 0.99, 0.99),  # user override wins
+        ("nuts_pymc", None, None),  # NUTS keeps PyMC's own default
+    ],
+)
+def test_pymc_target_accept_forwarded(method, target_accept, expected):
+    """target_accept from the user reaches the PyMC sampler with the right default."""
+    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)
+
+    mcmc_posterior = build_from_potential(potential_fn, prior)
+    mcmc_posterior.sample(
+        (10,),
+        method=method,
+        target_accept=target_accept,
+        num_chains=1,
+        warmup_steps=10,
+        thin=1,
+        show_progress_bars=False,
+    )
+    assert mcmc_posterior._posterior_sampler._target_accept == expected
+
+
 @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",
     [