# DRAGON: Distributional Rewards Optimize Diffusion Generative Models

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

## Key facts

- Cited by 4 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 58 of its 78 references that are in the map.
- First posted on arXiv 21 Apr 2025.

## Placed in

- Diffusion models / [Diffusion and flow policies for robotics](https://mlsanity.com/diffusion/diffusion-and-flow-policies-for-robotics/)
- LLM alignment / RLHF / preference optimisation / [Diffusion and multimodal model alignment](https://mlsanity.com/alignment/diffusion-and-multimodal-model-alignment/)

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [Denoising Diffusion Implicit Models](https://mlsanity.com/paper/2010.02502/) (2020)
- [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://mlsanity.com/paper/2402.03300/) (2024)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Direct Preference Optimization: Your Language Model is Secretly a Reward Model](https://mlsanity.com/paper/2305.18290/) (2023)
- [Deep Unsupervised Learning using Nonequilibrium Thermodynamics](https://mlsanity.com/paper/1503.03585/) (2015)
- [DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning](https://mlsanity.com/paper/2501.12948/) (2025)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [Deep reinforcement learning from human preferences](https://mlsanity.com/paper/1706.03741/) (2017)
- [Elucidating the Design Space of Diffusion-Based Generative Models](https://mlsanity.com/paper/2206.00364/) (2022)
- [Scaling Instruction-Finetuned Language Models](https://mlsanity.com/paper/2210.11416/) (2022)
- [Improved Techniques for Training GANs](https://mlsanity.com/paper/1606.03498/) (2016)
- [Video Diffusion Models](https://mlsanity.com/paper/2204.03458/) (2022)
- [DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps](https://mlsanity.com/paper/2206.00927/) (2022)
- [Mixtral of Experts](https://mlsanity.com/paper/2401.04088/) (2024)
- [Consistency Models](https://mlsanity.com/paper/2303.01469/) (2023)
- [Training Diffusion Models with Reinforcement Learning](https://mlsanity.com/paper/2305.13301/) (2023)
- [Pick-a-Pic: An Open Dataset of User Preferences for Text-to-Image Generation](https://mlsanity.com/paper/2305.01569/) (2023)
- [Diffusion Model Alignment Using Direct Preference Optimization](https://mlsanity.com/paper/2311.12908/) (2023)
- [SoundStream: An End-to-End Neural Audio Codec](https://mlsanity.com/paper/2107.03312/) (2021)
- [Large-scale Contrastive Language-Audio Pretraining with Feature Fusion and Keyword-to-Caption Augmentation](https://mlsanity.com/paper/2211.06687/) (2022)
- [High Fidelity Neural Audio Compression](https://mlsanity.com/paper/2210.13438/) (2022)
- [AudioLM: a Language Modeling Approach to Audio Generation](https://mlsanity.com/paper/2209.03143/) (2022)
- [AudioLDM: Text-to-Audio Generation with Latent Diffusion Models](https://mlsanity.com/paper/2301.12503/) (2023)
- [Simple and Controllable Music Generation](https://mlsanity.com/paper/2306.05284/) (2023)

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

- [CMI-RewardBench: Evaluating Music Reward Models with Compositional Multimodal Instruction](https://mlsanity.com/paper/2603.00610/) (2026)
- [V2M-Zero: Zero-Pair Time-Aligned Video-to-Music Generation](https://mlsanity.com/paper/2603.11042/) (2026)
- [Vendi Novelty Scores for Out-of-Distribution Detection](https://mlsanity.com/paper/2602.10062/) (2026)
- [Explore or Converge? Stage-Guided Per-Step Optimization for Diffusion Models](https://mlsanity.com/paper/2608.06768/) (2026)

## Abstract (arXiv)

We present Distributional RewArds for Generative OptimizatioN (DRAGON), a versatile framework for fine-tuning media generation models towards a desired outcome. Compared with traditional reinforcement learning with human feedback (RLHF) or pairwise preference approaches such as direct preference optimization (DPO), DRAGON is more flexible. It can optimize reward functions that evaluate either individual examples or distributions of them, making it compatible with a broad spectrum of instance-wise, instance-to-distribution, and distribution-to-distribution rewards. Leveraging this versatility, we construct novel reward functions by selecting an encoder and a set of reference examples to create an exemplar distribution. When cross-modal encoders such as CLAP are used, the reference may be of a different modality (text versus audio). Then, DRAGON gathers online and on-policy generations, scores them with the reward function to construct a positive demonstration set and a negative set, and leverages the contrast between the two finite sets to approximate distributional reward optimization. For evaluation, we fine-tune an audio-domain text-to-music diffusion model with 20 reward functions, including a custom music aesthetics model, CLAP score, Vendi diversity, and Frechet audio distance (FAD). We further compare instance-wise (per-song) and full-dataset FAD settings while ablating multiple FAD encoders and reference sets. Over all 20 target rewards, DRAGON achieves an 81.45% average win rate. Moreover, reward functions based on exemplar sets enhance generations and are comparable to model-based rewards. With an appropriate exemplar set, DRAGON achieves a 60.95% human-voted music quality win rate without training on human preference annotations. DRAGON is a new approach to designing and optimizing reward functions for improving human-perceived quality. Demos at https://ml-dragon.github.io/web