DiffusionBench：面向生成式扩散 Transformer 的综合评估基准

背景

随着生成式 AI 技术的飞速发展，扩散模型（Diffusion Models）已成为图像生成的主流架构。然而，传统的评估体系正面临严峻挑战。长期以来，ImageNet 分类准确率是衡量模型性能的核心指标，但对于生成式扩散 Transformer（Diffusion Transformers, DiTs）而言，仅依赖 ImageNet 上的表现已远远不够。

为了应对这一现状，研究人员推出了 DiffusionBench。这是一个旨在对生成式扩散 Transformer 进行全方位、综合性评估的基准测试平台。该项目的核心理念在于：不仅关注模型在特定数据集上的生成质量，更强调评估轴度的多样性、指标的全面性以及对已发表方法的忠实复现。DiffusionBench 提供了一个统一的代码库，支持通过单一接口在不同生成任务（如 ImageNet 生成、文本到图像 T2I 等）之间进行训练和评估，从而为社区提供一个标准化、可比较的性能衡量标准。

核心内容

DiffusionBench 不仅仅是一个数据集或一组指标，它是一套完整的、开箱即用的训练与评估基础设施。以下是其核心功能与技术细节的深度解读：

1. 统一的基础设施与工作流程

DiffusionBench 构建了一个统一的代码库，旨在简化复杂模型的开发流程。它支持从预训练模型下载、数据准备到模型训练及离线评估的全链路操作。

• 依赖管理：推荐使用 uv 项目管理器进行依赖安装与环境同步。
• 数据准备：通过 scripts/prepare.py 脚本，用户可以灵活选择准备全部数据、仅 ImageNet 数据、仅 T2I 数据或仅评估所需数据。
• 预训练模型：支持从 Hugging Face Hub 下载预训练模型权重，并排除不必要的元数据文件以节省空间。

2. 两阶段训练范式

该基准测试采用了一种清晰的两阶段训练流程，分别针对重建器（Reconstruction）和生成器（Generation）进行优化：

• Stage 1：训练 RAE Tokenizer（重建自编码器）

  • 此阶段的目标是训练一个高效的 Tokenizer，用于将图像转换为潜在空间表示。
  • 使用 src/train_stage1.py 进行训练，支持多卡并行（--nproc_per_node=8）。
  • 配置项包括 BF16 精度混合、PyTorch 2.0 编译加速（--compile）以及 Weights & Biases (W&B) 日志记录。
  • 评估指标包括 rFID、PSNR、SSIM 和 LPIPS，用于衡量重建质量。

• Stage 2：训练扩散模型

  • 在 Stage 1 获得的潜在空间（VAE/RAE/Pixel space）上训练扩散模型。
  • 使用 src/train.py 进行训练，同样支持多卡并行和加速选项。
  • 在线评估：训练配置文件包含 eval: block，允许在训练过程中实时进行在线评估。
  • 离线评估：对于已发布的检查点，可以使用 sampling/ 目录下的配置进行独立的离线评估。这些配置自动嵌入 stage_2.ckpt 路径和评估时的引导（guidance）参数，确保权重自动加载。

3. 多样化的模型架构支持

DiffusionBench 在 Stage 2 中涵盖了四大类模型家族，且在 ImageNet 和 T2I 任务中保持架构一致性：

• VAE (11 种变体)
• RAE (6 种变体)
• REG (4 种变体)
• Pixel (3 种变体)

这种设计使得研究人员可以通过简单的路径更改，在不同任务间交换配置文件，极大地提高了实验的灵活性和可比性。

4. 全面的评估指标体系

DiffusionBench 引入了多维度的评估指标，超越了传统的 FID：

• 重建质量 (Stage 1)：rFID, PSNR, SSIM, LPIPS。
• 生成质量 (Stage 2)：FID, IS (Inception Score)。
• 语义与内容评估：GenEval, DPGBench 等，用于评估生成图像的内容准确性和语义一致性。

5. 配置与执行细节

• ImageNet 任务：提供两种主要配置模式：
  • CFG-off 基线：无分类器引导（Classifier-Free Guidance）的基准配置。
  • Per-model best-CFG 变体：针对每个模型优化的最佳引导强度配置（如 cfg<scale>-t0.0-0.9.yaml）。
• 环境设置：通过设置 EXPERIMENT_NAME、ENTITY、PROJECT 和 WANDB_KEY 环境变量，实现实验名称自定义和 W&B 日志追踪。

关键要点

• 超越 ImageNet：DiffusionBench 的推出标志着图像生成评估从单一的分类准确率转向多维度的生成质量、语义一致性和重建保真度。
• 统一接口：通过单一代码库支持 ImageNet 和 T2I 等多种生成任务，降低了复现不同扩散 Transformer 变体的门槛。
• 两阶段训练标准化：明确区分了 Tokenizer 训练（Stage 1）和扩散模型训练（Stage 2），并提供了标准化的训练和评估脚本。
• 广泛的模型覆盖：支持 VAE、RAE、REG 和 Pixel 四大类模型家族，涵盖了当前主流的扩散 Transformer 架构。
• 多维评估指标：不仅包含传统的 FID 和 IS，还引入了 GenEval、DPGBench 等更细粒度的语义和内容评估指标，以及 PSNR、SSIM 等重建质量指标。
• 开源与社区驱动：项目欢迎社区贡献，鼓励用户添加新的评估轴度、新指标以及复现已发表的方法，旨在共同完善这一基准测试。
• 技术栈现代化：基于 PyTorch 2.0 编译加速、BF16 精度支持以及 W&B 实验追踪，确保训练和评估的高效性与可复现性。

意义与影响

DiffusionBench 的发布对生成式 AI 社区具有深远的影响：

1. 标准化评估基准：在扩散 Transformer 架构百花齐放的今天，DiffusionBench 提供了一个统一的“竞技场”，使得不同研究团队的工作可以在相同的条件下进行比较，消除了因实现细节差异导致的评估偏差。
2. 加速模型迭代：通过提供预训练模型、标准配置和自动化评估脚本，研究人员可以快速验证新架构或新训练策略的有效性，从而加速模型迭代和创新。
3. 促进可复现性研究：该项目强调对已发表方法的忠实复现，有助于解决 AI 研究中普遍存在的复现危机，增强学术研究的透明度和可信度。
4. 推动多模态发展：支持 ImageNet 和 T2I 等多种任务，表明 DiffusionBench 不仅关注纯图像生成，也为未来扩展到视频、3D 等多模态生成任务奠定了坚实的基础。
5. 社区协作平台：通过开源代码和贡献指南，DiffusionBench 成为了一个活跃的社区协作平台，促进了学术界和工业界在生成模型评估领域的交流与合作。

总之，DiffusionBench 不仅是一个技术工具包，更是推动生成式扩散模型走向成熟、标准化和可比较的重要里程碑。