DPOK: Reinforcement Learning for Fine-tuning Text-to-Image Diffusion Models

Recent advances in diffusion models, such as DDPM and Latent Diffusion, combined with large-scale text encoders like CLIP, have revolutionized text-to-image synthesis. Models like Stable Diffusion, Dalle-2, and Imagen have demonstrated remarkable capabilities in generating diverse, high-resolution images. However, these models often struggle with complex compositionality, attribute binding, and adherence to nuanced prompts. Supervised fine-tuning on fixed datasets can improve certain aspects but tends to degrade image quality and reduce diversity. Learning from human feedback (LHF) has emerged as a promising approach, training reward models to capture human preferences and guide model updates. Yet, existing methods primarily rely on supervised or reward-filtered learning, which limits adaptability and stability. The challenge remains to develop a scalable, stable, and reward-driven fine-tuning framework that can dynamically optimize for human preferences while maintaining high fidelity.

The core challenge is to enhance the alignment between generated images and complex textual prompts without sacrificing image quality. Supervised fine-tuning often leads to overfitting, mode collapse, and visual artifacts, especially when models are adapted to specific datasets or preferences. Reinforcement learning offers a pathway to optimize for human-centric rewards directly, but applying RL to large-scale diffusion models introduces issues like training instability, high computational costs, and difficulty in balancing reward maximization with model fidelity. The key problem is designing an RL-based fine-tuning method that is both efficient and stable, capable of handling the multi-step denoising process inherent in diffusion models, and ensuring that the model’s outputs remain diverse and high-quality.

This work introduces DPOK, a novel RL framework that formulates diffusion model fine-tuning as a policy optimization problem. It uniquely combines policy gradient methods with KL divergence regularization to stabilize training and prevent overfitting. The approach leverages a reward model trained on human preferences to guide the model towards better text-image alignment. It also employs LoRA for parameter-efficient adaptation, enabling scalable training. Theoretical analysis demonstrates how KL regularization acts as an implicit reward, balancing reward maximization and model fidelity. Empirically, DPOK outperforms supervised fine-tuning across multiple metrics, especially in unseen scenarios, showcasing its robustness and generalization capabilities.

• �� Model the denoising process of the diffusion model as a T-step Markov decision process (MDP), with states representing latent variables and actions as denoising steps. • Use policy gradient algorithms (REINFORCE) to estimate the gradient of the expected reward, which is computed based on the reward model’s output (ImageReward). • Incorporate KL divergence between the current and pre-trained models as a regularizer, approximated via upper bounds on the marginal distributions over images. • Fine-tune only low-rank parameters (LoRA) to improve efficiency. • During training, sample prompts z, generate images x0, evaluate rewards, compute KL divergence, and update the policy parameters θ accordingly. • Repeat iteratively, balancing reward maximization with model stability, to achieve optimal alignment and image quality.

Experiments utilized the Stable Diffusion v1.5 model, fine-tuned with LoRA on 20K images generated from the original model. The reward model, ImageReward, was trained on human preference datasets. Evaluation metrics included ImageReward scores, aesthetic scores, and human preference tests. Baselines included supervised fine-tuning with KL regularization (KL-O). The models were tested on prompts involving color, count, and location, with unseen prompts to assess generalization. Ablation studies varied the KL regularization strength, demonstrating its role in stabilizing training and improving quality. Human evaluations confirmed RL fine-tuning’s superiority in both alignment and visual fidelity.

RL fine-tuning improved ImageReward scores by approximately 12% and aesthetic scores by 8% over supervised methods. Human raters preferred RL-generated images 65% of the time, especially in complex prompts. The models maintained diversity and detail, avoiding artifacts common in supervised fine-tuning. KL regularization was shown to stabilize training, reduce saturation, and improve naturalness. Cross-scenario tests revealed better generalization to unseen prompts, confirming the robustness of the RL approach. These results demonstrate the effectiveness of combining reward-driven optimization with model stability techniques.

The proposed framework enables personalized content creation, virtual assistants, and AI-driven design, where user preferences guide image generation. It requires a reward model aligned with human judgments and computational resources for iterative sampling. Long-term, this approach could facilitate adaptive AI systems capable of real-time preference learning, supporting applications in entertainment, education, and creative industries. Its ability to optimize for nuanced human feedback makes it suitable for interactive AI content generation, enhancing user engagement and satisfaction.

Dependence on the quality of the reward model; biased or noisy rewards can impair training. Computational costs remain high, especially for high-resolution images. Potential overfitting to specific reward signals may limit diversity. Future work should focus on improving reward robustness, reducing training costs, and extending to multi-modal feedback for broader applicability.