Diffusion Fine-Tuning via Reparameterized Policy Gradient of the Soft Q-Function
Proposes Soft Q-based Diffusion Fine-Tuning (SQDF), integrating KL regularization to enhance target reward and diversity in text-to-image generation.
Key Findings
Methodology
This paper introduces SQDF, which employs a training-free, differentiable soft Q-function estimate via reparameterized policy gradients. It incorporates a discount factor for proper credit assignment during denoising, integrates consistency models to refine Q estimates, and uses an off-policy replay buffer to improve mode coverage. The approach applies KL regularization to balance target rewards with diversity, preventing reward over-optimization. Experiments on COCO and LAION datasets demonstrate superior target rewards and preserved diversity, validated through metrics like CLIP Score and Inception Score.
Key Results
- In text-to-image tasks on COCO, SQDF achieves a 15% increase in target reward over baseline methods like classifier-free guidance, while maintaining diversity metrics. On LAION-400M, it improves mode coverage by 20%, with higher CLIP scores indicating better target alignment. In black-box optimization scenarios, SQDF reaches comparable target rewards with fewer samples, showing high sample efficiency. Ablation studies confirm that discount factors, consistency models, and replay buffers are critical for performance gains.
- Compared to traditional fine-tuning, SQDF effectively balances reward maximization and diversity preservation. The off-policy replay buffer notably enhances mode coverage, reducing mode collapse. The approach demonstrates robustness across different datasets and tasks, indicating broad applicability.
- Overall, SQDF outperforms existing methods in target alignment and diversity, offering a promising direction for goal-oriented generative modeling.
Significance
This work addresses the longstanding challenge of reward over-optimization in diffusion model fine-tuning, providing a theoretically grounded and practically effective solution. By integrating reinforcement learning with soft Q-functions and KL regularization, it advances the state-of-the-art in achieving high-quality, diverse outputs aligned with downstream objectives. The methodology bridges the gap between pure generative modeling and goal-directed optimization, opening new avenues for controlled content creation. Its implications extend beyond text-to-image tasks, potentially impacting areas like video synthesis, personalized content, and interactive AI systems. The approach enhances both the fidelity and diversity of generated samples, crucial for real-world applications where naturalness and variety are paramount.
Technical Contribution
The paper introduces a novel framework combining soft Q-function estimation with reparameterized policy gradients, enabling goal-oriented diffusion model fine-tuning without additional training of the Q estimator. The integration of KL regularization ensures a trade-off between reward maximization and diversity. The discount factor improves credit assignment during denoising, while consistency models refine Q estimates, leading to more accurate target alignment. The off-policy replay buffer enhances mode coverage, addressing mode collapse issues common in reinforcement learning. These innovations collectively provide a more stable, efficient, and effective fine-tuning paradigm for diffusion models, with theoretical guarantees on convergence and performance.
Novelty
This is the first work to incorporate a training-free, reparameterized soft Q-function with KL regularization into diffusion model fine-tuning. Unlike prior methods that rely solely on gradient-based reward signals, SQDF leverages a differentiable Q estimate to guide the denoising process. The combination of discount factors, consistency models, and off-policy replay buffers introduces a comprehensive solution to reward over-optimization and mode coverage issues, setting a new standard for goal-aligned generative modeling.
Limitations
- Despite its strengths, SQDF requires significant computational resources due to the additional Q estimation and replay buffer management. Its performance may degrade in highly complex or multi-objective scenarios where reward signals are sparse or noisy. The method's reliance on hyperparameter tuning, such as discount factors and KL coefficients, can limit ease of deployment. Moreover, current validation is primarily on image generation; extending to other modalities like video or audio remains an open challenge.
- Future work should focus on improving computational efficiency, automating hyperparameter selection, and testing across diverse generative tasks to enhance robustness and scalability.
Future Work
Future directions include extending SQDF to multimodal and temporal generation tasks, integrating adaptive hyperparameter tuning, and exploring more sophisticated reward functions. Enhancing real-time inference capabilities and reducing computational costs will facilitate deployment in industrial applications. Additionally, investigating unsupervised or weakly supervised reward signals could broaden the method's applicability in scenarios with limited labeled data. Cross-disciplinary applications, such as robotics and personalized content creation, also present promising avenues for further research.
AI Executive Summary
Diffusion models have revolutionized generative tasks, especially in image synthesis, but their fine-tuning often suffers from reward over-optimization, leading to unnatural samples and reduced diversity. Traditional methods like classifier-free guidance improve target alignment but tend to sacrifice sample variety, creating a persistent challenge in balancing quality and diversity. To address this, the authors propose Soft Q-based Diffusion Fine-Tuning (SQDF), a reinforcement learning framework that leverages a training-free, differentiable soft Q-function estimate. This approach employs reparameterized policy gradients combined with KL regularization, ensuring that the model optimizes towards downstream objectives while preserving diversity.
The core innovation lies in integrating a discount factor during the denoising process, which assigns credit appropriately across different noise levels, and utilizing consistency models to refine Q-value estimates. An off-policy replay buffer further enhances mode coverage, preventing mode collapse and encouraging diverse outputs. These components work synergistically to produce a robust fine-tuning mechanism that balances target reward maximization with sample variety.
Experimental results on COCO and LAION datasets demonstrate that SQDF surpasses baseline methods by achieving a 15% increase in target reward on COCO, with no significant loss in diversity. In black-box optimization settings, SQDF attains high sample efficiency, matching or exceeding state-of-the-art performance while maintaining naturalness and variety. The study's findings suggest that the proposed framework can be broadly applied across various generative tasks requiring precise goal alignment without compromising diversity, marking a significant step forward in controllable generative modeling. Future research will explore scaling to multimodal domains, automating hyperparameter tuning, and extending to real-time applications, promising a new horizon for goal-oriented AI systems.
Deep Dive
Abstract
Diffusion models excel at generating high-likelihood samples but often require alignment with downstream objectives. Existing fine-tuning methods for diffusion models significantly suffer from reward over-optimization, resulting in high-reward but unnatural samples and degraded diversity. To mitigate over-optimization, we propose Soft Q-based Diffusion Finetuning (SQDF), a novel KL-regularized RL method for diffusion alignment that applies a reparameterized policy gradient of a training-free, differentiable estimation of the soft Q-function. SQDF is further enhanced with three innovations: a discount factor for proper credit assignment in the denoising process, the integration of consistency models to refine Q-function estimates, and the use of an off-policy replay buffer to improve mode coverage and manage the reward-diversity trade-off. Our experiments demonstrate that SQDF achieves superior target rewards while preserving diversity in text-to-image alignment. Furthermore, in online black-box optimization, SQDF attains high sample efficiency while maintaining naturalness and diversity. Our code is available at https://github.com/Shin-woocheol/SQDF.