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.

翻译：归一化流（NF）是一种强大的基于似然的生成模型，能够在表达能力和可追踪性之间权衡来建模复杂密度。现在已经建立了一个成熟的研究方向，利用最优输运（OT）来寻找Monge映射，即在源分布和目标分布之间具有最小代价的模型。本文介绍了一种基于Brenier的极极分解定理的方法，将任何训练过的NF转化为更具OT效率的版本，而不改变最终的密度。我们通过学习源（高斯）分布的重排列，最小化源分布和最终密度之间的OT代价来实现这一点。我们通过欧拉方程约束的路径，进一步限制了导致估计的Monge映射的空间保体變微分同构的测地线。所提出的方法可以在不影响模型性能的情况下，为几个现有模型提供平滑流和降低OT成本。