TO-FLOW: Efficient Continuous Normalizing Flows with Temporal Optimization adjoint with Moving Speed

TL;DR

This paper introduces TO-FLOW, a method that optimizes the evolution time in continuous normalizing flows to accelerate training while maintaining performance, using a novel temporal optimization approach with neural ODEs.

Contribution

It proposes a temporal optimization technique for neural ODEs in CNFs, improving training efficiency through evolutionary time optimization and temporal regularization.

Findings

Significantly faster training times compared to baseline models.

Maintains comparable performance with improved training efficiency.

Compatible with existing regularization methods.

Abstract

Continuous normalizing flows (CNFs) construct invertible mappings between an arbitrary complex distribution and an isotropic Gaussian distribution using Neural Ordinary Differential Equations (neural ODEs). It has not been tractable on large datasets due to the incremental complexity of the neural ODE training. Optimal Transport theory has been applied to regularize the dynamics of the ODE to speed up training in recent works. In this paper, a temporal optimization is proposed by optimizing the evolutionary time for forward propagation of the neural ODE training. In this appoach, we optimize the network weights of the CNF alternately with evolutionary time by coordinate descent. Further with temporal regularization, stability of the evolution is ensured. This approach can be used in conjunction with the original regularization approach. We have experimentally demonstrated that the proposed approach can significantly accelerate training without sacrifying performance over baseline models.