Mohammad Ghavamzadeh, Amazon AWS AI Labs

Abstract:

We present three recent results in fine-tuning text-to-image diffusion models.

(i) In the first work, we train a reward function using a human-labeled text-image dataset and fine-tune the model by maximizing reward-weighted likelihood to improve text-image alignment. We also analyze several design choices and show that their careful investigation is crucial when balancing the alignment-fidelity tradeoff.

(ii) In the second work published at NeurIPS-2023, we propose using online reinforcement learning (RL) to fine-tune text-to-image diffusion models. We first formulate the fine-tuning task as an RL problem, and then show how to update the pre-trained model using policy gradient. Our approach integrates policy optimization with KL regularization. We evaluate our method in terms of both text-image alignment and image quality.

(iii) Finally, in the last work accepted at ICLR-2024, we take on the problem of reward over-optimization in text-to-image models. We first introduce the Text-Image Alignment Assessment (TIA2) benchmark, a diverse collection of text prompts, images, and human annotations, for studying this issue. We evaluate several state-of-the-art reward models for text-to-image generation on our benchmark and find that they are often not well-aligned with human assessment. We introduce a simple method, TextNorm, for inducing confidence calibration in reward models by normalizing the scores across prompts that are semantically different from the original prompt, and demonstrate that using the confidence-calibrated scores in fine-tuning can effectively reduce the risk of over-optimization.