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

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-bystep 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. 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. Diffusion models (DMs) are the de facto state of the art in prompt-based generative modelling across various tasks including text-to-image, text-to-video, molecular graph modelling and medical image reconstruction (Ramesh et al., 2021; Rombach et al., 2022; Ho et al., 2022; Singer et al., 2023; Jing et al., 2022; Song et al., 2022). Most of these applications build on the original Denoising Diffusion Probabilistic Model (DDPM) by Ho et al. (2020), but the extension to other formulations and variants is fast growing.