Diamond Maps: Efficient Reward Alignment via Stochastic Flow Maps

TL;DR

Diamond Maps are a new stochastic flow model that allows for rapid, accurate reward alignment at inference time, improving adaptability of generative models to user preferences.

Contribution

We introduce Diamond Maps, a novel stochastic flow model that enables efficient reward alignment and scalable adaptation of generative models at inference time.

Findings

Diamond Maps can be learned efficiently via distillation from GLASS Flows.

They achieve stronger reward alignment performance than existing methods.

Diamond Maps scale better and enable rapid adaptation to arbitrary preferences.

Abstract

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", stochastic flow map models that enable 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 alignment. 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 methods. Our results point toward a practical route to generative models that can be rapidly adapted to arbitrary preferences and constraints at inference time.