扩散模型新架构：DiffiT 用时间依赖自注意力重构去噪网络

简介PDF版论文《DiffiT: Diffusion Vision Transformers for Image Generation》来自NVIDIA研究团队面向图像生成、扩散模型和视觉Transformer方向的研究者与高年级学生针对传统U-Net去噪网络在架构设计上的不充分问题正式提出DiffiT模型。其核心贡献包括混合U形编码器-解码器结构以及随时间步自适应调整的时间依赖自注意力模块同时给出用于高分辨率合成的latent DiffiT在类条件与无条件合成任务中均达SOTAImageNet-256上FID低至1.73。包内为1个PDF文件大小43.22MB已有336人浏览学习。完整下载后可直接精读论文获取方法设计、实验设置、结果对比、消融分析及开源项目地址文中对扩散模型、ViTs、去噪网络、FID等关键知识也有清晰阐释适合用于文献调研、论文复现或入门到进阶的系统学习。1. 用 Transformer 重构扩散模型去噪网络DiffiT 的思路值不值得跟扩散模型的生成质量这几年已经被验证得足够强但多数去噪网络依然是卷积 U-Net 加几层自注意力的组合架构上的探究远没有采样器、调度器那么热闹。这篇论文把视觉 Transformer 直接搬进扩散模型提出了 DiffiT核心改动是让自注意力模块的 Q、K、V 都随时间去适应而不是只把时间 embedding 加到特征上。在 CIFAR-10、FFHQ-64 上拿到当时最好的 FID在 ImageNet-256 的 latent 空间把 FID 做到了 1.73。对想深入扩散模型架构、或者要在高分辨率图像生成上换一个主力模型的人来说这份资源的价值在于它给了你一个完整的、可对照源码阅读的 Transformer 去噪网络设计范式而不是零散的小 trick。2. 时间依赖自注意力DiffiT 的核心设计以及它为什么能work2.1 传统去噪网络的时间建模方式到底卡在哪扩散模型的去噪网络在不同时间步承担的任务差异非常大。噪声强的时候图像的高频信息基本被破坏网络主要靠低频结构来判断方向到了生成后期轮廓已经定了网络要补的是高频细节。理想情况下网络在不同阶段应该有不同的感受野偏好和注意力分布。但传统 U-Net 里的时间 embedding 一般是通过加法或 channel-wise 的 scale 和 shift 注入到残差块里的卷积核本身没有时间概念。自注意力层虽然能建模长距离依赖但大多数实现只在低分辨率特征图上使用因为高分辨率下计算量承受不住。于是出现一个矛盾越是需要精细调整的后期阶段网络的感知能力反而越有限。DiffiT 的做法很直接把时间 token 作为输入和空间 token 一起线性投影到 Q、K、V。这样每个时间步的注意力权重就是不同的相当于注意力层本身变成了一个时间条件模型不再需要额外的 AdaLN 调制。2.2 TMSA 的数学形式以及它的实现细节TMSATime-dependent Multi-head Self-Attention的计算方式可以理解为对每个空间 token xs 和时间 token xt分别用两套权重投影到 query、key、value 空间然后相加形成最终的 Q、K、V。假设某层输入空间 token 是 xs时间 token 是 xt那么 query、key、value 的投影公式就是import torch import torch.nn as nn class TimeDependentProjection(nn.Module): 将空间 token 与时间 token 分别投影后相加作为 Q/K/V 的输入。 这也是 DiffiT 中 TMSA 的基础计算单元。 def __init__(self, dim, num_heads): super().__init__() self.num_heads num_heads head_dim dim // num_heads # 空间 token 的投影权重 self.wq_s nn.Linear(dim, dim) self.wk_s nn.Linear(dim, dim) self.wv_s nn.Linear(dim, dim) # 时间 token 的投影权重 self.wq_t nn.Linear(dim, dim) self.wk_t nn.Linear(dim, dim) self.wv_t nn.Linear(dim, dim) def forward(self, xs, xt): # xs: [B, N, C] 空间 token # xt: [B, 1, C] 时间 embedding q self.wq_s(xs) self.wq_t(xt) # 时间 token 广播加到每个空间 token k self.wk_s(xs) self.wk_t(xt) v self.wv_s(xs) self.wv_t(xt) return q, k, v这段代码对应公式 3 到 5。逻辑上等价于把空间 token 和时间 token 拼接后做一次线性投影但拆成两套权重的写法在初始化上更干净而且时间 token 对 Q、K、V 的影响是分别可调的。相比 AdaLN 方案里时间对整个特征图做统一的 scale 和 shift这里时间的信息可以直接影响注意力矩阵的元素所以对空间交互的调节粒度细得多。参数上要注意num_heads和head_dim的关系。DiffiT 的默认设置里模型宽度一般是 256 或 512head 数为 4 或 8。如果你在 128 分辨率的任务上复现dim256, num_heads4会比dim512, num_heads8更稳妥因为注意力矩阵的时间变化在低分辨率下更容易过拟合head 数多了反而让相对位置偏置的学习变得不稳定。我一般习惯先用小配置跑通训练流程再逐步加宽。这里的wq_t等时间投影是对所有空间 token 共享的所以参数量增加不多但带来的是每个时间步都有不同的注意力模式。从实验效果看这种设计的收益在生成后期特别明显因为此时网络需要关注局部细节而这恰恰是时间自适应注意力最容易发挥作用的地方。2.3 相对位置偏置一个容易被忽视的细节TMSA 的注意力计算里还有一项是相对位置偏置 B。大部分 Transformer 的 2D 相对位置偏置都是提前初始化好的一组可学习参数DiffiT 也遵循这个做法。但这里有个讨论点论文特别指出B 不应该直接受时间 embedding 影响否则空间和时间信息会在这个模块里纠缠导致训练困难。实际操作上B 的 shape 是[num_heads, L, L]其中 L 是 token 序列长度。对 256 分辨率、patch size 为 8 的 latent 模型来说L 是 1024这张偏置表会相当占内存。我有一次在 512 分辨率上复现时就因为这个偏置表加上注意力矩阵的缓存显存直接超了。解决办法是在低分辨率阶段共享偏置表或者用插值的方式从 256 分辨率初始化到 512 分辨率而不是重新随机初始化。3. 把 DiffiT 跑起来环境搭建、推理与参数配置3.1 代码仓库结构与关键文件DiffiT 的官方开源仓库结构比较清晰代码分为几个模块模型定义、训练逻辑、采样逻辑、数据集加载。核心文件包括 DiffiT 块定义、TMSA 实现、训练脚本和采样脚本。仓库里也提供了预训练模型的下载脚本。不建议直接从零开始训练先把预训练权重跑通推理再考虑在自己的数据集上做微调或从头训练。因为扩散模型的训练周期长如果你只是想知道这个架构效果如何跑一次采样就能判断。3.2 环境准备与依赖安装建议使用 Python 3.9 和 PyTorch 2.xCUDA 版本在 11.7 以上。安装依赖时注意 diffusers、einops、timm 这些包的位置某些旧版本组合会引发算子兼容问题。# 建议使用 conda 创建独立环境避免和已有训练框架冲突 conda create -n diffit python3.9 conda activate diffit # 按显卡驱动选择对应版本的 PyTorch pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118 # 安装其他依赖 pip install einops timm diffusers omegaconf安装依赖时最容易踩的坑是 diffusers 版本。新版本的 diffusers 改了部分 API导致官方代码里的采样循环跑不通。我一般会锁定 diffusers 版本在 0.21 左右这个版本和官方代码配合比较稳定。如果不想锁版本就要做好改调用接口的心理准备。3.3 加载预训练权重并生成图像官方仓库提供了两个main.py入口一个对应 image-space 模型一个对应 latent 模型。image-space 模型直接在像素空间做扩散适合 CIFAR-10 和 FFHQ-64 这样的中低分辨率数据集latent 模型则需要先有一个预训练的自编码器把图像压缩到 latent 空间再进行扩散适合 ImageNet-256/512。# 以 latent DiffiT 在 ImageNet-256 上采样为例 # 配置项ckpt 为权重路径cfg 为采样步数scale 为引导强度 import torch from PIL import Image from main_imagenet import DiffiT_Model # 初始化模型结构dim 为 Transformer 宽度heads 为注意力头数 model DiffiT_Model(dim256, heads8, img_size256, num_classes1000) # 加载预训练权重strictFalse 是为了兼容不同 patch_size 的插值状态 state_dict torch.load(checkpoints/diffit_imagenet256.pt, map_locationcuda:0) model.load_state_dict(state_dict, strictFalse) model.eval().cuda() # 指定类别索引50 步 DDIM 采样cfg 设为 1.0 表示不使用 classifier guidance with torch.no_grad(): z torch.randn(1, 4, 32, 32, devicecuda:0) # latent 是 256/832 分辨率 y torch.tensor([207], devicecuda:0) # 207 对应金毛寻回犬 latents model.sample(z, y, num_steps50, cfg_scale1.0) # 将 latent 送入自编码器解码为像素空间 with torch.no_grad(): image model.vae.decode(latents / 0.18215) # 保存结果 Image.fromarray((image.squeeze(0).permute(1, 2, 0).cpu().numpy() * 255).astype(uint8)).save(output.png)这段代码有三个参数值得说明。dim256是 Transformer 的嵌入维度直接影响注意力计算量和模型容量如果你的显存足够加宽到 384 会明显提升生成细节但训练和推理时间也会变长latents的形状需要和 VAE 的下采样倍数对齐ImageNet-256 对应 8 倍下采样所以是 32×32换数据集时这个尺寸必须改num_steps50是采样步数步数越少速度越快但质量掉得明显实际使用中 50 步到 100 步之间是性价比区间。cfg_scale1.0和 2.0 的差距在于类别信息的注入强度过高的引导会让生成多样性和保真度失衡这个参数和 FID 高度相关。3.4 推理性能与显存估算以单张 24GB 显存的显卡为例ImageNet-256 的 batch size 1 推理大约需要 6 到 8GB 显存50 步采样耗时约 15 到 30 秒。如果显存溢出优先检查 VAE decode 阶段的内存峰值这是大多数人忽略的地方。VAE 在解码 latent 时的中间特征图会被放大到输入尺寸的数倍16GB 显存的卡上很容易在这里炸掉。实际跑批处理时建议把采样和 VAE 解码拆开。先生成所有 latent 并存在硬盘上再统一解码这样即使中途出错也不需要重新采样。这一点在做大规模评估时尤其重要。4. 避坑记录复现 DiffiT 时最容易翻车的五个位置4.1 预训练权重加载进来后生成纯噪声损失函数不收敛现象加载官方预训练权重后采样结果不是清晰图像而是彩色噪点。训练时 loss 曲线停在很高的位置下不去。原因最常见的是 time embedding 的维度或归一化方式与权重不匹配。DiffiT 使用 Swish 激活的 MLP 编码时间和常规的 sinusoidal embedding 在维度拼接方式上有差异。如果你的模型定义里 time embedding 的维度差 2 倍权重加载时会静默失败因为 PyTorch 的load_state_dict对 shape 不匹配会直接报错但如果你用了strictFalse部分层加载失败是不会有提示的。解决加载后校验每一层的参数范数是否合理。更简单的方式是打印模型参数量和官方报告对比差 10% 以上就说明定义有问题。4.2 训练到一半 loss 变成 NaN重启后又能跑现象训练 10 万步之后loss 值突然跳到 NaN检查学习率和梯度都没发现异常。重新加载 checkpoint 又能继续训练但过不了多久又复现。原因混合精度训练时注意力分数经过 softmax 后乘以 V如果中间结果过大FP16 的表示范围会溢出。DiffiT 的时间自适应注意力在后期阶段会产生较大的 attention logits在 FP16 下更容易触发这个问题。解决在 TMSA 的输出位置加一个 FP32 的落盘点或者直接对 attention 分数做 scaled 处理。我一般会在 softmax 之前手动除以sqrt(d)并关掉该层的自动混合精度使用torch.cuda.amp.autocast(enabledFalse)。如果你用的是 bfloat16这个问题会好很多。4.3 FID 分数虚高但生成图像肉眼看起来不错现象生成的图视觉质量还行但 FID 一直比论文高不少甚至翻倍。原因FID 计算时 Inception 特征提取的版本不一致。PyTorch 官方权重和 TF 的 Inception 在统计分布上差异很大直接用不同的特征统计文件算出来的分数不可比。市面上很多复现报告的 FID 差异其实都来自这里而不是模型实现有 bug。解决统一使用 参考统计数据并且用同版本的pytorch-fid库。评估时把生成图片保存后统一用脚本计算不要在训练脚本里内置 FID 计算逻辑。4.4 显存波动巨大甚至出现随机溢出现象同一配置下前几次训练正常某次 forward 突然 OOM。原因DiffiT 的时间自适应注意力使得 attention 图在不同时间步的稀疏度不同但显存分配是固定的。真正的原因是 PyTorch 的缓存分配策略加上 CUDA graph 在某些算子上的兼容性问题。解决设置PYTORCH_CUDA_ALLOC_CONFmax_split_size_mb:256把显存分配粒度调大减少碎片化。同时把 batch size 降到显存余量的 80% 以下给波动留缓冲。如果有条件使用torch._C._set_cuda_stack_pool_allocator_enabled(True)优化小内存块的分配。4.5 生成图像发灰或对比度偏低现象图像细节对但整体色彩像蒙了一层雾对比度不够。原因v-prediction 和 epsilon-prediction 混用导致 latent 尺度不对。DiffiT 的采样代码默认使用 v-prediction 变体如果你的权重是在 epsilon-prediction 下训练的直接套采样脚本会出错。解决确认模型训练时用的目标函数对应选择采样方式。如果你不确定可以从 latent 的均值方差判断。正确尺度下latent 均值接近 0方差约为 1。如果方差明显小于 1采样时在最后一步做一个重新缩放能缓解问题。这五个坑基本覆盖了从权重加载到训练稳定性再到评估指标的全链路。建议并行跑实验时把权重校验写成一个独立的脚本每次换配置先跑一遍能省下大量排查时间。5. 训练自己的 DiffiT数据集适配、超参与分布式策略5.1 如何把你的数据集接进训练流程DiffiT 的官方训练代码默认支持 CIFAR-10、FFHQ 和 ImageNet 三种数据格式。换到自己数据集时需要做以下三件事写一个返回[image, label]的 Dataset 类、把图像统一缩放到目标分辨率、确认 latent 空间模型是否需要额外的 VAE 预处理。import torch from torch.utils.data import Dataset from PIL import Image import os class CustomDataset(Dataset): 适配 DiffiT 训练脚本的最小 Dataset 实现。 data_root 下的每个子目录代表一个类别。 def __init__(self, data_root, img_size256, is_trainTrue): self.img_size img_size self.classes sorted(os.listdir(data_root)) self.class_to_idx {c: i for i, c in enumerate(self.classes)} self.samples [] # 遍历扫描所有图片路径避免一次性全部读入内存 for cls in self.classes: cls_dir os.path.join(data_root, cls) for fname in os.listdir(cls_dir): if fname.lower().endswith((.png, .jpg, .jpeg)): self.samples.append((os.path.join(cls_dir, fname), self.class_to_idx[cls])) def __len__(self): return len(self.samples) def __getitem__(self, idx): path, label self.samples[idx] img Image.open(path).convert(RGB) img img.resize((self.img_size, self.img_size), Image.LANCZOS) arr torch.from_numpy(np.array(img)).permute(2, 0, 1).float() / 127.5 - 1.0 return arr, label这里的关键是resize方式。LANCZOS 重采样对图像细节保留最好但如果你的数据集本身是大图直接 resize 到 256 会丢掉大量细节。更规范的做法是中心裁剪后 resize但裁剪会损失边缘信息。两者权衡下我一般对结构类数据比如建筑、车辆用中心裁剪对纹理类数据用整体缩放。5.2 超参数怎么调学习率、Batch Size、EMA 与 Adim 优化器DiffiT 官方使用的是 Adim 优化器这是一个专为扩散模型设计的优化器它对时间步的分组方式做了优化让不同时间步的梯度更新速度更均衡。如果你用手头的数据微调把优化器换成 AdamW 并调整学习率也能工作但训练曲线会明显不如 Adim 平滑。常用配置参考表参数推荐值说明学习率1e-4 到 5e-4微调用 1e-4从头训练用 2e-4 起步Batch Size64 到 256显存不够就缩减分辨率而不是缩减 batch采样步数50 到 100训练时短步数评估时长步数EMA Decay0.995 到 0.999越大越稳定但收敛变慢Adim beta(0.9, 0.999)保持默认即可微调时有一个经验先冻结时间投影层wq_t、wk_t、wv_t只训练空间注意力权重。因为预训练模型的时间注意力已经学到了通用的阶段转换规律你的数据集规模不够时强行更新这部分容易把已有的好分布搞坏。等空间部分稳定后再解开冻结。5.3 多卡训练DDP 的配置与容易忽略的问题多卡训练时使用 PyTorch DDPDiffiT 的代码结构天然支持。但有个坑是它对 TMSA 中相对位置偏置的同步处理。位置偏置是 buffer 而不是参数DDP 默认不会同步 buffer导致每个进程的偏置值不一致。训练到后期会出现不同卡上 loss 明显偏离的现象。# 启动命令示例2 节点 × 4 卡 torchrun --nnodes2 --nproc_per_node4 \ --master_addr192.168.1.10 --master_port29500 \ main_imagenet.py \ --batch_size 32 \ --lr 2e-4 \ --ema_decay 0.995 \ --img_size 256在代码里需要在 DDP 初始化后显式 broadcast bufferdef broadcast_position_bias(model): DDP 初始化后同步相对位置偏置 buffer防止各进程副本不一致。 for name, buf in model.named_buffers(): if pos_bias in name: torch.distributed.broadcast(buf, src0)这个函数必须在模型放到 GPU 之后、训练循环开始之前调用。如果你用的是DistributedDataParallel构造器它默认的broadcast_buffersTrue只能保证初始化时同步一次训练中如果 buffer 被原地更新仍然会分叉。所以要么把偏置注册为参数要么在每步训练时手动同步。5.4 训练质量的下探如何判断模型有没有真正学会训练到中途loss 已经低到和论文相当但生成结果依然不理想。这时候要区分是欠拟合还是模块之间存在冲突。我习惯在每个 checkpoint 处保存一组固定时间步在固定噪声下的生成结果对比不同 epoch 的同一组输入能直观看出生成轨迹是否在稳定变化。如果某段时间生成结果几乎没有变化说明模型可能掉进了注意力退化的局部最优。此时把学习率临时调低一个数量级训练 2000 步再恢复往往能打破僵局。这种学习率重启的做法在扩散模型上比在 GAN 上效果还明显。6. 生成质量的验证技巧FID 只是起点还有两个手动检查必须做跑通之后验证生成质量是最后一个环节。很多人只看 FID但 FID 对两种常见问题并不敏感一是整体分布正确但单张图像明显不协调二是多样性收敛导致所有生成图都像同一个模板换颜色。我每次都会手动做两个检查。第一个检查是把同一类别的生成图按 FID 特征空间的距离排序看最远和最近的样本差异。如果差异极小说明模型产生了模式坍塌即使在 FID 分数不错的情况下。第二个检查是检查类别条件是否真的被使用。方法很简单生成同一类别的 100 张图统计它们的 MSE 差异再用不同类别生成 100 张统计类间差异。如果类内差异和类间差异差不多说明 classifier guidance 没有生效条件信息在生成时被丢掉了。具体实现上可以用pytorch-fid提取生成图片的 Inception 特征然后做聚类分析import numpy as np from pytorch_fid.inception import InceptionV3 from pytorch_fid.fid_score import calculate_frechet_distance # 提取特征的 Inception 模型输出维度设为 2048 block_idx InceptionV3.BLOCK_INDEX_BY_DIM[2048] inception InceptionV3([block_idx]).cuda().eval() def extract_features(img_loader): features [] for batch in img_loader: with torch.no_grad(): act inception(batch.cuda())[0] features.append(act.cpu().numpy().reshape(act.size(0), -1)) return np.concatenate(features, axis0) # 类内距离同类别生成图之间的平均余弦距离 # 类间距离不同类别之间的平均余弦距离如果发现条件信息丢失优先检查两个位置。一个是 time token 的注入方向DiffiT 把类别 embedding 和时间 embedding 是分开处理的如果你把类别信息拼进了时间 token注意力层会被时间变化带偏。另一个是 classifier guidance 的梯度缩放cfg_scale太高时梯度会把信息推离原始条件方向生成结果看起来过激但实际不遵循类别。这两个检查完成后再配合 FID、sFID、IS 三个指标做最终确认基本可以放心把结果写入报告或文档。从那以后我每次训练完生成模型都强制走一遍特征距离分析和条件检查这已经成了一个习惯动作。它帮我挡下了至少三次分数好看但实际不可用的翻车。希望帮到你。本文还有配套的精品资源点击获取