Machine learning is all about probability. To train a model, we typically tune its parameters to maximise the probability of the training dataset under the model. To do so, we have to assume some probability distribution as the output of our model. The two distributions most commonly used are Categorical for classification and Gaussian for regression. The latter case can be problematic, as the true probability density function (pdf) of real data is often far from Gaussian. If we use the Gaussian as likelihood for image-generation models, we end up with blurry reconstructions. We can circumvent this issue by adversarial training, which is an example of likelihood-free inference, but this approach has its own issues.
Gaussians are also used, and often prove too simple, as the pdf for latent variables in Variational Autoencoders (VAEs), which I describe in my previous post. Fortunately, we can often take a simple probability distribution, take a sample from it and then transform the sample. This is equivalent to change of variables in probability distributions and, if the transformation meets some mild conditions, can result in a very complex pdf of the transformed variable. Danilo Rezende formalised this in his paper on Normalizing Flows (NF), which I describe below. NFs are usually used to parametrise the approximate posterior in VAEs but can also be applied for the likelihood function.
Change of Variables in Probability Distributions
We can transform a probability distribution using an invertible mapping (i.e. bijection). Let be a random variable and an invertible smooth mapping. We can use to transform . The resulting random variable has the following probability distribution:
We can apply a series of mappings , , with and obtain a normalizing flow, first introduced in Variational Inference with Normalizing Flows,
This series of transformations can transform a simple probability distribution (e.g. Gaussian) into a complicated multi-modal one. To be of practical use, however, we can consider only transformations whose determinants of Jacobians are easy to compute. The original paper considered two simple family of transformations, named planar and radial flows.
with and and an element-wise non-linearity. Let . The determinant can be easily computed as
We can think of it as slicing the -space with straight lines (or hyperplanes), where each line contracts or expands the space around it, see figure 1.
with , and parameters and .
Similarly to planar flows, radial flows introduce spheres in the -space, which either contract or expand the space inside the sphere, see figure 1.
These simple flows are useful only for low dimensional spaces, since each transformation affects only a small volume in the original space. As the volume of the space grows exponentially with the number of dimensions , we need a lot of layers in a high-dimensional space.
Another way to understand the need for many layers is to look at the form of the mappings. Each mapping behaves as a hidden layer of a neural network with one hidden unit and a skip connection. Since a single hidden unit is not very expressive, we need a lot of transformations. Recently introduced Sylvester Normalising Flows overcome the single-hidden-unit issue of these simple flows; for more details please read the paper.
Simple flows are useful for sampling, e.g. as parametrisation of in VAEs, but it is very difficult to evaluate probability of a data point that was not sampled from it. This is because the functions in planar and radial flow are invertible only in some regions of the -space, and the functional form of their inverse is generally unknown. Please drop a comment if you have an idea how to fix that.
Enhancing expressivity of normalising flows is not easy, since we are constrained by functions, whose Jacobians are easy to compute. It turns out, though, that we can introduce dependencies between different dimensions of the latent variable, and still end up with a tractable Jacobian. Namely, if after a transformation, the dimension of the resulting variable depends only on dimensions of the input variable, then the Jacobian of this transformation is triangular. As we know, a determinant of a triangular matrix is equal to the product of the terms on the diagonal. More formally, let be the Jacobian of the mapping , then
I would like to draw your attention to three interesting flows that use the above observation, albeit in different ways, and arrive at mappings with very different properties.
R-NVPs are arguably the least expressive but the most generally applicable of the three. Let , element-wise multiplication and and two mappings (Note that is not the sigmoid function). R-NVPs are defined as:
It is an autoregressive transformation, although not as general as equation (7) allows. It copies the first dimensions, while shifting and scaling all the remaining ones. The first part of the Jacobian (up to dimension ) is just an identity matrix, while the second part is lower-triangular with on the diagonal. Hence, the determinant of the Jacobian is
R-NVPs are particularly attractive, because both sampling and evaluating probability of some external sample are very efficient. Computational complexity of both operations is, in fact, exactly the same. This allows to use R-NVPs as a parametrisation of an approximate posterior in VAEs, but also as the output likelihood (in VAEs or general regression models). To see this, first note that we can compute all elements of and in parallel, since all inputs () are available. We can therefore compute in a single forward pass. Next, note that the inverse transformation has the following form, with all divisions done element-wise,
Note that and are usually implemented as neural networks, which are generally not invertible. Thanks to equation (11), however, they do not have to be invertible for the whole R-NVP transformation to be invertible. The original paper applies several layers of this mapping. The authors also reverse the ordering of dimensions after every step. This way, variables that are just copied in one step, are transformed in the following step.
We can be even more expressive than R-NVPs, but we pay a price. Here’s why.
Now, let and . We can introduce complex dependencies between dimensions of the random variable by specifying it in the following way.
Since each dimension depends only on the previous dimensions, the Jacobian of this transformation is a lower-triangular matrix with on the diagonal; the determinant is just a product of the terms on the diagonal. We might be able to sample in parallel (if different dimensions are i.i.d.), but the transformation is inherently sequential. We need to compute all before computing , which can be time consuming, and is therefore expensive to use as a parametrisation for the approximate posterior in VAEs.
This is an invertible transformation, and the inverse has the following form.
Given vectors and , we can vectorise the inverse transformation, similar to equation (11), as
The Jacobian is again lower-triangular, with on the diagonal and we can compute probability in a single pass.
The difference between the forward and the inverse transofrmations is that in the forward transformation, statistics used to transform every dimension depend on all the previously transformed dimensions. In the inverse transformation, the statistics used to invert (which is the input), depend only on that input, and not on any result of the inversion.
MAF directly uses equations (12) and (13) to transform as random variable. Since this transformation is inherently sequential, MAF is terribly slow when it comes to sampling. To evaluate the probability of a sample, however, we need the inverse mapping. MAF, which was designed for density estimation, can do that efficiently by using equation (15).
In principle, we could use it to parametrise the likelihood function (a.k.a. the decoder) in VAEs. Training would be fast, but, if the data dimensionality is high (e.g. images), generating new data would take very long. For a colour image of size , we would need to perform sequential iterations of equation (13). This cannot be parallelised, and hence, we abandon the all powerful GPUs we otherwise use.
We could also use MAF as a prior in VAEs. Training requires only evaluation of a sample under the prior . The dimensionality of the latent variable is typically much smaller than that of the output; often below . Sampling can still be expensive, but at least doable.
What other applications would you use MAF in? Please write a comment if anything comes to mind.
IAF defines a pdf by using a reparametrised version of equations (14) and (15), which we derive later. In this case, the transformed variable is defined as an inverse autoregressive mapping of the following form.
Since all and depend only on but not on , they can be all computed in parallel, in a single forward pass.
To understand how IAF affects the pdf of , we can compute the resulting probability density function. Other types of flows admit similar derivations. Here, we assume that follows a unit Gaussian,
The final pdf can be comprised of IAFs. To take this into account, we now set and ; i.e. is the result of transforming . To factor in subsequent transformations, we need to compute all the Jacobians:
If and are implemented as autoregressive transformations (with respect to ), then the first two terms in the Jacobian above are lower triangular matrices with zeros on the diagonal. The last term is a diagonal matrix, with on the diagonal. Thus, the determinant of the Jacobian is just
Therefore, the final log-probability can be written as
Sampling from an IAF is easy, since we just sample and then forward-transform it into . Each of the transformations gives us the vector , so that we can readily evaluate the probability of the sample .
To evaluate the density of a sample not taken from , we need to compute the chain of inverse transformations , . To do so, we have to sequentially compute
This can be expensive, but as long as and are implemented as autoregressive transformations, it is possible.
MAF vs IAF
Both MAF and IAF use autoregressive transformations, but in a different way. To see that IAF really is the inverse of MAF and that the equation (16) is in fact a reparametrised version of equation (14), set , , and .
This reparametrisation is useful, because it avoids divisions, which can be numerically unstable. To allow the vectorised form of equations (15) and (17), and have to be implemented as autoregressive functions; and one efficient way to do so is to use MADE-type neural networks (nicely explained in this blog post by Ferenc). In fact, both original papers use MADE as a building block.
To understand the trade-offs between MAF and IAF, it is instructive to study equations (15) and (17) in detail. You will notice, that although the equations look very similar, the position of inputs and outputs is swapped. This is why for IAF, sampling is efficient but density estimation is not, while for MAF, sampling is inefficient while density estimation is very fast.
Parallel WaveNet introduced the notion of Distribution Distillation, which combines the advantages of both types of flows. It trains one model, which closely resembles MAF, for density estimation. Its role is just to evaluate probability of a data point, given that data point. Once this model is trained, the authors instantiate a second model parametrised by IAF. Now, we can draw samples from IAF and evaluate their probability under the MAF. This allows us to compute Monte-Carlo approximation of the KL-divergence between the two probability distributions, which we can use as a training objective for IAF. This way, MAF acts as a teacher and IAF as a student. This clever application of both types of flows allowed to improve efficiency of the original WaveNet by the factor of 300.
- Two-part practical tutorial on normalising flows by Eric Jang
- MAF paper explores theoretical links between R-NVP, MAF and IAF in great detail,
- Parallel WaveNet combines MAF and IAF in a very clever trick the authors call Distribution Distillation,
- Continuous-Time Flows, as an example of even more expressive transformation.
I would like to thank Adam Goliński for fruitful discussions as well as his detailed feedback and numerous remarks on how to improve this post.