Tricking nimble into doing polya-gamma for multicategorical logit

[Ethan A]

The polya-gamma sampler works great for logistic regression. However, it's limited in that it can only handle bernoulli or binomial likelihoods.

I'm trying to trick it using the stick-breaking representation of the categorical distribution (so it boils down to a product of Bernoullis). 

Ultimately, I need this to work where the categories are latent (e.g., a mixture model), so I figured I would try to see if I can get it to work in an i.i.d. setting first. The model compiles just fine, and Nimble recognizes conjugacy everywhere. However, the results don't reflect the truth at all.

I'm not sure if it's a tiny bug, a flaw in the logic, or if this is just a current limitation (that can hopefully be fixed in the future). If it can be figured out, however, it opens the possibility for Nimble to be used with e.g., logit stick breaking processes.

Below is my code:

library(nimble)
code <- nimbleCode({
  ## Stick-breaking probabilities
  for ( k in 1:(K-1) )
    v[k] ~ dbeta(0.5, 0.5)
 
  w[1:K] <- stick_breaking(v[1:(K-1)])
 
  ## Priors for mu
  for ( k in 1:K )
    mu[k] ~ dnorm(0, 1)
 
  for ( i in 1:n ) {
    size[i, 1] <- 1            ## There is always stick for first component
    z[i, 1] ~ dbern(v[1])      ## First stick always random
    for ( k in 2:(K-1) ) {
      size[i, k]  <- size[i, k-1] - z[i, k-1]       ## = 1 if no success drawn yet and 0 if there has
      ztilde[i, k] ~ dbern(v[k])                    ## Latent Bernoulli
      z[i, k]      <- size[i, k] * ztilde[i,k]      ## z[i, k] = 0 if there is no stick left
    }
    y[i] ~ dnorm( inprod(z[i, 1:K], mu[1:K]) , 1)
  }
})

n  = 300
mu = c(-10, 0, 10)
p  = c(0.25, 0.5, 0.25)
z  = rcat(n, p)
y  = rnorm(n, mu[z], 1)


constants = list(n = n, K = length(p))
data = list(y = rnorm(n))
inits = list(v = runif(K-1), mu = rnorm(K), z = cbind(1, matrix(0, n, K-1)) )
inits$ztilde = inits$z
model = nimbleModel(code, constants, data, inits)
cmodel = compileNimble(model)
mcmcConf = configureMCMC(model, monitors = c('mu', 'w'))
mcmc = buildMCMC(mcmcConf)
cmcmc = compileNimble(mcmc, project = cmodel)

fit = runMCMC(cmcmc, 12000, 2000)

[Paul vdb]

Hello,

I see a small bug in your data input:

data = list(y = rnorm(n))

it should be

data = list(y = rnorm(n, mu[z], 1))

With that changed everything runs as it should (minus label switching).

We are planning to include the polya-gamma sampler for categorical distributions in a future release but if you are able to hack it with the product of Bernoulli distributions, please include that solution below.

Thanks,

Paul

[Ethan A]

Ah. Sorry about that. The below code (i.i.d. case) seems to now work with the polya gamma:

code <- nimbleCode({
  ## Stick-breaking probabilities

  for ( k in 1:(K-1) ) {
    beta[k] ~ dnorm(0, 1)     ## variables for polya-gamma data augmentation
    v[k] <- expit(beta[k])    ## stick-breaking variables

  }
 
  w[1:K] <- stick_breaking(v[1:(K-1)])
 
  ## Priors for mu
  for ( k in 1:K )
    mu[k] ~ dnorm(0, 1)
 
  for ( i in 1:n ) {
    size[i, 1] <- 1            ## There is always stick for first component
    z[i, 1] ~ dbern(v[1])      ## First stick always random
    for ( k in 2:(K-1) ) {
      size[i, k]  <- size[i, k-1] - z[i, k-1]       ## = 1 if no success drawn yet and 0 if there has
      ztilde[i, k] ~ dbern(v[k])                    ## Latent Bernoulli
      z[i, k]      <- size[i, k] * ztilde[i,k]      ## z[i, k] = 0 if there is no stick left
    }
    y[i] ~ dnorm( inprod(z[i, 1:K], mu[1:K]) , 1)
  }
})

n  = 300

mu = c(-3, 0, 3)

p  = c(0.25, 0.5, 0.25)
z  = rcat(n, p)
y  = rnorm(n, mu[z], 1)

K  = length(p)

constants = list(n = n, K = length(p))

data = list(y = y)
inits = list(beta = logit(runif(K-1)), mu = rnorm(K), z = cbind(1, matrix(0, n, K-1)) )

inits$ztilde = inits$z
model = nimbleModel(code, constants, data, inits)
cmodel = compileNimble(model)
mcmcConf = configureMCMC(model, monitors = c('mu', 'w'))

mcmcConf$removeSampler('beta')

for ( k in 1:(K-1) )

  mcmcConf$addSampler(paste0('beta[', k, ']'), 'polyagamma')

mcmc = buildMCMC(mcmcConf)
cmcmc = compileNimble(mcmc, project = cmodel)

fit = runMCMC(cmcmc, 12000, 2000)

fit %>% summarize_draws()




HOWEVER, when I turn to a regression setting, I am getting errors. Here's the code:

set.seed(123)
remove(list = ls())
library(nimble)
library(tidyverse)
library(posterior)
code <- nimbleCode({

  ## Priors for mu
  for ( k in 1:K )
    mu[k] ~ dnorm(0, 1)

  for ( k in 1:(K-1) ) {
    beta[1:p, k] ~ dmnorm(beta_mean[1:p], cov = beta_cov[1:p,1:p])
  }
  logit(v[1:n, 1:(K-1)]) <- X[1:n, 1:p] %*% beta[1:p, 1:(K-1)]

  for ( i in 1:n ) {
   
    size[i, 1] <- 1               ## There is always stick for first component

    z[i, 1] ~ dbern(v[i, 1])      ## First stick always random

    for ( k in 2:(K-1) ) {
      size[i, k]   <- size[i, k-1] - z[i, k-1]        ## = 1 if no success drawn yet and 0 if there has

      ztilde[i, k] ~  dbern(v[i, k])                  ## Latent Bernoulli
      z[i, k]      <- size[i, k] * ztilde[i, k]       ## z[i, k] = 0 if there is no stick left

    }
    y[i] ~ dnorm( inprod(z[i, 1:K], mu[1:K]) , 1)
  }
})

n    = 300

p    = 3

mu   = c(-10, 0, 10)

K    = length(mu)
beta = replicate(K-1, rnorm(p))
X    = cbind(1, replicate(K-1, rnorm(n)))
eta  = X %*% beta
prob = t( apply(eta, 1, function(x) stick_breaking(expit(x))) )
z    = sapply(1:n, function(i) rcat(1, prob = prob[i, ]) )

y    = rnorm(n, mu[z], 1)

constants = list(n = n, K = K, p = p)
data = list(y = y, X = X, beta_mean = rep(0, 3), beta_cov = diag(3))
inits = list(

 mu = rnorm(K), z = cbind(1, matrix(0, n, K-1))

  , beta = beta

)
inits$ztilde = inits$z
model        = nimbleModel(code, constants, data, inits)
cmodel       = compileNimble(model)

mcmcConf     = configureMCMC(model, monitors = c('mu', 'beta'))
mcmcConf$removeSampler('beta')
for ( k in 1:(K-1) ) {
  varname = paste0('beta[', ', ', k, ']')
  mcmcConf$addSampler(varname, 'polyagamma', check = FALSE, control = list(fixedDesignColumns=TRUE))

}
mcmc         = buildMCMC(mcmcConf)
cmcmc        = compileNimble(mcmc, project = cmodel)

fit = runMCMC(cmcmc, 12000, 2000)

The error I get is: 

Error in samplerFunction(model = model, mvSaved = mvSaved, target = target,  :
  object 'inflationNodes' not found

If I remove the check = FALSE argument, I get

Error in samplerFunction(model = model, mvSaved = mvSaved, target = target,  :
  polyagamma sampler: response nodes should all be part of the same declaration. If your model is in a non-standard form and you are sure the Pólya-gamma sampler is appropriate, you can disable this check by setting the control argument `check=FALSE`.


I've tried specifying independent normal priors to beta[j,k] and that does not seem to work either.

[Paul vdb]

Hi Ethan,

I ran your code, and you have identified a bug. I've made it an issue on the github: 

Polya-Gamma Sampler unchecked zero-inflation · Issue #1536 · nimble-dev/nimble

In the sampler set up we do the zero inflation calculations only when check=TRUE. However, if check = FALSE and there is zero inflation then an error is thrown. We will try and get it fixed soon, and I will post an update here. Thanks for sharing the example.

Paul

[Paul vdb]

Hi Ethan,

A follow up on your error. I have started drafting the sampler to recognize the stick_breaking transformation for when using dmulti or dcat. But in the meantime, I got your example working.

This line here

logit(v[1:n, 1:(K-1)]) <- X[1:n, 1:p] %*% beta[1:p, 1:(K-1)]

results in our sampler including all logit probabilities in the probability nodes and it becomes falsely identified as zero inflated.

Defining the logit probabilities in a loop removes that issue and the example you sent runs without a problem.

  for( i in 1:(K-1) )

    logit(v[1:n, i]) <- X[1:n, 1:p] %*% beta[1:p, i]

Another small note is to put check = FALSE in the control list: 

mcmcConf$addSampler(varname, 'polyagamma', control = list(fixedDesignColumns=TRUE,  check = FALSE))

Let me know if this works for you or if you have any other questions.

Paul

[Ethan A]

That works! Thanks a lot.