There was a interesting paper released some time in the past couple of weeks that showed it’s possible to do lossless compression via flow-based models. Historically flow based models didn’t get as much attention as VAE and definitely not as much as GAN’s so I paid very little mind to them. But flow based generative models exhibit very interesting properties and now the ability to do lossless compression warrants a deeper look into them. This post is rather large and therefore is split up into three sections. We’ll first introduce flow-based models, then we’ll discuss the various coupling layers we can define and lastly we’ll discuss how we can do lossless compression via neural networks.
The fundamental problem of generative modeling is learning a density over some natural distribution. Furthermore a necessary condition for our generative model is that we have a efficient sampling strategy .
In order to be able to sample from many generative models attempt to learn a function from a known prior distribution to the natural distribution .
For example VAE’s explicitly parameterize the output of the encoder to be standard normal via the reparameterization trick. GAN’s on the other hand directly take a sample from a prior distribution (standard normal most of the time) as an input to the generator.
But what if we could devise a function that maps from our natural distribution to our prior but also has the extra property that it is bijective (permits an inverse for every sample in it’s range)?
Let’s assume that we can find such a function what would could we do with it? Let’s start off by formalizing some of the elements we’ve introduced so far.
Let , furthermore let’s define the respective probability density functions: . Lastly let’s say that function:
We can couple the two probability distributions with the continuous change of variable formula.
And we can use the coupling to get
where is the Jacobian of at . Recall that our Jacobian will be a real matrix of size
What we’ve done is made it possible to directly compute the log probability of our natural distribution through a differentiable bijective function. Therefore we can use standard negative log-likelihood to train our generative model.
where represents our training dataset. Furthermore sampling is straightforward via
Being able to directly optimize towards the log probability of our training set is an incredibly powerful tool. But the assumptions we made and the derived equations are rather intimidating. First, our coupling function needs to be bijective and easily invertible. Furthermore, equation 4 requires calculating an expensive determinant which has an asymptotic runtime of .
The good news is, is that by carefully constructing our coupling function we can not only satisfy the necessary conditions but reduce the computational complexity of calculating the determinant. Let’s enumerate the problems that we have to face.
- Calculating the jacobian determinant for a non-constrained function
- Selecting a family of functions that is bijective.
- Inverse of our functions must be easily computable.
is non trivial.
Bijective Transforms with Tractable Jacobian Determinants
The first paper I read that tackled this problem is Non-linear Independent Components Estimation by Dinh et al. So we’ll follow along with their work as well as introduce some other forms of transforms that satisfy the conditions laid out at the end of the previous section.
The key fact that we’ll use is that the determinant of a triangular matrix is equal to the product of diagonal elements.
Meaning that if we can construct our function in such a way that it’s Jacobian is a triangular (either upper or lower) matrix then the computationally expensive determinant becomes a fairly cheap operation.
General Coupling Layer
Recall that . Let’s also say that . Let’s also say that represents taking the first elements of and represents taking all elements starting from element .
The Jacobian and it’s determinant will then be
where is the identity matrix of size . The inverse of this function will then be
What’s cool is that we have no limitations on function and therefore we can parameterize it with a standard neural network with arbitrary complexity.
Additive Coupling Layer
For example let’s define
So our non-identity function will be:
even cooler is that our inverse function is as easy to compute as our forward function
Furthermore our Jacobian will be lower triangular. This follows from
Therefore the determinant of this transform will always be 1, since our full jacobian is a lower triangular with only 1’s on the diagonal.
Affine Coupling Layer
The additive coupling layer has a unit jacobian (i.e. volume preserving). This limits the level of expressiveness of this coupling layer, even if we stack multiple transforms on top of eachother. The original paper introducing additive coupling layers proposed a solution to this problem based on rescaling. But let’s see if we can construct a transform that is more expressive.
Our Jacobian for this transform will be
As we can see from here the Jacobian is lower triangular, non volume preserving and we can use our shortcut to calculate the determinant inexpensively. This coupling layer was introduced by a follow up paper by Dinh et al titled Density Estimation using Real NVP.
Combining Coupling Layers
Although a single coupling layer can be powerful, it makes sense that we might want to stack coupling layers to increase their expressiveness even more. Also note that the general coupling form keeps a portion of the input unchanged. But before going deeper on how to deal with this issue let’s formalize exactly what it means to stack coupling layers.
We need to verify 3 things. Composition of function, determinant, inverse function.
We’ll leave the proof for these to Google, but these three rules provide us with a way to efficiently calculate Jacobian and inverse given a composition of our coupling layers. Now getting back to the problem at hand, which is that coupling layers in the general form provided do not augment a set of the original input. One trivial solution is to alternate across layers which portions of the input are passed through and which are changed. If we’re working with images we can imagine more complicated masking schemes.
Another layer introduced by the Real NVP paper was the Squeeze layer which reshapes the the image on the left into the image on the right. Directly from the paper: “The squeezing operation reduces the 4 × 4 × 1 tensor (on the left) into a 2 × 2 × 4 tensor (on the right). Before the squeezing operation, a checkerboard pattern is used for coupling layers while a channel-wise masking pattern is used afterward.“
Hoogeboom et al released a paper recently titled Integer Discrete Flows and Lossless Compression. In this paper they introduce a new type of coupling that works directly off of discrete integers and because couplings are invertible it’s possible to invert a discrete representation learned back into the original source. Now because we have theoretic guarantees on invertibility we can start to work toward learned lossless compression.
Recall that all types of flow-based generative models produce an output that has the same dimensions as the input, so naturally no trivial compression happens. But what Hoogeboom et al do is use a rANS entropy encoder to compress the discrete output given by the discrete flow model. Since both the discrete flow based model and rANS encoder are invertible we have an learnable end to end loss-less compression pipeline.
The papers contributions are that an easily invertible discrete coupling joined with a smart choice of prior can lead to a representation that can be effectively compressed using an entropy encoder. They showed serious compression gains over other lossless method.
Integer Discrete Coupling
The first step is understanding the design behind the discrete coupling. For this we’ll have to update our definitions of variables.
Let , the set of integers. Still using the general coupling formula we have:
where is rounding to the nearest integer the output from a neural network . One problem that immediately jumps out is that the rounding operator is not differentiable. One way to get around this is to use the Straight-Through Estimator, which for the rounding operator will make us just skip the rounding operator in the backpropagation step:
where is the identity function.
Recall the additive coupling layer we introduced a couple sections back, integer discrete coupling is almost identical in the forward pass, and exactly identical in the backward pass. Interestingly enough Hoogeboom et al argue that even though the transform has a unit determinant (volume preserving) it’s representative power when stacked is good enough. In general I haven’t seen a consensus on what coupling layers work best in what scenario and I haven’t read anything on theoretic guarantees of expressive power of coupling layers. This is a good future research direction.
Choice Of Prior
The authors choice for the prior distribution was a factored discretized logistic distribution.
where is the standard sigmoid function. Furthermore instead of having a static mean and scale the authors parametrize them using a neural network.
Notice our choice of prior so far is unimodal. To fix this issue the authors use a mixture of discretized logistic distributions.
The authors note that a mixture of about 5 distributions was enough.
The image to the left is an example of the flow architecture with 2 levels. The squeeze layer reduces the spatial dimension while increasing the number of channels. After that integer flow layers are applied. An integer flow layer consists of a channel permutation of the image and an integer discrete coupling layer. In total there are D integer flows after each squeeze operation.
One important thing to notice is that the portion of the input that is used to compute mean and scale is not the same portion of the layer that is being factored out.
“The IDF architecture is split up into one or more levels. Each level consists of a squeeze operation, D integer flow layers, and a factor-out layer. Hence, each level defines a mapping from to , except for the final level L, which defines a mapping . Each of the D integer flow layers per level consist of a permutation layer followed by an integer discrete coupling layer.
In the case of compression using an IDF, the mapping is defined by the IDF. Subsequently, is encoded under the distribution to a bitstream using an entropy encoder. Note that, when using factor-out layers, is also defined using the IDF. Finally, in order to decode a bitstream , an entropy encoder uses to obtain . and the original image is obtained by using the map , i.e., the inverse IDF.”
The result for lossless compression are substantially better than our current methods and show the power of learned lossless compression.
In conclusion we went in depth into flow-based generative models, various coupling layers that we can use and then lightly covered how we can learn to do lossless compression. The future of flow-based models is really interesting and the various benefits of flow-based models, specifically the tractable log probability, make it very attractive for a set of generative tasks.
If you’re new to flow-based generative modeling I would highly recommend reading the papers in order as they are linked in this blog. But hopefully reading this provided you with the basic knowledge for understanding flow-based models!