告别U-Net统治？用PyTorch复现DiT，实测潜空间扩散模型的Transformer新范式

用PyTorch实战DiT：Transformer如何重塑潜空间扩散模型

当Stable Diffusion掀起AIGC革命时，U-Net作为扩散模型的标准骨架似乎已成定局。但Meta提出的DiT（Diffusion Transformer）向我们展示了另一种可能——用纯Transformer架构在潜空间完成扩散过程。本文将带您用PyTorch从零实现DiT核心模块，并通过CIFAR-10实验直观对比其与CNN架构的差异。

1. 环境准备与数据加载

在开始构建DiT前，我们需要配置适合Transformer训练的环境。建议使用PyTorch 2.0+和CUDA 11.7+环境，这对混合精度训练和Flash Attention有更好的支持：

import torch import torch.nn as nn from torch.utils.data import DataLoader from torchvision.datasets import CIFAR10 from torchvision.transforms import Compose, ToTensor, Normalize # 检查环境配置 print(f"PyTorch版本: {torch.__version__}") print(f"CUDA可用: {torch.cuda.is_available()}") # 数据预处理 transform = Compose([ ToTensor(), Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ]) dataset = CIFAR10(root='./data', train=True, download=True, transform=transform) dataloader = DataLoader(dataset, batch_size=128, shuffle=True)

关键依赖说明：

• timm：用于获取ViT风格的Patch Embedding实现
• xformers：可选，用于优化Attention计算
• einops：简化张量操作

提示：在Colab Pro上使用T4 GPU时，建议将batch_size设置为64-128以获得最佳内存利用率

2. DiT核心模块实现

2.1 Patch Embedding与位置编码

与传统ViT不同，DiT处理的是VAE编码后的潜空间特征。我们需要将4x64x64的潜变量转换为序列：

from timm.layers import PatchEmbed class LatentPatchEmbed(nn.Module): def __init__(self, img_size=32, patch_size=2, in_chans=4, embed_dim=768): super().__init__() self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size) self.num_patches = (img_size // patch_size) ** 2 self.pos_embed = nn.Parameter(torch.zeros(1, self.num_patches, embed_dim)) def forward(self, x): x = self.proj(x) # [B, C, H, W] -> [B, D, H/P, W/P] x = x.flatten(2).transpose(1, 2) # [B, D, N] -> [B, N, D] return x + self.pos_embed

参数对比表：

2.2 AdaLN-Zero调制模块

这是DiT最具创新性的设计，通过条件信息动态调整归一化参数：

class AdaLNZero(nn.Module): def __init__(self, dim): super().__init__() self.norm = nn.LayerNorm(dim, elementwise_affine=False) self.mlp = nn.Sequential( nn.SiLU(), nn.Linear(dim, 6 * dim, bias=True) ) nn.init.constant_(self.mlp[-1].weight, 0) nn.init.constant_(self.mlp[-1].bias, 0) def forward(self, x, c): shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = \ self.mlp(c).chunk(6, dim=1) x = x + gate_msa.unsqueeze(1) * self.attn( self.modulate(self.norm(x), shift_msa, scale_msa) ) x = x + gate_mlp.unsqueeze(1) * self.mlp( self.modulate(self.norm(x), shift_mlp, scale_mlp) ) return x def modulate(self, x, shift, scale): return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)

2.3 条件集成系统

DiT采用Classifier-free Guidance策略，需要特殊处理条件嵌入：

class LabelEmbedder(nn.Module): def __init__(self, num_classes, hidden_size, dropout_prob=0.1): super().__init__() self.embedding = nn.Embedding(num_classes + 1, hidden_size) self.num_classes = num_classes self.dropout_prob = dropout_prob def forward(self, labels, train=False): if train and self.dropout_prob > 0: mask = torch.rand(labels.shape[0]) < self.dropout_prob labels[mask] = self.num_classes # 使用unconditional token return self.embedding(labels)

3. 完整DiT模型组装

整合各组件构建完整DiT模型：

class DiT(nn.Module): def __init__(self, input_size=32, patch_size=2, in_chans=4, depth=12, embed_dim=768, num_heads=12): super().__init__() self.patch_embed = LatentPatchEmbed(input_size, patch_size, in_chans, embed_dim) self.t_embed = nn.Sequential( nn.Linear(embed_dim, embed_dim), nn.SiLU(), nn.Linear(embed_dim, embed_dim) ) self.y_embed = LabelEmbedder(1000, embed_dim) self.blocks = nn.ModuleList([ DiTBlock(embed_dim, num_heads) for _ in range(depth) ]) self.final_layer = FinalLayer(embed_dim, patch_size, in_chans * 2) def forward(self, x, t, y): x = self.patch_embed(x) t = self.t_embed(t) y = self.y_embed(y) c = t + y for block in self.blocks: x = block(x, c) return self.final_layer(x, c)

模型配置对照：

4. 训练与实验结果分析

4.1 训练配置要点

在CIFAR-10上的训练建议配置：

from diffusers import DDPMScheduler model = DiT(input_size=32, patch_size=2, in_chans=4) optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4) scheduler = DDPMScheduler( num_train_timesteps=1000, beta_schedule="linear" ) # 混合精度训练 scaler = torch.cuda.amp.GradScaler()

关键训练技巧：

• 使用梯度裁剪（max_grad_norm=1.0）
• 线性warmup（约5000步）
• EMA模型权重平均（decay=0.9999）

4.2 与U-Net的对比实验

我们在CIFAR-10上对比了DiT-S与同规模U-Net的表现：

虽然DiT展现出更好的生成质量，但其计算开销显著更高。实际部署时需要权衡：

• 适合DiT的场景：

  • 需要最高生成质量
  • 有条件使用大型GPU集群
  • 需要模型可扩展性

• 适合U-Net的场景：

  • 边缘设备部署
  • 实时生成需求
  • 小规模数据集

5. 进阶优化方向

对于希望进一步提升DiT性能的开发者，可以考虑以下优化：

内存优化技巧：

# 启用Flash Attention from torch.backends.cuda import sdp_kernel with sdp_kernel(enable_flash=True): output = model(input) # 梯度检查点 from torch.utils.checkpoint import checkpoint x = checkpoint(block, x, c)

架构改进建议：

• 尝试混合精度训练（AMP）
• 加入LoRA进行参数高效微调
• 实验不同的patch大小（1x1到4x4）

在ImageNet-256数据集上，经过优化的DiT-XL可以达到2.17 FID的顶尖水平，这证实了Transformer在扩散模型中的巨大潜力。不过值得注意的是，要达到最佳性能，通常需要：

• 更大的模型规模（数亿参数）
• 更长的训练时间（百万步级）
• 大规模数据增强

DiT的成功不仅在于架构创新，更展示了如何将Transformer的优势与扩散模型的理论基础完美结合。虽然目前计算成本较高，但随着硬件进步和算法优化，Transformer很可能成为下一代扩散模型的标准骨架。