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!).