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网络在显著降低计算与内存开销的前提下，取得了与现有流模型和扩散模型相当的性能表现。