Turning Normalizing Flows into Monge Maps with Geodesic Gaussian Preserving Flows
Abstract
Normalizing Flows (NF) are powerful likelihood-based generative models that are able to trade off between expressivity and tractability to model complex densities.
A now well established research avenue leverages optimal transport (OT) and looks for Monge maps, i.e. models with minimal effort between the source and target distributions.
This paper introduces a method based on Brenier’s polar factorization theorem to transform any trained NF into a more OT-efficient version without changing the final density.
We do so by learning a rearrangement of the source (Gaussian) distribution that minimizes the OT cost between the source and the final density.
We further constrain the path leading to the estimated Monge map to lie on a geodesic in the space of volume-preserving diffeomorphisms thanks to Euler’s equations.
The proposed method leads to smooth flows with reduced OT cost for several existing models without affecting the model performance.
Introduction
To keep the length of this extended abstract short we have removed some proofs.
All the proofs are given in the full version of the paper available here https://arxiv.org/abs/2209.10873 .
Normalizing flows.
A popular class of generative models is Normalizing flows (NF).
NF models transform a known probability distribution (Gaussian in most cases) into a complex one allowing for efficient sampling and density estimation.
To do so they use a smooth diffeomorphism f : R d → R d which maps a target probability distribution µ to the known source distribution ν = f # µ [1, 2].
In practice the flow must satisfy the change of variables formula log p µ ( x ) = log p
ν ( f ( x ))
+ log | det ∇ f ( x ) | .
There are many possible parameterizations of f , usually relying on automatic differentiation to train their parameters via first order optimization algorithms.
For density estimation applications, training is done by maximizing the likelihood of the observed data [1, 3, 2, 4].
Optimal transport.
The question of choosing the ”best” transformation among all existing ones, independently from how accurately µ models the target distribution is an important one.
One way to make the architecture unique (under appropriate conditions on the two distributions) is to use optimal transport [5, 6, 7, 8], that is to choose the one giving the Wasserstein distance between µ and ν , with a squared L 2 ground cost W 22 ( µ, ν ) =
min f (cid:82) R d | f ( x ) − x | 2 dµ ( x ) , ν = f # µ.
An optimal model minimizes the total mass displacement which can be a desirable property even if it is often a difficult task.
In particular Brenier’s theorem [9] states that the optimal function f is the gradient of a scalar convex function, which is widely used in practice.
1.1 Main contributions
Polar factorization.
An overlooked implication of Brenier’s theorem is the so-called polar factorization theorem, that states that the optimal transport map ∇ ψ can be factorized into the composition of two functions ∇ ψ
= s ◦ f , the function f being some arbitrary smooth map from µ to ν and s an associated measure preserving function of ν [9].
The idea we exploit is the possibility, from a given flow f (and its corresponding inverse g = f − 1 ), to rearrange the distribution ν using s to obtain a new map reducing the OT cost without changing the distribution given by the push-forward µ :=
g # ν =
( g ◦ s − 1 )
# ν .
The OT-improved map can then be obtained with the composition g ◦ s − 1 .
Interestingly this property is not as popular as the previous one and to our knowledge is not used when dealing with OT and NF models.
Yet normalizing flows can take advantage of this formulation mostly because the distribution ν is known and simple which make it possible to construct architectures preserving ν .
We consider the most common case where the known distribution is a standard normal and call such rearranging maps Gaussian Preserving (GP) flows.
An important point is that by construction our GP map will only change the OT cost of the model.
The target density and therefore the training loss given by the model will stay the same.
his allows us to take any pre-trained model and compute the associated Monge map, thus improving the model in terms of OT displacement from the source to the target distribution, without changing the modeled density see Figure 1.
Euler’s equations.
Since several GP flow models can solve the same OT problem, we also look for a way to find the ”best” GP flow.
In this work, we show that the geodesics associated with the OT problem are actually given by solutions to the Euler equations, following a celebrated result by Arnold [11].
The penalization of Euler’s equations in high dimensions and its practical implementation is therefore also considered.
Disentanglement preservation with optimal transport.
Finally we show one potential interest of GP flows by studying the preservation of the data structure experimentally.
More specifically we focus on the preservation of disentanglement on the dSprites dataset [12] in some variational auto-encoder (VAE) latent space.
On this particular example we show that OT allows to preserve the structure of the latent data points which is otherwise destroyed when applying the NF model.
Gaussian preserving flows
Polar factorization theorem.
The main idea is to use the Brenier’s polar factorization theorem to construct the Monge map with a rearrangement of the known probability distribution ν .
To preserve the conventions from Brenier’s paper [9], we study the OT problem defined from the known probability distribution ν to µ and therefore consider the function g :=
f − 1 .
Theorem 1 (Brenier’s polar factorization [9]) .
Let ( X , ν ) be a probability space, X ⊂ R d open bounded.
Then for each non-degenerate g ∈ L p ( X , ν, R d ) , there exists a unique convex function ψ : X → R and a measure preserving function s : X → X such that g ( x ) = ∇ ψ ( s ( x )) , and s ( x ) minimizes the cost (cid:82) X | g ( x ) − s ( x ) | 2 dν ( x ) .
Our goal is to leverage the polar factorization theorem in order to solve the OT problem between ν and µ :=
g
# ν where g is given and ν = N ( 0 , Id) , by looking for the rearrangement s via an optimization problem.
To do so we need to construct a class of measure preserving maps.
Gaussian preserving flows.
Consider two probability measures α and β with density h α and h β respectively.
A map s is measure preserving between α and β if it satisfies the change of variable equality h β ( x ) =
h α ( s ( x )) | det( ∇ s ( x )) | .
In our case, we want s to be Gaussian preserving therefore h
α = h β = e −∥ x ∥ 2 / 2 and one gets
[Equation]
It turns out that Lebesgue preserving functions (i.e. satisfying | det ∇ ϕ | = 1 ) can be used to construct maps satisfying (1).
In the following we will denote erf : R d → R d the distribution function of a one dimensional Gaussian (that is erf( x ) = 2 √ π (cid:82) x 0 e t 2 dt ) applied component wise.
Proposition 1.
Let s be a smooth preserving function (i.e. satisfying (1) ).
Then there exists ϕ : ( − 1 , 1) d → ( − 1 , 1) d such that | det ∇ ϕ | = 1 and s ( x ) = √ 2 erf − 1 ◦ ϕ ◦ erf ( x √ 2 ) , x ∈ R d .
From now on we will focus on the construction of volume and orientation preserving maps (i.e. satisfying det ∇ ϕ = 1 ) since functions satisfying det ∇ ϕ
= − 1 can be constructed from them.
It is also possible to show that under simple hypothesis on the Monge map and the NF architecture, the
GP flow s is C 1 and either det ∇ ϕ = 1 everywhere or det ∇ ϕ
= − 1 everywhere.
2.1 Volume-orientation preserving maps
First we introduce the space SDiff(Ω)
we will working with from now on.
Let Diff(Ω) be the set of all diffeomorphisms in Ω then SDiff(Ω) :
= (cid:8) ψ ∈ Diff(Ω) , det( ∇ ψ )( x ) = 1 , ∀ x ∈ Ω (cid:9) , where Ω = ( − 1 , 1) d .
That is we need a transformation which satisfies two properties: 1) the function must be volume and orientation preserving, 2) the solution must stay in the domain ( − 1 , 1) d .
Consider the following ODE:
[Equation]
We impose two conditions on the velocity v :
[Equation]
where n is the outward normal at the boundary of Ω .
We define ϕ to be the solution at the final time ϕ ( x ) :
= X ( T, x ) .
Property (3) implies that det ∇ ϕ
= 1 , and property (4) ensures that ϕ does not escape Ω .
Any function in SDiff(Ω) can be written as a solution to (2) for d ≥ 3 [13], for d =some pathological cases can be constructed [14].
The incompressible property (3) and the boundary conditions (4) can be exactly implemented in the network in any dimension.
Note however that in order to get all the incompressible vector fields, we need to construct at least d ( d − 1) / 2 arbitrary scalar functions. See Appendix B for the practical construction of divergence free functions in high dimensions.
2.2 Euler’s geodesics
GP flows give a way to compute the Monge map for any trained NF architecture.
Many transformations can achieve this goal and the question of finding the best flow among all volume preserving transformations need to be considered.
Arnold’s theorem. In 1966, Arnold [11] showed that the flow described by Euler’s equations coincides with the geodesic flow on the manifold of volume preserving diffeomorphisms.
This theoretical result therefore gives the reason why regularizing our flows with Euler’s equations is a desirable property.
Mainly that Euler’s equations take the path with the lowest energy to reach the final configuration.
Consider the Euler equations:
[Equation]
where v := v ( t, x ) is the velocity field, p :
= p ( t, x ) the pressure and n := n ( x ) the outward normal at the boundary of Ω .
We introduce the energy E ( X ) :
= (cid:82) T 0 (cid:82) Ω 12 | ∂ t X ( t, x ) | 2 d x dt.
Now assume ϕ ∈ SDiff(Ω) , Arnold’s problem’s consists in finding the path X ( t, · ) t ∈ [0 ,T ] in SDiff(Ω) joining the identity to ϕ which minimizes E :
[Equation]
In other words (6) is the geodesic in SDiff(Ω) between Id and ϕ .
Theorem 2 (Arnold [11]) .
Assuming the existence of a solution to Arnold’s problem, X is solution to (6) if and only if v ( t, x ) :
= ∂ t X ( t, x ) satisfies Euler’s equations (5) .
In practice, we choose to penalize Euler’s equations along the trajectories of our GP flows as detailed in Appendix A. We then have two terms in our loss function: the negative log-likelihood and the term related to Euler’s equations.
An example of application: improving disentanglement preservation with optimal transport
Disentangle representations allow to encode the data in a latent space where change in one direction result in the change over one generative factor in the data.
Recently the construction of variational auto-encoders (VAE) [15] with disentangle latent space has received much attention [16, 17, 18, 19].
Applying a NF architecture to such latent space may be needed for various tasks such as density estimation, generative process or general interpolation.
The latter requires to preserve as much of the data structure as possible.
We therefore propose to experimentally study disentanglement preservation of NF with and without OT.
To do this we apply a NF architecture (FFJORD [20] in our case) to the
VAE’s latent distribution and consider the addition of a GP flow.
For the disentangle interpolation in the NF target (gaussian) space we consider the same directions which are present in the VAE latent space and are aligned with the axis.
This may be a naive approach since these directions may change depending on the source and target distribution but it seems enough here to already see an improvement.
Our experiments are run with the β -TCVAE architecture [18] and the latent space dimension is 10 . dSprites test case.
The dSprites dataset [12] is made of 64 × 64 images of 2D shapes procedurally generated from 5 ground truth independent latent factors.
These factors are shape, scale, rotation, x and y positions of a sprite.
Since the factors are known we can compute a quantitative evaluation of disentanglement and we choose here to consider the metric from [21] on the continuous factors (i.e. all the factors except the shape) for the three criteria: disentanglement, completeness and informativeness.
Table 1 shows that FFJORD destroy the latent structure and give the worst disentanglement, completeness and informativeness scores.
Adding GP flows allow to recover the same disentanglement score as the initial latent space and get values closer both for completeness and informativeness.
Note that the disentanglement score is slightly better with FFJORD+GP than the initial one however this is probably only due to some approximation in the metric used here since there is no reason that GP flows improved the disentanglement compared to the initial latent space.
On Table 2 the
OT costs are compared and as expected GP flows allow to reduce the OT cost without changing the loss.
Interestingly GP flows with no additional regularization do not converge completely to the
Monge map because particles get out of the domain at some point making it impossible to continue the training process.
We conjecture that this may be due to non-smooth trajectories of our GP flows and a regularization is therefore needed.
As shown on Table 2 adding Euler regularization fixed this issue and allows to further reduce the OT cost.
To illustrate the preservation of disentanglement some interpolations are also presented in Figures 2 and 3 in Appendix D. As we can see on the initial latent space picture, since the dimensions are sorted with respect to their KL divergence only the first dimensions carry information and therefore each of the first 5 lines correspond to a generative factor.
On the contrary when interpolating on the last dimensions the image stays the same.
This structure is lost when mapping the latent space to a gaussian with the FFJORD architecture.
When adding GP flows we recover this structure and the interpolation better match the initial latent one. [Table] [Table]
A Penalization of Euler’s equations in high dimensions
Numerical schemes developed to efficiently solve the Euler equations [22, 23] (mainly for fluid mechanics problems, i e. for dimensions up to 3 ) scale badly when the dimension increases.
In this work, the solution to Euler’s equations is interpreted as the geodesic to reach the solution of the OT problem and the dimension can be arbitrary large.
Therefore we approach the equation (5) through a penalization procedure which can be carried out in any dimensions.
As explained before the second and third equations in (5) are satisfied by construction in the network.
Our remaining goal is to constrain the network to be a smooth solution to ∂ t v + ( v · ∇ ) v = −∇ p .
The left hand side can therefore be written as the gradient of a scalar function and we note that if a vector w t, x ∈ R
d satisfies w t, x = ∇ p ( t, x ) , then its Jacobian is symmetric ∇ w t, x = ( ∇ w t, x ) T .
In order to solve the first equation in (5), we propose to penalize the non-symmetric part of the Jacobian for the total derivative of v .
Since a Jacobian-vector product can be efficiently evaluated in high dimensions (unlike the calculation of the full Jacobian which is computationally expensive), we do not calculate directly the Jacobian and use instead the following property of symmetric matrices: if
M is symmetric then y T M z − z T M y = 0 , ∀ y , z ∈ R d .
The idea is to sample random vectors y , z during the training and to penalize this term for the total derivative, that is to minimize:
[Equation]
with w t, x = ∂ t v + ( v · ∇ ) v .
In practice, we do not compute the full time integral in (7) as it would be computationally too expensive but calculate the penalization only at our time steps discretization.
Approximation of the total derivative.
To reduce the computational burden, we do not calculate exactly the total derivative w t, x but use an approximation of its Lagrangian formulation instead.
More precisely, consider the variable X ( t, x ) from (2) that is the position of a particle at time t with initial position x .
We recall the equality DDt v ( t, X ( t, x )) =
∂ t v ( t, X ( t, x ))
+ ( v ( t, X ( t, x )) · ∇ ) v ( t, X ( t, x )) and therefore choose to approximate the right hand side by using a first order Taylor expansion of D v /Dt :
[Equation]
In practice, ∆ t is set to 2 √ ε where ε is the machine precision.
This approximation can be easily computed since it requires only the evaluation of the velocity at two positions of a particle.
B Construction of incompressible vector fields in high dimensions
We focus on the vector fields satisfying (3) for arbitrary large dimensions.
Property (4) can then be incorporated with very little additional work.
All the proofs are given in the full version of the paper.
Proposition 2.
Consider an arbitrary vector field v : R d → R d .
Then ∇ · v =
if and only if there exists smooth scalar functions ψ ij : R d → R , with ψ ij = − ψ ji such that
[Equation]
where v = ( v 1 , ..., v d ) .
To impose the boundary conditions (4) one can simply multiply each ψ ij by ( x 2 i − 1)( x 2 j − 1) .
Lemma 1.
Let Ω =
[ − 1 , 1] d and consider the functions ψ ij ( x ) =
( x 2 i − 1)( x 2 j − 1) (cid:102) ψ ij ( x ) where (cid:101) ψ ij ( x ) : R d → R are arbitrary scalar functions satisfying (cid:101) ψ ij = − (cid:101) ψ ji .
Then the function v defined in
Proposition 2 satisfies ∇ · v = 0 and v · n = 0 on ∂ Ω .
The goal here is to have a GPU-friendly construction of the incompressible vector fields given in
Proposition 2.
In the following we consider stationary divergence free functions but the time variable can be added with no additional work simply by considering functions in R d +1 instead of R d (the gradients are still taken only on the space variables though). Notations. Regarding the notations we will use the operator diag for two distinct cases: 1) When w is a vector diag( w ) denotes the diagonal matrix obtained from the vector w . 2) When W is a matrix diag( W ) denotes the vector obtained from the diagonal of W .
The operation · denotes the scalar product between two vectors. Finally we have adopted the convention that when a scalar multiplies a vector it multiplies each of its component. Practical construction.
Let u n : R d → R d − n . We construct a divergence free function with the functions ψ ij defined as [Equation] To construct this divergence free function we define the matrix ( ∇ u ( x )) n ∈ R d × d and the vector 1 n ∈ R d as
[Equation]
Lemma 2.
Let n ∈ N , n ≤ d − 2 and consider the function u n : R d → R d − n .
Then the vector field v n : R d → R d defined as
[Equation]
is divergence free.
To construct the functions (12) we need 1) to compute the product between the Jacobian of a vector valued function and a constant vector 2) sum the diagonal elements of the Jacobian matrix.
Both of these operations can be done efficiently on GPU.
Note that in order to satisfy the boundary conditions v n · n = 0 one can modify the equation (12) as in Lemma 1 to obtain
[Equation]
where M nij = u ni − u nj , if i, j ≥ n + 1 and M nij = 0 otherwise.
{+It is possible to recover all the incompressible functions from Proposition 2 by adding the blocks v 
+}
+ v 1 + ... + v d − 2 .
Proposition 3.
Let v be a divergence free function in R d .
Then there exists d − 1 functions v 0 , ..., v d − 2 constructed as in Lemma 2 such that
[Equation]
The attentive reader would have noticed that with the vector functions u n , n = 0 , ..., d − 2 we have a total of ( d + 2)( d − 1) / 2 independent scalar functions while Proposition 2 only requires the construction of d ( d − 1) / 2 scalar functions leaving d − 1 additional functions which are not strictly needed to obtain the divergence free vectors. This is due to the vectorized constructions (12)-(13) which allow a fast evaluation of the divergence free functions on GPU. Having d − 1 additional scalar functions in return is not a big issue since the general order remains O ( d 2 ) . Also note that the practical implementation of the equations (10)-(11) requires to find a pythonic way to efficiently pad a group of matrices with different dimensions.
We have not yet find such way and therefore have simply chosen in our applications to construct
d − 1 vector valued functions u n ∈ R d , n = 0 , ..., d − 2 and fill the appropriate dimensions in (10)-(11) with 0 .
Again even if not optimal this is not a big issue as it multiplies the number of independent scalar functions by a factor 2 , but the general order remains O ( d 2 ) .
In practice, we have written the functions u n as the output of a big function u : R d → R ( d − 1) ×
d allowing to evaluate all the functions u n in a single pass.
The vector u is written as the composition of linear functions with some simple non linearity
[Equation]
where M i are rectangular matrices, b i a vector field and typically we have taken σ = tanh .
One big advantage of the formulation (14) is that the Jacobian of u (and therefore of all the functions u n ) can be computed analytically
[Equation]
for i = 1 , ..., n − 1 .
The formulation (15) therefore allows a fast evaluation of the term (12) in particular when summing the diagonal elements of the Jacobian.
In our experiments we have noticed that the analytical formulation of the Jacobian (15) was faster than using torch.autograd .
C Experiment details
dSprites datasets.
The VAE architecture used in the experiments is taken from the github repository of Yann Dubois 2 .
The parameters used for the GP flows are given in Table 3.
For all test cases we consider a GP flow with 15 time steps and three intermediate layers of 50 parameters with a
Runge-Kutta 4 discretization at each layer.
For the Euler penalization we take an initial parameter λ = 5 × 10 − 5 which is then divided 10 times periodically by a factor 2 during the training.
D Interpolation examples
[Table]
[Table]
[Table]