Autoregressive Flows

Autoregressive flows are a generalization of coupling flows that treat each input dimension as a separate “block”.

The $d$-th dimension of the output $h^{\prime }$ is computed based on the first $d-1$ dimensions of the input $h$:

$$h_d^{\prime }=g[h_d,\varphi [h_{1:d-1}]]$$

The function $g[\bullet ,\bullet ]$ is termed the transformer (nothing to do with attention), and the parameters $\varphi ,\varphi [h_1],\varphi [h_1,h_2]$ are termed conditioners.

• Similar to coupling flows, the transformer $g[\bullet ,\varphi ]$ must be invertible
• The conditioners $\varphi [\bullet ]$ can take any form and are usually neural networks
• If the transformer and conditioner are sufficiently flexible, autoregressive flows are universal approximators that can represent any probability distribution.

It’s possible to compute all of the entries of the output $h^{\prime }$ in parallel using a network with appropriate masks so that the parameters $\varphi$ at position $d$ only depend on previous positions. This is known as masked autoregressive flow. The principle is very similar to masked self-attention; connections that relate inputs to previous outputs are pruned.

Inverting the transformation is less efficient. Consider the forward mapping:

$$\begin{array}{rl}{h}_1^{\prime }& =g[h_1,\varphi ]\\ h_2^{\prime }& =g[h_2,\varphi [h_1]]\\ h_3^{\prime }& =g[h_3,\varphi [h_{1:2}]]\\ h_4^{\prime }& =g[h_4,\varphi [h_{1:3}]]\end{array}$$

This must be inverted sequentially:

$$\begin{array}{rl}{h}_1& =g^{-1}[h_1^{\prime },\varphi ]\\ h_2& =g^{-1}[h_2^{\prime },\varphi [h_1]]\\ h_3& =g^{-1}[h_3^{\prime },\varphi [h_{1:2}]]\\ h_4& =g{r}^{-1}[h_4^{\prime },\varphi [h_{1:3}]]\end{array}$$

This can’t be done in parallel as the computation for $h_d$ depends on $h_{1:d-1}$ (i.e., the partial results so far). Hence, inversion is time-consuming when the input is large.

Inverse autoregressive flows

Masked autoregressive flows are defined in the normalizing (inverse) direction. This is required to evaluate the likelihood efficiently and hence learn the model. However, sampling requires the forward direction, in which each variable must be computed sequentially at each layer, which is slow. If we use an autoregressive flow for the forward (generative) transformation, then sampling is efficient, but computing the likelihood (and training) is slow. This is known as inverse autoregressive flow.

A trick that allows fast learning and also fast (but approximate) sampling is to build a masked autoregressive flow to learn the distribution (teacher) and then use this to train an inverse autoregressive flow from which we can sample efficiently (student). This requires a different formulation of normalizing flows that learns from another function rather than a samples (Normalizing Flows for Approximating Densities).