【论文阅读笔记】One Diffusion to Generate Them All

---

介绍

Paper：https://arxiv.org/abs/2411.16318
Code：https://github.com/lehduong/onediffusion
Authors：Allen Institute for AI（AI2）\ 加州大学欧文分校 \ 华盛顿大学

理解

利用成组的序列数据通过Attention自身去完成其内部一致性的学习，多种任务同时训练又共同促进了模型的能力。

同类文章还有 GDT https://arxiv.org/pdf/2410.15027

---

引言

• LLMs 能作为通用模型，用于多种任务，并进行 zero-shot 泛化，作者团队试图寻求同样功能的扩散模型
• 现有的扩散模型往往需要 add-ons 来实现特定功能，比如条件控制、身份保持、多视图生成、姿态检测、深度提取、分割等
• 受基于序列数据的扩散模型启发，我们在训练期间将所有条件和目标图像建模为一系列具有不同噪声水平的“视图”。（本质上就是将对应任务的数据成组送入到模型，并使用不同的时间不 t 进行预测）
• 提出 OneGen 数据集，可用于训练不同任务
• 对多种特定任务起到了匹配特定模型的效果

二、相关工作

1. 介绍了多种基于扩散模型的针对特定生成任务的方法，如 Controlnet、T2I、instruct pix2pix、IP-adapter、PuLiD等，多视图任务有syncdreamer、mvdream 等
2. 介绍多种基于扩散模型的预测任务，如光流估计、开放语义分割等。这些研究表明，扩散模型不仅对生成图像有用，而且对计算机视觉中的各种预测任务也非常有效
3. 介绍了相关的统一任务生成，如 unicontrolnet，OmniGen 等

三、方法

预备知识

flow matching： 通过学习在两个概率分布之间传输的时间相关向量场来训练连续时间生成模型的框架。
$${\mathcal{L}}_{\mathrm{CFM}}(\theta ):={\mathbb{E}}_{t,q(z),p_t(x\mid z)}{\Vert v_{\theta }(t,x)-u_t(x\mid z)\Vert }^2$$
根据图 2，我们将多模态条件下的图像生成问题视为序列建模，将不同任务的训练数据编组。

训练

训练：采用 SD3 的时间采样方式$t\in LogNorm(0,1)$,为同一序列的输入进行单独采样构建加噪后的潜变量，然后遵循其去噪损失：$$\mathcal{L}(\theta )=\mathbb{E}\left[{\Vert v_{\theta }\left(t_1,\dots ,t_N,{\mathbf{x}}_1,\dots ,{\mathbf{x}}_N\right)-u\Vert }^2\right]$$

推理

推理阶段将条件的噪声强度设置为 0，想要输出的目标图像（可能是多个）使用相同的去噪时间步 t，可以遵循正常的去噪推理流程对 ODE求解。

实现细节

• 网络架构：采用 Next-DiT 架构；设计了一个支持多种输入的变量长度定义方式；使用 3D RoPE 来处理位置编码
• 任务设定
  • 文本到图像（1 个视图）。仅使用单个“视图”，训练和推理遵循与标准文本到图像扩散模型相同的过程。我们将任务标签 [[text2image]]] 添加到标题中以指定任务。
  • 图像到图像（2 个视图）。我们将第一个视图设置为目标图像，第二个视图设置为条件输入。在推理过程中，我们可以使用一个或两个视图来生成，并训练模型生成目标图像。对于边界框或语义图生成等任务，我们将十六进制颜色代码和类标签添加到提示中。例如，要用黄色掩码分割鼠标，提示为：[[semantic2image]] <#FF00 黄色掩码：鼠标>一张照片。
  • ID定制(2-4个视图)。我们跨视图对同一个人的图像进行采样，为每个输入图像连接字幕并使用标记 [[[imgX]] 来表示每张图像。我们还将任务标签 [[faceid]]] 添加到标题中。在推理时，我们可以以任意数量的图像为条件并生成多个输出，从而产生更一致的结果。
  • 多视图生成（4-12 个视图）。受Cat3D的启发，我们使用Pl̈ucker射线嵌入来表示相机姿势。对于每个图像补丁，我们使用其射线原点o和方向d计算Pl̈ucker坐标asr = (o × d, d)。结果嵌入的维度 H/8×W/8×6，匹配潜在空间大小，并跨通道复制以形成 16 个通道嵌入。与 [18] 不同，我们将光线嵌入视为图像潜在变量的独立“视图”作为一个统一的序列，而不是通过通道连接。这种设计允许灵活的去噪，支持以相机姿势或采样光线嵌入为条件的多视图图像生成，以从图像条件下预测姿势，类似于 RayDiffusion [67]。我们将射线嵌入缩放到具有单位方差，如[47]所示。与其他任务一样，我们将任务标签 [[多视图]] 添加到标题中。在推理过程中，我们可以将图像或Pl̈ucker射线嵌入替换为高斯噪声，分别用于多视图生成和摄像机姿态估计。

训练细节

我们的模型是使用流匹配目标从头开始训练的。与之前的工作类似，我们使用三个阶段的训练配方。在第一阶段，我们以 256(500K 步) 和 512(500K 步) 的分辨率预训练文本到图像模型。在第二阶段，我们继续对混合任务进行训练，T2I 使用 512，其他任务使用 256，总共 1M 步。最后，在最后阶段，我们以 (1024) 的高分辨率对模型进行微调，用于 T2I。对于 ID 定制微调，我们使用 2-5 个视图。对于更少的视图（2-3），我们应用 512 的分辨率，而对于更多视图，我们使用 256 分辨率。

在训练期间，我们在每个阶段使用批量采样策略，以相等的概率对任务（T2I、图像到图像、ID 定制和多视图生成）进行采样。噪声调度器的移位值设置为 3，如SD3中所建议的。我们使用学习率为 η = 0.0005 的 AdamW 优化器。在第一阶段以 256 的全局批量大小为 256 的 TPU v3-256 pod 进行训练，最后的微调阶段使用相同的配置在 64 H100 GPU 上完成。

数据集构建

Text-to-Image 我们利用公共和内部（合成）数据集。公共数据集包括：PixelPenrose [52]、Unsplash、Coyo [6]、JourneyDB [40]。此外，我们使用 10M 内部合成数据集，该数据集由使用 LLAVA-NeXT [31] 和 Molmo [11] 重新描述的图像组成。每个图像的文本描述长度从 100 到 150 个单词不等。当原始提示可用时，我们使用 LLM 生成的标题和原始标题。

图像到图像 对于更简单的任务，如去模糊、修复、从canny边缘生成图像或放大，我们使用合成数据的1M样本子集，并为每个图像应用相关的预处理器来创建输入条件。对于更复杂的任务，我们按照下面概述的过程从 Midjourney、Stable Diffusion [15] 和 Flux-dev 生成的输出创建一个合成数据集：

1. 语义图和检测对于每个图像，我们使用 LLAAVA-NeXT 模型来识别实体或主题（例如，人、衬衫、狗、建筑物），每张图像最多有 10 个实体。基于 LlaVA-Next 的这些主题名称，我们使用 SAM 执行语义分割并提取边界框。每个类都从预定义的列表中分配一个随机颜色。该数据集包含 350K 三元组，由语义图、边界框和原始图像组成。
2. 深度图我们通过将 DepthAnything-v2 应用于从各种数据集中采样的 500K 图像来生成深度数据集，包括真实图像和合成图像。此外，我们将 Hypersim 数据集 中的 40K 图像与 LLAVA-NeXT 的标题，并将它们合并到训练集中。
3. 人体姿势我们收集了一个包含 50K 图像的不同子集，主要是人类进行姿势调节。我们使用YOLOv5检测感兴趣区域的边界框，并应用ViTPose进行位姿估计
   ID定制 我们通过从公开可用的图像中从游戏和电影中收集名人和人物的数据集。在过滤以确保每个受试者至少有四个图像并删除NSFW内容后，数据集包括大约60K个受试者，总共130万张图像。我们使用 LLAAVA-NeXT 对这些图像进行字幕。
   多视图生成 我们使用DL3DV-10K数据集[29]、Obj厌恶[10]、CO3D[45]。对于 Objverse 数据集，我们利用从 LGM [54] 过滤拆分的 80K 和 Cap3D [38] 提供的标题。在 DL3DV 数据集中，我们从每个场景中采样图像并使用 LLAVA-Next 对其进行标题。对于 CO3D，我们排除了字幕并仅包含文本输入中的任务标记。

实验分

1. Text2Image
   文本到图像任务的OneDiffusion的定性结果如图3所示。由于我们的One-Gen数据集的多样性，该模型可以处理各种艺术风格，跨越艺术和逼真的设计。继之前的工作 [15] 之后，我们评估了我们的模型在 GenEval 基准 [20] 上的文本到图像能力。对于每个提示，我们使用具有 100 步的欧拉求解器和 5 的引导尺度生成 4 张图像。 OneDiffusion 的结果以及基线模型的结果如表 1 所示。与类似大小的基线相比，我们的模型表现出强大的性能，尽管在相对较小的数据集上进行训练，但在多任务能力方面表现出色。这一性能很大程度上归功于数据集的多样性和每个样本提供的综合字幕。

2. 可控图像生成
   文本到图像任务的OneDiffusion的定性结果如图3所示。由于我们的One-Gen数据集的多样性，该模型可以处理各种艺术风格，跨越艺术和逼真的设计。继之前的工作 [15] 之后，我们评估了我们的模型在 GenEval 基准 [20] 上的文本到图像能力。对于每个提示，我们使用具有 100 步的欧拉求解器和 5 的引导尺度生成 4 张图像。 OneDiffusion 的结果以及基线模型的结果如表 1 所示。与类似大小的基线相比，我们的模型表现出强大的性能，尽管在相对较小的数据集上进行训练，但在多任务能力方面表现出色。这一性能很大程度上归功于数据集的多样性和每个样本提供的综合字幕。
   

3. 多视角生成
   
   我们使用 **Google Scanned Object 数据集 [13] 评估我们方法的多视图生成能力。**表 2 显示了我们的方法与最先进的方法（包括 Zero123 [32]、Zero123-XL [12] 和 EscherNet [15]）的比较。需要注意的是，这些模型是专门为多视图生成任务设计和训练的。相比之下，我们的模型支持可变数量的条件输入，由于其灵活的去噪框架，即使相机姿势未知，也可以合并额外的条件视图。如表2所示，OneDiffusion
   图 5 展示了从单个前视图图像生成多视图的两个定性示例。我们的模型成功地在不同的方位角和仰角之间生成了一致的视图。有关额外的可视化，请参阅附录中的图 10 和 11。由于我们框架的灵活性，我们还可以直接执行文本到多视图，通过“屏蔽”所有图像，并且只输入相机姿势，如附录中的图 12 所示
   

4. ID 保持
   我们进一步评估了 ID 定制任务的 OneDiffusion，该任务涉及使用一个或多个 ID 图像作为个性化生成的输入。为了评估性能，我们使用定性和定量分析与 STOA 方法（包括 InstantID [58]、PuLID [21] 和 PhotoMaker [27]）进行了比较。我们的评估超出了标准基准（unsplash-50 [16]）来测试 ID 定制任务的泛化，例如不同的表达式、视点甚至非人类图像。图 6 说明了改变面部表情和注视方向（第一行）、改变视点（第二行）和定制非人类 ID（第三行）的示例。
   

5. 深度估计
   对于图像理解任务，我们使用标准基准评估我们的模型在单目深度估计上的性能：NYUv2 [51] 和 DIODE [56]。我们在表4中报告了定量结果。定量结果如表4所示，与利用预先训练的文本到图像扩散模型(如Marigold[22])的基线相比，我们的模型获得了具有竞争力的性能。然而，如图 7 所示，我们的模型表现出比 Marigold 等基于扩散的深度估计器更强的鲁棒性。具体来说，它擅长处理开放世界图像，包括绘画、朦胧天气和非常规纹理。为了进一步定性比较，请参见图 15 和图 16。
   

6. 相机姿态估计
   我们使用 Google Scanned Object 数据集 [13] 在相机位姿估计上评估我们的模型。对于这项任务，我们使用每个合成对象的六个渲染图像，并通过去噪相应的射线嵌入来估计相机姿势。在RayDiffusion[67]之后，我们应用最小二乘优化来估计相机的中心和旋转。表 5 报告了以 0.3 阈值测量的相机中心精度。
   图 8 提供了我们的模型和 RayDiffusion 之间的定性比较。由于训练数据的偏差，RayDiffusion 始终预测上半球的相机姿势，例如 CO3D，它主要具有上半球视图。相比之下，由于我们的大规模训练数据集的多样性，OneDiffusion 实现了更高的准确度并避免这个限制。

7. 其他任务
   由于直接从原始输出图像中提取掩码、边界框和关键点并不简单，我们在附录图17和图18中分别为COCO数据集上的人体姿态估计和语义分割提供了几个定性结果。由于我们的模型在训练期间不区分条件和图像的任务，因此其在理解任务上的性能作为模型的附加指标。我们将对这方面的进一步探索留给未来的工作。

---

结论

我们的实验表明，OneDiffusion 在各种任务中取得了令人印象深刻的结果，包括条件 T2I 生成、深度估计、开放词汇语义分割、姿势估计、多视图生成、ID 定制和相机位姿估计。我们相信这项工作提高了扩散模型的能力，提供了一种通用且可扩展的解决方案，可与大型语言模型提供的灵活性相媲美。这代表着朝着开发通用视觉模型迈出的重要一步，该模型可以作为各种应用的主干。

---

附录

数据分布：

更多的结果查看原文吧

---