扩散模型在图像生成领域逐渐成为主流方案，而去噪网络的主干结构直接决定了训练成本、推理速度和生成质量。U-Net凭借卷积归纳偏置长期占据主导位置，但随着模型规模扩大，Diffusion Transformer（DiT）开始展现更强的可扩展性和并行效率。本文从计算复杂度、内存占用、训练效率等维度对两者进行对比，帮助读者在具体任务中做出合理选择。
架构差异与计算复杂度对比
U-Net采用编码器-解码器结构，通过多次下采样降低空间分辨率，在高分辨率特征图上使用标准卷积提取局部特征。卷积操作的FLOPs与像素数呈近似线性关系，这意味着在分辨率翻倍时计算量大致增加四倍。U-Net的跳跃连接保留了不同尺度的细节信息，对图像生成任务非常友好。但卷积核的感受野有限，为了捕捉长距离依赖需要堆叠很深，模型参数效率下降。
DiT首先将图像切分为固定大小的patch，每个patch经过线性投影得到token序列，然后送入标准的Transformer编码器。自注意力机制使每个token都能与全局其他token交互，感受野覆盖整张图像。但自注意力的计算复杂度为O(N^2)，其中N是token数量，即patch数量的平方。对于1024分辨率的输入，如果patch大小为16，则N为4096，注意力矩阵的元素数量超过1600万，显存和计算开销会显著上升。实践中常通过窗口注意力、稀疏注意力或FlashAttention来降低峰值显存。
import torch
import torch.nn as nn
class UNetBlock(nn.Module):
def __init__(self, in_ch, out_ch):
super().__init__()
self.conv1 = nn.Conv2d(in_ch, out_ch, 3, padding=1)
self.conv2 = nn.Conv2d(out_ch, out_ch, 3, padding=1)
self.relu = nn.ReLU()
def forward(self, x):
x = self.relu(self.conv1(x))
x = self.relu(self.conv2(x))
return x
class DiTBlock(nn.Module):
def __init__(self, dim, num_heads):
super().__init__()
self.norm1 = nn.LayerNorm(dim)
self.attn = nn.MultiheadAttention(dim, num_heads, batch_first=True)
self.norm2 = nn.LayerNorm(dim)
self.mlp = nn.Sequential(
nn.Linear(dim, dim * 4),
nn.GELU(),
nn.Linear(dim * 4, dim),
)
def forward(self, x):
# x: (B, N, C)
attn_out, _ = self.attn(self.norm1(x), self.norm1(x), self.norm1(x))
x = x + attn_out
x = x + self.mlp(self.norm2(x))
return x
从代码可以看出，U-Net块只涉及局部卷积，参数量相对固定，计算复杂度随特征图尺寸线性增长。DiT块包含自注意力和前馈网络，参数量更大，但能够建模全局依赖关系。需要注意的是，DiT中的自注意力在序列长度增加时计算量会快速膨胀，而卷积则不会出现这种平方级的增长。
内存占用与推理速度差异
推理阶段，U-Net需要缓存编码器各层特征用于跳跃连接，这导致显存峰值较高。不过卷积算子经过长期优化，GPU利用率很高，推理速度在低分辨率下通常优于Transformer结构。DiT不需要多尺度特征缓存，但注意力矩阵的存储随序列长度平方增长。以批量大小为1、分辨率为512的图像为例，patch大小为16时token数为1024，单个注意力矩阵约4MB；当分辨率提升到1024，token数变为4096，单个注意力矩阵约64MB，多层堆叠后显存增长非常明显。
实际测试中，256分辨率下U-Net与DiT的推理耗时差距不大，DiT可能略慢。但在1024分辨率下，U-Net由于卷积局部性和成熟的下采样策略，显存增长相对平缓；DiT如果不采用优化，显存会迅速超过消费级显卡容量。引入FlashAttention后，DiT的显存占用可以下降数倍，推理速度也接近U-Net。参数量方面，同等级别的DiT通常比U-Net参数更多，但参数利用率更高。
import torch
def count_params(model):
return sum(p.numel() for p in model.parameters() if p.requires_grad)
def estimate_attention_flops(N, dim):
# 自注意力FLOPs估算: 2 * N * N * dim + 其他线性项
return 2 * N * N * dim
# 示例：patch数量与注意力FLOPs的关系
for img_size, patch_size in [(256, 16), (512, 16), (1024, 16)]:
N = (img_size // patch_size) ** 2
flops = estimate_attention_flops(N, 768)
print(f"分辨率 {img_size}，patch {patch_size}，token数 {N}，注意力FLOPs约 {flops/1e6:.1f}M")
上述统计只计算了单个注意力层的理论FLOPs，尚未包含前馈网络和归一化层。但即使如此，也可以看出分辨率提升对DiT计算量的影响远大于对U-Net卷积操作的影响。因此，在硬件资源受限的场景中，针对DiT的优化策略往往比U-Net更加关键。
训练效率与可扩展性分析
U-Net由于卷积的归纳偏置，在小规模数据集上更容易收敛，训练初期损失下降更快。其局部连接和权重共享特性降低了对数据量的依赖。DiT去除了这些归纳偏置，完全依赖数据学习空间关系，因此在小数据集上容易过拟合或收敛缓慢。但在千万级甚至亿级图像数据集上，DiT的性能随模型参数量增加呈现更陡峭的提升曲线，而U-Net在参数量达到一定规模后收益递减。
分布式训练方面，Transformer结构对GPU之间的并行通信更加友好。自注意力可以高效切分到多个设备上，配合序列并行和张量并行，DiT能充分利用大规模集群。U-Net的卷积操作在多卡训练时通常采用数据并行，模型并行实现复杂，收益有限。此外，DiT的LayerNorm和残差结构使训练更稳定，可以使用更大的学习率，从而加快收敛速度。
# 配置对比：U-Net与DiT的基础参数设置
unet_config = {
"model_type": "unet",
"image_size": 256,
"in_channels": 3,
"base_channels": 128,
"channel_mult": [1, 2, 4, 8],
"attention_resolutions": [16, 8],
"num_res_blocks": 2,
}
dit_config = {
"model_type": "dit",
"image_size": 256,
"patch_size": 16,
"in_channels": 3,
"hidden_size": 768,
"depth": 12,
"num_heads": 12,
"mlp_ratio": 4.0,
}
从配置上看，U-Net的通道数随下采样逐级翻倍，参数量主要集中在中低层特征；DiT则通过统一的隐藏宽度和深度来控制容量。这种结构差异导致两者在训练时的数据需求、优化器选择以及分布式策略都有明显不同。
实际任务中的选型建议
如果任务分辨率较低（256或512）、可用数据量有限、或者需要快速迭代原型，U-Net仍然是性价比最高的选择。其实现成熟、推理速度快、对硬件要求低。如果是超分辨率、医学图像分割等需要精细局部特征的任务，U-Net的跳跃连接能提供更好的细节恢复能力。此外，已有大量预训练U-Net权重可供迁移，能显著降低开发成本。
当任务涉及高分辨率图像生成、大规模数据集训练、或者追求生成质量的极限时，DiT的扩展潜力更大。配合条件控制机制，DiT在多模态生成任务中表现突出。对于显存受限的环境，可以选择优化后的DiT变体，如使用窗口注意力或线性注意力，在效率和全局建模之间取得平衡。最终选型应基于实际任务的分辨率、数据规模、硬件条件以及对生成质量的要求综合判断，必要时可以尝试两者混合的架构方案。
Diffusion TransformerDiTU-Net架构修改时间：2026-09-24 12:26:13