Return to Answer

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.

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.

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_()
```

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 to obtain $\mathbf{b}$: $$ \mathbf{b} = \mu + \log(1 + \exp(\rho)) \circ \epsilon $$ where these are now column vectors instead of matrices. You can then compute $\hat{y}$ as shown above.

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.