After equation 2 in the paper, the author writes:
where $\mathbf{w}^{(i)}$ denotes the $i^{th}$ Monte Carlo sample drawn from the variational posterior $q(\mathbf{w}^{(i)}|\theta)$.
In section 3.2 titled "Gaussian variational posterior":
Suppose that the variational posterior is a diagonal Gaussian distribution, then a sample of the weights $\mathbf{w}$ can be obtained by sampling a unit Gaussian, shifting it by a mean $\mu$ and scaling by a standard deviation $\sigma$. We parameterise the standard deviation pointwise as $\sigma = \log(1 + \exp(\rho))$ and so $\sigma$ is always non-negative. The variational posterior parameters are $\theta = (\mu,\rho)$. Thus the transform from a sample of parameter-free noise and the variational posterior parameters that yields a posterior sample of the weights $\mathbf{w}$ is: $\mathbf{w} = t(\theta,\epsilon) = \mu + \log(1 + \exp(\rho)) \circ \epsilon$ where $\circ$ is point-wise multiplication.
However, in your question you wrote \begin{align} \text{prior} &= \log(q(\mathbf{w}|\mu,\rho)) = \sum_i \log(p(w_i | 0, 1)) \\ \text{posterior} &= \log(P(\mathbf{w})) = \sum_i \log(p(w_i | \mu, \sigma^2)) \\ \text{likelihood} &= \log(P(\mathcal{D}|\mathbf{w})) = y \cdot \log(\text{softmax}(\hat{y})) \end{align} which does not seem to be correct, since $q(\mathbf{w}|\mu,\rho)$ is the posterior according to the author. Here is what I think the authors meant \begin{align} \text{posterior} &= q(\mathbf{w}|\theta) \\ \text{prior} &= P(\mathbf{w}) \\ \text{likelihood} &= P(\mathcal{D}|\mathbf{w}) \end{align} Also, the authors did not specify the prior $P(\mathbf{w}$ until equation 7 in section 3.3, which is a mixture of Gaussians and not a standard Gaussian as you wrote. Just wanted to point this out.
Here is a rough outline of what you can do in torch to implement this:
- Create the dataset $\mathcal{D}$ consisting of the vectors $\mathbf{x}_1,\mathbf{x}_2,\dots,\mathbf{x}_N$ and the labels $y_1,y_2,...,y_N$. This can be done by sampling $\mathbf{x}$ vectors from some multivariate distribution, and sampling $y$ labels from a Bernoulli distribution.
- Sample $\epsilon$ from the multivariate normal distribution $N(0,I)$.
- Sample initial values for $\mu$ and $\rho$ from the multivariate normal distribution $N(0,I)$. You only need to do this once to start performing gradient descent. Note that $\mu$ is a column vector and $rho$ is a diagonal matrix. You can construct $rho$ by first sampling a vector of values from the multivariate normal distribution, and then putting these values in a diagonal matrix.
- Compute $\mathbf{w} = \mu + \log(1 + \exp(\rho)) \circ \epsilon$.
- Compute $\log(q(\mathbf{w}|\mu,\rho))$ by inputting the $\mathbf{w}$,$\mu$, and $\rho$ that you obtained in steps 3 and 4 into the multivariate normal probability density function.
- Compute $\log(P(\mathbf{w}))$ by inputting the $\mathbf{w}$ that you obtained in step 3 into the multivariate normal probability density function with mean $\mathbf{0}$ and covariance $I$. Note again that this does not follow what the authors did in the paper, and that you would instead need to implement equation 7 in the paper.
- Compute $\log(P(\mathcal{D}|\mathbf{w}))$. Note that you did not specify what $\hat{y}$ is. For the sake of simplicity, I will assume that $\hat{y} = \mathbf{w}^T \mathbf{x}$. This means that $\log(P(\mathcal{D}|\mathbf{w})) = -\sum_{i=1}^N y_i \cdot \log(\hat{y}_i)$. This is just the cross-entropy between the $y$ labels and the $\hat{y}$ labels. Note that each $y$ and $\hat{y}$ must range between 0 and 1.
- Compute $f(\mathbf{w},\theta) = \log(q(\mathbf{w}|\theta)) - \log(P(\mathbf{w})) - \log(P(\mathcal{D}|\mathbf{w}))$.
- Compute the gradients as given in steps 5 and 6 of the algorithm.
- Update $\mu$ and $\rho$ as given in step 7 of the algorithm.
- Repeat the following aforementioned steps until convergence: step 2 (sampling $\epsilon$) $\rightarrow$ step 4 (computing $\mathbf{w}$) $\rightarrow$ steps 5,6,7,8,9, and 10.
I am aware that I did not mention how to compute the gradients in torch. If that is still something that you are not sure about, let me know.