Mass matrix adaptation for chees-HMC

[jonny-so]

Hi,

I see in the implementation of chees-HMC adaptation the mass matrix is fixed to be the identity. I guess this is because the algorithm presented in the paper assumes the problem has been transformed such that the identity mass matrix is appropriate, but this was really just assumed for ease of exposition (see footnote 1 in the chees-HMC paper). In practice it will be necessary adapt the mass matrix most of the time, alongside the step-size and trajectory length adaptation.

As a first attempt it would seem reasonable to use windowed adaptation to learn a mass matrix and (initial) step size, followed by a chees-HMC adaptation to learn the trajectory length function (and possibly further refine the step size). But as far as I can tell this is not possible at the moment -- save for manually applying a linear change of variables to the problem -- as we can't pass a mass matrix in to the chees-HMC adaptation. This is also suboptimal, as we would like to make use of a recent estimate of the optimal trajectory length while learning the mass matrix, and vice-versa - it seems as though in an ideal world we would incorporate trajectory-length adaptation into the windowed adaptation somehow.

Does anyone have any thoughts on this, or examples of using mass-matrix adaptation with chees-HMC more generally?

Thanks!

Jonny

[junpenglao]

+1 that we should extend the windowed adaptation to have more general mass matrix adaptation.