Aflow-based generative model is agenerative model used inmachine learning that explicitly models aprobability distribution by leveragingnormalizing flow,^[1][2][3] which is a statistical method using thechange-of-variable law of probabilities to transform a simple distribution into a complex one.

The direct modeling of likelihood provides many advantages. For example, the negative log-likelihood can be directly computed and minimized as theloss function. Additionally, novel samples can be generated by sampling from the initial distribution, and applying the flow transformation.

In contrast, many alternative generative modeling methods such asvariational autoencoder (VAE) andgenerative adversarial network do not explicitly represent the likelihood function.

Let${\displaystyle z_0}$ be a (possibly multivariate)random variable with distribution${\displaystyle p_0(z_0)}$.

For${\displaystyle i=1,...,K}$, let${\displaystyle z_i=f_i(z_{i-1})}$ be a sequence of random variables transformed from${\displaystyle z_0}$. The functions${\displaystyle f_1,...,f_K}$ should be invertible, i.e. theinverse function${\displaystyle f_i^{-1}}$ exists. The final output${\displaystyle z_K}$ models the target distribution.

The log likelihood of${\displaystyle z_K}$ is (seederivation):

${\displaystyle \log p_K(z_K)=\log p_0(z_0)-\sum _{i=1}^K\log \left|det\frac{d{f}_i(z_{i-1})}{d{z}_{i-1}}\right|}$

Learning probability distributions by differentiating such log Jacobians originated in the Infomax (maximum likelihood) approach to ICA,^[4] which forms a single-layer (K=1) flow-based model. Relatedly, the single layer precursor of conditional generative flows appeared in.^[5]

To efficiently compute the log likelihood, the functions${\displaystyle f_1,...,f_K}$ should be easily invertible, and the determinants of their Jacobians should be simple to compute. In practice, the functions${\displaystyle f_1,...,f_K}$ are modeled usingdeep neural networks, and are trained to minimize the negative log-likelihood of data samples from the target distribution. These architectures are usually designed such that only the forward pass of the neural network is required in both the inverse and the Jacobian determinant calculations. Examples of such architectures include NICE,^[6] RealNVP,^[7] and Glow.^[8]

Consider${\displaystyle z_1}$ and${\displaystyle z_0}$. Note that${\displaystyle z_0=f_1^{-1}(z_1)}$.

By thechange of variable formula, the distribution of${\displaystyle z_1}$ is:

${\displaystyle p_1(z_1)=p_0(z_0)\left|det\frac{d{f}_1^{-1}(z_1)}{d{z}_1}\right|}$

Where${\displaystyle det\frac{d{f}_1^{-1}(z_1)}{d{z}_1}}$ is thedeterminant of theJacobian matrix of${\displaystyle f_1^{-1}}$.

By theinverse function theorem:

${\displaystyle p_1(z_1)=p_0(z_0)\left|det{\left(\frac{d{f}_1(z_0)}{d{z}_0}\right)}^{-1}\right|}$

By the identity${\displaystyle det(A^{-1})=det(A{)}^{-1}}$ (where${\displaystyle A}$ is aninvertible matrix), we have:

${\displaystyle p_1(z_1)=p_0(z_0){\left|det\frac{d{f}_1(z_0)}{d{z}_0}\right|}^{-1}}$

The log likelihood is thus:

${\displaystyle \log p_1(z_1)=\log p_0(z_0)-\log \left|det\frac{d{f}_1(z_0)}{d{z}_0}\right|}$

In general, the above applies to any${\displaystyle z_i}$ and${\displaystyle z_{i-1}}$. Since${\displaystyle \log p_i(z_i)}$ is equal to${\displaystyle \log p_{i-1}(z_{i-1})}$ subtracted by a non-recursive term, we can infer byinduction that:

${\displaystyle \log p_K(z_K)=\log p_0(z_0)-\sum _{i=1}^K\log \left|det\frac{d{f}_i(z_{i-1})}{d{z}_{i-1}}\right|}$

As is generally done when training a deep learning model, the goal with normalizing flows is to minimize theKullback–Leibler divergence between the model's likelihood and the target distribution to be estimated. Denoting${\displaystyle p_{\theta }}$ the model's likelihood and${\displaystyle p^{\ast }}$ the target distribution to learn, the (forward) KL-divergence is:

${\displaystyle D_{\text{KL}}[p^{\ast }(x)\Vert p_{\theta }(x)]=-\underset{p^{\ast }(x)}{\mathbb{E}}[\log p_{\theta }(x)]+\underset{p^{\ast }(x)}{\mathbb{E}}[\log p^{\ast }(x)]}$

The second term on the right-hand side of the equation corresponds to the entropy of the target distribution and is independent of the parameter${\displaystyle \theta }$ we want the model to learn, which only leaves the expectation of the negative log-likelihood to minimize under the target distribution. This intractable term can be approximated with a Monte-Carlo method byimportance sampling. Indeed, if we have a dataset${\displaystyle \{x_i{\}}_{i=1}^N}$ of samples each independently drawn from the target distribution${\displaystyle p^{\ast }(x)}$, then this term can be estimated as:

${\displaystyle -{\hat{\mathbb{E}}}_{p^{\ast }(x)}[\log p_{\theta }(x)]=-\frac{1}{N}\sum _{i=0}^N\log p_{\theta }(x_i)}$

Therefore, the learning objective

${\displaystyle \underset{\theta }{\mathrm{a}\mathrm{r}\mathrm{g}\mathrm{m}\mathrm{i}\mathrm{n}}{D}_{\text{KL}}[p^{\ast }(x)\Vert p_{\theta }(x)]}$

is replaced by

${\displaystyle \underset{\theta }{\mathrm{a}\mathrm{r}\mathrm{g}\mathrm{m}\mathrm{a}\mathrm{x}}\sum _{i=0}^N\log p_{\theta }(x_i)}$

In other words, minimizing theKullback–Leibler divergence between the model's likelihood and the target distribution is equivalent tomaximizing the model likelihood under observed samples of the target distribution.^[9]

A pseudocode for training normalizing flows is as follows:^[10]

• INPUT. dataset${\displaystyle x_{1:n}}$, normalizing flow model${\displaystyle f_{\theta }(\cdot ),p_0}$.
• SOLVE.${\displaystyle \underset{\theta }{max}\sum _j\log p_{\theta }(x_j)}$ by gradient descent
• RETURN.${\displaystyle \hat{\theta }}$

The earliest example.^[11] Fix some activation function${\displaystyle h}$, and let${\displaystyle \theta =(u,w,b)}$ with the appropriate dimensions, then$${\displaystyle x=f_{\theta }(z)=z+uh(⟨w,z⟩+b)}$$The inverse${\displaystyle f_{\theta }^{-1}}$ has no closed-form solution in general.

The Jacobian is${\displaystyle |det(I+h^{\prime }(⟨w,z⟩+b)u{w}^T)|=|1+h^{\prime }(⟨w,z⟩+b)⟨u,w⟩|}$.

For it to be invertible everywhere, it must be nonzero everywhere. For example,${\displaystyle h=\tanh }$ and${\displaystyle ⟨u,w⟩>-1}$ satisfies the requirement.

Let${\displaystyle x,z\in {\mathbb{R}}^{2n}}$ be even-dimensional, and split them in the middle.^[6] Then the normalizing flow functions are$${\displaystyle x=\left[\begin{array}{c}{x}_1\\ x_2\end{array}\right]=f_{\theta }(z)=\left[\begin{array}{c}{z}_1\\ z_2\end{array}\right]+\left[\begin{array}{c}0\\ m_{\theta }(z_1)\end{array}\right]}$$where${\displaystyle m_{\theta }}$ is any neural network with weights${\displaystyle \theta }$.

${\displaystyle f_{\theta }^{-1}}$ is just${\displaystyle z_1=x_1,z_2=x_2-m_{\theta }(x_1)}$, and the Jacobian is just 1, that is, the flow is volume-preserving.

When${\displaystyle n=1}$, this is seen as a curvy shearing along the${\displaystyle x_2}$ direction.

The Real Non-Volume Preserving model generalizes NICE model by:^[7]$${\displaystyle x=\left[\begin{array}{c}{x}_1\\ x_2\end{array}\right]=f_{\theta }(z)=\left[\begin{array}{c}{z}_1\\ e^{s_{\theta }(z_1)}\odot z_2\end{array}\right]+\left[\begin{array}{c}0\\ m_{\theta }(z_1)\end{array}\right]}$$

Its inverse is${\displaystyle z_1=x_1,z_2=e^{-s_{\theta }(x_1)}\odot (x_2-m_{\theta }(x_1))}$, and its Jacobian is${\displaystyle \prod _{i=1}^n{e}^{s_{\theta }(z_{1,})}}$. The NICE model is recovered by setting${\displaystyle s_{\theta }=0}$.Since the Real NVP map keeps the first and second halves of the vector${\displaystyle x}$ separate, it's usually required to add a permutation${\displaystyle (x_1,x_2)\mapsto (x_2,x_1)}$ after every Real NVP layer.

In generative flow model,^[8] each layer has 3 parts:

• channel-wise affine transform$${\displaystyle y_{cij}=s_c(x_{cij}+b_c)}$$with Jacobian${\displaystyle \prod _c{s}_c^{HW}}$.
• invertible 1x1 convolution$${\displaystyle z_{cij}=\sum _{c^{\prime }}{K}_{c{c}^{\prime }}{y}_{cij}}$$with Jacobian${\displaystyle det(K{)}^{HW}}$. Here${\displaystyle K}$ is any invertible matrix.
• Real NVP, with Jacobian as described in Real NVP.

The idea of using the invertible 1x1 convolution is to permute all layers in general, instead of merely permuting the first and second half, as in Real NVP.

An autoregressive model of a distribution on${\displaystyle {\mathbb{R}}^n}$ is defined as the following stochastic process:^[12]

where${\displaystyle {\mu }_i:{\mathbb{R}}^{i-1}\to \mathbb{R}}$ and${\displaystyle {\sigma }_i:{\mathbb{R}}^{i-1}\to (0,\mathrm{\infty })}$ are fixed functions that define the autoregressive model.

By thereparameterization trick, the autoregressive model is generalized to a normalizing flow:The autoregressive model is recovered by setting${\displaystyle z\sim N(0,I_n)}$.

The forward mapping is slow (because it's sequential), but the backward mapping is fast (because it's parallel).

The Jacobian matrix is lower-diagonal, so the Jacobian is${\displaystyle {\sigma }_1{\sigma }_2(x_1)\cdots {\sigma }_n(x_{1:n-1})}$.

Reversing the two maps${\displaystyle f_{\theta }}$ and${\displaystyle f_{\theta }^{-1}}$ of MAF results in Inverse Autoregressive Flow (IAF), which has fast forward mapping and slow backward mapping.^[13]

Instead of constructing flow by function composition, another approach is to formulate the flow as a continuous-time dynamic.^[14][15] Let${\displaystyle z_0}$ be the latent variable with distribution${\displaystyle p(z_0)}$. Map this latent variable to data space with the following flow function:

${\displaystyle x=F(z_0)=z_T=z_0+{\int }_0^{T}f(z_t,t)dt}$

where${\displaystyle f}$ is an arbitrary function and can be modeled with e.g. neural networks.

The inverse function is then naturally:^[14]

${\displaystyle z_0=F^{-1}(x)=z_T+{\int }_T^{0}f(z_t,t)dt=z_T-{\int }_0^{T}f(z_t,t)dt}$

And the log-likelihood of${\displaystyle x}$ can be found as:^[14]

${\displaystyle \log (p(x))=\log (p(z_0))-{\int }_0^T\text{Tr}\left[\frac{\mathrm{\partial }f}{\mathrm{\partial }{z}_t}\right]dt}$

Since the trace depends only on the diagonal of the Jacobian${\displaystyle {\mathrm{\partial }}_{z_t}f}$, this allows "free-form" Jacobian.^[16] Here, "free-form" means that there is no restriction on the Jacobian's form. It is contrasted with previous discrete models of normalizing flow, where the Jacobian is carefully designed to be only upper- or lower-diagonal, so that the Jacobian can be evaluated efficiently.

The trace can be estimated by "Hutchinson's trick":^[17][18]

Given any matrix${\displaystyle W\in {\mathbb{R}}^{n\times n}}$, and any random${\displaystyle u\in {\mathbb{R}}^n}$ with${\displaystyle E[u{u}^T]=I}$, we have${\displaystyle E[u^{T}Wu]=tr(W)}$. (Proof: expand the expectation directly.)

Usually, the random vector is sampled from${\displaystyle N(0,I)}$ (normal distribution) or${\displaystyle \{\pm n^{-1/2}{\}}^n}$ (Rademacher distribution).

When${\displaystyle f}$ is implemented as a neural network,neural ODE methods^[19] would be needed. Indeed, CNF was first proposed in the same paper that proposed neural ODE.

There are two main deficiencies of CNF, one is that a continuous flow must be ahomeomorphism, thus preserve orientation andambient isotopy (for example, it's impossible to flip a left-hand to a right-hand by continuous deforming of space, and it's impossible toturn a sphere inside out, or undo a knot), and the other is that the learned flow${\displaystyle f}$ might be ill-behaved, due to degeneracy (that is, there are an infinite number of possible${\displaystyle f}$ that all solve the same problem).

By adding extra dimensions, the CNF gains enough freedom to reverse orientation and go beyond ambient isotopy (just like how one can pick up a polygon from a desk and flip it around in 3-space, or unknot a knot in 4-space), yielding the "augmented neural ODE".^[20]

Any homeomorphism of${\displaystyle {\mathbb{R}}^n}$ can be approximated by a neural ODE operating on${\displaystyle {\mathbb{R}}^{2n+1}}$, proved by combiningWhitney embedding theorem for manifolds and theuniversal approximation theorem for neural networks.^[21]

To regularize the flow${\displaystyle f}$, one can impose regularization losses. The paper^[17] proposed the following regularization loss based onoptimal transport theory:$${\displaystyle {\lambda }_K{\int }_0^T{\Vert f(z_t,t)\Vert }^{2}dt+{\lambda }_J{\int }_0^T{\Vert {\mathrm{\nabla }}_{z}f(z_t,t)\Vert }_F^{2}dt}$$where${\displaystyle {\lambda }_K,{\lambda }_J>0}$ are hyperparameters. The first term punishes the model for oscillating the flow field over time, and the second term punishes it for oscillating the flow field over space. Both terms together guide the model into a flow that is smooth (not "bumpy") over space and time.

When aprobabilistic flow transforms a distribution on an${\displaystyle m}$-dimensionalsmooth manifold embedded in${\displaystyle {\mathbb{R}}^n}$, where, and where the transformation is specified as a function,${\displaystyle {\mathbb{R}}^n\to {\mathbb{R}}^n}$, the scaling factor between the source and transformedPDFs isnot given by the naive computation of thedeterminant of the${\displaystyle n\text{-by-}n}$ Jacobian (which is zero), but instead by the determinant(s) of one or more suitably defined${\displaystyle m\text{-by-}m}$ matrices. This section is an interpretation of the tutorial in the appendix of Sorrenson et al.(2023),^[22] where the more general case of non-isometrically embeddedRiemann manifolds is also treated. Here we restrict attention toisometrically embedded manifolds.

As running examples of manifolds with smooth, isometric embedding in${\displaystyle {\mathbb{R}}^n}$ we shall use:

• Theunit hypersphere:${\displaystyle {\mathbb{S}}^{n-1}=\{\mathbf{x}\in {\mathbb{R}}^n:{\mathbf{x}}^{\prime }\mathbf{x}=1\}}$, where flows can be used to generalize e.g.Von Mises-Fisher or uniform spherical distributions.
• Thesimplex interior:${\displaystyle {\mathrm{\Delta }}^{n-1}=\{\mathbf{p}=(p_1,\dots ,p_n)\in {\mathbb{R}}^n:p_i>0,\sum _i{p}_i=1\}}$, where${\displaystyle n}$-waycategorical distributions live; and where flows can be used to generalize e.g.Dirichlet, or uniform simplex distributions.

As a first example of a spherical manifold flow transform, consider thenormalized linear transform, which radially projects onto the unitsphere the output of an invertible linear transform, parametrized by the${\displaystyle n\text{-by-}n}$ invertible matrix${\displaystyle \mathbf{M}}$:

${\displaystyle f_{\text{lin}}(\mathbf{x};\mathbf{M})=\frac{\mathbf{M}\mathbf{x}}{\Vert \mathbf{M}\mathbf{x}\Vert }}$

In full Euclidean space,${\displaystyle f_{\text{lin}}:{\mathbb{R}}^n\to {\mathbb{R}}^n}$ isnot invertible, but if we restrict the domain and co-domain to the unitsphere, then${\displaystyle f_{\text{lin}}:{\mathbb{S}}^{n-1}\to {\mathbb{S}}^{n-1}}$is invertible (more specifically it is abijection and ahomeomorphism and adiffeomorphism), with inverse${\displaystyle f_{\text{lin}}(\cdot ;{\mathbf{M}}^{-1})}$. The Jacobian of${\displaystyle f_{\text{lin}}:{\mathbb{R}}^n\to {\mathbb{R}}^n}$, at${\displaystyle \mathbf{y}=f_{\text{lin}}(\mathbf{x};\mathbf{M})}$ is${\displaystyle \Vert \mathbf{M}\mathbf{x}{\Vert }^{-1}({\mathbf{I}}_n-{\mathbf{y}\mathbf{y}}^{\prime })\mathbf{M}}$, which has rank${\displaystyle n-1}$ and determinant of zero; whileas explained here, the factor (see subsection below) relating source and transformed densities is:${\displaystyle \Vert \mathbf{M}\mathbf{x}{\Vert }^{-n}\left|\det \mathbf{M}\right|}$.

For, let${\displaystyle \mathcal{M}\subset {\mathbb{R}}^n}$ be an${\displaystyle m}$-dimensional manifold with a smooth, isometric embedding into${\displaystyle {\mathbb{R}}^n}$. Let${\displaystyle f:{\mathbb{R}}^n\to {\mathbb{R}}^n}$ be a smooth flow transform with range restricted to${\displaystyle \mathcal{M}}$. Let${\displaystyle \mathbf{x}\in \mathcal{M}}$ be sampled from a distribution with density${\displaystyle P_X}$. Let${\displaystyle \mathbf{y}=f(\mathbf{x})}$, with resultant (pushforward) density${\displaystyle P_Y}$. Let${\displaystyle U\subset \mathcal{M}}$ be a small, convex region containing${\displaystyle \mathbf{x}}$ and let${\displaystyle V=f(U)}$ be its image, which contains${\displaystyle \mathbf{y}}$; then by conservation of probability mass:

${\displaystyle P_X(\mathbf{x})\mathrm{volume}(U)\approx P_Y(\mathbf{y})\mathrm{volume}(V)}$

where volume (for very small regions) is given byLebesgue measure in${\displaystyle m}$-dimensionaltangent space. By making the regions infinitessimally small, the factor relating the two densities is the ratio of volumes, which we term thedifferential volume ratio.

To obtain concrete formulas for volume on the${\displaystyle m}$-dimensional manifold, we construct${\displaystyle U}$ by mapping an${\displaystyle m}$-dimensional rectangle in (local) coordinate space to the manifold via a smooth embedding function:${\displaystyle {\mathbb{R}}^m\to {\mathbb{R}}^n}$. At very small scale, the embedding function becomes essentially linear so that${\displaystyle U}$ is aparallelotope (multidimensional generalization of a parallelogram). Similarly, the flow transform,${\displaystyle f}$ becomes linear, so that the image,${\displaystyle V=f(U)}$ is also a parallelotope. In${\displaystyle {\mathbb{R}}^m}$, we can represent an${\displaystyle m}$-dimensional parallelotope with an${\displaystyle m\text{-by-}m}$ matrix whose column-vectors are a set of edges (meeting at a common vertex) that span the paralellotope. Thevolume is given by the absolute value of the determinant of this matrix. If more generally (as is the case here), an${\displaystyle m}$-dimensional paralellotope is embedded in${\displaystyle {\mathbb{R}}^n}$, it can be represented with a (tall)${\displaystyle n\text{-by-}m}$ matrix, say${\displaystyle \mathbf{V}}$. Denoting the parallelotope as${\displaystyle /\mathbf{V}/}$, its volume is then given by the square root of theGram determinant:

${\displaystyle \mathrm{volume}/\mathbf{V}/=\sqrt{\left|\det ({\mathbf{V}}^{\prime }\mathbf{V})\right|}}$

In the sections below, we show various ways to use this volume formula to derive the differential volume ratio.

As a first example, we develop expressions for the differential volume ratio of a simplex flow,${\displaystyle \mathbf{q}=f(\mathbf{p})}$, where${\displaystyle \mathbf{p},\mathbf{q}\in \mathcal{M}={\mathrm{\Delta }}^{n-1}}$. Define theembedding function:

${\displaystyle e:\tilde{\mathbf{p}}=(p_1\dots ,p_{n-1})\mapsto \mathbf{p}=(p_1\dots ,p_{n-1},1-\sum _{i=1}^{n-1}{p}_i)}$

which maps a conveniently chosen,${\displaystyle (n-1)}$-dimensional representation,${\displaystyle \tilde{\mathbf{p}}}$, to the embedded manifold. The${\displaystyle n\text{-by-}(n-1)}$ Jacobian is${\displaystyle \mathbf{E}=\left[\begin{array}{c}{\mathbf{I}}_{n-1}\\ -{\mathbf{1}}^{\prime }\end{array}\right]}$.To define${\displaystyle U}$, the differential volume element at the transformation input (${\displaystyle \mathbf{p}\in {\mathrm{\Delta }}^{n-1}}$), we start with a rectangle in${\displaystyle \tilde{\mathbf{p}}}$-space, having (signed) differential side-lengths,${\displaystyle d{p}_1,\dots ,d{p}_{n-1}}$ from which we form the square diagonal matrix${\displaystyle \mathbf{D}}$, the columns of which span the rectangle. At very small scale, we get${\displaystyle U=e(\mathbf{D})=/\mathbf{E}\mathbf{D}/}$, with:

${\displaystyle \mathrm{volume}(U)=\sqrt{\left|\det ({\mathbf{D}\mathbf{E}}^{\prime }\mathbf{E}\mathbf{D})\right|}=\sqrt{\left|\det ({\mathbf{E}}^{\prime }\mathbf{E})\right|}\left|\det \mathbf{D})\right|=\sqrt{n}\prod _{i=1}^{n-1}\left|d{p}_i\right|}$

To understand the geometric interpretation of the factor${\displaystyle \sqrt{n}}$, see the example for the 1-simplex in the diagram at right.

The differential volume element at the transformation output (${\displaystyle \mathbf{q}\in {\mathrm{\Delta }}^{n-1}}$), is the parallelotope,${\displaystyle V=f(U)=/{\mathbf{F}}_{\mathbf{p}}\mathbf{E}\mathbf{D}/}$, where${\displaystyle {\mathbf{F}}_{\mathbf{p}}}$ is the${\displaystyle n\text{-by-}n}$ Jacobian of${\displaystyle f}$ at${\displaystyle \mathbf{p}=e(\tilde{\mathbf{p}})}$. Its volume is:

${\displaystyle \mathrm{volume}(V)=\sqrt{\left|\det ({\mathbf{D}\mathbf{E}}^{\prime }{{\mathbf{F}}_{\mathbf{p}}}^{\prime }{\mathbf{F}}_{\mathbf{p}}\mathbf{E}\mathbf{D})\right|}=\sqrt{\left|\det ({\mathbf{E}}^{\prime }{{\mathbf{F}}_{\mathbf{p}}}^{\prime }{\mathbf{F}}_{\mathbf{p}}\mathbf{E})\right|}\left|\det \mathbf{D})\right|}$

so that the factor${\displaystyle \left|\det \mathbf{D})\right|}$ cancels in the volume ratio, which can now already be numerically evaluated. It can however be rewritten in a sometimes more convenient form by also introducing therepresentation function,${\displaystyle r:\mathbf{p}\mapsto \tilde{\mathbf{p}}}$, which simply extracts the first${\displaystyle (n-1)}$ components. The Jacobian is. Observe that, since${\displaystyle e\circ r\circ f=f}$, thechain rule for function composition gives:${\displaystyle \mathbf{E}\mathbf{R}{\mathbf{F}}_{\mathbf{p}}={\mathbf{F}}_{\mathbf{p}}}$. By plugging this expansion into the above Gram determinant and then refactoring it as a product of determinants of square matrices, we can extract the factor${\displaystyle \sqrt{\left|\det ({\mathbf{E}}^{\prime }\mathbf{E})\right|}=\sqrt{n}}$, which now also cancels in the ratio, which finally simpifies to the determinant of the Jacobian of the "sandwiched" flow transformation,${\displaystyle r\circ f\circ e}$:

${\displaystyle R_f^{\mathrm{\Delta }}(\mathbf{p})=\frac{\mathrm{volume}(V)}{\mathrm{volume}(U)}=\left|\det (\mathbf{R}{\mathbf{F}}_{\mathbf{p}}\mathbf{E})\right|}$

which, if${\displaystyle \mathbf{p}\sim P_{\mathbf{P}}}$, can be used to derive the pushforward density after a change of variables,${\displaystyle \mathbf{q}=f(\mathbf{p})}$:

${\displaystyle P_{\mathbf{Q}}(\mathbf{q})=\frac{P_{\mathbf{P}}(\mathbf{p})}{R_f^{\mathrm{\Delta }}(\mathbf{p})},\text{where}\mathbf{p}=f^{-1}(\mathbf{q})}$

This formula is valid only because the simplex is flat and the Jacobian,${\displaystyle \mathbf{E}}$ is constant. The more general case for curved manifolds is discussed below, after we present two concrete examples of simplex flow transforms.

Acalibration transform,${\displaystyle f_{\text{cal}}:{\mathrm{\Delta }}^{n-1}\to {\mathrm{\Delta }}^{n-1}}$, which is sometimes used in machine learning for post-processing of the (class posterior) outputs of a probabilistic${\displaystyle n}$-class classifier,^[23][24] uses thesoftmax function to renormalize categorical distributions after scaling and translation of the input distributions in log-probability space. For${\displaystyle \mathbf{p},\mathbf{q}\in {\mathrm{\Delta }}^{n-1}}$ and with parameters,${\displaystyle a\ne 0}$ and${\displaystyle \mathbf{c}\in {\mathbb{R}}^n}$ the transform can be specified as:

${\displaystyle \mathbf{q}=f_{\text{cal}}(\mathbf{p};a,\mathbf{c})=\mathrm{softmax}(a^{-1}\log \mathbf{p}+\mathbf{c})⟺\mathbf{p}=f_{\text{cal}}^{-1}(\mathbf{q};a,\mathbf{c})=\mathrm{softmax}(a\log \mathbf{q}-a\mathbf{c})}$

where the log is applied elementwise. After some algebra thedifferential volume ratio can be expressed as:

${\displaystyle R_{\text{cal}}^{\mathrm{\Delta }}(\mathbf{p};a,\mathbf{c})=\left|\det (\mathbf{R}{\mathbf{F}}_{\mathbf{p}}\mathbf{E})\right|={\left|a\right|}^{1-n}\prod _{i=1}^n\frac{q_i}{p_i}}$

• This result can also be obtained by factoring the density of theSGB distribution,^[25] which is obtained by sendingDirichlet variates through${\displaystyle f_{\text{cal}}}$.

While calibration transforms are most often trained asdiscriminative models, the reinterpretation here as a probabilistic flow allows also the design ofgenerative calibration models based on this transform. When used for calibration, the restriction${\displaystyle a>0}$ can be imposed to prevent direction reversal in log-probability space. With the additional restriction${\displaystyle \mathbf{c}=\mathbf{0}}$, this transform (with discriminative training) is known in machine learning astemperature scaling.

The above calibration transform can be generalized to${\displaystyle f_{\text{gcal}}:{\mathrm{\Delta }}^{n-1}\to {\mathrm{\Delta }}^{n-1}}$, with parameters${\displaystyle \mathbf{c}\in {\mathbb{R}}^n}$ and${\displaystyle \mathbf{A}}$${\displaystyle n\text{-by-}n}$ invertible:^[26]

${\displaystyle \mathbf{q}=f_{\text{gcal}}(\mathbf{p};\mathbf{A},\mathbf{c})=\mathrm{softmax}(\mathbf{A}\log \mathbf{p}+\mathbf{c}),\text{subject to}\mathbf{A}\mathbf{1}=\lambda \mathbf{1}}$

where the condition that${\displaystyle \mathbf{A}}$ has${\displaystyle \mathbf{1}}$ as aneigenvector ensures invertibility by sidestepping the information loss due to the invariance:${\displaystyle \mathrm{softmax}(\mathbf{x}+\alpha \mathbf{1})=\mathrm{softmax}(\mathbf{x})}$. Note in particular that${\displaystyle \mathbf{A}=\lambda {\mathbf{I}}_n}$ is theonly allowed diagonal parametrization, in which case we recover${\displaystyle f_{\text{cal}}(\mathbf{p};{\lambda }^{-1},\mathbf{c})}$, while (for${\displaystyle n>2}$) generalizationis possible with non-diagonal matrices. Theinverse is:

${\displaystyle \mathbf{p}=f_{\text{gcal}}^{-1}(\mathbf{q};\mathbf{A},\mathbf{c})=f_{\text{gcal}}(\mathbf{q};{\mathbf{A}}^{-1},-{\mathbf{A}}^{-1}\mathbf{c}),\text{where}\mathbf{A}\mathbf{1}=\lambda \mathbf{1}⟹{\mathbf{A}}^{-1}\mathbf{1}={\lambda }^{-1}\mathbf{1}}$

Thedifferential volume ratio is:

${\displaystyle R_{\text{gcal}}^{\mathrm{\Delta }}(\mathbf{p};\mathbf{A},\mathbf{c})=\frac{\left|\det (\mathbf{A})\right|}{|\lambda |}\prod _{i=1}^n\frac{q_i}{p_i}}$

If${\displaystyle f_{\text{gcal}}}$ is to be used as a calibration transform, further constraint could be imposed, for example that${\displaystyle \mathbf{A}}$ bepositive definite, so that${\displaystyle (\mathbf{A}\mathbf{x}{)}^{\prime }\mathbf{x}>0}$, which avoids direction reversals. (This is one possible generalization of${\displaystyle a>0}$ in the${\displaystyle f_{\text{cal}}}$ parameter.)

For${\displaystyle n=2}$,${\displaystyle a>0}$ and${\displaystyle \mathbf{A}}$ positive definite, then${\displaystyle f_{\text{cal}}}$ and${\displaystyle f_{\text{gcal}}}$ are equivalent in the sense that in both cases,${\displaystyle \log \frac{p_1}{p_2}\mapsto \log \frac{q_1}{q_2}}$ is a straight line, the (positive) slope and offset of which are functions of the transform parameters. For${\displaystyle n>2,}$${\displaystyle f_{\text{gcal}}}$does generalize${\displaystyle f_{\text{cal}}}$.

It must however be noted that chaining multiple${\displaystyle f_{\text{gcal}}}$ flow transformations doesnot give a further generalization, because:

${\displaystyle f_{\text{gcal}}(\cdot ;{\mathbf{A}}_1,{\mathbf{c}}_1)\circ f_{\text{gcal}}(\cdot ;{\mathbf{A}}_2,{\mathbf{c}}_2)=f_{\text{gcal}}(\cdot ;{\mathbf{A}}_1{\mathbf{A}}_2,{\mathbf{c}}_1+{\mathbf{A}}_1{\mathbf{c}}_2)}$

In fact, the set of${\displaystyle f_{\text{gcal}}}$ transformations form agroup under function composition. The set of${\displaystyle f_{\text{cal}}}$ transformations form a subgroup.

Also see:Dirichlet calibration,^[27] which generalizes${\displaystyle f_{\text{gcal}}}$, by not placing any restriction on the matrix,${\displaystyle \mathbf{A}}$, so that invertibility is not guaranteed. While Dirichlet calibration is trained as a discriminative model,${\displaystyle f_{\text{gcal}}}$ can also be trained as part of a generative calibration model.

Consider a flow,${\displaystyle \mathbf{y}=f(\mathbf{x})}$ on a curved manifold, for example${\displaystyle {\mathbb{S}}^{n-1}}$ which we equip with the embedding function,${\displaystyle e}$ that maps a set of${\displaystyle (n-1)}$angular spherical coordinates to${\displaystyle {\mathbb{S}}^{n-1}}$. The Jacobian of${\displaystyle e}$ is non-constant and we have to evaluate it at both input (${\displaystyle {\mathbf{E}}_{\mathbf{x}}}$) and output (${\displaystyle {\mathbf{E}}_{\mathbf{y}}}$). The same applies to${\displaystyle r}$, the representation function that recovers spherical coordinates from points on${\displaystyle {\mathbb{S}}^{n-1}}$, for which we need the Jacobian at the output (${\displaystyle {\mathbf{R}}_{\mathbf{y}}}$). The differential volume ratio now generalizes to:

${\displaystyle R_f(\mathbf{x})=\left|\det ({\mathbf{R}}_{\mathbf{y}}{\mathbf{F}}_{\mathbf{x}}{\mathbf{E}}_{\mathbf{x}})\right|\frac{\sqrt{\left|\det ({\mathbf{E}}_{\mathbf{y}}^{\prime }{\mathbf{E}}_{\mathbf{y}})\right|}}{\sqrt{\left|\det ({\mathbf{E}}_{\mathbf{x}}^{\prime }{\mathbf{E}}_{\mathbf{x}})\right|}}}$

For geometric insight, consider${\displaystyle {\mathbf{S}}^2}$, where the spherical coordinates are co-latitude,${\displaystyle \theta \in [0,\pi ]}$ and longitude,${\displaystyle \varphi \in [0,2\pi )}$. At${\displaystyle \mathbf{x}=e(\theta ,\varphi )}$, we get${\displaystyle \sqrt{\left|\det ({\mathbf{E}}_{\mathbf{x}}^{\prime }{\mathbf{E}}_{\mathbf{x}})\right|}=\sin \theta }$, which gives the radius of the circle at that latitude (compare e.g. polar circle to equator). The differential volume (surface area on the sphere) is:${\displaystyle \sin \theta d\theta d\varphi }$.

The above derivation for${\displaystyle R_f}$ is fragile in the sense that when usingfixed functions${\displaystyle e,r}$, there may be places where they are not well-defined, for example at the poles of the 2-sphere where longitude is arbitrary. This problem is sidestepped (using standard manifold machinery) by generalizing tolocal coordinates (charts), where in the vicinities of${\displaystyle \mathbf{x},\mathbf{y}\in \mathcal{M}}$, we map from local${\displaystyle m}$-dimensional coordinates to${\displaystyle {\mathbb{R}}^n}$ and back using the respective function pairs${\displaystyle e_{\mathbf{x}},r_{\mathbf{x}}}$ and${\displaystyle e_{\mathbf{y}},r_{\mathbf{y}}}$. We continue to use the same notation for the Jacobians of these functions (${\displaystyle {\mathbf{E}}_{\mathbf{x}},{\mathbf{E}}_{\mathbf{y}},{\mathbf{R}}_{\mathbf{y}}}$), so that the above formula for${\displaystyle R_f}$ remains valid.

Wecan however, choose our local coordinate system in a way that simplifies the expression for${\displaystyle R_f}$ and indeed also its practical implementation.^[22] Let${\displaystyle \pi :\mathcal{P}\to {\mathbb{R}}^n}$ be a smooth idempotent projection (${\displaystyle \pi \circ \pi =\pi }$) from theprojectible set,${\displaystyle \mathcal{P}\subseteq {\mathbb{R}}^n}$, onto the embedded manifold. For example:

• The positive orthant of${\displaystyle {\mathbb{R}}^n}$ is projected onto thesimplex as:${\displaystyle \pi (\mathbf{z})=(\sum _{i=1}^n{z}_i{)}^{-1}\mathbf{z}}$
• Non-zero vectors in${\displaystyle {\mathbb{R}}^n}$ are projected onto theunitsphere as:${\displaystyle \pi (\mathbf{z})=(\sum _{i=1}^n{z}_i^2{)}^{-\frac{1}{2}}\mathbf{z}}$

For every${\displaystyle \mathbf{x}\in \mathcal{M}}$, we require of${\displaystyle \pi }$ that its${\displaystyle n\text{-by-}n}$ Jacobian,${\displaystyle {\mathbf{\Pi }}_{\mathit{x}}}$ has rank${\displaystyle m}$ (the manifold dimension), in which case${\displaystyle {\mathbf{\Pi }}_{\mathit{x}}}$ is anidempotent linear projection onto the local tangent space (orthogonal for the unitsphere:${\displaystyle {\mathbf{I}}_n-{\mathbf{x}\mathbf{x}}^{\prime }}$;oblique for the simplex:${\displaystyle {\mathbf{I}}_n-{\mathit{x}\mathbf{1}}^{\prime }}$). The columns of${\displaystyle {\mathbf{\Pi }}_{\mathit{x}}}$ span the${\displaystyle m}$-dimensional tangent space at${\displaystyle \mathbf{x}}$. We use the notation,${\displaystyle {\mathbf{T}}_{\mathbf{x}}}$ for any${\displaystyle n\text{-by-}m}$ matrix with orthonormal columns (${\displaystyle {\mathbf{T}}_{\mathbf{x}}^{\prime }{\mathbf{T}}_{\mathbf{x}}={\mathbf{I}}_m}$) that span the local tangent space. Also note:${\displaystyle {\mathbf{\Pi }}_{\mathit{x}}{\mathbf{T}}_{\mathbf{x}}={\mathbf{T}}_{\mathbf{x}}}$. We can now choose our local coordinate embedding function,${\displaystyle e_{\mathbf{x}}:{\mathbb{R}}^m\to {\mathbb{R}}^n}$:

${\displaystyle e_{\mathbf{x}}(\tilde{x})=\pi (\mathbf{x}+{\mathbf{T}}_{\mathbf{x}}\stackrel{\mathbf{~}}{\mathbf{x}}),\text{with Jacobian:}{\mathbf{E}}_{\mathbf{x}}={\mathbf{T}}_{\mathbf{x}}\text{at}\tilde{\mathbf{x}}=\mathbf{0}.}$

Since the Jacobian is injective (full rank:${\displaystyle m}$), a local (not necessarily unique)left inverse, say${\displaystyle r_{\mathbf{x}}^{\ast }}$ with Jacobian${\displaystyle {\mathbf{R}}_{\mathbf{x}}^{\ast }}$, exists such that${\displaystyle r_{\mathbf{x}}^{\ast }(e_{\mathbf{x}}(\tilde{x}))=\tilde{x}}$ and${\displaystyle {\mathbf{R}}_{\mathbf{x}}^{\ast }{\mathbf{T}}_{\mathbf{x}}={\mathbf{I}}_m}$. In practice we do not need the left inverse function itself, but wedo need its Jacobian, for which the above equation does not give a unique solution. We can however enforce a unique solution for the Jacobian by choosing the left inverse as,${\displaystyle r_{\mathbf{x}}:{\mathbb{R}}^n\to {\mathbb{R}}^m}$:

${\displaystyle r_{\mathbf{x}}(\mathbf{z})=r_{\mathbf{x}}^{\ast }(\pi (\mathbf{z})),\text{with Jacobian:}{\mathbf{R}}_{\mathbf{x}}={\mathbf{T}}_{\mathbf{x}}^{\prime }}$

We can now finally plug${\displaystyle {\mathbf{E}}_{\mathbf{x}}={\mathbf{T}}_{\mathbf{x}}}$ and${\displaystyle {\mathbf{R}}_{\mathbf{y}}={\mathbf{T}}_{\mathbf{y}}^{\prime }}$ into our previous expression for${\displaystyle R_f}$, thedifferential volume ratio, which because of the orthonormal Jacobians, simplifies to:^[28]

${\displaystyle R_f(\mathbf{x})=\left|\det ({{\mathbf{T}}_{\mathbf{y}}}^{\prime }{\mathbf{F}}_{\mathbf{x}}{\mathbf{T}}_{\mathbf{x}})\right|}$

For learning the parameters of a manifold flow transformation, we need access to the differential volume ratio,${\displaystyle R_f}$, or at least to its gradient w.r.t. the parameters. Moreover, for some inference tasks, we need access to${\displaystyle R_f}$ itself. Practical solutions include:

• Sorrenson et al.(2023)^[22] give a solution for computationally efficient stochastic parameter gradient approximation for${\displaystyle \log R_f.}$
• For some hand-designed flow transforms,${\displaystyle R_f}$ can be analytically derived in closed form, for example the above-mentioned simplex calibration transforms. Further examples are given below in the section on simple spherical flows.
• On a software platform equipped withlinear algebra andautomatic differentiation,${\displaystyle R_f(\mathbf{x})=\left|\det ({{\mathbf{T}}_{\mathbf{y}}}^{\prime }{\mathbf{F}}_{\mathbf{x}}{\mathbf{T}}_{\mathbf{x}})\right|}$ can be automatically evaluated, given access to only${\displaystyle \mathbf{x},f,\pi }$.^[29] But this is expensive for high-dimensional data, with at least${\displaystyle \mathcal{O}(n^3)}$ computational costs. Even then, the slow automatic solution can be invaluable as a tool for numerically verifying hand-designed closed-form solutions.

In machine learning literature, various complex spherical flows formed by deep neural network architectures may be found.^[22] In contrast, this section compiles fromstatistics literature the details of three very simple spherical flow transforms, with simple closed-form expressions for inverses and differential volume ratios. These flows can be used individually, or chained, to generalize distributions on the unitsphere,${\displaystyle {\mathbb{S}}^{n-1}}$. All three flows are compositions of an invertible affine transform in${\displaystyle {\mathbb{R}}^n}$, followed by radial projection back onto the sphere. The flavours we consider for the affine transform are: pure translation, pure linear and general affine. To make these flows fully functional for learning, inference and sampling, the tasks are:

• To derive the inverse transform, with suitable restrictions on the parameters to ensure invertibility.
• To derive in simple closed form thedifferential volume ratio,${\displaystyle R_f}$.

An interesting property of these simple spherical flows is that they don't make use of any non-linearities apart from the radial projection. Even the simplest of them, the normalized translation flow, can be chained to form perhaps surprisingly flexible distributions.

The normalized translation flow,${\displaystyle f_{\text{trans}}:{\mathbb{S}}^{n-1}\to {\mathbb{S}}^{n-1}}$, with parameter${\displaystyle \mathbf{c}\in {\mathbb{R}}^n}$, is given by:

The inverse function may be derived by considering, for${\displaystyle \ell >0}$:${\displaystyle \mathbf{y}={\ell }^{-1}(\mathbf{x}+\mathbf{c})}$ and then using${\displaystyle {\mathbf{x}}^{\prime }\mathbf{x}=1}$ to get aquadratic equation to recover${\displaystyle \ell }$, which gives:

${\displaystyle \mathbf{x}=f_{\text{trans}}^{-1}(\mathbf{y};\mathbf{c})=\ell \mathbf{y}-\mathbf{c},\text{where}\ell ={\mathbf{y}}^{\prime }\mathbf{c}+\sqrt{({\mathbf{y}}^{\prime }\mathbf{c}{)}^2+1-{\mathbf{c}}^{\prime }\mathbf{c}}}$