简介《DiffiT: Diffusion Vision Transformers for Image Generation》是NVIDIA团队关于扩散模型去噪网络架构的研究论文面向生成式AI研究者、计算机视觉开发者及相关方向学生帮助理解如何用视觉变换器改进传统卷积U-Net并提升图像生成质量。压缩包共1个PDF文件约43.22MB收录完整论文重点阐述时间依赖自注意力模块、混合层次U形编码器—解码器、latent DiffiT以及ImageNet-256实验。内容系统覆盖扩散模型、视觉变换器、去噪网络、FID评估、类条件与无条件合成等知识点并给出高分辨率图像生成结果其中latent DiffiT在ImageNet-256上取得1.73的FID分数。已有338人学习适合研读方法设计、实验设置与复现思路。1. 画质卡点的真正位置不是采样器是时间条件注入了多少很多人调 Diffusion 模型第一反应是换采样器、扫 guidance scale、堆步数再不行就怪数据。但把去噪网络拆开看会撞见一个不太舒服的事实同一个 3×3 卷积核在 t900 和 t50 两个时刻是同一组权重时间只以一个通道级的缩放平移量混进去。Vision Transformer 进入图像生成Image Generation这件事之所以值得认真对待不是因为注意力比卷积更“高级”而是它提供了另一条把时间塞进网络内部的路径——DiffiT 把时间 token 直接投影进 q、k、v让注意力本身随去噪阶段漂移。这一层改动的分量比换主干更值得深挖。2. 时间条件注入通道级仿射为什么撑不住整条去噪轨迹2.1 卷积去噪器的时间不敏感到底意味着什么扩散模型的前向过程可以写得很干净给定数据 z₀采样 z_t z₀ σ_t·εε 服从标准高斯σ_t 是一条单调递增的噪声日程。训练目标就是让一个带参数的噪声预测网络 ε_θ 在给定 z_t 和 t 的情况下把 ε 猜回来损失外面再乘一个随 t 变化的权重 λ(t) 做重加权。整个网络只有一份权重在 t 从 T 到 0 的整条轨迹上反复复用。问题就出在这个“复用”上。把常用的卷积残差 U-Net 展开看时间信息的流动路径先对 t 做正弦位置编码过一个小 MLP 得到时间向量然后在每个残差块里用 FiLM 式操作把时间向量变成逐通道的缩放与偏移作用在归一化的特征图上。卷积核 W 本身不参与这个过程。换句话说网络在不同时间步之间共享同一组空间算子改变的只是每个通道的增益和偏置。我一般会这么理解这件事的代价相当于一台相机在不同曝光条件下只调增益不换镜头。可扩散过程是有明确阶段分工的。t 很大时信噪比极低高频内容被完全扰动网络主要该预测低频与整体结构t 接近 0 时结构已经定下来网络该做的是抠高频细节。这两件事需要的感受野、频率响应、空间算子都不一样而卷积核不随时间变化就等于强行让同一套局部算子去兼任两个完全不同的任务。更麻烦的是注意力的存在让这个矛盾被放大了。传统 U-Net 只在低分辨率特征图上挂自注意力避免二次复杂度爆炸而在这些位置上注意力承担的恰恰是长程空间建模。长程依赖的形式在去噪早期拼全局构图和晚期修局部纹理连贯性应该是不一样的但 Qi、Ki、Vi 的投影矩阵却和时间无关。时间向量能做的只是间接地通过归一化层的缩放去微调投影后的尺度——这种影响非常绕也不具备方向性。2.2 AdaLN 的表达能力边界把 U-Net 主干整体换成 Transformer 的那一类方案把时间注入的方式也换了一套Adaptive LayerNorm。它把时间向量映射成三组逐通道参数 γ(t)、β(t)、α(t)然后写成“归一化、仿射、过 MLP、门控残差”的顺序。相比卷积网络里的 FiLMAdaLN 更彻底地让时间参与每一个 block 的计算而且把 shift 和 scale 都显式参数化了。但它仍然有一个结构性限制。把 AdaLN 之后接的那个线性投影展开就能看清归一化输出先逐通道乘上 (1γ)再加上 β然后统一送入同一组 W_q、W_k、W_v。乘性部分等价于对输入通道做一次对角重加权加性部分等价于在投影后得到一个对所有 token 都相同的常数偏移。前者会改变注意力的尺度可以理解为软性地调温度后者落到 softmax 里会被平移不变性吃掉一部分。于是结论比较清楚时间在这里能改写的是“每个通道的权重有多大”改不了“token 之间怎么交互”。注意力矩阵的形状主要由投影矩阵的行空间决定而 AdaLN 没有给投影矩阵本身引入任何随时间变化的自由度。对于一个需要在不同阶段切换空间建模策略的网络来说这个粒度太粗了。2.3 TMSA把时间 token 投影进 Q、K、VDiffiT 的核心动作可以一句话说完不让时间只作用于归一化后的激活而是把它作为一路 token和空间 token 一起送进线性投影直接参与 q、k、v 的构造。这个模块叫时间依赖多头自注意力TMSA。形式上有两套权重。空间侧有 W_qs、W_ks、W_vs时间侧有 W_qt、W_kt、W_vt最终每个空间 token 的查询向量是 x_s·W_qs x_t·W_qt键和值同理。时间 token x_t 由时间位置编码过一个小 MLP 得到并且在整条去噪网络的所有层之间共享——这一点很关键它保证了不同深度看到的“当前时刻”是同一个表示而不是每层各自重新编码一遍。这个式子有个很实用的等价解读把空间 token 和时间 token 在通道维拼起来再做一次线性投影数学上完全等价。也就是说TMSA 本质上是给每个空间 token 的表示里永久拼上了一段“时间上下文”然后让网络自己去学怎么用它。相比 AdaLN 那种“先归一化、再缩放偏移”的两段式这里时间与空间在同一个线性子空间里相加网络可以学到“时间偏置抵消空间偏置”这种补偿模式。注意力计算本身还保留了相对位置偏置 B加在 softmax 之前的 logits 上。这个细节后面会单独说因为它和时间项的关系处理不好是复现时最容易掉进去的坑之一。3. TMSA 源码级拆解一次前向里时间到底改了哪几个矩阵3.1 从两路 token 流到共享投影空间先把张量形状理清楚很多实现层面的错误都出在广播上。空间 token 是 (B, N, C)N 是当前分辨率的 token 数比如 16×16 的特征图展平成 256C 是通道数。时间 token 是 (B, C)在每个 block 里被扩成 (B, 1, C) 后参与计算。空间侧走一个Linear(C, 3C)输出拆成 q、k、v 三份再 reshape 成 (B, H, N, D)H 是头数D C/H。时间侧走另一个Linear(C, 3C)输出同样拆成三份reshape 成 (B, H, 1, D)。两个张量在 N 这一维上直接相加靠的是 PyTorch 的广播1 会被复制成 N等价于“把同一个时间偏置加到所有空间 token 上”。这一步的形状设计不是随意的。如果时间投影输出的是 (B, N, C) 而不是 (B, 1, C)那每个空间 token 拿到的时间信息就不同了网络会误以为同一时刻存在多个时间步训练很难收敛。反过来如果空间侧没保留 N 维而是错误地做了全局池化长程建模能力就直接废掉。3.2 一个最小可运行的 TMSA 实现下面这段是能直接跑的最小版本去掉了工程细节但保留了所有影响语义的结构。import torch import torch.nn as nn class TMSA(nn.Module): 时间依赖多头自注意力q/k/v 由空间 token 与时间 token 共同决定 def __init__(self, dim, num_heads, use_rel_pos_biasTrue): super().__init__() assert dim % num_heads 0 self.num_heads num_heads self.head_dim dim // num_heads self.scale self.head_dim ** -0.5 # 空间侧投影一次算出 qkv减少 kernel launch self.qkv_spatial nn.Linear(dim, dim * 3) # 时间侧投影独立一套权重输出维度必须与空间侧对齐 self.qkv_time nn.Linear(dim, dim * 3) self.proj nn.Linear(dim, dim) self.use_rel_pos_bias use_rel_pos_bias if use_rel_pos_bias: # 简化版每个头一份可学习标量偏移真实实现是按相对坐标建表 self.rel_pos_bias nn.Parameter(torch.zeros(num_heads, 1, 1)) def forward(self, x, t_token): B, N, C x.shape # 空间分支 - (3, B, H, N, D) qkv_s self.qkv_spatial(x).reshape(B, N, 3, self.num_heads, self.head_dim) qkv_s qkv_s.permute(2, 0, 3, 1, 4) # 时间分支 - (3, B, H, 1, D)随后靠广播加到所有空间 token 上 qkv_t self.qkv_time(t_token).reshape(B, 1, 3, self.num_heads, self.head_dim) qkv_t qkv_t.permute(2, 0, 3, 1, 4) q qkv_s[0] qkv_t[0] k qkv_s[1] qkv_t[1] v qkv_s[2] qkv_t[2] attn (q k.transpose(-2, -1)) * self.scale if self.use_rel_pos_bias: # 位置偏置只走空间侧不与时间项混在一起 attn attn self.rel_pos_bias attn attn.softmax(dim-1) out (attn v).transpose(1, 2).reshape(B, N, C) return self.proj(out)几个参数值得单独说明。dim是当前 stage 的通道数图像空间模型里会在 128 和 256 之间切换切换点就在下采样之后。num_heads决定每个头的维度头数不变时通道翻倍意味着单头维度翻倍注意力的表达粒度会变粗这也是为什么高层次低分辨率、大通道通常要相应调整头数。use_rel_pos_bias打开时会加一个相对位置项注意它加在 softmax 之前、缩放之后位置不能挪。整个 block 的残差结构很标准先TMSA(LN(x_s), x_t) x_s再MLP(LN(x̂_s)) x̂_s。这里没有 AdaLN 的 shift/scale 门控时间信息完全通过 TMSA 这一条路径进入。这个设计选择的直接后果是MLP 分支本身对时间不敏感只处理 TMSA 汇总后的空间信息。好处是参数省、路径清晰代价是如果某一层几乎不需要时间调制也没有开关可以关掉它。3.3 一个容易被忽略的数学性质K 上的时间偏置去哪了把 q、k 的表达式代进注意力 logits 里会看到一个挺有意思的现象。设空间项为 q_s、k_s时间项为 q_t、k_t都是 D 维向量且对所有 token 相同那么logits (q_s q_t)·(k_s k_t)ᵀ q_s·k_sᵀ q_s·k_tᵀ q_t·k_sᵀ q_t·k_tᵀ第二项 q_s·k_tᵀ 对每个 query 是一个标量常数第四项 q_t·k_tᵀ 是全局常数。这两项在 softmax 里会被平移不变性整块吃掉。真正改变注意力分布的只有第三项 q_t·k_sᵀ它给每个 key 加上了一个由时间决定、由 key 内容调制的偏置。也就是说k 上的时间偏置在数学上是冗余的。这不是实现 bug而是 softmax 的固有性质。理解这一点有两个直接用处一是复现时如果只给 k 加时间项、q 和 v 不加会发现训练曲线和基线几乎重合然后开始怀疑人生二是反过来如果想让时间条件更“锋利”把自由度集中在 q 和 v 上就够了。v 上的时间偏置则是另一回事。它相当于给每个位置的输出加一个由 t 决定的常数向量再经过输出投影等价于一个时间相关的输出偏置项。这部分是实打实生效的而且作用在值空间上影响的是“取回什么信息”不是“往哪看”。所以 TMSA 实际的有效自由度是 q 和 v 两路加 k 更多是为了形式对称和参数共享的便利。这个判断我在复现时验证过改动 k 的时间分支对最终 FID 的影响在噪声范围内。4. U 形混合主干图像空间与潜空间为什么用两套拓扑4.1 收缩-扩张路径上的分辨率与通道预算图像空间的 DiffiT 没有走纯 Transformer 路线。原因很实际在 64×64 甚至更高分辨率上直接做全局注意力token 数是 4096 起步二次复杂度顶不住。所以它保留了 U 形的骨架只在特定位置放 Transformer block。整体的形状变化可以列成一张表看得更清楚阶段特征图分辨率通道数主要操作Stem64×641283×3 卷积做 tokenizerStage 164×64128DiffiT ResBlock × L1Down 132×32256卷积下采样stride 2Stage 232×32256DiffiT ResBlock × L2Down 216×16256卷积下采样Stage 316×16256DiffiT ResBlock × L3Down 38×8256卷积下采样Stage 48×8256DiffiT ResBlock × L4Up ×316→32→64256→128卷积上采样 skip 拼接Head64×6433×3 卷积输出 RGB这张表里有几个值得琢磨的取舍。第一通道数在 Stage 1 之后就不再增长四个 stage 全是 256。这和常规 U-Net 越深通道越多的做法不同原因是 Transformer block 的参数量和通道数的平方成正比再往上堆会迅速吃掉显存。第二下采样和上采样用的是卷积而不是 patch merge这让局部归纳偏置在高分辨率段仍然存在减轻了小数据上的过拟合。第三每个分辨率都保留了 skip connection这是 U 形结构在扩散模型里被反复验证有效的原因早期层的低频信息可以绕过深层直接送到解码端避免在多次下采样中丢失。注意 Stage 4 只有 8×8token 数 64这才是全局注意力真正便宜的区间。但 DiffiT 并没有把注意力只放在这里——它在每个 stage 的 ResBlock 内部都嵌了 TMSA只是高分辨率段的 token 数大实际计算量占比更高。这个设计意味着高分辨率段的注意力实际上是一种“局部窗口之外还要看全局”的妥协靠 U 形结构的层次性来分摊成本。4.2 DiffiT ResBlock卷积与注意力的拼装顺序单个 ResBlock 内部的顺序不是随便排的。它先走一段“Swish 激活 → GroupNorm → 3×3 卷积”在这里注入时间嵌入然后接一个 DiffiT Transformer内部就是前面那个 TMSA block最后再走一次“GroupNorm → 3×3 卷积”和输入做残差相加。class DiffiTResBlock(nn.Module): def __init__(self, dim, num_heads): super().__init__() self.norm1 nn.GroupNorm(32, dim) self.conv1 nn.Conv2d(dim, dim, 3, padding1) self.time_mlp nn.Linear(dim, dim) # 时间嵌入 - 通道偏置 self.attn TMSA(dim, num_heads) # 注意 TMSA 在 conv 之后 self.norm2 nn.GroupNorm(32, dim) self.conv2 nn.Conv2d(dim, dim, 3, padding1) self.act nn.SiLU() def forward(self, x, t_emb): h self.conv1(self.act(self.norm1(x))) # 时间嵌入在这里以逐通道加性方式进入卷积分支 h h self.time_mlp(t_emb)[:, :, None, None] # Transformer 分支先把 2D 特征展平成 token 序列 B, C, H, W h.shape tokens h.flatten(2).transpose(1, 2) tokens tokens self.attn(tokens, t_emb) h tokens.transpose(1, 2).reshape(B, C, H, W) h self.conv2(self.act(self.norm2(h))) return x h这个顺序的含义是先由卷积把局部纹理整理成 token 表示再由注意力做全局混合最后由卷积把混合结果还原回空间特征。反过来排先注意力再卷积在早期实验里通常更差因为原始像素级噪声进注意力会让 softmax 分布过于平坦梯度信噪比低。还有一点时间嵌入在卷积分支里仍然保留了加性注入所以这个结构其实是“卷积侧用通道偏置、注意力侧用 token 投影”的双通道时间注入。这和论文里“时间依赖只集成到自注意力层”的表述看起来有出入实际上指的是“注意力的投影权重随时间变化”这件事只发生在 TMSA 里卷积侧那个加性偏置是通道级的、不改变卷积核属于常规做法。4.3 潜空间版本patch 化之后 TMSA 才真正发挥图像空间的模型受限于分辨率和显存注意力只能在有限的位置发挥作用。潜空间版本把这个约束解开了先用一个自编码器把图像压到低维潜变量再在潜空间上做扩散。此时同样一幅 256×256 的图潜变量可能只有 32×32token 数降到原来的 1/64纯 Transformer 主干才跑得动。潜空间 DiffiT 的结构就更接近标准的 ViT 了潜变量切 patch、线性投影、加上位置编码然后堆叠若干个带 TMSA 的 transformer block最后线性投影回 patch 维度、反 patch 化得到预测的噪声。类别条件通过额外的 embedding 加进去和扩散时间步的 embedding 一起构成条件信号。这里有个容易忽略的点潜空间模型里没有卷积下采样所有 stage 的分辨率相同因此 TMSA 在每一层看到的 token 数一致相对位置偏置的表结构也可以复用。相比图像空间版本要在四个分辨率上分别维护位置编码潜空间版本的实现要干净得多这也是复现时建议先从潜空间版本入手的原因。5. 复现排查时间依赖自注意力最容易翻车的五个地方5.1 只给 K 加时间项训练曲线纹丝不动现象把时间投影只接到 k 上q 和 v 保持纯空间投影训练损失和 FID 与不加时间条件的基线几乎完全重合跑几千步也看不出差异。原因softmax 对 logits 的整体平移不变。k 上叠加一个对所有 token 相同的时间向量后logits 里多出来的是 q_s·k_tᵀ 这一项它对每个 query 都是标量常数被 softmax 直接消掉。真正有效的是 q_t·k_sᵀ也就是时间项必须作用在 query 侧才能改变注意力分布。解决确认三路里至少 q 和 v 拿到了时间项。如果为了省参数想砍掉一路砍 k不要砍 q。5.2 时间 token 的广播维度写错静默算出错误结果现象代码不报错loss 能降但生成样本在不同时间步之间几乎没差别看起来像“时间条件没生效”。原因时间投影输出被 reshape 成了 (B, N, C) 而不是 (B, 1, C)广播时按 token 逐位对齐等于给每个空间位置喂了一个不同的“时间步”。更隐蔽的一种是把时间 token 在 batch 维上错误 broadcast导致同一批里所有样本共享了第一个样本的时间嵌入。解决在 forward 里加形状断言明确要求时间分支输出维度是 (B, 1, C) 或 (B, H, 1, D)。跑一次小 batch 的单元测试把 t 固定成两个差异极大的值看输出差异是否显著不显著就说明广播有问题。5.3 时间嵌入尺度没对齐训练早期出现 loss 尖峰现象前几百步 loss 剧烈震荡偶尔出现 NaN把学习率调小也只能缓解。原因时间位置编码的数值范围通常比较大正弦函数的组合过 MLP 之后如果没做归一化就直接加进 q/k/v会和空间 token 的尺度严重不匹配。注意力 logits 被时间项主导softmax 饱和成 one-hot梯度几乎消失。混合精度下这个问题更明显。解决时间 MLP 输出后接一层 LayerNorm或者在初始化时把时间投影的权重缩小若干倍让训练初期时间项只是一个小扰动。同时检查 λ(t) 重加权项信噪比极低的那些时间步权重过大时也会放大这个问题。5.4 把位置偏置和时间嵌入混在一起算现象模型能训但收敛后的 FID 比预期差一截换更大的模型也没有明显改善。原因相对位置偏置需要编码的是“两个 token 在空间上隔多远”时间嵌入编码的是“现在是第几步”这两类信息混在同一个参数张量里会互相干扰。位置偏置通常对每个头独立而时间项对所有头共享同一个来源混在一起后网络很难学出干净的空间先验。解决把位置偏置作为独立的加性项只加在 softmax 之前的 logits 上不参与时间投影的任何线性组合。如果确实想让位置编码也带时间信息就做成两个独立的偏置张量相加让网络自己学权重。5.5 潜空间版本的缩放因子没对上loss 正常但 FID 崩现象潜空间模型训练损失下降得很漂亮采样出来的图却是灰蒙蒙的一片或者色彩严重偏移。原因自编码器的潜变量通常做过标准化扩散过程在这个标准化空间里进行。采样结束后需要按同样的因子反标准化再解码。如果训练时用了缩放、采样时忘了或者前后用了不同的因子像素空间的统计量就完全错位而损失是在潜空间算的根本发现不了。解决把缩放因子做成配置项训练和采样共用同一份配置不要在两处各写一遍。采样脚本里加一步断言检查解码前潜变量的均值方差是否接近训练时的统计量。6. 设计取舍时间条件到底该放在哪一层把三种方案摆在一起取舍就很清楚了方案时间进入的位置单层新增参数对注意力分布的影响实现复杂度通道级仿射归一化后的激活O(C)间接仅调尺度低交叉注意力额外的 key/value 序列O(C²)直接但需额外注意力中TMSAq/k/v 的线性投影O(C²)直接作用在 query 侧低从表达力看TMSA 和交叉注意力处在同一档都能直接改写注意力分布从成本看TMSA 只是把一次线性投影拆成两次再相加额外开销基本可以忽略而交叉注意力要多算一次注意力矩阵。这就是它划算的地方用接近通道级仿射的代价买到了接近交叉注意力的表达力。但它也有明确的代价。时间信息完全走注意力这一条路意味着 MLP 分支对时间不敏感如果某一层需要的是“整体增益随时间变化”这种简单调制TMSA 反而是绕远路。另外q 和 v 上的时间项对所有 token 共享是一种全局的低秩调制如果某个任务需要“不同空间位置有不同的时间响应”这套机制就表达不了——那需要的是逐 token 的时间条件参数会再涨一个量级。所以我在复现时会做两件事来验证设计是否真的生效。第一件是固定随机种子和初始噪声只改时间步看同一层的注意力图是否出现结构性变化如果只是整体亮度变了、注意力的相对分布没动说明时间项退化了。第二件是在 q 和 v 上分别做消融看 FID 各自掉多少通常 v 的贡献比 q 更稳定因为它是加性的、不经过 softmax。从那以后我每次给扩散模型加时间条件都会先画一张“时间信息进入网络的路径图”把每一条路径上的张量形状和数学作用标出来再动手写代码。翻车最惨的那几次都是路径图画不出来、靠试出来的。希望帮到你。本文还有配套的精品资源点击获取
