DiC：重新思考扩散模型中的 Conv3x3 设计

paper是PKU发表在arxiv 2025的工作

Paper Title：DiC: Rethinking Conv3x3 Designs in Diffusion Models

Code:https://github.com/YuchuanTian/DiC

Abstract

扩散模型在视觉生成任务中表现出色。最近，这些模型已经从传统的 U 型 CNN-Attention 混合结构转变为完全基于 Transformer 的各向同性架构。虽然这些 Transformer 表现出强大的可扩展性和性能，但它们对复杂的自注意力操作的依赖导致推理速度缓慢。与这些工作相反，我们重新思考了深度学习中最简单但最快的模块之一 3x3 卷积，以构建一个放大的纯卷积扩散模型。我们首先发现，编码器-解码器沙漏设计优于 Conv3x3 的可扩展各向同性架构，但仍低于我们的预期。进一步改进架构，我们引入了稀疏跳过连接以减少冗余并提高可扩展性。基于该架构，我们引入了条件改进，包括特定于阶段的嵌入、中间块条件注入和条件门控。这些改进导致了我们提出的扩散 CNN (DiC)，它是一种快速而有竞争力的扩散架构基线。在不同规模和设置下的实验表明，DiC 在性能方面远远超越了现有的扩散变压器，同时保持了良好的速度优势。

1. Introduction

基于扩散的图像生成模型最近引起了广泛关注，稳定扩散 [27] 等著名作品展示了令人印象深刻的结果。随着这些模型的不断发展，它们现在能够生成越来越逼真的图像。包括与 Sora 相关的作品在内的最新进展已将这种能力扩展到视频生成，制作出可与现实世界镜头相媲美的高质量、时间一致的视频。

自注意力机制在最近的许多扩散模型中起着关键作用。早期的工作 [18, 31] 将卷积 U-Net 结构与自注意力融合在一起。最近，U-ViT [1] 和 DiT [26] 等架构已完全过渡到基于注意力的设计，完全放弃了卷积 U-Net。这些模型展示了卓越的生成能力，尤其是在扩大规模时。 较新的文本到图像和文本到视频模型，包括 Stable Diffusion 3 [13]、FLUX [21]、PixArt 系列 [2–4] 和 OpenSora [20, 39]，也采用了完全基于 Transformer 的架构，实现了令人印象深刻的生成质量。

不幸的是，在扩大规模时，包含自注意力机制的扩散模型会面临巨大的计算开销和延迟。复杂的 Transformer 架构对实时、资源受限的应用构成了相当大的挑战，在这些应用中，大规模 Transformer 因其巨大的时间成本而不切实际。

人们已经做出了努力来加速扩散模型中的自注意力。例如，ToDo [30] 和 PixArt-Σ [3] 实现了更有效的自注意力机制。然而，这些方法仍然局限于自注意力范式。其他研究 [32, 38] 旨在用基于状态空间模型的架构 [15] 取代自注意力，试图绕过其固有的低效率。尽管有这些创新，但典型潜在域图像分辨率的生成速度仍然不是最优的，无法满足实时或大规模应用的需求。

在这项工作中，我们重新思考了简单而高效的全卷积扩散模型架构，该架构具有显著的速度优势。作为自注意力的对应物，卷积被广泛认为是硬件友好的，在各种平台上都有出色的支持。在众多卷积运算中，标准的 3x3 stride1 卷积以其出色的速度脱颖而出，这在很大程度上要归功于 Winograd 加速 [22] 等广泛的硬件优化技术。例如，RepVGG [11] 等研究已经证明了简单的卷积如何实现高效率和具有竞争力的性能。

然而，3x3 卷积具有固有的局限性，特别是受限制的感受野，这可能会阻碍其在复杂生成任务中的可扩展性。我们首先进行了初步试验，在现有的可扩展框架内扩展纯 3x3 CNN，但这些实验的结果仍然不如 Diffusion Transformer 同类产品。

为了解决这个问题，我们提出了一系列对传统 ConvNets 的调整，形成了一个路线图（如图 1 所示），专门针对扩大的扩散模型对其进行了调整。 首先，我们专注于改进模型架构，以更好地利用 Conv3x3。现有的设计分为三类：各向同性架构，如 DiT [26]、具有跳跃的各向同性架构，如 U-ViT [1]，以及用于典型 CNN 和自注意混合模型的编码器-解码器沙漏结构。经验观察表明，沙漏架构对于纯 Conv3x3 模型更有效，因为编码器中的下采样和上采样扩大了弥补 Conv3x3 狭窄范围的感受野。跳跃连接也很重要，但用于廉价卷积块扩大的传统块式跳跃过于密集，无法有效用于大型 ConvNets。 为了解决这个问题，我们引入了稀疏跳过连接，减少了跳过的次数，同时确保必要信息有效地从编码器流到解码器。

此外，我们专注于改进条件以更好地适应沙漏卷积网络。在现有模型中，一组条件嵌入通常映射到不同阶段的不同块。然而，在编码器-解码器卷积网络中，每个阶段都在不同的特征空间中运行。因此，我们引入了特定于阶段的嵌入，确保每个阶段使用独立的、不重叠的条件嵌入表。我们还仔细检查了嵌入注入的位置，并建议在块中间注入条件以获得更好的性能。此外，我们借用了特征图上的条件门控。此外，我们改进了内部块结构，用 GELU 替换了所有激活函数。

借助这些增强功能，我们提出了 DiC，这是一种由 Conv3x3 组成的扩散模型。架构和条件改进共同使我们的模型能够实现出色的结果，同时保持其速度优势。在各种规模和设置上进行的实验证明了 DiC 相对于 Diffusion Transformers 的性能优势及其高吞吐量特性。

图 1. 3x3 卷积扩散模型的路线图。 性能通过经过 200K 次迭代训练的模型的 FID（↓）来衡量。一系列提出的模型改进使 DiC 比 Diffusion Transformers 更具优势。

2. Related Work

扩散架构中的自注意力。由于自注意力在捕获长距离依赖关系方面具有显著的有效性，它已成为现代扩散架构的基石。早期基于 UNet 的扩散模型 [18, 31] 在网络的较高阶段加入了自注意力层，在这些阶段捕获全局上下文最为关键。Dhariwal 等人通过 ADM [10] 进一步发展了这一技术，将自注意力扩展到所有阶段，这显著提升了模型的生成能力。

近期的作品，例如 DiT [26] 和 U-ViT [1]，已经完全采用了基于 Transformer 的架构，用完全自注意力设计取代了传统的 U-Net 主干。 这些模型不仅实现了卓越的性能，而且还表现出了出色的可扩展性，尤其是在模型尺寸和 FLOP 较大的情况下。大多数后续作品，例如 PixArt-α [2] 和 Stable Diffusion 3，都采用了完全 Transformer 主干设计，展示了出色的文本到图像生成能力。

除此之外，一些研究继续探索和改进自注意力和 Transformer 架构在扩散模型中的使用，进一步突破了基于 Transformer 的生成模型的效率界限。例如，[3、30、33] 研究了自注意力的下采样标记；[37] 使用线性注意力。其他研究 [32、38] 选择使用状态空间模型进行扩散，但 SSM 在普通形状图像的延迟方面存在问题 [32]。

与这些研究不同的是，我们通过使用 3x3 卷积设计整个扩散模型，将架构简单性推向极致。通过精心设计的结构优化，这个精简的网络不仅充分利用了 3x3 卷积的速度优势，而且还实现了与基于 Transformer 的模型相当的扩散生成性能。

高效卷积。从历史上看，AlexNet [19]、VGG [29] 和 ResNet [17] 等架构一直主导着计算机视觉任务。虽然 transformer [34] 后来成为许多视觉应用的最新技术，但它们的高计算成本和延迟使它们在资源受限的环境中不太实用，而 CNN 仍然擅长于此。最近的进展重新审视并改进了 CNN，使它们在性能和效率方面能够与 transformer 相媲美甚至超越它。

例如，ConvNeXt [24] 通过宏观和微观架构变化（例如增加内核大小和优化规范化）实现了 Transformer 级性能。可变形卷积 (DCN) [6, 35, 36] 通过学习空间自适应内核形状来增强感受野，而 RepLKNet [12] 利用大量 31x31 内核来捕获长距离依赖关系。类似地，Wavelet Conv [14] 使用小波变换在频域中运行以扩展 CNN 功能。

另一方面，有些工作致力于提高简单卷积的效率。RepVGG [11] 专注于简单 3x3 卷积的潜力，采用重新参数化技巧来提高其有效性。受其速度潜力的启发，我们探索在扩散中使用纯 3x3 卷积，结合架构和条件改进，以在基于扩散的生成模型中实现具有竞争力的结果。但是，我们不会使用重新参数化 [11] 之类的技巧。我们的方法表明，即使采用简单的设计，高质量的生成也可以既高效又有效。

3. Method

3.1. Preliminaries: Conv3x3 as the Ingredient

---

我们选择普通步长为1的全卷积 $3\times 3$ 作为模型的主要组成部分。步长为1的 $3\times 3$ 卷积运行速度极快，这主要得益于现代深度学习框架中的硬件和算法优化。操作优化得益于 Winograd 算法 [22]，该算法通过减少常见卷积操作所需的乘法次数，将计算量降低了 $\frac{5}{9}$。对于 $3\times 3$ 的卷积核，该方法通过高效的矩阵变换显著加速了计算。

此外，与深度卷积（depth-wise convolution）不同，$3\times 3$ 卷积实现了高度的计算并行性，使现代 GPU 能够充分发挥其计算能力。这种并行性与较低的内存访问开销相结合，使步长为 $1$ 的 $3\times 3$ 卷积尤其高效。相比于更大或更复杂的卷积操作，它们在计算成本和表示能力之间提供了理想的平衡，使其成为实时应用和大规模模型的实用选择。

图 2. 扩展架构的基本块。与 [26] 中广泛应用的复杂 Transformer 块相比，我们的 Conv 3x3 基本块更简单，硬件支持更好，这是实现高吞吐量的关键。

3.2. Architectures for Scale-Up

---

我们的基本块。我们重新审视了经典的扩散模型 [10, 31]，并发现 3x3 卷积也得到了广泛的应用。然而，这些模型并不是纯粹的 CNN：模型中的基本块单元由两个连续的 3x3 卷积层组成，后面跟着一个关注整个图像的自注意力模块。在每次卷积之前，使用 GroupNorm 对输入进行归一化，然后传递给 SiLU 的非线性激活函数。 值得注意的是，块内的中间特征保持相同数量的通道，确保特征传播稳定。整体块采用残差结构 [17]，其中输入通过快捷连接直接添加到输出。

我们提出的扩散模型 DiC（3x3 卷积降噪器）的路线图从上述重新审视的卷积基本块开始。为了满足路线图的最终目标，我们在卷积后删除了自注意力层，以形成由两个 3x3 卷积组成的最小单元块（基本块）。其他块设计暂时保持不变。

主流架构。以此 Conv 3x3 基本块为基础，我们通过堆叠多个这样的块并增加网络的深度和宽度来扩展模型，并且我们面临着众多架构选择（如图 3 所示）：

图 3. 主流扩散架构。我们尝试使用 Conv3x3 基本块的现有扩散架构 (a、b、c)。我们进一步使用稀疏跳跃改进 U-Net (d)。

1. 各向同性架构。受 DiT 及其后续作品 [2、5、16、26] 的启发，此设计具有简单的垂直堆叠结构。经过patch化后，中间特征图大小在整个模型中不会改变。它因其著名的可扩展性而被广泛采用，尤其是在扩展到大型模型时。

2. 具有跳跃连接的各向同性架构。这种方法通过在非相邻层之间引入长跳跃连接来扩展柱状各向同性架构，类似于 U-ViT [1] 中使用的设计。与传统的 U-Net 不同，尽管采用了跳跃设计，但该架构在整个网络中保持一致的空间分辨率，避免了任何上采样和下采样操作。

3. U-Net 沙漏架构。这种经典设计 [28] 由编码器和解码器组成，形成漏斗状。编码器逐步下采样特征维度，而解码器对称上采样空间维度。除了沙漏主干之外，编码器中的每个块都通过跳跃连接连接到其相应的解码器块。这项工作在早期的扩散工作中得到了广泛的应用 [10, 18, 31]。

实验分析。使用基础 Conv 3x3 基本块，我们对上述架构选择进行了扩展实验。结果总结在表 1 中，表明所有架构的性能通常都不理想。然而，这些实验揭示了几种模式，为设计有效的基于 3x3 卷积的扩散模型提供了宝贵的见解。

表 1. 在 ImageNet 256×256 条件生成上比较 Conv3x3 扩展的架构。由于扩大了感知范围，Conv3x3 充分利用了 U-Net Hourglass 架构的潜力；稀疏化跳过连接进一步改进了模型。

首先，对于仅由 3x3 卷积组成的扩散模型，编码器-解码器架构沙漏是必不可少的。这源于 3x3 卷积的一个基本限制：它们固有的接受场受限。在纯各向同性、标准变压器架构设置中，每个 3x3 卷积在每个方向上仅将接受场扩大一个像素。因此，实现全局接受场需要深度堆叠卷积，导致计算和参数效率低下。 相比之下，编码器-解码器架构显著缓解了这一限制。通过在编码器中逐步下采样并在解码器中上采样，更高阶段的 3x3 卷积可以有效地感知输入的更大区域。例如，在更深的阶段，单个 3x3 卷积可以覆盖原始图像空间中 6x6 或 12x12 的区域，从而大大增加了模型的接受场。这种分层方法使模型能够有效地捕捉局部和全局背景，从而使得生成任务更加有效。

其次，跳跃连接在基于 Conv3x3 的架构中至关重要。一方面，在编码器-解码器结构中，它们通过将特征图直接传递给解码器来减少下采样过程中的信息损失，从而丰富其表示。另一方面，跳跃通过提供额外的梯度路径并保留重要特征来加速训练并提高生成质量，从而实现更高效、更有效的扩散建模。

首先，我们采用 U-Net 作为我们使用的通用模型架构。

通过步进跳跃改进 U-Nets。然而，由于扩展我们基于 Conv3x3 的扩散模型需要堆叠大量卷积层，因此密集的块式跳跃连接成为瓶颈。大量块式跳跃特征必须连接到特征并在相应的解码器块内的通道维度上合并，这会消耗过多的计算成本。这种低效率妨碍了模型的可扩展性。为了解决这个问题，我们提出了步进跳跃连接，其中跳跃仅每隔几个块应用一次，而不是在每个块之后应用一次。这种方法提高了性能，我们的实验（见表 1）表明，与密集连接的跳跃设计相比，它确实产生了更好的结果。

3.3. Conditioning Improvements

---

除了架构改进之外，我们认为条件反射也可以改进。在模型规模扩大后，之前的条件反射设计需要根据架构进行改进。

特定阶段的嵌入。在传统的扩散模型 [10, 31] 中，以前的模型共享一个唯一的嵌入表。然而，在扩大的 Conv3x3 模型中，我们依靠编码器-解码器结构来扩展接受场。编码器-解码器的每个阶段都以不同的通道维度运行，反映了阶段之间显著的结构和功能差异。鉴于特征维度在不同阶段有所不同，将相同的条件嵌入均匀地应用于不同阶段的所有块可能并不理想。每个阶段的不同角色和表示表明，更适合阶段的条件嵌入方法可能更有效，因为单个嵌入集可能无法充分捕捉每个阶段的不同特征。

作为一项改进，我们建议对编码器-解码器结构的每个阶段使用特定于阶段的条件嵌入。每个阶段都有自己独立的嵌入，嵌入维度与该阶段的特征维度对齐。这确保了条件嵌入适合每个阶段的特定特征和规模。

值得一提的是，我们分析了阶段特定嵌入带来的开销。结果表明，增加的开销微乎其微。具体来说，添加新的阶段特定嵌入只会使模型大小增加 14.06M 个参数，这仅占整个模型大小的 2%。引入的计算开销甚至可以忽略不计，只增加了 12M FLOPs，对整体计算效率几乎没有影响。这表明，加入阶段特定嵌入不会显著影响模型的性能或计算成本。

此外，在训练过程中，跨不同阶段的嵌入上的标签删除操作应该同步。这可以防止阶段之间发生任何标签泄漏，并确保模型以一致且隔离的方式进行学习，从而提高扩散模型的稳定性和性能。

条件注入位置。我们设计的另一个关键方面是确定注入条件信息的最佳位置。在扩散模型中，有两种流行的条件注入策略。第一种方法是在每个块的最开始处注入条件，通常通过 LayerNorm，如 [5, 26] 等模型中所示。 第二种方法是在块的中间引入条件，如 [10, 31] 中所使用的。

通过实验，我们发现，对于我们的全卷积扩散架构，将条件注入每个块中的第二个卷积层可获得最佳性能。这种放置有效地调节了特征表示，增强了模型的生成质量，同时又不影响其效率。

条件门控。此外，我们采用了 DiT 的 AdaLN [26] 中的条件门控机制。除了执行通道级仿射映射到特征图的传统仿射条件之外，AdaLN 还引入了一个门控向量，该向量沿通道维度缩放特征，从而对条件过程提供更动态、更细粒度的控制。此修改增强了模型适应不同条件的能力，并提高了整体生成质量。我们的实验表明，这种增强可显著提高性能，进一步证明了这种条件策略在扩散模型中的有效性。

3.4. Minor Modification

---

我们重新审视了整个模型的激活函数。 受 ConvNeXt [24] 的启发，我们用 GELU 替换了常用的 SiLU 激活函数，这是 Transformer 架构中的标准选择。除了其他修改之外，这一变化还有助于性能的适度但持续的提升。