DiffiT
NVlabs/DiffiT加载项目详情…
本应用为开源项目，仅供学习研究，请遵守其开源协议。
微信扫码咨询客服
POMOAI — 让开源 AI 即刻运行
加载项目详情…
本应用为开源项目，仅供学习研究，请遵守其开源协议。
“让 AI 画出一张完美照片，需要多少步？”
这个问题曾困扰了图像生成领域多年。GAN 时代，生成一张人脸可能需要数十亿次参数调优；CLIP 引导模型虽好，却在细节上难以把控。直到扩散模型（Diffusion Model）的出现——通过逐步去噪的方式，从纯噪声中“雕刻”出逼真图像，生成质量才有了质的飞跃。
但扩散模型也有自己的短板：传统基于 U-Net 的架构在捕捉长距离依赖关系时能力有限，尤其是在高分辨率图像生成场景下，往往需要堆叠大量参数才能勉强维持质量。
2023 年年年底，NVIDIA Research 团队发布了一个令人眼前一亮的模型——DiffiT（Diffusion Vision Transformers），并在 ECCV 2024（欧洲计算机视觉会议）正式发表。与此前 Meta 的 DiT（Diffusion Transformer）一脉相承，DiffiT 进一步将 Vision Transformer（ViT）架构引入扩散模型的去噪核心，并创新性地提出了时序依赖多头自注意力机制（Time-dependent Multihead Self Attention，TMSA），让模型在每个去噪时间步都能自适应地调整注意力分布，从而更精准地控制图像细节的生成过程。
该模型最令人惊叹之处在于参数效率：DiffiT 相比同期 Transformer 扩散模型 MDT 参数量减少 19.85%，相比 Meta 的 DiT 减少 16.88%，却在生成质量上实现了全面超越。这说明 Vision Transformer 的全局建模能力与扩散模型的去噪范式天然契合，并非简单的堆参数换性能。
图1：DiffiT 在 ImageNet-256 上 FID 得分高达 1.73，刷新当时 SOTA 纪录。数据来源：项目 README。
DiffiT 提供两种工作模式，分别对应不同的生成分辨率和使用场景：
ImageNet-256 模式：生成 256×256 分辨率的高质量图像，在 ImageNet 数据集上达到 FID 1.73、Inception Score 276.49 的惊人成绩。这意味着模型能够根据给定的类别标签（如“金毛犬”、“钢琴”），从纯噪声中生成高度逼真、语义准确的照片级图像。
ImageNet-512 模式：扩展到 512×512 高分辨率，FID 2.67、IS 252.12，同样是当时的 SOTA 水平。高分辨率生成对显存和算力要求更高，但 DiffiT 的 Transformer 架构保证了长距离细节（如毛发、建筑纹理）的连贯性。
模型使用分类器无关引导（Classifier-Free Guidance，CFG）技术增强条件生成质量，256 场景下 guidance scale 为 4.4，512 场景下为 1.49。此外还采用 latent diffusion 范式，通过预训练的 VAE 将图像压缩到 latent 空间进行去噪，最后解码回像素空间，显著降低计算开销。
推理时，每张图像需要 250 步去噪迭代，在多块 NVIDIA A100/H100 GPU 上通过 NCCL 分布式通信并行加速。代码库提供了 Slurm 脚本 slurm_sample_256.sh 和 slurm_sample_512.sh，方便在 HPC 集群上一键提交 5 万张图像的批量采样任务，用于标准 FID 评估。
DiffiT 的代码结构清晰，分为核心模型层和工具层：
| 模块 | 文件 | 职责 |
|---|---|---|
| 扩散核 | diffit/diffit.py | Time-dependent MSA（TMSA）实现 |
| 扩散流程 | diffit/gaussian_diffusion.py | DDPM/DDIM 去噪调度器 |
| 模型基座 | diffit/nn.py | Transformer 块、LayerNorm、残差连接 |
| 位置编码 | diffit/pos_emb.py | 可学习的时序+空间位置编码 |
| VAE 解码 | sample.py | 集成 SD-VAE |
| 分布式 | diffit/dist_util.py | NCCL 多卡初始化与梯度同步 |
| 评估工具 | evaluator.py | TensorFlow Inception V3 FID 计算 |
TMSA（时序依赖多头自注意力） 是 DiffiT 区别于 DiT 的核心创新。在传统 ViT 中，自注意力的 Query-Key-Value 投影是静态的，与输入时间步无关；而 DiffiT 额外引入了与去噪时间步 t 相关联的调制（modulation）信号，使得注意力头的权重分布随 t 动态变化。这相当于让模型在不同“去噪阶段”自动调整关注重点——早期关注整体构图，后期聚焦局部纹理细节。
从依赖库来看，项目深度依赖 timm（PyTorch Image Models）提供预训练的 ViT 骨干网络，diffusers 提供 VAE 和扩散调度器封装，accelerate 管理分布式训练与混合精度，mpi4py 和 blobfile 负责 HPC 场景下的并行文件 IO 和云存储访问。
DiffiT 不提供 Web UI，所有操作通过命令行完成，这对非 ML 研究者来说有一定门槛。但对于 AI 开发者而言，这正是最纯粹、最可控的使用方式。
完整推理流程（以 ImageNet-256 为例）：
# 安装依赖
pip install accelerate timm diffusers tqdm mpi4py blobfile
# 下载预训练模型（需手动从 HuggingFace 下载）
MODEL=path/to/diffit_256.safetensors
LOG_DIR=./samples
# 运行采样
python sample.py \\
--log_dir $LOG_DIR \\
--cfg_scale 4.4 \\
--model_path $MODEL \\
--image_size 256 \\
--model Diffit \\
--num_sampling_steps 250 \\
--num_samples 50000 \\
--cfg_cond True
# 计算 FID
bash eval_run.sh
生成 5 万张 256×256 图像会占用大量磁盘空间（约 50GB+），建议使用高速 SSD。分布式推理依赖 SLURM 调度器，普通用户可在单机多卡环境下运行，但需手动调整 WORLD_SIZE 和 RANK 环境变量。
DiffiT 的定位是学术研究工具，存在几个显著限制：
1. 算力门槛高：单张 256×256 图像推理约需 8-12GB VRAM（FP16），批量生成 5 万张图像在多卡环境下仍需数小时。没有 NVIDIA GPU 的开发者几乎无法使用。
2. 缺乏灵活的条件控制：目前仅支持 ImageNet 类别标签（1000 类）作为条件，无法通过文本描述或草图引导生成图像。相比之下，Stable Diffusion 等产品化工具在控制粒度上更为丰富。
3. 预训练模型需额外下载：GitHub 代码库不包含模型权重，用户需手动从 HuggingFace 下载 .safetensors 文件。
4. 无 Docker 支持：代码未提供容器化方案，在不同 CUDA 版本环境下的依赖兼容性需自行解决。
DiffiT 的意义不仅在于刷新了 FID 纪录，更在于它验证了一个核心假设：Vision Transformer 是扩散模型的最佳拍档。
在此之前，扩散模型主要依赖 U-Net 架构（CNN-based），而 DiT 的出现首次证明了 Transformer 在这一领域同样有效。DiffiT 则在 DiT 的基础上更进一步，用 TMSA 机制解决了 Transformer 在处理不同去噪阶段时缺乏动态适应性的问题。
从更宏观的视角看，DiffiT 与 DiT、MDT、SiT 等工作共同标志着图像生成领域从“CNN 主导”向“Transformer 主导”的技术范式迁移。NVIDIA Research 通过 DiffiT 再次证明了其在生成式 AI 领域的前沿地位。对于 AI 研究者和工程师而言，DiffiT 的开源代码是深入理解 Transformer + Diffusion 融合的最佳学习素材。