DiffiT源码复现：时间依赖自注意力与U形Transformer扩散模型

简介这份资源是NVIDIA团队发表的《DiffiT: Diffusion Vision Transformers for Image Generation》论文PDF面向扩散模型、图像生成方向的研究者与算法工程师聚焦去噪网络架构选择这一常被忽视的问题。论文提出以视觉变换器替代传统卷积残差U-Net构建含U形编码器与解码器的混合层次结构并引入时间依赖自注意力模块使注意力层能随去噪阶段自适应调整行为其latent DiffiT在ImageNet-256上取得1.73的FID分数刷新当时最优水平。资源包共1个PDF文件约43.22MB完整收录摘要、引言、模型设计、实验结果与结论等章节并附相关知识点解析涵盖扩散模型原理、ViT机制、FID指标与类条件、无条件合成任务等便于读者系统理解方法脉络与实验细节。已有338人学习适合作为图像生成方向论文精读与架构选型的参考材料。1. 当扩散模型换掉 U-NetDiffiT 这份源码包到底能跑出什么扩散模型生成图像大家默认的骨架几乎都是卷积残差 U-Net自注意力只在低分辨率特征图上点缀几层。DiffiT 这份来自 NVIDIA 团队的官方实现做的是一件反直觉的事把去噪网络整体换成视觉 Transformer并且让注意力本身随去噪时间步变化。它解决的不是“能不能生成”而是“去噪网络架构该怎么设计”这个长期被忽视的问题。适合两类人一是想复现 ImageNet-256 上 FID 1.73 这类指标的研究型工程师二是想把 Transformer 去噪骨架接进自己训练管线、需要看清时间条件注入方式的从业者。源码包给的是可训练的模型定义与采样脚本不是开箱即用的产品读它之前最好对 DDPM/DDIM 的噪声调度有基本概念。2. 时间依赖自注意力TMSA 到底改了哪几个矩阵DiffiT 的核心不是“用了 Transformer”而是“Transformer 里的注意力怎么感知时间步”。传统做法是把时间嵌入通过 AdaLN 的 shift/scale 注入残差块卷积核本身与时间无关。DiffiT 把时间信息直接塞进 Q、K、V 的线性投影里让每一层的注意力在不同去噪阶段表现出不同行为。这一章把这个模块拆到矩阵级别再落到代码里怎么改。2.1 从 AdaLN 到 TMSA为什么时间要进 QKV先看常规 DiT 类模型的条件注入时间嵌入 t 经过 MLP 得到调制向量对 LayerNorm 的输出做逐通道的缩放和平移即 AdaLN。这种方式的问题在于它调制的是特征图的通道统计量注意力打分矩阵 QK^T 的结构并不直接受时间影响。去噪早期图像几乎全是高斯噪声网络该关注低频全局结构去噪后期结构已定该关注高频细节。这种“注意力该看哪里”的切换AdaLN 只能间接影响。DiffiT 的做法是给空间 token 和时间 token 各配一套投影权重然后相加q_s x_s · W_qs x_t · W_qtk_s x_s · W_ks x_t · W_ktv_s x_s · W_vs x_t · W_vt其中 x_s 是二维网格上排布的空间 tokenx_t 是时间嵌入经 swish 激活的小 MLP 得到的单个时间 token。因为 Q、K、V 都是空间与时间的线性函数注意力打分自然随 t 变化。论文里特别指出这等价于把时间 token 拼到每个空间 token 后面再做一次线性投影但分开写权重更利于实现时复用。一个容易忽略的细节相对位置偏置 B 是加在 softmax 之前的即 Softmax(QK^T/√d B)V。论文明确说 B 虽然隐式受时间嵌入影响但不要把时间直接融进 B否则它要同时编码空间和时间反而次优。这个结论在复现时值得记一笔很多人会想当然地把位置偏置也做成时间相关。2.2 DiffiT Block 的代码骨架与参数含义下面这段是 DiffiT Transformer Block 的最小可读实现保留了 TMSA 的关键投影拆分去掉了工程上的并行优化方便对照论文公式import torch import torch.nn as nn class TMSA(nn.Module): def __init__(self, dim, heads, dim_head): super().__init__() self.heads heads self.scale dim_head ** -0.5 inner heads * dim_head # 空间 token 的 QKV 投影 self.to_qkv_spatial nn.Linear(dim, inner * 3, biasFalse) # 时间 token 的 QKV 投影与空间投影维度对齐 self.to_qkv_time nn.Linear(dim, inner * 3, biasFalse) self.to_out nn.Linear(inner, dim) def forward(self, x, t_token, pos_biasNone): # x: (B, N, C) 空间 token; t_token: (B, 1, C) 时间 token qkv_s self.to_qkv_spatial(x).chunk(3, dim-1) qkv_t self.to_qkv_time(t_token).chunk(3, dim-1) # 空间与时间投影逐元素相加得到时间相关的 QKV q, k, v [s t for s, t in zip(qkv_s, qkv_t)] B, N, _ q.shape # 拆多头: (B, heads, N, dim_head) q, k, v [z.view(B, N, self.heads, -1).transpose(1, 2) for z in (q, k, v)] attn (q k.transpose(-2, -1)) * self.scale if pos_bias is not None: attn attn pos_bias # 相对位置偏置不参与时间融合 attn attn.softmax(dim-1) out (attn v).transpose(1, 2).reshape(B, N, -1) return self.to_out(out) class DiffiTBlock(nn.Module): def __init__(self, dim, heads, dim_head, mlp_ratio4.0): super().__init__() self.norm1 nn.LayerNorm(dim) self.tmsa TMSA(dim, heads, dim_head) self.norm2 nn.LayerNorm(dim) self.mlp nn.Sequential( nn.Linear(dim, int(dim * mlp_ratio)), nn.GELU(), nn.Linear(int(dim * mlp_ratio), dim), ) def forward(self, x, t_token, pos_biasNone): x self.tmsa(self.norm1(x), t_token, pos_bias) x x self.mlp(self.norm2(x)) x return x逻辑说明to_qkv_spatial和to_qkv_time分别对应论文里的 W_qs/W_ks/W_vs 与 W_qt/W_kt/W_vt两者输出维度必须一致才能相加这是实现时最容易写错的地方——如果时间投影只输出 dim 而不是 inner*3相加会直接报维度不匹配。pos_bias单独传入、在 softmax 前相加对应公式里的 B刻意不和时间 token 做任何运算。参数上dim_head一般取 64heads按 dim/64 推算mlp_ratio取 4 是 Transformer 的常规默认DiffiT 论文没有在这块做特殊设计照搬即可。2.3 时间 token 的生成与注入位置时间 token 不是每层重新算而是在网络入口算一次然后传给所有 DiffiT Block。常见做法是把连续时间步 t 先做正弦位置编码再过一个小 MLPclass TimeToken(nn.Module): def __init__(self, dim, freq_dim256): super().__init__() self.freq_dim freq_dim self.mlp nn.Sequential( nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim), ) def forward(self, t): # t: (B,) 连续时间步范围约 [0, 1000] half self.freq_dim // 2 freqs torch.exp(-torch.arange(half, devicet.device) * (9.21 / half)) args t[:, None].float() * freqs[None] emb torch.cat([torch.cos(args), torch.sin(args)], dim-1) return self.mlp(emb).unsqueeze(1) # (B, 1, dim)这里9.21约等于 ln(10000)是正弦编码的常用底数缩放保证不同频率覆盖足够宽的时间范围。SiLU就是论文里的 swish 激活。生成后unsqueeze(1)变成单 token 形状广播到所有空间 token 上做投影相加。参数上freq_dim取 256 是常见值太小会导致时间分辨率不足去噪后期细节阶段区分度下降太大则 MLP 第一层参数量上升收益不明显。提示时间 token 必须在整个去噪网络里共享同一份不要在每个 Block 里各自对 t 重新编码否则不同层看到的时间表示不一致训练时 loss 会抖得厉害。3. 图像空间与潜空间两套骨架U 形层次怎么搭DiffiT 给了两个变体图像空间的 DiffiT 直接在像素上做潜空间的 latent DiffiT 接在 VAE 编码后的特征上。两者共享 TMSA 和 DiffiT Block差别在整体拓扑和分辨率处理。这一章把 U 形编码器-解码器的搭建逻辑讲清楚并给出可复现的层级配置。3.1 图像空间 DiffiT 的层级与通道配置图像空间版本是混合层次结构不是纯 Transformer 的等分辨率堆叠而是像 U-Net 一样有下采样和上采样每个分辨率阶段堆若干 DiffiT ResBlock 或 DiffiT Transformer。论文图 2 给出的通道走向大致是 64×64×3 输入经 tokenizer 和若干阶段到 8×8×256 的瓶颈再对称上采样回 64×64×3。每个阶段之间用卷积下采样/上采样阶段内部用 DiffiT Block 做注意力。搭建时的一个关键选择低分辨率阶段用 Transformer Block高分辨率阶段保留卷积 ResBlock。原因是标准自注意力的计算量随 token 数平方增长64×64 分辨率下 token 数 4096全用注意力显存吃不消。常见做法是在 16×16 及以下用 TMSA32×32 和 64×64 用带时间嵌入的卷积残差块。下面是一个阶段配置的示意# 每个 stage: (分辨率, 通道, 块类型, 块数) image_stages [ (64, 128, conv_res, 2), # 高分辨率卷积为主 (32, 256, conv_res, 2), (16, 256, transformer, 4), # 进入 Transformer 阶段 (8, 256, transformer, 4), # 瓶颈 ]参数说明通道从 128 起到瓶颈 256和论文图 2 的 64→128→256→256 走向一致块数在 Transformer 阶段给 4卷积阶段给 2是算力和表达力的折中。如果显存紧张优先砍 64 分辨率的块数因为那一层 token 最多、注意力最贵如果误用了 Transformer 的话。跳连接在每个分辨率上把编码器特征拼到解码器对应层拼接后过 1×1 卷积对齐通道这一步和 U-Net 完全一致不要省。3.2 latent DiffiT接 VAE 时的分辨率与通道对齐latent 版本把像素先过 VAE 编码器压到潜空间常见下采样倍率 8即 256×256 图像变成 32×32 潜特征通道数由 VAE 决定常见 4 或 16。DiffiT 的 Transformer 直接在这个潜特征上跑不再做 U 形下采样到 8×8因为潜空间本身已经很小。论文在 ImageNet-256 上拿到 FID 1.73 的就是这个变体。接 VAE 时有两个对齐点必须检查。第一潜特征的通道数要和 DiffiT 的输入投影层dim匹配如果 VAE 输出 4 通道而模型dim256入口要有一个nn.Conv2d(4, 256, 1)或nn.Linear做升维。第二潜空间的噪声调度和像素空间不同VAE 的缩放因子常见 0.18215 这类常数要在加噪前乘进去否则训练时模型看到的噪声方差和采样时不一致生成结果会偏灰或过饱和。这两点在源码里通常写在数据预处理和noise_scheduler的初始化处复现时先翻这两处。# latent 入口对齐示意 class LatentPatchEmbed(nn.Module): def __init__(self, in_ch4, dim256, patch2): super().__init__() self.proj nn.Conv2d(in_ch, dim, kernel_sizepatch, stridepatch) def forward(self, z): # z: (B, 4, 32, 32) - (B, 256, 16, 16) - (B, 256, 256) x self.proj(z) return x.flatten(2).transpose(1, 2)patch2把 32×32 潜特征切成 16×16 个 patchtoken 数 256注意力完全扛得住。如果 patch 取 1token 数 1024计算量翻四倍但细节保留更好属于可调边界。in_ch必须和实际 VAE 输出一致这是接第三方 VAE 时最常见的翻车点。3.3 训练循环里时间步采样与 loss 的落点DiffiT 的训练目标和标准扩散模型一致都是预测噪声 εloss 是带权重的 MSE。差别只在网络前向时多传一个时间 token。下面是一个最小训练步def train_step(model, vae, x, scheduler, device): with torch.no_grad(): z vae.encode(x).latent_dist.sample() * 0.18215 noise torch.randn_like(z) t torch.randint(0, scheduler.num_train_timesteps, (z.shape[0],), devicedevice) z_t scheduler.add_noise(z, noise, t) t_token model.time_token(t.float()) pred model(z_t, t_token) # 简单均匀权重论文用 λ(t) 加权复现时可先跑通再调 loss ((pred - noise) ** 2).mean() return loss逻辑上add_noise按调度表把 z 加到 z_t模型预测噪声和真噪声算 MSE。参数上num_train_timesteps常见 1000t用均匀采样即可论文的 λ(t) 加权在初期可以先用常数 1 跑通确认 loss 能稳定下降后再加。0.18215是 VAE 缩放常数换 VAE 时要同步改。这一步跑通的标准是 loss 在前几千步单调下降、不出现 NaN如果 NaN 先查时间 token 有没有传对形状。4. 避坑与排查复现 DiffiT 时最容易翻车的五处这份源码包结构不算复杂但 TMSA 的投影拆分和 U 形拓扑叠在一起出错时往往不报错、只是指标难看。下面五条是按“现象→原因→解决”整理的踩坑记录都是复现这类 Transformer 扩散模型时高频出现的。4.1 生成图全是灰雾FID 卡在很高位现象采样出来的图有轮廓但整体发灰、对比度低FID 比论文差一大截。原因多半是潜空间缩放常数没对齐VAE 编码后没乘缩放因子或者采样时去噪调度和训练时的噪声表不是同一套。解决先打印训练时 z_t 的方差和理论 σ_t 对比偏差超过 10% 就说明缩放或调度有问题确认add_noise和采样器用的是同一个noise_scheduler实例不要训练用一个、采样新建一个。4.2 loss 前几百步就 NaN现象训练几步后 loss 变 NaN梯度爆炸。原因通常是时间 token 形状不对广播到了错误维度导致 QKV 相加后维度错乱但被后续 reshape 掩盖或者学习率对 Transformer 骨架偏高。解决在 TMSA 里打印 q、k、v 的 shape确认都是 (B, heads, N, dim_head)学习率先降到 1e-4 量级Transformer 去噪网络对初始学习率比 U-Net 敏感。4.3 高分辨率阶段显存直接爆现象64×64 阶段一开注意力就 OOM。原因是在高分辨率用了全注意力token 数 4096注意力矩阵 4096×4096 乘多头显存瞬间上去。解决按 3.1 的配置32×32 以上用卷积 ResBlock只在 16×16 及以下开 TMSA或者把 patch size 调大减少 token 数代价是细节损失。4.4 训练 loss 正常但采样图结构崩坏现象训练 loss 降得很好采样时早期步数正常、后期突然结构崩。原因常见于采样器步数和训练时间步分布不匹配比如训练用 1000 步均匀采样推理用 20 步 DDIM 但没做相应的步长映射。解决确认采样器的timesteps是从训练调度表里子采样出来的不是另起一套先用 50 步 DDIM 验证再压到 20 步看退化程度。4.5 换数据集后类别条件完全失效现象在自定义类别数据集上训练无条件生成还行带类别条件生成全是乱图。原因多是类别嵌入的维度或标签映射和模型里的num_classes不一致或者 classifier-free guidance 的 dropout 概率设成了 0训练时根本没学过无条件分支。解决核对num_classes与数据集标签数guidance dropout 保持在 0.1 左右采样时 guidance scale 从 1.5 起调别一上来就 7 以上。注意这五条里 4.1 和 4.4 经常同时出现排查时先固定随机种子、用同一批噪声对比训练和采样的中间张量比盲调超参快得多。5. 从跑通到可用用固定种子做逐阶段张量比对复现这类模型最省时间的习惯不是调参而是在跑通后立刻建立一套逐阶段张量比对流程。我的做法是固定一个种子构造一个极小的 batch比如 2 张图在训练前向和采样前向里各插一组 hook把每个 DiffiT Block 的输入输出、时间 token、注意力打分矩阵的统计量均值、方差、最大值打出来存成两份记录然后逐层对比。训练和采样在同一个时间步 t 上这些统计量应该高度接近如果某一层开始偏离问题就锁定在那层之前。具体操作上用 PyTorch 的register_forward_hook挂在每个 Block 上记录x.mean()、x.std()和t_token.norm()。采样时把采样器的当前 t 手动传给同一个 hook 逻辑保证两边 t 一致。对比表大致长这样层号训练 x.std采样 x.std偏差判定Block 11.021.011%正常Block 40.870.6130%查该层前的时间注入Block 80.550.1278%上游已崩回溯偏差超过 15% 就值得停下来查通常问题出在时间 token 的数值范围——训练时 t 是整数张量转 float采样时如果传的是归一化到 [0,1] 的 t正弦编码出来的频率完全不同注意力行为自然对不上。这个坑我在第一次接自定义采样器时踩过loss 一切正常采样图就是不对查了两天才发现是 t 的量纲问题。另一个实用技巧是给 TMSA 的注意力打分矩阵做可视化抽查取去噪早期t 接近 1000和后期t 接近 0各一张看注意力分布是否真的在变化。如果两个时间步的打分矩阵几乎一样说明时间投影权重没起作用多半是to_qkv_time的初始化太小或者被 weight decay 压没了。DiffiT 的价值就在这个时间依赖性上这一项不成立整个模型就退化成普通 DiT。从那以后我每次复现带时间条件的扩散模型都强制先跑一遍固定种子的逐层张量比对再谈训练和调参省下的返工时间远超这点准备工作。希望这套流程帮到你把 DiffiT 这份源码真正跑成自己能改的骨架。本文还有配套的精品资源点击获取