The Traceplot Thickens: MCMC Diagnostics for Non-Euclidean Spaces

Consider a parameter space which contains a large number of binary or discrete parameters. Common examples of these parameters include variable inclusion indicators from Bayesian models incorporating variable selection, edge-existence parameters in a graphical model, or co-association parameters in a clustering model. For each of these examples, realistic applications require a large number of these discrete parameters, and standard convergence diagnostics simply don’t hold up.

Consider a 15-node Bayesian network (BN), an example of which is visualized in Figure 1. One common field in BN analysis is structure discovery, in which the goal is to identify relationships among a set of variables which are each visualized by a node. In the example given in Figure 1, we see that node 1 (representing variable 1) has a direct relationship with node 2, but an indirect relationship with node 13, as there is no arrow drawn directly from 1 to 13. We represent these relationships with the parameters $e_{i,j}$, in which $e_{i,j}=1$ if there is an arrow drawn from node $i$ to node $j$, and $e_{i,j}=0$ otherwise. Note that in a 15-node BN space there are 210 such parameters.

To provide an example, we simulated data from a 15-node Bayesian network and used the partition-MCMC algorithm [15] to sample 1000 iterations of 5 chains, each from a posterior distribution defined on the space $\mathcal{X}_{B}$, the space of 15-node Bayesian networks. Figure 2 gives a univariate traceplot of the arrow-existence parameter $e_{1,10}.$

Because the parameter in question is binary, the information gained by viewing the traceplot is very limited. The situation $e_{1,10}$ is clearly very unlikely, but does the fact that the sampler doesn’t often include it in the model actually demonstrate that the sampler is not performing well?

Additionally, univariate measures of the effective sample size or potential scale reduction factor are equally useless as both assume there will be some level of variability per parameter. Multivariate versions of these diagnostics may fare better, but will still struggle as they assume a Euclidean paradigm in this high-dimensional non-Euclidean space.

A decent solution to this problem is to produce traceplots (and potentially calculate univariate diagnostics) using the posterior value or likelihood evaluated at each MCMC sample [11, 13]. However, this approach is flawed when the parameter space includes multiple modes of similar likelihood, as is common (and unpredictable) in many modern high-dimensional applications. In these situations, any likelihood-based diagnostics fail to identify the difference between the modes.

In Section 4.2 below, we present a generalized traceplot on the same MCMC samples which incorporates information from all $e_{i,j}$ variables, and clearly demonstrates that the sampler is failing to mix adequately.

To explore potential consequences of non-Euclidean MCMC moves, we consider a hypothetical scenario in which a statistician (we’ll call him Luke) wants to sample from a univariate distribution on the random variable $X$ using an MCMC algorithm. The distribution is an equal parts mixture of three normal distributions $N_{1}(\mu_{1},\sigma^{2}_{1}),N_{2}(\mu_{2},\sigma^{2}_{2}),N_{3}(\mu_{3},\sigma^{2}_{3})$, where $(\mu_{1},\mu_{2},\mu_{3})=(-3,0,3)$ and $\sigma_{1}=\sigma_{2}=\sigma_{3}=.1.$ Due to the small standard deviations, the three parts of the mixture essentially don’t overlap at all. We denote this sampling space and distribution using $\mathcal{X}_{3}$, as it is tri-modal.

In an effort to ensure good results, Luke uses two different MCMC samplers to draw from $\mathcal{X}_{3}$. $\mathcal{M}_{1}$ is a typical Metropolis-Hastings algorithm in which proposed draws $X^{*}$ follow the distribution $X^{*}\sim N(X_{t},1)$, where $X_{t}$ is the current draw. However, attempting to cleverly improve the rate at which mixing is achieved, Luke builds $\mathcal{M}_{2}$ as a Metropolis-Hastings algorithm with a proposal distribution that is an equal parts mixture of a normal distribution centered at the current draw with standard deviation .1, and a normal distribution centered at the negative of the current draw with standard deviation .1. That is, proposed draws follow the distribution

The thought behind $\mathcal{M}_{2}$ is that a lower standard deviation will improve the mixing within each mode, and that taking approximately half of the draws from a distribution centered at $-X_{t}$ instead of $X_{t}$ will allow for the algorithm to effectively move between the two outside modes (Luke apparently forgot about the mode centered at 0).

Both samplers were then run on 7 chains of 2000 draws each, starting at the points $(-6,-4,-2,0,2,4,6).$ Figures 3 and 4 show the traceplots of the draws from Samplers 1 and 2 respectively, while Table 1 presents the ESS and PSRF.

Looking at the diagnostics for $\mathcal{M}_{1}$, we see that because the traceplot is clearly not mixing well across modes, the ESS is very low, and the PSRF is much larger than the recommended 1.01, we should not use the results from $\mathcal{M}_{1}$. In contrast, for $\mathcal{M}_{2}$ the traceplot looks like it is mixing well, the ESS is high, and the PSRF is 1.01. Thus, based on these diagnostics, it appears that we should accept the results from $\mathcal{M}_{2}$.

However, as could be guessed by the strange construction of the proposal distribution, $\mathcal{M}_{2}$ is also failing to mix well and results in a very bad approximation of the target distribution. In fact, if we define $\mathcal{P}$ to be the target distribution, $\mathcal{Q}_{1}$ and $\mathcal{Q}_{2}$ to be the empirical distributions given by the draws from Samplers 1 and 2 respectively, and $KL$ to be the Kullback-Leibler divergence, then $KL(\mathcal{P}|\mathcal{Q}_{1})=12,129$ while $KL(\mathcal{P}|\mathcal{Q}_{2})=86,355$. That is, the $KL$ divergence between $\mathcal{P}$ and $\mathcal{Q}_{2}$ is 7-times worse than the divergence between $\mathcal{P}$ and $\mathcal{Q}_{1}$. This is mainly due to the fact that, while 1/3 of $\mathcal{P}$ is concentrated in the central mode, only 1/7 of the draws making up $\mathcal{Q}_{2}$ are from the central mode.

Clearly, $\mathcal{M}_{2}$ was a bad idea from the beginning as it avoided the mode centered at 0 almost entirely. However, the traceplot, ESS and PSRF all appeared to show good mixing for $\mathcal{M}_{2}$. If an ill-considered MCMC algorithm was able to fool our selected set of diagnostics so completely, we need to understand why. After all, most practical applications of MCMC algorithms do not operate in clearly-defined univariate spaces where poorly-constructed algorithms are easily spotted from the start.

The problem here lies in the distinction between the distance being used by the algorithm vs. the distance being used by the diagnostic. Traceplots, the univariate effective sample size, and the univariate potential scale reduction-factor all assume that distances between draws in an MCMC algorithm are Euclidean. This assumption works in the case of $\mathcal{M}_{1}$ as, defining the current draw as $X_{t}$, moving to a different value $Y$ becomes monotonically less likely as the Euclidean distance between $X_{t}$ and $Y$ increases.

However, $\mathcal{M}_{2}$ operates with a different understanding of distance. Given that the current draw is $X_{t}$, the probability of moving to $Y$ no longer monotonically decreases as the Euclidean distance increases. In fact, if $Y=-X_{t}$, we know for most of our sample space that the probability of moving from $X_{t}$ to $Y$ is greater than the probability of moving from $X_{t}$ to 0, despite the fact that $|X_{t}-0|<|X_{t}-Y|.$ This is represented using a heatmap in Figure 5, which shows (among other things) that the probability of moving from -3 to 3, is much higher than the probability of moving from -3 to 0.

This discrepancy in definitions of distance causes all of the problems in the convergence diagnostics shown for $\mathcal{M}_{2}$. For the traceplot, the fact that most of the chains regularly bounce between modes at -3 and 3 obscures the fact that most chains never visit the mode at 0. Here the traceplot’s notion of Euclidean distance (which assumes that 0 falls between -3 and 3) is in conflict with the distance used by $\mathcal{M}_{2}$ which essentially treats -3 and 3 as identical points, with nothing in between them.

Then, notice that both ESS and the PSRF are based on the sum of squared errors around chain means. Since the sample mean is a measure of central tendency with respect to the Euclidean distance, if a chain bounces back and forth between the modes at -3 and 3, the chain mean will be somewhere around 0. However, the chain never visits anywhere near 0, and so the sample mean fails as a true measure of central tendency in terms of the MCMC movements of $\mathcal{M}_{2}$.

In Section 4.3, we present generalized diagnostics that clearly demonstrate the lack of mixing in $\mathcal{M}_{2}$.

Selecting a distance function $d$ on $\mathcal{X}$ to use with a sampler $\mathcal{M}$ involves evaluating and prioritizing several important factors. First, in order to be relevant, $d$ must at least somewhat reflect the movement of $\mathcal{M}$ in the sample space $\mathcal{X}$. That is, if $X_{1},X_{2},Y\in\mathcal{X}$ with $d(X_{1},Y)<d(X_{2},Y)$, then if $\mathcal{M}$ is currently at $Y$, $X_{1}$ should generally be a more-likely next move than $X_{2}$. This is the most obvious requirement for $d$, and is fulfilled quite clearly by standard convergence diagnostics based on the Euclidean distance. If this requirement is not met (at least for the majority of sampling steps) the diagnostic measures will register frequent movements by the sampler that appear to be very large, but are in reality, much smaller (this was the fundamental problem in our second case-study). This requirement of relevance to $\mathcal{M}$ can be difficult to accomplish, particularly when a sampler has different ‘types’ of moves that it can choose from, like a split-merge step in a DPMM. However, we have found distance functions to be successful in identifying mixing issues as long as $d$ agrees with $\mathcal{M}$ in the large majority of draws.

Second, the distance function $d$ should be relevant to the diagnostic target. If $\mathcal{X}$ is a ten dimensional parameter space and $d$ is defined to ignore nine of those dimensions, the resulting diagnostics will only be applicable to the mixing of the one parameter included in the definition. Defining $d$ to ignore particular dimensions of the parameter space can be useful as it can allow for examination of sampler mixing on groups of related or important parameters, but should be done with caution knowing that information is being excluded.

Third, computational time should be considered when defining $d$. Depending on the proximity-map (which we will discuss in Section 3.2), $d$ may need to be evaluated as many as $N^{2}$ times (or as few as $N$ times), where $N$ is the total number of unique draws across all chains. If $d$ is slow to calculate and $N$ is large, this can lead to exceedingly long computation times, rendering the diagnostic process largely useless. With these considerations in mind we provide two examples of distance functions that we have found useful in application.

The proximity-map must, at least in part, be based on the selected distance function. Remaining fully faithful to a particular definition of distance is difficult when performing any sort of dimension reduction or mapping, especially in higher dimensions, but may also be difficult in lower dimensions. However, some effort must be made in order to ensure that the resulting diagnostics are informative. Standard dimensionality reduction techniques, such as random projections or principle component analysis, assume a Euclidean distance paradigm and would not serve as effective proximity-maps, unless your distance function is the Euclidean distance. As we will demonstrate in the examples below, the notion of ’remaining faithful’ to the distance function can be quite loose in application and still get excellent results.

The second consideration, also very important, is computation time. Generally speaking, proximity-maps that are more faithful to the distance function also tend to be more difficult to compute. Selecting a good proximity-map essentially comes down to a trade-off between the two. Diagnostic measures that take too long to compute are not particularly useful, but neither are diagnostic measures that fail to reveal mixing or convergence issues. We now present two examples of proximity-maps and discuss the costs and benefits associated with each.

We begin with a simple bi-modal distribution on the real line, and denote this sample space $\mathcal{X}_{2}$. $\mathcal{X}_{2}$ is constructed as an equal-parts mixture of two normal distributions, one centered at -3 and the other at 3, both with standard deviation 1. The posterior is visualized in Figure 6.

For this simulation we use two different MCMC samplers, denoted $\mathcal{M}_{3}$ and $\mathcal{M}_{4}$, both defined on $\mathcal{X}_{2}$. Both $\mathcal{M}_{3}$ and $\mathcal{M}_{4}$ are Metropolis-Hastings algorithms with normal proposal distributions, centered at the current draw. The only difference is that the proposal distribution for $\mathcal{M}_{3}$ has standard deviation .1, and so will have difficulty crossing between the two modes, while the proposal for $\mathcal{M}_{3}$ has standard deviation 2. For each algorithm we draw 7 chains beginning at points (-6, -4, -2, 0, 2, 4, 6) and run each chain for 2000 iterations.

Figure 7 shows the standard traceplots for algorithms $\mathcal{M}_{3}$ and $\mathcal{M}_{4}$, while Table 2 gives the ESS and PSRF, calculated in the typical way. Based on the separation of chains and high auto-correlation in the traceplot, as well as the low ESS and PSRF much higher than the recommended 1.01, $\mathcal{M}_{3}$ does not appear to be mixing well at all. In contrast, the (somewhat) well-mixed traceplot for $\mathcal{M}_{4}$, along with a higher ESS and much lower PSRF, suggest draws that are much closer to convergence.

Now consider the generalized diagnostics given in Figure 8 and Table 3. These are the traceplot, ESS and PSRF derived from using the Generalized Convergence Diagnostic algorithm with the Euclidean distance and Nearest Neighbor proximity-map. As we would hope, since standard diagnostics are also based on the Euclidean distance, these diagnostics are almost identical to the standard diagnostics. This demonstrates that our generalized diagnostics can be viewed as approximate extensions of standard diagnostics.

In Section 2.1 we presented a short example of a standard traceplot on a binary parameter from a 15-node Bayesian network. As discussed in Section 2.1, the standard traceplot was almost entirely useless, and was not informative about the mixing of the MCMC algorithm. While there are other methods of diagnosing convergence in such models, in order to demonstrate the capabilities of our approach, we show a generalized traceplot of the same MCMC samples based on the Hamming distance and Nearest Neighbor proximity-map in Figure 9. Using this traceplot it is immediately apparent that our sampler is mixing very poorly. This can be seen by the significant number of flat lines suggesting failures of the MCMC sampler to move from a single point, and by the chain represented by the green line which appears to be sampling from an entirely different portion of the sample space as the rest of the chains.

In Section 2.2 we introduced a tri-modal distribution on the real line $\mathcal{X}_{3}$. Draws using the simple algorithm $\mathcal{M}_{1}$ were clearly not mixing, while draws from a more complex, but worse algorithm $\mathcal{M}_{2}$ appeared to be mixing based on standard convergence diagnostics. Here we present Figure 10, a generalized traceplot for the draws from $\mathcal{M}_{2}$ created using the Metropolis-Hastings distance and the Nearest Neighbor proximity-map, and Table 4 which gives the generalized ESS and PSRF calculated using the same method. Note that both the traceplot and the PSRF clearly reveal a lack of between-chain convergence, while the ESS is much lower (and certainly more accurate) than the ESS calculated using the Euclidean distance in Section 2.2.

The Dirichlet Process Mixture Model (DPMM) may be used to cluster observations into latent groups, and is unique in that its construction does not require that the number of groups be specified a priori. The algorithm and its theoretical justification were introduced approximately 50 years ago [2, 8, 9].

A unique problem when using an MCMC sampler to perform inference with a DPMM is that the number of parameters at each iteration may change as latent mixture components are added or deleted. For example, at the start of a given MCMC iteration to fit a DPMM using a mixture of Gaussians, there may be three latent components. Since each component has its own group-specific mean and variance parameters, we start the iteration with three mean vectors $\mu_{1},\mu_{2},\mu_{3}$, and three covariance matrices $\Sigma_{1},\Sigma_{2},\Sigma_{3}$. However, as we iterate through the observed values assigned to each of the three components, the MCMC sampler may assign one or more observations to a previously nonexistent fourth group. Thus, at the start of the next iteration our parameter space will consist of $\mu_{1},\dots,\mu_{4}$ and $\Sigma_{1},\dots,\Sigma_{4}$.

In terms of MCMC diagnostics, there are two issues here: (1) the multivariate nature of the data, and (2) the expansion (and contraction) of the number of latent components, and thus the dimension of the parameter space over the course of the MCMC as described above. This second issue in particular is not handled well by current diagnostic tools. In order to make traceplots or calculate the values of MCMC diagnostics, one must condition on the number of latent components found at each MCMC iteration, and evaluate diagnostics on those results separately. These MCMC iterations are not necessarily contiguous after conditioning, so this process results in an incomplete or potentially misleading picture of the path of the MCMC. A common solution is to simply plot the log-likelihood at each iteration. However, as discussed in Section 2.1, this practice can obscure information when the posterior space is multi-modal.

To demonstrate the utility of the generalized diagnostics introduced in this paper, we simulated data from a mixture of three Gaussian densities, denoted $\mathcal{X}_{D}$, and used it to fit a DPMM. The true component means are $(-20,20)$, $(0,0)$, and $(20,-20)$, variances are all $10*I_{2}$, and the mixing proportions are $(0.4,0.3,0.3)$. We fit a DPMM with a diagonal covariance structure, setting a normal prior on the mean parameters $\mu_{k}\sim N(\mu_{0},\sigma^{2})$ and an inverse gamma prior on the diagonal variance $\sigma^{2}\sim IG(a,b)$. Using Algorithm 2 from Neal [19], we learned parameters $\mu_{1},\dots,\mu_{K}$, and $\sigma^{2}_{1},\dots,\sigma^{2}_{K}$ for latent groups $k=1,\dots,K$ (noting again that $K$ changes between iterations). We ran five chains with this MCMC sampler, denoted $\mathcal{M}_{5}$, for 12,000 iterations each including 2,000 burn-in iterations. Each chain was initialized at a different place, with either 1,2,3,4 or 5 latent components.

The generalized diagnostic tables and plots for sampling from this space are provided in Figures 11-12, and Table 5. These generalized diagnostics use the Hamming distance on the co-association matrices from each MCMC iteration, and the Lanfear proximity-map. Figure 11 shows major differences between the chains during the burn-in period, the scale of which distorts the remaining iterations. In Figure 12, we include only the iterations from after the burn-in, and note that the sampler appears to demonstrate some signs of lack of convergence between chains, although overall the chains appear to be moving in similar spaces. While the generalized ESS and PSRF look good, the DPMM in this instance does warrant some additional investigation to ensure that the occasional points when the MCMC chains appear trapped in different modes are not interfering with the overall inference.