Discover the power of JAX in deep learning. Gain insights into its ecosystem and learn about linear algebra, pseudo-random number generation, and optimization algorithms for cleaner, structured coding.

Doxx.tar.gz

NormalCode

JuPyter

JAX is a Python library designed for high-performance ML research. It is a powerful numerical computing library, just like Numpy, but with some key improvements.

In this course, you will learn all about JAX and its ecosystem of libraries (Haiku, Jraph, Chex, Flax, Optax). Addressing a wide range of audiences, you will cover several topics including linear algebra, random variables theory, pseudo-random number generation, and optimization algorithms.

By the end of this course, you will have a new set of skills that will make deep learning programming more intuitive, structured, and clean.

Introduction to JAX and Deep Learning

We talked about Wasserstein distance in the [last lesson](https://www.educative.io/courses/intro-jax-deep-learning/divergence-measures). It's quite important to get a background of the underlying theory of **optimal transport (OT)** first.

>**Note:** This lesson is a basic introduction to optimal transport (OT). It can be skipped if needed.

# Optimal transport

In 1781, French mathematician Gaspard Monge presented the following problem:

>_“A worker with a shovel in hand has to move a large pile of sand lying on a construction site. The worker's goal is to erect a target pile with a prescribed shape (for example, a giant sandcastle). Naturally, the worker wishes to minimize her total effort, quantified, for instance, as the total distance or time spent carrying shovelfuls of sand.” <br/>[G. Monge, Mémoire sur la théorie des déblais et des remblais (De l’Imprimerie Royale, 1781)]

So the problem is quite clear: 

* We have a distribution $\mathcal{X}$ (the raw pile of sand in this case).
* We want to **transport** it to another distribution $\mathcal{Y}$.
* We want to reduce the cost (be it _time_ or _cost_) of the transportation, $T:\mathcal{X} \to \mathcal{Y}$.

This problem applies in almost any area of daily life whenever we have to move a distribution rather than a single item. Common areas include computer vision, bioinformatics (especially in **single-cell sequencing**), and machine learning.





Actually, Monge's original motivation was to design forts for Napoleon. Monge was a strong supporter of the revolution, and Napoleon used this problem to optimize the development of forts during his Egypt campaign.

# Wasserstein distance

Having seen the optimal transport problem, we are now in a better position to understand the Wasserstein distance.

In 1939, Russian mathematician **Leonid Kantorovich** defined this metric to address the OT problem. However, his work remained obscure and was only rediscovered 30 years later by another Russian mathematician, **Leonid Vaseršteĭn** and hence named **Wasserstein** $-$ which is a bit unfair to Kantorovich. So some scholars encourage the use of his name).

>**Note:** Wasserstein distance is often referred to as **Earth-mover distance** in computer science literature because of the original Monge formulation.

Like norms, Wasserstein distance can be defined over a range of natural numbers. Generally,

$$ W_{p} (\mu, \nu) = ({inf E[(d(\mathcal X,\mathcal Y)^p]})^{1/p} $$

>**Remember:** Supremum is the maximum operator generalized to include infinite sets as well. Infimum is the counterpart for the _minimum_ operator.

The most commonly used Wasserstein distance is Wasserstein-1, which is used in WGANs as well. Its equation will reduce to a much simpler form:


$$ \mathcal W_{1} (\mu, \nu) = {\mathcal{inf}  \mathrm{E}[d(\mathcal X,\mathcal Y]} $$

The choice of distance function is up to us. If the distance function is L1 norm, $\mathcal W_1$ becomes:

$$ \mathcal W_{1} (\mu, \nu) = {inf \mathrm{E}[|\mathcal X-\mathcal Y|]} $$

Wasserstein distance has an advantage over other divergence measures in terms of optimization. Under mild conditions,  Wasserstein distance is:

* Smooth (Lipschitz continuous)
* Continuous everywhere
* Differentiable almost everywhere
* The best cost function on low dimensional data

Covering the mathematical details and explanations of these new topics would make this an unnecessarily lengthy read. We encourage the interested readers to please refer to the [original Wasserstein GAN paper (2017)](http://proceedings.mlr.press/v70/arjovsky17a/arjovsky17a.pdf) for the detailed theorems and their proofs.






This lesson will introduce Optimal Transport and the OTT library.

Introduction

JAX Programming Model

Linear Algebra

Random Variables and Distributions

JAX Ecosystem

Project: GAN Using the JAX ecosystem

Appendix

Introduction to Optimal Transport and OTT

Optimal transport

Wasserstein distance