Search⌘ K
AI Features
Home
Courses
Introduction to JAX and Deep Learning

Divergence Measures

Discover how to measure differences between probability distributions using divergence measures such as entropy, relative entropy, Kullback-Leibler divergence, Jensen-Shannon divergence, total variation distance, and Wasserstein distance. Learn their definitions, applications, and significance in machine learning contexts, including GANs and NLP.

We'll cover the following...

Here we’ll focus on a quite important aspect of statistical learning. This lesson is advanced and can be reasonably skipped if needed.

Introduction

We can easily compare a couple of scalar values by their difference or a ratio. Similarly, we can compare the two vectors by taking the L1 or L2 norm.

To extend this notion of divergence between a couple of distributions requires some better measures, though. There are several real-world applications where we need to find the similarity (or difference) between two distributions. For example, text comparison between two sequences in bioinformatics, text comparison in Natural Language Processing (NLP), comparison of generated images by Generative Adversarial Networks (GANs), and so on.

Entropy

Let’s begin with the fundamental measure. The entropy of an independent vector is defined as:

H(X)=−∑i=1nP(Xi)logP(Xi)H(X) = -\sum_{i=1}^n P(X_i) logP(X_i)

Usually, the base of the log is taken as either 22 or ee.

Relative entropy

The relative entropy between two vectors XX and YY is defined as:

D(X,Y)=−∑i=1nXilog(XiYi)D(X,Y) = -\sum_{i=1}^n X_i log(\frac{X_i}{Y_i})

Since the equation involves (element-wise) ratio as well as logarithm, we must make sure to check for the zeros.

Python 3.8
key = random.PRNGKey(0)
A = random.uniform(key,(1,100))
B = random.uniform(key,(1,100))
def RelativeEntropy(A,B):
    return jnp.sum(jnp.where(B != 0, A * jnp.log(A / B), 0))
print("Divergence between A and B is:", RelativeEntropy(A,B)) #Answer will be 0
key, subkey = random.split(key)
C = random.uniform(subkey,(1,100))
#Now C will be different from A or B.
print("Divergence between A and C is:",RelativeEntropy(A,C))

It’s a good use of splitting in the above example to sample diverse vectors from the distributions.

Kullback-Leibler (KL) divergence

Relative entropy is generalized as Kullback-Leibler divergence, denoted as DKL(X∣∣Y)D_{KL}(X||Y) or simply D(X∣∣Y)D(X||Y), defined as:

DKL(X∣∣Y)=∑i=1n(Xilog(XiYi)−Xi+Yi)D_{KL}(X ||Y) = \sum_{i=1}^n (X_i log(\frac{X_i}{Y_i})-X_i+Y_i )

Kullback-Leibler divergence finds many applications and is a frequently used term in machine learning publications.

Note: Some authors refer to relative entropy itself as KL divergence, but we’ll follow a more precise definition from Convex Optimization, by Stephen Boyd and Lieven Vandenberghe, Cambridge University Press, 2004.

Python 3.8
A = random.uniform(random.PRNGKey(0),(1,100))
B = random.uniform(random.PRNGKey(38),(1,100))
def KLDivergence(A,B):
    return jnp.sum(jnp.where(B != 0, (A * jnp.log(A / B))-A+B, 0))
print("KL Divergence between A and B is:", KLDivergence(A,B))
print("KL Divergence between B and A is:", KLDivergence(B,A))

Jensen-Shannon divergence

As we saw, KL divergence is asymmetric:

DKL(X∣∣Y)≠DKL(Y∣∣X)D_{KL}(X ||Y) \ne D_{KL}(Y ||X)

Thus it cannot be used as a distance measure. Jensen-Shannon divergence neatly overcomes this issue:

DJS(X∣∣Y)=12DKL(X∣∣Y)+12DKL(Y∣∣X)D_{JS}(X ||Y) = \frac {1}{2} D_{KL}(X||Y) + \frac {1} {2} D_{KL}(Y||X)

JS divergence is thus much more applicable due to its additional mathematical properties.

Python 3.8
A = random.uniform(random.PRNGKey(0),(1,100))
B = random.uniform(random.PRNGKey(57),(1,100))
def KLDivergence(A,B):
    return jnp.sum(jnp.where(B != 0, (A * jnp.log(A / B))-A+B, 0))
def JSDivergence(A,B):
    return (KLDivergence(A,B)+KLDivergence(B,A))/2
print("JSD between A and B is:", JSDivergence(A,B)) 
print("JSD between B and A is:", JSDivergence(B,A))

Total variation distance

Total variation (or TV) distance is another common measure. It is defined by the maximum possible difference between the probability distributions for a given event, AA. This is formally defined as:

δ(X,Y)=sup⁡A∈F∣X(A)−Y(A)∣\delta(X,Y)=\sup_{ A \in \mathcal{F}} |X(A) - Y(A)|

Note: A supremum operator or sup is nothing but a generalization of the maximum operator to infinite sets. We use it here for the sake of completeness.

Total variation norm is often used for regularization in computer vision applications.

Python 3.8
A = random.uniform(random.PRNGKey(0),(1,100))
B = random.uniform(random.PRNGKey(57),(1,100))
def TVDistance(A,B):
    return jnp.max(jnp.abs(A-B))
print("TV Distance between A and B:", TVDistance(A,B)) 
print("TV Distance between B and A:", TVDistance(B,A))

Wasserstein distance

In 2017, Martin Arjovsky, et al. set the machine learning community ablaze by introducing yet another type of GAN: Wasserstein GANs.

Wasserstein GAN has since then become a state-of-the-art model. At the heart of its architecture lies a new divergence measure called Wasserstein distance.

Understanding Wasserstein distance and its implementation require some background, which is well complemented by a dedicated library in JAX. We will cover it in the following lesson.

Remember: It’s quite okay to use the WGANs in applied areas like computer vision without much background, so you can skip much of the theory in the next lesson and focus purely on the applied part if needed.