PYMC3:NUTS 难以从分层零膨胀伽玛模型中采样

PYMC3: NUTS has difficulty sampling from a hierarchical zero inflated gamma model

我正在尝试复制 a paper from Richard McElreath 中的数据分析,他在其中使用分层零膨胀伽玛模型对数据进行了拟合。数据是关于 returns 大约 15000 次狩猎旅行的 150 猎人超过二十年。因为很多狩猎旅行都是零 returns,模型假设每次旅行都有 pi 为零的概率 returns,以及 1 - pi 为正 returns 的概率具有参数 alphabeta.

的 Gamma 分布

预测变量是年龄,该模型使用年龄多项式(高达 3 阶)对 pialpha 进行建模。由于 15000 次旅行属于 150 个猎人,每个猎人都有自己的系数,并且所有系数都遵循共同的多元正态分布。模型详情请参考以下代码。 model specification 貌似没问题,但是NUTS开始采样时遇到了问题:大约20分钟后只给出了大约10个样本,采样器就停在那里,并告诉我需要数百小时才能完成采样。我想知道是什么导致了这些问题。

通常的进口

import pymc3 as pm
import numpy as np
from pymc3.distributions import Continuous, Gamma
import theano.tensor as tt

数据可以从github

获取
n_trip = len(d)
n_hunter = len(d['hunter.id'].unique())
idx_hunter = d['hunter.id'].values

y = d['kg.meat'].values
age = d['age.s'].values
age2 = (d['age.s'].values)**2
age3 = (d['age.s'].values)**3

零膨胀 Gamma 的对数概率密度函数。

class ZeroInflatedGamma(Continuous):
    def __init__(self, alpha, beta, pi, *args, **kwargs):
        super(ZeroInflatedGamma, self).__init__(*args, **kwargs)
        self.alpha = alpha
        self.beta = beta
        self.pi = pi = tt.as_tensor_variable(pi)
        self.gamma = Gamma.dist(alpha, beta)

    def logp(self, value):
        return tt.switch(value > 0,
                         tt.log(1 - self.pi) + self.gamma.logp(value),
                         tt.log(self.pi))

这是一个索引9X9矩阵之前的相关矩阵的矩阵,pymc3中的LKJ先验是作为一维向量给出的

dim = 9
n_elem = dim * (dim - 1) / 2
tri_index = np.zeros([dim, dim], dtype=int)
tri_index[np.triu_indices(dim, k=1)] = np.arange(n_elem)
tri_index[np.triu_indices(dim, k=1)[::-1]] = np.arange(n_elem)

这是模型

with pm.Model() as Vary9_model:

    # hyper-priors
    mu_a = pm.Normal('mu_a', mu=0, sd=100, shape=9)
    sigma_a = pm.HalfCauchy('sigma_a', 5, shape=9)

    # build the covariance matrix
    C_triu = pm.LKJCorr('C_triu', n=2, p=9)    
    C = tt.fill_diagonal(C_triu[tri_index], 1)
    sigma_diag = tt.nlinalg.diag(sigma_a)
    cov = tt.nlinalg.matrix_dot(sigma_diag, C, sigma_diag)

    # priors for each hunter and all the linear components, 9 dimensional Gaussian  
    a = pm.MvNormal('a', mu=mu_a, cov=cov, shape=(n_hunter, 9))

    # linear function  
    mupi = a[:,0][idx_hunter] + a[:,1][idx_hunter] * age + a[:,2][idx_hunter] * age2 + a[:,3][idx_hunter] * age3
    mualpha = a[:,4][idx_hunter] + a[:,5][idx_hunter] * age + a[:,6][idx_hunter] * age2 + a[:,7][idx_hunter] * age3

    pi = pm.Deterministic('pi', pm.math.sigmoid(mupi))
    alpha = pm.Deterministic('alpha', pm.math.exp(mualpha))
    beta = pm.Deterministic('beta', pm.math.exp(a[:,8][idx_hunter]))

    y_obs = ZeroInflatedGamma('y_obs', alpha, beta, pi, observed=y)

    Vary9_trace = pm.sample(6000, njobs=2)

这是模型的状态:

Auto-assigning NUTS sampler...
Initializing NUTS using advi...
Average ELBO = -28,366: 100%|██████████| 200000/200000 [15:36<00:00, 213.57it/s]
Finished [100%]: Average ELBO = -28,365
  0%|          | 22/6000 [15:51<63:49:25, 38.44s/it]

我对这个问题有一些想法,但不确定可能是哪个原因。

或者是其他原因?

附带说明一下,我尝试使用 Metropolis 采样器,我的计算机 运行 内存不足并且链仍然没有收敛。

ZeroInflatedGamma 看起来不错。密度函数关于 pi、alpha 和 beta 是可微分的。这就是观察变量所需的全部。如果您要估算值,则只需要相对于值的导数。

LKJCorr 的实现出现问题: https://github.com/pymc-devs/pymc3/pull/1863 你可以在master上再试一次。遗憾的是,pymc3 不支持在 cholesky 分解参数化中使用 MVNormal 和 LKJCorr。这也可能有帮助。 github 上有一个正在进行的拉取请求: https://github.com/pymc-devs/pymc3/pull/1875

为了提高收敛性,您可以尝试 a 的非中心参数化。类似于

a_raw = pm.Normal('a_raw', shape=(9, n_hunter))
a = mu_a[None, :] + tt.dot(tt.slinalg.cholesky(cov), a_raw)

当然,如果我们有胆小的 LKJCorr,这会更快...