View source on GitHub |
The ChEES criterion from [1].
tfp.experimental.mcmc.chees_criterion(
previous_state,
proposed_state,
accept_prob,
trajectory_length,
validate_args=False,
experimental_shard_axis_names=None,
experimental_reduce_chain_axis_names=None
)
ChEES stands for Change in the Estimator of the Expected Square.
ChEES = 1/4 E[(||x' - E[x]||**2 - ||x - E[x]||**2)**2],
where x
is the previous chain state, x'
is the next chain state, and
||.||
is the L2 norm. Both expectations are with respect to the chain's
stationary distribution. In practice, the inner expectation is replaced by the
empirical mean across chains, so computing this criterion requires that at
least 2 chains are present. The outer expectation is computed by the caller
(e.g. in the GradientBasedTrajectoryLengthAdaptation
kernel).
This can be thought of as the standard expected squared jump distance (ESJD) criterion, except that the jump distance is computed in the space of centered squared L2 norms.
Unlike ChEES, regular ESJD is maximized by perfectly anticorrelated proposals, which can give excellent mean estimates but terrible variance estimates; maximizing ChEES should give good estimates across a wider range of types of posterior expectations.
Returns | |
---|---|
chees
|
The value of the ChEES criterion. |
Raises | |
---|---|
ValueError
|
If accept_prob indicates that there are fewer than 2 chains.
|
References
[1]: Hoffman, M., Radul, A., & Sountsov, P. (2020). An Adaptive MCMC Scheme for Setting Trajectory Lengths in Hamiltonian Monte Carlo. <https://proceedings.mlr.press/v130/hoffman21a>