How do i calculate the gradients for the bayes-by-backprop algorithm in R?

Question

UPDATE - edits with my code progress are added at the end of this post.

I'm trying to create a neural network in R based on the bayes-by-backprop model as described in this paper. As inspiration for the code i'm using the python description on this website. However, in this paper they use an autograd function and i'm trying to calculate the gradient manually. The important algorithm from the paper is this one:

I think i've got most of the code ready, but i'm getting stuck on calculating the partial derivatives needed to calculate the gradients for mu (mean) and rho (parameter of the standard deviation).

This is the backpropagation function that i'm trying to make. I'm trying to adapt it from a regular neural network that i created before. The parameters are:

• y = ground truth
• Wn = list, every entry represents a layer of the network, containing a sample of the probability distributions in that layor
• Un = list similar to Wn, the mus (mean) of the probability distributions in the posterior
• Rn = list similar to Wn, the rhos (parameter used to define the standard deviation)
• sigma = list similar to Wn, the standard deviation
• eps = list similar to Wn, the random variable that is used for the reparametrization trick
• ff = a list of the outputs of each feedforward layer, the last output is the model prediction.

backpropagate <- function(y, Wn, Un, Rn, sigma, ff, eps, learn_rate = 0.01) {
  
  y_hat <- ff[[3]]
  
  loss <- list()
  dWn <- list()
  
  loss[[2]] <- loss_function(y_hat, y, Wn, Un, sigma)
  
  dUn[[2]] <- ???
  dRn[[2]] <- ???
  
  loss[[1]] <- ???
  
  dUn[[1]] <- ???
  dRn[[1]] <- ???
    
  Un[[1]] <- Un[[1]] - learn_rate * dUn[[1]]
  Un[[2]] <- Un[[2]] - learn_rate * dUn[[2]]
  
  Rn[[1]] <- Rn[[1]] - learn_rate * dRn[[1]]
  Rn[[2]] <- Rn[[2]] - learn_rate * dRn[[2]]
  
  return(list(Un, Rn))
}

How do i calculate the gradients of mu and rho? I can do partial derivatives on polynomial and exponential functions but mu and rho are defined here only as probability distributions. The loss function is defined as the $$posterior - prior * likelihood$$ but how do i take the derivative w.r.t. mu and rho of those?

$$posterior = \log(P(w)) = \sum_i \log(N(w_i | \mu, \sigma^2)) $$ $$prior = \log(q(w|\theta)) = \sum_i \log(N(w_i | 0, 1)) $$ $$likelihood = \log(P(D|w)) = y * log(softmax(\hat{y})) $$

For reference:

My main training function looks like this:

train <- function(x, y, neurons = 32, layers = 3, rate = 0.01, iterations = 10000) {
  
  d <- ncol(x) + 1
  
  Un <- list()
  Rn <- list()
  
  Un[[1]] <- matrix(rnorm(d * neurons), d, neurons)        # generate mus for the gaussian distribution of each neuron. matrix d (rows) x neurons (columns)
  Un[[2]] <- matrix(rnorm(neurons + 1), neurons + 1, 1)
  
  Rn[[1]] <- matrix(-3, d, neurons)        # generate the rhos to calculate the standard deviation later
  Rn[[2]] <- matrix(-3, neurons + 1, 1)
  
  for (i in 1:iterations) {
    
    Wn <- list()
    sigma <- list()
    eps <- list()
    
    sigma[[1]] <- log(1 + exp(Rn[[1]]))     # calculate the standard deviation from rho
    sigma[[2]] <- log(1 + exp(Rn[[2]]))
    
    eps[[1]] <- matrix(rnorm(d * neurons, 0, 0.1), nrow = d, ncol = neurons)     # generate a random number epsilon for every neuron
    eps[[2]] <- matrix(rnorm(neurons + 1, 0, 0.1), neurons + 1, 1)
    
    Wn[[1]] <- Un[[1]] + sigma[[1]] * eps[[1]]     # take one sample of the posterior
    Wn[[2]] <- Un[[2]] + sigma[[2]] * eps[[2]]

    ff <- feedforward(x, Wn)
    backlist <- backpropagate(y, Wn, Un, Rn, sigma, ff, eps, learn_rate = rate)
    
    Un <- backlist[[1]]
    Rn <- backlist[[2]]
  }
  return(Wn)
}

And then my feedforward plus helper functions:

feedforward <- function(x, Wn) {
  
  ff <- list(x)
  Zn <- cbind(1, ff[[1]]) %*% Wn[[1]]
  ff[[2]] <- sigmoid::relu(Zn)
  Zn <- cbind(1, ff[[2]]) %*% Wn[[2]]
  ff[[3]] <- Zn

  return(ff)
}

log_softmax <- function(y_hat, y){
  y * log(exp(y_hat) / sum(exp(y_hat)))
}

log_gaussian <- function(x, mu, sigma){
  return(-0.5 * log(2*pi) - log(sigma) - (x - mu)**2 / (2 * sigma**2))
}

loss_function <- function(y_hat, y, Wn, Un, sigma){
  
  posterior <- sum(log_gaussian(Wn[[i]], Un[[i]], sigma[[i]]))
  
  prior <- sum(log_gaussian(Wn[[i]], 0, 1))
  
  likelihood <- log_softmax(y_hat, y)
  
  return(posterior - prior * likelihood)
}

---

EDIT: here's the code that i've tried in Torch for R so far. I think i have everything correct up to the point where i need to calculate the posterior and prior for the loss function. I don't know how i can call the gaussian function in torch to generate the distributions for the posterior and the prior.

For the posterior, i need to generate a distribution for every neuron with mean $\mu$ and standard deviation $\sigma$, which are the parameters for each neuron.

For the prior, i need to generate a distribution for every neuron with mean $0$ and standard deviation $1$.

The code is basically adapted from this page on the torch website.

library(torch)

### generate training data -----------------------------------------------------

# input dimensionality (number of input features)
d_in <- 3
# output dimensionality (number of predicted features)
d_out <- 1
# number of observations in training set
n <- 100


# create random data
x <- torch_randn(n, d_in)
y <- x[, 1, NULL] * 0.2 - x[, 2, NULL] * 1.3 - x[, 3, NULL] * 0.5 + torch_randn(n, 1)


### initialize weights ---------------------------------------------------------

# dimensionality of hidden layer
d_hidden <- 32
# mean values of the neurons in the first layer
mu1 <- torch_randn(d_in, d_hidden, requires_grad = TRUE)
# mean values of the neurons in the second layer
mu2 <- torch_randn(d_hidden, d_out, requires_grad = TRUE)

# rho values (parameters for the standard deviation) for the first layer
rho1 <- torch_randn(d_in, d_hidden, requires_grad = TRUE)
# rho values for the second layer
rho2 <- torch_randn(d_hidden, d_out, requires_grad = TRUE)

### network parameters ---------------------------------------------------------

learning_rate <- 1e-4

### training loop --------------------------------------------------------------

for (t in 1:200) {
  
  # random error
  eps1 <- torch_randn(d_in, d_hidden, requires_grad = TRUE)
  eps2 <- torch_randn(d_hidden, d_out, requires_grad = TRUE)
  
  ### -------- Sample variational posterior --------
  w1 <- torch_add(mu1, torch_multiply(torch_log(torch_add(torch_exp(rho1), 1))), eps1)
  w2 <- torch_add(mu2, torch_multiply(torch_log(torch_add(torch_exp(rho2), 1))), eps2)
  
  ### -------- Forward pass --------
  y_pred <- x$mm(w1)$clamp(min = 0)$mm(w2)$clamp(min = 0)
  
  ### -------- Likelihood -------- 
  likelihood <- torch_multiply(y, nnf_log_softmax(y_hat))
  
  ### -------- Prior -------- 
  
  
  ### -------- Posterior -------- 
  
  
  ### -------- compute loss -------- 
  loss <- (posterior - prior * likelihood)$sum()
  if (t %% 10 == 0)
    cat("Epoch: ", t, "   Loss: ", loss$item(), "\n")
  
  ### -------- Backpropagation --------
  
  # compute gradient of loss w.r.t. all tensors with requires_grad = TRUE
  loss$backward()
  
  ### -------- Update weights -------- 
  
  # Wrap in with_no_grad() because this is a part we DON'T 
  # want to record for automatic gradient computation
  with_no_grad({
    mu1 <- mu1$sub_(learning_rate * mu1$grad)
    mu2 <- mu2$sub_(learning_rate * mu2$grad)
    rho1 <- rho1$sub_(learning_rate * rho1$grad)
    rho2 <- rho2$sub_(learning_rate * rho2$grad)  
    
    # Zero gradients after every pass
    mu1$grad$zero_()
    mu2$grad$zero_()
    rho1$grad$zero_()
    rho2$grad$zero_()  
  })
  
}

---

EDIT2: This is the R code that i have right now, using torch.

the log_gaussian function is copied directly from this website.

I create 100 observations with 3 features, with a hidden layer of 32 neurons and a single output. the output is either 1 or 0.

It seems to be working, but after 20 or so iterations the loss suddenly returns NaN and the network stops learning. I have some debugging still to do.

library(torch)
library(Rlab)

### generate training data -----------------------------------------------------

# input dimensionality (number of input features)
d_in <- 3
# output dimensionality (number of predicted features)
d_out <- 1
# number of observations in training set
n <- 100


# create random data
x <- torch_randn(n, d_in)
y <- torch_tensor(rbern(n, 0.6))


### initialize weights ---------------------------------------------------------

# dimensionality of hidden layer
d_hidden <- 32

# mean values of the neurons in the first and second layers
mu1 <- torch_randn(d_in, d_hidden, requires_grad = TRUE)
mu2 <- torch_randn(d_hidden, d_out, requires_grad = TRUE)
bmu1 <- torch_randn(1, d_hidden, requires_grad = TRUE)
bmu2 <- torch_randn(1, d_out, requires_grad = TRUE)

# rho values (parameters for the standard deviation) for the first and second layers
rho1 <- torch_randn(d_in, d_hidden, requires_grad = TRUE)
rho2 <- torch_randn(d_hidden, d_out, requires_grad = TRUE)
brho1 <- torch_randn(1, d_hidden, requires_grad = TRUE)
brho2 <- torch_randn(1, d_out, requires_grad = TRUE)

### network parameters ---------------------------------------------------------

learning_rate <- 1e-11

### training loop --------------------------------------------------------------

for (t in 1:200) {
  
  # random error
  eps1 <- torch_randn(d_in, d_hidden)
  eps2 <- torch_randn(d_hidden, d_out)
  eps3 <- torch_randn(d_in, d_hidden)
  eps4 <- torch_randn(d_hidden, d_out)
  
  beps1 <- torch_randn(1, d_hidden)
  beps2 <- torch_randn(1, d_out)
  beps3 <- torch_randn(1, d_hidden)
  beps4 <- torch_randn(1, d_out)
  
  s1 <- torch_log(torch_add(torch_exp(rho1), 1))
  s2 <- torch_log(torch_add(torch_exp(rho2), 1))
  
  bs1 <- torch_log(torch_add(torch_exp(brho1), 1))
  bs2 <- torch_log(torch_add(torch_exp(brho2), 1))
  
  ### -------- Sample variational posterior --------
  w1 <- rho1$exp()$add(1)$log()$multiply(eps1)$add(mu1)
  w12 <- rho1$exp()$add(1)$log()$multiply(eps3)$add(mu1)
  w2 <- rho2$exp()$add(1)$log()$multiply(eps2)$add(mu2)
  w22 <- rho2$exp()$add(1)$log()$multiply(eps4)$add(mu2)
  
  b1 <- brho1$exp()$add(1)$log()$multiply(beps1)$add(bmu1)
  b12 <- brho1$exp()$add(1)$log()$multiply(beps3)$add(bmu1)
  b2 <- brho2$exp()$add(1)$log()$multiply(beps2)$add(bmu2)
  b22 <- brho2$exp()$add(1)$log()$multiply(beps4)$add(bmu2)
  
  ### -------- Forward pass --------
  softmax <- nn_softmax(1)
  y_hat <- softmax(x$mm(w1)$add(b1)$clamp(min = 0)$mm(w2)$add(b2))
  y_hat2 <- softmax(x$mm(w12)$add(b12)$clamp(min = 0)$mm(w22)$add(b22))
  
  ### -------- Likelihood -------- 
  likelihood <- y_hat$log()$multiply(y)$sum()
  likelihood2 <- y_hat2$log()$multiply(y)$sum()
  
  ### -------- Prior -------- 
  
  prior1 <- torch_tensor(0)$subtract(w1)$square()$divide(2)$subtract(torch_log(2*pi))$sum()
  prior2 <- torch_tensor(0)$subtract(w2)$square()$divide(2)$subtract(torch_log(2*pi))$sum()
  prior3 <- torch_tensor(0)$subtract(b1)$square()$divide(2)$subtract(torch_log(2*pi))$sum()
  prior4 <- torch_tensor(0)$subtract(b2)$square()$divide(2)$subtract(torch_log(2*pi))$sum()
  
  prior <- prior1$add(prior2)$add(prior3)$add(prior4)
  
  prior12 <- torch_tensor(0)$subtract(w12)$square()$divide(2)$subtract(torch_log(2*pi))$sum()
  prior22 <- torch_tensor(0)$subtract(w22)$square()$divide(2)$subtract(torch_log(2*pi))$sum()
  prior32 <- torch_tensor(0)$subtract(b12)$square()$divide(2)$subtract(torch_log(2*pi))$sum()
  prior42 <- torch_tensor(0)$subtract(b22)$square()$divide(2)$subtract(torch_log(2*pi))$sum()
  
  priort2 <- prior12$add(prior22)$add(prior32)$add(prior42)
  
  ### -------- Posterior -------- 
  posterior1 <- w1$subtract(mu1)$square()$divide(s1$square()$multiply(2))$subtract(s1$log())$subtract(torch_log(2*pi))$sum()
  posterior2 <- w2$subtract(mu2)$square()$divide(s2$square()$multiply(2))$subtract(s2$log())$subtract(torch_log(2*pi))$sum()
  posterior3 <- b1$subtract(bmu1)$square()$divide(bs1$square()$multiply(2))$subtract(bs1$log())$subtract(torch_log(2*pi))$sum()
  posterior4 <- b2$subtract(bmu2)$square()$divide(bs2$square()$multiply(2))$subtract(bs2$log())$subtract(torch_log(2*pi))$sum()
  
  posterior <- posterior1$add(posterior2)$add(posterior3)$add(posterior4)
  
  posterior12 <- w12$subtract(mu1)$square()$divide(s1$square()$multiply(2))$subtract(s1$log())$subtract(torch_log(2*pi))$sum()
  posterior22 <- w22$subtract(mu2)$square()$divide(s2$square()$multiply(2))$subtract(s2$log())$subtract(torch_log(2*pi))$sum()
  posterior32 <- b12$subtract(bmu1)$square()$divide(bs1$square()$multiply(2))$subtract(bs1$log())$subtract(torch_log(2*pi))$sum()
  posterior42 <- b22$subtract(bmu2)$square()$divide(bs2$square()$multiply(2))$subtract(bs2$log())$subtract(torch_log(2*pi))$sum()
  
  posteriort2 <- posterior12$add(posterior22)$add(posterior32)$add(posterior42)
  
  ### -------- compute loss -------- 
  loss <- posterior$subtract(prior)$subtract(likelihood)$add(posteriort2$subtract(priort2)$subtract(likelihood2))
  if (t %% 10 == 0 | t < 10)
    cat("Epoch: ", t, "   Loss: ", loss$item(), "\n")
  
  ### -------- Backpropagation --------
  
  # compute gradient of loss w.r.t. all tensors with requires_grad = TRUE
  loss$backward()
  
  ### -------- Update weights -------- 
  
  # Wrap in with_no_grad() because this is a part we DON'T 
  # want to record for automatic gradient computation
  with_no_grad({
    mu1 <- mu1$sub_(learning_rate * mu1$grad)
    mu2 <- mu2$sub_(learning_rate * mu2$grad)
    rho1 <- rho1$sub_(learning_rate * rho1$grad)
    rho2 <- rho2$sub_(learning_rate * rho2$grad)
    bmu1 <- bmu1$sub_(learning_rate * bmu1$grad)
    bmu2 <- bmu2$sub_(learning_rate * bmu2$grad)
    brho1 <- brho1$sub_(learning_rate * brho1$grad)
    brho2 <- brho2$sub_(learning_rate * brho2$grad)  
    
    # Zero gradients after every pass
    mu1$grad$zero_()
    mu2$grad$zero_()
    rho1$grad$zero_()
    rho2$grad$zero_()
    bmu1$grad$zero_()
    bmu2$grad$zero_()
    brho1$grad$zero_()
    brho2$grad$zero_() 
  })

Answer 1

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:

1. 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.
2. Sample $\epsilon$ from the multivariate normal distribution $N(0,I)$.
3. 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 the torch_randn function. 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 also a column vector because the authors use the element-wise multiplication operation $\circ$ in their paper to multiply $\epsilon$ by $\log(1 + \exp(\rho))$.
4. Compute $\mathbf{w} = \mu + \log(1 + \exp(\rho)) \circ \epsilon$.
5. 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.
6. 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.
7. 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.
8. Compute $f(\mathbf{w},\theta) = \log(q(\mathbf{w}|\theta)) - \log(P(\mathbf{w})) - \log(P(\mathcal{D}|\mathbf{w}))$.
9. Compute the gradients as given in steps 5 and 6 of the algorithm.
10. Update $\mu$ and $\rho$ as given in step 7 of the algorithm.
11. 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. You can (almost) translate quickly between the two by substituting . in Python with $ in R. 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_()
```