# Diamond Maps: Efficient Reward Alignment via Stochastic Flow Maps

arXiv: https://arxiv.org/abs/2602.05993 · HTML: https://mlsanity.com/paper/2602.05993/

## Key facts

- Cited by 29 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion and flow policies for robotics, one of 14 directions in Diffusion models.
- Builds on 69 of its 84 references that are in the map.
- First posted on arXiv 5 Feb 2026.

## Placed in

- Diffusion models / [Diffusion and flow policies for robotics](https://mlsanity.com/diffusion/diffusion-and-flow-policies-for-robotics/)

## Builds on (most cited first, 30 of 69)

- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Denoising Diffusion Implicit Models](https://mlsanity.com/paper/2010.02502/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Deep Unsupervised Learning using Nonequilibrium Thermodynamics](https://mlsanity.com/paper/1503.03585/) (2015)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Scaling Rectified Flow Transformers for High-Resolution Image Synthesis](https://mlsanity.com/paper/2403.03206/) (2024)
- [Generative Modeling by Estimating Gradients of the Data Distribution](https://mlsanity.com/paper/1907.05600/) (2019)
- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [Deep Learning Face Attributes in the Wild](https://mlsanity.com/paper/1411.7766/) (2014)
- [CLIPScore: A Reference-free Evaluation Metric for Image Captioning](https://mlsanity.com/paper/2104.08718/) (2021)
- [Consistency Models](https://mlsanity.com/paper/2303.01469/) (2023)
- [Diffusion Posterior Sampling for General Noisy Inverse Problems](https://mlsanity.com/paper/2209.14687/) (2022)
- [ImageReward: Learning and Evaluating Human Preferences for Text-to-Image Generation](https://mlsanity.com/paper/2304.05977/) (2023)
- [Pick-a-Pic: An Open Dataset of User Preferences for Text-to-Image Generation](https://mlsanity.com/paper/2305.01569/) (2023)
- [GenEval: An Object-Focused Framework for Evaluating Text-to-Image Alignment](https://mlsanity.com/paper/2310.11513/) (2023)
- [Stochastic Interpolants: A Unifying Framework for Flows and Diffusions](https://mlsanity.com/paper/2303.08797/) (2023)
- [Human Preference Score v2: A Solid Benchmark for Evaluating Human Preferences of Text-to-Image Synthesis](https://mlsanity.com/paper/2306.09341/) (2023)
- [Diffusion Model Alignment Using Direct Preference Optimization](https://mlsanity.com/paper/2311.12908/) (2023)
- [Flow-GRPO: Training Flow Matching Models via Online RL](https://mlsanity.com/paper/2505.05470/) (2025)
- [Mean Flows for One-step Generative Modeling](https://mlsanity.com/paper/2505.13447/) (2025)
- [DPOK: Reinforcement Learning for Fine-tuning Text-to-Image Diffusion Models](https://mlsanity.com/paper/2305.16381/) (2023)
- [Improved Techniques for Training Consistency Models](https://mlsanity.com/paper/2310.14189/) (2023)
- [Universal Guidance for Diffusion Models](https://mlsanity.com/paper/2302.07121/) (2023)
- [DanceGRPO: Unleashing GRPO on Visual Generation](https://mlsanity.com/paper/2505.07818/) (2025)
- [Analyzing and Improving the Training Dynamics of Diffusion Models](https://mlsanity.com/paper/2312.02696/) (2023)
- [One Step Diffusion via Shortcut Models](https://mlsanity.com/paper/2410.12557/) (2024)
- [Flow Matching Guide and Code](https://mlsanity.com/paper/2412.06264/) (2024)
- [FreeDoM: Training-Free Energy-Guided Conditional Diffusion Model](https://mlsanity.com/paper/2303.09833/) (2023)
- [Inference-Time Scaling for Diffusion Models beyond Scaling Denoising Steps](https://mlsanity.com/paper/2501.09732/) (2025)

## Built on by (most cited first, 29 of 29)

- [Reinforce Adjoint Matching: Scaling RL Post-Training of Diffusion and Flow-Matching Models](https://mlsanity.com/paper/2605.10759/) (2026)
- [Scaling Categorical Flow Maps](https://mlsanity.com/paper/2605.07820/) (2026)
- [Reward Score Matching: Unifying Reward-based Fine-tuning for Flow and Diffusion Models](https://mlsanity.com/paper/2604.17415/) (2026)
- [How to Guide Your Flow: Few-Step Alignment via Flow Map Reward Guidance](https://mlsanity.com/paper/2604.27147/) (2026)
- [Strong Stochastic Flow Maps](https://mlsanity.com/paper/2606.01086/) (2026)
- [Aligning Few-Step Generative Models by Amortizing Sample-based Variational Inference](https://mlsanity.com/paper/2605.26552/) (2026)
- [Are we really tilting? The mechanics of reward guidance in flow and diffusion models](https://mlsanity.com/paper/2606.02884/) (2026)
- [QPILOTS: Efficient Test-Time Q-Steering for Flow Policies](https://mlsanity.com/paper/2606.14801/) (2026)
- [Flow-Map GRPO: Reinforcement Learning for Few-Step Flow-Map Generators via Anchored Stochastic Composition](https://mlsanity.com/paper/2607.00535/) (2026)
- [Advances in GRPO for Generation Models: A Survey](https://mlsanity.com/paper/2603.06623/) (2026)
- [Accelerating Diffusion Sampling via Speculative Draft Trees](https://mlsanity.com/paper/2609.17691/) (2026)
- [Aligning Flow Map Policies with Optimal Q-Guidance](https://mlsanity.com/paper/2605.12416/) (2026)
- [StressDream: Steering Video World Models for Robust Policy Evaluation and Improvement](https://mlsanity.com/paper/2606.00267/) (2026)
- [Few-step Cofolding with All-Atom Flow Maps](https://mlsanity.com/paper/2606.08375/) (2026)
- [BFMT: Enhancing Search Capabilities of Tree Sampler via Bootstrap Flow-Map Tree](https://mlsanity.com/paper/2607.02915/) (2026)
- [Expanding Flow Maps](https://mlsanity.com/paper/2607.21585/) (2026)
- [VGAS: Variance-Reduced Guidance and Adaptive Selection for Training-Free Reward Alignment in Discrete Diffusion](https://mlsanity.com/paper/2608.26585/) (2026)
- [OMatG-flash: An All-Atom Flow Map with Reinforce Adjoint Matching for Scalable Materials Discovery](https://mlsanity.com/paper/2609.26402/) (2026)
- [WTF?! Simulation-Free Reinforcement Learning with Wasserstein-Tilted Flow Maps](https://mlsanity.com/paper/2609.27033/) (2026)
- [FluxLite: Inference-Time Proposal Control for Discrete Diffusion Models](https://mlsanity.com/paper/2609.35947/) (2026)
- [Stochastic Transition-Map Distillation for Fast Probabilistic Inference](https://mlsanity.com/paper/2605.07661/) (2026)
- [Follow the Mean: Reference-Guided Flow Matching](https://mlsanity.com/paper/2605.10302/) (2026)
- [Sequentially-Controlled Interactive Multi-Particle Flow-Maps for Online Feedback-Driven Search](https://mlsanity.com/paper/2607.01144/) (2026)
- [A Mean-Field Framework for Inference-Time Distributional Control of Diffusion Models](https://mlsanity.com/paper/2608.08770/) (2026)
- [Interacting particle guidance for sampling reward-tilted generative priors](https://mlsanity.com/paper/2609.37227/) (2026)
- [Steepest Guidance: A Practical and Principled Approach to Inference-Time Alignment of Flow and Diffusion-based Models](https://mlsanity.com/paper/2609.39091/) (2026)
- [Langevin Flow Maps: Efficient Molecular Dynamics and Transition Path Sampling](https://mlsanity.com/paper/2610.05998/) (2026)
- [Noise, Denoise, Correct: MCMC Posterior Sampling with Diffusion Priors in Three Steps](https://mlsanity.com/paper/2610.09407/) (2026)
- [On the Design Space of Continuous Diffusion Language Models](https://mlsanity.com/paper/2610.13834/) (2026)

## Abstract (arXiv)

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.