引言：高分辨率图像合成的挑战与进展
在人工智能领域，高分辨率图像的合成一直是一个极具挑战性的课题。随着技术的不断进步，我们已经能够通过各种模型，如稳定扩散模型（Stable Diffusion）、Imgen、eDiff-I、DALL·E 和 Emu，从简单的文本提示中生成引人入胜的图像和设计。这些模型在多样化的开放网络数据上进行训练，展现出了令人惊叹的能力。然而，当涉及到生成高分辨率图像时，巨大的计算成本成为了一个障碍，尤其是在互动应用中，高分辨率导致的延迟问题尤为突出。
为了解决这一问题，研究人员提出了各种加速扩散模型推理的方法，包括减少采样步骤和优化神经网络推理。尽管计算资源的快速增长使得利用多个GPU加速推理变得可行，但在生成单个图像时，通常只涉及到一个GPU的使用，这限制了加速的潜力。因此，开发一种能够利用多个GPU加速单图像生成的方法变得尤为重要。
在这篇论文中，我们提出了一种名为DistriFusion的方法，它通过在多个GPU上并行运行扩散模型来减少单样本生成的延迟，而不会损害图像质量。DistriFusion的核心思想是利用相邻扩散步骤之间的输入相似性，通过重用前一时间步的预计算特征图来为当前步骤提供上下文。这种方法支持异步通信，可以通过计算来流水线化。广泛的实验表明，我们的方法可以应用于最新的Stable Diffusion XL，无需降低质量，且在八个NVIDIA A100 GPU上相比单个GPU可实现高达6.1倍的加速。我们的代码已在GitHub上公开。
论文标题: DistriFusion: Distributed Parallel Inference for High-Resolution Diffusion Models
项目地址: https://github.com/mit-han-lab/distrifuser
公众号【AI论文解读】后台回复“论文解读” 获取论文PDF!
DistriFusion方法概述：分布式并行推理的动机与基本原理
随着AI生成内容（AIGC）的兴起，如Adobe Firefly、Midjourney和Sora等工具展示了从简单文本提示生成引人入胜的图像和设计的惊人能力。这些成就得益于扩散模型的进步，这些模型通过迭代去噪过程合成内容。然而，生成高分辨率图像仍然具有挑战性，因为它需要巨大的计算成本，导致交互式应用程序的延迟过高。为了解决这个问题，我们提出了DistriFusion方法，通过利用多个GPU上的并行性来降低单样本生成的延迟，同时不损害图像质量。DistriFusion的核心思想是将模型输入分割成多个补丁，并将每个补丁分配给一个GPU。然而，简单地实现这样的算法会破坏补丁之间的交互并且失去保真度，而引入这种交互将会产生巨大的通信开销。为了克服这个困境，我们观察到相邻扩散步骤之间的输入具有高度相似性，并提出了位移补丁并行主义，它利用扩散过程的顺序性质，通过重用前一个时间步的预计算特征图为当前步骤提供上下文。因此，我们的方法支持异步通信，可以通过计算进行流水线处理。广泛的实验表明，我们的方法可以应用于最新的Stable Diffusion XL，无质量退化，并在八个NVIDIA A100s上与一个相比实现高达6.1倍的加速。
实现高效并行的关键：位移补丁并行主义
1. 位移补丁并行主义的定义与工作原理
位移补丁并行主义是一种新的并行化机会，它基于对扩散模型顺序特性的观察，通过重用前一步骤的特征作为当前步骤的上下文来重叠通信和计算。在DistriFusion中，我们首先将输入图像分割成多个补丁，每个补丁分配给不同的设备。对于每个设备，我们使用从上一步骤预计算的激活来提供全局上下文和补丁交互。我们进一步共同设计了一个推理框架来实现我们的算法。具体来说，我们的框架通过异步通信有效地隐藏了通信开销，并且仅在分配的区域上稀疏地运行卷积和注意力层，从而按比例减少了每个设备的计算量。
2. 通过异步通信优化计算与通信的重叠
在DistriFusion中，每个设备都有一个模型副本，并将独立地、并行地处理单个补丁。每个设备上的补丁首先被分散到上一步骤的陈旧激活中，然后选择性地应用层操作（线性、卷积或注意力层）到这些新鲜区域，从而生成相应区域的输出。这个过程对每一层重复。最后，所有层的输出被同步在一起，以近似ϵθ(xt)。通过这种方法，每个设备只负责总计算量的1/N，实现了高效的并行化。
为了获得上一步骤的陈旧激活，每个设备在每个时间步骤t获取Al,(i)，然后将激活广播到所有其他设备并执行AllGather操作。现代GPU通常支持异步通信和计算，这意味着AllGather过程不会阻塞正在进行的计算。到达下一个时间步骤的层l时，每个设备应该已经接收到Al t的副本。这种方法有效地将通信开销隐藏在计算阶段内。然而，有一个例外：第一步（即xT）。在这种情况下，每个设备简单地执行标准的同步通信，并缓存中间激活以备下一步使用。
稀疏操作与校正异步GroupNorm
1. 稀疏操作在减少每个设备计算量中的作用
稀疏操作是DistriFusion方法中的一个关键技术，它通过在每个设备上仅对新鲜区域执行计算来减少每个设备的计算量。在这种设置中，如果Fl是卷积、线性或跨注意力层，操作符仅应用于新鲜区域，而不是完整的特征图。对于自注意力层，它被转换为跨注意力层，只保留来自新鲜区域的查询令牌，而键和值令牌则覆盖整个特征图。这样，Fl的计算成本与新鲜区域的大小成正比，从而实现了稀疏操作。
2. 校正异步GroupNorm以保持图像质量
在DistriFusion中，为了保持图像质量，采用了校正异步GroupNorm技术。由于扩散模型通常采用组归一化（GN）层，需要聚合激活以恢复其完整的空间形状。研究发现，仅归一化新鲜补丁或重用陈旧特征会降低图像质量。为了解决这一问题，引入了一个校正项来近似全局均值和方差。这种方法通过异步通信计算精确的统计数据，从而将通信成本管道化到计算中。在极少数情况下，如果近似方差为负，则回退使用新鲜补丁的局部方差。
实验设置与基线比较
1. 实验模型与数据集介绍
实验中主要使用了最先进的公开文本到图像模型Stable Diffusion XL（SDXL）。SDXL首先使用预训练的自动编码器将图像压缩为原始大小的1/8，然后在这个潜在空间中应用扩散模型。此外，SDXL还采用了多个跨注意力层以促进文本条件。与原始的Stable Diffusion相比，SDXL采用了更多的注意力层，导致模型计算更加密集。实验使用了COCO Captions 2014数据集的HuggingFace版本进行基准测试。
2. 与现有方法的比较基准
DistriFusion与以下基线进行了比较：Naïve Patch方法将输入图像分成多个补丁，每个补丁独立处理，然后拼接在一起，但由于缺乏补丁间的交互，导致明显的边界缝隙。ParaDiGMS技术通过并行降噪多个步骤来加速预训练的扩散模型，使用Picard迭代来猜测未来步骤的解，并迭代细化直至收敛。实验使用了PSNR、LPIPS和FID等标准指标来评估图像质量。DistriFusion在保持与原始模型相当的图像质量的同时，实现了与使用设备数量成比例的延迟减少。
主要结果：质量与效率的评估
1. DistriFusion在图像质量保持上的表现
DistriFusion是一种新颖的并行计算方法，它通过将图像分割成多个补丁并在多个GPU上并行处理，旨在降低单样本生成的延迟，同时不损害图像质量。实验表明，DistriFusion能够在不降低视觉保真度的前提下，复制原始Stable Diffusion XL (SDXL)的性能。在生成1024×1024图像的设置中，即使使用8个设备，DistriFusion也能够在PSNR、LPIPS和FID评分上与原始模型的结果相媲美。
2. 在不同分辨率下的延迟减少与加速比
DistriFusion在多个NVIDIA A100 GPUs上的性能测试显示，与单个GPU相比，使用2、4和8个GPU时，生成单张图像的延迟分别降低了1.8倍、3.4倍和6.1倍。当分辨率提高到2048×2048和3840×3840时，GPU设备的利用率得到了更好的提升，从而进一步减少了延迟。特别是在生成3840×3840图像时，DistriFusion实现了显著的加速，达到了使用2、4和8个A100s分别减少1.8倍、3.4倍和6.1倍的延迟。
消融研究：验证设计选择的有效性
1. 通信成本分析与比较
在使用8个A100 GPUs进行的通信成本比较中，与同步张量并行（Sync. TP）和同步补丁并行（Sync. PP）相比，DistriFusion进一步减少了50%至60%的通信开销。这一结果表明，DistriFusion通过异步通信和计算重叠，有效地隐藏了通信开销。
2. 输入相似性的量化验证
DistriFusion的设计基于一个关键观察：相邻去噪步骤的输入具有高度相似性。通过量化计算使用50步DDIM采样器的所有连续步骤的模型输入差异，平均差异仅为0.02（输入范围在[-4, 4]内，约占0.3%）。这一发现支持了DistriFusion使用来自前一步骤的稍微过时的激活来促进补丁间交互的假设。
3. 少步采样与预热步骤的影响
在使用10步DPM-Solver的配置中，DistriFusion在没有预热步骤的情况下难以保持图像质量。然而，通过增加额外的两步预热步骤，显著提高了性能，同时仅略微增加了延迟。这表明预热步骤对于维持图像质量，尤其是在减少步骤数的采样器中，是非常重要的。
关注公众号 AI论文解读
获取最新AI论文解读
