Re-Align: Structured Reasoning-guided Alignment for In-Context Image Generation and Editing

AI summary

Overview

Research area: Computer Vision — unified multimodal models for in-context image generation and editing (ICGE), combining chain-of-thought reasoning with reinforcement learning for visual generation.

Technical level: Advanced. The paper assumes familiarity with diffusion/rectified-flow image generation, multimodal foundation models, chain-of-thought prompting, and policy optimization (GRPO/PPO).

Scope: This paper proposes Re-Align, a framework that connects a unified multimodal model's understanding capability to its image generation capability by having the model produce a structured reasoning trace (In-Context Chain-of-Thought) before generating an image, and then aligning the two using a CLIP-based surrogate reward inside a GRPO training loop.

What This Paper Is About

In-context image generation and editing lets users describe what they want using interleaved images and text — for example, "replace the hat in the first image with the cup in the second image." Models that are strong at understanding such prompts often fail to execute them: their internal reasoning may be correct, but the final image does not reflect it. The paper's goal is to close that reasoning-to-generation gap by structuring the reasoning output and then explicitly training the model to make its generated image consistent with that reasoning.

Key Contributions

1. Re-Align framework: A unified framework for in-context image generation and editing built on the BAGEL multimodal foundation model architecture, reported to achieve state-of-the-art performance among methods of comparable model scale and computational resources on both ICGE tasks.

2. In-Context Chain-of-Thought (IC-CoT): A structured reasoning paradigm that decouples two components — semantic guidance (an explicit predicted caption bracketed by <out_caption>/</out_caption>) and reference association (a per-reference-image role description bracketed by <relation_i>/</relation_i>, with the number of associations matching the number of reference images).

3. A surrogate reward plus a reasoning-induced diversity (RID) strategy: Instead of task-specific reward models, the paper uses CLIP image-text similarity between the predicted caption and the generated image as the reward signal, optimized with Group Relative Policy Optimization (GRPO). RID generates distinct IC-CoT reasoning chains for each sample within a group to raise reward variance in a controlled way and stabilize training.

4. Re-Align-410K dataset: A 410K-sample ICGE dataset with IC-CoT annotations, produced by an automated pipeline (reference image preparation, adaptive instruction generation with Gemini 2.5, reasoning text generation, target image generation with GPT-4o, and multidimensional filtering) covering three generation task types and three editing task types.

Main Findings

• OmniContext benchmark (Table 2): Re-Align attains an average score of 8.21, the highest overall average among models with comparable scale and resources. For comparison, the paper reports BAGEL at 5.73, OmniGen2 at 7.18, Qwen-Image-Edit-2509 at 7.69, DreamOmni2 at 6.31, and Gemini 2.5 Flash Image at 7.84. GPT-4o and Emu3.5 are reported at 8.80 and 8.82 respectively, above Re-Align's average; the paper frames its claim specifically around comparable model scale and resources.

• Per-task breakdown on OmniContext: Re-Align scores 8.25 (single character), 8.55 (single object), 8.25 (multiple character), 8.07 (multiple object), 8.28 (multiple character + object), 8.21 (scene character), 8.25 (scene object), and 7.82 (scene character + object). It ranks second only to Qwen-Image-Edit (2509) on the SINGLE task and reports the best overall performance on the MULTIPLE and SCENE tasks.

• DreamOmni2Bench (Table 3): Re-Align leads on editing tasks — Add (PF 9.27, SC 9.27, Overall 9.27), Replace (Overall 8.61), Global (Overall 7.85), and Local (Overall 6.35) — and on generation (PF 7.74, SC 7.67, Overall 7.24). DreamOmni2, which uses separate parameters for generation and editing, reports Overall scores of 6.87 (Add), 7.05 (Replace), 7.76 (Global), 5.44 (Local), and 6.56 (Generation).

• Impact of reference image count (Table 4): Using two references for all editing tasks and one to four for generation, Re-Align reports Overall scores of 7.19 (editing), 6.37 (1 reference), 7.93 (2), 7.70 (3), and 7.39 (4), frequently ranking first or second across metrics.

• Ablation of training stages (Table 5): Starting from PF 6.92 / SC 5.47 / Overall 5.80 / CLIP_out 32.44 with no training stages, supervised fine-tuning (SFT) raises these to 7.51 / 6.46 / 6.77 / 33.32. Adding reasoning-generation alignment (RGA) yields 7.46 / 6.54 / 6.80 / 33.50 — a CLIP_out gain but no significant PF gain, which the paper attributes to low sample diversity hurting RL training. Adding reasoning-induced diversity (RID) produces the best results: PF 7.61, SC 6.57, Overall 6.89, CLIP_out 33.90.

• IC-CoT versus alternatives (GSB evaluation, Figure 7): The structured IC-CoT outperforms a no-chain-of-thought variant and an unstructured reasoning variant following BAGEL, with win rates 20% and 16.25% higher, respectively.

• Qualitative observations (Figure 5): OmniGen2 often incorporates irrelevant elements from reference images; BAGEL, Qwen-Image-Edit, and DreamOmni2 show weaker subject consistency in some cases; most existing models fail on material replacement and object addition edits.

• Failure cases (Figure 9): The model occasionally generates incorrect reasoning text (for example, for the complex action semantics of "come here"), and on editing tasks without dedicated training (such as editing based on referenced text styles or object color schemes) it shows semantic comprehension but low reference consistency.

Methodology in Plain English

The researchers start from BAGEL, a multimodal model that both understands interleaved image-text input and generates images. Their first idea is that the model should write down a structured "plan" before drawing. That plan, called IC-CoT, has two mandatory parts. The first part is a plain caption of what the output image should look like, delimited by <out_caption> tags — this effectively turns a messy multi-image request into an ordinary text-to-image problem. The second part is one short statement per reference image describing what role that image plays, delimited by <relation_i> tags — this prevents the model from mixing up which reference contributes what.

They train in two stages. Stage one is ordinary supervised fine-tuning on the Re-Align-410K dataset, so the model learns to produce IC-CoT text followed by an image. Stage two is where the alignment happens: they use Group Relative Policy Optimization (GRPO), but instead of building an expensive task-specific reward model, they use a cheap surrogate — they extract the caption the model wrote between the <out_caption> tags and measure its CLIP similarity to the image the model actually produced. If the model draws something matching its own stated caption, it gets rewarded. Because ICGE prompts already constrain outputs heavily, generated samples in a group tend to look similar, which flattens the reward signal that GRPO relies on; their RID strategy addresses this by producing a different reasoning chain for each sample in the group, naturally diversifying outputs without the image-quality degradation that comes from simply increasing noise.

The dataset itself is built automatically: sample reference images from a pool of characters, objects, and scenes; have Gemini 2.5 write instructions for those images; have the same MLLM write the IC-CoT reasoning (deliberately withholding the target image to reduce hallucination); have GPT-4o generate the target image; then filter with image-text similarity, aesthetics and human-preference metrics, and OmniContextScore, discarding roughly 20% of the data to leave 410K samples.

Training details: SFT runs for 100,000 steps on 64 NVIDIA H20 GPUs at a learning rate of 5×10⁻⁶; the alignment stage runs for 200 steps with group size 32 at a learning rate of 1×10⁻⁶. Images are generated at 1024×1024 with 50 denoising steps.

Why This Matters

Impact on research: The paper argues that strong multimodal understanding does not automatically transfer to generation, and that unstructured reasoning (as in BAGEL) does not actually help the downstream image. It offers a concrete recipe — structured reasoning plus a cheap, structure-enabled surrogate reward — for turning a unified model's reasoning into a controllable intermediate signal. It also contributes a 410K annotated dataset and an automated pipeline covering both generation and editing tasks, which prior video-frame-extraction approaches (such as OmniGen2) largely could not handle.

Real-world applications (plausible given the demonstrated tasks):

• Subject-driven generation, keeping a specific character or product consistent across new scenes.
• Compositional generation that places multiple referenced subjects, or subjects plus a referenced scene, into one new image.
• Reference-based editing such as adding or replacing objects, transferring texture/pose/style attributes locally or globally, and scene replacement.
• Multi-reference inputs (the paper reports generation results with one to four reference images), useful when a request depends on several source images at once.

Industry relevance: The paper comes from researchers affiliated with IIE, CAS, UCAS, and Tencent Hunyuan, and it benchmarks against closed systems such as GPT-4o, Gemini 2.5 Flash Image, and Nano Banana. Re-Align is built on BAGEL and stated to be compatible with other models providing unified understanding and generation (Janus-Pro, Transfusion, and others cited), so the alignment recipe is positioned as reusable rather than tied to one backbone.

Future Directions

• Scaling model size and data: The conclusion explicitly notes that model size and data scale are limited compared with production-level work like GPT-4o, which may constrain performance in diverse scenarios; scaling both is suggested as a remedy for the failure cases.
• Visual chain-of-thought: The current IC-CoT operates "purely at the textual level," and the authors name extending reasoning to visual Chain-of-Thought as a promising future direction.
• Broader editing coverage: Failure cases include editing based on referenced text styles or object color schemes, which the model handles with low reference consistency because these tasks lack dedicated training data — suggesting expanded task coverage in future data construction.
• Reward design for ICGE: The paper states that existing RL-based approaches lack effective reward design and comprehensive experimental validation for ICGE, and its own surrogate reward is an indirect proxy (CLIP similarity against the predicted caption) rather than a direct task-quality reward, leaving room for better reward formulations.

Target Audience

Researchers and engineers working on unified multimodal generation models, controllable image generation and editing, and reinforcement learning for visual generation. It will be most valuable to readers who already understand diffusion/rectified-flow training and policy optimization, and to practitioners building subject-driven or reference-based editing systems who want a concrete data-construction and alignment pipeline. Readers looking for a beginner-level introduction to image generation will find the method sections demanding, since the paper relies on equations for language modeling, rectified flow, and reward computation without introductory background.

Authors’ abstract

In-context image generation and editing (ICGE) enables users to specify visual concepts through interleaved image-text prompts, demanding precise understanding and faithful execution of user intent. Although recent unified multimodal models exhibit promising understanding capabilities, these strengths often fail to transfer effectively to image generation. We introduce Re-Align, a unified framework that bridges the gap between understanding and generation through structured reasoning-guided alignment. At its core lies the In-Context Chain-of-Thought (IC-CoT), a structured reasoning paradigm that decouples semantic guidance and reference association, providing clear textual target and mitigating confusion among reference images. Furthermore, Re-Align introduces an effective RL training scheme that leverages a surrogate reward to measure the alignment between structured reasoning text and the generated image, thereby improving the model's overall performance on ICGE tasks. Extensive experiments verify that Re-Align outperforms competitive methods of comparable model scale and resources on both in-context image generation and editing tasks.