如何在PyMC3中禁用NUTS调参以实现MCMC方法公平对比

Great question—totally get wanting a strict apples-to-apples comparison between your custom MCMC method and the original NUTS implementation, without PyMC3's extra tuning features muddling the results. Here's exactly how to dial back PyMC3's defaults to match Algorithm 6: Efficient No-U-Turn Sampler with Dual Averaging from Hoffman & Gelman's 2014 paper:

Key Adjustments to Disable Extra Tuning #

PyMC3's default NUTS includes adaptive mass matrix estimation (a post-paper addition) and other minor tweaks. To stick to the standard algorithm, focus on these settings:

1. Match Algorithm 6 (With Dual Averaging Step Size Adaptation) #

The original Algorithm 6 includes dual averaging for step size adaptation—this is part of the standard method, so we keep it, but disable the non-standard mass matrix adaptation:

import pymc3 as pm

with pm.Model() as your_model:
    # Define your model variables/priors/likelihood here...
    
    # Initialize standard NUTS step
    standard_nuts = pm.NUTS(
        adapt_step_size=True,  # Keep dual averaging (core to Algorithm 6)
        adapt_mass_matrix=False,  # Disable non-standard mass matrix adaptation
        target_accept=0.8,  # Original paper's target acceptance rate
        max_treedepth=10,  # Default from the paper
        # Optional: Explicitly set dual averaging params to match the paper
        gamma=0.05,
        k=0.75,
        t0=10
    )
    
    # Run sampling: Tune phase is for step size adaptation (Algorithm 6 requirement)
    trace = pm.sample(
        draws=2000,  # Number of post-tune samples to collect
        tune=1000,  # Step size adaptation runs during this phase
        step=standard_nuts,
        chains=4,
        cores=4,
        discard_tuned_samples=True  # Only keep samples after adaptation is done
    )

2. Fully Disable All Adaptation (Fixed Step Size/Unit Mass Matrix) #

If you need a completely "static" NUTS run (no adaptive steps at all) for your comparison, disable both step size and mass matrix adaptation, and manually set a step size:

with pm.Model() as your_model:
    # Define your model here...
    
    # First, find a reasonable fixed step size (run a quick adaptive NUTS first!)
    # temp_trace = pm.sample(tune=1000, draws=1, chains=1)
    # fixed_step_size = temp_trace.get_sampler_stats("step_size")[0][-1]
    
    fixed_step_size = 0.1  # Replace with a value tuned to your model
    static_nuts = pm.NUTS(
        adapt_step_size=False,
        adapt_mass_matrix=False,
        step_size=fixed_step_size,
        target_accept=0.8,
        max_treedepth=10
    )
    
    # No tune phase needed—start sampling immediately
    trace = pm.sample(
        draws=2000,
        tune=0,
        step=static_nuts,
        chains=4,
        cores=4
    )

Critical Notes for Fair Comparison #

• Consistent Initialization: Make sure both your custom method and NUTS start from the exact same parameter values. You can use pm.sample(start=your_initial_params) to enforce this.
• Tune Step Size First: If using fixed step size, always run a short adaptive NUTS run first to get a reasonable step size—randomly picking a value will likely lead to poor performance.
• Verify Settings: Check the sampler stats with trace.get_sampler_stats() to confirm mass matrix isn't being adapted (look for mass_matrix staying as the identity matrix) and step size behavior matches your setup.

内容的提问来源于stack exchange，提问作者KiaSh