Pith. sign in

REVIEW 5 major objections 5 minor 1 cited by

Harnessing the Power of Reinforcement Learning for Adaptive MCMC

T0 review · 5 major / 5 minor · reviewed 2026-08-06 · deepseek-v4-flash

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗

Pith's one-line read Reinforcement learning can tune a gradient-based MCMC sampler's position-dependent step size, beating an optimised constant step size in 89% of benchmark tasks.

desk verdict Useful new reward and honest empirical study of RL-tuned RMALA, but the headline win rate rests on a confounded and under-analysed comparison. read the letter →

arxiv 2507.00671 v1 pith:AQUP24K4 submitted 2025-07-01 stat.CO cs.LGstat.ML

Congye Wang , Matthew A. Fisher , Heishiro Kanagawa , Wilson Chen , Chris. J. Oates This is my paper
classification stat.COcs.LGstat.ML MSC 65C0562F1560J22
keywords reinforcementlearningadaptiveMarkovchainMonteCarloMetropolis-adjustedLangevinalgorithmcontrastivedivergenceposition-dependentstepsizedeepdeterministicpolicygradientBayesianinferenceposteriorsampling
verification ladder T0 review T1 audit T2 compute T3 formal

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

This paper argues that the parameters of a gradient-based Markov chain Monte Carlo sampler can be tuned by reinforcement learning, and that the choice of reward function is what makes such tuning work. It shows that two natural rewards—the acceptance rate and the expected squared jump distance—give too little signal when sampling with a Metropolis–Hastings kernel that is cast as a decision process, and it proposes a contrastive-divergence lower bound (CDLB) reward that uses only quantities available inside a single transition. On a benchmark of 44 real posterior distributions, a position-dependent step size trained with this reward beats an optimally tuned constant step size in 89% of tasks. The paper also reports that CDLB avoids catastrophic training failures observed with an earlier reward in 7% of tasks. If correct, this turns a substantial part of MCMC tuning into a learning problem that does not require second derivatives of the target.

What carries the argument

The mechanism that carries the argument is the RLMH Markov decision process, in which the state is the pair $(X_n, X^*_{n+1})$ of current and proposed states and the action is the pair of step sizes $(\epsilon_\theta(X_n), \epsilon_\theta(X^*_{n+1}))$. The load-bearing identity is Proposition 1, a lower bound on the contrastive divergence $D_{\mathrm{KL}}(P_n\|P) - D_{\mathrm{KL}}(P_{n+1}\|P)$ in terms of an expected log-ratio plus an expected transition entropy; Proposition 2 then bounds the exploration term using only the acceptance probability and the proposal density. Combining these gives the CDLB reward in equation (11), which can be computed from data available inside a single Metropolis–Hastings step. The policy is trained with deep deterministic policy gradient (DDPG), and after training the learning rate is set to zero so the sampler becomes an ordinary Riemannian MALA kernel with a fixed learned step-size function.

What would settle it

Evaluate a trained step-size network $\epsilon_\theta$ on a fine grid that extends beyond the region covered by gold-standard samples for each of the 44 benchmark targets; if any learned network is unbounded, discontinuous, or violates the Lipschitz and boundedness conditions used in the ergodicity theory for position-dependent MALA, then the post-training ergodicity claim is not sustained.

Watch

Extended reading notes

Core claim

The central claim is that a gradient-based sampler—specifically Riemannian MALA equipped with a position-dependent step size $\epsilon_\theta(x)$—can be trained end-to-end inside the reinforcement-learning Metropolis–Hastings (RLMH) formulation, and that the learned step-size function mixes faster than the best constant step size. To make training feasible, the paper replaces the usual rewards with the contrastive divergence lower bound (CDLB) of equation (11): a per-transition reward combining the log-density change from current to proposed state, the entropy of the accept/reject decision, and the log of the proposal density. Using deep deterministic policy gradient, the paper trains a two-layer neural-network step-size policy and evaluates it on 44 real posterior distributions, reporting the smallest maximum mean discrepancy to gold-standard samples in 89% of tasks. It also reports that training with CDLB eliminates the catastrophic failures that occurred in 7% of tasks with the log-expected-squared-jump-distance reward, and that the learned policies are interpretable: large steps in the tails of a Laplace target, near-constant steps for a Gaussian, and small steps in the narrow tails of a banana target.

Load-bearing premise

After training, the learning rate is turned off, and the final sampler is assumed to sample the target correctly because it is a standard position-dependent Langevin algorithm; this assumes the learned neural network that sets the step size stays smooth and bounded everywhere the chain can go, which is never checked.

Editorial extensions

If this is right

  • If the results hold, tuning gradient-based samplers no longer requires hand-picking a step size against acceptance-rate or squared-jump-distance criteria; a learned policy can do it in a way that adapts to the target's geometry.
  • The CDLB reward is computed from one transition, so it can be used as an objective for adaptive MCMC beyond RLMH, including for random-walk and Barker samplers that the paper also tests.
  • Because the method avoids second-order derivatives of the target, it offers a cheap alternative to Hessian-based Riemannian MALA for targets with varying curvature or tail behaviour.
  • The absence of improvement on the highest-dimensional tasks in the benchmark implies that the practical gain of learning a position-dependent step size is currently concentrated in low- and moderate-dimensional posteriors.
  • Setting the learning rate to zero after training produces a fixed position-dependent sampler whose targeting is inherited from standard ergodicity results; this is the basis for the claim that the final sampler samples the correct posterior.

Reading between the lines

Editorial extensions of the paper, not claims the author makes directly.

  • A testable extension is to learn a diagonal or low-rank preconditioner in addition to the step-size function; the paper reports that its own attempts to learn the full covariance failed, but a factored action space could make that tractable.
  • The CDLB reward's entropy-of-accept/reject term suggests a principled way to balance exploration and exploitation in adaptive MCMC generally, which might also improve stochastic-gradient samplers that lack a Metropolis correction.
  • The reported 89% figure depends on benchmark tasks where the preconditioner is set from gold-standard samples; the appendix's check with a Hessian-based preconditioner suggests the advantage is robust, but the practical gain in fully unsupervised settings remains an open question.
  • A direct consequence not explored in the paper is that the same training pipeline could tune Hamiltonian Monte Carlo step size and trajectory length, where the action space is two-dimensional and likely easier to learn.
Share X Bluesky LinkedIn Reddit HN

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Referee Report

5 major / 5 minor

Summary. The paper extends the Reinforcement Learning Metropolis-Hastings (RLMH) framework of Wang et al. (2025) to gradient-based MCMC. Concretely, it learns a position-dependent step-size function εθ for Riemannian MALA (RMALA) using DDPG, and proposes a new reward called the contrastive divergence lower bound (CDLB) that combines an exploitation term, an accept/reject entropy term, and a proposal-entropy term. Two lower-bound propositions are proved and used to motivate CDLB. The empirical section compares RMALA with RLMH-learned step sizes against RMALA with constant step sizes tuned by AAR or ESJD on 44 posteriordb tasks, reporting that RLMH attains the smallest average MMD in 89% of tasks and that CDLB eliminates the catastrophic training failures occasionally seen with the LESJD reward. Appendices provide proofs, implementation details for both RMALA and RLMH, full benchmark tables, a sensitivity analysis for G0, and additional results for the Barker proposal.

Significance. If the empirical claims are sustained, the paper would be a useful step toward practical RL-based tuning of gradient-based samplers. The CDLB reward is a novel and potentially transferable objective, and the paper is unusually transparent in reporting catastrophic failures, full benchmark tables, and sensitivity to the preconditioner choice. The code is made available and the computational budget is stated. However, the headline 89% claim is currently supported only by point estimates, the RLMH initialization uses gold-standard samples, and the ergodicity guarantee after freezing the policy is asserted rather than verified for the trained networks. These issues directly affect the paper's central claims and need to be addressed before the results can be taken at face value.

major comments (5)
  1. [Section 3, Appendix D.2] The headline comparison is confounded by oracle-informed initialization. Appendix D.2 states that εθ is pretrained by fitting against the gold-standard samples {y_j} with m=10^4, and that the initial step size ε† is selected through a polynomial regression whose coefficients were tuned "through experimentation" on the posteriordb benchmark. The RMALA baselines, by contrast, start from ε=0.1 and adapt online (Appendix C). Thus RLMH begins with target-specific information that the baselines do not receive. To support the claim that the learned step-size policy outperforms an optimized constant step size, the authors should report an ablation without gold-standard pretraining, or give both methods identical initialization information.
  2. [Section 3, Table 2] The 89% win rate is a count of which method has the smallest average MMD over ten replicates, with no significance testing. Many of the reported differences are within one standard error, e.g. gppoisregr-gpregr (CDLB 2.4(0.1) vs AAR 2.2(0.1)) and earnings-log10earnheight (CDLB 1.3(0.0) vs AAR 1.4(0.0)). The claim that RLMH "outperformed" the optimized constant step size is therefore not established by the data as presented; paired tests, bootstrap confidence intervals for the differences, or a multilevel model over tasks should be reported.
  3. [Section 2.1] The assertion that setting the learning rate to zero after training makes ergodicity "immediate from ergodicity of RMALA [Roy and Zhang, 2023]" is not justified. Roy and Zhang's conditions on the position-dependent preconditioner involve regularity properties such as boundedness, positive definiteness, and smoothness, which are not verified for the trained ReLU network εθ. Without checking or enforcing these conditions, the post-adaptation chain need not be P-invariant. Please state the exact conditions required and verify them for the trained policies, or modify the architecture or add clipping to guarantee them.
  4. [Section 2.3 after Proposition 1] The statement that "at stationarity the second term in (8) does not depend on the choice of Markov transition kernel P" is incorrect. The term E[log pn(Xn)/(1+pn(Xn+1))] is an expectation under the joint law of (Xn, Xn+1); even when Pn=Pn+1=P, the joint law depends on the transition kernel. This undermines the stated justification for maximizing only the first term of the lower bound when designing CDLB. The claim should be corrected or the heuristic reformulated.
  5. [Section 3 vs Appendix D.2] The iteration counts are inconsistent. Section 3 states that a total of n=30,000 iterations was performed for each posterior, with the final 5,000 used for assessment, while Appendix D.2 specifies 100 episodes of 500 MCMC iterations each, which amounts to 50,000 iterations for RLMH. This discrepancy also calls into question the statement that the computational cost of all gradient-based algorithms is equal. Please clarify the exact protocol and confirm that all methods use identical numbers of target and gradient evaluations.
minor comments (5)
  1. [Appendix D.2, Parametrisation of εθ] The stated parameter count p=(8+d)(d+1) does not match a fully-connected two-layer network with 8 features per hidden layer; please provide the correct count or the exact architecture used.
  2. [Algorithm 4] Line 6 contains the duplicated phrase "set set Xi+1"; this should be corrected.
  3. [Appendix D.2] There is a typo: "initilaisation" should be "initialisation".
  4. [Propositions 1 and 2] Calling the bounds "tight" because they follow from a single application of Jensen's inequality is stronger than what is demonstrated; the combined bounds may be loose, so a more cautious wording would be preferable.
  5. [Throughout] The benchmark name is written inconsistently as "PosteriorDB" and "posteriordb"; please use the official capitalization consistently.

Circularity Check

1 steps flagged · score 2.0 of 10

Mild circularity from gold-standard pretraining of the RLMH step size; the CDLB derivation and external MMD evaluation are otherwise self-contained.

  1. fitted input called prediction [Appendix D.2, 'Pre-training of εθ' (and Appendix C baseline tuning)]
    "The parameters θ of the neural network εθ were initialised by pre-training against the loss θ ↦ (1/m) Σ_j ||εθ(y_j) − ε†||^2, computed over the gold-standard samples {y_j} with m=10^4 ... This heuristic was further optimised for experiments on posteriordb via the nonlinear regression ε† = a1·ε0^3 + a2·ε0^2 + a3·ε0 + a4, with numerical coefficients determined through experimentation."

    The RLMH policy that underlies the reported 89% win rate is initialized by regressing εθ onto the gold-standard samples of the same posteriordb tasks used for evaluation, and its scalar anchor ε† is itself tuned by regression over that benchmark. The constant-step RMALA baselines instead start at the generic value ε=0.1 and adapt online with no gold-standard information (Appendix C). Thus the headline comparison is not an independent test of learned position dependence: part of RLMH's advantage is built into its target-informed starting point. Because subsequent DDPG updates do alter εθ and MMD is measured against external gold-standard samples, this is a partial confound rather than a complete reduction-by-construction.

full rationale

The paper's theoretical contribution—the CDLB reward in Eq. (11)—is a genuine derived lower bound on contrastive divergence, not a restatement of the performance measure (MMD) used for evaluation, so the main methodological derivation is self-contained. The RLMH framework is self-cited from Wang et al. (2025), with overlapping authorship, but it is used as a foundation and as a LESJD baseline, not as a uniqueness theorem that forces the present conclusions; the central empirical content is evaluated against external posteriordb gold-standard samples. The one notable circularity-adjacent issue is that the learned step-size network is pretrained on those same gold-standard samples and its initial scale ε† is fitted to the posteriordb benchmark, while the constant-step comparators are tuned from a generic initialization; this makes the headline 89% figure partly a fitted input. The paper's ergodicity claim, obtained by setting the learning rate to zero and appealing to Roy and Zhang (2023), rests on unverified regularity of the trained εθ; that is a correctness gap, not circularity. Overall, the derivation is not circular, but the oracle-informed initialization adds a mild circularity burden to the empirical comparison, giving score 2.

Assumptions & free parameters 3 free parameters · 4 assumptions · 0 invented entities

The central claim rests on standard probability assumptions, the heuristic adoption of the CDLB objective, unverified regularity of the learned step-size function, and the off-policy DDPG training on a non-stationary environment. The main free parameter is the benchmark-fitted initialization step size.

free parameters (3)
  • epsilon_dagger polynomial coefficients = a1=29, a2=-26, a3=3.0, a4=1.3
    Initial step size for pretraining epsilon_theta is a cubic polynomial of epsilon_0 fitted through experimentation on the same posteriordb benchmark used for evaluation (Appendix D.2).
  • DDPG hyperparameters = alpha_pi=1e-6, alpha_Q=1e-2, buffer=2.5e4, minibatch=48, episodes=100, episode length=500, centering gain=1e-3
    Chosen without systematic search; these affect the learned policy and are free implementation parameters (Appendix D.2).
  • Baseline RMALA step-size interval and adaptation frequency = [1e-4, 2], adaptation frequency 5000, initial step size 0.1
    Chosen to ensure reasonable performance over posteriordb (Appendix C); affects the baseline comparison.
assumptions (4)
  • standard math Jensen's inequality and existence of densities for P_n and the transition kernel
    Used in proofs of Propositions 1 and 2; the chain's marginal law is assumed to have a density and the transition kernel is assumed to have a density against mu_x = lambda + delta_x.
  • ad hoc to paper Maximizing the CDLB reward improves sampling quality
    The reward is a lower bound on a component of contrastive divergence; the paper adopts this as a heuristic, with no proof that optimizing the reward optimizes mixing (Section 2.3).
  • domain assumption Learned epsilon_theta satisfies conditions for RMALA ergodicity
    Invoked when stopping adaptation and claiming p-invariance is immediate from ergodicity of RMALA, without verifying the neural-network step-size function meets the required regularity and boundedness conditions (Section 2.1).
  • domain assumption DDPG policy gradient theorem applies to the non-stationary MDP environment
    The deterministic policy gradient theorem assumes expectations under the stationary distribution of the MDP, but during RLMH training the chain is not stationary; the paper uses DDPG off-the-shelf (Section 2.1 and Appendix D.1).

how reviews work

0 comments
Cite this review

Pith. "Pith review of Harnessing the Power of Reinforcement Learning for Adaptive MCMC." pith.science (2026). https://pith.science/paper/AQUP24K4

@misc{pith2026250700671,
  author       = {Pith},
  title        = {Pith review of: Harnessing the Power of Reinforcement Learning for Adaptive MCMC},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/AQUP24K4}},
  note         = {Machine review of arXiv:2507.00671}
}
read the original abstract

Sampling algorithms drive probabilistic machine learning, and recent years have seen an explosion in the diversity of tools for this task. However, the increasing sophistication of sampling algorithms is correlated with an increase in the tuning burden. There is now a greater need than ever to treat the tuning of samplers as a learning task in its own right. In a conceptual breakthrough, Wang et al (2025) formulated Metropolis-Hastings as a Markov decision process, opening up the possibility for adaptive tuning using Reinforcement Learning (RL). Their emphasis was on theoretical foundations; realising the practical benefit of Reinforcement Learning Metropolis-Hastings (RLMH) was left for subsequent work. The purpose of this paper is twofold: First, we observe the surprising result that natural choices of reward, such as the acceptance rate, or the expected squared jump distance, provide insufficient signal for training RLMH. Instead, we propose a novel reward based on the contrastive divergence, whose superior performance in the context of RLMH is demonstrated. Second, we explore the potential of RLMH and present adaptive gradient-based samplers that balance flexibility of the Markov transition kernel with learnability of the associated RL task. A comprehensive simulation study using the posteriordb benchmark supports the practical effectiveness of RLMH.

Figures

Figures reproduced from arXiv: 2507.00671 by the authors.

Figure 1
Figure 1. Training the Riemannian Metropolis-adjusted Langevin algorithm (RMALA) [PITH_FULL_IMAGE:figures/full_fig_p009_1.png] view at source ↗
Figure 2
Figure 2. Comparison of training based on CDLB (top panel) and LESJD (bottom panel). Moving av￾erages of the rewards rn are dis￾played. Training with CDLB oscil￾lates during the initial exploration phase, before stabilising once suffi￾cient information has been gained. In contrast, LESJD provides insuf￾ficient signal and can lead to catas￾trophic failure of RLMH. A secondary, but still important result was that, although the … view at source ↗

Discussion (0). Sign in to comment.

Forward citations

Cited by 1 Pith paper

Reviewed papers in the Pith corpus that reference this work. Sorted by Pith novelty score. Full citation record

  1. Stop the Sampler! Classifier-Based Adaptive Stopping for Sampling Kernels

    cs.LG 2026-06 conditional novelty 6.0 of 10

    Classifier-based adaptive stopping, trained as a non-acyclic GFlowNet, shortens MCMC trajectories while preserving or improving sample quality relative to ULA and diffusion-sampler baselines.

Reference graph

Works this paper leans on

16 extracted references · 11 canonical work pages · cited by 1 Pith paper

  1. [1]

    Here S + d denotes the set of symmetric positive definite d × d matrices and |M | denotes the determinant of a matrix M ∈ S+ d . The experiments that we report in Section 3 of the main text concern either the case where the preconditioner matrix G(·) = ϵ−1G0 is position-independent, or G(·) = ϵ(·)−1G0 with a position-dependent scale ϵ(·) learned using RLM...

    work page 1998

  2. [2]

    To begin, we condition on the value ofXn and recall the definition of the transition density p(Xn+1|Xn) from the main text: E[− log p(Xn+1|Xn) | Xn] = − Z log(p(x′|Xn))p(x′|Xn) dµXn(x′) = −r(Xn) logr(Xn) − Z u(x′|Xn) logu(x′|Xn) dλ(x′) = − 1 − Z u(x′|Xn) dλ(x′) log 1 − Z u(x′|Xn) dλ(x′) | {z } (15.1) − Z u(x′|Xn) logu(x′|Xn) dλ(x′) | {z } (15.2) . (15) 18...

    work page 2011

  3. [11]

    Drawing on the reward-centring framework of Naik et al

    for further detail. Drawing on the reward-centring framework of Naik et al. [2024], we replace the vanilla TD approximation Qn with the centralised TD approximation ¯Qn = (rn − ¯R) + γQϑ′ (sn+1, πθ(sn+1)) (18) where the running average of past rewards is denoted ¯R. This centring of the immediate reward term around ¯R has the capacity to reduce variance i...

    work page 2024

  4. [12]

    parameters to be inferred. Pre-training of ϵθ The parameters θ of the neural network ϵθ were initialised by pre- training against the loss θ 7→ 1 m mX j=1 ϵθ(yj) − ϵ† 2 , computed over the gold-standard samples {yj}m j=1 with m = 10 4, so that the proposal corresponding to ϵθ is initially a position-independent proposal with step size ϵ†. This was achieve...

    work page 1951

  5. [13]

    (soft-update rate) 1: initialise actor parameters θ and critic parameters ϑ 2: πθ′ ← πθ Qϑ′ ← Qϑ ▷ initialise target networks 3: ¯R ← 0 ▷ initialise average reward 4: for each episode do 5: for t = 0, 1, 2, . . .until episode ends do 6: at = πθ(st) + Nt ▷ action with exploration noise 7: execute at, observe (rt, st+1) 8: store transition ( st, at, rt, st+...

    work page 2012

  6. [15]

    E.4 Exploring the Sensitivity to G0 The results that we report for posteriordb in the main text set G0 based on 10 4 gold- standard samples from the target

    Instances where RMALA failed with AAR or ESJD were because proposal covariance became numerically zero; could be avoided but in principle this could introduce bias and so we simply report the failure in these cases. E.4 Exploring the Sensitivity to G0 The results that we report for posteriordb in the main text set G0 based on 10 4 gold- standard samples f...

    work page 2008

  7. [16]

    It is based on the accept-reject rule of Barker [1965], stated for d-dimensional distributions in Algorithm

    as an alternative to MALA that trades efficiency for being more easily tuned. It is based on the accept-reject rule of Barker [1965], stated for d-dimensional distributions in Algorithm

    work page 1965

  8. [1951]

    doi: 10.1214/aoms/1177729586. G. O. Roberts and J. S. Rosenthal. Optimal scaling of discrete approximations to Langevin diffusions. Journal of the Royal Statistical Society Series B: Statistical Methodology, 60(1):255–268,

    arXiv

Show all 16 references
  1. [1996]

    K. Kim, Z. Xu, J. R. Gardner, and T. Campbell. Tuning sequential Monte Carlo samplers via greedy incremental divergence minimization. arXiv preprint arXiv:2503.15704 ,

  2. [2000]

    T. Liu, N. Surjanovic, M. Biron-Lattes, A. Bouchard-Cˆ ot´ e, and T. Campbell. AutoStep: Locally adaptive involutive MCMC. arXiv preprint arXiv:2410.18929 ,

    arXiv

  3. [2012]

    Y. Chen, D. Z. Huang, J. Huang, S. Reich, and A. M. Stuart. Sampling via gradient flows in the space of probability measures. arXiv preprint arXiv:2310.03597 ,

    arXiv

  4. [2016]

    for full detail. The deterministic policy gradient theorem states that ∇θJ(πθ) = Eπ h ∇θπ(s) ∇aQπ(s, a)|a=π(s) i , (16) where the expectation here is taken with respect to the stationary distribution of the MDP when the policy π is fixed. The action-value function Qπ(s, a) giv...

    2014

  5. [2018]

    Garreau, W

    D. Garreau, W. Jitkrittum, and M. Kanagawa. Large sample analysis of the median heuristic. arXiv preprint arXiv:1707.07269 ,

    arXiv

  6. [2022]

    E.1 Implementation Details The initial state x0 ∈ Rd of all Markov chains was taken to be the arithmetic mean of 104 gold-standard samples from the target

    are contained and discussed in Appendix E.5. E.1 Implementation Details The initial state x0 ∈ Rd of all Markov chains was taken to be the arithmetic mean of 104 gold-standard samples from the target. For the experiments reported in Section 3 of the main text, the static compo...

    2008

  7. [2023]

    Doig and L

    R. Doig and L. Wang. A unified framework for multiple-try Metropolis algorithms. arXiv preprint arXiv:2503.11583,

    arXiv

  8. [2024]

    URL https: //arxiv.org/abs/2405.09999. C. Pasarica and A. Gelman. Adaptively scaling the Metropolis algorithm using expected squared jumped distance. Statistica Sinica, pages 343–364,

    arXiv

Pith tools

Reviewed August 6, 2026 · model on record in the stance chip above.