# Avoiding mode collapse in diffusion models fine-tuned with reinforcement learning

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

## Key facts

- Cited by 12 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 25 of its 27 references that are in the map.
- First posted on arXiv 10 Oct 2024.

## Placed in

- Diffusion models / [Diffusion and flow policies for robotics](https://mlsanity.com/diffusion/diffusion-and-flow-policies-for-robotics/)
- Deep reinforcement learning methods / [RL with diffusion and flow models](https://mlsanity.com/deeprl/rl-with-diffusion-and-flow-models/)

## Builds on (most cited first, 25 of 25)

- [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)
- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [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)
- [Zero-Shot Text-to-Image Generation](https://mlsanity.com/paper/2102.12092/) (2021)
- [Generative Modeling by Estimating Gradients of the Data Distribution](https://mlsanity.com/paper/1907.05600/) (2019)
- [Progressive Growing of GANs for Improved Quality, Stability, and Variation](https://mlsanity.com/paper/1710.10196/) (2017)
- [Improved Techniques for Training GANs](https://mlsanity.com/paper/1606.03498/) (2016)
- [Trust Region Policy Optimization](https://mlsanity.com/paper/1502.05477/) (2015)
- [Imagen Video: High Definition Video Generation with Diffusion Models](https://mlsanity.com/paper/2210.02303/) (2022)
- [Make-A-Video: Text-to-Video Generation without Text-Video Data](https://mlsanity.com/paper/2209.14792/) (2022)
- [Training Diffusion Models with Reinforcement Learning](https://mlsanity.com/paper/2305.13301/) (2023)
- [Solving Inverse Problems in Medical Imaging with Score-Based Generative Models](https://mlsanity.com/paper/2111.08005/) (2021)
- [DPOK: Reinforcement Learning for Fine-tuning Text-to-Image Diffusion Models](https://mlsanity.com/paper/2305.16381/) (2023)
- [Understanding Diffusion Models: A Unified Perspective](https://mlsanity.com/paper/2208.11970/) (2022)
- [The Vendi Score: A Diversity Evaluation Metric for Machine Learning](https://mlsanity.com/paper/2210.02410/) (2022)
- [Torsional Diffusion for Molecular Conformer Generation](https://mlsanity.com/paper/2206.01729/) (2022)
- [Efficient Diffusion Training via Min-SNR Weighting Strategy](https://mlsanity.com/paper/2303.09556/) (2023)
- [A Phase Transition in Diffusion Models Reveals the Hierarchical Nature of Data](https://mlsanity.com/paper/2402.16991/) (2024)
- [PRDP: Proximal Reward Difference Prediction for Large-Scale Reward Finetuning of Diffusion Models](https://mlsanity.com/paper/2402.08714/) (2024)
- [Towards Faster Training of Diffusion Models: An Inspiration of A Consistency Phenomenon](https://mlsanity.com/paper/2404.07946/) (2024)

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

- [Probing the Latent Hierarchical Structure of Data via Diffusion Models](https://mlsanity.com/paper/2410.13770/) (2024)
- [Beyond the Dirac Delta: Mitigating Diversity Collapse in Reinforcement Fine-Tuning for Versatile Image Generation](https://mlsanity.com/paper/2601.12401/) (2026)
- [ResidualDroppath: Enhancing Feature Reuse over Residual Connections](https://mlsanity.com/paper/2411.09475/) (2024)
- [Generative diffusion model surrogates for mechanistic agent-based biological models](https://mlsanity.com/paper/2505.09630/) (2025)
- [Learning To Sample From Diffusion Models Via Inverse Reinforcement Learning](https://mlsanity.com/paper/2602.08689/) (2026)
- [Policy-based Tuning of Autoregressive Image Models with Instance- and Distribution-Level Rewards](https://mlsanity.com/paper/2603.23086/) (2026)
- [Primitive-based Truncated Diffusion for Efficient Trajectory Generation of Differential Drive Mobile Manipulators](https://mlsanity.com/paper/2604.04166/) (2026)
- [Self-Improving Tabular Language Models via Iterative Reward-Guided Post-Training](https://mlsanity.com/paper/2604.18966/) (2026)
- [An Analysis of Posterior Collapse, Parameterization and Initialization in Variational Deep Gaussian Processes](https://mlsanity.com/paper/2606.25882/) (2026)
- [iADD: Improving Alignment and Diversity in Diffusion Policy Optimization](https://mlsanity.com/paper/2610.01789/) (2026)
- [CFG-Policy: Unified Diffusion Policy Learning Framework via Classifier-free Guidance](https://mlsanity.com/paper/2610.14936/) (2026)
- [Weighted Adjoint Matching: Finetuning Flow Matching Models with Time-Dependent Control Cost](https://mlsanity.com/paper/2610.20035/) (2026)

## Abstract (arXiv)

Fine-tuning foundation models via reinforcement learning (RL) has proven promising for aligning to downstream objectives. In the case of diffusion models (DMs), though RL training improves alignment from early timesteps, critical issues such as training instability and mode collapse arise. We address these drawbacks by exploiting the hierarchical nature of DMs: we train them dynamically at each epoch with a tailored RL method, allowing for continual evaluation and step-by-step refinement of the model performance (or alignment). Furthermore, we find that not every denoising step needs to be fine-tuned to align DMs to downstream tasks. Consequently, in addition to clipping, we regularise model parameters at distinct learning phases via a sliding-window approach. Our approach, termed Hierarchical Reward Fine-tuning (HRF), is validated on the Denoising Diffusion Policy Optimisation method, where we show that models trained with HRF achieve better preservation of diversity in downstream tasks, thus enhancing the fine-tuning robustness and at uncompromising mean rewards.