Adding Stochastic Gradient Hamiltonian Monte Carlo (SGHMC) to TFP

[gergely-flamich]

Hello there!

I noticed that TFP still lacks an implementation of SGHMC [1]. I would like to use it in some of my projects, and since I am already relying heavily on the TFP ecosystem, it would be very convenient if it was there. There is definitely interest in this, see Issue #361.

I am wondering whether the reason that it is not implemented yet is simply due to commitments towards more burning issues on the internal development side, or there are deeper reasons involving e.g. compatibility or planned refactoring? In case it is the former reason, I would be very happy to have a stab at implementing it, most likely following up on the suggestions made in #361. Concretely, it appears that it would be more appropriate to implement it as a tf.optimizers.Optimizer as it is done with tfp.optimizer.StochasticGradientLangevinDynamics.

Cheers,

Greg

[1] Chen, T., Fox, E., and Guestrin, C. Stochastic Gradient Hamiltonian Monte Carlo. ICML 2014.

[gergely-flamich]

Based on the implementation of tf.optimizers.SGD and tf.optimizer.StochasticGradientLangevinDynamics, I created a bare-bones implementation of SGHMC without the estimation of the gradient noise covariance matrix (beta_hat = 0 to use the notation from [1]). You can inspect it here. Could you comment on whether you find this a reasonable approach?

[1] Chen, T., Fox, E., and Guestrin, C. Stochastic Gradient Hamiltonian Monte Carlo. ICML 2014.

[brianwa84]

The tfp.optimizer package is a bit neglected, at least the tf.optimizers-like ones.(We have focused more on functional forms like L/BFGS.)I think we'd probably be open to putting an SGHMC impl in there, WDYT @csuter ?

[mhavasi]

+1 I'd be interested too.