Autoregressive Flows: MAF and IAF

An autoregressive model with Gaussian conditionals is already a normalising flow, and noticing that is the whole idea. Write \(x_i = \mu_i(x_{<i}) + \sigma_i(x_{<i}) \cdot u_i\) with \(u \sim \mathcal{N}(0, I)\). The map \(u \mapsto x\) is invertible, and its Jacobian is triangular, so the log-determinant is just \(\sum_i \log \sigma_i\). Stack a few of these, permuting the variable order between layers, and you have a flexible density model whose likelihood is exact and cheap (Papamakarios, Pavlakou and Murray, 2017, Masked Autoregressive Flow for Density Estimation, arXiv:1705.07057).

The conditioners are computed with masking rather than recurrence, so all \(\mu_i\) and \(\sigma_i\) come from one pass of a masked network (Germain et al., 2015, MADE, arXiv:1502.03509). That is what makes the direction asymmetry sharp instead of merely inconvenient.

The asymmetry

Evaluating \(u = (x - \mu(x_{<i}))/\sigma(x_{<i})\) needs only \(x\), which is fully available, so one masked forward pass produces every \(u_i\) at once. That is the density-evaluation direction, and it is parallel.

Sampling reverses it. To compute \(x_1\) you need \(\mu_1\), which needs nothing; to compute \(x_2\) you need \(\mu_2(x_1)\), which needs \(x_1\). \(D\) sequential passes for \(D\) dimensions. For a 3,072-dimensional CIFAR image that is 3,072 network evaluations per sample.

Inverse autoregressive flow is the same transformation with the roles swapped: it makes the autoregression run over the noise variables, \(z_i = \mu_i(u_{<i}) + \sigma_i(u_{<i}) \cdot u_i\) (Kingma et al., 2016, Improving Variational Inference with Inverse Autoregressive Flow, arXiv:1606.04934). Now sampling is one parallel pass and evaluating the density of an externally supplied point is sequential. Neither direction was made cheap; the expense was moved.

Which asymmetry fits which job

The VAE case is the subtle one and explains why IAF exists. The ELBO needs \(\log q(z \mid x)\) only at points the encoder itself sampled, and those come with their \(u\) already in hand, so the sequential inverse is never invoked. IAF gets a rich posterior at one parallel pass per flow step. Hand it someone else's \(z\) and it becomes as slow as MAF is at sampling.

Real NVP sits at the compromise point: coupling layers split the variables in half and transform one half conditioned on the other, which makes both directions a single pass at the cost of a much weaker per-layer transformation. MAF is a strict generalisation of it, obtained by letting the "conditioning set" be every earlier variable rather than a fixed half.

Getting both directions

Probability density distillation trains a parallel IAF student to match a slow autoregressive teacher, which is how neural vocoders reached real-time speech synthesis while keeping teacher-quality samples: Parallel WaveNet distilled a sequential WaveNet into a feedforward IAF and generated audio more than twenty times faster than real time (van den Oord et al., 2018, Parallel WaveNet, arXiv:1711.10433). The student never evaluates external densities, so its slow direction is never exercised. This pattern, train the accurate sequential model then distil into the parallel one, recurs throughout generative modelling and is worth recognising in its general form.

When it breaks

Variable order is a modelling choice with teeth. A triangular Jacobian requires a fixed ordering, and for images or tabular data no ordering is natural. Random permutations between layers help; Glow's learned \(1 \times 1\) convolution generalises the permutation to a rotation, at the cost of a determinant that must be computed rather than read off.

Depth compounds the slow direction. \(L\) stacked autoregressive layers make the sequential direction \(L \times D\) passes, not \(D\).

Numerical inversion is not a way out. The inverse can be found by bisection per dimension for monotone transformations, but it is still sequential in \(i\) and adds iterations per step. It buys flexibility, not speed.

Exact likelihoods on discrete data are a trap. Flows are continuous densities, and fitting one to integer pixel values without dequantisation lets the density concentrate on the lattice and diverge. The reported bits per dimension is then meaningless.