Nice Jacobians for normalizing flows
We cover some nice Jacobians used in various implementations of normalizing flows. Here "nice" means the determinant of these Jacobians is easy to calculate. The key paradigm of normalizing flows is learning a transformation of input data such that the distribution of is a known probability distribution (usually Gaussian or uniform).
We cover three main transformations the Jacobians of which are simply the product of the diagonal (possibly block) elements of the Jacobian.
1. Autoregressive transformations
This is arguably the simplest and most likely the earliest transformation used in normalizing flows. The Jacobian of an autoregressive transformation is a lower triangular matrix since .
Jacobian of an autoregressive transformation is a lower triangular matrix.
The determinant of this Jacobian is the product of the diagonal elements, which is easy to calculate.
2. Affine transformations
The key idea here is to partition the vector into two parts of equal length. These two parts are transformed to in the following way:
This leads to a Jacobian that is even sparser than a lower triangular matrix. The determinant of this Jacobian is the product of the scaling parameters (which is also the product of the diagonal elements).
Determinant of Jacobian is just the product of scaling parameters .
This transformation was used in the paper Density estimation using Real NVP. In order to make sure that the entire input space was covered by the transformation, the sets were switched in the next step (i.e was used to calculate the affine parameters to scale and shift ).
3. Channel-wise transformations
In this transformation, the input vector is partitioned into channels of equal length and each channel is linearly transformed by multiplying with a matrix .
Each channel is linearly transformed by a matrix .
The Jacobian of this transformation is a block diagonal matrix where each block is the matrix . In the example below, is partitioned into four channels.
Determinant of Jacobian is .
This transformation was used in the paper Glow: Generative Flow with Invertible 1x1 Convolutions. In that paper, the linear transformation was applied to each location of the input tensor of shape . In this case, the number of "channels" is .