[Paper Review] $R^*$: A robust MCMC convergence diagnostic with uncertainty using decision tree classifiers

Ben Lambert, Aki Vehtari|arXiv (Cornell University)|Mar 17, 2020

Markov Chains and Monte Carlo Methods31 references4 citations

TL;DR

This paper introduces $R^{*}$, a novel MCMC convergence diagnostic that uses machine learning classifiers—specifically gradient-boosted regression trees and random forests—to detect lack of mixing across chains by measuring how well they can predict chain identity from posterior samples. Unlike traditional $˂{R}$, $R^{*}$ leverages the full joint distribution and provides uncertainty estimates, offering a more robust and comprehensive assessment of convergence, especially in high-dimensional or complex posterior geometries.

ABSTRACT

Markov chain Monte Carlo (MCMC) has transformed Bayesian model inference over the past three decades: mainly because of this, Bayesian inference is now a workhorse of applied scientists. Under general conditions, MCMC sampling converges asymptotically to the posterior distribution, but this provides no guarantees about its performance in finite time. The predominant method for monitoring convergence is to run multiple chains and monitor individual chains' characteristics and compare these to the population as a whole: if within-chain and between-chain summaries are comparable, then this is taken to indicate that the chains have converged to a common stationary distribution. Here, we introduce a new method for diagnosing convergence based on how well a machine learning classifier model can successfully discriminate the individual chains. We call this convergence measure $R^*$. In contrast to the predominant $\widehat{R}$, $R^*$ is a single statistic across all parameters that indicates lack of mixing, although individual variables' importance for this metric can also be determined. Additionally, $R^*$ is not based on any single characteristic of the sampling distribution; instead it uses all the information in the chain, including that given by the joint sampling distribution, which is currently largely overlooked by existing approaches. We recommend calculating $R^*$ using two different machine learning classifiers - gradient-boosted regression trees and random forests - which each work well in models of different dimensions. Because each of these methods outputs a classification probability, as a byproduct, we obtain uncertainty in $R^*$. The method is straightforward to implement and could be a complementary additional check on MCMC convergence for applied analyses.

Motivation & Objective

• To address the limitations of existing MCMC convergence diagnostics, particularly their focus on marginal distributions and failure to detect joint distribution convergence.
• To develop a convergence metric that uses all available information in the MCMC chains, including joint sampling patterns, rather than relying on summary statistics like variance or rank-based measures.
• To provide a diagnostic that quantifies uncertainty in convergence assessment, enabling more reliable inference in applied Bayesian analysis.
• To offer a complementary, stringent check to existing diagnostics like $˂{R}$ and effective sample size (ESS), especially in high-dimensional models where traditional methods may fail.

Proposed method

• Train a binary classifier (e.g., gradient-boosted trees or random forests) to predict which chain a given MCMC draw originated from, using posterior samples as input features.
• Use the classifier’s out-of-sample accuracy as the $R^{*}$ statistic: high accuracy indicates poor mixing (chains are distinguishable), while low accuracy suggests good mixing (chains are indistinguishable).
• Calculate $R^{*}$ as the mean classification error across all chains, with values close to 0.5 indicating good convergence (random guessing), and values significantly below 0.5 indicating poor mixing.
• Estimate uncertainty in $R^{*}$ using the variance of the classifier’s prediction probabilities across multiple runs or bootstrap resampling.
• Apply chain splitting (e.g., splitting each chain into two parts) before computing $R^{*}$ to improve sensitivity to intra-chain mixing issues.
• Incorporate additional covariates such as iteration block indicators to detect non-stationary behavior or random walk-like dynamics in chains.

Experimental results

Research questions

• RQ1 Can a machine learning classifier effectively detect lack of mixing in MCMC chains by learning to distinguish between chains based on their posterior samples?
• RQ2 How does $R^{*}$ perform in detecting convergence issues in the joint posterior distribution when marginal distributions appear converged?
• RQ3 Does $R^{*}$ provide a more stringent and reliable convergence diagnostic than existing methods like $˂{R}$, especially in high-dimensional models?
• RQ4 Can $R^{*}$ quantify uncertainty in convergence assessment, and how does this improve interpretability compared to point estimates?
• RQ5 How does the inclusion of auxiliary features like iteration blocks affect the performance and interpretation of $R^{*}$?

Key findings

• In multiple test cases, $R^{*}$ successfully diagnosed poor convergence in the tails of marginal distributions and in the joint distribution, even when $˂{R}$ indicated convergence.
• The use of gradient-boosted trees and random forests as base classifiers yielded consistent and robust $R^{*}$ values across diverse models, including high-dimensional ones with thousands of parameters.
• Splitting chains before computing $R^{*}$ improved its ability to detect poor intra-chain mixing, suggesting this should be standard practice.
• The uncertainty estimates derived from classifier probabilities provided a meaningful measure of confidence in the $R^{*}$ result, enhancing interpretability.
• Despite higher computational cost than $˂{R}$, $R^{*}$ remained feasible when using thinned draws or subset dimensions, and it detected convergence issues missed by traditional diagnostics.
• In models with complex posterior geometry, $R^{*}$ outperformed $˂{R}$ in identifying non-converged chains, particularly in detecting non-mixing in the joint distribution.