How to Guide Your Flow: Few-Step Alignment via Flow Map Reward Guidance
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 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 deterministic optimal control problem, yielding a hierarchy of algorithms that subsumes existing approaches at the coarsest level. We show that the 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 Flow Map Reward Guidance (FMRG): a training-free, 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 as few as 3 NFEs, giving at least an order-of-magnitude speedup in comparison to prior state of the art.
Lay Summary
Generative AI models can now create high-quality images in only a few model calls, making them much faster than earlier systems. However, steering these models toward a desired goal—such as reconstructing an image from degraded measurements, improving visual quality, or better matching a prompt—often still requires many additional computation steps. We introduce Flow Map Reward Guidance (FMRG), a method for guiding fast flow-based generative models without retraining them. FMRG uses the model’s own learned path from noise to image to look ahead at the final output, measure how well it satisfies the desired goal, and steer the current sample accordingly. This allows guidance to happen along a single trajectory, rather than requiring many competing samples or expensive search. Across image reconstruction tasks and reward-guided text-to-image generation, FMRG produces high-quality results using far fewer model calls than prior methods. More broadly, our work suggests that guidance for generative AI should be designed for the same fast, few-step setting that modern flow-based generators are moving toward.