tfp.mcmc.HamiltonianMonteCarlo Not working in Tensorflow Probability - tensorflow

Question

I have the following code, which basically tries to fit a simple regression model using tensorflow probability. The code runs without error, but the MCMC sampler doesn't seem to be doing anything in that it returns a trace of the initial states.
import tensorflow.compat.v2 as tf
import tensorflow_probability as tfp
from tensorflow_probability import distributions as tfdimport warnings
tf.enable_v2_behavior()
plt.style.use("ggplot")
warnings.filterwarnings('ignore')
ru=4
N=102
N = 102 # number of data points
t = np.linspace(0, 4*np.pi, N)
data = 3+np.sin(t+0.001) + 0.5 + np.random.randn(N)
media_1 = ((data-min(data))/(max(data)-min(data)) ) #+ np.random.normal(0,.05, N)
y = np.repeat(ru, N) + np.random.normal(.3,.01,N) * media_1 + np.random.normal(0, .005, N)
# model
model = tfd.JointDistributionNamed(dict(
beta1 = tfd.Normal(0,1) ,
intercept = tfd.Normal(0,5 ) ,
var = tfd.Normal(0.05, 0.0005) ,
y = lambda intercept,beta1,var:
tfd.Independent(tfd.Normal(loc=intercept + beta1 * media_1, scale=var),
reinterpreted_batch_ndims=1
)
))
def target_log_prob_fn(intercept, beta1, var):
return model.log_prob({'intercept':intercept, 'beta1':beta1, 'var':var, 'y': y })
s = model.sample()
init_states = [ tf.fill([1], s['intercept'].numpy(), name='init_intercept'),
tf.fill([1], s['beta1'].numpy(), name='init_beta1'),
tf.fill([1], s['var'].numpy(), name='init_var'),]
num_results = 5000
num_burnin_steps = 3000
# Improve performance by tracing the sampler using `tf.function`
# and compiling it using XLA.
#tf.function(autograph=False, experimental_compile=True)
def do_sampling():
return tfp.mcmc.sample_chain(
num_results=num_results,
num_burnin_steps=num_burnin_steps,
current_state=init_states,
kernel=tfp.mcmc.HamiltonianMonteCarlo(
target_log_prob_fn=target_log_prob_fn,
step_size=0.1,
num_leapfrog_steps=3)
)
states, kernel_results = do_sampling()
The trace that is returned in states is exactly the same as the initial values in initial_states... Any ideas?

Answer 1

I can confirm that this MCMC sampler is not mixing by printing the acceptance rate with a snippet I found here
print("Acceptance rate:", kernel_results.is_accepted.numpy().mean())
That same page provides some hints about how to make your HMC kernel adaptive, which means it will automatically reduce the step size if too many proposed steps are rejected (and increase it if too many are accepted, too):
# Apply a simple step size adaptation during burnin
#tf.function
def do_sampling():
adaptive_kernel = tfp.mcmc.SimpleStepSizeAdaptation(
tfp.mcmc.HamiltonianMonteCarlo(
target_log_prob_fn=target_log_prob_fn,
step_size=0.1,
num_leapfrog_steps=3),
num_adaptation_steps=int(.8 * num_burnin_steps),
target_accept_prob=np.float64(.65))
return tfp.mcmc.sample_chain(
num_results=num_results,
num_burnin_steps=num_burnin_steps,
current_state=init_states,
kernel=adaptive_kernel,
trace_fn=lambda cs, kr: kr)
samples, kernel_results = do_sampling()
print("Acceptance rate:", kernel_results.inner_results.is_accepted.numpy().mean())
This produces a non-zero acceptance rate for me.