Normalizing flow is a class of deep generative models for efficient sampling and density estimation. In practice, the flow often appears as a chain of invertible neural network blocks; to facilitate training, existing works have regularized flow trajectories and designed special network architectures. The current paper develops a neural ODE flow network inspired by the Jordan-Kinderleherer-Otto (JKO) scheme, which allows efficient block-wise training of the residual blocks without sampling SDE trajectories or inner loops of score matching or variational learning. As the JKO scheme unfolds the dynamic of gradient flow, the proposed model naturally stacks residual network blocks one by one, reducing the memory load and difficulty in performing end-to-end deep flow network training. We also develop adaptive time reparameterization of the flow network with a progressive refinement of the trajectory in probability space, which improves the model training efficiency and accuracy in practice. Using numerical experiments with synthetic and real data, we show that the proposed JKO-iFlow model achieves similar or better performance in generating new samples compared with the existing flow and diffusion models at a significantly reduced computational and memory cost.

翻译：归一化流是一类用于高效采样和密度估计的深度生成模型。实际应用中，该流通常表现为可逆神经网络模块的链式结构；为便于训练，现有研究已对流轨迹进行正则化并设计特殊网络架构。本文受Jordan-Kinderleherer-Otto (JKO) 格式启发，开发了一种神经ODE流网络，能够对残差模块实现高效的逐块训练，而无需采样SDE轨迹或进行分数匹配与变分学习的内循环。由于JKO格式揭示了梯度流的动力学特性，所提模型自然地逐层堆叠残差网络模块，从而降低端到端深度流网络训练中的内存负载与难度。我们还开发了流网络的自适应时间重参数化技术，通过概率空间中轨迹的渐进细化来提升实际训练效率与模型精度。基于合成数据与真实数据的数值实验表明，与现有流模型和扩散模型相比，所提出的JKO-iFlow模型在生成新样本时能以显著降低的计算与内存开销实现相当甚至更优的性能。