I'm trying to understand the Bayes by Backprop algorithm from the paper Weight Uncertainty in Neural Networks, the idea is to make a NN in which each weight has it's own probability distribution. I get the theory, but I don't undertsand how to update the mean and variance in the learning part. I found a code in Pytorch which simply does:

class BayesianLinear(nn.Module):
def __init__(self, in_features, out_features):
    # Weight parameters
    self.weight_mu = nn.Parameter(torch.Tensor(out_features, in_features).uniform_(-0.2, 0.2))
    self.weight_rho = nn.Parameter(torch.Tensor(out_features, in_features).uniform_(-5,-4))
    self.weight = Gaussian(self.weight_mu, self.weight_rho)

How does the optimizer knows how to update mu (the mean) and rho (the variance)? It is the Parameter function which allows it? Or it is happening somewhere in the code? (it looks like a normal neural net from there!).

  • $\begingroup$ In similar cases, it is always highly recommended to include also a link to the relevant paper (edited). $\endgroup$
    – desertnaut
    Apr 15 at 13:11

The PyTorch remembers a computation graph. It means that for each value you get after the forward pass it remembers which operation you have applied and for all basic operations, it knows how to calculate the gradients. So, that's why in PyTorch you get the optimization for free -- at least, for common neural nets. The nn.Parameter method specifically tells to track gradients and change them during optimization.

The approach works well in the particular case for the most part, but there are two not-so-common steps in bayes by backprop:

  1. For each neuron we sample weights. Technically, we start with sampling from $N(0, 1)$ and then we apply the trainable params. The specific values we get from $N(0, 1)$ are kind of extra inputs and for some operations, they change the gradients.
  2. Loss usually depends only on outputs and true labels, but for the KL term, you need to calculate the closed-form value based on weights themselves. You would probably need to explicitly collect the KL loss term from each layer.

Everything else should be handled fine by pytorch

  • $\begingroup$ Thank you for your answer. I'm using a code I found which works well, my confusion came from it looking so simple to implement. I mean, Pytorch realizes not only to update the weights, but to update the things which generate the weights. Is this affirmation correct then? If so, Pytorch is quite amazing! $\endgroup$
    – Cosapocha
    Apr 15 at 17:28
  • $\begingroup$ It updates everything that marked as parameters, yes. $\endgroup$ Apr 16 at 10:03

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.