In this blog, we shall discuss on how to implement probabilistic deep learning models using Tensorflow. The problems to be discussed in this blog appeared in the exercises / projects in the coursera course “Probabilistic Deep Learning“, by Imperial College, London, as a part of TensorFlow 2 for Deep Learning Specialization. The problem statements / descriptions are taken from the course itself.

Naive Bayes and logistic regression with Tensorflow Probability

In this problem, we shall develop a Naive Bayes classifier model to the Iris dataset using Distribution objects from TensorFlow Probability. We shall also explore the connection between the Naive Bayes classifier and logistic regression.

The Iris dataset

In this problem, we shall use the Iris dataset. It consists of 50 samples from each of three species of Iris (Iris setosa, Iris virginica and Iris versicolor). Four features were measured from each sample: the length and the width of the sepals and petals, in centimeters. For a reference, see the following papers:

• R. A. Fisher. “The use of multiple measurements in taxonomic problems”. Annals of Eugenics. 7 (2): 179–188, 1936.

Our goal will be to construct a Naive Bayes classifier model that predicts the correct class from the sepal length and sepal width features. Under certain assumptions about this classifier model, we shall explore the relation to logistic regression.

The following figures show the 3 categories of flowers from the Iris dataset, namely, setosa, versicolor and verginica, respectively.

Let’s start by importing the required libraries.

Load and prepare the data

We will first read in the Iris dataset, and split the dataset into training and test sets, using the following code snippet.

Naive Bayes classifier

We will briefly review the Naive Bayes classifier model. The fundamental equation for this classifier is Bayes’ rule:

In the above, dd is the number of features or dimensions in the inputs X (in our case d=2), and K is the number of classes (in our case K=3). The distribution P(Y) is the class prior distribution, which is a discrete distribution over K classes. The distribution P(X|Y) is the class-conditional distribution over inputs.

The Naive Bayes classifier makes the assumption that the data features Xi are conditionally independent give the class Y (the ‘naive’ assumption). In this case, the class-conditional distribution decomposes as follows:

This simplifying assumption means that we typically need to estimate far fewer parameters for each of the distributions P(Xi|Y=yk) instead of the full joint distribution P(X|Y=yk).

Once the class prior distribution and class-conditional densities are estimated, the Naive Bayes classifier model can then make a class prediction for a new data input according to

Define the class prior distribution

We will begin by defining the class prior distribution. To do this we will simply take the maximum likelihood estimate, given by

where the superscript (n) indicates the n-th dataset example, and N is the total number of examples in the dataset. The above is simply the proportion of data examples belonging to class k.

Let’s now define a function that builds the prior distribution from the training data, and returns it as a Categorical Distribution object.

• The input to your function y will be a numpy array of shape (num_samples,)
• The entries in y will be integer labels k=0,1,…,K−1
• The function should build and return the prior distribution as a Categorical distribution object
  • The probabilities for this distribution will be a length-KK vector, with entries corresponding to P(Y=yk) for k=0,1,…,K−1
  • Your function should work for any value of K≥1
  • This Distribution will have an empty batch shape and empty event shape

Define the class-conditional densities

Let’s now turn to the definition of the class-conditional distributions P(Xi|Y=yk) for i=0,1 and k=0,1,2. In our model, we will assume these distributions to be univariate Gaussian:

with mean parameters μik and standard deviation parameters σik, twelve parameters in all. We will again estimate these parameters using maximum likelihood. In this case, the estimates are given by

Note that the above are just the means and variances of the sample data points for each class.

Let’s now implement a function the computes the class-conditional Gaussian densities, using the maximum likelihood parameter estimates given above, and returns them in a single, batched MultivariateNormalDiag Distribution object.

• The inputs to the function are
  • a numpy array x of shape (num_samples, num_features) for the data inputs
  • a numpy array y of shape (num_samples,) for the target labels
• The function should work for any number of classes K≥1 and any number of features d≥1

We can visualise the class-conditional densities with contour plots by running the cell below. Notice how the contours of each distribution correspond to a Gaussian distribution with diagonal covariance matrix, since the model assumes that each feature is independent given the class.

Make predictions from the model

Now the prior and class-conditional distributions are defined, you can use them to compute the model’s class probability predictions for an unknown test input, according to

The class prediction can then be taken as the class with the maximum probability:

Let’s now implement a function to return the model’s class probabilities for a given batch of test inputs of shape (batch_shape, 2), where the batch_shape has rank at least one.

• The inputs to the function are the prior and class_conditionals distributions, and the inputs x
• The function should use these distributions to compute the probabilities for each class k as above
  • As before, the function should work for any number of classes K≥1
• It should then compute the prediction by taking the class with the highest probability
• The predictions should be returned in a numpy array of shape (batch_shape)

Binary classifier

We will now draw a connection between the Naive Bayes classifier and logistic regression.

First, we will update our model to be a binary classifier. In particular, the model will output the probability that a given input data sample belongs to the ‘Iris-Setosa’ class: P(Y=y0|X1,…,Xd). The remaining two classes will be pooled together with the label y1.

We will also make an extra modelling assumption that for each class kk, the class-conditional distribution P(Xi|Y=yk) for each feature i=0,1, has standard deviation σi, which is the same for each class k.

This means there are now six parameters in total: four for the means μik and two for the standard deviations σiσi (i,k=0,1).

We will again use maximum likelihood to estimate these parameters. The prior distribution will be as before, with the class prior probabilities given by

We will use our previous function get_prior to redefine the prior distribution.

For the class-conditional densities, the maximum likelihood estimate for the means are again given by

However, the estimate for the standard deviations σi is updated. There is also a closed-form solution for the shared standard deviations, but we will instead learn these from the data.

Let’s now implement a function that takes the training inputs and target labels as input, as well as an optimizer object, number of epochs and a TensorFlow Variable. This function should be written according to the following spec:

• The inputs to the function are:
  • a numpy array x of shape (num_samples, num_features) for the data inputs
  • a numpy array y of shape (num_samples,) for the target labels
  • a tf.Variable object scales of length 2 for the standard deviations σiσi
  • optimiser: an optimiser object
  • epochs: the number of epochs to run the training for
• The function should first compute the means μik of the class-conditional Gaussians according to the above equation
• Then create a batched multivariate Gaussian distribution object using MultivariateNormalDiag with the means set to μik and the scales set to scales
• Run a custom training loop for epochs number of epochs, in which:
  • the average per-example negative log likelihood for the whole dataset is computed as the loss
  • the gradient of the loss with respect to the scales variables is computed
  • the scales variables are updated by the optimiser object
• At each iteration, save the values of the scales variable and the loss
• The function should return a tuple of three objects:
  • a numpy array of shape (epochs,) of loss values
  • a numpy array of shape (epochs, 2) of values for the scales variable at each iteration
  • the final learned batched MultivariateNormalDiag distribution object

print("Class conditional means:")
print(class_conditionals_binary.loc.numpy())
print("\nClass conditional standard deviations:")
print(class_conditionals_binary.stddev().numpy())
Class conditional means:
[[4.9692307 3.3820512]
 [6.2172837 2.8814814]]

Class conditional standard deviations:
[[0.5590086  0.34253535]
 [0.5590086  0.34253535]]
# Plot the loss and convergence of the standard deviation parameters

fig, ax = plt.subplots(1, 2, figsize=(14, 5))
ax[0].plot(nlls)
ax[0].set_title("Loss vs epoch")
ax[0].set_xlabel("Epoch")
ax[0].set_ylabel("Average negative log-likelihood")
for k in [0, 1]:
    ax[1].plot(scales_arr[:, k], color=label_colours_binary[k], label=labels_binary[k])
ax[1].set_title("Standard deviation ML estimates vs epoch")
ax[1].set_xlabel("Epoch")
ax[1].set_ylabel("Standard deviation")
plt.legend()
plt.show()

We can also plot the contours of the class-conditional Gaussian distributions as before, this time with just binary labelled data. Notice the contours are the same for each class, just with a different centre location.

We can also plot the decision regions for this binary classifier model, notice that the decision boundary is now linear.

The following animation shows how we can learn the standard deviation parameters for the class-conditional distributions for Naive Bayes using tensorflow, with the above code snippet.

Link to logistic regression

In fact, we can see that our predictive distribution P(Y=y0|X) can be written as follows:

With our additional modelling assumption of a shared covariance matrix Σ, it can be shown (using the Gaussian pdf) that a is in fact a linear function of X:

The model therefore takes the form P(Y=y0|X)=σ(wTX+w0), with weights w∈R2 and bias w0∈R. This is the form used by logistic regression, and explains why the decision boundary above is linear.

In the above we have outlined the derivation of the generative logistic regression model. The parameters are typically estimated with maximum likelihood, as we have done.

Finally, we will use the above equations to directly parameterize the output Bernoulli distribution of the generative logistic regression model.

Let’s now write the following function, according to the following specification:

• The inputs to the function are:
  • the prior distribution prior over the two classes
  • the (batched) class-conditional distribution class_conditionals
• The function should use the parameters of the above distributions to compute the weights and bias terms w and w0 as above
• The function should then return a tuple of two numpy arrays for w and w0

We can now use these parameters to make a contour plot to display the predictive distribution of our logistic regression model.

Probabilistic generative models

Let’s start with generative models, using normalizing flow networks and the variational autoencoder algorithm. We shall create a synthetic dataset with a normalizing flow with randomised parameters. This dataset will then be used to train a variational autoencoder, and the trained model will be used to interpolate between the generated images. The concepts to be used will be

• Distribution objects
• Probabilistic layers
• Bijectors
• ELBO optimization
• KL divergence Regularizers.

The next figure represents the theory required for the implementation:

Let’s start by running importing the required libraries first. 

We shall create our own image dataset from contour plots of a transformed distribution using a random normalizing flow network and then use the variational autoencoder algorithm to train generative and inference networks, and synthesise new images by interpolating in the latent space.

The normalising flow

• To construct the image dataset, you will build a normalizing flow to transform the 2-D Gaussian random variable z = (z1, z2), which has mean 0 and covariance matrix Σ=σ^2.I2, with σ=0.3.
• This normalizing flow uses bijectors that are parameterized by the following random variables:
  • θ ∼ U[0,2π)
  • a ∼ N(3,1)

The complete normalising flow is given by the following chain of transformations:

• f1(z)=(z1,z2−2)
• f2(z)=(z1,z2/2),
• f3(z)=(z1,z2+a.z1^2),
• f4(z)=R.z, where R is a rotation matrix with angle θ,
• f5(z)=tanh(z), where the tanh function is applied elementwise.

The transformed random variable x is given by x=f5(f4(f3(f2(f1(z))))).

• We need to use or construct bijectors for each of the transformations fi, i=1,…,5 and use tfb.Chain and tfb.TransformedDistribution to construct the final transformed distribution.
• Ensure to implement the log_det_jacobian methods for any subclassed bijectors that you write.
• Display a scatter plot of samples from the base distribution.
• Display 4 scatter plot images of the transformed distribution from your random normalizing flow, using samples of θ and a. Fix the axes of these 4 plots to the range [−1,1].

The following code block shows how to implement the above steps:

Create the image dataset

• Let’s now use your random normalizing flow to generate an image dataset of contour plots from our random normalising flow network.
• First, let’s display a sample of 4 contour plot images from our normalizing flow network using 4 independently sampled sets of parameters, using the following get_densities function: this function calculates density values for a (batched) Distribution for use in a contour plot.
• The dataset should consist of at least 1000 images, stored in a numpy array of shape (N, 36, 36, 3). Each image in the dataset should correspond to a contour plot of a transformed distribution from a normalizing flow with an independently sampled set of parameters. It will take a few minutes to create the dataset.
• As well as the get_densities function, the following get_image_array_from_density_values function will help to generate the dataset. This function creates a numpy array for an image of the contour plot for a given set of density values Z. Feel free to choose your own options for the contour plots.
• Let’s display a sample of 20 images from your generated dataset in a figure.

Create tf.data.Dataset objects

• Let’s now split your dataset to create tf.data.Dataset objects for training and validation data.
• Using the map method, ;et’s normalize the pixel values so that they lie between 0 and 1.
• These Datasets will be used to train a variational autoencoder (VAE). Use the map method to return a tuple of input and output Tensors where the image is duplicated as both input and output.
• Randomly shuffle the training Dataset.
• Batch both datasets with a batch size of 20, setting drop_remainder=True.
• Print the element_spec property for one of the Dataset objects.

Build the encoder and decoder networks

• Let’s now create the encoder and decoder for the variational autoencoder algorithm.
• Let’s design these networks, subject to the following constraints:
  • The encoder and decoder networks should be built using the Sequential class.
  • The encoder and decoder networks should use probabilistic layers where necessary to represent distributions.
  • The prior distribution should be a zero-mean, isotropic Gaussian (identity covariance matrix).
  • The encoder network should add the KL divergence loss to the model.
• Print the model summary for the encoder and decoder networks.

Train the variational autoencoder

• Let’s now train the variational autoencoder. Build the VAE using the Model class and the encoder and decoder models. Print the model summary.
• Compile the VAE with the negative log likelihood loss and train with the fit method, using the training and validation Datasets.
• Plot the learning curves for loss vs epoch for both training and validation sets.

Use the encoder and decoder networks

• Let’s now put your encoder and decoder networks into practice!
• Randomly sample 1000 images from the dataset, and pass them through the encoder. Display the embeddings in a scatter plot (project to 2 dimensions if the latent space has dimension higher than two).
• Randomly sample 4 images from the dataset and for each image, display the original and reconstructed image from the VAE in a figure.
  • Use the mean of the output distribution to display the images.
• Randomly sample 6 latent variable realisations from the prior distribution, and display the images in a figure.
  • Again use the mean of the output distribution to display the images.

The following animation of latent space interpolation shows the decoder’s generations, depending on the latent space.

To be continued…