简介：这份资源是NVIDIA团队发表的《DiffiT: Diffusion Vision Transformers for Image Generation》论文PDF，面向从事扩散模型、图像生成与视觉变换器研究的算法工程师、研究生及AI爱好者，帮助读者理解如何用ViT替代传统卷积U-Net作为去噪网络。压缩包内仅含1个PDF文件，约43.22MB，完整呈现论文正文、图表与实验数据。内容围绕混合层次U形编码器解码器结构、时间依赖自注意力模块展开，并给出latent DiffiT在ImageNet-256上取得1.73 FID的SOTA结果，同时覆盖类条件与无条件合成任务的评估分析。读者可从中获取去噪网络架构设计的完整思路、自注意力在扩散过程中的动态调整机制，以及高分辨率图像生成的实验细节与对比结论，适合作为研究与复现的参考材料。目前已有338人学习。

1. 时间步不是外挂参数：DiffiT 为什么要把时间嵌进注意力里

扩散模型生成一张图，要反复去噪几十上百步。绝大多数实现里，时间步 t 的处理方式相当“敷衍”：算一个时间嵌入向量，然后通过加法或 AdaLN 的 shift/scale 注入到残差块里。卷积核本身跟时间无关，时间只做了一次通道级的平移和缩放。这在工程上够用，但有个反直觉的结论——去噪网络真正需要的是“在不同时间步改变注意力行为”，而不是给特征图加一个偏置。DiffiT 的核心主张就一句话：把时间 token 和空间 token 一起投影成 Q、K、V，让注意力权重本身随 t 变化。这解决了扩散模型里长期被忽视的一个问题：早期去噪该看低频全局结构，后期该看高频细节，而固定卷积核加时间偏置根本表达不了这种切换。搞懂这个设计，你就理解了为什么 Transformer 能在扩散模型里替代 U-Net 的卷积主干，而不是简单“换个 backbone”。

2. 时间依赖自注意力：Q、K、V 怎么被时间步改写

2.1 从“时间加在特征上”到“时间参与投影”

传统做法可以写成h = Conv(x) + time_emb，时间嵌入只影响通道偏移。DiffiT 换了个思路：每个 Transformer block 同时收到空间 token 集合{x_s}和一个时间 tokenx_t，然后对 Q、K、V 分别做双线性投影。论文里的公式是：

# 时间依赖多头自注意力的 Q/K/V 构造（对应论文 Eq.3-5） # x_s: 空间 token, shape [B, N, D] # x_t: 时间 token, shape [B, 1, D] # W_qs/W_qt 等为空间/时间两套线性投影权重 q_s = x_s @ W_qs + x_t @ W_qt # 查询同时含空间与时间信息 k_s = x_s @ W_ks + x_t @ W_kt # 键同理 v_s = x_s @ W_vs + x_t @ W_vt # 值同理 # 注意力计算，B 为相对位置偏置 attn = softmax((q_s @ k_s.transpose(-1, -2)) / sqrt(d) + B) @ v_s

关键在参数说明：W_qs、W_ks、W_vs是空间投影，W_qt、W_kt、W_vt是时间投影，两组权重独立。代码后要理解的是——这等价于把空间 token 和时间 token 拼接后做一次线性投影，但拆开写让时间分支可以单独初始化、单独分析。因为 Q、K、V 都是时间和空间的线性函数，注意力打分矩阵QK^T里就隐含了t的二次项，注意力模式会随 t 连续变化。这比 AdaLN 的 shift/scale 表达能力强一个量级：后者只能整体平移缩放特征，前者能改变“哪些 token 之间互相关注”。

2.2 相对位置偏置为什么不跟时间耦合

论文 Eq.6 里有个细节：B是相对位置偏置，但它不直接乘时间嵌入。作者在 5.4 节专门分析过，如果把时间信息硬塞进位置偏置，反而次优，因为位置偏置要同时编码空间相对关系和时序变化，两个目标打架。常见的实现思路是让位置偏置只负责空间，时间依赖全部交给 Q/K/V 投影去学。我一般会这么理解：位置偏置是“结构先验”，时间投影是“动态调制”，两者职责分离，梯度不会互相污染。这也是 DiffiT 跟 DiT 的一个分水岭——DiT 用 AdaLN 把时间条件做成 shift/scale，DiffiT 则让时间进入注意力核本身。

2.3 DiffiT Block 的残差结构

一个 DiffiT Transformer block 的完整前向是：

# DiffiT block 前向（对应论文 Eq.7-8） # LN: LayerNorm, MLP: 两层全连接 + GELU x_hat = TMSA(LN(x_s), x_t) + x_s # 时间依赖自注意力 + 残差 x_out = MLP(LN(x_hat)) + x_hat # 前馈网络 + 残差

注意 TMSA 的输入是LN(x_s)和原始x_t，时间 token 不经过 LayerNorm 再进注意力，而是直接参与投影。这个顺序在源码里很关键：如果先把时间 token 也 LN 掉，它的尺度会被归一化，时间步之间的区分度下降，去噪早期和后期的注意力差异会被抹平。血泪经验是，复现时把x_t误放进 LN 分支，FID 会明显变差，但训练 loss 看不出来——这就是黑匣子式的隐性 bug。

3. U 形层次结构：Transformer 怎么搭出编码器-解码器

3.1 图像空间 DiffiT 的分辨率阶梯

DiffiT 不是纯 Transformer 平铺，而是保留了 U-Net 的层次骨架：编码器逐级下采样，解码器逐级上采样，同分辨率之间加 skip connection。图像空间版本的通道/分辨率走向大致是 64×64×3 → 64×64×128 → 32×32×256 → 16×16×256 → 8×8×256，再对称回升。每个阶段堆叠若干 DiffiT ResBlock 和 DiffiT Transformer block，下采样/上采样用卷积完成。

为什么不全用 Transformer？因为高分辨率特征图上的全局自注意力计算量是 token 数的平方，64×64 就有 4096 个 token，直接做全局注意力显存吃不消。U 形结构让注意力主要作用在低分辨率层，高分辨率层用卷积残差块扛，这是计算复杂度和长程建模之间的折中。常见的实现思路是：浅层高分辨率用卷积抓局部纹理，深层低分辨率用 TMSA 抓全局结构和时间动态。

3.2 latent DiffiT：把 Transformer 搬到潜空间

高分辨率生成走的是 latent 路线：先用一个 tokenizer（自编码器）把图像压到潜空间，再在潜空间上跑 DiffiT。论文里 latent DiffiT 在 ImageNet-256 上拿到 FID 1.73，这个数字的意义在于——它证明 Transformer 主干在潜空间扩散里不仅能 work，还能刷到 SOTA。潜空间的 token 数远少于像素空间，全局注意力变得可行，U 形层次可以简化成更纯粹的 Transformer 堆叠。

这里的设计取舍值得说透：图像空间 DiffiT 保留完整 U 形，是因为像素级去噪需要多尺度细节；latent DiffiT 的 tokenizer 已经承担了部分多尺度压缩，主干可以更“Transformer 化”。两者共享同一套 TMSA 和 DiffiT block，只是宏观排布不同。这种“统一设计模式”是论文的贡献之一——去噪网络不再是一堆卷积残差块的拼凑，而是有明确模块边界的分层 Transformer。

3.3 skip connection 在 Transformer U 形里的作用

U 形结构里，编码器每级输出通过 skip 连到解码器同级。在卷积 U-Net 里这很自然，因为两边通道数对齐。放到 Transformer 里要注意：skip 连接的是特征图，不是 token 序列，所以同分辨率阶段的 token 网格必须一致。实现上通常在同分辨率阶段之间做通道拼接或相加，再进下一级 DiffiT block。如果 skip 路径的通道数和主干不匹配，需要 1×1 卷积对齐——这个细节论文正文没展开，但源码里一定有，复现时漏掉会直接 shape 报错。

4. 训练目标与采样：score 网络和 Transformer 怎么对接

4.1 去噪目标函数没有变

DiffiT 再花哨，训练目标还是标准的去噪自编码器损失：

# 扩散训练目标（对应论文 Eq.1） # z0: 干净样本, sigma_t: 噪声强度, eps: 采样噪声 z_t = z0 + sigma_t * eps loss = lambda_t * || eps - eps_theta(z_t, t) ||_2^2

eps_theta就是 DiffiT 网络，输入带噪样本和时间步，输出预测噪声。lambda_t是时间相关的损失权重。这里要强调的是：架构换了，扩散的数学框架没换。TMSA 只是让eps_theta对 t 的表达更细粒度，训练信号和采样方程完全沿用。所以从 U-Net 迁移到 DiffiT，数据管线、噪声调度、损失权重都可以复用，改动集中在网络定义。

4.2 采样器：ODE 和 SDE 的选择

采样阶段解的是反向随机微分方程（论文 Eq.2）。beta_t = 0时退化成概率 ODE，可以用 DDIM 这类 ODE 积分器，步数少、收敛快；beta_t > 0时是 SDE，用 DDPM 类求解器，对不准确的 score 模型更鲁棒。DiffiT 论文没有绑定特定采样器，这意味着它的收益来自网络表达力，而不是采样技巧。实际用的时候，少步数高质量优先选 ODE 路线，稳定性优先选 SDE 路线。踩坑点是：换了主干后，原来 U-Net 上调好的采样步数和beta_tschedule 不一定最优，需要重新扫一遍。

4.3 时间嵌入的构造

时间 token 的生成方式跟 Ho 等人的做法一致：位置时间嵌入过一个带 swish 激活的小 MLP，得到x_t，然后广播到所有层。注意这个x_t是每一层都重新参与投影的，不是只在输入端注入一次。这跟 AdaLN 每层做 shift/scale 在形式上类似，但作用点不同——AdaLN 作用在归一化后的特征上，TMSA 作用在 Q/K/V 投影上。前者调制特征幅度，后者调制 token 间的关联结构。

5. 避坑与排查：复现 DiffiT 时最容易翻车的几处

5.1 时间 token 被 LayerNorm 误伤

现象：训练 loss 正常下降，但生成样本在不同采样步之间差异很小，早期和后期去噪看起来一个样。原因：把x_t也送进了 LayerNorm，时间步的尺度信息被归一化掉，TMSA 退化成近似时间无关的注意力。解决：确认 TMSA 内部只对空间 token 做 LN，时间 token 直接进线性投影；检查代码里LN(x_s)和x_t是分开传入的。

5.2 相对位置偏置维度对不上

现象：多头注意力里位置偏置广播报错，或者训练不报错但 FID 异常高。原因：B的形状要能广播到[B, heads, N, N]，如果按单头形状构造，多头时会静默广播错。解决：显式检查B的 shape，确保每个 head 有独立或正确共享的位置偏置；对照论文 Eq.6 确认B加在 softmax 之前、缩放之后。

5.3 skip connection 通道未对齐

现象：U 形解码器拼接时 shape mismatch，或者用相加代替拼接后特征表达变弱。原因：编码器和解码器同分辨率阶段的通道数在 DiffiT 里不一定相等，直接相加会丢信息。解决：skip 路径加 1×1 卷积把通道对齐到解码器输入通道，再拼接或相加；确认下采样/上采样倍数和 token 网格严格对应。

5.4 潜空间模型误用像素空间配置

现象：latent DiffiT 训练显存爆炸或收敛极慢。原因：把图像空间的分辨率阶梯和通道配置直接套到潜空间，token 数和通道数不匹配 tokenizer 的输出。解决：先确认 tokenizer 的 latent 通道数和下采样倍率，再按潜空间 token 数重新排布 U 形各级；潜空间通常不需要那么深的编码器。

5.5 采样步数与噪声 schedule 未重调

现象：换 DiffiT 后生成质量不如预期，但训练指标正常。原因：沿用旧 U-Net 的采样步数和beta_t配置，跟新主干的 score 精度不匹配。解决：固定模型后单独扫采样步数（如 20/50/100）和 ODE/SDE 切换，观察 FID 曲线拐点；不要用训练 loss 判断采样配置好坏。

6. 设计取舍：为什么是“时间进注意力”而不是“时间进归一化”

DiffiT 和 DiT 的分歧点，本质是时间条件该放在哪一层。DiT 用 AdaLN，把时间嵌入变成每层的 shift 和 scale，计算便宜、实现简单，但它是全局通道级调制——同一时间步下，所有空间位置被同样的平移缩放。DiffiT 的 TMSA 把时间放进 Q/K/V 投影，时间影响的是 token 之间的注意力权重，是位置对级别的调制。表达力更强，代价是每个注意力层多了三组时间投影矩阵，参数量和计算量都上升。

这个取舍在什么场景下划算？当去噪过程需要“空间上不均匀”的时间行为时——比如早期某些区域已经出现结构、另一些还是纯噪声，全局 shift/scale 没法区别对待，TMSA 可以。反过来，如果任务简单、分辨率低，AdaLN 的性价比更高。论文在 ImageNet-256 上刷到 FID 1.73，说明在中高分辨率、复杂分布上，TMSA 的额外开销换来了实打实的质量提升。

验证这套设计是否真的起作用，我一般会做消融对照：把 TMSA 换成时间无关注意力加 AdaLN，保持参数量接近，看 FID 差距；再单独冻结时间投影权重，观察不同采样步的注意力图是否退化。如果冻结后生成质量明显掉，说明时间确实在改变注意力行为，而不是被其他路径吸收。从那以后我每次复现这类“时间条件进注意力”的模型，都强制先跑一遍时间投影冻结对照，再决定要不要投入调参。希望帮到你。

本文还有配套的精品资源，点击获取