Non-centered parameterization for hierarchical mixture model?

Modeling
1

Background

I’m reproducing the code from https://onlinelibrary.wiley.com/doi/10.1002/pst.1730 in Stan (originally written with BUGS).

The basic data setup is that there are J = 10 clinical trials of varying sample size, each with an unknown log-odds \theta_j of success. We observe r_j successes out of n_j total patients in each trial.

The model is a hierarchical mixture model where with probability p_\text{exch}, \theta_j is exchangeable with the other strata (partial pooling) and with probability 1 - p_\text{exch}, \theta_j is non-exchangeable with the other strata (no pooling).

By setting p_\text{exch} to 0 for each strata we obtain a completely stratified analysis, the exchangeable/non-exchangeable (EXNEX) mixture with 0.5, and the standard hierarchical model partial pooling with when setting p_\text{exch} to 1.

I am able to reproduce the results of the original paper with my below code. However, I am running into some convergence issues. The model runs fine with p_\text{exch} = 0 and p_\text{exch} = 0.5 but the difficulties arise when p_\text{exch} is close to 1. I rewrote the model for the special case when p_\text{exch} = 1 using the non-centered parameterization (code also included below), and the model ran much better.

My question is, using the mixture model formulation shown in the code below, how can I implement the non-centered parameterization? Any guidance would be very helpful.

EXNEX Stan Code

data {
  int<lower=1> J; // Number of trial strata
  array[J] int n; // Number of subjects per strata
  array[J] int r; // Number of successes per strata
  array[J] real<lower=0, upper=1> p_exch; // Mixture probability of exchangeability for each strata
}
parameters {
  real mu;               // Mean of hierarchical log-odds distribution 
  real<lower=0> tau;     // SD of hierarchical log-odds distribution
  vector[J] theta;       // Log-odds for each strata
}
model {
  // Priors
  mu ~ normal(-1.73, 2.616);     // Centered at logit(0.15)
  tau ~ normal(0, 1);            // Half-normal prior because of the constraint

  for (j in 1:J) {  // Loop over the trial strata
    target += log_mix(p_exch[j],                             // Strata-specific mixing parameter
                      normal_lpdf(theta[j] | mu, tau),       // EX
                      normal_lpdf(theta[j] | -1.734, 2.801)  // NEX
                      );
      
    r[j] ~ binomial_logit(n[j], theta[j]);                 // Binomial likelihood with built-in expit transform
  }
}

Special case for p_exch = 1 with non-centered parameterization

data {
  int<lower=1> J; // Number of trial strata
  array[J] int n; // Number of subjects per strata
  array[J] int r; // Number of successes per strata
}
parameters {
  real mu;             // Mean of hierarchical log-odds distribution 
  real<lower=0> tau;   // SD of hierarchical log-odds distribution 
  vector[J] theta_std; // Log-odds of response for each strata (standardized for non-centered parameterization)
}
model {
  for (j in 1:J) {                                  // Loop over the trial strata
    theta_std[j] ~ normal(mu, 1);                   // Hierarchical structure + use non-centered parameterization
    r[j] ~ binomial_logit(n[j], theta_std[j]*tau);  // Binomial likelihood with built-in expit transform
  }
  // Priors
  mu ~ normal(-1.73, 2.616);     // Centered at logit(0.15)
  tau ~ normal(0, 1);            // Half-normal prior because of the constraint
}