How to improve ESS and R_hat for DPMM?

4

Glancing over the references, it looks like they use a Gibbs sampler based on a Chinese Restaurant Process for sampling DPs in general (check algorithm 1 from the 2nd paper). Specifically for the profile regression, they do “metropolis-in-gibbs”, which is an old school way to say they have a BlockedStep, where certain variables are updated by gibbs, and others are updated by metropolis.

Focusing on the cluster labels, means, and covarainces, it appears the setup is conditionally conjugate for each data point’s label assignment, given all other label assignments and fixed observation noise( \sigma_y in the paper). Playing around with the sampler written in R in Appendix B of that paper (I just threw it into GPT and asked for a python version). It seems like it does not suffer from label switching, so having this as a joint sampling step over the implicated variables would be nice.

There has been some work on DPs here and here by @larryshamalama . I think another approach though, given the conditional conjugacy, would be to add a DP case to the machinery proposed in this PR by @ricardoV94, which allows “optimized sampling” by first exploiting conjugate relationships. For this, we’d have to first write the model in an ‘unmarginalized’ way:

with pm.Model(coords=coords) as model:
    x_data = pm.Data("x_data", simdat, dims="obs_num")

    alpha = pm.Gamma("alpha", alpha=2, beta=0.75)
    w = pm.StickBreakingWeights("w", alpha=alpha, K=K_minus_1, dims="cluster")
    mu = pm.Normal("mu", mu=9, sigma=6, dims="cluster")
    sigma = pm.Gamma("sigma", alpha=2, beta=1, dims="cluster")

    class_idx = pm.Categorical('class_idx', p=w, dims='cluster')
    obs = pm.Normal("obs", mu=mu[class_idx], sigma=sigma['class_idx'] observed=x_data, dims="obs_num")

We could then look for the DP in this model, defined as a:

  1. Normal likelihood, with;
  2. Mu and Sigma indexed, and;
  3. The index is a categorical variable, and;
  4. The weights of the categorical variable are stick breaking weights

This is what would signal a DP, and we could then sample all of mu, sigma, components, obs jointly with the CRP.