Elucidating Optimal Reward-Diversity Tradeoffs in Text-to-Image Diffusion Models (2409.06493v1)

Abstract: Text-to-image (T2I) diffusion models have become prominent tools for generating high-fidelity images from text prompts. However, when trained on unfiltered internet data, these models can produce unsafe, incorrect, or stylistically undesirable images that are not aligned with human preferences. To address this, recent approaches have incorporated human preference datasets to fine-tune T2I models or to optimize reward functions that capture these preferences. Although effective, these methods are vulnerable to reward hacking, where the model overfits to the reward function, leading to a loss of diversity in the generated images. In this paper, we prove the inevitability of reward hacking and study natural regularization techniques like KL divergence and LoRA scaling, and their limitations for diffusion models. We also introduce Annealed Importance Guidance (AIG), an inference-time regularization inspired by Annealed Importance Sampling, which retains the diversity of the base model while achieving Pareto-Optimal reward-diversity tradeoffs. Our experiments demonstrate the benefits of AIG for Stable Diffusion models, striking the optimal balance between reward optimization and image diversity. Furthermore, a user study confirms that AIG improves diversity and quality of generated images across different model architectures and reward functions.

Elucidating Optimal Reward-Diversity Tradeoffs in Text-to-Image Diffusion Models

Text-to-Image (T2I) diffusion models have gained traction as premier generative models capable of producing high-quality images from textual descriptions. However, their reliance on large, unfiltered image-caption datasets harvested from the internet presents inherent risks, such as the generation of unsafe or misaligned content. Furthermore, the aesthetic quality often deviates from human preferences, necessitating optimization strategies incorporating these preferences into model training. Recent endeavors have involved fine-tuning T2I models using human preference datasets or optimizing bespoke reward functions. Nonetheless, these methodologies often succumb to reward hacking, where the model overly optimizes to the reward function, consequently diminishing the diversity of generated images.

The paper “Elucidating Optimal Reward-Diversity Tradeoffs in Text-to-Image Diffusion Models” by Jena et al. explores this critical issue. The authors highlight the inevitability of reward hacking when optimizing reward functions directly. They discuss established regularization techniques such as KL divergence and LoRA scaling, pointing out their limitations. Building on this exploration, the authors introduce Annealed Importance Guidance (AIG), an inference-time regularization technique inspired by Annealed Importance Sampling (AIS), which balances reward optimization and diversity preservation effectively.

Methodology and Contributions

The authors provide a rigorous analysis to demonstrate that reward hacking is an intrinsic outcome of the reward function maximization framework. Using a straightforward proof, they show that in the non-parametric setting, the optimal probability distribution is a Dirac-delta function, indicating a complete lack of diversity. This observation necessitates the introduction of regularization.

Existing approaches using KL divergence and LoRA scaling are critically examined. KL divergence introduces a trade-off hyperparameter (λ) that must be fine-tuned to balance reward and diversity. However, this hyperparameter's effectiveness varies across model architectures, as evidenced by distinct behaviors observed between SDv1.4 and SDXL models in the experiments. LoRA scaling, on the other hand, mitigates substantial parameter changes by interpolating between the base and fine-tuned model parameters, effectively implementing implicit regularization akin to Tikhonov regularization under mild assumptions.

AIG emerges as a novel approach that dynamically anneals the influence between the base and DRaFT (Differentiable Reward Fine-Tuning) score functions during the reverse diffusion process. By varying the mixing coefficient (γ), which transitions from 0 to 1 over time, AIG leverages the base model's diverse mode recovery capabilities in the initial phases and the DRaFT model’s reward optimization prowess in the later stages.

Experimental Analysis

The effectiveness of AIG is evaluated against both SDv1.4 and SDXL diffusion models trained on PickScore and HPSv2 reward models. Comprehensive ablation studies involving over 149,600 generated images per configuration reveal that AIG achieves a Pareto-optimal balance between reward maximization and image diversity without necessitating multiple hyperparameter-specific training runs. Traditional metrics such as Frechet Inception Distance (FID) and precision-recall evaluations substantiate the quantitative merits of AIG. Moreover, the authors introduce the Spectral Distance metric to address reference mismatches, focusing on distributional spreads rather than differences in means or rotations.

Qualitative assessments are bolstered by a user paper involving 36 participants who validate that images generated using AIG demonstrate superior diversity and aesthetics compared to those from DRaFT or other regularizations. This user-centric evaluation underscores AIG’s practical benefits in producing varied yet high-quality images that better align with human aesthetic preferences.

Implications and Future Developments

The implications of this research are multifaceted. Practically, the introduction of AIG facilitates more reliable deployment of T2I models by ensuring generated images maintain requisite diversity while adhering to specified reward functions. Theoretically, this work paves the way for deeper inquiries into the nature of reward optimization in generative models, highlighting the necessity of novel regularization schemes that adaptively balance competing objectives.

Future research directions could involve integrating more sophisticated scheduling for the annealing process, potentially facilitating user-customizable trade-offs between diversity and reward optimization. Another promising avenue is leveraging Large Multimodal Models (LMMs) to refine alignment correction mechanisms, particularly for intricate textual attributes. Given the foundational insights provided by this paper, subsequent explorations can adopt a more nuanced approach to understanding and mitigating the complexities of reward hacking in generative modeling.

In summary, this paper presents a thorough examination and innovative solution to one of the pivotal challenges in T2I diffusion models—ensuring diverse and high-quality image generation through strategically balanced regularization techniques.