Is it possible to create a NN with combined probabilistic and nonprobabilistic weights?


#1

Is it possible to create a NN with combined probabilistic and nonprobabilistic weights? For example, I have a neural network which consists of 2 layers. I want to sample weights for the first layer (from Normal distribution) and use certain weights for the second layer. Is it possible in pyro?
P.S.: now it seems impossible because of different optimizers (probabilistic part needs ELBO, but nonprobabilistic part may need MSE (for regression task)


Bayesian RNN (NaN loss issue)
#2

It is possible. You can use pyro.random_module with only priors for probabilistic weights.


#3

now it seems impossible because of different optimizers

I believe you can encode both ELBO and MSE in a single objective. First observe that MSE loss can be encoded as maximum likelihood estimation wrt Gaussian observation likelihood. Then as @fehiepsi points out, the priors over the first layer will result in an ELBO term, and the parameters in the second layer will be optimized wrt maximum likelihood. As long as your final observation is declared with Gaussian errors, you can use a vanilla Trace_ELBO():

pyro.sample("obs", dist.Normal(neural_net_output, 1.0), obs=data)

#4

Thank you, it works. I was confused in the theoretical part but @fritzo explained it.