DiffusionGemma: The Developer Guide

DiffusionGemma：开发者指南 - Google Developers Blog

Google Tag Manager (noscript)

结束 Google Tag Manager (noscript)

HTML

DiffusionGemma：开发者指南

2026年6月10日

Ian Ballantyne

高级开发者关系工程师

Omar Sanseviero

技术团队成员

分享

• Facebook

• Twitter

• LinkedIn

• 邮件

在我们发布文章之后，我们发布了这份开发者指南，帮助您理解、使用并自定义这个实验性模型。

基于Gemma 4的架构，DiffusionGemma为开发者工作流程引入了几个里程碑：

• 计算受限的并行生成：通过将瓶颈从内存带宽转移到计算，绕过内存带宽限制，使GPU上的令牌生成速度提高高达4倍（在NVIDIA GeForce RTX 5090上每秒生成700多个令牌，在单个NVIDIA H100上每秒生成1000多个令牌）。

• 双向上下文与自我修正：在生成过程中，使用双向注意力同时评估整个文本块，实现实时错误修正和并行上下文传播。

• 开发者友好的模型大小：设计为26B专家混合（MoE）模型，在推理过程中仅激活3.8B参数，允许在18GB显存限制内进行量化部署。

架构

对于使用传统LLM在GPU上构建的开发者来说，主要瓶颈是内存带宽。自回归语言模型必须反复从内存中加载模型权重，以逐个令牌生成文本。DiffusionGemma通过将瓶颈从内存带宽转移到计算，绕过这一限制，同时生成并优化一个256个令牌的画布。通过为GPU提供一个大规模的并行工作负载，它利用了在本地服务期间通常处于空闲状态的张量核心。

• 均匀状态扩散：DiffusionGemma不按顺序预测令牌，而是从一个随机占位符令牌的画布开始，并行迭代地优化它们。经过多次去噪过程，高度自信的令牌有助于解决相邻位置，使整个序列清晰聚焦。

• 可变长度生成的块自回归扩散：对于超过256个令牌的序列，一旦一个256个令牌的块完全去噪，模型会处理并将其提交到KV缓存。然后模型转移到下一个块，初始化一个以先前提交的历史为条件的新鲜256个令牌画布。这结合了并行块的速度和自回归模型的顺序稳定性。

展示：使用并行去噪解决数独

传统的自回归模型在像数独这样的严格、多变量约束问题上存在困难。因为它们严格从左到右生成文本，无法评估未来的占位符或回溯。

为了展示DiffusionGemma的自定义功能，我们发布了使用Hackable Diffusion（一个模块化的JAX研究工具箱）的微调配方和结果。这个训练设置专注于一个经典的多变量网格任务：数独求解器。

为什么数独对扩散来说很有趣

在一个81个字符的数独字符串表示中（其中空单元格用句点标记），每个数字都受到严格的水平、垂直和9x9网格约束。

双向上下文传播：与自回归模型不同，DiffusionGemma 的去噪步骤允许每个画布查询同时关注所有位置。信息在棋盘上对称流动，每一步都能解决全局依赖关系。

• 通过重新噪声进行错误纠正：在均匀状态扩散下，模型同时评估整个棋盘。如果置信度下降，采样器会用随机数字替换现有数字，从而实现持续的自我纠正。

• 高效的早期停止：在数独上的微调表明，适配器增强了早期停止的效果。经过SFT微调的模型比基础模型更快稳定，使引擎能够更早停止，从而减少延迟和计算成本。

抱歉，您的浏览器不支持此视频的播放。

左图：DiffusionGemma生成数独输出。基础模型在48步后仍无法解决数独。右图：经过微调（SFT）的DiffusionGemma在12步后解决了数独。由于自适应停止，它能够提前完成。

性能影响：虽然基础的DiffusionGemma模型并未专门训练以解决数独谜题（成功率约为0%），但通过对数独数据集应用简单的JAX SFT方法，正确率可提高到80%，同时减少了整体推理步骤的数量。

块自回归去噪

为了实现块自回归去噪，DiffusionGemma在推理过程中在逐步预填充和去噪之间交替进行：

• 预填充 / 逐步预填充（因果）：使用因果注意力来摄入提示上下文并写入KV缓存。该步骤首先运行一次以预填充初始上下文，然后每次处理一个块时运行一次，将每个最终确定的256个标记的画布追加到KV缓存中，然后再进行下一个画布的去噪。

• 去噪（双向）：使用双向注意力对画布进行迭代去噪。画布上任意位置的查询标记都可以关注所有其他画布标记（以及KV缓存），使模型能够双向处理上下文。

这种架构选择使得以下成为可能：

• 全局上下文感知：与只能“回看”的自回归（AR）模型不同，去噪器的双向注意力使画布上的每个标记都可以关注其他所有标记。这使得模型在解决非顺序问题（如数独）时更加有效，因为第一单元格中的数字必须满足最后一个单元格的约束。

• 自我纠正：由于模型对整个画布进行迭代优化，它可以“修正”早期的错误。如果在一次遍历过程中某个标记的置信度下降，采样器可以重新噪声并替换它。这是自回归模型所缺乏的能力，因为一旦生成一个标记，它们就“被困”在该标记中，尤其是在长输出序列中。

• 高效的长上下文扩展：这种“块自回归”方法使模型能够处理长序列。它结合了扩散模型在块上的并行速度和自回归模型在长文本上的已验证的顺序稳定性。

• 简化的部署：使用与Gemma 4 26B A4B模型相同的架构意味着开发人员只需实现一个去噪步骤，这使得将其集成到现有的服务框架（如vLLM）中更加容易。

服务DiffusionGemma

为了高效地服务这种实验性架构，我们与vLLM团队合作，将DiffusionGemma集成到vLLM中。这种集成使引擎能够高效地在批处理请求流中运行迭代并行去噪循环。

开发者可以使用 vLLM 的标准 OpenAI 兼容本地服务器直接部署 DiffusionGemma。

vllm serve google/diffusiongemma-26B-A4B-it \
  --max-model-len 262144 \
  --max-num-seqs 4 \
  --gpu-memory-utilization 0.85 \
  --attention-backend TRITON_ATTN \
  --generation-config vllm \
  --hf-overrides '{"diffusion_sampler": "entropy_bound", "diffusion_entropy_bound": 0.1}' \
  --diffusion-config '{"canvas_length": 256}' \
  --enable-chunked-prefill

Shell

已复制

今天开始使用

准备好探索非自回归文本生成的前沿吗？查看以下资源以了解更多：

• 下载权重：直接在 Hugging Face 上获取实验模型的权重（该模型在 Apache 2.0 许可证下发布）。

• 集成与学习：查看 DiffusionGemma 的视觉指南，了解基于文本的扩散模型的工作原理。在 Gemma 文档中了解更多关于 DiffusionGemma 的信息。

• 使用您喜欢的推理框架：使用 vLLM、Hugging Face Transformers、SGLang 和 MLX 高效运行模型。

• 适配与微调：为了快速实验，我们发布了使用 Hackable Diffusion 的官方训练配方。您还可以使用 Unsloth 或 NVIDIA NeMo 探索高效的微调方法。

• 按您的方式部署：使用 Model Garden 或 NVIDIA NIM 立即在 Google Cloud 上部署模型。该模型在从消费级 RTX 4090 和 5090 显卡到企业级 Hopper 和 Blackwell 服务器的整个硬件堆栈上进行了原生优化。

发布在：

• AI

• 公告

• 探索

上一篇

下一篇

相关文章

列表

AI

公告

探索

Gemma 4 12B：开发人员指南

2026 年 6 月 3 日

云

学习

介绍 Google Colab CLI

2026 年 6 月 5 日

导航点