简介《DiffiT: Diffusion Vision Transformers for Image Generation》是NVIDIA研究团队发表的学术论文PDF专门面向关注扩散模型、视觉Transformer与图像生成的研究者和开发者。论文提出DiffiT模型将视觉Transformer引入扩散模型的去噪网络并设计了时间依赖的自注意力模块使其能在不同去噪阶段自适应调整行为同时给出了用于高分辨率生成的latent DiffiT方案在ImageNet-256上取得FID 1.73的最优结果。资源为单文件PDF压缩包大小约43.22MB内容完整包含摘要、引言、方法描述、实验分析及结论适合快速获取论文核心思路、模型架构与基准数据。目前已有336人学习下载可作为该方向入门或对比实验的重要参考。1. DiffiT 是什么扩散模型换掉 U-Net 主干之后图像生成为什么反而更稳把扩散模型的去噪网络从 U-Net 换成纯 Vision TransformerViT这个想法听起来有点反直觉——U-Net 的多尺度跳跃连接明明是行业默认配置为什么还要换成全局注意力结构DiffiTDiffusion Vision Transformers做的就是这件事用带时间感知的 Transformer 作为扩散模型主干在 ImageNet 和同类高分辨率生成任务上把 FID 压了下去。它解决的问题很具体U-Net 的感受野受限于卷积核尺寸对长距离空间依赖的表达要靠加深网络来补而 Transformer 天然有全局建模能力。适合谁正在做扩散模型训练、觉得 U-Net 越调越复杂想尝试主干替换路线或者想搞懂 ViT 在生成任务里怎么处理时间步信息的人。这篇笔记不讲论文复现神话只讲我照着 DiffiT 思路改主干时踩过的坑和能直接落地的参数。2. 原理ViT 主干 time-aware 注意力为什么能解决 U-Net 的瓶颈2.1 U-Net 在扩散任务里的两个硬伤扩散模型的训练目标是让网络在不同噪声强度下都能预测出干净的噪声或原图这意味着同一个网络要同时处理「几乎全是噪声」和「几乎全是原图」两种极端输入。U-Net 靠下采样扩大感受野但下采样会丢失高频细节所以它必须额外设计跳连把浅层信息传到深层。问题出在两点第一U-Net 的感受野本质上还是局部的。卷积核是 3x3、5x5堆到几十层之后理论感受野很大但实际有效感受野远小于理论值网络对「画面左上角和右下角是否属于同一个物体」这类长距离关系建模不足。生成大尺寸图像时这种缺陷会表现为结构不连贯——比如人像的左右手风格不一致。第二U-Net 里时间步信息的注入方式主要是加在残差块旁边的 scale-shift 操作也就是把时间嵌入变成归一化层的 gamma 和 beta。这个机制有效但本质上是对特征做逐通道仿射变换信息容量有限。当类别数很多或者图像分辨率很高时单纯靠 scale-shift 很难把「第 500 步该做什么」和「第 5 步该做什么」区分得足够干净。DiffiT 的思路是让时间步直接参与注意力权重的计算而不是只调通道。2.2 time-aware attention把时间表征塞进注意力权重而不是通道里标准 ViT 的自注意力对每一层、每一步都用同一套注意力逻辑。但扩散过程里有个特点早期步噪声大需要关注全局结构晚期步噪声小只需要修复局部纹理。如果所有时间步用同样的注意力模式网络要么偏向全局细节要么偏向局部准确很难两头兼顾。time-aware attention 的做法是先算出时间步的任务起因再把任务起因与图像 patch token 的空间位置联合映射到注意力图上迫使每个位置决定「这一步我该看多远」。具体来说在自注意力的 QK 计算之外增加了一个从时间嵌入映射出来的偏置项这个偏置项按 patch 位置给出一个匹配关系。于是不同时间步自然学到了不同的注意力距离——高噪声步看全图低噪声步看邻域。这个设计与 U-Net 中的 scale-shift 注入本质区别在于scale-shift 只改变特征的幅值而 time-aware 偏置直接改变信息的流动路径。实现上并不复杂核心就是在原始注意力 logits 上加一个由时间步生成的位置偏置这部分代码下节会给出。2.3 单尺度与多尺度 token 结构分辨率翻倍时的选择DiffiT 里有两种配置一种是纯单尺度 token 流——整张图切成固定大小的 patch全程在同一个分辨率下做全局注意力另一种是多尺度潜空间 token 流——模仿 U-Net 的多分辨率习惯把 token 序列做几次下采样在低分辨率分支做长距离建模再恢复回原分辨率。后者兼顾了 Transformer 的全局建模和 U-Net 式的多尺度层次适合 256 甚至 512 分辨率生成。我自己的经验是单尺度结构在 64x64、128x128 这种中等分辨率下够用训练省显存收敛也快一旦目标分辨率到 256 以上单尺度 token 序列很长自注意力复杂度是 O(N^2)训练速度会明显下降且梯度不稳定。多尺度结构把 N 缩小了几倍再算注意力实际训练速度能快 30% 以上。这不是说单尺度不能跑 256只是时间和显存成本划不来。分辨率和 batch size 之间的权衡通常的做法是先定显存上限再定多尺度下采样次数最后才调注意力头数。3. 实现从注意力模块到训练循环的可复制代码3.1 先写 time-aware 注意力模块核心就三行改动我一般是先在一个最小配置上把注意力模块跑通再套进整个模型。这里给出一个可直接运行的 time-aware 注意力实现省去位置编码等外围代码只看核心逻辑import torch import torch.nn as nn import torch.nn.functional as F class TimeAwareAttention(nn.Module): def __init__(self, dim, num_heads, time_dim, max_patches1024): super().__init__() self.num_heads num_heads self.head_dim dim // num_heads self.scale self.head_dim ** -0.5 # 关键缩放因子不写会梯度爆炸 self.qkv nn.Linear(dim, dim * 3, biasFalse) # 把时间嵌入映射成每个位置对应的偏置bias 维度是 [B, heads, 1, N] self.time_proj nn.Sequential( nn.Linear(time_dim, dim), nn.SiLU(), nn.Linear(dim, num_heads * max_patches) ) # 头数太多或 patch 数太多时可以拆成逐头逐位置计算避免一次生成大 tensor def forward(self, x, t_emb): B, N, C x.shape qkv self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim) q, k, v qkv.permute(2, 0, 3, 1, 4).unbind(0) attn_logits torch.matmul(q, k.transpose(-2, -1)) * self.scale # 时间偏置对每个 patch 位置给一个与时间步相关的偏移量 time_bias self.time_proj(t_emb).reshape(B, self.num_heads, 1, N) attn_logits attn_logits time_bias # 广播到所有 query 位置 attn F.softmax(attn_logits, dim-1, dtypetorch.float32) out torch.matmul(attn, v) out out.transpose(1, 2).reshape(B, N, C) return out参数说明time_proj的输出维度是num_heads * max_patches意思是每个头都给 N 个 patch 位置生成一个独立的标量偏置。这样不同头可以关注不同距离——某些头学全局某些头学局部。max_patches需要设为输入序列长度的上限超过会报错实际运行时如果你的 patch 数小于这个值把多余偏置截掉即可。实现里有两个细节值得注意。第一个是attn_logits必须乘self.scale不乘的话当 head_dim 较大时 softmax 输入过大梯度很容易在反向传播时变成 NaN这是最常见的翻车点。第二个是softmax里显式指定dtypetorch.float32因为混合精度训练时注意力权重如果用 fp16 计算小概率位置的梯度会直接下溢成 0生成结果会出现随机黑点。3.2 组装 DiffiT 块时间嵌入、条件归一化和残差流有了注意力模块下一步是把前馈网络、归一化层和时间嵌入组装成完整的 Transformer 块。这里我参考了常见做法的结构——先做 AdaLN 条件归一化、再走注意力、再走 MLP、最后残差连接class DiffiTBlock(nn.Module): def __init__(self, dim, num_heads, time_dim, class_dim, max_patches1024): super().__init__() self.norm1 nn.LayerNorm(dim, elementwise_affineFalse) self.norm2 nn.LayerNorm(dim, elementwise_affineFalse) self.attn TimeAwareAttention(dim, num_heads, time_dim, max_patches) self.ff nn.Sequential( nn.Linear(dim, dim * 4), nn.GELU(), nn.Linear(dim * 4, dim), ) # AdaGN时间和类别共同生成归一化的 scale 和 shift self.ada_gn nn.Sequential( nn.SiLU(), nn.Linear(time_dim class_dim, dim * 2) ) def forward(self, x, t_emb, class_emb): # 时间与类别拼接成一个条件向量 cond torch.cat([t_emb, class_emb], dim-1) scale, shift self.ada_gn(cond).chunk(2, dim-1) h self.norm1(x) * (1 scale.unsqueeze(1)) shift.unsqueeze(1) h self.attn(h, t_emb) x x h h self.norm2(x) * (1 scale.unsqueeze(1)) shift.unsqueeze(1) x x self.ff(h) return x逻辑说明AdaGN把时间和类别向量融合后生成scale和shift施加在归一化后的特征上。这样的设计让类别信息通过仿射变换注入而时间信息同时走两条路——一条是注意力偏置一条是 scale-shift。两条路各司其职注意力偏置负责调感受野和关注模式归一化参数负责调特征幅值。1 scale的写法是参考常见条件归一化的做法直接乘原始 scale 会导致训练初期输出幅值剧烈波动。这里的class_emb就是类别的可学习嵌入一般用nn.Embedding(num_classes, class_dim)生成。如果没有类别标签可以用一个固定的可学习向量代替或者直接拼一个全零向量。需要注意time_dim、class_dim和dim三者没必须相等实际项目中常见的是time_dim256、class_dim512、dim512这样的配置。3.3 训练循环EMA、噪声采样和损失函数的选择主干结构搭好之后训练循环反而是最容易出细节问题的部分。扩散模型的训练看起来就几行但损失函数选择、EMA 策略、随机时间步采样这三个点直接决定生成质量# 简化训练循环省略数据加载和日志部分 def train_step(model, x, y, optimizer, ema_model): x x.to(device) y y.to(device) t torch.randint(0, num_train_steps, (x.size(0),), devicedevice).long() noise torch.randn_like(x) # 前向加噪x_t sqrt(alpha_bar) * x0 sqrt(1 - alpha_bar) * noise x_t sqrt_alpha_bar[t].view(-1, 1, 1, 1) * x sqrt_one_minus_alpha_bar[t].view(-1, 1, 1, 1) * noise pred model(x_t, t, y) # 用 L1 loss 而不是 L2扩散模型收敛更稳 loss F.l1_loss(pred, noise, reductionmean) optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() ema_model.update(model) return loss.item()参数说明num_train_steps是扩散总步数常见范围 1000 到 4000。时间步t是均匀随机采样而不是顺序采样原因是要让网络学会处理任意噪声强度否则某些步数的噪声永远见不到。sqrt_alpha_bar和sqrt_one_minus_alpha_bar是预设的噪声调度表可以用余弦调度或线性调度生成差异在中间步数的噪声密度分布——我一般用余弦调度它在训练后期的高噪声步数上效果更稳定。损失函数选L1而不是L2是反复试出来的经验。L2 对离群点更敏感生成图像容易边缘模糊L1 对每个位置误差的惩罚更均匀图像更锐利。如果你用的是预测v参数而不是预测噪声的训练目标损失函数保持不变只是目标值从 noise 换成 v。梯度裁剪设在 1.0 能有效避免时间步接近 0 时 loss 突然变大导致的训练崩溃。EMA 的更新策略也值得单独说EMA 模型不是每步都重新复制参数而是做滑动平均。常见的做法是ema_model.update(model)内部维护一份影子权重每次按衰减率更新。衰减率设到 0.9999 到 0.99999 之间太小则 EMA 模型与训练模型差别不大平滑作用不明显太大则指标变化很慢。采样时必须用 EMA 模型的权重而不是训练权重这个很多人会忘。4. 参数选择训练 DiffiT 时决定成败的关键设置4.1 优化器与学习率AdamW 之外的另一个选择DiffiT 这类 ViT 主干的模型对优化器比较挑剔。ViT 的注意力层权重初始化方差较大如果学习率太高第一轮训练 loss 会直接飞到 NaN。我这边常用的配置是参数项推荐范围备注优化器AdamWbetas(0.9, 0.999) 不要动学习率1e-4 到 3e-4线性 warmup 5000 步Weight decay0.03 到 0.1作用于所有参数包括 biasBatch size256 以上小于 128 时生成质量明显下降训练步数50 万到 100 万跟数据集大小有关EMA 衰减0.9999采样用 EMA 权重学习率是最敏感的超参数。我做过几次对照实验学习率从 1e-4 提到 5e-4训练 loss 表面下降更快但生成样本的 FID 反而差了。原因是 ViT 的注意力权重在训练后期需要精细调整学习率太大会破坏已经学到的位置编码结构。线性 warmup 是标配5000 步不算多因为扩散模型训练步数动辄几十万步前 1% 的步数用于稳定优化器状态很划算。Weight decay 0.03 是我在 ImageNet 这类大数据集上的经验值如果换成小数据集weight decay 要降到 0.01 甚至 0否则模型欠拟合。这里一个容易忽略的点AdamW 的 weight decay 和 L2 正则不是一回事weight decay 是直接对参数做衰减不参与梯度计算对小批量训练更友好。4.2 采样器与步数DDIM、DPM-Solver 和 CFG 的配合训练好模型之后采样参数直接决定你看到的生成效果。DiffiT 论文和社区实践里常用的是 DPM-Solver它比 DDIM 在相同步数下更准确。我实际用下来两边的主要差异在步数少的时候更明显采样器推荐步数适用场景注意事项DDIM100 到 250需要确定性的采样轨迹步数太少会出现条纹伪影DPM-Solver20 到 50追求速度且画质要求高对噪声调度表较敏感二阶 DPM-Solver10 到 20快速出图验证需要配合 CFG 共同调参CFG无分类器引导的尺度也是关键参数。DiffiT 配合 CFG 的典型范围是 2.0 到 4.0小于 2.0 生成的图像多样性好但忠实度差大于 4.0 会出现颜色过饱和和边缘伪影。CFG 的计算方式是pred pred_uncond scale * (pred_cond - pred_uncond)注意这里要求模型在训练时以一定概率丢掉类别条件常见做法是 10% 概率把类别嵌入换成空向量。如果你训练的模型没有做这个 dropout那 CFG 完全无效这是我在某次实验中踩过的坑。采样器步数和 CFG 是一对需要联合调参的量。步数太少时CFG 尺度要相应降低否则每一步的引导误差会被放大。我常用的快速验证组合是20 步 DPM-Solver CFG 2.5看到整体效果之后再尝试 50 步 CFG 3.0 的组合。每步训练时间有限不建议一开始就上 250 步采样纯浪费时间。4.3 分辨率、patch size 与显存的三角关系ViT 主干的显存占用和 patch size 直接相关。Patch size 决定 token 序列长度 N而注意力复杂度是 O(N^2)。同样一张 256x256 的图patch size 16 时 token 数 256patch size 8 时 token 数 1024显存差距将近 4 倍。实际操作时要在这三者间做取舍生成要求精细纹理比如人脸细节时patch size 越小越好但要接受显存和训练时间翻倍生成大场景风景、建筑时patch size 16 就够注意力偏重全局结构局部纹理靠多尺度分支补多尺度结构可以缓解这个问题——低分辨率分支用大 patch 保证全局一致性高分辨率分支用相对小的 patch 修复局部我个人的经验配比256x256 图像用 patch size 16 做主分支下采样一次到 128x128 用 patch size 8 做多尺度分支显存开销比单尺度 patch 8 少一半生成质量接近。如果你的显存只有 24G256x256 batch size 64 是能跑的极限再大就得开梯度累积或者混合精度。梯度累积的 step 数等于目标 batch size 除以实际 batch size需要注意与学习率的配合——增大有效 batch size 时学习率要按比例微调否则收敛变慢。5. 避坑排查DiffiT 训练与采样中最常见的 4 个翻车场景5.1 现象训练 loss 下降正常但生成图像全是模糊色块这是我第一次训练 ViT 主干扩散模型时遇到的最诡异问题训练 loss 从 0.3 降到 0.05看起来收敛良好但生成的图像像打码一样。排查了几天最后发现问题出在位置编码上。原因我的实现把位置编码和 patch embedding 相加后直接进 Transformer 块但如果位置编码的数值范围远大于 patch embedding注意力在初始阶段完全被位置信息主导网络学到了「位置相近的 patch 特征应该相似」的归纳偏置所以生成结果模糊。解决把位置编码初始化范围缩小到 patch embedding 标准差的 0.1 倍或者改用可学习位置编码并在训练初期加一个很小的权重衰减。调整之后生成的图像立刻锐利起来。检查方法很简单打印第一层 attention 权重的分布——如果大部分注意力都集中在对角线附近说明位置编码主导了信息流。5.2 现象EMA 模型采样效果好但训练模型采样效果完全不能看这是扩散模型训练里的经典现象不算 bug但很容易让人误判训练出了问题。EMA 的作用是把训练过程中参数的快速波动平滑掉得到的是一套更稳定的权重。训练模型每步都在大幅更新采样时直接用它生成结果会有明显的噪声残留。原因扩散模型的 loss 曲面非常崎岖单步最优参数在采样时并不稳定。EMA 衰减率 0.9999 意味着权重更新有 99.99% 的保留率累积几千步后 EMA 权重趋于平滑收敛而训练权重还在震荡。解决采样时必须加载 EMA 权重。如果你的训练框架没有自动保存 EMA 权重可以每 5000 步手动把 EMA 模型的状态字典存一次采样时加载这个文件。另外EMA 衰减率不要设到 0.99999 以上否则 EMA 权重更新过于缓慢初期采样效果反而比训练模型差——我遇到过这种情况调低衰减率到 0.9999 之后立刻正常。5.3 现象DPM-Solver 采样时出现周期性条纹伪影同样的模型权重DDIM 采样出来的图像正常换 DPM-Solver 就出现条纹尤其在背景平滑区域特别明显像打印机的横纹。原因DPM-Solver 对噪声调度表的一致性要求比 DDIM 高。我用的噪声调度表是线性调度前向加噪在低噪声区间分布过于密集DPM-Solver 的数值求解在高曲率区域误差放大出现周期性的过冲。解决换成余弦调度表或者用 DPM-Solver 自带的噪声调度校正逻辑。还有一个简单的方法把采样步数从 20 提到 40条纹基本消失但速度减半。我最后的方案是余弦调度 40 步 DPM-Solver生成质量和速度都能接受。如果你已经训练完模型不想重新训也有个取巧办法——在采样前把初始噪声做一次小幅度高斯平滑能压掉一部分高频条纹但这是治标不治本的做法。5.4 现象增加 CFG 尺度后图像过饱和、出现彩色噪点CFG 从 2.0 调到 4.0图像对比度变强是正常的但如果颜色完全过饱和边缘出现类似故障艺术的彩色噪点说明 CFG 实际上加过头了。原因CFG 的本质是把条件预测和无条件预测的差值放大。当模型在某个像素位置对条件信息不确信时差值会产生很大的随机分量放大后就成了噪点。ViT 主干的全局注意力更容易在某些高频位置产生这种不确定所以这个问题在 DiffiT 上比 U-Net 更明显。解决有两个方向。一是降低 CFG 到 3.0 以下同时增加采样步数用计算换质量二是在 CFG 计算后对差值做动态截断比如限制差值超过某个阈值时按比例缩小。我常用的是后者阈值设 2.0 左右既保留 CFG 增强细节的效果又不至于过饱和。另外训练时的类别 dropout 概率也可以从 10% 调到 15%让模型更适应无条件推理CFG 的有效区间会变大。6. 进阶技巧用 FID 验证 DiffiT 生成质量时的 3 个自我检查习惯训练结束后怎么判断模型真的能用只看几组肉眼样本不够FID 指标才是客观标准。但 FID 本身也有坑我养成了三个固定检查习惯。第一采样 5 万张以上再算 FID少于这个数 FID 方差太大。之前有次只采了 2000 张FID 显示很好换 5 万张直接涨了三倍。第二FID 计算器要求输入图像 resize 到 299x299不同采样方式的 resize 插值算法会影响 FID 结果——统一用双线性插值我见过 LANCZOS 和双线性混用导致的 FID 虚高。第三算 FID 时用同一个噪声种子也就是用固定的随机种子生成 5 万张图这样两次调整参数之间的对比才有意义否则 FID 的波动会把你的调参结论完全淹没。这三点做完还有一个更细的习惯把 FID 拆成「训练集 FID 和验证集 FID」对比。如果训练集 FID 低但验证集 FID 明显高说明模型记住训练集了这时要加大 dropout 或增强数据增强。对 DiffiT 来说ViT 的全局注意力建模能力强过拟合速度比 U-Net 快不少这类检查尤其必要。我现在每轮训练结束都会跑一遍这套流程花的时间不多但能避免「感觉模型不错实际一测不行」的尴尬。最后说一个经验之谈DiffiT 这类纯 Transformer 主干和 U-Net 调的参数完全不是一套逻辑不要拿 U-Net 的成功经验直接套。U-Net 上设为 0.1 的注意力 dropout在 ViT 上大概率会抑制生成多样性U-Net 上有效的 5e-4 学习率在 ViT 上很可能训崩。我自己第一次做主干替换时交了不少学费希望这篇笔记里的参数和避坑记录能帮你少走这些弯路希望帮到你。本文还有配套的精品资源点击获取