DiffusionGemma: 通过并行扩散实现 4 倍速文本生成

DiffusionGemma 通过将顺序 Token 生成替换为并行文本扩散，实现了 4 倍速的推理速度

DiffusionGemma 是一款实验性的开源权重模型，旨在消除传统自回归大语言模型 (LLM) 的延迟瓶颈。通过同时生成整个文本块而非逐个 Token 生成，它在专用 GPU 上实现了高达 4 倍的文本生成速度提升。该模型基于 Apache 2.0 协议发布，采用 26B Mixture of Experts (MoE) 架构，在推理时仅激活 3.8B 参数，使其在量化后能够适配高端消费级 GPU 的 18GB VRAM 限制。

解决内存带宽瓶颈

传统的 LLM 运行方式类似于“打字机”，从左到右顺序生成 Token。虽然这种方式对于可以进行批处理请求的高并发云端服务非常高效，但对于本地、单用户推理而言却效率低下。在本地环境中，GPU 往往处于利用不足的状态，因为系统将权重从 RAM 移动到处理器的耗时比实际执行计算的时间还要长。

DiffusionGemma 将瓶颈从内存带宽转移到了计算能力。它不再是预测单个下一个 Token，而是同时起草一个包含 256 个 Token 的段落。这种“印刷机”式的方法使硬件的计算能力达到饱和，从而在本地加速器上实现了显著的速度提升：

  • NVIDIA H100: 1000+ tokens per second.
  • NVIDIA GeForce RTX 5090: 700+ tokens per second.

技术架构与能力

DiffusionGemma 集成了基于 Gemma 4 系列和 Gemini Diffusion 研究构建的新型扩散头 (diffusion head)。其运行逻辑类似于 AI 图像生成：它从一个包含随机占位符 Token 的画布开始，通过多次迭代优化，直到文本收敛为最终输出。

关键技术优势

  • 双向注意力机制 (Bi-directional Attention): 由于 256 个 Token 是并行生成的，每个 Token 都可以关注到块内的所有其他 Token。这使得该模型特别适合非线性任务，例如代码填充 (code infilling)、行内编辑、数学图表和氨基酸序列。
  • 智能自我修正: 模型会一次性评估整个文本块，从而允许它在优化过程中实时迭代修复错误。
  • 硬件优化: 模型支持 NVFP4 (4-bit floating-point) 内核，这可以在 Hopper 和 Blackwell 架构上以接近无损的精度加速计算吞吐量。

权衡与生产环境使用

DiffusionGemma 优先考虑速度和并行布局，而非原始输出质量。其整体质量低于标准的自回归 Gemma 4 模型。因此，Google 建议在需要最高质量的应用中使用标准的 Gemma 4，而 DiffusionGemma 则适用于对速度要求极高、需要交互式的本地工作流。

开发者生态与集成

DiffusionGemma 旨在进行快速实验，并与多个主流 AI 框架兼容：

  • 推理服务 (Serving): 支持通过 MLX、vLLM (包含 Red Hat 集成) 和 Hugging Face Transformers。对 llama.cpp 的官方支持尚在等待中。
  • 微调 (Fine-tuning): 开发者可以使用 Hackable Diffusion JAX 工具箱、Unsloth 或 NVIDIA NeMo。一个显著的微调案例是使用 Unsloth 使模型能够解决 Sudoku，在这种任务中，双向注意力机制比自回归模型具有明显优势。
  • 部署 (Deployment): 可通过 Hugging Face、Gemini Enterprise Agent Platform Model Garden 和 NVIDIA NIM 获取。

社区见解与分析

开发者之间的技术讨论强调了扩散模型通过减少提示词与响应之间的等待时间，从而改变“结对编程”体验的潜力。

"It was more of a pair-programming experience instead of the SOTA agentic experience of prompting and waiting... It felt less of a slot machine where you prompt, wait, and hope it went in the right direction." — @vineyardmike

其他贡献者强调了边缘设备上内存带宽问题的关键性：

"On edge you have a different problem: your inference accelerator is starved while sloshing GB of weights back and forth from RAM... Diffusion can compute tokens in parallel which relieves the memory bandwidth bottle neck." — @samuelknight

关于模型在特定领域（如工具调用和推理）的表现仍存疑问。一些用户建议，该模型重新编辑先前行的能力使其更适合“变更流 (change streams)”或跨多个文件的系列编辑操作，而非单一的工具调用。

Sources