python - PYMC3:NUTS 难以从分层零膨胀 Gamma 模型中采样

Question

我正在尝试从 a paper from Richard McElreath 复制数据分析，其中他用分层零膨胀 Gamma 模型拟合数据。该数据是关于 15000 次狩猎之旅的狩猎返回，来自大约 150 猎人超过 20 年。由于很多狩猎旅行的返回为零，该模型假设每次旅行都有 pi 的零返回概率和 1 - pi 的正返回概率，其服从 Gamma 分布参数 alpha 和 beta。

预测变量是年龄，该模型使用年龄多项式(最多 3 阶)来建模 pi 和 alpha。由于 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]

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

• 九维高斯太难采样了吗？我以前只将 mualpha 和 mupi 的截距建模为双变量高斯分布，它很慢但有效(模型拟合大约需要 20 分钟)

• 是概率密度导致了问题吗？我自己编写了密度函数，但不确定它是否运行良好。我认为密度函数在零处不可微，这会给坚果采样器带来麻烦吗？

• 是因为预测变量高度相关吗？该模型中的线性模型组件是年龄的三次多项式，自然预测变量高度相关。

还是因为其他原因？

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

Answer 1

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，这会更快......