ADVI, NUTS and Metropolis produce significantly different results

[kayhan-batmanghelich]

Dear PyMC developers,

I am trying to develop simple mixed effect model using pymc (see the code below). I have tried NUTS, ADVI, Metropolis for inference. Aside from variations in speed for this model (Time : AVDI ~= Metrolopis << NUTS ), the results are significantly different. To be more specific, I am interested to estimate $h^2$ which is basically proportion of variance explained by the random effect.

This is the code:

import numpy as np, pymc3 as pm, theano.tensor as T


mixedEffect_model = pm.Model()

with pm.Model() as mixedEffect_model:
    ### hyperpriors
    tau = pm.Gamma('tau', 1.e-3, 1.e-3)
    sigma2 = 1.0/tau
    h2 = pm.Uniform('h2')
    beta_0 = pm.Uniform('beta_0', lower=-1000, upper=1000)   # a replacement for improper prior
    w = pm.Uniform('w', lower=-1000, upper=1000, shape=(M,1))   # a replacement for improper prior

    z = pm.Normal('z', mu=0, sd= (h2*sigma2)**0.5 , shape=(N,1))
    g = T.dot(L,z)

    y = pm.Normal('y', mu = beta_0+g + T.dot(X,w), sd= ((1-h2)*sigma2)**0.5 , shape=(N,1) , observed=y )

ADVI:

with mixedEffect_model:
    v_params = pm.variational.advi(n=5000)

with mixedEffect_model:
    trace = pm.variational.sample_vp(v_params, draws=5000)

_ = hist(trace['h2'],100)

Metropolis:

with mixedEffect_model:
    start = pm.find_MAP()
    step = pm.Metropolis()
    trace = pm.sample(5000, step, start=start)

_ = hist(trace['h2'],100)

NUTS (slow):

with mixedEffect_model:
    start = pm.find_MAP()
    step = pm.NUTS()
    trace = pm.sample(5000, step, start=start)

_ = hist(trace['h2'],100)    

Given that w and beta_0 are treated as parameter with (approximately) improper prior, I expect the results to be very close to maximum likelihood estimation. Using maximum likelihood estimation (more specifically Restricted maximum Likelihood estimation), I expect values around 0.5 for $h^2$ with standard error of around 0.07. The results is very different. It is very different from ML and also very different across different inference engines. NUTS is closer to ML but given that the results are drastically different, I don't know how to explain it.

Any idea?

[fonnesbeck]

This is entirely possible. Metropolis can mix really poorly sometimes, and ADVI can give really bad approximations. You would have a better idea if something is wrong or not if you looked at some convergence diagnostics. If the ELBO in ADVI has not converged to a stationary value and the MCMC samplers have not converged, then it is premature to compare them.

Also, in general, you should not compare NUTS and Metropolis based on sampling speed. NUTS is vastly more efficient in terms of effective sample size in the resulting trace than Metropolis is.

BTW, you appear to be overwriting your data (y) with the likelihood itself (also y) in the last line of the model. You should probably avoid that.

[kayhan-batmanghelich]

The ELBO in ADVI stabilizes very quickly:

%time means, sds, elbos = pm.variational.advi( \
    model=mixedEffect_model, n=5000, learning_rate=1e-1)


Iteration 0 [0%]: ELBO = -55115425.64
Iteration 500 [10%]: ELBO = -20841.0
Iteration 1000 [20%]: ELBO = -19739.33
Iteration 1500 [30%]: ELBO = -28763.83
Iteration 2000 [40%]: ELBO = -55537.09
Iteration 2500 [50%]: ELBO = -20212.82
Iteration 3000 [60%]: ELBO = -20057.14
Iteration 3500 [70%]: ELBO = -20062.19
Iteration 4000 [80%]: ELBO = -20067.08
Iteration 4500 [90%]: ELBO = -20674.95
Finished [100%]: ELBO = -20232.54
CPU times: user 44min 19s, sys: 47.7 s, total: 45min 7s
Wall time: 1min 18s

Metropolis seems to suffer from bad initialization. h2 stays at zero and does not move at all:

with mixedEffect_model:
    start = pm.find_MAP()
    step = pm.Metropolis()
    trace_mp = pm.sample(5000, step, start=start)

Is there anyway to use the mean of VB to initialize the metropolis?

[fonnesbeck]

What do Metropolis and NUTS do when you do not initialize it with the MAP? Any better? I don't usually have much success with using it to initialize. If you are going to use it, you might pass scaling=start to NUTS as well.

You may also get better mileage out of specifying sigma directly rather than tau, unless you are truly interested in the precision. See what you get out of a HalfCauchy(5) on sigma.

[kayhan-batmanghelich]

Not using start didn't improve the results for Metropolis. Meanwhile I realized that I had intercept in the X, so I removed beta_0:

This is a new code (I followed your advice ):

import numpy as np, pymc3 as pm, theano.tensor as T


mixedEffect_model = pm.Model()

with pm.Model() as mixedEffect_model:
    ### hyperpriors
    sigma2 = pm.InverseGamma('sigma2', 1.e-3, 1.e-3)
    #tau = pm.Gamma('tau', 1.e-3, 1.e-3)
    #sigma2 = 1.0/tau
    h2 = pm.Uniform('h2')
    #beta_0 = pm.Uniform('beta_0', lower=-1000, upper=1000)   # a replacement for improper prior
    w = pm.Uniform('w', lower=-1000, upper=1000, shape=(M,1))   # a replacement for improper prior

    z = pm.Normal('z', mu=0, sd= (h2*sigma2)**0.5 , shape=(N,1))
    g = T.dot(L,z)

    #y = pm.Normal('y', mu = beta_0+g + T.dot(X,w), sd= ((1-h2)*sigma2)**0.5 , shape=(N,1) , observed=Pheno )
    y = pm.Normal('y', mu = g + T.dot(X,w), sd= ((1-h2)*sigma2)**0.5 , shape=(N,1) , observed=Pheno )

NUTs is on the right path (the value matches ML estimation) but judging based on the traceplot, it seems that it is not quite converged yet. Perhaps more iteration helps but it is quite slow (5000 iteration takes more an hour):

with mixedEffect_model:
    start = pm.find_MAP()
    step = pm.NUTS()
    trace_nuts = pm.sample(5000, step, start=start)


figure()
_ = hist(trace_nuts['h2'],100)

pm.traceplot(trace_nuts, varnames=['h2'])

Metropolis still does not move:

with mixedEffect_model:
    #start = pm.find_MAP()
    step = pm.Metropolis()
    #trace_mp = pm.sample(5000, step, start=start)
    trace_mp = pm.sample(5000, step)

figure()
_ = hist(trace_mp['h2'],100)

pm.traceplot(trace_mp, varnames=['h2'])

ADVI is pretty off too:

v_params = pm.variational.advi( \
    model=mixedEffect_model, n=5000, learning_rate=1e-1)

with mixedEffect_model:
    trace = pm.variational.sample_vp(v_params, draws=5000)

_ = hist(trace['h2'],100)
print mean(trace['h2'])    

[junpenglao]

@kayhan-batmanghelich
I am also currently experimenting with different ways to fit a mixed-effect model in PyMC3. In my case, I found that NUTS is very sensitive to starting value, and using the value return from find_MAP() doesn't seem to work. Moreover, for my model I need around 50000 samples and the trace start to convert around 20000 samples in.
You might find https://github.com/paul-buerkner/brms helpful, it's a R package for mixed-effect model using STAN. One of the tricks it used is to center the design matrix X, remove the intercept from X, and model the population-level intercept as an additional term:
y = pm.Normal('y', mu = intercept + g + T.dot(X_centered,w), sd= ((1-h2)*sigma2)**0.5 , shape=(N,1) , observed=Pheno )

[kayhan-batmanghelich]

@junpenglao
Thank you for your note. I will try that idea (removing the intercept and adding independently). Theoretically, the results should be the same so I will be surprised if it is different. I will also try replacing sigma2 ~ inv-gamma with something like sigma ~ half-cauchy as suggest by Andrew Gelman (what is half cauchy in PyMC3 by the way?).

I noticed that if I run NUTS for about 20,000 iterations, the posterior value is close the maximum likelihood estimation but that takes very very long. I couldn't get ADVI to produce a reasonable result. Metropolis also does not produces a good results; I tried initializing it with ADVI values, increasing iteration, no success.

Overall it is very surprising. This is a basic mixed effect model perhaps a bit big but really nothing special about it. My final goal is not to solve this specific mixed effect model but rather using it as stepping stone to train more complicated model. I am very interested to get PyMC3 working since it helps me to focus more on the model rather than inference algorithm.

[junpenglao]

@kayhan-batmanghelich

Theoretically, the results should be the same so I will be surprised if it is different.

Actually, it makes a difference at least in the way brm models it, and in PyMC the model estimation is much closer to a classical mixed model as return from statsmodels. However, I still having difficulty to hack it to produce the same result. On the other hand the brm estimation is nearly identical to lme4...

And for me, only Metropolis works but not NUTS and ADVI...

My final goal is not to solve this specific mixed effect model but rather using it as stepping stone to train more complicated model. I am very interested to get PyMC3 working since it helps me to focus more on the model rather than inference algorithm.

I am trying to do the same thing - if you want we can share the codes and try to solve this together.

[kayhan-batmanghelich]

@junpenglao
I would be happy to discuss it, shoot me email and we can skype.

I found STAN very slow for this problem. This is a basic model, there are not much to change in the model. This is my model:

Other than numerical issue, I don't have any other explanation why it doesn't work.

[datnamer]

How did STAN's result compare?

[kayhan-batmanghelich]

@datnamer For STAN, I tried both NUTS and ADVI. NUTS stopped at the warmup iteration (iteration 1) for 8 hours, so I stopped it. Full-rank ADVI is relatively fast (couple of hours) but does not converge to a correct answer.

[junpenglao]

@kayhan-batmanghelich
Your model actually runs fine using some generated data, I tried with the following code and it works with both Metropolis and NUTS:

import numpy as np, pymc3 as pm, theano.tensor as T, matplotlib.pyplot as plt

M    = 6  # number of columns in X - fixed effect
N    = 10 # number of columns in L - random effect
nobs = 10

# generate design matrix using patsy
from patsy import dmatrices
import pandas as pd
predictors = []
for s1 in range(N):
    for c1 in range(2):
        for c2 in range(3):
            for i in range(nobs):
                predictors.append(np.asarray([c1+1,c2+1,s1+1]))
tbltest             = pd.DataFrame(predictors, columns=['Condi1', 'Condi2', 'subj'])
tbltest['Condi1']   = tbltest['Condi1'].astype('category')
tbltest['Condi2']   = tbltest['Condi2'].astype('category')
tbltest['subj']     = tbltest['subj'].astype('category')
tbltest['tempresp'] = np.random.normal(size=(nobs*M*N,1))

Y, X    = dmatrices("tempresp ~ Condi1*Condi2", data=tbltest, return_type='matrix')
Terms   = X.design_info.column_names
_, L    = dmatrices('tempresp ~ -1+subj', data=tbltest, return_type='matrix')
X       = np.asarray(X) # fixed effect
L       = np.asarray(L) # mixed effect
Y       = np.asarray(Y) 
# generate data
w0 = [5,1,2,3,1,1]
z0 = np.random.normal(size=(N,))
Pheno   = np.dot(X,w0) + np.dot(L,z0) + Y.flatten()
#%%
with pm.Model() as mixedEffect_model:
    ### hyperpriors
    h2     = pm.Uniform('h2')
    sigma2 = pm.HalfCauchy('eps', 5)
    #beta_0 = pm.Uniform('beta_0', lower=-1000, upper=1000)   # a replacement for improper prior
    w = pm.Normal('w', mu = 0, sd = 100, shape=M)
    z = pm.Normal('z', mu = 0, sd= (h2*sigma2)**0.5 , shape=N)
    g = T.dot(L,z)
    y = pm.Normal('y', mu = g + T.dot(X,w), 
                  sd= ((1-h2)*sigma2)**0.5 , observed=Pheno )

with mixedEffect_model:
    # means, sds, elbos = pm.variational.advi(n=10000)
    # trace = pm.sample(50000,step=pm.Metropolis())
    trace = pm.sample(3000)

I think you model might be ill-condition for your data, you should try to reparameterize the model or perform some preprocessing.
In any case, @kayhan-batmanghelich i will send you an email.

[datnamer]

Glad we are atleast keeping up with stan.

[hvasbath]

I would be interested how atmcmc performs with your model. As it is the perfect case with the multiply peaked distribution. Would you be interested to try it?
You would need to add a likelihood variable to your model that contains the model likelihood and pass it to the input. Also you would need an output folder where to store the traces. You can follow the example script ATMIP_2gaussians in the example folder in how to call it:

import pymc3 as pm
import numpy as np
from pymc3.step_methods import ATMCMC as atmcmc
import theano.tensor as tt
from matplotlib import pylab as plt

test_folder = ('ATMIP_TEST')

n_chains = 500
n_steps = 100
tune_interval = 25
njobs = 1

n = 4

mu1 = np.ones(n) * (1. / 2)
mu2 = -mu1

stdev = 0.1
sigma = np.power(stdev, 2) * np.eye(n)
isigma = np.linalg.inv(sigma)
dsigma = np.linalg.det(sigma)

w1 = stdev
w2 = (1 - stdev)


def two_gaussians(x):
    log_like1 = - 0.5 * n * tt.log(2 * np.pi) \
                - 0.5 * tt.log(dsigma) \
                - 0.5 * (x - mu1).tt.dot(isigma).dot(x - mu1)
    log_like2 = - 0.5 * n * tt.log(2 * np.pi) \
                - 0.5 * tt.log(dsigma) \
                - 0.5 * (x - mu2).T.dot(isigma).dot(x - mu2)
    return tt.log(w1 * tt.exp(log_like1) + w2 * tt.exp(log_like2))

with pm.Model() as ATMIP_test:
    X = pm.Uniform('X',
                   shape=n,
                   lower=-2. * np.ones_like(mu1),
                   upper=2. * np.ones_like(mu1),
                   testval=-1. * np.ones_like(mu1),
                   transform=None)
    like = pm.Deterministic('like', two_gaussians(X))
    llk = pm.Potential('like', like)

with ATMIP_test:
    step = atmcmc.ATMCMC(n_chains=n_chains, tune_interval=tune_interval,
                         likelihood_name=ATMIP_test.deterministics[0].name)

trcs = atmcmc.ATMIP_sample(
                        n_steps=n_steps,
                        step=step,
                        njobs=njobs,
                        progressbar=True,
                        trace=test_folder,
                        model=ATMIP_test)

pm.summary(trcs)
Pltr = pm.traceplot(trcs, combined=True)
plt.show(Pltr[0][0])

[junpenglao]

@hvasbath Sure I will give it a go. However I think there is some mistakes in the ATMIP_2gaussians:

def two_gaussians(x):
    log_like1 = - 0.5 * n * tt.log(2 * np.pi) \
                - 0.5 * tt.log(dsigma) \
                - 0.5 * (x - mu1).tt.dot(isigma).dot(x - mu1)
    log_like2 = - 0.5 * n * tt.log(2 * np.pi) \
                - 0.5 * tt.log(dsigma) \
                - 0.5 * (x - mu2).T.dot(isigma).dot(x - mu2)
    return tt.log(w1 * tt.exp(log_like1) + w2 * tt.exp(log_like2))

As 0.5 * (x - mu1).tt.dot(isigma).dot(x - mu1) is not allowed by Theano