具有 Tensorflow 概率的贝叶斯线性回归

Bayesian Linear Regression with Tensorflow Probability

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

!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

有什么想法吗?

首先,这个问题提得很好——我很容易就能 copy/paste 你的代码并重现 :)

第一个问题(与 HMC 一样)是 step_size/num_leapfrog_steps 设置。通过减小步长 (~.007) 和更少的越级步数 (~3),我能够获得不错的接受率 (~.7)。一般来说,你希望越级步数越小越好,因为每一步都会产生昂贵的梯度计算。

下一个问题是,看起来链条与这些参数的混合仍然不太好。我怀疑这可能是居中与非居中(参见 https://discourse.mc-stan.org/t/centered-vs-noncentered/1502)参数化的问题;你可能更幸运地改变模型的参数化,虽然我不是很熟悉这个的细节......

为了自动找到更好的步长参数,您可以尝试 2 个步长自适应内核之一——SimpleStepSizeAdaptation 或 DualAveragingStepSizeAdaptation——它们可以环绕 HMC 内核实例。此外,您可以尝试使用 NoUTurnSampler 作为 vanilla HMC 的替代品,以避免必须调整 leapfrog steps 参数。对于一个简单的贝叶斯线性回归示例来说,所有这些都是相当重量级的,但我想这个示例突出了让 MCMC 正确的难度:)

希望对您有所帮助!请随时通过 tfprobability@tensorflow.org 向我们发送消息以进行更多讨论——也许其他对参数化问题有更多了解的人会插话。再次感谢您提出的好问题。

这么简单的问题居然这么难,真让人吃惊!原始 HMC 在设置推理时可能对相对较小的细节极为敏感。像 Stan 这样的系统试图通过为您进行大量调整来解决这个问题,但 TFP 的自动调整现在要简单得多。

我发现了一些似乎可以使推理在此处运行良好的更改。简而言之,它们是:

  • 在不受约束的情况下重新参数化 space。
  • 使用非均匀步长(等价于:添加对角线预处理器)
  • 增加 HMC 越级步数。
  • 使用 DualAveragingStepSizeAdaptation 代替 SimpleStepSizeAdaptation

第一个技巧是使用 TransformedTransitionKernel 重新参数化,使比例参数处于不受约束的 space 中。例如,我们可以使用 Exp 双射器来定义对数尺度参数:

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

这确保推理只考虑比例的正值,因此它不必拒绝每一个使比例低于零的动作。当我这样做时,接受率会上升,尽管混合仍然不是很好。

第二个变化是对三个变量使用非均匀步长(这相当于对角线预处理)。看起来这个模型中的后验条件是病态的:二十个数据点确定斜率比确定截距或比例更精确。 "Step size adaptation" 在 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 的一种方法是,在 leapfrog 积分器中,它类似于具有动量的优化器,但每次停止到 accept/reject 时它都会失去动量(通过重新采样)。因此,在我们每一步都这样做的极端情况下(num_leapfrog_steps=1,即 Langevin 动力学),没有任何动量,增加 leapfrog 步骤的数量往往会提高导航棘手几何的能力,类似于动量改进优化器。我没有严格调整任何东西,但是设置 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)