Planar and Radial Flows

Stacking multiple invertible functions creates highly expressive probability distributions. To make this work in practice, specific mathematical functions are required that are both invertible and have easily computable Jacobian determinants. Two foundational architectures that meet these strict requirements are planar flows and radial flows. These models provide an excellent starting point for understanding how to warp simple probability densities into complex geometric shapes.

Planar Flows

A planar flow modifies the input space by applying a transformation along a specific hyperplane. You can think of it as taking a flat sheet of rubber and stretching or compressing it along a single straight line. Everything perpendicular to that line remains relatively unchanged, while the space along the line expands or contracts.

Let $z \in \mathbb{R}^D$ be our continuous input vector. The planar flow transformation $f(z)$ is defined mathematically as:

$f(z) = z + u h(w^T z + b)$

Here, $w \in \mathbb{R}^D$ and $u \in \mathbb{R}^D$ are learnable parameter vectors that define the orientation and scale of the transformation. The term $b \in \mathbb{R}$ is a scalar bias that shifts the hyperplane. The function $h$ is a smooth, differentiable non-linear activation function. The hyperbolic tangent function, denoted as $\tanh$, is most commonly used for $h$ because its derivative is well-behaved and bounded.

To use this transformation in a normalizing flow, we must compute the determinant of its Jacobian matrix. Taking the derivative of $f(z)$ with respect to the input vector $z$ yields:

$\frac{\partial f}{\partial z} = I + u h'(w^T z + b) w^T$

This result is an identity matrix $I$ plus the outer product of two vectors, scaled by the derivative of the activation function. We can compute the determinant of this specific matrix structure highly efficiently using the matrix determinant lemma. The lemma states that $\det(A + u v^T) = (1 + v^T A^{-1} u) \det(A)$. Applying this linear algebra identity to our Jacobian gives a much simpler scalar equation:

$\det\left(\frac{\partial f}{\partial z}\right) = 1 + u^T w h'(w^T z + b)$

This scalar result is exceptionally computationally cheap. It evaluates in $O(D)$ time instead of the standard $O(D^3)$ time required to compute a general $D \times D$ matrix determinant. This efficiency allows us to stack hundreds of planar flow layers without prohibitive computational costs during training.

Computational graph of a single planar flow transformation showing the skip connection adding the original input to the scaled activation.

For the flow to be valid, the transformation must be strictly invertible. This requires that the Jacobian determinant never equals zero and maintains a consistent sign. In practice, we constrain the parameters to ensure that $1 + u^T w h'(w^T z + b) > 0$. When using the $\tanh$ activation function, its derivative $h'$ is strictly bounded between $0$ and $1$. We enforce invertibility by modifying the vector $u$ slightly during the forward pass to satisfy the geometric condition $u^T w \geq -1$.

Radial Flows

While planar flows apply transformations along a straight hyperplane, radial flows introduce distortions radiating outward from a specific center point. You can imagine a radial flow as placing a magnifying glass over a specific coordinate in the space. The space is either stretched outward from that point or compressed inward toward it.

For an input vector $z$, the radial flow transformation is defined as:

$f(z) = z + \beta h(\alpha, r)(z - z_0)$

Here, $z_0 \in \mathbb{R}^D$ is the learnable reference point or center of the flow. The scalar $r = ||z - z_0||$ represents the Euclidean distance between the input and the reference point. The parameters $\alpha \in \mathbb{R}^+$ and $\beta \in \mathbb{R}$ dictate the spread and magnitude of the distortion. The function $h(\alpha, r)$ is designed to decay as the distance $r$ increases, typically defined as $1 / (\alpha + r)$.

Putting these components together gives the explicit formulation:

$f(z) = z + \beta \frac{1}{\alpha + ||z - z_0||} (z - z_0)$

The Jacobian determinant for a radial flow also benefits from a specialized computation that avoids building the full $D \times D$ matrix. Using similar matrix identities, the determinant is evaluated as:

$\det\left(\frac{\partial f}{\partial z}\right) = \left( 1 + \beta h(\alpha, r) \right)^{D-1} \left( 1 + \beta h(\alpha, r) + \beta h'(\alpha, r) r \right)$

Just like planar flows, this determinant calculation operates in linear time $O(D)$. To guarantee invertibility for a radial flow, we mathematically restrict the parameter $\beta$ such that $\beta \geq -\alpha$.

Probability density function before and after applying an invertible transformation layer.

A single planar or radial layer is quite limited in what it can model. A single planar layer only compresses or expands along one direction, and a single radial layer only warps around a single point. To approximate highly irregular probability densities, we apply these transformations sequentially. By chaining $K$ layers together, the overall transformation becomes heavily non-linear and fully capable of mapping a simple isotropic base distribution into highly complex target shapes.