苹果重磅开源俄罗斯套娃扩散模型！MDM：多任务高分辨率生成又快又好！

文章链接：https://arxiv.org/pdf/2310.15111
项目链接：https://github.com/apple/ml-mdm

亮点直击

• 提出了Matryoshka Diffusion Models (MDM)，通过联合处理多个分辨率的输入，避免了级联或潜在扩散方法的复杂性，并引入了Nested UNet架构，嵌套不同尺度的特征和参数，提升了高分辨率生成的效果。
• 采用了多分辨率损失，显著加速了高分辨率去噪的收敛速度，同时使用渐进式训练策略，从低分辨率开始逐步引入高分辨率内容，实现了训练成本和生成质量的良好平衡。
• 在多个生成任务中表现优异，包括类条件图像生成、文本到图像生成和视频生成，尤其是在1024×1024像素分辨率下，使用较小的数据集（CC12M）依然取得了高质量生成结果。
• MDM具备强大的泛化能力，不仅在图像生成中表现突出，还自然扩展到视频生成，展现了广泛的应用潜力。

总结速览

解决的问题：

• 生成高质量图像和视频的扩散模型面临高维度学习的计算和优化挑战，尤其是在处理高分辨率图像时。

提出的方案：

• 引入Matryoshka Diffusion Model (MDM)，一种新颖的高分辨率图像和视频合成框架，采用联合降噪的扩散过程，在多个分辨率下处理输入。
• 使用NestedUNet架构，在大尺度输入的特征和参数中嵌套小尺度的特征和参数。
• 设计了一种从低分辨率到高分辨率的渐进式训练方式，优化高分辨率生成的效果。

应用的技术：

• MDM通过多分辨率联合降噪的扩散过程，使用NestedUNet架构，在像素空间直接进行高分辨率生成。
• 使用渐进式训练策略，逐步从低分辨率到高分辨率，解决了高分辨率生成的优化问题。

达到的效果：

• MDM在多个基准测试中表现出色，包括类条件图像生成、高分辨率文本到图像生成和文本到视频生成任务。
• 成功训练了单一的像素空间模型，最高分辨率可达1024×1024像素。
• 在使用仅包含1200万张图像的CC12M数据集上展示了强大的零样本泛化能力。

Matryoshka扩散模型

本节介绍了Matryoshka扩散模型（MDM），一个新型的扩散模型类别，在高分辨率空间中进行训练，同时利用数据生成的层次结构。MDM首先在扩展空间中推广了标准扩散模型，为此提出了专门的嵌套架构和训练过程。

扩展空间中的扩散模型

与级联或潜在方法不同，MDM通过在扩展空间中引入多分辨率扩散过程，学习一个具有层次结构的单一扩散过程。下图2展示了该过程的示意图。

NestedUNet架构

除了相对于其他层次化方法的简洁性，NestedUNet 还允许以最有效的方式分配计算。如下图 3 所示，早期探索发现，当将大部分参数和计算分配到最低分辨率时，MDM 在可扩展性方面表现得更好。

训练

研究者们在多个分辨率上联合使用正常的去噪目标来训练 MDM，具体如下：

渐进式训练

实验

MDM 是一种通用技术，适用于任何输入维度可以逐步压缩的问题。考虑了两个超出类条件图像生成的应用，展示了本文方法的有效性——文本到图像和文本到视频生成。

实验设置

数据集

在论文中广泛依赖 CC12M 作为文本到图像生成模型的数据集，显著不同于以往依赖极大且有时无法获取的数据集的研究。CC12M 足以构建高质量的文本到图像模型，并具备强大的zero-shot 能力，训练时间相对较短。这使得社区能够对方法进行更一致的比较，因为该数据集是免费提供的，且训练时间是可行的。因此 CC12M 更适合作为该问题研究的共同训练和评估基线。

评估
根据以往的研究，使用 Fréchet Inception Distance（ImageNet, CC12M）和 CLIP 分数（CC12M）来评估本文的图像生成模型。为了检查它们的zero-shot 能力，我们还报告了使用 COCO验证集生成图像的 FID/CLIP 分数，使用 CC12M 训练的模型。我们还在补充材料中提供了图像和视频合成的其他定性样本。

实现细节

基线模型除了与现有最先进方法的比较外，还在控制设置下对 MDM 与三个基线模型进行了详细分析：

1. 简单 DM：将标准 UNet 架构直接应用于高分辨率输入；我们还考虑了 Nested UNet 架构，但忽略低分辨率损失；这两种情况本质上与最近的端到端扩散模型如 Hoogeboom 等（2023）相同。
2. 级联 DM：遵循 Saharia 等（2022）的实现细节，训练一个直接与 MDM 可比较的 CDM，其中上采样器的配置与我们的 NestedUNet 相同。我们还对低分辨率条件图像应用噪声增强，并在推理过程中遍历最佳噪声水平。
3. 潜在 DM：利用来自 Rombach 等（2022）自动编码器的潜在编码，随后训练与 MDM UNet 维度匹配的扩散模型。

主要结果

比较简单 DM 和 MDM，可以清晰地看到 MDM 收敛速度更快，最终性能更好。这表明，多分辨率扩散过程结合多分辨率损失有效地改善了模型的收敛性，同时带来的复杂性微乎其微。当遵循渐进式训练计划时，看到 MDM 的性能和收敛速度进一步提高。作为直接比较，发现级联 DM 基线显著低于 MDM 的表现，尽管两者都从相同的64X64模型开始。需要注意的是，这一点非常显著，因为级联 DM 的参数总数大于 MDM（因为 MDM 在不同分辨率之间有广泛的参数共享），且推理步骤是其两倍。假设级联 DM 性能较差的主要原因在于64X64模型没有经过严格训练，这导致训练与推理之间在条件输入方面存在较大差距。最后，与 LDM 相比，MDM 的性能也更好。尽管这不是一个直接的对比，因为 LDM 确实由于其小输入尺寸而更有效，但 MDM 的训练和推理管道更为简单。

与文献的比较在下表 1 中，MDM 与现有文献中的方法进行了比较，报告了 ImageNet256X256的 FID-50K 和 MSCOCO 的zero-shot FID-30K。对于 ImageNet，我们的架构和超参数并没有经过优化，MDM 能够达到 3.51 的竞争性 FID 值，与 CFG 相比。我们的 FID 结果与文献相当，而 MDM 在训练时使用的数据量明显少于基线模型，如 Imagen 和 DALL·E 2。

定性结果下面展示训练后的 MDM 随机样本，用于图像生成（ImageNet 256X256 ，下图 5）、文本到图像生成（CC12M 1024X1024， 下图 6）和文本到视频生成（WebVid-10M，下图 7）。尽管在相对较小的数据集上进行训练，MDM 展现出强大的zero-shot 能力，能够生成高分辨率的图像和视频。值得注意的是，对所有三个任务使用相同的训练流程，表明其处理各种数据类型的多样化能力。

消融研究

渐进式训练的效果实验了渐进式训练计划，在该计划中，改变了低分辨率模型在继续训练目标分辨率之前的训练迭代次数（下图 8a）。看到更多的低分辨率训练明显有利于高分辨率的 FID 曲线。需要注意的是，在低分辨率输入上进行训练在内存和时间复杂度方面更为高效，因此渐进式训练为在训练过程中寻找最佳计算权衡提供了一种直接的选择。

嵌套层数的效果接下来，比较了在 CC12M 上使用不同数量嵌套分辨率的性能。结果如上图 8b 所示。从两个分辨率层增加到三个分辨率层始终改善了模型的收敛性。值得注意的是，增加嵌套层数仅带来了微不足道的成本。

CLIP-FID 权衡最后，在上图 8c 中展示了 COCO 的zero-shot 评估中 CLIP-FID 的帕累托曲线，这是通过改变无分类器引导（CFG）权重实现的。MDM 对 CFG 的响应与其他扩散模型变体类似。作为比较，叠加了 Imagen 报告的相同图。Imagen 通常展示出更小的 FID，这归因于其在大数据集上训练所导致的更高多样性。然而，MDM 展示出强大的 CLIP 得分，而在实践中发现，这种高 CLIP 得分与生成图像的视觉质量有很好的相关性。

讨论与未来方向

本文展示了跨不同分辨率共享表示可以加快训练速度并获得高质量结果，尤其是当低分辨率首先被训练时。相信这是因为模型能够更有效地利用不同分辨率之间的相关性，既在空间上也在时间上。尽管在这里仅探讨了一小部分架构，预计通过对权重共享架构的更详细探索，以及在当前架构中不同分辨率间参数分配的新方法，可以实现更多改进。

本文工作的另一个独特方面是使用了扩展空间，其中在多个分辨率上同时进行去噪。在这种形式下，时间和空间上的分辨率以相同的方式处理，时间和空间中的相关结构差异由权重共享模型的不同参数学习。对多分辨率的联合优化进行更一般的概念化的方法是将不同分辨率的损失解耦，并给予它们不同的权重。可以设想，在训练低分辨率到高分辨率的过程中可以实现平滑的过渡。

还注意到，尽管本文将方法与 LDM 进行了比较，这些方法是互补的。可以在自编码器代码的基础上构建 MDM。尽管并没有声称基于 MDM 的模型达到了最先进的水平，但将对 MDM 在大规模数据集和模型规模上的评估留作未来工作。

本文转自 AI生成未来，作者：AI生成未来

原文链接:​​https://mp.weixin.qq.com/s/BLkteFzrZJq1-tCf678DOA​​