摘要

苹果研究人员开源了一种高清扩散模型，通过多分辨率去噪和NestedUNet架构提高训练效率和生成质量。

MDM引入了多分辨率扩散过程，允许模型在不同分辨率间共享信息，从而更好地捕捉数据的内在结构和特征。

MDM使用的NestedUNet架构允许模型在处理多分辨率输入时，通过嵌套结构共享特征，提高生成的连贯性和一致性。

MDM采用渐进式训练方法，先从低分辨率开始，逐步增加分辨率，减轻计算负担，帮助模型学习分辨率间的关联和转换规则。

扩散模型已经成为文本生成图像、视频的主要框架之一，其基本原理是通过逐步引入噪声并在反向过程中去除噪声来生成数据。但在生成2K、4K高分辨率时，面临巨大的算力需求和时间消耗。

为了解决这一难题，苹果的研究人员开源了高清扩散模型Matryoshka Diffusion Models（简称“MDM”）， 通过引入一种特殊的扩散过程可在多个分辨率下同时进行去噪，并使用NestedUNet架构，将小尺度输入的特征和参数嵌套在大尺度的结构中，允许模型在不同分辨率间共享信息，从而提高训练效率和生成质量。

开源地址：https://github.com/apple/ml-mdm?tab=readme-ov-file

传统的扩散模型通常在单一分辨率下进行去噪操作，而MDM引入了多分辨率扩散过程，将不同分辨率的潜在变量联合起来进行处理。

例如，当生成一张图像时，MDM不仅考虑最终的高分辨率输出，还同时关注中间的低分辨率信息。这就好比在制作一幅精美的刺绣作品时，不仅要关注整体的图案设计，还要注重每一针每一线的细节，而低分辨率信息就像是刺绣中的底层线条，为高分辨率的呈现提供了基础。

在这个扩展空间中，MDM为数据点定义了时间相关的潜在变量，包含了多个不同分辨率的潜在变量。每个都有其特定的分布，并且与其他分辨率的潜在变量相互关联。

这种多分辨率的联合处理使得模型能够更好地捕捉数据的内在结构和特征，丰富了模型对数据分布的建模能力。由于在推理过程中关注了多个分辨率的信息，MDM可以将中间分辨率视为额外的隐藏变量，这些隐藏变量有助于模型更好地理解和生成高分辨率的输出

MDM使用的NestedUNet架构类似于传统的UNet架构，但在处理多分辨率输入时有一些区别。传统的UNet架构通过跳过连接和计算块来保留细粒度输入信息，而NestedUNet架构在此基础上，进一步将所有分辨率的潜在变量分组在一个去噪函数中形成嵌套结构。

这种嵌套结构使得低分辨率潜在变量的计算能够自然地为高分辨率潜在变量的计算提供有益的信息，就像套娃游戏一样，小的套娃嵌套在大的套娃中，彼此相互关联共同构成一个完整的整体。

在图像生成任务中，当模型处理不同分辨率的图像时，NestedUNet架构能够有效地整合这些信息。低分辨率的图像特征会被逐渐传递到高分辨率的部分，为高分辨率图像的生成提供重要的线索和指导。

例如，我们想生成一幅“海滩日落”的高分辨率图像。在MDM生成过程中会从一幅低分辨率的草图开始。这幅草图捕捉了海滩日落的基本元素，如天空、海面和太阳的大致轮廓。

随后，模型会在这一草图的基础上逐步增加细节，比如太阳的光辉、海浪的波纹和沙滩上的纹理。在每一级分辨率的提升过程中，模型都会利用NestedUNet架构中的特征共享机制，确保生成的图像在细节上的连贯性和一致性。

渐进式多阶段训练是MDM的另外一大优势。在传统的训练方法中，模型往往需要一次性处理所有分辨率的数据，这在高分辨率情况下会导致巨大的计算负担。

MDM则通过渐进式训练，先从低分辨率开始训练，逐步增加分辨率，直至达到所需的高分辨率输出。

这种方法不仅减轻了训练初期的计算压力，还有助于模型更好地学习到不同分辨率之间的关联和转换规则。