python - 具有 tensorflow 概率的贝叶斯线性回归

Question

我无法让贝叶斯线性回归与 Tensorflow Probability 一起使用。这是我的代码:

!pip install tensorflow==2.0.0-rc1
!pip install tensorflow-probability==0.8.0rc0

import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp
tfd = tfp.distributions
N = 20
std = 1
m = np.random.normal(0, scale=5, size=1).astype(np.float32)
b = np.random.normal(0, scale=5, size=1).astype(np.float32)
x = np.linspace(0, 100, N).astype(np.float32)
y = m*x+b+ np.random.normal(loc=0, scale=std, size=N).astype(np.float32)

num_results = 10000
num_burnin_steps = 5000

def joint_log_prob(x, y, m, b, std):
  rv_m = tfd.Normal(loc=0, scale=5)
  rv_b = tfd.Normal(loc=0, scale=5)
  rv_std = tfd.HalfCauchy(loc=0., scale=2.)

  y_mu = m*x+b
  rv_y = tfd.Normal(loc=y_mu, scale=std)

  return (rv_m.log_prob(m) + rv_b.log_prob(b) + rv_std.log_prob(std)
          + tf.reduce_sum(rv_y.log_prob(y)))

# Define a closure over our joint_log_prob.
def target_log_prob_fn(m, b, std):
    return joint_log_prob(x, y, m, b, std)

@tf.function(autograph=False)
def do_sampling():
  kernel=tfp.mcmc.HamiltonianMonteCarlo(
          target_log_prob_fn=target_log_prob_fn,
          step_size=0.05,
          num_leapfrog_steps=3)
  kernel = tfp.mcmc.SimpleStepSizeAdaptation(
        inner_kernel=kernel, num_adaptation_steps=int(num_burnin_steps * 0.8))
  return tfp.mcmc.sample_chain(
      num_results=num_results,
      num_burnin_steps=num_burnin_steps,
      current_state=[
          0.01 * tf.ones([], name='init_m', dtype=tf.float32),
          0.01 * tf.ones([], name='init_b', dtype=tf.float32),
          1 * tf.ones([], name='init_std', dtype=tf.float32)
      ],
      kernel=kernel,
      trace_fn=lambda _, pkr: [pkr.inner_results.accepted_results.step_size,
                               pkr.inner_results.log_accept_ratio])

samples, [step_size, log_accept_ratio] = do_sampling()
m_posterior, b_posterior, std_posterior = samples

p_accept = tf.reduce_mean(tf.exp(tf.minimum(log_accept_ratio, 0.)))
print('Acceptance rate: {}'.format(p_accept))

n_v = len(samples)
true_values = [m, b, std]

plt.figure()
plt.title('Training data')
plt.plot(x, y)

plt.figure()
plt.title('Visualizing trace and posterior distributions')
for i, (sample, true_value) in enumerate(zip(samples, true_values)):
  plt.subplot(2*n_v, 2, 2*i+1)
  plt.plot(sample)
  plt.subplot(2*n_v, 2, 2*i+2)
  plt.hist(sample)
  plt.axvline(x=true_value)
>>> Acceptance rate: 0.006775229703634977

有什么想法吗？

Answer 1

令人惊讶的是，这样一个简单的问题竟然如此困难!原始 HMC 在设置推理时可能对相对较小的细节极其敏感。像 Stan 这样的系统尝试通过为您进行大量调整来解决这个问题，但 TFP 的自动调整目前更加基础。

我发现一些更改似乎使推理在这里运行良好。简而言之，它们是:

• 在不受约束的空间中重新参数化。
• 使用非均匀步长(相当于:添加对角预处理器)
• 增加 HMC 蛙跳步数。
• 使用 DualAveragingStepSizeAdaptation 而不是 SimpleStepSizeAdaptation。

第一个技巧是使用 TransformedTransitionKernel 重新参数化，以便比例参数位于不受约束的空间中。例如，我们可以使用 Exp 双射器来定义对数刻度参数:

tfb = tfp.bijectors
kernel = tfp.mcmc.TransformedTransitionKernel(
      inner_kernel=kernel,
      bijector=[tfb.Identity(), tfb.Identity(), tfb.Exp()]
  )

这确保推理仅考虑比例的正值，因此不必拒绝比例低于零的每一个移动。当我这样做时，接受率会大幅上升，尽管混合效果仍然不是很好。

第二个变化是对三个变量使用非均匀步长(这相当于对角线预处理)。看起来这个模型中的后验条件是病态的:二十个数据点确定的斜率比确定截距或尺度更精确。 TFP 中的“步长自适应”只是找到一个步长，以便接受指定百分比的样本，这通常由后验最严格约束的组件控制:如果其他组件具有更宽的后验，小步长将阻止防止它们混合。估计合理步长的一种方法是使用变分推理的标准差和因子正态代理后验:

surrogate_posterior = tfp.experimental.vi.build_factored_surrogate_posterior(
    event_shape=[[], [], []],
    constraining_bijectors=[None, None, tfb.Exp()])
losses = tfp.vi.fit_surrogate_posterior(
    target_log_prob_fn, surrogate_posterior,
    optimizer=tf.optimizers.Adam(
        learning_rate=0.1,
        # Decay second-moment estimates to aid optimizing scale parameters.
        beta_2=0.9),
    num_steps=1000)
approximate_posterior_stddevs = [np.std(x) for x in surrogate_posterior.sample(50)]

另一个通用技巧是增加蛙跳步数。考虑 HMC 的一种方法是，在蛙跳积分器中，它类似于具有动量的优化器，但每次停止接受/拒绝时，它都会失去动量(通过重新采样)。因此，在极端情况下，我们每一步都这样做(num_leapfrog_steps=1，即朗之万动力学)，没有任何动量，并且增加蛙跳步数往往会提高导航棘手几何图形的能力，类似于动量如何改进优化器。我没有严格调整任何内容，但设置 num_leapfrog_steps=16 而不是 3 似乎在这里有很大帮助。

这是我对您的代码的修改版本，其中包含了这些技巧。它似乎在大多数执行中都混合得很好(尽管我确信它并不完美):

!pip install tensorflow==2.0.0-rc1
!pip install tensorflow-probability==0.8.0rc0

import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp

tfd = tfp.distributions
tfb = tfp.bijectors

N = 20
std = 1
m = np.random.normal(0, scale=5, size=1).astype(np.float32)
b = np.random.normal(0, scale=5, size=1).astype(np.float32)
x = np.linspace(0, 100, N).astype(np.float32)
y = m*x+b+ np.random.normal(loc=0, scale=std, size=N).astype(np.float32)

num_results = 1000
num_burnin_steps = 500

def joint_log_prob(x, y, m, b, std):
  rv_m = tfd.Normal(loc=0, scale=5)
  rv_b = tfd.Normal(loc=0, scale=5)
  rv_std = tfd.HalfCauchy(0., scale=1.)

  y_mu = m*x+b
  rv_y = tfd.Normal(loc=y_mu, scale=rv_std[..., None])

  return (rv_m.log_prob(m) + rv_b.log_prob(b)
          + rv_std.log_prob(std)
          + tf.reduce_sum(rv_y.log_prob(y)))

# Define a closure over our joint_log_prob.
def target_log_prob_fn(m, b, std):
    return joint_log_prob(x, y, m, b, std)

# Run variational inference to initialize per-variable step sizes.
surrogate_posterior = tfp.experimental.vi.build_factored_surrogate_posterior(
    event_shape=[[], [], []],
    constraining_bijectors=[None, None, tfb.Exp()])
losses = tfp.vi.fit_surrogate_posterior(
    target_log_prob_fn,
    surrogate_posterior,
    optimizer=tf.optimizers.Adam(
        learning_rate=0.1,
        # Decay second-moment estimates to aid optimizing scale parameters.
        beta_2=0.9),
    num_steps=1000)
approximate_posterior_stddevs = [np.std(z) for z in surrogate_posterior.sample(50)]

@tf.function(autograph=False)
def do_sampling():
  kernel=tfp.mcmc.HamiltonianMonteCarlo(
          target_log_prob_fn=target_log_prob_fn,
          step_size=approximate_posterior_stddevs,
          num_leapfrog_steps=16)
  kernel = tfp.mcmc.TransformedTransitionKernel(
      inner_kernel=kernel,
      bijector=[tfb.Identity(),
                tfb.Identity(),
                tfb.Exp()]
  )
  kernel = tfp.mcmc.DualAveragingStepSizeAdaptation(
        inner_kernel=kernel,
        num_adaptation_steps=int(num_burnin_steps * 0.8))
  return tfp.mcmc.sample_chain(
      num_results=num_results,
      num_burnin_steps=num_burnin_steps,
      current_state=[
          0.01 * tf.ones([], name='init_m', dtype=tf.float32),
          0.01 * tf.ones([], name='init_b', dtype=tf.float32),
          1. * tf.ones([], name='init_std', dtype=tf.float32)
      ],
      kernel=kernel,
      trace_fn=lambda _, pkr: [pkr.inner_results.inner_results.accepted_results.step_size,
                               pkr.inner_results.inner_results.log_accept_ratio])

samples, [step_size, log_accept_ratio] = do_sampling()
m_posterior, b_posterior, std_posterior = samples

p_accept = tf.reduce_mean(tf.exp(tf.minimum(log_accept_ratio, 0.)))
print('Acceptance rate: {}'.format(p_accept))

n_v = len(samples)
true_values = [m, b, std]

plt.figure(figsize=(12, 12))
plt.title('Visualizing trace and posterior distributions')
for i, (sample, true_value) in enumerate(zip(samples, true_values)):
  plt.subplot(2*n_v, 2, 2*i+1)
  plt.plot(sample)
  plt.subplot(2*n_v, 2, 2*i+2)
  plt.hist(sample.numpy())
  plt.axvline(x=true_value)