Diamond Maps: Efficient Reward Alignment via Stochastic Flow Maps

Flow and diffusion models produce high-quality samples, but adapting them to user preferences or constraints post-training remains costly and brittle, a challenge commonly called reward alignment. We argue that efficient reward alignment should be a property of the generative model itself, not an afterthought, and redesign the model for adaptability. We propose Diamond Maps, a stochastic flow-map model that enables efficient and accurate alignment to arbitrary rewards at inference time. Diamond Maps amortize many simulation steps into a single-step sampler, like flow maps, while preserving the stochasticity required for optimal reward adaptation. This design makes search, Sequential Monte Carlo, and guidance scalable by enabling efficient and consistent estimation of the value function. Our experiments show that Diamond Maps can be learned efficiently via distillation from GLASS Flows, achieve stronger reward-alignment performance, and scale better than existing alignment methods. Overall, our results point toward a practical route to generative models that can be rapidly adapted to arbitrary preferences and constraints at inference time.

We introduce Diamond Maps, a new kind of generative model that lets us steer pretrained diffusion and flow models toward high-reward outputs at inference time, without expensive retraining. Modern fast samplers known as flow maps can jump from noise to a clean sample in just a few steps, but they are deterministic: given the same starting point, they always produce the same output. This is a problem for reward alignment, because deciding which direction is "promising" requires looking ahead and averaging over the many possible outcomes a noisy state could lead to. Diamond Maps fix this by making flow maps stochastic, supplying exactly the kind of exploration needed to estimate the value of partially generated samples cheaply and accurately. Building on flow maps and GLASS Flows as our underlying machinery, we develop two complementary designs. The first, Posterior Diamond Maps, are one-step samplers distilled to draw from the posterior over clean images, giving a simple and consistent way to estimate the value function and its gradient. The second, Weighted Diamond Maps turn ordinary off-the-shelf flow maps into reliable value estimators with no retraining and direct reuse of existing distilled models through a lightweight renoising procedure applied purely at inference time. Finally, we show empirically that Diamond Maps achieve fast, high-quality reward alignment that scales better than strong baselines such as Best-of-N and regular Guidance, across tasks ranging from inverse problems and prompt alignment to high-resolution text-to-image generation with models like FLUX and SANA-Sprint.