For convenience, we denote the set of all joint distribution matrices with marginals \(\mu, \nu\) as

\(U(\mu, \nu) = \{P(x, y) \in M(d, d)^+, \mu = P \cdot 1_d, \nu = P^T 1_d\}\).

In fact this space is a convex polytope and we call it the Transport Polytope.

Information Theory Revisited

You might also recall our discussion of Shannon entropy and KL divergence from the last post. I might have given the impression that these concepts have been displaced by the new exciting world of optimal transport, but as it turns out, they are very important for computing the Sinkhorn metric!

As we can readily gather from the problem setting, there can potentially be a lot of joint distributions with marginals \(\mu, \nu\), even in the discrete case. However, we can get a constraint on exactly which joint distributions are possible through computing its Shannon entropy.

Theorem¹: \(H(P) \le H(\mu) + H(\nu)\), and the inequality is sharp.

In other words, the Shannon entropy of the joint distribution is always bounded by the sum of the Shannon entropies of the marginal distributions. Moreover, this bound is sharp since we can always construct the independent joint distribution matrix \(\mu\nu^T\).

With this constraint in mind, we define a subset of the transport polytope:

Intuitively, this subset contains all joint distributions whose Shannon entropy is less than \(\alpha\) away from the maximum possible entropy of \(H(\mu) + H(\nu)\) (which is attained by the joint distribution \(\mu\nu^T\)). I.e. all joint distributions which have sufficient entropy with respect to \(H(\mu), H(\nu)\).

Sinkhorn Metric

Now we are ready to define the Sinkhorn Metric:

We see that this distance is similar to the Wasserstein-1 distance except we are only looking for the minimum distance among joint distributions with sufficient entropy with respect to \(H(\mu), H(\nu)\).

Why only consider this restricted subset of joint distributions? The authors of ² explain that in solutions to the optimal transport problem, the optimum joint distributions are often on the corners of the transport polytope, and thus are statistical outliers. By setting this entropic restriction, the space of joint distributions is smoothened, which allows for easier computation.

Special cases

In one extreme, let us taken \(\alpha\) to infinity. Then as \(\alpha\) grows larger, the set of allowable entropies gets larger, so \(U_{\alpha}(\mu, \nu)\) gets larger until it converges to \(U(\mu, \nu)\) proper. So we can see that \(U_{\alpha}(\mu, \nu)\) is a good approximation of the transport polytope with large enough \(\alpha\).

In the other extreme, if we set \(\alpha = 0\), we can solve the optimal transport problem in closed form given any distance matrix. In fact we have:

Theorem: \(d_{D, 0} = \mu^T D \nu\).

Beyond these special cases, the author proves the symmetry and triangle inequality properties. The proof is similar to C. Villani’s proof in ³.

Duality

In many optimization problems, solving a problem amounts to solving its dual problem. Perhaps the most popular example is with Lagrange multipliers. Given a (constrained) optimization problem

its dual problem turns the problem into an unconstrained one, i.e.

where \(\mathcal{L}(x, \lambda) = f_0(x) + \displaystyle\sum\limits_{i=1}^m \lambda_i f_i(x)\) is called the Lagrangian and the parameters \(\lambda_i\) are called Lagrange multipliers.

It may seem at first glance that the optimal transport problem \(\min\limits_{P(x, y) \in U(\mu, \nu)} \langle D(x, y), P(x, y)\rangle_F\) lacks constraints, but by considering entropy, we do have a hard constraint, namely that \(H(P) \le H(\mu) + H(\nu)\)! Therefore we can transform the optimal transport problem into the dual problem

(the author uses a different convention of putting the Lagrange multiplier in the denominator, this does not change the optimization problem)

Notice that by introducing the sufficient entropy bound \(\alpha\), we get a dual problem for each value of \(\alpha\) and indeed a corresponding Lagrange multipler for each. Therefore we can also dualize the Sinkhorn distance \(\min\limits_{P(x, y) \in U_{\alpha}(\mu, \nu)} \langle D(x, y), P(x, y)\rangle_F\).

It turns out that computing the dual problem is computationally simpler.

Algorithm

The algorithm starts with computing the matrix \(K = \exp(-\lambda D)\), which is the elementwise exponential of the distance matrix \(D\). Then it uses the Sinkhorn fixed point iteration algorithm⁴ to compute the dual joint distribution matrix \(P_{\lambda}\). Finally, after some convergence threshold, the final iteration of the matrix is used to compute the Sinkhorn distance by taking a Frobenius inner product with \(D\).

Below is pseudocode of the algorithm (all dots before operators indicate elementwise operations):

function dual_sinkhorn(D, lambda, mu, nu){
    K = .exp(-\lambda D) #Elementwise exponentation of the matrix
    u = ones(size(mu))
    K_tilde = diag(1. / r) * K  
    while u changes more than some threshold:
        u = 1. / K_tilde * (nu ./ (K_tilde.transpose() * u)) #Sinkhorn fixed-point iteration
    v = (nu ./ (K_tilde.transpose() * u)
    d = sum(u .* ((K .* D)v)) #Frobenius inner product of K and D
}

The author does in fact extend the algorithm to simultaneously calculate sinkhorn distances for \(N\) target distributions simultaneously. But this is just a matter of concatenating vectors and will be omitted in this post.

Computation Package

Now that we’ve explained and summarized the approach. You can implement this yourself in your favorite language/computational software. However I have found a very good python package that already has Sinkhorn metrics built in. The package is Geomloss written by machine learning researcher Jean Feydy. This project is open source and available on Github. It is also compatible with Pytorch tensors. I highly recommend it if you need to compute Wasserstein distances.

Some example code you can use with this package (after installation):

from geomloss import SamplesLoss

torch.random.manual_seed(0)
X = torch.rand((3,100))
Y = torch.rand((3,100))

Loss =  SamplesLoss("sinkhorn", blur=0.05)
print(Loss( X, Y ).item())

Comparing Datasets

In many applications of machine learning, one might find it useful to compare two different datasets. For example, one might want to compare the distributions of training data to testing data. Or one might want to compare different training datasets for the purposes of transfer learning. Whatever the case, it is important to have a metric to determine how “different” (in terms of distribution) the datasets are to each other, assuming we don’t know their distributions beforehand.

An Information-theoretic Perspective

Traditionally the answer to the problem lay in information theory. Given two random variables \(X\) and \(Y\), we may treat \(X\) as data that’s already observed (i.e. training data) and \(Y\) as a model that aims to approximate \(X\) (i.e. testing data).

Shannon Entropy

From an information-theoretic perspective, a (discrete) random variable \(X\) can be thought of containing information, depending on what its probability distribution is like. This information can be encoded in bits of 0s and 1s.

Now we want to ask, is there a quantifiable difference between a fair coin (50/50 heads or tails) and a rigged coin? It turns out that the amount of information needed to encode the two types of coins is very different.

Suppose for the moment that you know the coin is fair. And someone flips the coin 10 times without your knowledge. You want to know how the flips landed. Out of these 2 flips we have a total of \(2^{10}\) possible outcomes of ordered sequences of heads or tails. Suppose the 2 flips were done without your knowledge, and you wanted to know exactly which outcomes occurred on the tosses. You must ask the tosser a series of yes/no questions in order to find out which coins were heads and which coins were tails. How many questions would you need to ask on average, in order to figure out the configurations of the coin tosses?

The first thing that may come into your mind is to ask sequentially:

• Is the first coin heads?
• Is the second coin heads?

Then we would need to ask 2 questions for any particular experiment, in order to know exactly how the 2 coin flips landed. In terms of information, we can encode this as two bits, the leftmost bit encodes the answer to the first question and the rightmost bit encodes the answer to the second question. E.g. the string 10 encodes the answers yes to the first question, and no to the second question. Since this uniquely determines the configuration, the string 10 corresponds to heads-tails.

Since the coin is fair, it means that every one of the \(2^2\) configurations is equally likely, since we needed to ask 2 questions for each, the average number of yes/no questions requires is 2. The average number of questions required to get a single coin flip is 1. That’s how much information we needed to figure out the outcomes of the coin tosses. So for a fair coin, we may say that it requires 1 bit of information to encode it’s distribution.

Now what if you knew that the coin was unfair? Let’s take a very extreme example, the coin is magnetized on one side and there’s a hidden magnet underneath the table so that it always comes up heads, and you know this. How many questions (on average) would we need to ask in this case?

The answer is 0. Because you already know that the coin will come up heads every time. So no questions required. Therefore a coin which always comes up heads requires 0 information to encode.

Now what about something in between? Suppose that \(P(Heads) = \frac{3}{4}\) and \(P(Tails) = \frac{1}{4}\).

It may still seem like we need \(n\) questions for \(n\) tosses to figure out the exact configuration. But it turns out that in this biased case, we can do better.

Consider the following series of questions we ask for a 2-toss experiment:

1. Is there at least one tails? -> If answer is no, then we must have head-head.
2. If 1. is yes, is the first coin tails? -> If answer is no, then we have head-tail (since we know there is at least one tail and it’s not the first coin).
3. If 2. is yes, is the second coin heads? -> If answer is yes, then we have tail-head. If answer is no, then we must have tail-tail.

In terms of encoding, we encode each question as a bit once again. So for example, the string 110 represents: yes to question 1 and 2, and no to question 3. So 110 corresponds to the configuration tail-tail, for example.

So as we can see, here we use a strategy that is not uniform, in the sense that some configurations will lead us to only ask one question (head-head), whereas some others require 3 questions (tail-tail). Since we are interested in the average number of questions, this allows us to surpass the naive strategy of asking 2 sequential questions (which is 2 on average as well).

Why is this? Let’s look at the probabilities of each configuration and the number of corresponding questions:

• Head-head: probability = \((\frac{3}{4})^2 = \frac{9}{16}\), questions = 1
• Head-tail: probability = \(\frac{3}{4} \cdot \frac{1}{4} = \frac{3}{16}\), questions = 2
• Tail-head: probability = \(\frac{1}{4} \cdot \frac{3}{4} = \frac{3}{16}\), questions = 3
• Tail-tail: probability = \(\frac{1}{4} \cdot \frac{1}{4} = \frac{1}{16}\), questions = 3

Then the expected (or average) number of questions is given by

and since this gives us 2 tosses, the average number of questions per toss is 0.8435. Lower than the 1 bit (or question) required for a fair coin. In terms of information, we needed an average of 0.8435 bits to encode this probability distribution.

Now it’s easy to see why this strategy won’t work for a fair coin, since each configuration has equal probability, we would be taking the average of \(\{1, 2, 3, 3\}\) which gives us a number greater than 2. So for the fair coin, asking the value of each coin sequentially is actually the best we can do. This gives us a meaningful and quantitative way of distinguishing between coins of different biases. I.e. the fair coin needs more information to “figure out” or encode, and the more biased a coin is, the less information is needed to “figure it out”. So a fair coin contains more information content.

In the case of the biased coin, how can we be sure that we’ve asked the best set of questions (i.e. most efficient encoding)? Formally, Claude Shannon has proven an achievable lower bound for the amount of information needed. Given a discrete random variable \(X\), its Shannon entropy is given by:

The base in the logarithm indicates the unit of information we’re dealing with. Usually base 2 is chosen because we quantify information in terms of bits. A quick calculation shows that the Shannon entropy of a fair coin is 1, however the Shannon entropy of the biased coin we looked at from earlier is in fact around 0.81. Which means that in our example, we didn’t use the best possible encoding.

Kullback-Leibler Divergence

The Kullback-Leibler (KL) Divergence is a relative measure of how much information is lost when trying to approximate the known information \(P\) with the new model \(Q\) (defined on the same sample space). More precisely, it measures the number of extra bits required to code samples from \(P\) based on the new code \(Q\). Given densities \(p(x), q(x)\) for \(P, Q\) respectively. Its formula is as follows

Intuitively, the log factor measures the difference in information between the two distributions, and therefore the sum becomes the expected difference in information between \(P\) and \(Q\).

To take the example of our two coins from before. Let \(P\) denote the random variable of the fair coin and \(Q\) denote the random variable of the coin with \(P(Heads) = \frac{3}{4}\) and \(P(Tails) = \frac{1}{4}\). The KL divergence from \(Q\) to \(P\) is given by

We interpret this as saying – if we have an optimized set of questions which tells us configurations of the biased coin. We would need an extra 0.143 questions (on average) per coin in order to figure out the configurations of the fair coin. This makes sense since the fair coin holds more information, as we’ve already seen.

Caveats

Now this seems like a perfect way to measure differences between distributions. For example, one might have a training set and wonder how much extra information is needed to classify the testing set, knowing an optimal function to classify the training set. However, there’s one huge caveat: KL divergence is not a distance function!

To see this, note that for any reasonable definition of distance, going from a point A to a point B should yield the same distance as going back from B to A. However KL divergence doesn’t have this property! Lets take the two coins in the previous section, now we want to compute the KL divergence from \(P\) to \(Q\):

So given our sequential strategy of asking one coin at a time for the fair coin, if we applied the same strategy to the biased coin, it would require on average additional questions! This happens even though the biased coin contains less information, because of course, the optimal encoding for one distribution might not be the optimal encoding for the other, and vice-versa.

Another property that the KL divergence does not satisfy is the triangle inequality, formally it takes the form

for points \(x, y, z\) in the same space and a distance function \(d\). Intuitively, it says that going from a point A to another point B is always at least as short as going to an intermediate point in between.

To show an example of this failure, consider again the fair and biased coins \(P, Q\) from before, with an additional biased coin \(R\) with \(P(Heads) = \frac{9}{10}\) and \(P(Tails) = \frac{1}{10}\). We leave you to show that \(KL(P||Q) > KL(P||R) + KL(R||Q)\).

This somehow implies that switching directly to a new dataset requires more information (e.g. extra training) than going to some third dataset, and then switching to the new dataset, which seems paradoxical.

Costs of Moving

Sometimes people will use the Jansen-Shannon divergence as a symmetrized version of the KL divergence. Unfortunately this only takes care of the symmetry property but not the triangle inequality property (in general).

More recently, machine learning researchers have started to think about this problem of comparing distributions using a different perspective: rather than think about some notion of intrinsic difference between two distributions, instead think about how much it would cost (this could be computing power, physical energy, monetary cost, etc.) to move the data from one distribution to the other.

In a way this is even more useful than KL divergence. Even though KL divergence tells us how much information it takes to move from one distribution to another, it’s agnostic on how much each bit of information costs. In real-world applications, we care very much about costs, and some forms of information may be cheaper than others. As we will see shortly, this interpretation involving costs of moving also gives us a bona fide distance function to work with.

Wasserstein Distance

The original motivation behind Wasserstein distance, and indeed the field of optimal transport, comes from transporting goods and resources over distances. For example, a copper mine is located some distance away from the factories, and manufacturers often needed to calculate the most cost-efficient way to transport raw copper to the factories for processing. This is often a hard problem because there are many paths to choose from (different roads have different costs) and also because there are multiple mines and factories (which mine should transport to which factory).

By this analogy, given distributions \(P\) and \(Q\), we can think of samples in \(P\) as locations of the mines and samples in \(Q\) as the locations of the factories. Then what we want to do is to calculate the most cost-efficient way to move the copper (samples of \(P\)) to the factories (samples of \(Q\)). For now let’s declare the cost of moving from point \(x\) to point \(y\) the Euclidean distance between them: \(||x - y||\).

Now let us set up the definition of Wasserstein distance. Suppose we have two random variables \(X\) and \(Y\) jointly distributed on a probability space \(\mathcal{X}\), such that the marginal densities of \(X\) and \(Y\) are given by \(\mu, \nu\) respectively. Moreover suppose we have a distance function \(d(x, y)\).

Then we define the Wasserstein-p distance as:

where the infimum is taken over all possible joint distributions \(\pi\) of \(X, Y\) with marginals \(\mu\) and \(\nu\).

Above is the general definition for continuous random variables, but since we are mostly concerned with doing computations and optimizations, we are more concerned with the discrete case. In the discrete setting, the probability space is finite with size \(d = |\mathcal{X}|\). The distance function \(d(x, y)\) becomes a distance matrix \(D(x, y)\) of size \(d \times d\), and the joint distribution also becomes a matrix \(P(x, y)\) of the same size. Here the marginal distributions become the vectors \(\mu = P \cdot 1_d, \nu = P^T 1_d\), the row and column sums of the joint distribution matrix.

So the discrete Wasserstein-p distance is given by:

where \(\langle \cdot, \cdot \rangle_F\) denotes the Frobenius inner product for matrices, and \(D(x, y)^{(p)}\) denotes the elementwise \(p\)-th power of \(D(x, y)\).

It’s easy to see that under both definitions, the Wasserstein-p distance is symmetrical. The proof of triangle inequality is more involved, see ¹.

Conclusion

In this post we have summarized the long-standing topic of calculating differences between distributions. We have presented an information-theoretic perspective and defined KL divergence and Wasserstein(-p) distance. In the next post we will cover a practical way of computing these quantities.

The format will be as follows: I will post two Youtube videos of the music side by side and indicate the time snippets during which the similarities happen. For better clarification I will also include a picture of the motif or idea that’s common in both compositions alongside the videos.

Important:

• I am in no way implying any sort of plagiarism with these comparisons! That would defeat the purpose of this series. So I will not be comparing Sergei Rachmaninoff’s “Piano Concerto No. 2” to Eric Carmen’s “All by Myself”. The point is to bring attention to the absurd coincidences we sometimes encounter in music, and in life!
• Samples, tributes, and homages don’t count. It has to be a genuine coincidence! So I will not be comparing King Crimson’s “21st Century Schizoid Man” with Kanye West’s “Power”, nor will I be comparing Martika’s “Toy Soldiers” with Eminem’s “Like Toy Soldiers”.
• Standards and folk tunes don’t count, since music that use them are all referring to a common starting point–not a coincidence.

To kick off the series, I present a cross-genre coincidence:

Common Motif

Branford Marsalis Quartet - Dance of the Evil Toys (motif at 1:41, repeats throughout)

and…

Sergei Prokofiev - Piano Concerto No. 3 (motif at 15:50, repeats a few times in movement 2)

As a bonus, here’s an intra-genre coincidence:

A basic reduction of the common passage, different keys and meters are used in the pair of examples below

Frederic Chopin - Etude Op. 10 No. 7 (passage at 0:23, only appears once)

and…

Johannes Brahms - Hungarian Dance No. 5 (passage at 0:48, appears multiple times)

I will continue the series as I find more examples (which can be very hard!). So updates on a regular interval is probably out of the question.

In this post, we will put these ideas into practice by training a simple neural network on the MNIST dataset using mixed point precision. The machine learning framework used will be Pytorch with the AMP (automatic mixed precision) package.

As the name suggests, the AMP package will automatically carry out the mixed precision algorithm (master copy, loss scaling, arithmetic precision) without the user having to do floating point operations manually.

Here is the basic code format taken from the Pytorch AMP examples page:

# Packages required for automatic mixed precision
from torch.cuda.amp import GradScaler, autocast

# Creates model and optimizer in default precision
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)

# Creates a GradScaler once at the beginning of training.
scaler = GradScaler()

for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()

        # The forward pass and loss function need to be encapsulated under autocast().
        with autocast():
            output = model(input)
            loss = loss_fn(output, target)

        # The scale() function from the GradScaler object is called instead of
        # the usual loss.backward()
        scaler.scale(loss).backward()

        # Scaling back the gradients is necessary if there are gradient operations
        # (such as clipping) to be done

        # torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=max)

        # scaler.step() first unscales the gradients of the optimizer's assigned params.
        # If these gradients do not contain infs or NaNs, optimizer.step() is then called,
        # otherwise, optimizer.step() is skipped.
        scaler.step(optimizer)

        # A new line that's needed: updates the scaler for next iteration.
        scaler.update()

notice that the main differences are the precision casting in the feedforward and loss function calculation, and the loss scaling.

Experiments

We test out mixed precision training on an NLP classification task, with a pretrained BERT (Bidirectional Encoder Representations from Transformers) (cased) model fine-tuned on IMDB Review Data from HuggingFace Datasets. The dataset consists of reviews of films or TV series by IMDB users, and each review is labelled with a postive (1) sentiment or negative (0) sentiment. Our aim is to train the BERT model to classify reviews as positive or negative.

The Colab notebook can be found here. Much of the code was taken from the excellent tutorial on Curiousily, as well as a HuggingFace pages.

Warning: Mixed precision only works with GPUs! So if you’re not using GPUs, go to Runtime > Change Runtime Type > Hardware Accelerator, and select GPU from the dropdown menu.

The reason for picking BERT lies both in its effectiveness and for its vast number of parameters (~110 million). Therefore we would expect that the techniques in mixed precision training would help us save a significant amount of training time.

We took a small sample of 1000 reviews from the dataset and split them into 800 training samples and 200 testing samples. We then trained 10 epochs with a batch size of 64 with the AdamW optimizer using both regular training and mixed precision training.

The results are summarized as follows:

Notice that with mixed precision, we had almost cut down half the training time! Moreover, the test accuracy was comparable between the two models.

Therefore it can be concluded that for large models such as BERT, mixed precision training is all but essential.

Smaller Models

One small caveat we should mention is that mixed precision is best done with bigger models. In another experiment, we have attempted mixed precision on the Resnet18 model, which has around 11 million parameters. The task was to train Resnet18 on the MNIST dataset under the two training methods. Mixed precision did not offer any advantage in this case.

(*) Technically speaking, the running time also included time spent on validation after each epoch, but validation of 200 samples takes so little time it can safely be ignored from the analysis.

Precision Formats

In machine learning and computing in general, more often than not that we must do computations outside the realm of integers. However, since most real numbers are irrational and our computers have finite memory, we can never represent a real number exactly on a computer.

Therefore computer engineers have developed hardware around the convention that we will store a real number up to a certain number of significant digits (in binary), and each significatn digit is stored as a bit in memory. We will soon touch on terms such as float32 and float16.

Representing a number in binary

Our familiar number system uses base 10, i.e. we use the convention of representing every number as a sum of powers of 10, i.e.

In binary the same number is represented as

Since now we use powers of 2 instead of powers of 10, the only possible digits are 0’s and 1’s. Each binary digit is also called a bit. Already we notice that representing a number in binary instead of decimal yields more significant digits.

Single precision

One convention we use for storing (an approximation) of real numbers on a computer is single precision or float32. This means that we will use 32 bits of memory to represent a number. However, this doesn’t just mean we shove 32 significant digits in there. Each bit serves a very specific purpose.

Image from Wikipedia ²

The first bit \(s\) is the sign, it stores whether the number is negative or positive. The next 8 bits store the exponent \(\epsilon\) in binary. The remaining 23 bits store the significand, call this \(b = (b_1, \cdots, b_{23})\). The number associated with these 32-bits is given by

\((-1)^s \cdot 2^{\epsilon - 127} (1 + b)\).

Let’s calculate the range of such a configuration. The maximum and minimum exponents are 255 and 0 respectively, so with the offset of -127 we have a max value of \(128\) and minimum value of \(-127\) for the exponent bits. However the maximum and minimum exponents are reserved for special numbers. So the actual max and min exponents we can use are \(127\) and \(-126\) respectively. For the significand, the maximum value is \(1 - 2^23\). With a plus sign in front, we get that the maximum possible value is

with the minimum possible value the negative of that number.

For precision, the smallest positive number is \(2^{-149} \approx 1.4 * 10^{-45}\) and the biggest negative number is the negative of that.

This range is good enough for almost all computational purposes.

As for the reserved special numbers, those have to do with overflow, underflow, and indeterminates. A very nice summary of floating point is given in these course notes.

Half Precision

Half precision or float16 uses 16 bits: 1 for the sign, 5 for the exponent, and 10 for the significand.

This format results in a smaller range and more coarse precision than both single and double precision. Positive numbers range from \(2^{−24} \approx 5.96 * 10^{-8}\) to \((2 - 2^{-10}) * 2^{15} = 65504\).

Double Precision

There is another format called Double precision, or float64. Here we use 64 bits: 1 for the sign, 11 for the exponent, and 52 for the significand.

Tradeoff

Here we can see a tradeoff between the amount of memory used and the precision of the computation we are doing. The more precise we make our calculations, the more memory we need to use and vice-versa. Moreover, since we are using less digits, we also speed up computation time. For machine learning, we would obviously like to achieve the dream of using low memory to achieve high precision.

At first this seems like a logistical problem: How can we allocate more memory? How can we reduce the number of parameters in a deep neural network? But it turns out that we can exploit a bag of tricks of floating-point arithmetic we achieve our dream scenario.

The idea (and the tricks)

To fully explain the idea, let us first think about what kind of arithmetic operations are done when we train a neural network. No matter what the type of network it is, be it fully-connected, convolutional, or recurrent, the main operations are one of the following: matrix multiplication, reduce operations (e.g. summing up all entries of a vector, mostly comes up when computing means), or activation functions.

In the old days we carried out all these operations in float32. Now that we realized float16 saves memory and speed, we would ideally like to carry out all these operations in float16.

However, a potential pitfall is the significant loss in precision. Consider the following float32 : \(2^{-16} = 0 01101111 00000000000000000000000 \approx 0.00001525 = 1.525 * 10^{-5}\). I have written the 32-bit representation on the right, but clearly since the exponent is -16, it cannot be represented as a float16 number. Which means that such numbers will be treated as 0 in float16. And of course as we know, quantities to the scale of \(10^{-5}\) or smaller is very common in machine learning (small learning rates, small gradients, etc.). In fact, experiments have shown that over half of activation gradients involved in training certain models are too small to be represented by float16 numbers, and therefore get rounded down to 0.

Coupled with the fact that each pass through the neural network involves potentially millions of arithmetic operations, a seemingly small error such as this can compound into bigger errors. Ultimately affecting model accuracy. So the key idea is to identify when exactly we can get away with using float16 instead of float32. Let’s use base-10 to give some examples.

• Addition: Does not affect precision, if we add two numbers in base-10 with 2 significant digits for example: \(1.43 + 55.09\), the result still has at most 2 significant digits. Likewise, adding 2 float16’s stays in float16.
• Multiplication: Of course when we use float16 to multiply, every language we use that supports float16 will give us the product back in float16. So it may seem like multiplication doesn’t change precision either, but there is secretly a casting operation done under the hood. Again let us consider a base-10 example: \(0.2 * 0.3 = 0.06\). The multiplicands both had 1 significant digit but the product now has 2. If we were hypothetically working in the world of float1’s, allowing only 1 significant digit, then we will have to get rid of the 6 digit somehow. Depending on the implementation it could be rounding-up or rounding-down. Either way some precision is lost here.

Trick 1: Master Copy

The mixed precision approach keeps a master copy of the weights in float32, with the float16 conversion being used in the actual model.

What this means is that in the beginning, before the first pass through the model, the weights are initialized in float16 and a float32 master copy is then copied from the initial weights are stored separately. Then before each forward and backward pass, the weights are converted from the master copy into float16 and used for the model (this has no effect for the first training iteration but has an impact later on). After carrying out all relevant computations in float 16 (with performance improvements outlined in the next 2 tricks), the weight update is done in float32 on the master copy.

In other words, we use float16 copies of the weights purely as a funnel for time-saving computations, whereas the actual weights we store for the model are in the float32 master copy. This ensures that we don’t lose precision.

But what about the precision lost when converting from the master copy into float16? Let’s do an example in base-10 to see what happens:

Consider doing gradient descent on the convex function \(f(w) = w^2\) with learning rate \(0.0001\), starting from the initialization \(w = 1\). Initially the weight \(w = 1\) is stored in “float2” and a master copy is made in “float6” (*). After calculating gradients, we get that we should subtract \(2 * 0.0001 = 0.0002\) from the weight for the update (note that the weight updates are in the higher precision). This leaves us with the new weight being \(0.9998\). So far so good.

For the next iteration, the weight \(0.9998\) gets cast down to \(0.99\) before gradient computations, then the gradient would be \(0.99\) and the weight update is given by \(- 0.000198\). After updating, the new weight becomes \(0.999602\). But in the cast iteration the weight gets cast down to \(0.99\) again! So it seems like we got stuck in a local minimum with the weights.

However, we did get a nonzero gradient in the previous iteration, namely \(- 0.000198\), and the weight is continually being updated in the master copy. Which means that after more iterations, the master copy will eventually become \(0.98...\), and the “float2” version of the weight will change to \(0.98\), then \(0.97\), and so on, until \(f(w)\) reaches the global minimum of \(w = 0\).

Trick 2: Loss Scaling

We briefly mentioned that most of the activation gradients during training are too small to be represented by float16. However, it is rare to see huge gradients and loss functions(to the order of \(> 2^{12}\)). This means that we can preserve all the information of a float32 gradient in a float64 format simply by shifting bits!.

The procedure is as follows:

1. Compute loss function in float32
2. Convert loss value to float16 by shifting bits up.
3. Backpropagate with the float16 loss value, notice that by the rules of differentiation, the (constant) scaling is preserved in this step.
4. Scale the bits back down into float32.
5. Do gradient clipping, weight decay, etc. (optional)

The final gradient will be in float32 after this procedure, but we have saved a lot of computation time during the computationally intensive backpropagation phase. Moreover, we don’t lose any precision by doing this!

How is the scaling factor determined? A very simple trial-and-error approach: we start with the largest possible factor \(2^{15}\) and check if the result overflows, if it does then we drop the scaling factor by a factor of 2, and so on until overflow no longer happens. The authors of the paper have found that a scaling factor of \(2^3\) (3 bits) is good enough for most purposes.

Trick 3: Arithmetic Precision

In the training of neural networks, we often need to multiply matrices. When multiplying two matrices \(A_{m \times n} \cdot B_{n \times k}\), the \((i, j)\)-th entry of their product is given by \(\displaystyle\sum\limits_{\ell=1}^n a_{i, \ell} b_{\ell, j}\). Notice that such a computation consists of a series of scalar products added together.

As we have seen earlier, float16 numbers naturally multiply to float32 numbers. The strategy in the paper is to store the scalars of the matrices in float16, then when they are multiplied, keep them in float32 until all the addition operations are done, and only then cast them back to float16.

To see the upshot of this approach, let us again consider an example in base-10. We want to compute the quantity \(1.23 * 0.45 + 2.12 * (-5.94)\). These scalars are stored in “float2” at present. The result ends up being \(-12.0393\), which has 4 significant digits (in the real world).

Let us first carry out the usual computation of our computers in “float2”, which is to multiply into a “float4” number and then cast down to “float2” immediately. Then we get \(0.5535 + (-12.5928)\), each of which gets casted down to \(0.55 + (-12.60) = -12.05\).

Now let us accumulate the sums in “float4” before we cast back down. By using the round-down convention, we get \(-12.04\) when we convert back to “float2”. This means that accumulating the sums before casting back down preserves higher precision. This effect becomes compounded when the matrices are very large, resulting in more addition operations.

Conclusion

The authors note that training models with mixed point precision has accuracy comparable to training in float32, with the added benefit of saving computation time for many models.

In the next post in this series, we will talk about how to implement mixed precision training using Pytorch.

(*) It should be noted that the base-10 examples are purely for demonstration, and that the terms “float2”, “float4”, and so forth are meaningless for the floating point arithmetic of most CPU architectures, which are in binary. And of course, float32 does not mean that the maximum precision is exactly 32 digits for each number, since numbers can go as low as \(2^{-149}\).