Lei Zhang and colleagues at Meta, University of California San Diego, Worcester Polytechnic Institute, and Northwestern University presented a fine-tuning method that trains an image generator to plan, create elements, inspect them, and refine the composition rather than composing the whole image at once.
Goal
The authors aimed to reduce mismatches between text prompts and generated images — particularly errors in spatial relations (for example, which object is in front of or behind another) and certain object attributes (such as numbers of fingers, arms, or legs). Instead of producing a full image in a single pass, the model learns an iterative process: plan, sketch, inspect, and refine.
The staged process (example)
For a prompt like “a bear hovering above a silver spoon,” the stages are:
- Plan: write an instruction for the next change (iteration 1: “draw a bear”; iteration 2: “add a spoon under the bear”) and describe the image after the change.
- Sketch: generate the updated image so far (iteration 1: a bear; iteration 2: the bear and a spoon).
- Inspect: compare the instruction and description to the prompt and the image to the instruction.
- Refine: if needed, issue a corrective command (e.g., “the spoon was in front of the bear, draw it under the bear”) and produce a new image.
This approach teaches the model to incrementally construct and correct a scene rather than attempting to capture everything at once.
Baseline model and fine-tuning procedure
The authors began with BAGEL-7B, a pretrained multimodal model that accepts images and text and outputs images and text. They fine-tuned it so it could cycle through the plan, sketch, inspect, and refine stages.
-
Fine-tuning plan and sketch: they generated 32,000 examples, each containing about three to five intermediate images plus a final image. They prompted GPT-4o to convert prompts from two datasets into text-based scene graphs whose nodes were objects or object attributes and whose edges encoded relationships. From each graph they randomly selected parts containing objects, attributes, and relations; GPT-4o turned these parts into incremental prompts to add those items to an image. For each incremental change they used FLUX.1 Kontext to generate an image, and they kept only results GPT-4o judged consistent with the incremental prompts. The model thus learned (i) to predict the next text tokens of the examples and (ii) to generate the next image by adjusting pixel values of the current image over multiple flow-matching steps.
-
Fine-tuning inspect: using the plan+sketch fine-tuned model, the authors generated examples of incremental instructions and images and then asked GPT-4o to judge whether an intermediate text description conflicted with the original prompt; GPT-4o was also prompted to produce critiques and corrective instructions. This produced nearly 7,000 consistent examples and about 8,300 inconsistent ones. By learning to reproduce GPT-4o’s critiques and instructions, the model learned either to accept the current plan as consistent with the original prompt or to describe how to fix inconsistencies.
-
Fine-tuning refine: the authors also fine-tuned on a dataset of images, text reflections about how an image could be improved, and the improved images.
Finally, they fine-tuned the model on all three datasets together using the same loss terms.
Results and metrics
The fine-tuning improved BAGEL-7B on tasks that require object relationships to match a text prompt (for example, placing a bear on a spoon instead of behind it). It also improved the model’s ability to produce images consistent with real-world knowledge such as time of day or historical era.
- GenEval (percentage of prompt-mentioned details that appear in the generated image): BAGEL-7B rose from 77% to 83% after fine-tuning on 62,000 examples; the model used 131 flow-matching steps.
- Comparison: PARM, a method that improves image generation by critiquing intermediate noisy diffusion states, reached 77% after fine-tuning on 688,000 examples and used 1,000 flow-matching steps.
- WISE (GPT-4o ratings for realism, aesthetics and prompt-consistency on a 0–1 scale): BAGEL-7B’s average score increased from 0.7 to 0.76.
The fine-tuned model more frequently placed scenes in the correct historical or temporal context; on a chemistry dataset it more often generated chemically plausible structures, substances, and laboratory scenes.
Why this matters
Modern image generators often produce visually convincing images, but the content can disagree with the prompt: objects may be misplaced or have incorrect attributes. The staged planning, sketching, inspection and refinement pipeline provides a way to improve prompt fidelity without relying solely on scaling up training data.
Commentary
The staged image-composition approach resembles chain-of-thought style stepwise reasoning in large language models: both encourage the model to break a task into parts, and both yield more reliable outputs.
Conclusion
Training an image generator to compose in stages appears to be an effective route to improving spatial relations, attribute correctness and real-world consistency in generated images while using fewer data and computation steps than some prior approaches.
