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

TL;DR

This paper introduces Hierarchical Reward Fine-tuning (HRF), a novel method for stabilizing diffusion model training with reinforcement learning, improving diversity preservation and robustness during fine-tuning.

Contribution

The paper proposes HRF, a hierarchical RL-based fine-tuning approach that dynamically trains diffusion models and regularizes parameters to prevent mode collapse and enhance robustness.

Findings

HRF improves diversity in downstream tasks.

HRF achieves higher robustness without sacrificing mean rewards.

Dynamic, step-wise fine-tuning stabilizes diffusion models.

Abstract

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.