# Bayesian Logistic Regression

Contact: mihaelacr@google.com

Aim of this colab:
Logistic regression with a twist! We will use variational inference to learn a distribution over the learned parameters.

Througout this work we will be focusing on a binary classification task, on 
a small dataset (the UCD breast cancer dataset). 


## Logistic regression

Learn the weights which best classify the $(x, y)$ pairs in the dataset. 
The weights ($w$) and biases $b$, form the parameters which are learned via the following maximization problem:

\begin{equation}
\mathbb{E}_{p^*(x, y)} \log p(y|x, \theta) = \\
 \mathbb{E}_{p^*(x, y)} \left[y \log \sigma(w x + b) + (1 - y) \log(1- \sigma(w x + b)\right] 
\end{equation}

Here $\sigma$ denotes the sigmoid function.

## Bayesian logistic regression - computing the evidence lower bound

Learn a distribution over the weights which best classify the $(x, y)$ pairs in the dataset.

We now want to learn the distributional parameters of $q(w)$ in order to maximize:

\begin{equation}
\mathbb{E}_{p^*(x, y)} \log p_w(y| x) = \\
\mathbb{E}_{p^*(x, y)} \log \int p(y|x, w) p(w) \delta w = \\
\mathbb{E}_{p^*(x, y)} \log \int p(y|x, w) p(w) \frac{q(w)}{q(w)} \delta w \ge \\
\mathbb{E}_{p^*(x, y)} \mathbb{E}_{q(w)} \log \left[p(y| x, w) \frac{p(w)}{q(w)} \right]= \\
\mathbb{E}_{q(w)} \mathbb{E}_{p^*(x, y)} \log p(y|x, w)  -  KL(q(w)||p(w)) = \\
 \mathbb{E}_{q(w)}\mathbb{E}_{p^*(x, y)} \left[y \log \sigma(w x + b) + (1 - y) \log(1- \sigma(w x + b)\right] -  KL(q(w)||p(w))
\end{equation}

In Bayesian Logistic Regression, we thus learn a distribution over parameters
$w$ which can explain the data, while staying close to a chosen prior $p(w)$.

Throughout this lab, we will use Gaussian Distributions for $p(w)$ and $q(w)$.

For more details, see 
[this paper](https://pdfs.semanticscholar.org/e407/ea7fda6d152d2186f4b5e27aa04ec2d32dcd.pdf) 
or this tutorial
[https://www.ece.rice.edu/~vc3/elec633/logistic.pdf].

### Gradient estimation

In Bayesian Logistic Regression, we aim to learn the parameters of the distribution $q(w)$. We will call these parameters $\theta$. In the case of a 
Gaussian distribution, these will be the mean and covariance matrix of a
multivariate distribution. Since we will use stochastic gradient descent to learn $\theta$, we need to be able to compute the gradients with respect to 
$\theta$ of our objective. Since our objective contains an expectation with respect to $q(w)$, this can be challenging. To illustrate this, from now on we will denote $q(w)$ as $q_{\theta}(w)$.

Specifically, we are interested in:
\begin{equation}
  \nabla_{\theta} \left[\mathbb{E}_{q_{\theta}(w)} \mathbb{E}_{p^*(x, y)} \log p(y|x, w)  -  KL(q_{\theta}(w)||p(w))\right]
\end{equation}

Since in we will be using Gaussian distributions for $q_{\theta}(w)$ and $p(w)$ 
$KL(q_{\theta}(w)||p(w))$ can be computed in closed form and we can rely on 
TensorFlow's automatic differentiation to compute 
$\nabla_{\theta} KL(q_{\theta}(w)||p(w))$.

Let's now turn attention to computing $\nabla_{\theta} \mathbb{E}_{q_{\theta}(w)} \mathbb{E}_{p^*(x, y)} \log p(y| x, w)$. Unlike the KL, the integral 
$\mathbb{E}_{q_{\theta}(w)} \mathbb{E}_{p^*(x, y)} \log p(y|x, w)$
is not tractable, and thus cannot be computed in closed form. Thus, we are interested in other approaches to compute the gradient.

### Reinforce gradient estimation

A first approach to compute an estimate of the gradient is to use the
REINFORCE gradient estimator:

\begin{equation}
  \nabla_{\theta} \mathbb{E}_{q_{\theta}(w)} \mathbb{E}_{p^*(x, y)} \log p(y| x, w)  = \\
  \nabla_{\theta} \int q_{\theta}(w) \mathbb{E}_{p^*(x, y)} \log p(y| x, w) \delta w = \\
  \int \nabla_{\theta} \left[q_{\theta}(w) \mathbb{E}_{p^*(x, y)} \log p(y| x, w) \right] \delta w = \\
 \int \nabla_{\theta} \left[q_{\theta}(w)\right] \mathbb{E}_{p^*(x, y)} \log p(y| x, w) \delta w = \\
 \int q_{\theta}(w) \nabla_{\theta} \left[\log q_{\theta}(w)\right] \mathbb{E}_{p^*(x, y)} \log p(y| x, w)  \delta w = \\
 \mathbb{E}_{q_{\theta}(w)} \left[ \nabla_{\theta} \log q_{\theta}(w) \mathbb{E}_{p^*(x, y)} \log p(y| x, w) \right]
\end{equation}

We can use samples from $q_{\theta}(w)$ to compute the Monte Carlo estimate of
the last integral, to get an unbiased estimator of the true gradient. The more 
samples we use to estimate the integral, the more accurate the gradient estimate
will be.


### Pathwise gradient estimation (reparametrization)

In the REINFORCE estimator, we did not use any knowledge of the variational
distribution $q_{\theta}(w)$.
For a Gaussian distribution, we can use the reparametrization trick to obtain
an unbiased gradient estimator of  $\nabla_{\theta} \mathbb{E}_{q_{\theta}(w)} \mathbb{E}_{p^*(x, y)} \log p(y| x, w)$.

To do so, we will use the fact that 
\begin{equation}
    z \sim N(\mu, \sigma), z = \mu + \epsilon \sigma, {\text {with }} \epsilon \sim N(0, 1)
\end{equation}

Namely:
\begin{equation}
  \nabla_{\theta} \mathbb{E}_{q_{\theta}(w)} \mathbb{E}_{p^*(x, y)} \log p(y| x, w)  = \\
  \nabla_{\theta} \mathbb{E}_{p(\epsilon)} \mathbb{E}_{p^*(x, y)} \log p(y| x, \mu + \epsilon \sigma)  = \\
  \mathbb{E}_{p(\epsilon)} \nabla_{\theta} \mathbb{E}_{p^*(x, y)} \log p(y| x, \mu + \epsilon \sigma)
\end{equation}

Note that we were able to move the gradient inside the integral since 
$p(\epsilon)$ does not depend on $\theta$.
To estimate the last integral, we can use a Monte Carlo estimate using samples
of $p(\epsilon)$.


## Your tasks


*   Define the Gaussian posterior distribution.
*   Fill in the code for the ELBO
*   Fill in the reinforce function to create the surrogate loss
*   Fill in the loss function for reparametrization
*   Visualize the effects of batch size, learning rates and number of posterior samples on the learned model, as well as gradient variance.




```
import enum
import numpy as np

import math
from sklearn import datasets
from sklearn import linear_model
import collections

from matplotlib import pyplot as plt
import tensorflow as tf
import tensorflow_probability as tfp
import seaborn as sns

tfd = tfp.distributions

```


```
sns.set(rc={"lines.linewidth": 2.8}, font_scale=2)
sns.set_style("whitegrid")
```

## Get the dataset and visualize it


```
def get_data(normalize = True):
    data = datasets.load_breast_cancer()

    print(data.target_names)
    print(data.feature_names)

    features = np.array(data.data, dtype=np.float32)
    targets = np.array(data.target, dtype=np.float32)

    if normalize:
        # Note: the data dimensions have very different scales.
        # We normalize by the mean and scale of the entire dataset.
        features = features - features.mean(axis=0)
        features = features / features.std(axis=0)

    # Add a dimension of 1 (bias).
    features = np.concatenate((features, np.ones((features.shape[0], 1))), axis=1)
    features = np.array(features, dtype=np.float32)
  
    return features, targets, data.feature_names
```

### Dataset

This datasets has meaningful features: we are trying to classify based on these features whether someone has breast cancer or not. We can look at the
features we are using to classify the data below.



```
features, targets, feature_names = get_data()
```

### Trivial baseline

How well can we do with a constant classifier?

If we look at the average label, we can see how many 1s are in the data, so we know that our classifier has to get more than that accuracy


```
np.mean(targets)
```


```
features.shape
```


```
dataset_size = features.shape[0]
data_dims = features.shape[-1]
```

## Logistic regression baseline - how well can we do on this dataset?

Let's look at a simple, non Bayesian classification approach to get an idea of how well we can do on the dataset. 


```
lr = linear_model.LogisticRegression()
```


```
lr.fit(features, targets)
lr.score(features, targets)
```

## Utilities


```
def get_shape_list(tensor):
  return tensor.shape.as_list()
```


```
def get_sklearn_data_as_tensors(
      features, targets, batch_size, dataset_name='breast_cancer'):
  """Read sklearn datasets as tf.Tensors.

  Args:
    batch_size: Integer or None. If None, the entire dataset is used.
    dataset_name: A string, the name of the dataset.

  Returns:
    A tuple of size two containing two tensors of rank 2 `[B, F]`, the data
    features and targets.
  """
  dataset = tf.data.Dataset.from_tensor_slices((features, targets))
  if batch_size:
    # Shuffle, repeat, and batch the examples.
    batched_dataset = dataset.shuffle(1000).repeat().batch(batch_size)
  else:
    batch_size = features.shape[0]
    batched_dataset = dataset.repeat().batch(batch_size)

  iterator = batched_dataset.make_one_shot_iterator()
  batch_features, batch_targets = iterator.get_next()

  data_dim = features.shape[1]
  batch_features.set_shape([batch_size, data_dim])
  batch_targets.set_shape([batch_size])

  return batch_features, batch_targets
```

## Bayesian Logistic Regression - define the prior and the posterior

We will learn a distribution over parameters, so we will first define prior distribution, and then a posterior distribution over parameters which we 
will learn.

Both of these distibutions will be Gaussian.

## Task: implement the posterior distribution

You have to implement two functions 


1.   multi_normal, which takes as arguments the mean and the log scale of the normal distribution, and returns a tensorflow distribution object. Have a look at the TensorFlow distibutions package. 
2.   diagonal_gaussian_posterior which returns a tuple of two elements, one being the posterior distribution (a call to multi_normal) and the learned variables. Your task here is to define the variables we will learn. The argument `data_dims` tells you the size of the learned parameters, and thus the size of the variables of the distribution).




```
def multi_normal(loc, log_scale):
  raise NotImplementedError(
        'Implement the definition of multi normal')

def diagonal_gaussian_posterior(data_dims):
  mean = None
  log_scale = None
  learned_vars = [mean, log_scale]
  return multi_normal(loc=mean, log_scale=log_scale), learned_vars
```

## Bayesian Logistic Regression - the model

### Hypers


```
BATCH_SIZE = 500
NUM_POSTERIOR_SAMPLES = 100
```


```
prior = multi_normal(loc=tf.zeros(data_dims), log_scale=tf.zeros(data_dims))
posterior, learned_vars  = diagonal_gaussian_posterior(data_dims)
```


```
def _predictions(logits):
  return tf.to_float(logits >= 0)

def _accuracy(targets, predictions):
  # `predictions` have rank 2: batch_size, num_posterior samples.
  # We expand dims the targets to compute the accuracy per sample.
  targets = tf.expand_dims(targets, axis=1)
  return tf.reduce_mean(tf.to_float(tf.equal(targets, predictions)))

def linear_model(data, targets, posterior_samples):
  num_posterior_samples = tf.shape(posterior_samples)[0]
  logits = tf.matmul(data, posterior_samples, transpose_b=True)
  # Make targets [B, 1] to use broadcasting.
  targets = tf.expand_dims(targets, axis=1)
  targets = targets * tf.ones([1, num_posterior_samples])

  log_probs = - tf.nn.sigmoid_cross_entropy_with_logits(
      labels=targets, logits=logits)
  return log_probs, logits
```

## Task: Run the model and return the components of the elbo


```
def run_model(features, targets, prior, posterior, num_samples):
  # Sample from the posterior to obtain a set of parameters used for
  # prediction. Use `num_samples` for the number of samples. 
  posterior_samples = None

  # Compute the log probs of the data under all the models we have sampled.
  # These tensors are [num_samples, B]. Use the features and targets tensors.
  log_probs, logits = None

  # Compute the KL between the posterior and the prior.
  # Note: since we use Gaussian distributions, the closed form of the 
  # KL can be computed. Remember that distributions are objects in TensorFlow!
  kl = None

  # Sum over data. Normalize by dataset size to ensure that the gradients
  # are not very different in magnitude when we change batch size.  
  param_log_probs = tf.reduce_mean(log_probs, axis=0) * dataset_size
  param_log_probs.shape.assert_is_compatible_with([num_samples])

  return param_log_probs, kl, posterior_samples, logits 
```


```
param_log_probs, kl, posterior_samples, logits = run_model(
      features, targets, prior, posterior, NUM_POSTERIOR_SAMPLES)
predictions = _predictions(logits)
accuracy = _accuracy(targets=targets, predictions=predictions)
accuracy = tf.reduce_mean(accuracy)

```

## Task: compute the elbo from the log probs and the KL 

`per_sample_elbo` should have shape [num_samples], since we obtain a model prediction for each of the sampled parameters.  


```
per_sample_elbo = None
elbo = tf.reduce_mean(per_sample_elbo)
```

## Training loop

## Training the model via reparametrization


```
# Stochastic loss depends on the output of a sampling operation  (posterior samples)
# but tensorflow implements reparametrization by default.
per_sample_reparametrization_loss = - per_sample_elbo
reparametrization_loss = tf.reduce_mean(per_sample_reparametrization_loss)
```


```
optimizer = tf.train.GradientDescentOptimizer(0.0001)
reparam_min_op = optimizer.minimize(reparametrization_loss)
```


```
NUM_ITERATIONS = 1000
```


```
sess = tf.Session()

# Initialize all variables
sess.run(tf.initialize_all_variables())

reparam_accuracies = []
reparam_kls = []
reparam_elbos = []

for i in xrange(NUM_ITERATIONS):
  sess.run(reparam_min_op)

  if i % 10 == 0:
    reparam_acc, reparam_kl, reparam_elbo = sess.run([accuracy, kl, elbo])
    reparam_accuracies += [reparam_acc]
    reparam_kls += [reparam_kl]
    reparam_elbos += [reparam_elbo]

    print('Iteration {}. Elbo {}. KL {}'.format(
        i, reparam_elbo, reparam_kl))

reparam_learned_mean, reparam_learned_log_scale = sess.run(learned_vars)
```


```
fig, axes = plt.subplots(1, 3, figsize=(3*8,5))

axes[0].plot(reparam_elbos, label='ELBO')
axes[0].set_title('Time', fontsize=15)
axes[0].set_ylim((-500, -50))
axes[0].legend()

axes[1].plot(reparam_kls, label='KL')
axes[1].set_title('Time', fontsize=15)
axes[1].legend()

axes[2].plot(reparam_accuracies, label='Accuracy')
axes[2].set_title('Time', fontsize=15)
axes[2].legend()
```

## Training the model via reinforce

## Task: define the surrogate reinforce loss for the log prob term.

Note: we do not need to use reinforce the KL term, since we can compute the KL analytically for the two Gaussians, and can use standard backprop. 

For the reinforce implementation, instead of changing the gradients, we will change the loss such that TensorFlow's automatic differentiation does the right thing. This is called a surrogate loss. 

The gradient we want to obtain is the following (since we want to minimize the loss, we add a minus):
\begin{equation}
- \mathbb{E}_{q_{\theta}(w)} \left[ \nabla_{\theta} \log q_{\theta}(w) \mathbb{E}_{p^*(x, y)} \log p(y| x, w) \right]
\end{equation}

so we need to construct a loss which has this gradient, by ensuring that the gradients will flow through the $\log q_{\theta}(w)$ term, but not the rest. 


```
# Since TensorFlow does automatic differentation, we have to use 
# tf.stop_gradient to ensure the gradient only flows to
# the log_{q_\theta}(w) term.
per_sample_reinforce_loss =  None   # Shape [num_samples, batch size].
# Reduce now over the number of samples.
reinforce_loss = tf.reduce_mean(per_sample_reinforce_loss) + kl
```


```
optimizer = tf.train.GradientDescentOptimizer(0.0001)
reinforce_min_op = optimizer.minimize(reinforce_loss)
```


```
NUM_ITERATIONS = 1000
```


```
reinforce_accuracies = []
reinforce_kls = []
reinforce_elbos = []

sess = tf.Session()

# Initialize all variables
sess.run(tf.initialize_all_variables())

for i in xrange(NUM_ITERATIONS):
  sess.run(reinforce_min_op)

  if i % 10 == 0:
    reinforce_acc, reinforce_kl, reinforce_elbo = sess.run([accuracy, kl, elbo])
    reinforce_accuracies += [reinforce_acc]
    reinforce_kls += [reinforce_kl]
    reinforce_elbos += [reinforce_elbo]

    print('Iteration {}. Elbo {}. KL {}'.format(
        i, reinforce_elbo, reinforce_kl))

reinforce_learned_mean, reinforce_learned_log_scale = sess.run(learned_vars)
```


```
fig, axes = plt.subplots(1, 3, figsize=(3*8,5))

axes[0].plot(reinforce_elbos, label='ELBO')
axes[0].set_title('Time', fontsize=15)
axes[0].set_ylim((-500, -50))
axes[0].legend()

axes[1].plot(reinforce_kls, label='KL')
axes[1].set_title('Time', fontsize=15)
axes[1].legend()

axes[2].plot(reinforce_accuracies, label='Accuracy')
axes[2].set_title('Time', fontsize=15)
axes[2].legend()
```

## Task: Analysis

* How sensitive are the different estimators to the number of posterior samples used? Try 1, 10, 100.
* How sensitive are the different estimators to the learning rate used? Try 1e-4, 1e-3 and 1e-2 for both estimators.
* How sensitive are the different estimators to batch size? Try batch size 10 for both estimators.

# Visualize learned parameters

* Let's visualize the learned parameters, as well as check the different uncertanties around different parameters
* Can we now see which parameters are most important in the classification problem?


```
learned_mean = reparam_learned_mean
learned_scale = np.exp(reparam_learned_log_scale)
```


```
PARAM_INDEX = 10
```


```
def gaussian(x, mu, sig):
    return np.exp(-np.power(x - mu, 2.) / (2 * np.power(sig, 2.)))
```


```
mean = learned_mean[PARAM_INDEX]
scale = learned_scale[PARAM_INDEX]
x_values = np.linspace(-3, 3, num=1000)
y_values = [gaussian(x, mean, scale) for x in x_values]

plt.figure(figsize=(7,5))
plt.plot(x_values, y_values)
plt.vlines(
      learned_mean[PARAM_INDEX], 
      ymin=0, 
      ymax=gaussian(mean, mean, scale),
      colors='g',
      alpha=0.8)

plt.grid('off')
x_std_values = np.linspace(mean - scale, mean + scale, num=1000)
y_std_values = [gaussian(x, mean, scale) for x in x_std_values]
plt.fill_between(x_std_values, y_std_values, 0, alpha=0.3, color='b')

plt.title('Parameter name= {}'.format(feature_names[PARAM_INDEX]))
```

## Visualize gradient variance - which estimator has lower variance?

* Let's assess the gradient variance per of each estimator, by looking at how much gradients vary for each estimators. 


```
def per_batch_grads(ys, xs):
  """Computes the gradients of ys[i] wrt xs."""
  grads = []

  for y in tf.unstack(ys, axis=0):
    grads.append(tf.squeeze(tf.gradients(y, xs)[0]))
  return tf.stack(grads, axis=0)
```


```
def compute_grad_var(grads, param_index):
  return np.mean(grads[:, param_index]), np.std(grads[:, param_index])**2
```


```
per_batch_reinforce_grads = per_batch_grads(per_sample_reinforce_loss, learned_vars[0])
per_batch_reinforce_grads.shape
```


```
per_batch_reparam_grads = per_batch_grads(per_sample_reparametrization_loss, learned_vars[0])
per_batch_reparam_grads.shape
```


```
reinforce_grads = []

sess = tf.Session()

# Initialize all variables
sess.run(tf.initialize_all_variables())

for i in xrange(NUM_ITERATIONS):
  sess.run(reinforce_min_op)

  if i % 10 == 0:
    reinforce_grads.append(sess.run(per_batch_reinforce_grads))
```


```
reparam_grads = []

sess = tf.Session()

# Initialize all variables
sess.run(tf.initialize_all_variables())

for i in xrange(NUM_ITERATIONS):
  sess.run(reparam_min_op)

  if i % 10 == 0:
    reparam_grads.append(sess.run(per_batch_reparam_grads))
```


```
PARAM_INDEX = 19
```


```
reinforce_grad_vars = []
for timestep_grad in reinforce_grads:
  _, reinforce_grad_var = compute_grad_var(timestep_grad, PARAM_INDEX)
  reinforce_grad_vars.append(reinforce_grad_var)
```


```
reparam_grad_vars = []
for timestep_grad in reparam_grads:
  _, reparam_grad_var = compute_grad_var(timestep_grad, PARAM_INDEX)
  reparam_grad_vars.append(reparam_grad_var)
```


```
plt.figure(figsize=(7,5))
plt.plot(reinforce_grad_vars, label='reinforce')
plt.plot(reparam_grad_vars, label='reparametrization')
plt.yscale('log')
plt.xlabel('It')
plt.ylabel('Gradient variance')
plt.legend()
```

### Task: Analysis

* What do you notice about gradient variance as training progresses?
* What do you notice about the variance of the different estimators?
* How do batch size and number of posterior samples affect the variance of the different estimators?
* How does gradient variance affect convergence?
* Which gradient estimator would you use in practice for this problem?

## More

You can find similar experiments with more analysis and more gradient estimators in [this paper](https://arxiv.org/pdf/1906.10652.pdf). See Section 8.3.

03e2d0dbe0db5fdf6239d207bd27cd6db7c1801f
ipynb
Jupyter Notebook
unsupervised/[Exercises]_1_bayesian_logistic_regression.ipynb
hazratali/PracticalSessions
95c10a71267f6fb1f61a6bd6f8ea33041365ceb5
MIT
2020-02-11T23:34:05.000Z
2021-08-20T02:58:08.000Z
unsupervised/[Exercises]_1_bayesian_logistic_regression.ipynb
hazratali/PracticalSessions
95c10a71267f6fb1f61a6bd6f8ea33041365ceb5
MIT
unsupervised/[Exercises]_1_bayesian_logistic_regression.ipynb
hazratali/PracticalSessions
95c10a71267f6fb1f61a6bd6f8ea33041365ceb5
MIT
2019-06-27T12:15:38.000Z
2019-10-21T12:29:54.000Z
Qwen/Qwen-72B
1. YES
2. YES