新加坡国立大学提出 OminiControl: 扩散变换器 (Diffusion Transformer) 的最小通用控制 ！

在本文中，作者提出了OminiControl，这是一个高度灵活且参数高效的框架，它将图像条件集成到预训练的Diffusion Transformer（DiT）模型中。 OminiControl的核心是利用参数重用机制，使DiT能够使用自己作为强大的 Backbone 来编码图像条件，并使用其灵活的多模态注意力处理器来处理这些条件。 与现有方法不同，OminiControl（1）仅需额外约0.1%的参数，就能有效地且高效地将注入的图像条件集成进去，并（2）以统一的方式解决广泛的图像条件任务，包括由主题驱动的生成以及空间对齐的条件，如边缘、深度等。 值得注意的是，这些能力是通过在DiT生成的图像上进行训练实现的，这对于由主题驱动的生成特别有利。 大量评估表明，OminiControl在由主题驱动的生成和空间对齐条件生成方面，都优于现有的UNet基础模型和DiT调整模型。 此外，作者还发布了作者的训练数据集Subjects200K，这是一个包含超过20万张身份一致的图像的多样化集合，以及一个高效的数据合成 Pipeline ，以推动主体一致性生成的研究。

1 Introduction

扩散模型[9, 25, 28]在视觉生成领域实现了革命性的变革，展示了在图像质量和多样性方面显著优于传统方法如生成对抗网络（GANs）[6]的卓越能力。尽管这些模型在生成高度逼真的图像方面表现出色，但仍然存在一个关键挑战：实现对生成过程的精准和灵活控制，以满足多样化和复杂化的用户需求。

文本为基础的条件作用在推进可控制生成方面发挥了关键作用，为用户提供了一种直观的方式来指定他们想要的输出。然而，仅使用文本 Prompt 往往无法传达用户希望控制的精确空间细节和结构属性。因此，最近的研究探索了用于指导扩散模型的补充条件作用模式，其中基于图像的控制作为一种特别有效的方法出现。这种多模态条件作用策略使得对生成过程的控制更加详细和准确，解决了纯文本界面固有的限制。

当前的图像预处理方法可以广泛地分为空间对齐和非空间对齐两大类。空间对齐任务如从草图到图像生成和修复，要求条件和输出图像之间存在直接的对应关系，通常通过在保持空间不变的方式下将条件特征注入，例如使用ControlNet[41]等方法。相比之下，非空间对齐的应用包括基于主题的生成和风格迁移，例如由IP-Adapter[39]所展示，通常会使用预训练的编码器如CLIP[27]通过交叉注意力机制来整合全局特征。

尽管现有的基于图像的方法具有有效性，但它们存在一些局限性，这些局限性阻碍了它们的效率和灵活性。首先，大多数现有方法都是专门针对UNet结构设计的，正如Stable Diffusion模型所展示的那样。这些方法在UNet的编码器-解码器结构中表现良好，但可能无法有效地转化为更先进的Diffusion Transformer（DiT）模型[23]，这些模型已经展示了卓越的图像生成质量。此外，当前的方法通常专注于空间对齐[22, 41, 43]或非空间对齐任务，缺乏一种统一的架构来有效处理两种控制类型。这种专业化通常需要实践者针对不同的控制场景采用不同的方法，增加了系统的复杂性和实现开销。此外，这些方法依赖于额外的网络结构，这引入了大量的参数开销。

为了克服这些限制，作者提出了一种参数高效的将基于图像的控制方法集成到DiT架构[23]中的方法。作者的方法利用模型现有的VAE解码器[28]来处理条件图像。遵循与噪声图像 Token 相同的 Token 处理 Pipeline ，作者使用可学习的位置嵌入[34]增强编码特征，并将它们与潜在噪声无缝地集成到去噪网络中。这种设计使得条件 Token 和生成 Token 在整个DiT的 Transformer 块之间可以进行直接的多模态注意力交互[23, 30]，从而促进有效信息交换和控制信号传播。

作者在具有120亿参数的大规模模型FLUX.1-dev[13]上实现了作者的方法。这是一种性能强大的基于分层的扩散模型。在边缘引导生成、深度感知合成、区域特定编辑和身份保护生成等方面的广泛实验表明，与基于UNet的实现[7, 41]及其在FLUX.1模型[14, 37]上的社区调整相比，作者的基于DIT的方法取得了更好的结果。

为了保持身份一致的生成，作者开发了一种新颖的数据合成 Pipeline ，该 Pipeline 可以生成高质量的身份一致图像对。利用这种 Pipeline ，作者创建了一个包含超过20万个多样化图像的综合数据集。为了促进这一方向的未来研究，作者将发布作者的数据集和完整的 Pipeline 实现作为开源资源2。

总的来说，作者在此突出作者的贡献如下：

1. 作者提出了一种参数高效的在扩散 Transformer （DiT）模型中实现图像条件控制的方法，实现了在统一的框架内实现空间对齐和非空间对齐控制。
2. 作者通过在各种控制任务上的广泛实验来证明作者方法的有效性，包括边缘引导生成、深度感知合成、区域特定编辑和身份保护生成，在这些任务上始终优于现有方法，无论是UNet实现还是其DiT改编。
3. 作者发布了Subjects200K，这是一个包含超过20万张主题一致图像的高质量数据集，以及一个高效的数据合成流水线。这些资源将有助于研究界进一步探索主题一致生成任务。

2 Related works

Diffusion-based models

扩散基础的方法已经崛起为一个强大的图像生成框架[9, 28]，在各种任务上取得了成功，包括文本到图像合成[2, 28, 33]，图像到图像转换[32]，以及图像编辑[1, 20]。最近的进展显著提高了质量和效率，特别是在引入潜在扩散模型[28]方面。为了进一步提高生成能力，大规模的 Transformer 架构已经被整合到这些框架中，导致了像DiT[2, 3, 13, 23]这样的高级模型。在这些架构创新的基础上，FLUX[13]将基于 Transformer 的结构与 Stream 匹配目标[18]相结合，实现了最先进的生成性能。

Controllable generation with diffusion models

可控生成在扩散模型背景下得到了广泛研究。文本到图像模型[25, 28]奠定了条件生成的基础，同时开发了各种方法以包含额外的控制信号，如图像。显著的方法包括ControlNet [41]，使扩散模型实现空间对齐控制，以及T2I-Adapter [22]，使用轻量级 Adapter 提高效率。UniControl [26]使用Mixture-of-Experts (MoE)统一不同的空间条件，进一步减小模型大小。然而，这些方法依赖于将条件信息空间添加到去噪网络的隐藏状态，固有地限制了它们在非空间任务（如主题驱动生成）的有效性。IP-Adapter [39]通过引入额外的编码器通过交叉注意力解决这一问题，而SSR-Encoder [42]在图像条件任务中进一步提高了身份保护。尽管取得了这些进展 [5, 15, 19]，但同时实现空间对齐和非对齐任务的一体化解决方案仍然是一个未解之谜。

3 Methods

Preliminary

每个 Transformer 块处理两种类型的 Token ：噪声图像 Token 和文本条件 Token ，其中是嵌入维度，和分别表示图像和文本 Token 的数量（如图2所示）。这些 Token 被嵌入到隐藏状态和中，在 Transformer 块的整个过程中保持一致的形状。

在每个DiT块中，首先对和进行归一化处理，然后通过核心MM-Attention模块[30]进行处理。该模块利用旋转向量位置嵌入（RoPE）[34]来整合跨 Token 位置的依赖关系。对于二维网格中位置为的 Token ，RoPE对 Query 和键投影应用旋转变换矩阵：

在位置处的旋转矩阵为。同样地，文本条件 Token 的 Query 和键投影也按照相同的方式定义，所有文本 Token 位置在FLUX中均设置为。1.

在应用RoPE之后，两种token类型的 Query （queries）、键（keys）和值（values）被拼接形成统一的矩阵、和，代表组合的token集合。接下来计算MM-Attention操作：

借助注意力机制实现图像与条件 Token 之间的互动。

Image condition integration

作者的方法首先通过模型的VAE对条件图像进行编码，将其映射到与噪声图像 Token 相同的潜在空间，从而形成 。

之前的方法如ControlNet [41] 和 T2I-Adapter [22] 通过对条件图像进行空间对齐并将其隐藏状态直接添加到噪声图像 Token 中：

其中表示进一步处理的联合隐藏状态，其中来自条件图像的隐藏状态。虽然这种方法对于空间对齐任务是有效的，但它面临两个关键限制：(1) 在非对齐场景下，它处理起来缺乏灵活性，(2) 即使在空间对齐的情况下，直接添加隐藏状态也会限制 Token 之间的交互，可能限制模型的性能。

相比之下，为了实现非对齐控制任务并提供更大的条件灵活性，作者的方法将条件图像 Token 、文本 Token 和噪声图像 Token 统一处理，并将其集成到统一序列中：

代表噪声图像 Token 、文本 Token 和条件图像 Token 的连接序列。这种统一方法使得可以直接参与多模态注意力 [30]，而无需专门的处理路径（如图2 所示）。

比较结果表明，与直接添加方法相比，作者的方法在生成质量和与条件的对齐方面实现了更高的性能，如图3所示。此外，图3(a)中的训练曲线显示，多模态注意力方法始终实现了比直接添加方法更低的损失值。此外，这种统一的序列方法在空间对齐和非空间对齐任务上都得到了验证（图5），突显了其在处理多样化条件生成场景的灵活性。

Adaptive position embedding

作者的统一序列设计允许条件图像 Token 的灵活集成，但这需要融入位置信息以确保与噪声图像 Token 的有效互动。这些 Token 的相对位置至关重要，因为它直接影响着模型的学习效率和泛化能力。

在FLUX.1的Transformers中，每个 Token 都被分配一个相应的位置索引以编码空间信息。对于一个512×512的目标图像，VAE[11]编码器首先将其映射到潜在空间，然后将潜在表示分为一个32×32的 Token 网格，其中每个 Token 都分配一个唯一的二维位置索引（i,j），其中i,j∈[0,31]。这种索引方案保留了原始图像在潜在空间中的空间结构，而文本 Token 保持固定的位置索引（0,0）。

对于空间对齐任务，作者的初始方法是将条件 Token 分配与噪声图像中相应 Token 相同的位置嵌入。然而，对于非空间对齐任务，如主题驱动生成，作者的实验表明，将条件 Token 的位置索引进行调整可以实现更快的收敛（图3(b)）。具体而言，作者将条件图像 Token 移动到索引，其中且，确保与原始图像 Token 之间没有空间重叠。

Condition strength factor

我国采用的统一注意力机制不仅实现了灵活的词元交互，还使作者能够精确控制条件图像的影响。具体而言，作者设计了一种方法，在推理过程中可以手动调整条件图像的效果。对于给定的强度因子γ，当γ等于0时，条件图像的影响被移除，输出仅基于原始输入。当γ等于1时，输出完全反映了条件图像，随着γ的增加超过1，条件的影响变得更加明显。

为了实现可控性，作者在原始的MM-Attention操作中引入了一个偏置项。具体来说，作者对公式3进行了修改：

在这里，偏差γ被设计为根据强度因子γ调整条件 Token 和其他 Token 之间的注意力权重。偏差项被构建为一个（M+2N）×（M+2N）的矩阵，其中M表示文本 Token 的数量，N表示每个噪声图像 Token 和条件图像 Token 的数量。该矩阵具有以下结构：

该设计确保了强度因子γ仅影响噪声图像 Token 与条件图像 Token 之间的注意力权重，同时保持文本 Token 和同质性交互的原有注意力模式。

Subjects200K datasets

为了训练具有主题一致生成的模型，通常需要成对的图像，这些图像在保持身份一致的同时，表现出姿势、光照和其他属性的变化。以前的方法，如IP-Adapter [39]，使用相同的图像进行条件和目标对齐，这对其方法是有效的。然而，在作者的框架中，这种设置导致过拟合，使得模型生成的输出几乎与输入相同。

为了克服这些限制，作者开发了一个数据集，其中包含保持主体身份的同时融入自然差异的图像。尽管现有的数据集[12, 15, 17, 31]解决了类似的需求，但它们往往在质量或规模上受到限制。因此，作者提出了一种新颖的合成 Pipeline ，利用FLUX固有的从精心设计的 Prompt 中生成视觉上相关的图像对的能力。

作者使用的 Pipeline 利用ChatGPT-4o生成超过20,000种不同的图像描述，这些描述引导FLUX生成超过200,000张图像（图6）。生成的图像通过ChatGPT-4o的视觉能力进行质量评估，以确保最终数据集的一致性和多样性。精选的数据集和完整的生成 Pipeline 都已公开3。

4 Experiment

Setup

任务与基础模型： 作者在两个条件生成任务类别上评估作者的方法：空间对齐任务（包括Canny-to-image，depth-to-image，masked-based inpainting和colorization）和主题驱动生成。作者将作者的方法构建在FLUX.1 [13]上，这是一个用于图像生成的潜在正则化流 Transformer 模型。默认为在空间对齐任务中使用FLUX.1-dev生成图像。在主题驱动生成任务中，作者切换到FLUX.1-schnell，因为作者观察到它倾向于产生更好的视觉质量。

实现细节作者的方法利用LoRA[4]对基础模型进行微调，默认为4级排序。为了保留模型的原始能力和实现灵活性，在处理非条件 Token 时，默认为将LoRA尺度设置为0。

训练。 作者的模型以批量大小为1进行训练，梯度累积为8步（实际批量大小为8）。作者使用Prodigy优化器[21]，启用安全加温和偏置校正，并将权重衰减设置为0.01。实验在2个NVIDIA H100 GPU（每个80GB）上进行。对于空间对齐任务，模型进行50,000次迭代训练，而受主题驱动的生成模型进行15,000次迭代训练。

Baseline 。 对于空间对齐任务，作者将方法与原始ControlNet [41]和T2I-Adapter [22]在Stable Diffusion 1.5上进行比较，以及与ControlNetPro [14]和ControlNet的FLUX.1实现进行比较。对于主题驱动生成，作者与IP-Adapter [39]进行比较，评估其实现FLUX.1 [37]。

评价指标。作者在空间对齐任务和基于主题的生成上评估作者的模型。对于空间对齐任务，作者评估两个方面：生成质量和可控性。生成质量通过FID [8], SSIM, MAN-IQA [38], 和MUSIQ [10]对视觉保真度进行测量，同时使用CLIP Score [27]对语义一致性进行评估。对于可控性，作者在边缘条件生成中计算提取和输入边缘图之间的F1得分，以及对于其他任务（如深度 Anything 和颜色通道分离用于颜色增强等）计算提取和原始条件图之间的均方误差。对于基于主题的生成，作者提出一个五标准框架来评估主题特性的保持（身份保持，材质质量，颜色保真度，自然外观）以及请求修改的准确性，所有评估都通过GPT-4o的视觉能力进行以确保系统的评估。详细的评估方法论见附录B.1。

评估协议：作者在两个数据集上进行了评估。对于空间对齐任务，作者使用COCO 2017验证集（5,000张图像），将其重置为512512，使用任务特定的条件和关联的描述作为 Prompt ，固定种子为42。对于主导性生成，作者在DreamBooth[31]数据集的750个文本条件对（30个主题×25个 Prompt ）上进行测试，使用5个不同的种子，每个主题选择一张图像作为条件。

Main result

空间对齐任务如表1所示，作者在五个空间对齐任务上全面评估了作者的方法与现有方法的性能。在深度到图像生成任务上，作者的方法获得了最高的F1分数0.38，显著优于基于SD1.5的控制网络（ControlNet）[41]、T2I-Adapter[22]和基于FLUX.1的控制网络Pro（ControlNetPro）[14]。在总体质量指标方面，作者的方法在大多数任务上展现出一致的优势，尤其是在SSIM[36]、MAN-IQA[38]和MUSIQ[10]得分上，表现尤为出色。对于具有挑战性的任务如去模糊和色彩化，作者的方法取得了显著的改进：与ControlNetPro相比，均方误差（MSE）分别降低了77%和93%，而FID分数[8]从30.38提高到11.49（去模糊）。

CLIP-Score指标[27]表明，作者的方法在所有任务上保持了高文本到图像的一致性，说明在实现更好的控制和视觉质量的同时，有效保留了语义对齐。如图7所示，作者的方法在色彩化任务中产生了更锐利的细节和更忠实于原图的色彩再现，同时保持了边缘引导生成的结构和去模糊场景的更高保真度。

主题驱动生成 如图8所示，作者对现有 Baseline 进行了全面的比较。作者的方法在身份保护和修改准确性方面表现优越。对随机种子求平均，作者的方法与IP-Adapter（FLUX）相比，修改准确性达到了75.8%，而IP-Adapter（SD 1.5）的57.7%则保持50.6%的身份保护。在最佳种子场景下，作者的方法将修改准确性提高到90.7%，并将身份保护保持在82.3%，比最强的 Baseline 高出15.8和18.0个百分点，这证明了有效的主题忠诚度编辑。这些定量结果得到了附录B.1中呈现的用户研究的支持。

比较参数效率。如表2所示，与现有方法相比，作者的方法在参数效率方面取得了显著的提升。对于12B参数的FLUX.1模型，作者的方法只需要14.5M个可训练参数（约占0.1%），这明显低于ControlNet（27.5%）和IP-Adapter（7.6%）。即使使用FLUX.1的原始VAE编码器，作者的方法在额外增加0.4%的参数时仍能保持高效率，这充分展示了作者参数高效设计的有效性。

Empirical studies

训练数据的影响。 对于主题驱动的生成，作者的模型以一个主题的参考图像（例如，一个毛绒玩具或一个物体）和文本描述作为输入，试图根据文本指导生成相同主题的新图像，同时保留其关键特征。

为了验证作者在第3.5节中描述的Subjects200K数据集的有效性，作者对比了两种训练策略。第一种方法依赖于传统数据增强，其中作者对原始图像进行随机裁剪、旋转、缩放和对比度、饱和度和颜色的调整。第二种方法利用了作者 Subjects200K 数据集。如图9所示，仅使用数据增强训练的模型仅学习在最小更改的情况下复制输入条件。在第一行中，它只是将薯条布偶放在一个明亮的房间环境中，同时保持其精确的外观和姿势。同样，在第二行中，尽管黄色闹钟被放在窗户一侧，但模型仍几乎以相同的细节复制了它。相比之下，作者的 Subjects200K 训练的模型展示了在遵循文本 Prompt 的同时生成多样且一致的主体视图的能力。

条件强度分析。 作者通过定性实验评估了条件强度控制。图10显示了不同强度因子γ（0.25，0.5，0.75，1.0）生成的结果。结果表明，γ有效地控制了生成过程，无论是空间对齐的任务（如深度到图像生成）还是非空间对齐的任务（如主题驱动生成），都能实现条件影响的灵活控制。

LoRA rank的Impact 作者针对Canny图像任务进行了大量实验，实验中使用了不同的LoRA排名（1，2，4，8，16）。如表3所示，作者的实验表明，随着LoRA排名的增加，模型性能总体上得到提高，排名16在多个方面取得了最佳结果：图像质量（通过FID和SSIM衡量），条件控制能力（通过F1分数衡量），同时保持了竞争力的文本图像一致性（通过CLIP-Score衡量）。值得注意的是，即使使用较小的排名（如1），模型也展示了竞争力的性能，尤其是在文本图像对齐方面，它达到了最高的CLIP-Score 0.765，展示了即使参数有限，作者方法的有效性。

条件深度。 FLUX.1 的 Transformer 架构具有两种不同的块：早期块采用分别针对不同模态 Token （文本和图像）的归一化模块，而后期块则采用所有 Token 的统一归一化。如表3所示，实验表明仅限制条件信号整合到这些早期块中会导致生成过程的可控性不足。这表明允许条件信号影响整个 Transformer 堆栈对于实现所需程度的控制至关重要。值得注意的是，这一发现表明，在早期块中主要插入条件信号的预览方法[14, 22, 37, 39, 41]，在基于 UNet 的架构中有效，可能不完全适用于 DiT 基础模型，如 FLUX.1。

5 Conclusion

Omniscient Control提供了一种参数高效的图像条件控制方法，用于在各种任务中使用扩散 Transformer ，该方法采用统一的 Token 方法，无需额外的模块。

作者的方法超越了传统方法，新发布的Subjects200K数据集，该数据集包含超过20万个高质量、主题一致的图像，支持主题一致生成的进步。

结果证实了Omniscient-Control在扩散模型中的可扩展性和有效性。

参考文献

[0]. OminiControl: Minimal and Universal Control for Diffusion Transformer.