Normalizing flow is a class of deep generative models for efficient sampling and likelihood estimation, which achieves attractive performance, particularly in high dimensions. The flow is often implemented using a sequence of invertible residual blocks. Existing works adopt special network architectures and regularization of flow trajectories. In this paper, we develop a neural ODE flow network called JKO-iFlow, inspired by the Jordan-Kinderleherer-Otto (JKO) scheme, which unfolds the discrete-time dynamic of the Wasserstein gradient flow. The proposed method stacks residual blocks one after another, allowing efficient block-wise training of the residual blocks, avoiding sampling SDE trajectories and score matching or variational learning, thus reducing the memory load and difficulty in end-to-end training. We also develop adaptive time reparameterization of the flow network with a progressive refinement of the induced trajectory in probability space to improve the model accuracy further. Experiments with synthetic and real data show that the proposed JKO-iFlow network achieves competitive performance compared with existing flow and diffusion models at a significantly reduced computational and memory cost.

翻译：正则化流是一类用于高效采样和似然估计的深度生成模型，在高维场景中展现出优异性能。该流通常通过序列化的可逆残差块实现。现有研究采用特殊网络架构和流路径正则化方法。本文受Jordan-Kinderleherer-Otto(JKO)格式启发，提出一种名为JKO-iFlow的神经ODE流网络，该网络解构Wasserstein梯度流的离散时间动力学。所提方法通过逐块堆叠残差块实现高效分块训练，避免采样SDE轨迹与得分匹配或变分学习，从而降低内存负担和端到端训练难度。我们还开发了流网络的自适应时间重参数化方法，通过概率空间中诱导轨迹的渐进细化进一步提升模型精度。合成数据与真实数据实验表明，与现有流模型和扩散模型相比，所提JKO-iFlow网络在显著降低计算和内存成本的同时实现了具有竞争力的性能。