Learning from Mistakes: Iterative Prompt Relabeling for Text-to-Image Diffusion Model Training (2024)

Abstract: Diffusion models have shown impressive performance in many domains. However, the model’s capability to follow natural language instructions (e.g., spatial relationships between objects, generating complex scenes) is still unsatisfactory. In this work, we propose Iterative Prompt Relabeling (IPR), a novel algorithm that aligns images to text through iterative image sampling and prompt relabeling with feedback. IPR first samples a batch of images conditioned on the text, then relabels the text prompts of unmatched text-image pairs with classifier feedback. We conduct thorough experiments on SDv2 and SDXL, testing their capability to follow instructions on spatial relations. With IPR, we improved up to 15.22% (absolute improvement) on the challenging spatial relation VISOR benchmark, demonstrating superior performance compared to previous RL methods. Our code is publicly available at https://github.com/cxy000000/IPR-RLDF.

## introduction recent advancements in the field of image generation have been notably driven by diffusion models, especially in the area of text-to-image conversion @xcite . however, a significant challenge arises when these models are tasked with interpreting and executing complex instructions, particularly those involving spatial relationships @xcite . one simple example is that a prompt like "a dog left to a car" often results in images where the spatial relationship "left to" is not accurately depicted. this limitation underscores a crucial gap in the current models' ability to understand and render intricate spatial relationships. in response to these challenges, we investigate how models can self-improve without using extensive training data. some studies have found that self-supervised learning is useful for training diffusion models @xcite . however, these methods still lack an effective self-correction mechanism, making the self-generated dataset either noisy or being trimmed off too much. recent advances in large language models (llm) have found that learning from mistakes and receiving language feedback can effectively enhance the model's reasoning capacity @xcite . inspired by this, in our research, we study whether we can effectively use the originally incorrect text-image pairs to train the diffusion model. in particular, we propose iterative prompt relabeling (ipr), a novel technique designed to enhance the alignment of images with text through an iterative process of image sampling and prompt relabeling. our approach uses rich language feedback for all the images and a simple reward design. a high-level overview of our approach is shown in figure 1 . we begin by establishing a reward function based on an external detection model that automatically classifies the correctness of an image and assigns rewards based on the detection results. this enables a straightforward reward design and neglects the complicated training of reward models. then, we use prompt relabeling to relabel the input prompt of the mismatched image-text pairs based on the results of the detection model. this allows the models to use the mis-generated images with the correct version. finally, we adopt iterative training that continuously trains the model with its self-generated images. iterative training allows dynamically scaling up the dataset and receiving additional feedback through multiple training rounds, which can further enrich the feedback and progressively refine the model's performance. we demonstrate the efficacy of ipr across sdv2 @xcite and sdxl @xcite , trained with and w/o lora, and we test the performance on the challenging spatial relation task, where we observe a substantial improvement of up to 15.22% (absolute improvement) on the vi-sor @xcite benchmark. this performance underlines the potential of ipr in pushing the boundaries of text-to-image generation models, especially in terms of understanding and rendering complex spatial relationships. ## related work text-to-image models have facilitated the generation of high-resolution, multi-styled images @xcite @xcite . however, training diffusion models from scratch demands a substantial amount of data and time. therefore, various fine-tuning strategies have been explored. these approaches include associating a unique identifier with a particular subject @xcite , introducing new embeddings to represent user-provided concepts @xcite , adapting compositional generation @xcite , and implementing low-rank adaptation @xcite . our work utilizes lora to fine-tune the diffusion model and follow the reinforcement learning line of work to rescale loss for effective training. recent research has demonstrated the efficacy of prompt relabeling techniques. tempera @xcite offers interpretable prompts tailored to various queries through the creation of an innovative action space, enabling flexible adjustments to the initial prompts. furthermore, the hindsight instruction relabeling (hir) @xcite approach conceptualizes the instruction alignment problem as a goal-reaching problem within the context of decision-making. it entails the conversion of feedback into instructions by re-labeling the original instructions. our approach represents a significant advancement as it employs prompt relabeling techniques with diffusion models for the first time. ## experiments in this section, we assess ipr's effectiveness through multiple settings. we first introduce our experiment settings. then, we present the quantitative results and qualitative results of irp. finally, we present ablation studies to fully analyze our method. ## conclusion in this work, we present ipr in response to the challenging task of spatial location generation. ipr is a novel algorithm that designs rich language feedback and incorporates it with the detection model rewards. then it trains the diffusion model by iteratively receiving rewards and language feedback. this algorithm is a plug-and-play method that is applicable to a range of diffusion models. extensive results show the model's effectiveness on the challenging spatial relationship benchmark across various settings.