IAF guide with 2D parameters

I want to use an inverse autoregressive flow (IAF) for the guide distribution in a Bayesian neural network. I’m not sure how to specify the IAF distribution for the weight matrices because these are 2D. Previously I was using a Normal distribution, and I can just create a location and scale with the same shape as the parameter (e.g. 512 x 100). However, InverseAutoregressiveFlow only accepts an int for the input dimension. I assume I’ll need to have input_dim=512 * 100, but then the samples won’t be the right shape, so I expect there will be an error in SVI (or will it reshape for me?). Should I define a new Transform that does the reshaping? Or is there a better way to do this?

i’m not sure off the top of my head what the best approach would be. however, i don’t think you actually want to implement what you’re describing. for one thing the IAF would contain weight matrices of size at least (512 * 100) x (512 * 100), which corresponds to something like 10gb of memory.

1 Like

I was just drawing my model up on a whiteboard and noticing the same thing. I guess I hadn’t thought it through enough initially to realize how many additional parameters I would be introducing this way. I’m going to try using the approach from the Multiplicative Normalizing Flows paper instead so I can still have NFs but of a much smaller dimensionality. Thanks!