DiffusionGemma: 通过文本扩散实现4倍更快的文本生成

DiffusionGemma: 通过文本扩散实现4倍更快的文本生成

Google DeepMind 已发布 DiffusionGemma，这是一个实验性的 26B MoE 模型，利用文本扩散同时生成文本块，相比自回归模型在专用 GPU 上实现最高 4 倍更快的推理。

高速推理和硬件优化

DiffusionGemma 将解码瓶颈从内存带宽转移到计算，使其在本地低并发推理期间能够更高效地利用 GPU 硬件。与传统自回归模型相比，此方法可实现显著更高的令牌输出速度。

  • 性能基准： 该模型在单个 NVIDIA H100 上可达到每秒 1000+ 个令牌，在 NVIDIA GeForce RTX 5090 上可达到每秒 700+ 个令牌。
  • 硬件占用： 作为一个在推理过程中仅激活 3.8B 参数的 26B 专家混合（MoE）模型，DiffusionGemma 在量化后可以适应高端消费级 GPU 的 18GB VRAM 限制。
  • 企业和消费者支持： 该模型针对 NVIDIA 硬件进行了优化，包括使用 NVFP4（4 位浮点）内核的 Hopper 和 Blackwell 架构，以实现近乎无损精度的加速吞吐量。它与消费级 GPU 如 RTX 4090 和 5090 兼容。

技术架构：文本扩散 vs. 自回归生成

DiffusionGemma 与标准的“打字机”方法不同——该方法从左到右逐个生成令牌。相反，它的工作方式类似于“印刷机”，同时草拟整个 256 令牌段落。

扩散过程

与图像生成模型类似，DiffusionGemma 通过迭代细化过程生成文本：

  1. 画布： 过程从一个随机占位符令牌的画布开始。
  2. 迭代细化： 模型进行多次遍历，锁定正确的令牌并将其作为上下文来细化剩余的占位符。
  3. 最终润色： 文本收敛为最终的高质量输出。

双向注意力和非线性生成

由于模型并行生成 256 个令牌，每个令牌可以关注所有其他令牌。这种双向注意力为令牌依赖于未来上下文的非线性任务提供了明显优势，例如：

  • 行内编辑和代码填充。
  • 生成氨基酸序列或数学图。
  • 解数独谜题（自回归模型通常在这项任务上遇到困难）。

用例和生产权衡

尽管 DiffusionGemma 提供了极高的速度，但它在输出质量和部署环境方面涉及特定的权衡。

质量 vs. 速度

DiffusionGemma 在最大输出质量方面优先考虑速度和并行布局生成。对于需要最高可能质量的应用，Google DeepMind 推荐使用标准的 Gemma 4 模型。

本地 vs. 云部署

DiffusionGemma 专门针对本地和低并发推理进行了优化。在高 QPS（每秒查询数）云环境中，自回归模型更高效，因为它们可以批量处理数千个请求以饱和计算。在这种情况下，DiffusionGemma 的并行解码收益递减，可能会增加服务成本。

可用性和集成

DiffusionGemma 在 Apache 2.0 许可证下发布，可通过 Hugging Face 获得。它得到多种开发工具和框架的支持，包括：

  • 服务框架： MLX, vLLM（带 Red Hat 集成），以及 Hugging Face Transformers。
  • 微调工具： Hackable Diffusion（一个 JAX 工具箱），Unsloth，以及 NVIDIA NeMo。
  • 部署平台： Gemini Enterprise Agent Platform Model Garden 和 NVIDIA NIM。

Sources