OmniGen：统一图像生成模型

一、研究背景与动机

随着大型语言模型（LLMs）的兴起，自然语言处理领域已经实现了多种语言任务的统一处理，极大地推动了人机交互的发展。然而在图像生成领域，尚未出现能够在一个框架内处理多种任务的通用模型。当前的图像生成模型大多专注于特定任务，如文本到图像生成、图像编辑等，每种任务都需要设计专门的模块并进行微调，这不仅限制了图像生成技术的发展，还导致了繁琐的工作流程。

北京人工智能研究院的研究团队提出 OmniGen 模型，旨在实现图像生成任务的统一处理。OmniGen 基于扩散模型构建，能够在单一框架内支持多种图像生成任务，包括文本到图像生成、图像编辑、主体驱动生成和视觉条件生成等。通过简化架构并采用统一的训练格式，OmniGen 实现了不同任务间的知识迁移，能够处理未见过的任务和领域，并展现出全新的能力。

二、模型架构与训练方法

1. 架构设计

OmniGen 的设计遵循通用性和简洁性原则，主要由变分自编码器（VAE）和预训练大型变换器模型组成。其中，VAE 用于从图像中提取连续的视觉特征，而变换器模型基于指令条件生成图像。与现有扩散模型不同，OmniGen 不依赖额外的编码器预处理条件信息，而是自行编码条件输入，显著简化了处理流程。

OmniGen 支持任意交错的文本和图像输入作为条件，以指导图像生成。文本通过 Phi-3 的分词器处理，图像则通过 VAE 和线性层提取潜在表示，再嵌入到序列中。每个图像序列被特殊标记包裹后插入文本标记序列，并附加时间步嵌入，无需任务特定标记。

2. 注意力机制

考虑到图像的整体性，OmniGen 对常见的因果注意力机制进行修改，将双向注意力整合其中。具体来说，序列中的每个元素应用因果注意力，而每个图像序列内的元素应用双向注意力。这种设计允许图像中的每个 patch 关注同一图像中的其他 patch，同时确保图像只能关注之前出现的其他图像或文本序列。

3. 推理过程

推理时，OmniGen 随机采样高斯噪声，通过流匹配方法预测目标速度，经过多次迭代获得最终潜在表示，再由 VAE 解码器解码为预测图像。得益于注意力机制，OmniGen 可以像 LLMs 一样通过使用 kv-cache 加速推理，存储输入条件的键和值状态以避免冗余计算。

4. 训练策略

OmniGen 使用修正流优化模型参数，与 DDPM 不同，修正流通过线性插值在噪声和数据之间进行前向过程。在训练过程中，模型被训练为直接回归目标速度，以最小化均方误差损失。

对于图像编辑任务，输入图像和目标图像之间的差异通常较小，可能导致模型学习到直接复制输入图像的捷径。为此，研究团队在发生更改的图像区域放大损失权重，使这些区域的权重显著高于未更改区域，引导模型专注于修改区域。

训练过程中，图像分辨率逐步增加，初始阶段使用低分辨率提高数据效率，后期使用高分辨率提升生成图像的美学质量。每个训练阶段的详细信息如下表所示：

阶段	分辨率	步骤（K）	批量大小	学习率
1	256×256	500	1040	1e-4
2	512×512	300	520	1e-4
3	1024×1024	100	208	4e-5
4	2240×2240	30	104	2e-5
5	多分辨率	80	104	2e-5

三、X2I 数据集构建

为了实现多任务处理能力，研究团队构建了大规模统一图像生成数据集 X2I（“anything to image”），包含约 1 亿张图像，涵盖多种任务并统一为标准格式。

1. 文本到图像任务数据

从多个开源数据集收集纯文本输入的数据，包括 RecapDataComp、SAM-LLaVA、ShareGPT4V 等。这些数据用于学习广泛的图像-文本匹配关系和多样化的知识。后期使用内部收集的 1600 万张高质量图像提升生成图像的美学质量。

2. 多模态到图像任务数据

OmniGen 支持灵活的多模态指令输入，适应多种数据和任务类型。

（1）通用混合模态提示

数据来源包括图像编辑、人体运动、虚拟试穿和风格转移等任务。通过 MultiGen 数据集学习利用额外的视觉条件进行更精细的空间控制，并选择六种代表性视觉条件：Canny、HED、深度、骨架、边界框和分割图。

（2）主体驱动图像生成

构建了大规模基础数据集（GRIT 实体数据集）和高质量高级数据集（网络图像数据集）。利用 GRIT 数据集的标注信息，结合 GroundingDINO 和 SAM 技术获取对象掩码，并通过 MS-Diffusion 提升数据质量和多样性。同时，从网络收集同一人物的不同图像对，以增强主体驱动图像生成能力。

（3）计算机视觉任务

引入经典计算机视觉任务扩展模型的生成能力。对于低级视觉任务，如弱光图像增强、去雨、去模糊等，仅添加文本指令；对于高级视觉任务，将所有标注表示为图像，如人体姿态、深度映射等。

3. 少样本到图像任务数据

构建少样本到图像数据集，以激发模型的上下文学习能力。对于每个任务，随机选择几个示例与原始输入组合形成新输入。

四、实验结果与分析

1. 图像生成任务

（1）定性结果

OmniGen 在多种图像生成任务中表现出色，包括文本到图像生成、图像编辑和视觉条件生成等。例如，在主体驱动生成任务中，OmniGen 能够从参考图像中提取所需对象并生成新图像，无需额外预处理步骤，如图像裁剪或人脸识别。

（2）文本到图像生成评估

在 GenEval 基准测试中，OmniGen 与当前流行的图像生成模型（如 SD3）性能相当，尽管其参数规模较小（仅 38 亿参数，而 SD3 模型总参数达 127 亿）。OmniGen 的架构显著简化，消除了额外编码器的成本，极大地提升了参数利用效率。

2. 多模态到图像任务

在 EMU-Edit 数据集和 DreamBench 数据集上的评估表明，OmniGen 在图像编辑和主体驱动生成任务中表现优异，与当前最先进的模型相当。例如，在主体驱动生成任务中，OmniGen 使用特定对象的一张图像作为输入，无需像 DreamBooth 那样对模型进行微调。

3. 计算机视觉任务

OmniGen 在多种传统视觉任务中展现出能力，如去雨、去模糊、修复等低级视觉任务，以及人体姿态识别等高级视觉任务。尽管 OmniGen 在这些任务中的目标并非超越专门开发的现有模型，但其能够将从传统计算机视觉任务中学到的知识迁移到图像生成任务中，释放更大潜力。

五、进一步分析

1. 新兴能力

通过将所有任务标准化为统一格式并在 X2I 数据集上训练，OmniGen 获得了通用知识，能够跨不同场景和任务进行知识迁移，从而实现对未见任务和领域的生成能力。

（1）任务组合

OmniGen 能够同时处理多个指令，包括不同任务的指令以及同一任务的多个指令，凸显了其多功能性和广泛应用潜力。

（2）端到端工作流

OmniGen 具备出色的多模态理解和图像生成能力，无需依赖外部模型即可完成多种任务，显著简化工作流程并节省成本。例如，OmniGen 可直接从参考图像提取条件信息，并基于捕获的条件在单一步骤中生成新图像。

（3）未见领域的上下文学习

在未见领域的评估中，OmniGen 展现了泛化能力。当提供示例后，模型能够成功完成未见分割任务。

2. 推理能力

研究团队探索了 OmniGen 的推理能力。当给定未明确指定对象的指令时，模型能够识别图像内容并推断所需对象。例如，模型可在图像中识别水槽区域并进行标记，为具身智能领域提供潜在应用，如协助智能代理理解多模态指令、定位必要对象并规划后续动作。

3. 链式思维

受人类绘画基本方式的启发，研究团队尝试通过逐步生成的方式模拟人类艺术家的绘画过程。尽管最终生成图像的质量未超越原始模型，但这种方法为图像生成提供了新思路。团队认为，对图像绘制过程进行监督是一个有前景的方向，可能帮助模型处理更复杂和多样化的场景。

六、局限性与未来方向

尽管 OmniGen 实现了多种任务的图像生成能力，但仍存在一些局限性：

1. 文本渲染能力有限，无法准确生成长文本。

2. 输出图像可能包含不期望的细节，如异常的手部。

3. 对于未见过的图像类型（如表面法线图），处理效果不佳。

研究团队认为，通过在更多相关数据上训练模型，大多数局限性可以得到解决。此外，由于 OmniGen 支持多模态输入，用户无需为特定任务设计网络，只需准备数据即可对模型进行微调，这为下游任务的应用提供了广泛的可能性。

七、核心技术总结表