如何使用 MCMC 分解混合分布

How to decompose a mixed distribution using MCMC

我的数据是 50:50 正态分布和常数值的混合:

numdata = 10000
data = np.random.normal(0.0,1.0,numdata).astype(np.float32)
data[int(numdata/2):] = 0.0
plt.hist(data,30,density=True)

我的任务是为该数据拟合混合密度。 我正在使用 tfd.Mixture 与 tfd.Normal 和 tfd.Deterministic Normal 与 Deterministic 的已知(在示例数据的情况下)比率为 0.5 我的 MCMC 相反 returns 0.83 的配给比正常。

是否有更好的方法来以正确的比例拟合此分布?

完整的示例代码如下:

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
import tensorflow as tf
import tensorflow_probability as tfp
import matplotlib.pyplot as plt
tfd = tfp.distributions
tfb = tfp.bijectors

import numpy as np
from time import time

numdata = 10000
data = np.random.normal(0.0,1.0,numdata).astype(np.float32)
data[int(numdata/2):] = 0.0
_=plt.hist(data,30,density=True)

root = tfd.JointDistributionCoroutine.Root
def dist_fn(rv_p,rv_mu):
    rv_cat = tfd.Categorical(probs=tf.stack([rv_p, 1.-rv_p],-1))
    rv_norm  = tfd.Normal(rv_mu,1.0)
    rv_zero =  tfd.Deterministic(tf.zeros_like(rv_mu))
    
    rv_mix = tfd.Independent(
                tfd.Mixture(cat=rv_cat,
                            components=[rv_norm,rv_zero]),
                reinterpreted_batch_ndims=1)
    return rv_mix


def model_fn():
    rv_p    = yield root(tfd.Sample(tfd.Uniform(0.0,1.0),1))
    rv_mu   = yield root(tfd.Sample(tfd.Uniform(-1.,1. ),1))
    
    rv_mix  = yield dist_fn(rv_p,rv_mu)
    
jd = tfd.JointDistributionCoroutine(model_fn)
unnormalized_posterior_log_prob = lambda *args: jd.log_prob(args + (data,))

n_chains = 1

p_init = [0.3]
p_init = tf.cast(p_init,dtype=tf.float32)

mu_init = 0.1
mu_init = tf.stack([mu_init]*n_chains,axis=0)

initial_chain_state = [
    p_init,
    mu_init,
]

bijectors = [
    tfb.Sigmoid(),  # p
    tfb.Identity(),  # mu
]

step_size = 0.01

num_results = 50000
num_burnin_steps = 50000


kernel=tfp.mcmc.TransformedTransitionKernel(
    inner_kernel=tfp.mcmc.HamiltonianMonteCarlo(
    target_log_prob_fn=unnormalized_posterior_log_prob,
    num_leapfrog_steps=2,
    step_size=step_size,
    state_gradients_are_stopped=True),
    bijector=bijectors)

kernel = tfp.mcmc.SimpleStepSizeAdaptation(
    inner_kernel=kernel, num_adaptation_steps=int(num_burnin_steps * 0.8))

#XLA optim
@tf.function(autograph=False, experimental_compile=True)
def graph_sample_chain(*args, **kwargs):
  return tfp.mcmc.sample_chain(*args, **kwargs)


st = time()
trace,stats = graph_sample_chain(
      num_results=num_results,
      num_burnin_steps=num_burnin_steps,
      current_state=initial_chain_state,
      kernel=kernel)
et = time()
print(et-st)


ptrace, mutrace = trace
plt.subplot(121)
_=plt.hist(ptrace.numpy(),100,density=True)
plt.subplot(122)
_=plt.hist(mutrace.numpy(),100,density=True)
print(np.mean(ptrace),np.mean(mutrace))

p 和 mu 的结果分布如下:

显然,p 的平均值应该为 0.5 我怀疑 model_fn() 可能有问题。 我尝试在不同的 p 值下评估模型的 log_prob,实际上“最佳”约为 0.83 我只是不明白为什么以及如何修复它以重建原始混合物。

[编辑] 使用 pymc3 的“更简单”的演示代码。仍然是相同的行为,结果是 0.83 而不是 0.5

import pymc3 as pm
import numpy as np
import arviz as az
import matplotlib.pyplot as plt


numdata = 1000
data1 = np.random.normal(0.0,1.0,numdata).astype(np.float32)
data2 = np.zeros(numdata).astype(np.float32)
data = np.concatenate((data1,data2))


_=plt.hist(data,30,density=True)

with pm.Model() as model:
    norm = pm.Normal.dist(0.0,1.0)
    zero = pm.Constant.dist(0.0)
    
    components = [norm,zero]
    w = pm.Dirichlet('p', a=np.array([1, 1]))  # two mixture component weights.
    like = pm.Mixture('data', w=w, comp_dists=components, observed=data)
    
    posterior = pm.sample()
    
    idata = az.from_pymc3(posterior)
    az.plot_posterior(posterior)

概率密度与质量的不可公度性

这里的问题是来自每个模型的可能性涉及高斯的概率密度和离散的质量,这是不相称的。具体来说,用于比较零观察来自何处的计算将涉及似然

P[x=0|Normal[0,1]] = 1/sqrt(2*pi) = 0.3989422804014327
P[x=0|   Zero    ] = 1

这将比较这些(由 p 加权),就好像它们具有相同的单位一样。然而,前者是密度,因此相对于后者来说是无穷小的。如果忽略这种不可通约性,那么实际上就像高斯有 40% 的机会生成零,而实际上它 almost never 恰好生成一个零。

解决方法:伪离散分布

我们需要以某种方式转换单位。一种简单的方法是用连续分布来近似离散分布,这样它生成的可能性就以密度单位表示。例如,使用以离散值为中心的高精度(窄)高斯或拉普拉斯分布会产生以 0.5:

为中心的 p 后验
with pm.Model() as model:
    norm = pm.Normal.dist(0.0, 1.0)
    pseudo_zero = pm.Laplace.dist(0.0, 1e-16)
    
    components = [norm, pseudo_zero]
    w = pm.Dirichlet('p', a=np.array([1, 1]))  # two mixture component weights.
    like = pm.Mixture('data', w=w, comp_dists=components, observed=data)
    
    posterior = pm.sample()
    
    idata = az.from_pymc3(posterior)
    az.plot_posterior(posterior)


为什么 p=0.83

我们在混合离散和连续时观察到的后验不是任意的。这里有几种获取它的方法。下面我们就用一个p来表示来自高斯的概率

MAP 估计值

忽略不可通约性,我们可以推导出 p 的 MAP 估计如下。让我们用 D = { D_1 | D_2 } 表示组合观测值,其中 D_1 是来自高斯等的子集,而 n 是来自每个子集的观测值的数量。然后我们可以写出似然

P[p|D] ~ P[D|p]P[p]

由于 Dirichlet 是均匀的,我们可以忽略 P[p] 并扩展我们的数据

P[D|p] = P[D_1|p]P[D_2|p]
       = (Normal[D_1|0,1]*(p^n))(Normal[0|0,1]*p + 1*(1-p))^n
       = Normal[D_1|0,1]*(p^n)(0.3989*p + 1 - p)^n
       = Normal[D_1|0,1]*(p - 0.6011*(p^2))^n

取导数w.r.t。 p 并设置为零,我们有

0 = n*(1-1.2021*p)(p-0.6011*p^2)^(n-1)

p = 1/1.2021 = 0.8318669 处取一个(非平凡的)零。

抽样思维实验

解决这个问题的另一种方法是通过抽样实验。假设我们使用下面的方案来采样p.

  1. 从给定的 p.
  2. 开始
  3. 对于每个观察值,使用两个模型的可能性绘制伯努利样本,并由先前的 p 值加权。
  4. 计算一个新的 p 作为所有这些伯努利平局的平均值。
  5. 转到第 1 步。

本质上,p 的吉布斯采样器,但对不可能的观察模型分配具有鲁棒性。

对于第一次迭代,让我们从 p=0.5 开始。对于真正来自高斯的所有观察结果,它们对于离散模型的可能性为零,因此,至少,我们所有伯努利抽取的一半将为 1(对于高斯)。对于来自离散模型的所有观察结果,这将是对每个模型中观察到零的可能性的比较。离散模型为 1,高斯模型为 0.3989422804014327。对此进行归一化,意味着我们有伯努利抽签的概率为

p*0.3989/((1-p)*1 + p*0.3989)
# 0.2851742248343187

赞成高斯分布。现在我们可以更新 p,这里我们只使用预期值,即:

p = 0.5*1 + 0.5*0.2851742248343187
# 0.6425871124171594

如果我们开始迭代这个会发生什么?

# likelihood for zero from normal
lnorm = np.exp(pm.Normal.dist(0,1).logp(0).eval())

# history
p_n = np.zeros(101)

# initial value
p_n[0] = 0.5

for i in range(100):
    # update
    p_n[1+i] = 0.5 + 0.5*p_n[i]*lnorm/((1-p_n[i])+p_n[i]*lnorm)

plt.plot(p_n);
p_n[100]
# 0.8318668635076404

同样,预期值收敛到我们的后验平均值 p = 0.83

因此,抛开 PDF 和 PMF 的代码域单位不同这一事实,TFP 和 PyMC3 似乎都在正常运行。