How to build a consistency model: Learning flow maps via self-distillation

TL;DR

This paper introduces a unified, systematic framework for training flow map models in generative modeling, leveraging self-distillation to improve efficiency and stability, and presents novel Lagrangian methods that outperform existing schemes.

Contribution

It provides a unified algorithmic framework for learning flow maps via self-distillation, introducing Lagrangian methods that enhance stability and performance.

Findings

Lagrangian methods outperform Eulerian and Progressive schemes in stability and performance.

The framework unifies and extends existing distillation and training schemes.

Code implementation is publicly available for reproducibility.

Abstract

Flow-based generative models achieve state-of-the-art sample quality, but require the expensive solution of a differential equation at inference time. Flow map models, commonly known as consistency models, encompass many recent efforts to improve inference-time efficiency by learning the solution operator of this differential equation. Yet despite their promise, these models lack a unified description that clearly explains how to learn them efficiently in practice. Here, building on the methodology proposed in Boffi et. al. (2024), we present a systematic algorithmic framework for directly learning the flow map associated with a flow or diffusion model. By exploiting a relationship between the velocity field underlying a continuous-time flow and the instantaneous rate of change of the flow map, we show how to convert any distillation scheme into a direct training algorithm via self-distillation, eliminating the need for pre-trained teachers. We introduce three algorithmic families based on different mathematical characterizations of the flow map: Eulerian, Lagrangian, and Progressive methods, which we show encompass and extend all known distillation and direct training schemes for consistency models. We find that the novel class of Lagrangian methods, which avoid both spatial derivatives and bootstrapping from small steps by design, achieve significantly more stable training and higher performance than more standard Eulerian and Progressive schemes. Our methodology unifies existing training schemes under a single common framework and reveals new design principles for accelerated generative modeling. Associated code is available at https://github.com/nmboffi/flow-maps.