简介：《DiffiT: Diffusion Vision Transformers for Image Generation》是NVIDIA团队围绕扩散模型去噪网络架构创新撰写的论文PDF，面向生成式AI、计算机视觉与图像生成方向的研究者和进阶学习者。内容提出DiffiT混合层次U形编解码结构，引入时间依赖自注意力模块，使注意力层能随去噪阶段动态调整，并给出latent DiffiT在高分辨率生成中的实现与实验。资源包共1个PDF文件，大小约43.22MB，便于离线阅读与专题研读。已有338人学习或浏览，读者可从中获取扩散模型与视觉变换器结合的设计思路、类条件与无条件合成任务评估结果，以及ImageNet-256上FID 1.73的SOTA表现和相关知识点解析，适合论文复现、组会分享与方案调研时参考。

1. 一个反直觉的结论：去噪网络的时间感，应该写进注意力的权重里

Diffusion Vision Transformer 这一类模型最反直觉的一点是：明明整个生成过程是一条长长的去噪时间轴，主流的做法却只给网络一个小的 time embedding，用通道级 shift/scale 的方式注入卷积块——卷积核本身从 t=0 到 t=T 一模一样。DiffiT 把这套做法推翻了，它把时间令牌直接拼进 Q、K、V 的线性投影，让每个时间步的注意力权重都不同。改动只是在矩阵乘上加了一项，却让扩散模型在 ImageNet-256 的潜空间图像生成上拿到 1.73 的 FID。一条时间向量凭什么能从「调节旋钮」升级成「改变注意力行为」的开关？这件事值得从机制拆到公式，再拆到代码。

2. 时间嵌入的妥协：为什么给卷积块做通道缩放不够用

2.1 早期去噪看低频、晚期去噪看高频，这个时间动态从哪来

扩散模型的前向过程是一族随机变量：给定数据分布 q0(z0)，按 qt(zt)=∫q(zt|z0)q0(z0)dz0 注入高斯噪声，其中 q(zt|z0)=N(z0, σt²I)。σt 取非降序列，σ0=0，σT 远大于数据方差，这就是方差爆炸式的噪声调度；只要用一个与 t 相关的缩放项重新标定数据，它也能等价地转成方差保持式的调度。网络训练目标写成去噪自编码器形式，就是让 εθ 在 zt = z0 + σtε 上回归注入的噪声 ε，损失带一个 λ(t) 权重。换个角度看，这也是在拟合数据分数 sθ(zt,t) = −εθ(zt,t)/σt。

真正有意思的是这条时间轴上的统计特性并非均匀。去噪刚开始时图像里塞满强高斯噪声，高频内容被彻底扰动掉了，网络能做的只有把低频结构先摆正；到了末尾，图像骨架已经成型，剩下的活是补高频细节。这意味着同一个网络在不同 t 上承担的是性质完全不同的任务。采样阶段也是同一个道理：采样可以写成从 t=T 积分回 t=0 的随机微分方程，方程里 βt 描述采样过程中注入的随机噪声量。βt 恒为 0 时退化成概率 ODE，可以用 DDIM 这类积分器少步数收敛；βt 不为 0 就是 SDE 求解，步数多但对不准的分数模型更鲁棒。步数少和鲁棒这两件事互相牵制，也侧面说明网络对时间的敏感度必须足够细。

2.2 通道级 shift/scale 的表达上限

传统去噪网络的骨架是卷积残差 U-Net，时间条件通常这样注入：把位置时间嵌入过一个小 MLP，得到时间向量，再用空间加法或通道级的 shift/scale 作用到各个残差块上。放到公式里看，某个卷积块对特征 h 的作用大致是 conv(h)·γ(t) + β(t)，γ 和 β 都是逐通道的标量。这相当于对特征图做一次仿射变换，改变的是每个通道的整体强弱和偏置。

这种做法的表达上限很清楚：它只能调节「通道有多活跃」，不能调节「哪些位置之间该建立联系」。而去噪早期到晚期真正变化最大的恰恰是后者——早期要把分散的低频块拼成整体，晚期要在局部相邻像素间对齐细节。仿射变换对这种「看哪里」的变化几乎无能为力，因为卷积核的权重 ω 从头到尾没有变过，与时间无关，时间只改动了偏置和缩放。原论文对此的判断很直接：卷积滤波器和时间无关，时间嵌入只施加通道级 shift 与 scale，这种简单机制未必能在整个去噪过程中最优地刻画时间依赖。

2.3 自注意力的位置为什么更关键

有人会问：U-Net 里不是也有自注意力层吗，时间条件照样能传进去。问题在于，这些自注意力层通常只放在低分辨率特征图上，因为它的计算复杂度随 token 数平方增长，放全分辨率代价太高；而且怎么放、放几层、放在哪个阶段，长期没有标准设计模式。更关键的是，即便注意力层吃到了时间嵌入，注入方式仍然是「先算完注意力再加时间偏置」，注意力的打分矩阵 QKᵀ 本身不含 t。

这就是要解决的问题：时间条件必须进到注意力打分的内部，而不是停留在注意力的输出之后。DiffiT 的切入点正在这里——把时间令牌和空间令牌投影到同一个共享空间，再分别合成 Q、K、V。这样一来，每个时间步的注意力分布形状都变了，而不只是幅度变了。

3. TMSA 的实现：时间令牌如何与空间令牌拼出 Q、K、V

3.1 时间令牌怎么造出来

每个 DiffiT Transformer 块的输入有两样东西：一组按二维网格空间排布的令牌 xs，以及一个代表时间步的时间令牌 xt。xt 的造法和 Ho 等人的做法一致——把位置时间嵌入喂给一个带 swish 激活的小 MLP，得到一个向量。这个时间令牌会被送到去噪网络的每一层，和空间令牌一起参与计算。

这一步看着平淡，但它决定了一件重要的事：xt 是一个全局共享的标量时刻的向量表示，同一时间步下所有空间 token 用的是同一个 xt。这带来两个后果。第一，时间项在注意力里对所有 token 是共享的，实现上可以先算一次再广播，几乎不增加多少计算。第二，时间对注意力的影响只能通过「对所有 token 施加同一个偏移」来实现，它改变的是 token 之间的相对打分，而不是给每个 token 单独定制一个时间向量。理解这两点，后面的公式和代码就都顺了。

3.2 QKV 的双线性投影：从公式到代码

时间依赖多头自注意力（TMSA）的核心是三行公式。对空间令牌 xs 和时间令牌 xt，分别用空间侧权重 W^s 和时间侧权重 W^t 做线性投影再相加：

qs = xs·Wq_s + xt·Wq_t ks = xs·Wk_s + xt·Wk_t vs = xs·Wv_s + xt·Wv_t

这三行等价于把每个空间令牌与时间令牌拼接后做一次线性投影，区别只在参数的组织方式。注意力计算本身还是标准形式：Softmax(QKᵀ/√d + B)·V，d 是键的缩放因子，B 是相对位置偏置矩阵。注意 B 只编码空间相对位置，原论文明确指出，把时间嵌入直接混进相对位置偏置里可能得到次优结果，因为这一个矩阵要同时承担空间和时间两种信息。

值得自己推一遍的是展开后的四项。以 qs·ksᵀ 为例，展开后包含：纯空间项 xs·Wq_s·Wk_sᵀ·xsᵀ，两个交叉项 xs·Wq_s·Wk_tᵀ·xtᵀ 与 xt·Wq_t·Wk_sᵀ·xsᵀ，以及纯时间项 xt·Wq_t·Wk_tᵀ·xtᵀ。纯时间项对所有 token 对都是同一个标量，softmax 归一化会把它整体吸收掉，等于没起作用。真正让注意力分布随时间改变的是两个交叉项——它们相当于给每个空间 token 的打分向量加了一个由 t 决定的偏置。这就是「时间改变了注意力行为」的数学来源，也是它比通道 shift/scale 更强的根本原因。

落到代码上，常见的实现思路是拆成两组线性层再相加，而不是真的去拼张量：

import torch import torch.nn as nn class TimeDependentMSA(nn.Module): """TMSA：空间令牌与时间令牌在共享空间做双线性投影后算多头注意力""" def __init__(self, dim, heads, dim_head, bias=True): super().__init__() inner = heads * dim_head self.heads = heads self.scale = dim_head ** -0.5 # 空间侧投影 W^s self.to_q = nn.Linear(dim, inner, bias=bias) self.to_k = nn.Linear(dim, inner, bias=bias) self.to_v = nn.Linear(dim, inner, bias=bias) # 时间侧投影 W^t，输入是同一个时间令牌 xt self.to_q_t = nn.Linear(dim, inner, bias=bias) self.to_k_t = nn.Linear(dim, inner, bias=bias) self.to_v_t = nn.Linear(dim, inner, bias=bias) self.to_out = nn.Linear(inner, dim) def forward(self, x, t_token, rel_pos_bias=None): # x: (B, N, C) 展平后的空间令牌；t_token: (B, C) b, n, _ = x.shape # 时间项与 t 有关、对所有 token 共享，先算一次再沿 token 维广播 q = self.to_q(x) + self.to_q_t(t_token)[:, None, :] k = self.to_k(x) + self.to_k_t(t_token)[:, None, :] v = self.to_v(x) + self.to_v_t(t_token)[:, None, :] # 拆多头 q, k, v = [t.view(b, n, self.heads, -1).transpose(1, 2) for t in (q, k, v)] attn = (q @ k.transpose(-1, -2)) * self.scale if rel_pos_bias is not None: attn = attn + rel_pos_bias # B 项只编码空间相对位置，不掺时间 attn = attn.softmax(dim=-1) out = (attn @ v).transpose(1, 2).reshape(b, n, -1) return self.to_out(out)

参数含义：dim是令牌通道数，heads与dim_head决定多头切分，inner = heads * dim_head是投影后的总宽度，rel_pos_bias是形状可广播到注意力矩阵的空间相对位置偏置，bias控制线性层是否带偏置项。[:, None, :]那一步是关键，它把形状 (B, C) 的时间投影扩成 (B, 1, C)，再与 (B, N, C) 相加完成广播；漏掉它就会变成逐通道相加，时间信息直接被空间特征淹没，后面第 5 章会专门说这个坑。

设计上为什么不直接拼张量：拼接会让中间张量宽度翻倍再投影，显存和访存都不划算；拆成两组线性层后，时间项只依赖 t，推理或训练时同一步内可以复用一次计算结果，等于用一点参数换了一次广播。这是那种「懂原理才敢改」的地方——数学上等价，工程上差一截。

3.3 相对位置偏置为什么不跟时间嵌入混在一起

相对位置偏置 B 的作用是给每个注意力头编码 token 之间的空间相对关系，让注意力在不依赖绝对坐标的情况下也能区分「相邻」和「远处」。它与时间令牌有一个本质区别：B 是空间结构先验，训练中基本是稳定的；xt 是逐时间步变化的输入。把两者塞进同一个矩阵，等于要求一个矩阵同时拟合两种变化频率完全不同的信号。

原论文的说法是这种做法可能次优，因为该组件需要同时捕获空间和时间信息。我一般会这么理解：这本质上是一个容量分配问题。相对位置偏置的形状只与 token 排布有关，如果让它承担时间调制，每个时间步都想要不同的偏置，模型要么学一个折中解、要么把容量浪费在这种耦合上。把它拆开，让时间走 QKV、位置走 B，两条通路各管一件事，调参和排查也都更好定位——出问题时你能确定是时间分支没生效，还是位置先验配错了。

4. U 形层次与分辨率调度：DiffiT 怎么把计算量压在低分辨率上

4.1 图像空间 DiffiT：对称 U 形与四个分辨率层级

图像空间的 DiffiT 用的是对称的 U 形编码器-解码器，收缩路径和扩张路径在每个分辨率上用 skip connection 相连。入口是一个 3×3 卷积做 tokenizer，把 64×64×3 的图升到 64×64×128；随后每经过一个阶段就做一次卷积下采样，通道翻倍、分辨率减半，依次走到 32×32×256、16×16×256、8×8×256；解码端对称地上采样回 64×64×128，最后用 3×3 卷积头出 64×64×3。

表里的 L1 到 L4 是各阶段的重复次数，具体取值以论文附录和仓库配置为准，不同数据集规模下并不相同。每个 DiffiT ResBlock 内部由两部分组成：一组 Swish + GroupNorm + 3×3 卷积的卷积残差，以及一个 DiffiT Transformer 块。时间嵌入在 ResBlock 这一层仍然保留传统的注入方式，TMSA 则负责让注意力本身随时间变化。

4.2 下采样、上采样与 skip：为什么要在低分辨率上堆 Transformer

把 64×64 的特征图展平成 token 序列是 4096 个 token，注意力矩阵是 4096×4096，无论显存还是算力都不适合在这个尺度上堆 Transformer。所以架构上的取舍是：分辨率变化交给卷积完成，卷积的局部性先验在上下采样里天然合适；Transformer 主体压在 32×32、16×16、8×8 这些 token 数小得多的层级上，在那里做全注意力才划算。

这和纯 Transformer 的思路不同。以 DiT 为代表的方案把 U-Net 骨架整个换成 Transformer，在潜空间里做全注意力；DiffiT 的图像空间模型保留了 U 形骨架，因为它要直接在像素空间生成 64×64 的图，全程用注意力代价太高。skip connection 的作用也不只是梯度通路——编码端在 64×64 上捕获的细节，解码端需要拿回来补高频，否则上采样出来的结果会糊。

4.3 latent DiffiT：把像素空间换成潜空间

高分辨率生成走的是另一条路。latent DiffiT 先用自编码器把图像压到潜空间，再在潜空间里跑一个由 TMSA 层堆成的 Transformer，去掉了 U 形骨架和卷积上下采样。潜空间的好处是 token 数大幅下降：同样生成 256×256 或 512×512 的图，直接在像素空间做注意力几乎不可行，压到潜空间后在 32×32 或 16×16 的网格上做注意力就现实得多，而且可以在整个网络里统一用 Transformer 块，不必再设计卷积和注意力混排的规则。

代价是引入了自编码器的重建损失，最终画质上限受潜空间表达能力约束。原论文的实测里，latent DiffiT 在 ImageNet-256 上拿到 1.73 的 FID；图像空间的 DiffiT 则在 CIFAR-10 和 FFHQ-64 上取得同类最优。两类任务、两套架构，用同一套 TMSA 打底，这正是论文强调的「统一去噪网络设计模式」的含义——不是强行一套结构打所有分辨率，而是把时间条件这一个关键接口统一掉。

还有一个容易被忽略的工程细节：时间令牌在 U 形结构的每一层都要用，但生成它的 MLP 只有一份。这意味着时间表示的容量是全局共享的，不同分辨率层级看到的是同一个 xt；层级之间的差异靠各层自己的 W^t 来消化。这种设计省参数，但也意味着如果某个层级对时间的敏感度和别处差很远，只能靠该层的投影权重去补偿。

5. 排查：复现 TMSA 时最容易翻车的四个地方

5.1 四条踩坑记录

现象一：训练 loss 平稳下降，但采样结果在切换时间步时几乎没差别。原因通常是时间令牌的广播维度写错了，比如把t_token[:, None, :]写成了逐通道相加，或者干脆在某个 reshape 之后把时间维挤掉，导致时间项退化成常数偏置；这时网络实际学成了一个时间无关的回归器，loss 当然能降，因为它只要拟合平均意义上的噪声方向就够了。解决办法是加一个最小单测：固定同一个空间输入 x，分别喂 t=0.1 和 t=0.9 的时间令牌，断言输出的 q、k、v 差异超过阈值，同时断言注意力矩阵本身也变了。这个检查十行代码，能省掉几天排查。

现象二：小分辨率下正常，切到高分辨率阶段 FID 明显变差。多半是把时间信息塞进了相对位置偏置，或者让 B 的形状随 t 变化。相对位置偏置本来只该编码空间关系，一旦要求它同时表达时间，高分辨率下 token 对数量暴涨，这种耦合带来的误差会被放大。解决方式是把两条通路彻底分开：B 只由 token 坐标生成，与 t 无关；时间只走 QKV 的线性投影。改完之后再看注意力图，应该能看到注意力分布随着 t 推进从弥散走向局部，而不是整体幅度缩放。

现象三：训练到中途显存爆掉，报错位置在注意力矩阵。这是把 TMSA 用到了全分辨率层级上。64×64 展平是 4096 个 token，注意力矩阵元素量是千万级，多头之后还要再乘一份。正确的做法遵循架构本意：全分辨率阶段用卷积残差承担主要计算，TMSA 只放在 token 数可控的层级。如果确实需要全分辨率的全局建模，退一步的做法是窗口化注意力或降低该层的头数与宽度，而不是硬扛。

现象四：训练前期 loss 抖动剧烈，后期才慢慢稳住。常见原因是时间侧投影和时间侧投影用了同一套初始化尺度，训练初期时间项的量级压过空间项，注意力分布被打乱，网络在「看空间」和「看时间」之间反复横跳。比较稳的处理是把时间侧的线性投影做小初始化甚至零初始化，让模型先学空间注意力，再由梯度逐步释放时间调制能力；配合短 warmup，抖动基本能压下去。这个技巧在条件生成里很常见，本质是让新增的条件通路以「无操作」的状态起步。

5.2 上线前的最小验证清单

第一，验证时间确实进了注意力，而不是停在注意力的输出之后：关掉时间侧投影，FID 应该显著变差，如果几乎不变，说明你的时间条件在别处被旁路了。第二，验证相对位置偏置没有被时间污染：把 B 换成零矩阵重训一版，性能应该下降但不应崩掉，如果崩掉说明 B 承担了过多职责。第三，验证潜空间和像素空间不要混用一套超参：两者的特征尺度不同，噪声调度的缩放、学习率、批大小都得单独调，直接套用图像空间那套配置，latent 模型往往收敛很快但样本发糊。第四，固定随机种子做一次跨时间步的可视化，看同一张初始噪声在 t 递减过程中结构是怎么长出来的，这比看曲线更能发现「时间没生效」这类玄学问题。

注意：上面四条里，前两条属于实现错误，后两条属于配置错误。实现错误会让模型「看起来能跑但没学到东西」，配置错误则表现为「能学到但上限低」，排查时先分清是哪一类，能省掉大量无效调参。

6. 设计取舍：TMSA 与 AdaLN 的边界，以及怎么验证你选对了

把时间条件放进注意力内部，和把条件放进归一化层，是两条不同的路。以 DiT 为代表的方案用自适应 LayerNorm 做条件注入：时间嵌入生成 shift 和 scale 两组向量，作用在 LayerNorm 的输出上。参数量极小，实现简单，条件信号作用在特征归一化的位置，对整层特征做统一调制。DiffiT 选择在 QKV 上做双线性投影，代价是每层多出三组线性投影，参数和计算都更重。

这个取舍不该被理解成「谁更先进」。参数量敏感、层数很深的场景，AdaLN 的性价比更高；需要精细控制注意力行为、生成结构复杂的场景，TMSA 的表达力更值这份参数。两者也不互斥——把 TMSA 和 AdaLN 叠在一起，让归一化层管特征尺度、注意力层管空间关系，是我个人比较看好的组合方向，代价是调参空间变大，得小心两者互相抢功。

验证自己有没有选对，我习惯做两个消融。第一个是冻结时间：训练好的模型，在推理时把所有时间令牌替换成同一个常数，看 FID 掉多少。掉得越多，说明时间条件的利用率越高；如果几乎不掉，那这套时间机制在你的任务上就是摆设。第二个是打乱时间：把相邻时间步的时间令牌互换，观察生成过程是否出现结构错乱——如果打乱后结果基本正常，说明网络其实只在依赖少数几个时间点，中间段的精细调制没学到位。这两个实验都不需要重训，跑几次采样就有结论。

顺着这条线还能往下走：时间条件是否需要按分辨率分层，让不同层级看到不同粒度的时间表示；或者把时间令牌换成连续的时间嵌入而非离散索引，让采样步数变化时不必重新对齐。这些问题在原论文里没有展开，但都是这套机制自然引出的下一步。

说到底，我最初也踩过「时间嵌入加上了就算完事」的坑，模型能训、指标能看，但换个时间步采样就发现结果几乎一模一样，那种感觉很像发现黑匣子里根本没接线的后悔药。从那以后我每次做条件生成，都强制走一遍「固定输入、只变条件、看输出是否真的变」的检查，再谈调参。希望帮到你。

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