DiCo：重振卷积网络以实现可扩展且高效的扩散建模 原创 精华

摘要 

扩散 Transformer（DiT）作为一种前景广阔的视觉生成扩散模型，展现出了令人瞩目的性能，但同时也带来了巨大的计算开销。有趣的是，对预训练 DiT 模型的分析表明，全局自注意力往往存在冗余，主要捕捉局部模式，这凸显了寻找更高效替代方案的潜力。在本文中，我们重新审视卷积，将其作为构建高效且富有表现力的扩散模型的替代构建块。然而，直接用卷积替换自注意力通常会导致性能下降。我们的研究发现，这种性能差距归因于卷积网络相比 Transformer 存在更高的通道冗余。为解决这一问题，我们引入了一种紧凑通道注意力机制，该机制能够促进更多样化通道的激活，从而增强特征多样性。由此产生了扩散卷积网络（DiCo），这是一系列完全由标准卷积网络模块构建的扩散模型，具有强大的生成性能和显著的效率提升。在类条件 ImageNet 基准测试中，DiCo 在图像质量和生成速度方面均优于以往的扩散模型。值得注意的是，DiCo-XL 在 256×256 分辨率下的 FID 为 2.05，在 512×512 分辨率下为 2.53，相比 DiT-XL/2 分别实现了 2.7 倍和 3.1 倍的加速。此外，我们最大的模型 DiCo-H，参数规模达 10 亿，在 ImageNet 256×256 上的 FID 达到 1.90，且在训练过程中无需任何额外监督。 

1、引言 

扩散模型 [68,70,27,69,71] 在生成式学习领域引发了变革性的进步，在合成高度逼真的视觉内容方面展现出了卓越的能力。其多功能性和有效性使其在广泛的实际应用中得到了广泛采用，包括文本到图像生成 [60,63,61]、图像编辑 [53,40,7]、图像修复 [39,1,2]、视频生成 [30,82,4] 和 3D 内容创建 [58,81,78]。 

早期的扩散模型（如 ADM [11] 和 Stable Diffusion [61]）主要采用混合 U-Net [62] 架构，该架构将卷积层与自注意力相结合。最近，Transformer [77] 已成为一种更强大且可扩展的主干架构 [56,3]，促使人们转向完全基于 Transformer 的设计。因此，扩散 Transformer（DiTs）正逐渐取代传统的 U-Nets，如领先的扩散模型 Stable Diffusion 3 [16]、FLUX [44] 和 Sora [6] 所示。然而，自注意力的二次计算复杂度带来了巨大挑战，尤其是在高分辨率图像合成中。最近的研究 [90,73,19,57,84] 探索了更高效的替代方案，聚焦于线性复杂度的类 RNN 架构，如 Mamba [20] 和门控线性注意力 [85]。尽管这些模型提高了效率，但其因果设计与视觉生成的双向特性 [23,49] 本质上存在冲突，限制了它们的有效性。此外，如图 3 所示，即使采用高度优化的 CUDA 实现，在高分辨率设置下，它们相比传统 DiTs 的运行时优势仍然有限。这引发了我们的关键问题：是否有可能设计一种硬件高效的扩散主干架构，同时保留如 DiTs 般强大的生成能力？ 

为解决这一问题，我们首先分析 DiTs 生成能力背后的特性。在视觉识别任务中，视觉 Transformer [15] 的成功通常归因于自注意力捕捉长距离依赖的能力 [36,18,24]。然而，在生成任务中，我们观察到不同的动态。如图 4 所示，对于预训练的类条件（DiT-XL/2 [56]）和文本到图像（PixArt-α[9] 和 FLUX [44]）DiT 模型，当用锚定标记查询时，注意力主要集中在附近的空间标记上，很大程度上忽略了远处的标记。这一发现表明，计算全局注意力对于生成任务可能是冗余的，强调了局部空间建模的重要性。与需要长距离交互进行全局语义推理的识别任务不同，生成任务似乎更注重细粒度的纹理和局部结构的保真度。这些观察揭示了 DiTs 中注意力的固有局部性，并促使我们寻求更高效的架构。 

在这项工作中，我们重新审视卷积神经网络（ConvNets），并提出扩散卷积网络（DiCo），这是一种专为扩散模型设计的简单而高效的卷积主干架构。与自注意力相比，卷积操作对硬件更友好，在大规模和资源受限的部署中具有显著优势。虽然用卷积替代自注意力可大幅提高效率，但通常会导致性能下降。如图 5 所示，这种直接替换会引入明显的通道冗余，许多通道在生成过程中保持不活跃状态。我们假设，这种性能下降源于自注意力相比卷积具有更强的动态内容依赖表示能力。为解决这一限制，我们引入了紧凑通道注意力（CCA）机制，该机制通过轻量级线性投影动态激活信息丰富的通道。作为一种逐通道的全局建模方法，CCA 在保持低计算开销的同时，增强了模型的表示能力和特征多样性。与依赖大型昂贵内核的现代识别卷积网络 [12,21] 不同，DiCo 采用了完全基于高效 1×1 逐点卷积和 3×3 深度卷积的简化设计。尽管架构简单，DiCo 仍实现了强大的生成性能。 

如图 2 和图 3 所示，DiCo 模型在 ImageNet 256×256 和 512×512 基准测试中均优于最先进的扩散模型。值得注意的是，我们的 DiCo-XL 模型在 256×256 和 512×512 分辨率下分别实现了令人印象深刻的 FID 分数 2.05 和 2.53。除了性能提升外，DiCo 模型相比基于注意力 [77]、基于 Mamba [20] 和基于线性注意力 [38] 的扩散模型还展现出显著的效率优势。具体而言，在 256×256 分辨率下，DiCo-XL 的 Gflops 减少了 26.4%，速度比 DiT-XL/2 [56] 快 2.7 倍。在 512×512 分辨率下，DiCo-XL 的运行速度分别比基于 Mamba 的 DiM-H [73] 和 DiS-H/2 [19] 模型快 7.8 倍和 6.7 倍。我们最大的模型 DiCo-H 拥有 10 亿参数，进一步将 ImageNet 256×256 上的 FID 降低至 1.90。这些结果共同凸显了 DiCo 在基于扩散的生成建模中的强大潜力。 

总体而言，这项工作的主要贡献可总结如下： 

•我们分析了预训练 DiT 模型，揭示了其全局注意力机制中存在的显著冗余和局部性。这些发现可能会启发研究人员开发更高效的策略来构建高性能扩散模型。 

•我们提出了 DiCo，这是一种简单、高效且强大的扩散模型卷积主干架构。通过引入紧凑通道注意力，DiCo 在不牺牲效率的情况下显著提高了表示能力和特征多样性。 

•我们在 ImageNet 256×256 和 512×512 基准测试上进行了广泛实验。DiCo 在生成质量和速度方面均优于现有扩散模型。据我们所知，这是第一项证明精心设计的全卷积主干架构可以在基于扩散的生成建模中实现最先进性能的工作。 

2、相关工作 

2.1 扩散模型架构 

早期的扩散模型通常采用 U-Net [62] 作为基础架构 [11,28,61]。最近，越来越多的研究探索将视觉 Transformer（ViTs）[15] 作为扩散模型的替代主干架构，取得了显著成果 [56,3,52,88,59,48]。值得注意的是，DiT [56] 展示了基于 Transformer 架构的卓越性能，在 ImageNet 生成任务中实现了最先进的性能。然而，ViTs 固有的二次计算复杂度在长序列建模的效率方面带来了巨大挑战。为缓解这一问题，最近的研究探索了使用具有线性复杂度的类 RNN 架构，如 Mamba [20] 和线性注意力 [38]，作为扩散模型的主干架构 [19,90,73,84,57]。DiS [19] 和 DiM [73] 采用 Mamba 来降低计算开销，而 DiG [90] 则利用门控线性注意力 [85] 来实现具有改进效率的竞争性能。在这项工作中，我们重新审视卷积网络作为扩散模型的主干架构。我们表明，通过适当设计，纯卷积架构可以实现最先进的生成性能，为 DiTs 提供了一种高效且强大的替代方案。 

2.2 卷积网络设计 

在过去的十年中，卷积神经网络（ConvNets）在计算机视觉领域取得了显著成功 [42,67,25,35,83]。为实际部署开发了许多轻量级卷积网络 [33,65,32,13]。尽管 Transformer 已逐渐成为广泛任务的主导架构，但其巨大的计算开销仍然是一个重大挑战。许多现代卷积网络设计在保持高效率的同时实现了竞争性能。ConvNeXt [51] 探索了标准卷积网络的现代化，并取得了优于基于 Transformer 模型的结果。RepLKNet [12] 研究了大内核卷积的使用，将内核大小扩展至 31×31。UniRepLKNet [14] 进一步将大内核卷积网络推广到音频、点云、时间序列预测等领域。在这项工作中，我们探索了纯卷积网络在基于扩散的图像生成中的潜力，并表明简单高效的卷积网络设计也可以实现优异性能。 

3、方法 

3.1 预备知识 

扩散公式 

我们首先回顾扩散模型的基本概念 [27,71]。扩散模型的特点是前向加噪过程，该过程将噪声逐步注入数据样本x_0。具体而言，前向过程可表示为： 

其中\overline{\alpha}_t是预定义的超参数。扩散模型的目标是学习反向过程：，其中神经网络对过程的均值和协方差进行参数化。训练涉及优化x_0对数似然的变分下界，该下界可简化为： 

为简化训练，模型的预测均值\mu_\theta可重新参数化为噪声预测器\epsilon_\theta。目标函数随后简化为预测噪声与真实噪声\epsilon_t之间的均方误差：。遵循 DiT [56]，我们使用简化损失L_{simple}训练噪声预测器\epsilon_\theta，而协方差\sum_\theta则使用完整损失进行优化。 

无分类器引导 

无分类器引导（CFG）[29] 是一种在条件扩散模型中提高样本质量的有效方法。它通过将采样过程引导至与给定条件c密切相关的输出来实现这一点。具体而言，它修改预测噪声以获得高p(x|c)，如下所示： 

其中s\geq1控制引导强度，\epsilon_\theta(x_t,\emptyset)是在训练期间通过随机省略条件信息获得的无条件预测。遵循先前的工作 [56,90]，我们采用该技术来提高生成样本的质量。 

3.2 网络架构 

当前，扩散模型主要分为三种架构类型：（1）无任何下采样层的各向同性架构，如 DiT [56]；（2）具有长跳跃连接的各向同性架构，以 U-ViT [3] 为例；（3）U 形架构，如 U-DiT [76]。受图像去噪中多尺度特征关键作用的启发 [89,80]，我们采用 U 形设计来构建分层模型。我们还进行了广泛的消融研究，以系统比较表 3 中这些不同架构选择的性能。 

如图 6（a）所示，DiCo 采用由堆叠 DiCo 块组成的三阶段 U 形架构。模型以 VAE 编码器生成的空间表示z为输入。对于大小为 256×256×3 的图像，对应的z维度为 32×32×4。为处理该输入，DiCo 应用 3×3 卷积，将z转换为具有D通道的初始特征图z_0。对于条件信息，即时间步t和类别标签y，我们采用多层感知机（MLP）和嵌入层，分别作为时间步和标签嵌入器。在 DiCo 的每个块l中，特征图z_{l-1}通过第l个 DiCo 块生成输出z_l。 

在每个阶段内，编码器和解码器之间的跳跃连接促进了中间特征之间的高效信息流。 concatenation 后，应用 1×1 卷积以降低通道维度。为实现跨阶段的多尺度处理，我们利用像素反洗牌操作进行下采样，利用像素洗牌操作进行上采样。最后，输出特征z_L经过归一化，并通过 3×3 卷积头传递，以预测噪声和协方差。 

3.3 DiCo 块 

动机 

如图 4 所示，DiT 模型中的自注意力计算（无论是类条件还是文本到图像生成）均表现出明显的局部结构和显著冗余。这一观察促使我们用硬件效率更高的操作取代 DiT 中的全局自注意力。自然的替代方案是卷积，其以高效建模局部模式而闻名。我们首先尝试用 1×1 逐点卷积和 3×3 深度卷积的组合替代自注意力。 

然而，直接替换会导致生成性能下降。如图 5 所示，与 DiT 相比，修改后的模型中许多通道保持不活跃状态，表明存在显著的通道冗余。我们假设，这种性能下降源于自注意力（动态且依赖内容）相比卷积（依赖静态权重）具有更强的表示能力。为解决这一限制，我们引入紧凑通道注意力机制，以动态激活信息丰富的通道。我们在下面详细描述完整设计。 

块设计 

DiCo 的核心设计围绕卷积模块展开，如图 6（c）所示。我们首先应用 1×1 卷积来聚合逐像素的跨通道信息，然后应用 3×3 深度卷积来捕捉逐通道的空间上下文。采用 GELU 激活进行非线性变换。为进一步解决通道冗余问题，我们引入紧凑通道注意力（CCA）机制，以激活更多信息丰富的通道。如图 6（d）所示，CCA 首先通过跨空间维度的全局平均池化（GAP）聚合特征，然后应用可学习的 1×1 卷积，接着进行 sigmoid 激活以生成逐通道的注意力权重。一般来说，卷积模块的整个过程可描述为： 

3.4 架构变体 

我们建立了四种模型变体 ——DiCo-S、DiCo-B、DiCo-L 和 DiCo-XL—— 其参数数量分别与 DiT-S/2、DiT-B/2、DiT-L/2 和 DiT-XL/2 一致。与对应的 DiT 模型相比，我们的 DiCo 模型在计算成本上显著降低，Gflops 仅为 DiT 的 70.1% 至 74.6%。此外，为探索我们设计的潜力，我们将 DiCo 扩展至 10 亿参数，得到 DiCo-H。这些模型的架构配置详见附录表 5。 

4、实验 

4.1 实验设置 

数据集和指标 

遵循先前的工作 [56,90, *75]，我们在类条件 ImageNet-1K [10] 生成基准上进行了 256×256 和 512×512 分辨率的实验。我们使用 Fréchet Inception Distance（FID）[26] 作为评估模型性能的主要指标。此外，我们还报告了 Inception Score（IS）[64]、Precision 和 Recall [43] 作为次要指标。所有这些指标均使用 OpenAI 的 TensorFlow 评估工具包 [11] 计算。 

实现细节 

对于 DiCo-S/B/L/XL，我们采用与 DiT 完全相同的实验设置。具体来说，我们使用恒定学习率1×10^{-4}，不使用权重衰减，批量大小为 256。唯一应用的数据增强是随机水平翻转。在训练过程中，我们对 DiCo 的权重进行指数移动平均（EMA），衰减率为 0.9999。使用预训练的 VAE [61] 提取潜在特征。对于我们最大的模型 DiCo-H，我们遵循 U-ViT [3] 的训练设置，将学习率提高到2×10^{-4}，并将批量大小扩展到 1024 以加速训练。更多细节见附录 B 节。 

4.2 主要结果 

4.2.1 在 DiT 设置下的比较 

除了 DiT [56]，我们还选择了最近的最先进扩散模型 DiG [90] 和 DiC [75] 作为基线，因为它们同样遵循 DiT 的实验设置。表 1 展示了在 ImageNet 256×256 上的比较结果。在不同模型规模下训练 40 万次迭代后，我们的 DiCo 在所有指标上始终实现最佳或次佳性能。此外，当使用无分类器引导（CFG）时，我们的 DiCo-XL 实现了最先进的 FID 2.05 和 IS 282.17。除了性能提升，DiCo 相比基线和基于 Mamba 的模型还展示了显著的效率提升。 

表 2 展示了在 ImageNet 512×512 上的结果。在更高分辨率下，我们的模型在性能和效率上都有更大的提升。具体来说，DiCo-XL 实现了最先进的 FID 2.53 和 IS 275.74，同时将 Gflops 减少了 33.3%，并实现了比 DiT-XL/2 快 3.1 倍的加速。这些结果表明，我们的卷积架构在高分辨率图像生成中仍然高效且有效。

4.2.2 模型扩展 

为了进一步探索我们模型的潜力，我们将其扩展到 10 亿参数。如表 1 所示，与 DiCo-XL 相比，更大的 DiCo-H 模型在 FID 上进一步提升（1.90 vs. 2.05），展示了我们架构的强大可扩展性。我们在附录 C 节和 D 节提供了额外的可扩展性分析和比较结果。更多生成样本可在附录 E 节和提交的补充材料中找到。 

4.3 消融研究 

对于消融研究，我们使用小规模模型并在 ImageNet 256×256 基准上评估性能。所有模型均训练 40 万次迭代。在本节中，DiT 中的自注意力未使用 FlashAttention-2 加速，以确保与其他高效注意力机制公平比较。我们分析了 DiCo 的整体架构和各个组件的贡献，以更好地理解它们对模型性能的影响。 

4.3.1 架构消融 

我们评估了 DiCo 在各种架构设计下的性能，并与 DiT 进行了公平比较。如表 3 所示，DiCo 在所有结构中始终优于 DiT，同时还实现了显著的效率提升。这些结果凸显了 DiCo 作为 DiT 的强大高效替代方案的潜力。 

4.3.2 组件消融 

我们对 DiCo 进行了组件级分析，考察了激活函数、卷积核大小、紧凑通道注意力（CCA）和卷积模块（CM）的影响。整体消融结果总结在表 4 中。增大卷积核尺寸会带来进一步的性能提升，但以降低效率为代价，突显了性能和计算成本之间的权衡。引入 CCA 导致 FID 提升 4.81 点。如图 7 的特征可视化所示，CCA 显著增强了特征多样性，证明了其在提高模型表示能力方面的有效性。我们还将 CCA 与 SE 模块 [34] 和通道自注意力 [89] 进行了比较；尽管简单，CCA 仍实现了更优的性能和更高的效率。对于卷积模块，我们将其与几种先进的高效注意力机制（窗口注意力 [50]、聚焦线性注意力 [22]、代理注意力 [59]）进行了基准测试。结果表明，我们的 CM 在性能和计算效率上均更优。 

5、结论 

我们提出了一种新的扩散模型主干架构 —— 扩散卷积网络（DiCo），作为扩散 Transformer（DiT）的有力替代方案。DiCo 用 1×1 逐点卷积和 3×3 深度卷积的组合取代了自注意力，并引入了紧凑通道注意力机制以减少通道冗余并增强特征多样性。作为全卷积网络，DiCo 在 ImageNet 256×256 和 512×512 基准上超越了最先进的扩散模型，同时实现了显著的效率提升。我们期待进一步扩展 DiCo 并将其应用于更广泛的生成任务，如文本到图像生成。 

​

本文转载自​​AIRoobt​​ ，作者：Yuang Ai1等