在这个大模型不断创造新成就的时代,我们通常对机器学习模型有一个直观认知:越大越好。但事实果真如此吗?
近日,Google Research 一个团队基于隐扩散模型(LDM)进行了大量实验研究,得出了一个结论:更大并不总是更好(Bigger is not Always Better),尤其是在预算有限时。- 论文标题:Bigger is not Always Better: Scaling Properties of Latent Diffusion Models
- 论文地址:https://arxiv.org/pdf/2404.01367.pdf
近段时间,隐扩散模型和广义上的扩散模型取得的成就不可谓不耀眼。这些模型在处理了大规模高质量数据之后,可以非常出色地完成多种不同任务,包括图像合成与编辑、视频创建、音频生成和 3D 合成。尽管这些模型可以解决多种多样的问题,但要想在真实世界应用中大规模使用它们,还需要克服一大障碍:采样效率低。这一难题的本质在于,为了生成高质量输出,LDM 需要依赖多步采样,而我们知道:采样总成本 = 采样步骤数 × 每一步的成本。具体来说,目前人们首选的方法需要使用 50 步 DDIM 采样。这个过程虽能确保输出质量,但在具备后量化(post-quantization)功能的现代移动设备上却需要相当长的延迟才能完成。因此,为了促进 LDM 的实际应用,就需要优化其效率。事实上,这一领域已经出现了一些优化技术,但对于更小型、冗余更少的模型的采样效率,研究社区还未给予适当关注。在这一领域,一个重大障碍是缺少可用的现代加速器集群,因为从头开始训练高质量文生图 LDM 的时间和资金成本都很高 —— 往往需要几周时间和数十万美元资金。该团队通过实验研究了规模大小的变化对 LDM 的性能和效率的影响,其中关注重点是理解 LDM 的规模扩展性质对采样效率的影响。他们使用有限的预算从头开始训练了 12 个文生图 LDM,参数量从 39M 到 5B 不等。图 1 给出了一些结果示例。所有模型都是在 TPUv5 上训练的,使用了他们的内部数据源,其中包含大约 6 亿对已过滤的文本 - 图像。他们的研究发现,LDM 中确实存在一个随模型规模变化的趋势:在同等的采样预算下,较小模型可能有能力超越较大模型。此外,他们还研究了预训练文生图 LDM 的大小会如何影响其在不同下游任务上的采样效率,比如真实世界超分辨率、主题驱动的文生图( 即 Dreambooth)。 对于隐扩散模型在文生图和其它多种下游任务上的规模扩展性质,该团队得到了以下重要发现:- 预训练的性能会随训练计算量而扩展。通过将模型的参数量从 39M 扩展到 5B,该团队发现计算资源和 LDM 性能之间存在明显联系。这表明随着模型增大,还有潜力实现进一步提升。
- 下游性能会随预训练而扩展。该团队的实验表明:预训练性能与在下游任务上的成功之间存在很强的关联。较小模型即使使用额外的训练也无法完全赶上较大模型的预训练质量所带来的优势。
- 较小模型的采样效率更高。当给定了采样预算时,较小模型的图像质量一开始会优于较大模型,而当放松计算限制时,较大模型会在细节生成上胜过较小模型。
- 采样器并不会改变规模扩展效率。无论使用哪种扩散采样器,较小模型的采样效率总是会更好一点。这对确定性 DDIM、随机性 DDPM 和高阶 DPM-Solver++ 而言都成立。
- 在步数更少的下游任务上,较小模型的采样效率更高。当采样步数少于 20 步时,较小模型在采样效率上的优势会延伸到下游任务。
- 扩散蒸馏不会改变规模扩展趋势。即使使用扩散蒸馏,当采样预算有限时,较小模型的性能依然能与较大蒸馏模型竞争。这说明蒸馏并不会从根本上改变规模扩展趋势。
该团队基于广被使用的 866M Stable Diffusion v1.5 标准,开发了一系列强大的隐扩散模型(LDM)。这些模型的去噪 UNet 具有不同的规模,参数数量从 39M 到 5B 不等。该团队通过逐渐增大残差模块中过滤器的数量,同时维持其它架构元素不变,实现了可预测的受控式规模扩展。表 1 展示了这些不同大小模型的架构差异。其中也提供了每个模型相较于基线模型的相对成本。图 2 展示了规模扩展过程中的架构差异。这些模型的训练使用了他们的内部数据源,其中有 6 亿对经过过滤的文本 - 图像。所有模型都训练了 50 万步,批量大小为 2048,学习率为 1e-4。这让所有模型都能到达收益递减的程度。图 1 表明这些不同大小的模型都具有稳定一致的生成能力。对于文生图任务,他们设置的采样步数为常用的 50 步,采样器为 DDIM,无分类器指导率为 7.5。可以看到,随着模型规模增大,所得结果的视觉质量明显提升。实验中,各种大小的 LDM 的生成性能相对于训练计算成本都有类似的趋势,尤其是在训练稳定之后 —— 通常是在 20 万次迭代之后。这些趋势表明不同大小的模型的学习能力具备明显的扩展趋势。具体来看,图 3 展示了参数量从 39M 到 5B 的不同模型的运行情况,其中的训练计算成本是表 1 中给出的相对成本和训练迭代次数的积。评估时,使用了相同的采样步数和采样参数。在训练计算量适中(即 < 1G,见图 3)的场景中,文生图模型的生成性能可在额外计算资源的帮助下很好地扩展。基于在文本 - 图像数据上预训练的模型,该团队又针对真实世界超分辨率和 DreamBooth 这两个下游任务进行了微调。表 1 给出了这些预训练模型的性能。图 4 左图给出了在超分辨率(SR)任务上的生成性能 FID 与训练计算量的对应情况。可以看出来,相比于训练计算量,超分辨率的性能更依赖模型大小。实验结果表明较小模型有一个明显的局限性:不管训练计算量如何,它们都无法达到与较大模型同等的性能。图 4 右图给出了失真度指标 LPIPS 的情况,可以看到其与生成指标 FID 有一些不一致。虽如此,还是可以从图 5 明显看出:较大模型比较小模型更擅长恢复细粒度的细节。基于图 4 能得到一个关键见解:相比于较小的超分辨率模型,较大模型即使微调时间更短,也能取得更好的结果。这说明预训练性能(由预训练模型大小主导)对超分辨率 FID 分数的影响比对微调的持续时间(即用于微调的计算量)的影响大。此外,图 6 比较了不同模型上 DreamBooth 微调的视觉结果。可以看到视觉质量和模型大小之间也有相似的趋势。 分析 CFG 率的影响。文生图生成模型需要超过单一指标的细致评估。采样参数对定制化来说非常重要,而无分类器引导(CFG)率可以直接影响视觉保真度以及与文本 prompt 的语义对齐之间的平衡。Rombach 等人的论文《High-resolution image synthesis with latent diffusion models》通过实验表明:不同的 CFG 率会得到不同的 CLIP 和 FID 分数。而这项新研究发现 CFG 率(一个采样参数)会在不同的模型大小上得到不一致的结果。因此,使用 FID 或 CLIP 分数以定量方式确定每个模型大小和采样步骤的最佳 CFG 率是很有趣的。该团队使用不同的 CFG 率(即 1.5、2.0、3.0、4.0、5.0、6.0、7.0、8.0)对不同规模的模型进行了采样,并以定量和定性方式比较了它们的结果。图 7 便是两个模型在不同的 CFG 率下的视觉结果,从中可以看出其对视觉质量的影响。该团队观察到,相比于 prompt 语义准确度,CFG 率的变化对视觉质量的影响更大,因此为了确定最佳 CFG 率,他们选取的评估指标是 FID 分数。图 8 给出了不同的 CFG 率对文生图任务的 FID 分数的影响。规模扩展效率趋势。使用每个模型在不同采样步骤下的最佳 CFG 率,该团队分析了最优性能表现,以理解不同 LDM 大小的采样效率。具体来说,图 9 比较了不同采样成本下(归一化成本 × 采样步数)的不同模型及其最优性能。通过追踪不同采样成本下的最优性能点(竖虚线),可以看到一个趋势:在一个采样成本范围内,较小模型的 FID 分数通常优于较大模型。图 10 则给出了较小和较大模型结果的定性比较,从中可以看到在相似的采样成本条件下,较小模型是可以匹敌较大模型的。 为了评估采样效率趋势在不同模型规模下的普遍性,该团队评估了不同大小的 LDM 使用不同扩散采样器的性能。他们使用的采样器有三种:DDIM、随机性 DDPM、高阶 DPM-Solver++。可以看出,当采样步数较少时,DDPM 采样器得到的质量通常低于 DDIM,而 DPM-Solver++ 则在图像质量上胜过 DDIM。另一个发现也很重要,即三种采样器都有一致的采样效率趋势:采样成本一样时,较小模型的性能会优于较大模型。由于 DPM-Solver++ 采样器的设计并不适合用于超过 20 步的采样,因此这也是其采样范围。结果表明:不管使用什么采样器,LDM 的规模扩展性质始终保持一致。这里关注的重点下游任务是超分辨率。这里是直接使用超分辨率采样结果,而不使用 CFG。受图 4 启发(在下游任务上,不同大小的 LDM 在采样 50 步时性能差距较大),该团队从两个方面调查了采样效率:较少采样步数和较多采样步数。如图 12 左图所示,当采样步数不超过 20 步时,不同大小模型的采样效率趋势在超分辨率任务上依然成立。但图 12 右图又表明,一旦超过这个范围,较大模型的采样效率就会超过较小模型。这一观察结果说明,在文生图和超分辨率等任务上,不同大小模型在采样步数较少时的采样效率趋势是一致的。 虽然之前的实验结果说明较小模型的采样效率往往更高,但需要指出,较小模型的建模能力也往往更差一些。对于近期那些严重依赖建模能力的扩散蒸馏方法来说,这就成了一大难题。人们可能会预测出一个矛盾的结论:经过蒸馏的大模型的采样速度快于经过蒸馏的小模型。为了展示经过蒸馏的不同大小模型的采样效率,该团队使用条件一致性蒸馏方法在文生图数据上对之前的不同大小模型进行了蒸馏操作,然后比较了这些已蒸馏模型的最佳性能。详细来说,该团队在采样步数 = 4(这已被证明可以实现最优的采样性能)的设定下测试了所有已蒸馏模型;然后在归一化的采样成本上比较了每个已蒸馏和未蒸馏模型。图 13 左图表明,在采样步数 = 4 时,蒸馏可以提升所有模型的生成性能,并且 FID 全面提升。而在右图中,可以看到在同等的采样成本下,已蒸馏模型的表现优于未蒸馏模型。但是,在特定的采样成本下(即采样成本≈8),较小的未蒸馏 83M 模型依然能取得与较大已蒸馏 866M 模型相近的性能。这一观察进一步支持了该团队提出的不同大小 LDM 的采样效率趋势,其在使用蒸馏时也依然成立。