Poor performance of BlackJax SgLD for Bayesian CNN on MNIST dataset.

[HitarthGandhi]

We are trying to use BlackJAX SgLD for BCNN on the MNIST dataset to solve blackjax-devs/sampling-book#14. We used the SgMCMCJax library for SgLD, which gives an accuracy of 94.4%. BlackJax is unable to reproduce the results in the same setting. The detailed code and our findings can be found in this notebook.

You can see in the below plot that BlackJax SGLD does not give the desired results.

[albcab]

Try installing the latest developments with pip install blackjax-nightly. Notice that you'll need to change your code a bit for it to work with the latest SgLD kernel. Specifically, you won't be initializing a state and instead passing the previous position directly to the step function, you'll also be passing the learning_rate/step_size to the step function instead of the kernel initializer:

sgld = blackjax.sgld(grad_fn)
new_position = sgld.step(rng_key, position, minibatch, step_size)

Hope this helps, let us know if it works. We're working on a new release soon.