How to handle divergent chains in Stan mixture models

Modeling
, ,
1

Hello, I want to use a mixture model to decide whether my data is better explained by two components or by a single component.

For example, let me take (simulated) data to be the following. The model behaves well (Reff close to 1; Neff comparable with iter) when the two components in the simulated data are equal in proportion. So, I will take proportion to be 1:2. In this case, the chains seem to be forming two clusters, with the two clusters having widely different lp__. If I understand correctly, this means the chains with lower lp__ are finding a posterior estimate with a much lower posterior density than the other chains.

However, how should I proceed after this? Should I discard the chains with lower mean lp__? With lower HDI of lp__? From some places (including Claude and ChatGPT), I read that I should proceed by computing Rhat for individual chains. However, Rhat seems to be computed across multiple chains rather than individual; so I’m missing something. But let’s say I do compute it (see below), I see that all Rhat are close to 1 (and even below 1?).

Thanks a lot in advance!

To elaborate:

my.data_ <- c(rbinom(40, size=8, prob=rbeta(40, 51, 951)),
              rbinom(80, size=8, prob=rbeta(80, 951, 51)));

My model is:

data{
  int<lower=0> num_participants;
  int<lower=1> num_components;
  array[num_participants] int<lower=0> x;
}
transformed data{
  int N = num_participants;
  int K = num_components; // number of groups
}
parameters{
  simplex[K] w;       // mixing proportions
  ordered[K] mode_raw;   // locations of mixture components
  vector<lower=0>[K] c;  // scales or concentrations of mixture components
}
transformed parameters{
  vector<lower=0, upper=1>[K] mode;
  mode = inv_logit(mode_raw);
  vector<lower=1>[K] a, b;
  for(i in 1:K){
    a[i] = mode[i] * c[i] + 1;
    b[i] = (1 - mode[i]) * c[i] + 1;
  }
}
model{
  vector[K] log_w = log(w);  // cache log calculation
  mode_raw ~ normal(0,10);
  c ~ lognormal(0, 2);
  for (n in 1:N) {
    vector[K] lps = log_w;
    for (k in 1:K){
      lps[k] += beta_binomial_lpmf(x[n] | 8, a[k], b[k]);
    }
    target += log_sum_exp(lps);
  }
}
generated quantities{
  // The two groups unmixed
  vector[K] x_group;
  for (k in 1:K)
    x_group[k] = beta_rng(a[k], b[k]);
}

And I fit it:

my.model <- rstan::stan(model_code = my.mixture_model,
                        data = list(num_participants = length(my.data_),
                                    x = my.data_,
                                    num_components = 2),
                        iter = 10000,
                        cores = 4,
                        chains = 10,
                        refresh = 0,
                        save_warmup = TRUE);

The diagnostics suggest that the chains are diverging:

my.model

Inference for Stan model: anon_model.
10 chains, each with iter=10000; warmup=5000; thin=1; 
post-warmup draws per chain=5000, total post-warmup draws=50000.

               mean se_mean    sd    2.5%     25%     50%     75%   97.5% n_eff  Rhat
w[1]           0.37    0.03  0.09    0.26    0.31    0.35    0.40    0.59     7  1.98
w[2]           0.63    0.03  0.09    0.41    0.60    0.65    0.69    0.74     7  1.98
mode_raw[1]   -6.70    3.13  8.98  -26.52  -11.53   -6.19   -3.02   10.17     8  1.68
mode_raw[2]    8.87    0.76  5.76    3.09    4.00    7.09   12.20   23.02    58  1.07
c[1]          24.83    6.79 59.10    6.44   10.21   13.91   22.82  106.93    76  1.06
c[2]          24.98    5.57 25.92    0.01   13.89   19.68   30.00   77.47    22  1.22
mode[1]        0.21    0.18  0.39    0.00    0.00    0.00    0.05    1.00     5 25.23
mode[2]        0.99    0.00  0.01    0.96    0.98    1.00    1.00    1.00    68  1.10
a[1]          13.24    9.89 57.51    1.00    1.00    1.02    2.41   94.38    34  1.11
a[2]          25.47    5.47 24.85    1.01   14.86   20.63   30.54   75.37    21  1.23
b[1]          13.60    2.62 14.90    1.00    8.28   11.90   15.96   44.09    32  1.11
b[2]           1.51    0.09  1.15    1.00    1.00    1.01    1.52    4.19   173  1.08
x_group[1]     0.26    0.16  0.36    0.00    0.03    0.08    0.21    0.99     5  5.52
x_group[2]     0.86    0.08  0.22    0.13    0.89    0.95    0.98    1.00     8  1.67
lp__        -194.04    6.39 14.37 -224.64 -189.76 -187.03 -185.99 -184.89     5  9.57

However, when I look at individual chains, I see them forming two clusters:

plot_chain_params <- function(model, param){

    params <- as.array(model, pars=param);
    if (length(dim(params)) > 2){
        dim(params) <- c(dim(params)[1] * dim(params)[3], dim(params)[2])
    }
    dimnames(params) <- list(iteration = NULL,
                             chain = paste0("chain_", 1:ncol(params)))

    df <- as.data.frame.table(params, responseName = "param") |>
        rename(chain = chain)
    print(str(df))

    ggplot(df, aes(x = param, fill=chain)) +
        geom_density() +
        labs(x = param, y = "Chain") +
        theme_minimal() +
        theme(legend.position = "none") +
        facet_wrap(facets="chain", ncol=1)
}

plot_chain_params(my.model, "x_group[2]")

## ggsave(here("plots", "stanforums_x_group_2.png"))

stanforums_x_group_2
stanforums_x_group_22096×2099 149 KB

Their lp__:

stanforums_lp
stanforums_lp2096×2099 173 KB

Computing Rhat for individual chains (does it make sense?)

for(i in 1:10){
    print(rstan::Rhat(as.array(my.model, pars="x_group[2]")[,i,1]))
}

[1] 1.000232
[1] 1.003585
[1] 1.002601
[1] 1.000395
[1] 1.001119
[1] 0.9998907
[1] 1.000066
[1] 1.018789
[1] 1.000938
[1] 1.000387

Assuming it makes sense, one explanation I see for my data is that the model is explaining the smaller component as noise. Both ess_tail as well as ess_bulk are also comparable across all the 10 chains.

2

One thing I see is that the between-chain divergences are much unlikely if I fix the modes to the values obtained by running kmeans over the data. (Though, merely passing them through init is not enough as the models can drift.) Post that, I can run the model with 1 vs 2 components and compare them using loo or something.

3

If you calculate an R-hat value for a single chain in Stan, it will split it in half and apply R-hat to the two halves. I assume RStan will do that in the usual way, but Id on’t use Rstan, so I’m not 100% sure.

This is a general problem with mixture models when the indexes are not identified, so different chains will lead to different results. There are two solutions. One is to only measure convergence for index-invariant properties. The other is to try to identify the model. A standard way to do this is to order the coefficients but you’ve already done that with the raw mode, so I’m guessing the modes are not going to be well separated.

I’m not sure what’s going wrong with your model as you’re ordering the coefficients and they appear well separated in the posterior. It looks like what may be happening is that the sampler is initializing and finding minor modes with very low log density judging from the lp__ plot.

Another thing you can do if you know some of the modes don’t matter to the posterior is to initialize in the “good” modes. Typically the sampler won’t be able to escape. It’s cheating a bit here, but with the gap in lp__ you see between the modes, the approximation of dropping the smaller mode shouldn’t really matter.

4

I will need to check what rstan is doing under the hood sometime, and perhaps also follow the references that documentation of rstan::Rhat mentions

 Aki Vehtari, Andrew Gelman, Daniel Simpson, Bob Carpenter, and
 Paul-Christian Burkner (2019). Rank-normalization, folding, and
 localization: An improved R-hat for assessing convergence of MCMC.
 _arXiv preprint_ ‘arXiv:1903.08008’.

My understanding of the model is that it is explaining the minority component as noise some of the times. Basically, the data can be explained as (i) two components, or (ii) one component with a wide spread but a peak at the major component. Checking the posterior predictive distributions could perhaps be one way to check.

I also tried initializing the modes to the good modes (kmeans) by passing them through the init parameter to rstan::stan. But the model drifts away in this case too.

What worked was removing the distribution over modes of the beta-binomials, and making them into fixed parameters (again initialized through kmeans) that I pass in as data. I’m probably going ahead with this, and coupling it with model comparison against the one-component version.