Radial BNN variational approximation


I’m trying to fit a radial BNN posterior variational approximation as per this paper.

However, since I’ll be training a BNN, I don’t want to have to write a custom guide and define this variational approximation for all of my layers, and so was trying to implement a custom AutoGuide which automatically puts a radial BNN approximation on all of my weights.

The radial approximation is defined as follows:
where I just need to sample all epsilon_MFVI from an independent standard normal distribution, normalize them, and multiply them by r, which is a scalar sampled from a standard normal.

How could I go about implementing this custom AutoGuide?
Is there a smarter way of implementing this variational approximation?

Thanks in advance.

might be easier to fork some of the code in TyXe rather than forking an AutoGuide in pyro

But in TyXe don’t you still mainly use the built-in autoguides from Pyro?

yes but afaik tyxe takes care to make use of the so-called “local reparameterization trick” which pyro autoguides do not implement (as it’s special to linear layers). this can be very important in practice