DiffiT

“让 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 的代码结构清晰，分为核心模型层和工具层：

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 版本环境下的依赖兼容性需自行解决。

行业意义：扩散模型进入 Transformer 时代

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 融合的最佳学习素材。