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

Paper: arXiv:2601.05124 Code: hrz2000/realign Code reference: master @ 14f5f074 (2026-01-19)

1. Motivation（研究动机）

ICGE（In-Context Image Generation and Editing）的核心难点不是“会不会生成图”，而是模型能否在一个交错的 image-text prompt 中先理解多个 reference images 的角色，再把这种理解稳定地传递到生成过程。普通 Text-to-Image 或单图编辑任务里，文本 prompt 已经显式描述了目标；ICGE 里用户经常只说“把第一张图里的帽子换成第二张图里的杯子”“让他们站在一起”，目标对象、属性、背景和 reference image index 之间的关系需要模型自己解析。

现有 unified multimodal models（尤其是 BAGEL）已经表现出不错的 multimodal understanding 和 reasoning 能力，但论文指出一个关键断层：模型可以写出看似正确的 CoT，却生成出没有遵循 CoT 的图像。也就是说，“理解”并没有自然转化成“生成控制”。Figure 2 中 BAGEL 对 barn owl 的例子能说明这一点：reasoning 里知道猫头鹰应该在椅子上、办公室背景应该保留，但生成图像并没有忠实反映这些关系。

Figure 2 解读：图中对比 BAGEL 和 Re-Align 的 reasoning paradigm。BAGEL 的 <think> 是一段长文本扩写，语义上看似合理，但没有把“reference image 1 是 home office，reference image 2 是 barn owl”这种结构化关系绑定到生成结果；Re-Align 则把输出 caption 和每个 reference 的 role 拆成 <out_caption> 与 <relation_i>，让生成端得到更清晰、可对齐的中间表示。

这篇论文要解决的问题是：如何让 unified model 的理解能力真正帮助 in-context image generation/editing，而不是停留在“会解释但不按解释画”的状态。解决这个问题的价值在于，如果 reasoning 可以成为生成的可优化中间状态，ICGE 就能从 prompt hacking/手工模板转向一种可训练、可评估、可做 RL alignment 的统一流程。

2. Idea（核心思想）

核心 insight 是：不要让 multimodal model 输出一段自由散漫的 CoT，而是把 reasoning 变成一个结构化、可解析、可打分的中间对象。Re-Align 的 IC-CoT 把 reasoning 拆成两部分：<out_caption> 给出目标图像的显式语义描述，<relation_i> 逐一说明每张 reference image 在目标图像中的作用。

这个结构化 reasoning 一方面把复杂的 image-text interleaved condition 部分转化成更接近 text-conditioned generation 的目标 caption，另一方面提供了一个 surrogate reward：用 CLIP 计算 <out_caption> 与生成图之间的 image-text similarity，从而不必为每类 ICGE 子任务单独训练昂贵的 reward model。与 BAGEL 式 prompt expansion 相比，Re-Align 的根本差异是：reasoning 不只是给用户看的解释，而是被训练成生成器的 conditioning signal，并在 RL 阶段被直接用于衡量 reasoning-generation consistency。

3. Method（方法）

3.1 Overall framework：BAGEL backbone + 两阶段训练

Re-Align 基于 BAGEL 的 unified understanding/generation 架构，输入是包含若干 reference images 和文本指令的交错 prompt $P$ 。模型先自回归地产生 structured reasoning $R = {r_{1}, r_{2}, \dots, r_{M}}$ ，然后在 image generation branch 中以 $P$ 和 $R$ 为条件生成图像 $I$ 。

Figure 3 解读：Re-Align 的 pipeline 有两个训练阶段。Stage 1 是 SFT，让模型学会在 IC-CoT 指导下生成图像；Stage 2 是 reasoning-generation alignment，用 GRPO 对齐 structured IC-CoT 和生成图。关键变化是，reasoning 不再只是 prompt 前缀，而是进入 image generation objective，并在 RL reward 中被重新读出。

训练目标分成 reasoning token likelihood 和 Rectified Flow image objective。对 reasoning，论文使用标准 LM objective： $L_{cot} (θ) = \sum_{i} lo g p_{θ} (r_{i} ∣ P, r_{< i}) .$ 对图像生成，沿用 BAGEL/Rectified Flow 形式：令 $x_{0} \sim p_{0}$ 为真实数据， $x_{1} \sim p_{1}$ 为 Gaussian noise， $x_{t} = (1 - t) x_{0} + t x_{1}$ ，目标速度 $v = x_{1} - x_{0}$ ，则 $L_{img} (θ) = E_{t, x_{0}, x_{1}} [∥ v - v_{θ} (x_{t}, t, P, R) ∥^{2}] .$ 直观上， $L_{cot}$ 让模型会“说清楚目标图和 reference 的关系”， $L_{img}$ 让生成过程在这个结构化说明上作画；后续 RL 再把“说清楚”和“画出来”之间的 gap 作为优化对象。

3.2 IC-CoT：semantic guidance + reference association

IC-CoT 的第一部分是 semantic guidance：模型在 <out_caption> 与 </out_caption> 之间预测最终图像 caption。这个 caption 的作用不是普通 captioning，而是把复杂的多图、多对象、多属性约束压缩成一个明确的生成目标。例如“第一张图里的男人和第二张图里的女人坐在第三张图里的沙发上，风格参考第四张图”会被转写成目标图描述。

第二部分是 reference association：模型为每张 reference image 生成 <relation_i>...</relation_i>，说明该图片在目标图中的 role。这个设计针对 ICGE 中常见的歧义：用户往往不会显式说 image index，而会用“them”“same style”“the object”这类表达；relation 字段强迫模型逐一绑定 “image_1 提供背景 / image_2 提供 subject / image_3 提供 style”。

Figure 1 解读：上半部分展示 Re-Align 同时支持 in-context image generation 和 in-context image editing；下半部分给出推理例子，右侧 JSON 化的 reasoning 中包含 out_caption 和多条 relations。这正对应 IC-CoT 的两类字段：caption 负责“目标是什么”，relation 负责“目标中每个来源是什么”。

IC-CoT 相比 BAGEL 式 unstructured prompt expansion 的优势在于结构紧凑、可解析、可作为 reward 输入。它不是简单增加推理长度，而是刻意把生成最需要的变量暴露出来：最终 caption 和 reference-to-target mapping。若去掉 semantic guidance，image branch 要直接从复杂 prompt 中猜目标；若去掉 reference association，多 reference 之间会出现 subject/style/background 混淆。

3.3 Reasoning-Generation Alignment：用 structured reasoning 做 surrogate reward

ICGE 的 RL 难点是 reward 设计：任务同时涉及 interleaved image-text 输入、generation/editing、多 reference consistency、prompt following 等维度。如果为每个子任务训练 reward model，成本很高且不容易泛化。Re-Align 的做法是利用 IC-CoT 的结构化输出，从 <out_caption> 中抽取 caption $c$ ，再用 CLIP image encoder $E$ 和 text encoder $T$ 计算 caption-image similarity： $s (x, c) = \frac{E ( x ) ^{⊤} T ( c )}{∥ E ( x )∥ \cdot ∥ T ( c )∥} .$ 这个 reward 不直接评价所有 reference consistency，但它抓住了一个可优化代理目标：如果模型生成图不能匹配自己先写出的目标 caption，说明 reasoning 与 generation 没对齐。因为 caption 是由复杂 prompt 和 reference associations 推出来的，提升 caption-image similarity 会间接推动模型把理解结果注入生成过程。

论文还提出 Reasoning-Induced Diversity（RID）。GRPO 需要一个 group 内多个 samples 的 reward 有足够方差，才能计算有信息量的 relative advantage；ICGE 的 reference constraints 很强，直接采样往往生成差异小，reward variance 低。以往视觉 GRPO 通过增大 SDE noise scale 来扩大差异，但可能损害画质。RID 改为让 group 内每个 sample 使用不同 IC-CoT reasoning chain，自然诱导不同生成轨迹，在不粗暴加噪的情况下增加 reward variance。

3.4 Re-Align-410K 数据构造

Figure 4 解读：Re-Align-410K 的数据流包括 reference images preparation、adaptive instruction generation、reasoning text generation、target image generation、data filtering 和任务分布统计。重点是 target image 不是先验存在的人工标注，而是由 reference group + instruction 输入 GPT-4o 生成；reasoning text 由 MLLM 基于 reference images 和 instruction 生成，且刻意不看 target image，以避免 reasoning 对已生成图的后验幻觉。

数据任务覆盖两大类。In-context image generation 包含 subject-driven generation、subject-subject compositional generation、subject-scene compositional generation；in-context image editing 包含 reference subject editing、reference attribute editing、reference scene editing。reference image pool 覆盖 characters、objects、scenes，并按不同任务类型采样。

Adaptive instruction construction 使用 Gemini 2.5 根据 reference group 生成可执行指令，并鼓励关注次要视觉细节以提升复杂度。Reasoning text generation 再用 MLLM 生成结构化 IC-CoT；Target Image Generation 使用 GPT-4o 合成目标图。Filtering 有三类指标：1）IC-CoT caption 与 target image 的 image-text similarity；2）LAION aesthetics、ImageReward、PickScore 等视觉质量/偏好分；3）OmniContextScore 的 instruction following 和 semantic consistency。低于阈值的样本被移除约 20%，最终得到 410K high-quality samples。

3.5 Pseudocode（基于论文流程；released repo 未提供实现代码）

公开 GitHub repo hrz2000/realign 在 master@14f5f074 仅包含 README 和展示资产，未发布训练/推理源码。因此下面 pseudocode 是按论文公式和 pipeline 还原的 PyTorch-style 逻辑，不声称对应 released code 的具体函数实现。

def build_ic_cot_prompt(reference_images, user_instruction, mlm):
    # reference_images: list[Image], user_instruction: str
    # Output: structured reasoning with <out_caption> and <relation_i> fields
    cot = mlm.generate(
        images=reference_images,
        text=user_instruction,
        schema={
            "out_caption": "target image caption",
            "relations": ["role of image_i in target for each reference"],
        },
    )
    caption = parse_tag(cot, "out_caption")
    relations = [parse_tag(cot, f"relation_{i+1}") for i in range(len(reference_images))]
    return cot, caption, relations

def sft_step(model, batch, optimizer):
    # batch contains interleaved prompt P, structured reasoning R, target image x0
    prompt, reasoning_tokens, target_image = batch["P"], batch["R"], batch["image"]
    noise = torch.randn_like(target_image)
    t = torch.rand(target_image.shape[0], device=target_image.device)
    x_t = (1 - t[:, None, None, None]) * target_image + t[:, None, None, None] * noise
    velocity_target = noise - target_image
 
    lm_loss = model.language_modeling_loss(prompt, reasoning_tokens)
    velocity_pred = model.flow_predict(x_t, t, prompt=prompt, reasoning=reasoning_tokens)
    img_loss = F.mse_loss(velocity_pred, velocity_target)
    loss = lm_loss + img_loss
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    return {"loss": loss.item(), "lm_loss": lm_loss.item(), "img_loss": img_loss.item()}

def surrogate_reward(generated_images, captions, clip_model):
    image_feat = F.normalize(clip_model.encode_image(generated_images), dim=-1)
    text_feat = F.normalize(clip_model.encode_text(captions), dim=-1)
    return (image_feat * text_feat).sum(dim=-1)
 
 
def rga_grpo_step(policy, old_policy, prompts, clip_model, group_size=32, eps=0.2):
    # RID: sample different structured reasoning chains inside each group
    samples = []
    for _ in range(group_size):
        reasoning, caption, image, logprob = policy.sample_with_ic_cot(prompts)
        reward = surrogate_reward(image, caption, clip_model)
        old_logprob = old_policy.logprob(prompts, reasoning, image)
        samples.append((logprob, old_logprob, reward))
 
    rewards = torch.stack([s[2] for s in samples])
    advantage = (rewards - rewards.mean()) / (rewards.std(unbiased=False) + 1e-6)
 
    losses = []
    for i, (logprob, old_logprob, _) in enumerate(samples):
        ratio = torch.exp(logprob - old_logprob)
        clipped = torch.clamp(ratio, 1 - eps, 1 + eps)
        losses.append(-torch.minimum(ratio * advantage[i], clipped * advantage[i]))
    return torch.stack(losses).mean()

3.6 Code-to-paper mapping

Code reference: master @ 14f5f074 (2026-01-19) — released repository inspection based on this commit; repo currently lacks training/inference implementation.

Paper Concept	Source File	Key Class/Function
Project metadata / paper links / abstract	`README.md`	README-only project description
Teaser / ICGE examples	`assets/README/1768824629720.png`	visual asset only
IC-CoT, SFT, RGA, RID training pipeline	not released	论文公式与 released code 实现差异：paper reports method and training settings, but released repo contains no implementation/config files to verify
Training config numbers	not released	paper text only: SFT 100,000 steps, 64 NVIDIA H20, LR $5 \times 1 0^{- 6}$ ; RGA 200 steps, group size 32, LR $1 \times 1 0^{- 6}$

4. Experimental Setup（实验设置）

4.1 Training data and config

训练数据分两层：一是论文构造的 Re-Align-410K，最终 410K high-quality ICGE samples，覆盖 3 类 in-context generation 与 3 类 in-context editing；二是实验设置中还提到一个约 2.0M samples 的 custom dataset，用于缓解 task-specific data scarcity，相比初始 1.8M 加入更多 high-quality instruction-reference image pairs。注意：released repo 没有数据构造代码或配置，以上数字来自论文 source/PDF，而不是代码复现配置。

训练以 BAGEL 为 backbone，并声明兼容 Janus-Pro、Transfusion、HunyuanImage-3 等具备 unified understanding/generation 的模型。SFT stage：100,000 steps，64 NVIDIA H20 GPUs，learning rate $5 \times 1 0^{- 6}$ 。Reasoning-generation alignment stage：200 steps，group size 32，learning rate $1 \times 1 0^{- 6}$ 。默认图像输出分辨率 1024×1024，50 denoising steps。

4.2 Baselines and benchmarks

Baselines 包括 BAGEL、OmniGen2、Echo-4o、Qwen-Image-Edit(2509)、DreamOmni2；OmniContext 表里还报告了 FLUX.1 Kontext [Max]、Gemini 2.0/2.5 Flash Image、GPT-4o、Emu3.5、OmniGen、InfiniteYou、UNO 等参考结果。

评测基准包括 OmniContext 和 DreamOmni2Bench。OmniContext 主要评估 in-context image generation；DreamOmni2Bench 同时覆盖 generation 和 editing，输入 reference images 从 1 到 5 张，编辑任务涵盖 Add、Replace、Global、Local。评测指标沿用 OmniContext 的 GPT-4.1 evaluator：Prompt Following（PF，是否符合指令）、Subject Consistency（SC，生成图与 reference visual concepts 是否一致）、Overall（PF 和 SC 的 geometric mean）。由于 DreamOmni2Bench 官方评测代码尚未发布，论文用 OmniContext metric framework 进行评估。

5. Experimental Results（实验结果）

5.1 Qualitative results

Figure 5 解读：前两行是 in-context image generation，后三行是 in-context image editing。论文观察到 OmniGen2 容易把 reference 中无关元素带入生成图，BAGEL、Qwen-Image-Edit、DreamOmni2 在 subject consistency 上不足；编辑任务中，material replacement 和 object addition 更容易暴露“理解了指令但执行不对”的问题。Re-Align 的优势体现在 reference subject、背景、属性迁移之间更少混淆。

Figure 6 解读：更多案例强调 Re-Align 可以处理多 reference 的生成和编辑。每组最后一张是输出，其余为输入 reference images；这类图的价值在于展示 relation binding 是否稳定，而不只是单张图片质量。

5.2 Quantitative results on OmniContext

OmniContext 上，Re-Align 在 comparable scale/resource 的开源/可比模型里取得最高 average 8.21，高于 Qwen-Image-Edit-2509 的 7.69、OmniGen2 的 7.18、DreamOmni2 的 6.31、BAGEL 的 5.73。按子类看，Re-Align 在 MULTIPLE Character/Object/Char+Obj 分别为 8.25/8.07/8.28，在 SCENE Character/Object/Char+Obj 为 8.21/8.25/7.82；SINGLE Character/Object 为 8.25/8.55，仅 Object 不如 Qwen-Image-Edit-2509 的 9.13。

Model	SINGLE Char	SINGLE Obj	MULTI Char	MULTI Obj	MULTI Char+Obj	SCENE Char	SCENE Obj	SCENE Char+Obj	Avg
BAGEL	5.48	7.03	5.17	6.64	6.24	4.07	5.71	5.47	5.73
OmniGen2	8.05	7.58	7.11	7.13	7.45	6.38	6.71	7.04	7.18
Qwen-Image-Edit-2509	8.35	9.13	7.65	8.85	7.90	5.16	7.75	6.73	7.69
DreamOmni2	7.36	7.43	6.10	6.73	6.66	5.20	5.34	5.64	6.31
Re-Align	8.25	8.55	8.25	8.07	8.28	8.21	8.25	7.82	8.21

5.3 Quantitative results on DreamOmni2Bench

DreamOmni2Bench 上，Re-Align 在 editing 和 generation 两端都表现强。Editing：Add 的 PF/SC/Overall 为 9.27/9.27/9.27；Replace 为 8.44/8.81/8.61；Global 为 7.47/8.57/7.85；Local 为 6.11/8.54/6.35。Generation：PF/SC/Overall 为 7.74/7.67/7.24。相较 DreamOmni2*，Re-Align 在 Add、Replace、Local 和 Generation Overall 上均更高，Global Overall 7.85 也略高于 DreamOmni2* 的 7.76。

Model	Add Overall	Replace Overall	Global Overall	Local Overall	Gen Overall
BAGEL	4.58	1.37	2.46	1.63	5.25
OmniGen2	7.52	5.60	6.88	2.99	4.99
Echo-4o	8.51	4.51	5.16	2.41	6.59
Qwen-Image-Edit(2509)	6.51	2.79	3.21	2.73	5.45
DreamOmni2*	6.87	7.05	7.76	5.44	6.56
Re-Align	9.27	8.61	7.85	6.35	7.24

按 reference image 数量拆分时，Re-Align 的 editing with 2 refs 为 PF/SC/Overall 6.94/8.62/7.19；generation with 1/2/3/4 refs 的 Overall 分别为 6.37、7.93、7.70、7.39。它在 2 refs 和 3 refs generation 上均为 best overall；4 refs overall 7.39 低于 Echo-4o 的 7.75，但仍高于 DreamOmni2* 的 6.61。

5.4 Ablations: IC-CoT, RGA, RID

Figure 7 解读：GSB 评估比较 w/o CoT、BagelCoT 和 IC-CoT。论文文字给出的结论是 IC-CoT 相比 w/o CoT 和 BagelCoT 的 win rate 分别高 20% 和 16.25%，说明结构化 reasoning 比“无 reasoning”或“非结构化扩写”更适合 ICGE。

训练阶段/策略消融如下：无 SFT/RGA/RID 时 PF/SC/Overall/CLIP $_{o u t}$ 为 6.92/5.47/5.80/32.44；加入 SFT 后提升到 7.51/6.46/6.77/33.32；再加入 RGA 后 SC、Overall、CLIP $_{o u t}$ 到 6.54/6.80/33.50，但 PF 从 7.51 降到 7.46，说明低 sample diversity 会限制 RL；加入 RID 后达到 7.61/6.57/6.89/33.90，为最佳。

Figure 8 解读：可视化消融展示 RGA+RID 后图像更能反映 intended instruction。它对应 Table ablation 的解释：RGA 让 caption-image alignment 更强，RID 则避免 group rewards 过于接近导致优化信号不稳定。

5.5 Reasoning capability and limitations

论文还单独报告 reasoning capability：BAGEL 的 Score $_{e d i t}$ /Score $_{re f}$ 为 6.88/6.27，Re-Align (SFT) 为 7.05/6.92，Re-Align (SFT+RL) 为 7.02/6.91。这说明主要 reasoning 能力来自 IC-CoT SFT；RL alignment 并没有显著进一步提高 reasoning score，更多是把 reasoning 和 generation 对齐。

Figure 9 解读：失败案例显示两类限制：第一，复杂动作语义（如 “come here”）有时会导致 reasoning text 本身错误，进而影响图像；第二，对未专门训练的编辑类型，如 referenced text styles 或 object color schemes，模型能理解语义但生成图 reference consistency 低。作者认为扩大模型规模和引入更全面的数据可缓解这些问题。

总体结论：Re-Align 提供了一个比较直接的“理解帮助生成”证据链——把理解阶段的 IC-CoT 结构化后用于条件生成，再用 caption-image surrogate reward 对齐 reasoning 与 image output；实验上 SFT、RGA、RID 均有可量化贡献。不过它还没有充分证明“生成帮助理解”：reasoning score 在 SFT+RL 后略低于 SFT（7.02 vs 7.05，6.91 vs 6.92），说明 RL alignment 的主要收益在生成执行，而非反向提升理解能力。

Paper Notes

探索

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

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

1. Motivation（研究动机）

2. Idea（核心思想）

3. Method（方法）

3.1 Overall framework：BAGEL backbone + 两阶段训练

3.2 IC-CoT：semantic guidance + reference association

3.3 Reasoning-Generation Alignment：用 structured reasoning 做 surrogate reward

3.4 Re-Align-410K 数据构造

3.5 Pseudocode（基于论文流程；released repo 未提供实现代码）

3.6 Code-to-paper mapping

4. Experimental Setup（实验设置）

4.1 Training data and config

4.2 Baselines and benchmarks

5. Experimental Results（实验结果）

5.1 Qualitative results

5.2 Quantitative results on OmniContext

5.3 Quantitative results on DreamOmni2Bench

5.4 Ablations: IC-CoT, RGA, RID

5.5 Reasoning capability and limitations

目录