Tensorflow Probability 中多项式模型的规范

Question

我正在使用 Tensorflow Probability 中的混合多项式离散选择模型。该函数应在 3 个备选方案中进行选择。所选择的备选方案由 CHOSEN（#observationsx3 张量）指定。下面是对代码的更新，以反映我在问题上的进展（但问题仍然存在）。

我目前收到错误：

tensorflow.python.framework.errors_impl.InvalidArgumentError: Incompatible shapes: [6768,3] vs. [1,3,6768] [Op:Mul]

回溯表明问题出在对联合分发的最终组件的 log_prob() 的调用中（即 tfp.Independent(tfp.Multinomial(...) )

主要组件是（感谢 Padarn Wilson 帮助修复联合分布定义）：

@tf.function
def affine(x, kernel_diag, bias=tf.zeros([])):
  """`kernel_diag * x + bias` with broadcasting."""
  kernel_diag = tf.ones_like(x) * kernel_diag
  bias = tf.ones_like(x) * bias
  return x * kernel_diag + bias

def mmnl_func():
    adj_AV_train = (tf.ones(num_idx) - AV[:,0]) * tf.constant(-9999.)
    adj_AV_SM = (tf.ones(num_idx) - AV[:,1]) * tf.constant(-9999.)
    adj_AV_car = (tf.ones(num_idx) - AV[:,2]) * tf.constant(-9999.)

    return tfd.JointDistributionSequential([
        tfd.Normal(loc=0., scale=1e5),  # mu_b_time
        tfd.HalfCauchy(loc=0., scale=5),  # sigma_b_time
        lambda sigma_b_time,mu_b_time: tfd.MultivariateNormalDiag(  # b_time
        loc=affine(tf.ones([num_idx]), mu_b_time[..., tf.newaxis]),
        scale_diag=sigma_b_time*tf.ones(num_idx)),
        tfd.Normal(loc=0., scale=1e5), # a_train
        tfd.Normal(loc=0., scale=1e5), # a_car
        tfd.Normal(loc=0., scale=1e5), # b_cost
        lambda b_cost,a_car,a_train,b_time: tfd.Independent(tfd.Multinomial(
          total_count=1,
          logits=tf.stack([
              affine(DATA[:,0], tf.gather(b_time, IDX[:,0], axis=-1), (a_train + b_cost * DATA[:,1] + adj_AV_train)),
              affine(DATA[:,2], tf.gather(b_time, IDX[:,0], axis=-1), (b_cost * DATA[:,3] + adj_AV_SM)),
              affine(DATA[:,4], tf.gather(b_time, IDX[:,0], axis=-1), (a_car + b_cost * DATA[:,5] + adj_AV_car))
          ], axis=1)
        ),reinterpreted_batch_ndims=1)
    ])

@tf.function
def mmnl_log_prob(mu_b_time, sigma_b_time, b_time, a_train, a_car, b_cost):
    return mmnl_func().log_prob(
      [mu_b_time, sigma_b_time, b_time, a_train, a_car, b_cost, CHOICE])

# Based on https://www.tensorflow.org/tutorials/customization/performance#python_or_tensor_args
# change constant values to tf.constant()
nuts_samples = tf.constant(1000)
nuts_burnin = tf.constant(500)
num_chains = tf.constant(1)
## Initial step size
init_step_size= tf.constant(0.3)
# Set the chain's start state.
initial_state = [
    tf.zeros([num_chains], dtype=tf.float32, name="init_mu_b_time"),
    tf.zeros([num_chains], dtype=tf.float32, name="init_sigma_b_time"),
    tf.zeros([num_chains, num_idx], dtype=tf.float32, name="init_b_time"),
    tf.zeros([num_chains], dtype=tf.float32, name="init_a_train"),
    tf.zeros([num_chains], dtype=tf.float32, name="init_a_car"),
    tf.zeros([num_chains], dtype=tf.float32, name="init_b_cost")
]

## NUTS (using inner step size averaging step)
##
@tf.function
def nuts_sampler(init):
    nuts_kernel = tfp.mcmc.NoUTurnSampler(
      target_log_prob_fn=mmnl_log_prob, 
      step_size=init_step_size,
      )
    adapt_nuts_kernel = tfp.mcmc.DualAveragingStepSizeAdaptation(
  inner_kernel=nuts_kernel,
  num_adaptation_steps=nuts_burnin,
  step_size_getter_fn=lambda pkr: pkr.step_size,
  log_accept_prob_getter_fn=lambda pkr: pkr.log_accept_ratio,
  step_size_setter_fn=lambda pkr, new_step_size: pkr._replace(step_size=new_step_size)
       )

    samples_nuts_, stats_nuts_ = tfp.mcmc.sample_chain(
  num_results=nuts_samples,
  current_state=initial_state,
  kernel=adapt_nuts_kernel,
  num_burnin_steps=tf.constant(100),
  parallel_iterations=tf.constant(5))
    return samples_nuts_, stats_nuts_

samples_nuts, stats_nuts = nuts_sampler(initial_state)

Answer 1

这很可能是您的初始状态和链数的问题。您可以尝试在采样器调用之外初始化您的内核：

nuts_kernel = tfp.mcmc.NoUTurnSampler(
      target_log_prob_fn=mmnl_log_prob, 
      step_size=init_step_size,
      )
    adapt_nuts_kernel = tfp.mcmc.DualAveragingStepSizeAdaptation(
  inner_kernel=nuts_kernel,
  num_adaptation_steps=nuts_burnin,
  step_size_getter_fn=lambda pkr: pkr.step_size,
  log_accept_prob_getter_fn=lambda pkr: pkr.log_accept_ratio,
  step_size_setter_fn=lambda pkr, new_step_size: pkr._replace(step_size=new_step_size)
       )

然后

nuts_kernel.bootstrap_results(initial_state)

并调查 logL 的形状，并返回提案状态。

另一件事是将你的初始状态输入你的 log-likelihood/posterior 并查看返回的对数似然的维度是否与你认为应该的相匹配（如果你正在做多个链那么也许它应该返回 # chains log likelihoods）。

据我了解，批量维度（# chains）必须是所有向量化计算中的第一个维度。

最后一部分 of my blog post on tensorflow and custom likelihoods 包含执行此操作的示例的工作代码。

Answer 2

我能够从我的模型中获得合理的结果。感谢大家的帮助！以下几点帮助解决了各种问题。

使用 JointDistributionSequentialAutoBatched() 生成一致的批处理形状。您需要安装 tf-nightly 才能访问。
超参数的更多先验信息。 Multinomial() 分布中的指数变换意味着无信息的超参数（即 sigma = 1e5）意味着您很快就会有大量正数进入 exp()，导致无穷大。
设置步长等也很重要。
我发现 answer Christopher Suter 对 Tensorflow 概率论坛上最近的一个问题提出了一个问题，该问题指定了来自 STAN 的模型很有用。我使用先验样本作为有用的初始似然参数的起点。
尽管 JointDistributionSequentialAutoBatched() 修正了批量形状，我回去修正了我的联合分布形状，以便打印 log_prob_parts() 给出一致的形状（即 [10,1] for 10 链）。我在不使用 JointDistributionSequentialAutoBatched() 的情况下仍然出现形状错误，但组合似乎有效。
我将 affine() 分成了两个函数。他们做同样的事情，但删除了回溯警告。基本上， affine() 能够广播输入，但它们不同，并且编写两个函数来设置具有一致形状的输入更容易。不同形状的输入导致 Tensorflow 多次跟踪函数。

Tensorflow Probability 中多项式模型的规范

Specification of Multinomial model in Tensorflow Probability

python

tensorflow

tensorflow-probability

discrete