苹果提出UniVG：首个统一扩散模型！单一权重覆盖7大图像生成任务，性能竟超专项模型！

🌐 社群导航

🔗 点击加入➡️【AIGC/LLM/MLLM/3D/自动驾驶】 技术交流群

🔗 点击加入➡️【2-3月CCF投稿】交流群

最新论文解读系列

论文名：UniVG: A Generalist Diffusion Model for Unified Image Generation and Editing                   

论文链接：https://arxiv.org/pdf/2503.12652 

项目链接：暂无

导读

文本到图像（T2I）扩散模型在根据用户提示生成视觉效果引人注目的图像方面取得了令人印象深刻的成果。在此基础上，各种方法进一步针对特定任务对预训练的T2I模型进行微调。然而，这需要不同的模型架构、训练设计和多组参数来处理不同的任务。在本文中，我们介绍了UniVG，这是一个通用扩散模型，能够用一组权重支持多种图像生成任务。UniVG将多模态输入视为统一条件，以实现各种下游应用，从文本到图像生成、图像修复、基于指令的编辑、身份保留生成和布局引导生成，到深度估计和指称分割。通过对数据混合和多任务训练进行全面的实证研究，我们详细深入地了解了训练过程和决策，这些为我们的最终设计提供了依据。例如，我们表明文本到图像生成和其他任务（如基于指令的编辑）可以在不牺牲性能的情况下共存，而深度估计和指称分割等辅助任务可以增强图像编辑效果。值得注意的是，我们的模型在各自的基准测试中甚至可以超越一些特定任务的模型，这标志着向统一图像生成模型迈出了重要一步。

简介

扩散模型，特别是为文本到图像生成开发的模型，已经取得了重大进展。像Stable Diffusion、DALL - E和Imagen这样的模型已经展示了根据文本提示生成高质量、逼真图像的能力。同时，各种努力将扩散模型扩展到专门任务，产生了如InstructPix2Pix、ControlNet和InstandID等模型。然而，特定任务模型数量的增加导致了有效管理这些系统和优化计算资源方面的挑战。一个更具可扩展性的解决方案是一个单一的、统一的模型，能够处理多种图像生成任务，从而简化开发和部署。这种需求推动了社区对开发通用扩散模型的兴趣不断增长。

图2. 我们的通用视觉生成模型（UniVG）概述。UniVG包含一个文本编码器，用于从输入文本中提取提示嵌入，以及一个多模态DiT（MM-DiT），用于进行跨模态融合以实现潜在扩散。在这个过程中，所有视觉引导信息（潜在噪声、输入图像和输入掩码）都沿着通道维度连接成一个固定长度的序列，以提高效率。此外，可以通过嵌入替换注入外部条件，以实现进一步的控制。因此，通用的UniVG可以支持多种任务，如文本到图像生成、内外绘画、基于指令的编辑、布局引导生成和保留身份信息的生成。我们还考虑了辅助任务，包括深度估计、姿态估计和指代表分割，以增强其视觉场景感知能力。

方法

1. 背景

多模态扩散变压器（Multi-modal Diffusion Transformer，MM-DiT）。与通过U-Net实现的原始交叉注意力机制不同，MM-DiT利用扩散变压器（Diffusion Transformer，DiT）的结构化注意力机制来融合来自多个模态的特征，然后执行扩散模型的去噪过程。具体而言，它将所有输入连接成单个序列，并捕捉复杂的跨模态建模，以提高多模态生成的保真度和可控性。

流匹配（Flow Matching）。经典的去噪扩散模型，如去噪扩散概率模型（Denoising Diffusion Probabilistic Model，DDPM）和去噪扩散隐式模型（Denoising Diffusion Implicit Model，DDIM），在对复杂数据分布进行建模（如图像生成）方面表现出了良好的能力。流匹配不是逐步添加高斯噪声，而是学习连续时间变换。具体来说，一个与时间相关的向量场  处理噪声和数据之间的转换，这由从时间0到时间  的常微分方程（Ordinary Differential Equation，ODE）在  上控制。由于  通常难以处理，条件流匹配（Conditional Flow Matching，CFM）可以学习一个模型  来模拟理想的变换：

其中  是时间  时的示例 ， 是条件。因此，常微分方程求解器不需要大量离散的时间步，并且可以适应最有效的轨迹，从而以更少的步骤实现更高效的采样。

2. 统一视觉生成器（UniVG）

图2展示了统一视觉生成器（UniVG）的总体架构，它包含用于为输入提示  提取提示嵌入  的文本编码器，以及用于扩散建模的MM-DiT 。遵循潜在扩散学习，我们对输入图像  应用变分自编码器（Variational Autoencoder，VAE），并相应地调整二进制输入掩码  的大小。在训练期间，对输出图像  应用带有高斯噪声  的线性插值调度：

其中目标速度场为 。我们将潜在噪声  与视觉输入沿通道维度连接成等长序列，这是即使考虑多种引导也能实现高效率并减轻上下文感知干扰的关键。我们将提示嵌入结合起来，并根据流匹配损失优化 ：

为了添加额外的条件  以进行进一步控制，我们可以利用外部编码器  来提取特定领域的特征 ，这些特征应与  具有相同的隐藏维度大小。然后，我们通过替换预先设计的占位符标记的提示嵌入来注入这个外部条件。例如，面部特征  将替代 "

" 作为我们新的提示嵌入。根据公式3，MM-DiT可以考虑来自提示 、输入图像 、输入掩码  和外部条件  的所有引导，以实现多样化的控制。请注意， 不限于图像。任何格式的引导都可以通过其编码器进行条件设置。只要有多个占位符标记， 的长度也是灵活的。

推理。在我们的统一视觉生成器（UniVG）推理过程中，我们遵循无分类器引导（Classifier-Free Guidance，CFG）：

其中  是  的潜在特征， 是调整大小后的输入掩码，

是引导尺度。在去噪回到  后，我们利用  得到实际的图像生成结果。

3. 多任务训练

为了支持各种图像生成应用，我们考虑不同的任务，并为统一视觉生成器（UniVG）的多任务训练和推理制定每种输入格式如下（图2）。

文本到图像与内/外绘画。我们为文本到图像任务添加一个特殊的任务标记，其中输入图像  是一张空白（黑色）图像，输入掩码  全为真（白色），这意味着我们必须在此次生成中填充所有区域。对于内/外绘画，我们复用 ，但  是一张带有黑色块的图像， 有一个对应的白色块，用于控制模型绘制指定区域。在训练过程中，我们随机采样一个区域来遮挡图像，并将其标题作为引导提示 [43]。完整的图像即为理想的输出图像 。我们进一步考虑背景内/外绘画，在这种情况下，提示将被视为空字符串而被舍弃。

表1. UniVG多任务训练所用的混合数据。

基于指令的编辑。对于基于指令的编辑，和分别为输入图像和编辑后的图像。提示信息是前面带有 ie>的指令，其中为空白，因为所有区域均可编辑。

辅助任务。为了增强UniVG对视觉场景的理解，我们将深度估计、姿态估计和指称分割作为所用的辅助任务。我们没有采用结构化输出，而是遵循OneDiff的方法，直接将每个任务的可视化结果作为，通过图像生成进行学习。我们使用进行深度估计，进行姿态估计，进行指称分割（提示信息中带有目标：颜色），其中为输入图像，均为真。

布局引导的生成。关于更细粒度的控制，布局引导的生成要求模型在指定区域生成对象，给定的布局包含每个对象的边界框。我们将布局可视化为，并将对象信息注入提示信息中，例如“... 蓝色方块中的女孩。绿色方块中的足球。”通过这种方式，UniVG可以在全为真的情况下，为布局引导的生成提供足够的空间引导。这表明，在我们的设计中，输入图像不一定需要在视觉上与输出图像相似。

身份保留的生成。我们采用CLIP图像编码器为输入人脸提取面部嵌入。然后，我们用其替换占位符标记的提示嵌入，并输入到MM - DiT中。因此，UniVG会同时遵循输入人脸和字幕进行身份定制。具体来说，我们使用CLIP的最后一层，后面接一个两层的多层感知机（MLP）作为所用的外部编码器。

训练方案。基于我们的经验观察，我们介绍UniVG所用的多阶段训练方案：

• 阶段I（基础训练）：我们在文本到图像任务上从零开始训练MM - DiT，使用和批量大小，训练步；

• 阶段II（多任务训练）：我们进行内外绘画、基于指令的编辑、辅助任务、布局引导的生成以及文本到图像的多任务训练。详细的混合数据见第4.1节，我们同样使用和批量大小，训练步；

• 阶段III（进一步微调）：在发现如果在阶段II涉及身份保留生成会出现灾难性遗忘问题后，我们随后以1:1的比例将此身份定制任务与所有其他多任务数据一起进行训练。所用的外部图像编码器也与MM - DiT一起训练，训练，批量大小为，训练步。

表2. 在GenEval、T2I-CompBench、DSG和HPSv2上的文本到图像生成结果。

我们在4.3节进行了全面的消融研究。

实验

1. 实验设置

数据集。我们构建了一个数据集集合，以构建一个支持多样化任务的通用模型。对于文本到图像任务，我们有内部的20亿个文本 - 图像对、JourneyDB - 400万和DALLE3 - 100万。我们考虑图像内/外绘画的两种场景：文本引导和背景。我们利用我们的文本 - 图像对进行文本引导的内/外绘画；我们使用我们内部的场景图像、OSV - 500万和Places365 - 100万进行背景内/外绘画。

我们纳入了开源数据集，包括IPr2Pr - 100万、UltraEdit - 400万、SeedEdit - 300万、OmniEdit - 和StyleBooth - 11 用于我们基于指令的编辑。所有这些数据集都包含（输入图像、指令、输出图像）三元组。大多数图像对是从Prompt - to - Prompt或内绘画合成的。在我们的辅助任务中，我们考虑使用COCO - 11.8万、KITTI - 7000和Hypersim - 7.5万进行深度估计，使用COCO - 2.7万进行姿态估计，使用COCO - 21.3万、RefCOCO和PhraseCut - 29.8万进行指称分割。

为了支持更多的引导，我们纳入Flickr - 14.8万和SBU - 84万 [41]用于布局引导生成，并遵循GLIGEN [30]中的预处理方法来获取每个对象和相应的边界框。我们收集我们内部具有清晰人脸的图像用于身份保留生成，其中裁剪的人脸是输入身份，原始图像是目标输出。我们使用整个图像的标题作为输入提示。UniVG多任务训练（第二阶段）使用的混合数据集如表1所示。在第三阶段，我们将身份保留生成与所有其他任务的比例设置为1:1，以进一步学习身份定制并保持生成和编辑的原始能力。

实现细节。我们将内部的CLIP - bigG作为文本编码器，UniVG包含38层的MM - DiT，隐藏维度大小为2432，有38个注意力头，总共形成一个37亿参数的模型。我们应用一个内部的8通道变分自编码器（VAE）来提取图像的潜在特征以进行扩散建模。对于身份保留生成，我们使用CLIP图像编码器作为人脸的外部编码器。我们遵循相关方法，使用Adafator在512 - v5p张量处理单元（TPU）上训练UniVG。在推理过程中，我们将引导尺度设置为(4.0, 1.5)。所有实现均使用AXLearn框架完成。

表3. 基于指令的编辑在MagicBrush和EmuEdit上的结果。

2. 评估结果

文本到图像生成。由于研究文本到图像质量的方面众多，我们采用GenEval和T2I - CompBench来评估组合性，如物体、纹理、颜色和相对位置。我们还纳入了DSG，它利用视觉问答来验证生成的图像是否与提示相符。从HPSv2中，我们根据预训练的人类偏好分数来研究视觉质量。我们将当代统一模型OneDiff和OmniGen作为主要基线。此外，我们还将SDXL、FLUX.1和SD3作为特定的文本到图像方法进行比较。表2显示，我们的UniVG在所有基准测试的组合性和语义方面都超越了那些统一基线。这些结果表明，我们可以精确遵循提示，同时生成视觉上吸引人的文本到图像。值得注意的是，UniVG的参数甚至比OmniGen更少，这凸显了我们模型设计和多任务训练的优势。此外，我们与特定任务的SD3取得了具有竞争力的性能，而SD3的参数是我们的两倍。这也证明了通用模型的潜力，即我们可以保持良好处理各种任务的能力。

基于指令的编辑。为了评估基于指令的编辑，我们考虑了MagicBrush和EmuEdit的测试集。我们使用目标标题和编辑后图像之间的CLIP - T进行文本 - 视觉对齐。我们还应用CLIP - I来研究输入保留情况。我们将InsP2P、MGIE、EmuEdit作为特定的编辑模型。表3将我们的UniVG与基线进行了比较。令人惊讶的是，我们的CLIP - T得分甚至明显高于特定任务方法，这表明我们在遵循指令修改现有图像方面具有强大的能力。与统一的OmniGen相比，UniVG在CLIP - T和CLIP - I方面具有全面优势，这两个指标是基于指令编辑的两个权衡目标。虽然OneDiff的CLIP - I得分较高，但其修改通常有限，导致CLIP - T得分较低，无法满足编辑期望。例如，在图4中，OneDiff未能在天花板上添加圆形灯光。MGIE和OmniGen进行了尝试，但结果在视觉上并不吸引人。相比之下，UniVG遵循相同的视觉流程进行编辑。此外，我们是唯一能够实现复杂修改的模型，如用“绿草包装纸”或“黑色边框”进行替换。我们的UniVG还支持多种用途，包括移除、面部表情和整体艺术风格化。这些定性结果说明了UniVG在通用基于指令的图像编辑方面的优势。

图3. 文本到图像生成的定性示例。请注意，为了更好地展示，我们简化了提示。

表4. 在Unsplash - 50数据集上的身份保留生成结果。

表5. 一次性进行训练会导致身份信息保留性能变差。

身份信息保留生成。我们采用Unsplash - 50数据集[12]来评估身份信息保留生成能力，该数据集提供人脸图像和描述信息，以使模型生成个性化图像。然后，我们遵循CurricularFace方法[25]来计算用于衡量身份相似度的人脸嵌入特征，并使用CLIP - T来计算提示 - 图像得分。我们将PhotoMaker[31]、InstantID[56]和PuLID[17]作为特定任务模型。在表4中，UniVG再次超越了统一基线模型，在人脸一致性方面有显著提升，这凸显了我们注入灵活引导并控制后续生成过程的设计的有效性。此外，较高的CLIP - T得分表明我们在身份定制方面也能更好地遵循提示信息。与特定任务方法相比，InstantID和PuLID依赖预训练的人脸编码器来实现较强的身份信息保留。然而，它们生成的人脸严重受限于输入，无法支持复杂的操作[28]。表5强调了我们精心设计的多阶段训练的重要性，一次性训练会导致身份信息保留出现灾难性遗忘（身份相似度显著降低）。

3. 消融实验

表6展示了我们多任务训练的消融实验结果。与(a)行相比，(b)行显示出强大的基于指令的编辑能力，同时在文本到图像生成任务中也保持了有竞争力的性能。这表明同时学习这两项任务不会对任何一项造成损害，反而能使单个模型具备这两种能力。(c)行接着进行身份信息保留生成的训练。然而，编辑能力出现了显著下降（例如，在MagicBrush上CLIP - T得分从28.3降至26.9），因为模型更加侧重于身份定制。为了克服这个问题，我们引入辅助任务来增强UniVG的视觉理解能力。(d)行在编辑能力上比(b)行有进一步提升，这凸显了辅助任务的作用。这一次，即使进行身份信息保留生成训练，(e)行在各项性能上达到了最佳平衡。

图4. 基于指令的编辑的定性比较。

表6. 多任务训练的消融研究。 文本到图像生成与内外绘画；♥ 基于指令的编辑； 辅助任务与布局引导生成； 身份保留生成。我们的方案：阶段I（♠）；阶段II（♠+♥+♦）；阶段III（♠+♥+♦+♦+·）。

表7. 推理期间时间（秒）和GPU成本（MB）的效率比较结果 。

推理效率。除了生成质量，我们还在表7中研究了在单个NVIDIA A100 GPU上的推理时间和GPU成本。此比较是在图像分辨率  和模型精度BFloat16的条件下进行的。

图5. 模型扩展结果，包括B（4.16亿参数）、L（18亿参数）和XL（37亿参数）；MagicBrush使用CLIP-T，Unsplash - 50使用ID。

OneDiff和OmniGen都不得不将额外的图像与噪声序列拼接，这显著增加了MM - DiT的计算开销，导致编辑效率显著下降（例如，OmniGen需要36秒以上）。另一方面，我们的UniVG在通道维度上同时考虑潜在噪声、输入图像和掩码，这可以在生成和编辑过程中保持相同的总序列长度。因此，尽管模型更大，但我们在编辑方面甚至为OneDiff带来了更高的效率。这些观察结果也凸显了

图6. 布局引导生成的定性示例，其中方块的颜色与提示中的对象相对应。

图7. 背景内/外绘画和文本引导内/外绘画的定性示例。

UniVG不仅能实现卓越的图像生成，而且在各种任务中始终保持高效率。

模型扩展。我们研究了统一图像生成的模型扩展性能。我们考虑了三种模型大小，包括B（4.16亿参数，18层）、L（18亿参数，30层）和XL（37亿参数，38层）。图5表明，随着模型规模的扩大，性能不断提升，并且适用于生成和编辑任务。这显示了未来更强大、更有能力的模型的巨大潜力。

更多可视化结果。图6展示了布局引导生成，其中实际提示与布局相结合（例如，“一瓶酒旁边有一杯酒。蓝色方块里是酒。绿色方块里是瓶子。”）。UniVG遵循每个对象的空间引导来生成图像。我们还可以处理计数（例如，两个甜甜圈）并将它们呈现在指定位置。

图7显示了内/外绘画的结果。即使没有提示，我们的UniVG仍然可以恢复缺失的区域（例如，汽车和人脸）。我们甚至可以从一个小方块想象出整个视觉场景，并进行外绘画以将其扩展为合理的图像（例如，鹦鹉和街景）。UniVG还可以通过文本进行控制，我们可以在内绘画中绘制特定的文本（例如，“限时优惠”）或对象（例如，侦探）。同样，我们可以引导和外绘画整个图像以支持创意图像完成。

图8. 辅助任务的定性示例，包括深度估计、姿态估计和指称分割。

图8表明，我们的UniVG还支持计算机视觉在辅助任务中的应用，如深度估计和姿态估计。关于指称分割，我们可以识别精确的对象（例如，汽车和带球的人），甚至可以分割出背景（例如，大海）。这一发现也展示了将我们的模型用于统一视觉理解和生成的潜力。

总结

在本文中，我们介绍了UniVG，这是一个通用扩散模型，它在单一框架内统一了多种图像生成任务。我们进行了全面的研究，并对关键的建模和数据选择进行了消融实验，采用了简约的模型架构设计，并引入了三阶段训练流程。我们的结果表明，单个模型可以有效地处理所有任务，而不会影响核心的文本到图像生成性能。对于未来的工作，我们计划将更多的视觉感知任务纳入该框架，并进一步扩大模型规模，以训练出更强大的通用模型。致谢。我们感谢陆家森、江一方和许多其他人提供的宝贵帮助和反馈。