Stable Training of the Lagrangian Flow-Map Objective

Develop a stable training procedure for the Lagrangian objective that enforces consistency between the time derivative of the discrete flow map X_u(x_s, s, t) and the average-velocity field u_θ evaluated at the mapped state and target time t within the FALCON few-step flow framework for Boltzmann Generators, so that this objective can be optimized reliably in the authors’ setting.

Background

To enable few-step generation with tractable likelihoods, FALCON trains flow maps using average-velocity objectives, and the authors discuss several formulations that have been proposed in recent work. Among these, a Lagrangian objective seeks to match the time derivative of the transported state to the velocity field evaluated at the mapped point and target time, providing a potentially strong consistency condition.

In their experiments, the authors attempted to use this Lagrangian objective but found optimization to be unstable in their Boltzmann Generator setting. A method that makes this training stable would expand the set of viable objectives for FALCON and related few-step, invertible flow models.