简介：DiffiT是NVIDIA团队提出的一种将视觉变换器引入扩散模型的图像生成新架构，面向生成式AI研究者、计算机视觉方向学生及扩散模型开发者，用于解决传统去噪网络过度依赖卷积残差U-Net、注意力机制设计缺乏统一范式的问题。资源包内含1个PDF文件，压缩包约43.22MB，即论文原文，涵盖摘要、引言、模型概述、时间依赖自注意力机制、latent DiffiT高分辨率生成方案、实验结果与相关知识点解析等完整章节。读者可从中了解DiffiT如何在U形编码器解码器中嵌入时间依赖自注意力模块，使注意力层随去噪阶段自适应调整行为，并掌握latent DiffiT在ImageNet-256上取得1.73 FID分数的技术路径。论文还系统梳理了扩散模型、视觉变换器、去噪网络、FID指标与类条件/无条件合成任务等基础概念，适合作为生成模型架构选型与实验复现的参考材料。目前已有336人学习。
1. 当扩散模型遇上视觉 Transformer：DiffiT 到底在解决什么问题
如果你最近在折腾图像生成，大概率已经被各种扩散模型刷屏了。但真正动手训过的人都知道，U-Net 那套骨干网络在低分辨率和高分辨率之间反复上下采样，显存吃得凶，长程依赖还经常建模不到位。DiffiT 这个方向做的事情，说白了就是把扩散模型的去噪骨干从 U-Net 换成纯视觉 Transformer，让注意力机制直接在潜空间里处理全局关系。它解决的不是“能不能生成”的问题，而是“生成质量、训练稳定性和结构一致性能不能同时兼顾”的问题。适合谁看？适合已经跑通过 DDPM 或 DDIM 基础流程、想换骨干网络做对比实验的从业者，也适合需要在高分辨率场景下控制显存和采样步数的工程同学。这一章先把概念立住，后面几章直接上代码和参数。
2. DiffiT 的骨干设计：为什么用 Transformer 替代 U-Net 做去噪
2.1 从 U-Net 的归纳偏置说起
传统扩散模型用 U-Net 做去噪网络，核心原因是卷积的局部归纳偏置和跳跃连接能很好地保留空间细节。但 U-Net 有个硬伤：感受野增长依赖下采样层数，层数一多，显存和计算量就上去了。而且卷积核在建模全局结构时，比如生成一张人脸，左右眼的对称关系需要堆很多层才能捕捉到。DiffiT 的思路是直接用 ViT 的 self-attention 来建模全局依赖，每个 patch 都能和所有其他 patch 交互，理论上一步到位。
但直接套 ViT 会翻车。扩散模型的时间步信息需要注入网络，U-Net 通常用 FiLM 或 AdaGN 把时间嵌入加到每个残差块里。ViT 没有这种天然结构，所以 DiffiT 设计了一个时间依赖的注意力机制，让注意力权重本身随时间步变化。常见做法是在 QKV 投影前把时间嵌入拼接到 token 上，或者用自适应层归一化（AdaLN）调制注意力输出。我一般会选 AdaLN，因为实现简单，训练也稳。
2.2 最小可跑的 DiffiT 去噪网络实现
下面这段代码是一个简化版的 DiffiT 骨干，保留了时间条件注入和 Transformer 块的核心逻辑。依赖 PyTorch 2.0 以上，torchvision 可选。
import torch import torch.nn as nn import math class TimeEmbedding(nn.Module): """把标量时间步映射成向量，和 Transformer 的 hidden dim 对齐""" def __init__(self, dim): super().__init__() self.dim = dim self.proj = nn.Sequential( nn.Linear(dim, dim * 4), nn.SiLU(), nn.Linear(dim * 4, dim) ) def forward(self, t): # t: (B,) 的整数时间步 half = self.dim // 2 freqs = torch.exp(-math.log(10000) * torch.arange(half, device=t.device) / half) args = t[:, None].float() * freqs[None] emb = torch.cat([torch.sin(args), torch.cos(args)], dim=-1) return self.proj(emb) class DiffiTBlock(nn.Module): """时间条件化的 Transformer 块，用 AdaLN 注入时间信息""" def __init__(self, dim, num_heads, mlp_ratio=4.0): super().__init__() self.norm1 = nn.LayerNorm(dim, elementwise_affine=False) self.attn = nn.MultiheadAttention(dim, num_heads, batch_first=True) self.norm2 = nn.LayerNorm(dim, elementwise_affine=False) self.mlp = nn.Sequential( nn.Linear(dim, int(dim * mlp_ratio)), nn.GELU(), nn.Linear(int(dim * mlp_ratio), dim) ) # AdaLN 的调制参数，由时间嵌入生成 self.adaLN_modulation = nn.Sequential( nn.SiLU(), nn.Linear(dim, 6 * dim) ) def forward(self, x, t_emb): # x: (B, N, D) 的 patch token shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = \ self.adaLN_modulation(t_emb).chunk(6, dim=-1) # 注意力分支 h = self.norm1(x) * (1 + scale_msa.unsqueeze(1)) + shift_msa.unsqueeze(1) h, _ = self.attn(h, h, h) x = x + gate_msa.unsqueeze(1) * h # MLP 分支 h = self.norm2(x) * (1 + scale_mlp.unsqueeze(1)) + shift_mlp.unsqueeze(1) h = self.mlp(h) x = x + gate_mlp.unsqueeze(1) * h return x class DiffiT(nn.Module): """简化版 DiffiT：patch embed + 多个时间条件 Transformer 块 + 输出投影""" def __init__(self, img_size=32, patch_size=4, in_chans=3, dim=384, depth=12, num_heads=6): super().__init__() self.patch_embed = nn.Conv2d(in_chans, dim, kernel_size=patch_size, stride=patch_size) num_patches = (img_size // patch_size) ** 2 self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, dim)) self.time_embed = TimeEmbedding(dim) self.blocks = nn.ModuleList([ DiffiTBlock(dim, num_heads) for _ in range(depth) ]) self.final_norm = nn.LayerNorm(dim, elementwise_affine=False) self.final_adaLN = nn.Sequential(nn.SiLU(), nn.Linear(dim, 2 * dim)) self.out_proj = nn.Linear(dim, patch_size * patch_size * in_chans) self.patch_size = patch_size self.in_chans = in_chans def forward(self, x, t): # x: (B, C, H, W) B = x.shape[0] x = self.patch_embed(x) # (B, D, H/P, W/P) x = x.flatten(2).transpose(1, 2) # (B, N, D) x = x + self.pos_embed t_emb = self.time_embed(t) for block in self.blocks: x = block(x, t_emb) shift, scale = self.final_adaLN(t_emb).chunk(2, dim=-1) x = self.final_norm(x) * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) x = self.out_proj(x) # (B, N, P*P*C) # 还原成图像形状 h = w = int(math.sqrt(x.shape[1])) x = x.transpose(1, 2).reshape(B, self.in_chans, h * self.patch_size, w * self.patch_size) return x逻辑说明：TimeEmbedding用正弦位置编码把时间步变成向量，再经过两层 MLP 对齐到 hidden dim。DiffiTBlock里的adaLN_modulation从时间嵌入生成六个调制参数，分别控制注意力和 MLP 分支的缩放与偏移。DiffiT主体先做 patch embed，加上可学习的位置编码，然后堆叠多个块，最后用 AdaLN 做输出前的归一化，再投影回像素空间。
参数说明：img_size和patch_size决定 token 数量，32x32 图像配 patch 4 得到 64 个 token，显存友好。dim是隐藏维度，384 在单卡 12GB 上跑 batch size 32 没问题。depth控制 Transformer 块数，12 层是常见起点。num_heads一般设成 dim 的 1/64，384 对应 6 个头。如果显存吃紧，优先降depth而不是dim，因为注意力计算量随 token 数平方增长，降维度对显存帮助更直接。
2.3 训练循环里必须对齐的噪声调度
骨干换成了 Transformer，但扩散过程本身还是 DDPM 那套。下面是一个最小训练步，重点看噪声采样和损失计算。
import torch.nn.functional as F def train_step(model, x0, optimizer, device): """单步训练：采样时间步、加噪、预测噪声、算 MSE""" B = x0.shape[0] model.train() optimizer.zero_grad() # 均匀采样时间步，实际训练中常用重要性采样 t = torch.randint(0, 1000, (B,), device=device).long() # 噪声调度：线性 beta，从 1e-4 到 0.02 betas = torch.linspace(1e-4, 0.02, 1000, device=device) alphas = 1.0 - betas alphas_cumprod = torch.cumprod(alphas, dim=0) sqrt_alpha_cumprod = torch.sqrt(alphas_cumprod[t]) sqrt_one_minus_alpha_cumprod = torch.sqrt(1 - alphas_cumprod[t]) # 加噪 noise = torch.randn_like(x0) x_t = sqrt_alpha_cumprod[:, None, None, None] * x0 + \ sqrt_one_minus_alpha_cumprod[:, None, None, None] * noise # 预测噪声 pred_noise = model(x_t, t) loss = F.mse_loss(pred_noise, noise) loss.backward() optimizer.step() return loss.item()逻辑说明：先随机采时间步，然后按 DDPM 的闭式公式把干净图像加噪成x_t，让模型预测注入的噪声，损失用 MSE。参数说明：betas从 1e-4 到 0.02 是 DDPM 原论文的线性调度，对 32x32 和 64x64 图像都适用。如果生成质量发灰，优先检查alphas_cumprod的计算有没有用错维度。t的类型必须是 long，否则索引会报错。训练初期 loss 在 0.8 到 1.2 之间波动正常，降到 0.3 以下基本就收敛了。
3. 从零跑通 DiffiT 训练：数据、超参和采样验证
3.1 数据准备和增强策略
DiffiT 对数据格式没特殊要求，但 patch embed 对图像尺寸敏感。常见做法是把图像统一 resize 到 32x32 或 64x64，再做随机水平翻转。不要用随机裁剪，因为裁剪会破坏全局结构，而 Transformer 恰恰依赖全局注意力。下面是一个简单的 Dataset 实现。
from torch.utils.data import Dataset, DataLoader from torchvision import transforms from PIL import Image import os class ImageFolderFlat(Dataset): """读取一个目录下的所有图片，统一 resize 和归一化""" def __init__(self, root, img_size=32): self.paths = [os.path.join(root, f) for f in os.listdir(root) if f.lower().endswith(('.png', '.jpg', '.jpeg'))] self.transform = transforms.Compose([ transforms.Resize((img_size, img_size)), transforms.RandomHorizontalFlip(), transforms.ToTensor(), # 自动归一化到 [0,1] transforms.Normalize([0.5]*3, [0.5]*3) # 转到 [-1,1] ]) def __len__(self): return len(self.paths) def __getitem__(self, idx): img = Image.open(self.paths[idx]).convert('RGB') return self.transform(img) # 使用示例 dataset = ImageFolderFlat('./data/train', img_size=32) loader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=4, drop_last=True)逻辑说明：Resize保证所有图像尺寸一致，RandomHorizontalFlip是唯一推荐的增强，Normalize把像素值从 [0,1] 转到 [-1,1]，和扩散模型的噪声分布匹配。参数说明：batch_size32 是 12GB 显存的保守值，如果换 64x64 图像，降到 16 或 8。num_workers根据 CPU 核数设，一般 4 到 8。drop_last=True避免最后一个 batch 尺寸不一致导致 reshape 报错。
3.2 学习率、优化器和 EMA 的配置
Transformer 训练对学习率敏感，DiffiT 也不例外。我一般用 AdamW，学习率 1e-4，权重衰减 0.01，配合余弦退火。EMA 不是必须的，但加上之后采样质量明显更稳。
import torch.optim as optim from torch.optim.lr_scheduler import CosineAnnealingLR model = DiffiT(img_size=32, patch_size=4, dim=384, depth=12, num_heads=6).cuda() optimizer = optim.AdamW(model.parameters(), lr=1e-4, weight_decay=0.01) scheduler = CosineAnnealingLR(optimizer, T_max=100000, eta_min=1e-6) # EMA 模型，用于采样时取参数平均 class EMA: def __init__(self, model, decay=0.9999): self.decay = decay self.shadow = {k: v.clone().detach() for k, v in model.state_dict().items()} def update(self, model): for k, v in model.state_dict().items(): if v.dtype.is_floating_point: self.shadow[k] = self.shadow[k] * self.decay + v * (1 - self.decay) else: self.shadow[k] = v.clone() ema = EMA(model)逻辑说明：AdamW 的权重衰减是解耦的，比 Adam 更适合 Transformer。余弦退火让学习率从 1e-4 平滑降到 1e-6，避免后期震荡。EMA 的 decay 设 0.9999，训练步数越多越稳。参数说明：T_max设成总训练步数，比如 100k 步。如果 loss 在 10k 步内不降，先检查学习率是不是太大，降到 5e-5 试试。EMA 的update要在optimizer.step()之后调用。
3.3 采样验证：用 DDIM 加速出图
训练过程中要定期采样看效果。DDPM 采样 1000 步太慢，用 DDIM 50 步就能看出大概。下面是一个最小采样函数。
@torch.no_grad() def sample_ddim(model, shape, device, num_steps=50, eta=0.0): """DDIM 采样，eta=0 是确定性采样，eta=1 退化成 DDPM""" model.eval() betas = torch.linspace(1e-4, 0.02, 1000, device=device) alphas = 1.0 - betas alphas_cumprod = torch.cumprod(alphas, dim=0) # 选取 DDIM 的时间步子序列 step_indices = torch.linspace(0, 999, num_steps, device=device).long() x = torch.randn(shape, device=device) for i in reversed(range(num_steps)): t = step_indices[i] t_prev = step_indices[i-1] if i > 0 else torch.tensor(0, device=device) alpha_t = alphas_cumprod[t] alpha_prev = alphas_cumprod[t_prev] pred_noise = model(x, t.unsqueeze(0).expand(shape[0])) # DDIM 更新公式 pred_x0 = (x - torch.sqrt(1 - alpha_t) * pred_noise) / torch.sqrt(alpha_t) pred_x0 = pred_x0.clamp(-1, 1) sigma = eta * torch.sqrt((1 - alpha_prev) / (1 - alpha_t) * (1 - alpha_t / alpha_prev)) c = torch.sqrt(1 - alpha_prev - sigma ** 2) x = torch.sqrt(alpha_prev) * pred_x0 + c * pred_noise if eta > 0: x = x + sigma * torch.randn_like(x) return x逻辑说明：DDIM 的核心是跳步采样，用子序列代替全部 1000 步。pred_x0是从x_t和预测噪声反推的干净图像，然后按 DDIM 公式更新到x_{t-1}。参数说明：num_steps50 是质量和速度的平衡点，再低到 20 会明显糊。eta=0是确定性采样，同一噪声输入每次出图一样，方便对比。clamp(-1,1)防止pred_x0数值爆炸。采样时模型要eval()，否则 BatchNorm 或 Dropout 会捣乱。
4. 避坑与排查：DiffiT 训练中最容易翻车的五个点
4.1 损失不降反升，采样全是噪声
现象：训练前 1000 步 loss 从 1.0 涨到 3.0 以上，采样出来是纯噪声。原因：学习率太大，或者时间嵌入的维度没对齐。Transformer 对初始化敏感，如果adaLN_modulation的最后一层没有零初始化，训练初期调制参数会很大，直接把注意力输出打飞。解决：把adaLN_modulation最后一层的权重和偏置初始化为零，学习率从 1e-4 降到 5e-5。我一般在DiffiTBlock的__init__里加一行nn.init.zeros_(self.adaLN_modulation[-1].weight)，这个血泪经验能省半天调参时间。
4.2 显存溢出，batch size 上不去
现象：跑 64x64 图像时，batch size 16 就 OOM。原因：patch size 设太小，token 数量爆炸。64x64 图像用 patch 4 得到 256 个 token，注意力矩阵是 256x256，比 32x32 的 64x64 大了 16 倍。解决：把 patch size 从 4 改成 8，token 数降到 64，显存直接降一个数量级。或者用梯度检查点，在DiffiTBlock的 forward 里包一层torch.utils.checkpoint.checkpoint，代价是训练速度慢 20% 左右。
4.3 采样图像有网格状伪影
现象：生成的图像上能看到明显的方块拼接痕迹。原因：patch embed 的卷积核重叠或者位置编码插值方式不对。如果patch_embed的 stride 小于 kernel_size，patch 之间会重叠，反推回图像时边界对不齐。解决：确保kernel_size == stride == patch_size。另外位置编码用可学习的绝对编码，不要用插值，因为插值会引入高频噪声。如果已经有伪影，检查out_proj后的 reshape 有没有用错h和w的顺序。
4.4 训练后期 loss 震荡，采样质量忽好忽坏
现象：loss 降到 0.3 左右开始上下跳，采样有时清晰有时糊。原因：学习率没有退火到底，或者 EMA 的 decay 太小。解决：把余弦退火的eta_min设到 1e-7，EMA decay 从 0.999 提到 0.9999。另外检查数据加载器有没有shuffle=True，如果数据顺序固定，模型会记住批次顺序，导致周期性震荡。我一般会在训练脚本里打印每个 epoch 的 loss 均值和方差，方差超过均值 30% 就说明有问题。
4.5 时间步采样不均匀导致某些步学不好
现象：采样时前几步正常，后面突然崩掉。原因：均匀采样时间步让模型在 t 接近 0 和 1000 的样本太少，这两个区域恰恰是去噪最难的部分。解决：用重要性采样，让 t 的分布偏向两端。常见做法是从 logit-normal 分布采样，或者简单点，把均匀采样和两端各 10% 的采样混合。代码上就是t = torch.randint(0, 1000, (B,))改成先采样一个 0 到 1 的均匀数，再用t = (torch.cos(u * math.pi) + 1) / 2 * 999映射，这样两端密度更高。
5. 进阶技巧：用分类器无关引导和注意力图做可控生成
5.1 分类器无关引导的工程实现
DiffiT 训练完之后，直接采样往往多样性够但保真度一般。分类器无关引导（CFG）是标配，做法是在训练时随机把条件（比如类别标签）置空，采样时同时算有条件噪声和无条件噪声，再线性外推。下面是一个带 CFG 的采样片段。
@torch.no_grad() def sample_with_cfg(model, shape, device, labels, guidance_scale=3.0, num_steps=50): """带分类器无关引导的 DDIM 采样，labels 是类别索引""" model.eval() betas = torch.linspace(1e-4, 0.02, 1000, device=device) alphas_cumprod = torch.cumprod(1 - betas, dim=0) step_indices = torch.linspace(0, 999, num_steps, device=device).long() x = torch.randn(shape, device=device) null_labels = torch.full_like(labels, 10) # 假设 10 是空类别 for i in reversed(range(num_steps)): t = step_indices[i] t_batch = t.unsqueeze(0).expand(shape[0]) # 有条件预测 noise_cond = model(x, t_batch, labels) # 无条件预测 noise_uncond = model(x, t_batch, null_labels) # CFG 外推 pred_noise = noise_uncond + guidance_scale * (noise_cond - noise_uncond) alpha_t = alphas_cumprod[t] alpha_prev = alphas_cumprod[step_indices[i-1]] if i > 0 else torch.tensor(1.0, device=device) pred_x0 = (x - torch.sqrt(1 - alpha_t) * pred_noise) / torch.sqrt(alpha_t) pred_x0 = pred_x0.clamp(-1, 1) x = torch.sqrt(alpha_prev) * pred_x0 + torch.sqrt(1 - alpha_prev) * pred_noise return x逻辑说明：模型需要支持类别条件输入，训练时以 10% 概率把标签替换成空类别。采样时同时算两次前向，按guidance_scale外推。参数说明：guidance_scale设 3.0 是常用起点，调到 7 以上颜色会过饱和，调到 1 以下退化成无条件生成。空类别索引要和训练时一致，我一般用num_classes作为空类别的索引。两次前向可以拼成一个 batch 减少开销，但显存会翻倍。
5.2 用注意力图做局部编辑
DiffiT 的注意力图是可解释的，最后一层的注意力权重能看出模型在生成某个区域时关注了哪些 patch。利用这一点可以做局部重绘：固定其他区域的噪声，只对目标区域重新采样。具体做法是在采样循环里，每步更新完之后，把非编辑区域的x替换成原始采样轨迹的值。这个技巧在 64x64 图像上效果明显，再大就需要更精细的掩码。我一般会导出注意力图看看模型有没有关注到正确的区域，如果注意力散乱，说明训练还不够或者 patch size 太大。
5.3 验证 DiffiT 是否真的学到了全局结构
一个简单的验证方法：生成一张左右对称的图，比如正面人脸，然后把左半边的 token 注意力权重可视化。如果模型真的学到了全局依赖，右半边的 patch 应该对左半边的 patch 有高注意力权重。具体操作是取最后一层注意力的平均头权重，reshape 成N x N矩阵，再按 patch 位置还原成图像。如果矩阵接近对角占优，说明模型只关注局部，全局建模没学好，这时候要检查位置编码有没有正确加载，或者 depth 是不是太浅。我习惯在训练 20k 步后做一次这个检查，比只看 loss 曲线靠谱得多。
希望帮到你。
本文还有配套的精品资源，点击获取
