Is it possible that Bayesian Neural Network sampled in Turing.jl is not computable in reasonable time?

[Mitnika]

Good afternoon, I have the following code

I would like to use Variational Inference in order to gain samples from the posterior of this Bayesian Neural Network of this simple NN on the MNIST dataset. This code however, does not run in reasonable time, meaning on my machine it cannot even find initial epsilon within 5 hours. Is it possible that variational inference does not scale well for such high dimensional inputs and this computation is actually prohibitive?

[torfjelde]

I don't think this should be beyond what you can do with ADVI, but unfortunately Turing.jl isn't really not targeting Bayesian NNs, and so it's going to be incredibly slow. There are optimizations you can do to your model to make it faster, e.g. the reconstruct is likely type-unstable which will have a signfiicant performance impact, the for-loop over the observations will make things slower, you're using the default ad backend in Turing.jl (ForwardDiff.jl) which scales much worse with dimensionality than, say, ReverseDiff.jl. For the latter two issues, have a look at https://turinglang.org/dev/docs/using-turing/performancetips.

In effect, if you do NUTS(0.99; adtype=AutoReverseDiff(true)), things should be much, much faster. I also wouldn't be using 0.99 as an acceptance rate; that's too high, and is probably also the reason why you're not even able to find an initial epsilon. Use something like 0.65 or 0.8.