Autoregressive Flows: MAF & IAF

Table of Contents

1. Learning Objectives
2. Notation
3. Core Intuition
4. Autoregressive Models as Flows
5. Masked Autoregressive Flow (MAF)
6. Inverse Autoregressive Flow (IAF)
7. The Density-Sampling Tradeoff
8. Neural Spline Flows
9. Masked Networks (MADE)
10. Common Pitfalls
11. Summary
12. Exercises

Learning Objectives

1. Derive MAF from autoregressive density estimation.
2. Explain the density-vs-sampling speed tradeoff between MAF and IAF.
3. Construct a MADE network for parallel autoregressive evaluation.
4. Explain neural spline flows and their advantage over affine transforms.
5. Choose between MAF and IAF for different applications.

Notation

• $x_i$ — $i$-th dimension of the input
• $\mu_i(\cdot), \sigma_i(\cdot)$ — autoregressive parameters
• $\pi$ — ordering/permutation of dimensions

Core Intuition

Autoregressive flows exploit the chain rule of probability: any joint distribution can be decomposed as $p(\mathbf{x}) = \prod_i p(x_i | x_{<i})$. By parameterizing each conditional as a simple transform of a base variable, we get a flow with a FULL triangular Jacobian (not block-triangular like coupling flows). This gives maximum expressiveness per layer — but at the cost of sequential computation in one direction.

Autoregressive Models as Flows

Any autoregressive model $p(x_i | x_{<i}) = \mathcal{N}(\mu_i(x_{<i}), \sigma_i^2(x_{<i}))$ defines an invertible transformation:

Forward (sampling from noise):

Inverse (mapping data to noise):

Jacobian (lower triangular):

Masked Autoregressive Flow (MAF)

MAF (Papamakarios et al., 2017): Use the INVERSE direction as the flow:

Density evaluation: Given data $\mathbf{x}$, compute ALL $z_i$ in ONE parallel pass (all $x_{<i}$ are known):

Sampling: Must compute sequentially: $x_1 \to x_2 \to \ldots \to x_d$ (each depends on previous).

Properties:

• Density: $O(1)$ passes (parallel). Fast for training/evaluation.
• Sampling: $O(d)$ sequential steps. Slow for generation.

Inverse Autoregressive Flow (IAF)

IAF (Kingma et al., 2016): Use the FORWARD direction:

Sampling: Given noise $\mathbf{z}$, compute ALL $x_i$ in ONE pass... wait, no — $x_i$ depends on $z_{<i}$, and $z_i$ is the input. So sampling IS parallel.

Actually: sample $\mathbf{z}$, then $x_i = f_i(z_1, \ldots, z_i)$. This is sequential in $z$'s ordering, but with MADE can be parallelized.

Density evaluation: Given data $\mathbf{x}$, must invert sequentially to find $\mathbf{z}$.

Properties:

• Sampling: $O(1)$ passes (parallel given noise). Fast for generation.
• Density: $O(d)$ sequential steps. Slow for evaluation.

The Density-Sampling Tradeoff

Implication:

• Training (needs density): MAF is preferred.
• Generation (needs sampling): IAF is preferred.
• Both fast: Coupling flows — but less expressive per layer.

Combining: Use MAF for training the density model; distill into IAF for fast sampling (Parallel WaveNet).

Neural Spline Flows

Problem with affine transforms: Each $x_i \to z_i$ mapping is linear (given conditioning). Limited expressiveness per dimension.

Solution (Durkan et al., 2019): Replace affine with monotone rational-quadratic splines:

• Divide each dimension into $K$ bins.
• Within each bin: a rational-quadratic function (monotone, smooth).
• Parameters: bin widths $w_k$, heights $h_k$, and derivatives $d_k$ at knots.

Jacobian: $f'(x)$ is the spline derivative (analytic, cheap).

Advantage: Much more expressive per layer. Fewer layers needed. State-of-the-art density estimation.

Masked Networks (MADE)

MADE (Germain et al., 2015): A single neural network that computes ALL autoregressive conditionals simultaneously.

Trick: Mask the weight matrices to enforce the autoregressive property:

where $m^{(l)}(k)$ is the "degree" of unit $k$ in layer $l$ — the maximum input index it's allowed to depend on.

Result: One forward pass through MADE computes $\mu_i(x_{<i})$ and $\sigma_i(x_{<i})$ for ALL $i$ simultaneously. This makes MAF density evaluation truly $O(1)$ network evaluations.

Common Pitfalls

Pitfall 1. Assuming MAF can generate samples quickly. Each sample requires $d$ sequential network evaluations — impractical for high dimensions.

Pitfall 2. Using fixed dimension ordering. Random or learned orderings can significantly improve performance. Use multiple random orderings (ensemble).

Pitfall 3. Using affine transforms when splines are available. Spline flows achieve the same density quality with 5-10x fewer layers.

Summary

• MAF: Fast density, slow sampling; ideal for training and evaluation.
• IAF: Fast sampling, slow density; ideal for generation (after training via other means).
• Fundamental tradeoff arises from the direction of the autoregressive decomposition.
• MADE networks enable parallel computation of autoregressive parameters.
• Neural spline flows dramatically improve per-layer expressiveness over affine transforms.
• Coupling flows sacrifice per-layer expressiveness for bidirectional speed.

Exercises

Exercise 1. For a 784-dimensional MAF (MNIST flattened): compute the number of sequential steps needed to generate one sample.

Exercise 2. Derive the Jacobian determinant of a rational-quadratic spline element.

Exercise 3. Design the MADE mask matrices for a 4-dimensional input with 2 hidden layers of width 8.

Exercise 4. Prove that MAF with a single MADE layer is equivalent to a Gaussian autoregressive model.

Exercise 5. Compare BPD (bits per dimension) of MAF vs RealNVP vs Neural Spline Flow on MNIST. Explain the gap.