How to Guide Your Flow: Few-Step Alignment via Flow Map Reward Guidance

TL;DR

This paper introduces Flow Map Reward Guidance (FMRG), a training-free, single-trajectory method for guidance in generative modeling that achieves high-quality results with significantly fewer neural function evaluations.

Contribution

It reformulates guidance as a deterministic optimal control problem and leverages the flow map for efficient, single-trajectory guidance in generative models.

Findings

FMRG matches or surpasses baselines in text-to-image tasks.

FMRG requires as few as 3 NFEs, greatly reducing computation.

Provides an order-of-magnitude speedup over prior methods.

Abstract

In generative modeling, we often wish to produce samples that maximize a user-specified reward such as aesthetic quality or alignment with human preferences, a problem known as \textit{guidance}. Despite their widespread use, existing guidance methods either require expensive multi-particle, many-step schemes or rely on poorly understood approximations. We reformulate guidance as a \textit{deterministic optimal control problem}, yielding a hierarchy of algorithms that subsumes existing approaches at the coarsest level. We show that the \textit{flow map}, an object of significant recent interest for its role in fast inference, arises naturally in the optimal solution. Based on this observation, we propose \textbf{Flow Map Reward Guidance (FMRG)}: a training-free, \textit{single-trajectory} framework that uses the flow map to both integrate and guide the flow. At text-to-image scale, FMRG matches or surpasses baselines across inverse problems and reward-guided generation with \textbf{as few as 3 NFEs}, giving at least an order-of-magnitude speedup in comparison to prior state of the art.