Custom Tensorflow optimizer cSGLD (Stochastic Langevin Dynamics) in TF2: correct update ops?

[Jordy-VL]

System information

Have I written custom code (as opposed to using a stock example script provided in TensorFlow): yes
OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Linux Ubuntu 16.04
TensorFlow installed from (source or binary): binary
TensorFlow version (use command below): v2.2
Python version: 3.6.9
CUDA/cuDNN version: v10.2
GPU model and memory: GeForce GTX 1070 - 8117MiB

Describe the current behavior:

I implemented Cyclical Stochastic Gradient MCMC for Bayesian Deep Learning as a custom optimizer in TF2. While it does run, I am not exactly sure about my implementation since the documentation for Tensorflow optimizers is almost non-existent... I would appreciate any help!
For your reference the pytorch implementation (WAY simpler) is here: https://github.com/ruqizhang/csgmcmc/blob/master/experiments/cifar_csgmcmc.py

Describe the bug

It's not so much a bug, rather asking advice on how to implement the noisy update in TF2 for stochastic gradient langevin dynamics without preconditioning. In Addons there is the TF1 only example of preconditioned SGLD, whereas I want to make it work for TF2.

Code to reproduce the issue

See the implementation with a running example here:
https://colab.research.google.com/drive/16kwhGfiat-SkK0RWvf4EADZC_FDVudL8?usp=sharing

If you use momentum (i.e. alpha < 1), then the implementation does not work.
I think the community would benefit from having access to this example (once working). Since there is no Tensorflow implementation yet of Cyclical Stochastic Gradient MCMC for Bayesian Deep Learning.

[tilakrayal]

I was able to reproduce the code in TF v2.4,2.5.0rc1,nightly.Please find the gist here

[sachinprasadhs]

@Jordy-VL does this document help you?

[Jordy-VL]

@Jordy-VL does this document help you?

Hi,

Thanks for the suggestion. It clears up some points, yet it still only treats preconditioned (c)SGLD, which only works in TF1.
I think my current implementation works fine :)

There is still an issue with SGHMC, where I cannot do the momentum sparse update due to an embedding lookup operation, any guidance there would be appreciated :)

Kind regards,

Jordy

[ymodak]

Have you tried using https://www.tensorflow.org/probability/api_docs/python/tfp/optimizer/StochasticGradientLangevinDynamics
from tensorflow-probabilty module for your case?

[Jordy-VL]

@ymodak That's the one I'm referring to in the intro.
It does not enable TF2 for the lines mentioned here: https://github.com/tensorflow/probability/blob/4683836be3701e75f459494bd3de12edb7560a2d/tensorflow_probability/python/optimizer/sgld.py#L277

I did find https://github.com/google-research/google-research/blob/master/cold_posterior_bnn/core/sgmcmc.py, but I still need to figure out if this implementation supports TF2 out-of-the-box.