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)$. The authors do not say anything about whether $\mu$ and $\rho$ are themselves random variables with an associated distribution $p(\theta)$ or are treated as unknown constants. For simplicity, I will assume they are treated as unknown constants. In that case, their initial values can be sampled from any distribution. For simplicity, you can sample them using thetorch_randnfunction. You only need to do this once to start performing gradient descent. Note that $\mu$ is a column vector and $\rho$ isa diagonal matrixalso a column vector because the authors use the element-wise multiplication operation $\circ$ in their paper to multiply $\epsilon$ by $\log(1 + \exp(\rho))$. - 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.
Response to comments
What is still unclear to me is how to design the hidden layers of neurons. Am i correct in understanding that if i have 100 observations, for (2) this means i sample 100 $\epsilon$'s from a mv normal distribution with $\mu=0$ and a 100x100 identity matrix? for (3) i sample 100 $\mu$'s as a column vector and $\rho$ is 100x100 diagonal matrix of 100 samples of the same mv normal distribution? and for (4) i will get 100 samples $\mathbf{w}$? For my hidden layer i have specified 32 neurons. Where do these come into play?
Suppose you have $N$ observations in your dataset $\mathcal{D}$ of $(\mathbf{x},y)$ pairs, as discussed in step 1 above, and suppose that each $\mathbf{x}$ is $K$-dimensional such that
$$
\mathbf{x} =
\begin{bmatrix}
x_1 \\
x_2 \\
\vdots \\
x_K
\end{bmatrix}
$$
Also, for the sake of simplicity, suppose that your network consists of a single layer and the output of your network is
$$
\hat{y} = \phi(\mathbf{W}^T \mathbf{x} + \mathbf{b})
$$
where $\mathbf{W}$ is a $K \times M$ matrix of weights, where $M$ is the number of possible classes, $\mathbf{b}$ is a $M \times 1$ column vector, and $\phi(\cdot)$ is the softmax function. Recall that the authors wrote
$$
\mathbf{w} = \mu + \log(1 + \exp(\rho)) \circ \epsilon
$$
This be generalized to a $K \times M$ matrix $\mathbf{W}$ by sampling a $K \times M$ matrix $\epsilon$ using torch_randn, sampling initial values (only need to sample once) of $K \times M$ matrices $\mu$ and $\rho$ using torch_randn. Since $\epsilon \sim N(0,I)$, then you can also sample a $K \times M$ matrix $\epsilon$ using torch_randn. You can then obtain $\mathbf{W}$ as shown above. Each row of $\mathbf{W}$ represents a single weight vector. You can repeat this process using column vectors for $\mu,\rho,$ and $\epsilon$ to obtain $\mathbf{b}$:
$$
\mathbf{b} = \mu + \log(1 + \exp(\rho)) \circ \epsilon
$$
You can then compute $\hat{y}$ as shown above.
For step 7, $\hat{y}$ is supposed to represent the predicted values, which are compared to the ground truth $y$. So i think we're talking about the same thing here. You mention both $y$ and $\hat{y}$ have to range between 0 and 1, does that mean i should apply a sigmoid activation function after the output layer?
I think we have slightly different definitions of $\hat{y}$. I am defining $\hat{y}$ as $$ \hat{y} = \phi(\mathbf{W}^T \mathbf{x} + \mathbf{b}) $$
I'm not experienced in R, so here is my attempt using Python and PyTorch. I don't get NaN values, but the f loss does not decrease consistently. I commented my code as much as possible for clarity. Hope this helps.
import torch
import math
# define pi
pi = torch.tensor(math.pi)
def log_likelihood(x,mu,rho):
# need to compute negative log-likelihood and NOT the
# log-likelihood. This is because we are performing gradient descent.
# In gradient descent, we are trying to minimize a function.
# Minimizing the negative log-likelihood function is equivalent to
# maximizing the log-likelihood function.
# need to re-parameterize sigma to keep it positive. See the first
# paragraph of section 3.2 in the paper and section 3.1 in the paper
# for details
sigma = torch.log(1 + torch.exp(rho))
return -torch.log(sigma * torch.sqrt(2 * pi)) - (0.5 * (torch.div(x-mu,sigma) ** 2))
# uncomment one or the other to choose your activation function
act_func = torch.nn.Sigmoid()
# act_func = torch.nn.ReLU()
# the logarithm of the softmax function. See
# https://pytorch.org/docs/stable/generated/torch.nn.LogSoftmax.html
# for details
log_softmax = torch.nn.LogSoftmax(dim = 0)
learning_rate = 1e-5
# STEP 1
# create the dataset D consisting of vectors x_1,x_2,...,x_N and labels
# y_1,y_2,...,y_N
# input dimensionality (number of input features)
d_in = 3
# dimensionality of the hidden layer
d_hidden = 32
# output dimensionality (number of predicted features)
d_out = 1
# number of observations in training set
n = 100
# number of training iterations
num_iter = 100
# create input data. This is a (N x d_in) matrix, where N is the number of
# observations and d_in is the dimensionality of the input
X = torch.randn(n,d_in)
# generate the ground truth with a bernoulli distribution with probability
# p = 0.6
p = 0.6
y = torch.bernoulli(p * torch.ones((n,)))
# In STEP 2, I will sample the required mu and rho matrices and column
# vectors. Since the authors in the paper do not specify p(mu,rho)in the
# paper, then mu and rho are both treated here as unknown constants and
# not as random variables. In this case, I can sample initial values for
# mu and rho from any distributuon. For simplicity, I will use the uniform
# distribution, which can be done using the torch.rand (not torch.randn)
# function.
#
# Here are the weight matrices and bias vectors I am trying to construct
# in STEP 2:
#
# weight matrix for the first layer:
# W1 = mu_W1 + log(1 + exp(rho_W1)) * epsilon_W1
# bias vector for the first layer:
# b1 = mu_b1 + log(1 + exp(rho_b1)) * epsilon_b1
# weight matrix for the second layer:
# W2 = mu_W2 + log(1 + exp(rho_W2)) * epsilon_W2
# bias vector for the second layer:
# b2 = mu_b2 + log(1 + exp(rho_b2)) * epsilon_b2
# STEP 2a
# sample mu matrix for the first layer with dimensions d_in x d_hidden,
# where d_hidden is the dimensionality of the hidden layer. This mu
# matrix will be used to construct W1, which is the d_in x d_hidden weight
# matrix for the first layer. Note that since I need to compute the
# gradient of f(W,theta) with respect to mu_W1, then I will need to
# set requires_grad = True for mu_W1.
mu_W1 = torch.randn(d_in,d_hidden,requires_grad = True)
# STEP 2b
# sample rho matrix for the first layer with dimensions d_in x d_hidden,
# where d_hidden is the dimensionality of the hidden layer. This rho
# matrix will be used to construct W1, which is the d_in x d_hidden weight
# matrix for the first layer. Note that since I need to compute the
# gradient of f(W,theta) with respect to rho_W1, then I will need to
# set requires_grad = True for rho_W1.
rho_W1 = torch.randn(d_in,d_hidden,requires_grad = True)
# STEP 2c
# sample mu column vector for the first layer with dimensions d_hidden x 1,
# where d_hidden is the dimensionality of the hidden layer. This mu column vector
# will be used to construct b1, which is the d_hidden x 1 bias
# vector for the first layer. Note that since I need to compute the
# gradient of f(W,theta) with respect to mu_b1, then I will need to
# set requires_grad = True for mu_b1.
mu_b1 = torch.randn(d_hidden,1,requires_grad = True)
# STEP 2d
# sample rho column vector for the first layer with dimensions d_hidden x 1,
# where d_hidden is the dimensionality of the hidden layer. This rho column vector
# will be used to construct b1, which is the d_hidden x 1 bias
# vector for the first layer. Note that since I need to compute the
# gradient of f(W,theta) with respect to rho_b1, then I will need to
# set requires_grad = True for rho_b1.
rho_b1 = torch.randn(d_out,1,requires_grad = True)
# sample the rest of the matrices and column vectors
mu_W2 = torch.randn(d_hidden,d_out,requires_grad = True)
rho_W2 = torch.randn(d_hidden,d_out,requires_grad = True)
mu_b2 = torch.randn(d_out,1,requires_grad = True)
rho_b2 = torch.randn(d_out,1,requires_grad = True)
for i in range(num_iter):
# STEP 3a
# sample epsilon matrix for the first layer with dimensions d_in x d_hidden,
# where d_hidden is the dimensionality of the hidden layer. This epsilon
# matrix will be used to construct W1, which is the d_in x d_hidden weight
# matrix for the first layer. Note that since I do not need to compute the
# gradient of f(W,theta) with respect to epsilon_W1, then I don't need to
# set requires_grad = True for epsilon_W1. This will save memory.
epsilon_W1 = torch.randn(d_in,d_hidden)
# STEP 3b
# sample epsilon column vector for the first layer with dimensions d_hidden x 1,
# where d_hidden is the dimensionality of the hidden layer. This epsilon
# matrix will be used to construct b1, which is the d_hidden x 1 bias vector
# for the first layer. Note that since I do not need to compute the
# gradient of f(W,theta) with respect to epsilon_b1, then I don't need to
# set requires_grad = True for epsilon_b1. This will save memory.
epsilon_b1 = torch.randn(d_hidden,1)
# STEP 3c
# sample epsilon matrix for the second layer with dimensions d_hidden x d_out,
# where d_out is the dimensionality of the output layer. This epsilon
# matrix will be used to construct W2, which is the d_hidden x d_out weight
# matrix for the second layer. Note that since I do not need to compute the
# gradient of f(W,theta) with respect to epsilon_W2, then I don't need to
# set requires_grad = True for epsilon_W2. This will save memory.
epsilon_W2 = torch.randn(d_hidden,d_out)
# STEP 3d
# sample epsilon column vector for the second layer with dimensions d_out x 1,
# where d_out is the dimensionality of the output layer. This epsilon
# column vector will be used to construct b2, which is the d_out x 1 bias
# vector for the second layer. Note that since I do not need to compute the
# gradient of f(W,theta) with respect to epsilon_b2, then I don't need to
# set requires_grad = True for epsilon_b2. This will save memory.
epsilon_b2 = torch.randn(d_out,1)
# STEP 4a
# compute W1 = mu_W1 + log(1 + exp(rho_W1)) * epsilon_W1
W1 = mu_W1 + torch.mul(
torch.log(1 + torch.exp(rho_W1)),
epsilon_W1
)
# STEP 4b
# compute b1 = mu_b1 + log(1 + exp(rho_b1)) * epsilon_b1
b1 = mu_b1 + torch.mul(
torch.log(1 + torch.exp(rho_b1)),
epsilon_b1
)
# STEP 4c
# compute W2 = mu_W2 + log(1 + exp(rho_W2)) * epsilon_W2
W2 = mu_W2 + torch.mul(
torch.log(1 + torch.exp(rho_W2)),
epsilon_W2
)
# STEP 4d
# compute b2 = mu_b2 + log(1 + exp(rho_b2)) * epsilon_b2
b2 = mu_b2 + torch.mul(
torch.log(1 + torch.exp(rho_b2)),
epsilon_b2
)
# STEP 5a
# compute log(q(W1|mu_W1,rho_W1))
posterior_W1 = log_likelihood(W1,mu_W1,rho_W1).sum()
# STEP 5b
# compute log(q(b1|mu_b1,rho_b1))
posterior_b1 = log_likelihood(b1,mu_b1,rho_b1).sum()
# STEP 5c
# compute log(q(W2|mu_W2,rho_W2))
posterior_W2 = log_likelihood(W2,mu_W2,rho_W2).sum()
# STEP 5d
# compute log(q(b2|mu_b2,rho_b2))
posterior_b2 = log_likelihood(b2,mu_b2,rho_b2).sum()
# STEP 6a
# compute log(P(W1)). Note that since sigma = log(1 + exp(rho)), and
# if sigma = 1, then rho = log(exp(1) - 1)
rho = torch.log(torch.exp(torch.tensor(1.)) - 1)
prior_W1 = log_likelihood(W1,0,rho).sum()
# STEP 6b
# compute log(P(b1))
prior_b1 = log_likelihood(b1,0,rho).sum()
# STEP 6c
# compute log(P(W2))
prior_W2 = log_likelihood(W2,0,rho).sum()
# STEP 6d
# compute log(P(b2))
prior_b2 = log_likelihood(b2,0,rho).sum()
# STEP 7
# compute log(P(D|W1,W2,b2,b1)), which is the negative of the
# cross-entropy between all y's and y_hat's in the dataset D. Recall
# that the cross-entropy between all y's and y_hat's in a dataset of
# size N is:
#
# - \sum_{i=1}^N y_i \cdot log(\hat{y}_i)
cross_entropy = 0
# iterte over each row of the X matrix to compute the cross-entropy,
# since each row is a single x vector. Note that this method is
# inefficient, since batch matrix multiplication would be more
# efficient. However, this is for illustrative purposes only.
for j in range(X.shape[0]):
# extract the input vector x and convert it into a column vector
x = X[j,:].unsqueeze(-1)
# output of first layer. Note that W1.T means transpose of W1
out1 = act_func(torch.matmul(W1.T,x) + b1)
# output of second layer. Note that the logarithm of the softmax
# function is computed here. The .squeeze() method is used to
# remove any extra dimensions
log_y_hat = log_softmax(torch.matmul(W2.T,out1) + b2).squeeze()
# accumulate cross entropy
cross_entropy = cross_entropy + (y[j] * log_y_hat)
# STEP 8
# compute f(w,theta). Note that since W1,b1,W2, and b2 are assumed to
# be conditionally independent given their corresponding parameters mu
# and rho, then:
#
# log(q(W1,b1,W2,b2|mu_W1,rho_W1,mu_W2,rho_b2)) =
# log(q(W1|mu_W1,rho_W1)) + log(q(b1|mu_b1,rho_b1)) +
# log(q(W2|mu_W2,rho_W2)) + log(q(b2|mu_b2,rho_b2))
#
# Note that log(P(D|W1,W2,b2,b1)) is the cross-entropy computed above.
f = (posterior_W1 + posterior_b1 + posterior_W2 + posterior_b2
- prior_W1 - prior_b1 - prior_W2 - prior_b2
+ cross_entropy)
# STEP 9
# compute gradients as shown in steps 5 and 6
f.backward()
# Delta_mu = (W1.grad + b1.grad + W2.grad + b2.grad
# + mu_W1.grad + mu_b1.grad + mu_W2.grad + mu_b2.grad)
# Delta_rho = ((W1.grad * (epsilon_W1 / (1 + torch.exp(-rho_W1)))
# + b1.grad * (epsilon_b1 / (1 + torch.exp(-rho_b1)))
# + W2.grad * (epsilon_W2 / (1 + torch.exp(-rho_W2)))
# + b2.grad * (epsilon_b2 / (1 + torch.exp(-rho_b2))))
# + rho_W1.grad + rho_b1.grad + rho_W2.grad + rho_b2.grad)
# STEP 10
# update mu and rho using gradient descent
with torch.no_grad():
mu_W1 -= mu_W1.grad * learning_rate
mu_b1 -= mu_b1.grad * learning_rate
mu_W2 -= mu_W2.grad * learning_rate
mu_b2 -= mu_b2.grad * learning_rate
mu_W1.grad.zero_()
mu_b1.grad.zero_()
mu_W2.grad.zero_()
mu_b2.grad.zero_()
rho_W1 -= rho_W1.grad * learning_rate
rho_b1 -= rho_b1.grad * learning_rate
rho_W2 -= rho_W2.grad * learning_rate
rho_b2 -= rho_b2.grad * learning_rate
rho_W1.grad.zero_()
rho_b1.grad.zero_()
rho_W2.grad.zero_()
rho_b2.grad.zero_()
```