tfp.experimental.mcmc.remc_thermodynamic_integrals

Estimate thermodynamic integrals using results of ReplicaExchangeMC.

tfp.experimental.mcmc.remc_thermodynamic_integrals(
    inverse_temperatures, potential_energy, iid_chain_ndims=0
)

Write the density, when tempering with inverse temperature b, as p_b(x) = exp(-b * U(x)) f(x) / Z_b. Here Z_b is a normalizing constant, and U(x) is the potential energy. f(x) is the untempered part, if any.

Let E_b[U(X)] be the expected potential energy when X ~ p_b. Then, -1 * integral_c^d E_b[U(X)] db = log[Z_d / Z_c], the log normalizing constant ratio.

Let Var_b[U(X)] be the variance of potential energy whenX ~ p_b(x). Then,integral_c^d Var_b[U(X)] db = E_d[U(X)] - E_c[U(X)]`, the cross entropy difference.

Integration is done via the trapezoidal rule. Assume E_b[U(X)] and Var_b[U(X)] have bounded second derivatives, uniform in b. Then, the bias due to approximation of the integral by a summation is O(1 / K^2).

Suppose U(X), X ~ p_b has bounded fourth moment, uniform in b. Suppose further that the swap acceptance rate between every adjacent pair is greater than C_s > 0. If we have N effective samples from each of the n_replica replicas, then the standard error of the summation is O(1 / Sqrt(n_replica * N)).

inverse_temperatures Tensor of shape [n_replica, ...], used to temper n_replica replicas. Assumed to be decreasing with respect to the replica index.
potential_energy The potential_energy field of ReplicaExchangeMCKernelResults, shape [n_samples, n_replica, ...]. If the kth replica has density p_k(x) = exp(-beta_k * U(x)) * f_k(x), then potential_energy[k] is U(X), where X ~ p_k.
iid_chain_ndims Number of dimensions in potential_energy, to the right of the replica dimension, that index independent identically distributed chains. In particular, the temperature for these chains should be identical. The sample means will be computed over these dimensions.

ReplicaExchangeMCThermodynamicIntegrals namedtuple.