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-Net：DiffiT 这份源码包到底能跑出什么

扩散模型生成图像，大家默认的骨架几乎都是卷积残差 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_qt
• k_s = x_s · W_ks + x_t · W_kt
• v_s = x_s · W_vs + x_t · W_vt

其中 x_s 是二维网格上排布的空间 token，x_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, bias=False) # 时间 token 的 QKV 投影，与空间投影维度对齐 self.to_qkv_time = nn.Linear(dim, inner * 3, bias=False) self.to_out = nn.Linear(inner, dim) def forward(self, x, t_token, pos_bias=None): # 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_ratio=4.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_bias=None): 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一般取 64，heads按 dim/64 推算，mlp_ratio取 4 是 Transformer 的常规默认，DiffiT 论文没有在这块做特殊设计，照搬即可。

2.3 时间 token 的生成与注入位置

时间 token 不是每层重新算，而是在网络入口算一次，然后传给所有 DiffiT Block。常见做法是把连续时间步 t 先做正弦位置编码，再过一个小 MLP：

class TimeToken(nn.Module): def __init__(self, dim, freq_dim=256): 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, device=t.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 及以下用 TMSA，32×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 通道而模型dim=256，入口要有一个nn.Conv2d(4, 256, 1)或nn.Linear做升维。第二，潜空间的噪声调度和像素空间不同，VAE 的缩放因子（常见 0.18215 这类常数）要在加噪前乘进去，否则训练时模型看到的噪声方差和采样时不一致，生成结果会偏灰或过饱和。这两点在源码里通常写在数据预处理和noise_scheduler的初始化处，复现时先翻这两处。

# latent 入口对齐示意 class LatentPatchEmbed(nn.Module): def __init__(self, in_ch=4, dim=256, patch=2): super().__init__() self.proj = nn.Conv2d(in_ch, dim, kernel_size=patch, stride=patch) 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)

patch=2把 32×32 潜特征切成 16×16 个 patch，token 数 256，注意力完全扛得住。如果 patch 取 1，token 数 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],), device=device) 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常见 1000，t用均匀采样即可，论文的 λ(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 一致。对比表大致长这样：

偏差超过 15% 就值得停下来查，通常问题出在时间 token 的数值范围——训练时 t 是整数张量转 float，采样时如果传的是归一化到 [0,1] 的 t，正弦编码出来的频率完全不同，注意力行为自然对不上。这个坑我在第一次接自定义采样器时踩过，loss 一切正常，采样图就是不对，查了两天才发现是 t 的量纲问题。

另一个实用技巧是给 TMSA 的注意力打分矩阵做可视化抽查：取去噪早期（t 接近 1000）和后期（t 接近 0）各一张，看注意力分布是否真的在变化。如果两个时间步的打分矩阵几乎一样，说明时间投影权重没起作用，多半是to_qkv_time的初始化太小或者被 weight decay 压没了。DiffiT 的价值就在这个时间依赖性上，这一项不成立，整个模型就退化成普通 DiT。

从那以后我每次复现带时间条件的扩散模型，都强制先跑一遍固定种子的逐层张量比对，再谈训练和调参，省下的返工时间远超这点准备工作。希望这套流程帮到你，把 DiffiT 这份源码真正跑成自己能改的骨架。

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