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 y=f(x),x,y∈Rny = f(x), x,y \in R^n such that the distribution of yy is a known probability distribution (usually Gaussian or uniform).

p(x)=p(y)×∣det⁡(∂y/∂x⏞Jacobian)∣⏟easy to calculatep(x) = p(y) \times \lvert \underbrace{\det(\overbrace{\partial y / \partial x}^{\text{Jacobian}}) \rvert}_{\textcolor{#b45309}{\text{easy to calculate}}}

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 ∂yi/∂xj=0 if i<j\textcolor{#b45309}{\partial y_i / \partial x_j = 0} \ \text{if} \ i \lt j.

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 x∈Rnx \in R^n into two parts xA,xB∈Rn/2x_A, x_B \in R^{n/2} of equal length. These two parts are transformed to yA,yB∈Rn/2y_A, y_B \in R^{n/2} in the following way:

yA=xAyB=γ(xA)⊙xB⏞elementwise scaling⊕ β(xA)⏟elementwise shifting\begin{align} y_A &= x_A \\ y_B &= \overbrace{\textcolor{#b45309}{\gamma(x_A) \odot x_B}}^{\textcolor{#b45309}{\text{elementwise scaling}}} \underbrace{\textcolor{#e11d48}{\oplus \ \beta(x_A)}}_{\textcolor{#b45309}{\text{elementwise shifting}}} \end{align}

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).

1 1 1 1 𝛾1 𝛾2 𝛾3 𝛾4

Determinant of Jacobian is just the product of scaling parameters ∏iγi\prod_i \gamma_i.

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 xA,xBx_A, x_B were switched in the next step (i.e xBx_B was used to calculate the affine parameters to scale and shift xAx_A).

3. Channel-wise transformations

In this transformation, the input vector x∈Rnx \in R^n is partitioned into cc channels of equal length Ci∈Rn/c,i∈{1,…,c}C_i \in R^{n/c}, i \in \{1, \dots, c\} and each channel is linearly transformed by multiplying with a matrix W∈R(n/c)×(n/c)W \in R^{(n/c) \times (n/c)}.

x1,x2⏟C1, x3,x4⏟C2, …, xn−1,xn⏟Cc↓ ×Wy1,y2⏟WC1, y3,y4⏟WC2, …, yn−1,yn⏟WCc\begin{array}{c} \underbrace{x_1, x_2}_{\textcolor{#b45309}{C_1}}, \ \underbrace{x_3, x_4}_{\textcolor{#e11d48}{C_2}}, \ \dots, \ \underbrace{x_{n-1}, x_n}_{\textcolor{#b45309}{C_c}} \\[1ex] \Big\downarrow {\scriptstyle \ \times W} \\[1ex] \underbrace{\textcolor{#b45309}{y_1, y_2}}_{{WC_1}}, \ \underbrace{\textcolor{#e11d48}{y_3, y_4}}_{{WC_2}}, \ \dots, \ \underbrace{\textcolor{#b45309}{y_{n-1}, y_n}}_{{WC_c}} \end{array}

Each channel Ci∈R2C_i \in R^2 is linearly transformed by a matrix WW.

The Jacobian of this transformation is a block diagonal matrix where each block is the n/c×n/cn/c \times n/c matrix WW. In the example below, x∈R8x \in R^8 is partitioned into four channels.

W W W W

Determinant of Jacobian is (det⁡(W))c(\det(W))^{c}.

This transformation was used in the paper Glow: Generative Flow with Invertible 1x1 Convolutions. In that paper, the linear transformation was applied to each x,yx,y location of the input tensor of shape c,h,wc, h, w. In this case, the number of "channels" is h×wh \times w.

Related articles: