DiffusionGemma：用离散文本扩散和双向注意力，把推理瓶颈从内存带宽转移到算力

DiffusionGemma 的发布在 AI 研究圈子里引发了相当多的讨论，原因在于它触及了 LLM 文本生成方式的根本性问题。

标准 Gemma 4 系列依赖自回归解码（autoregressive decoding），从前到后逐个预测 Token。DiffusionGemma 则是离散文本扩散（discrete text diffusion）的实验性模型，可以同时生成并精炼整个文本块，绕过了历史上制约本地 AI 性能的主要硬件瓶颈。

瓶颈的转移：从内存带宽到算力

理解 DiffusionGemma 得先看清传统自回归模型卡在哪里。

单 GPU 本地单用户运行时，标准 LLM 的瓶颈通常是内存带宽。生成每个 Token 都要从内存加载完整的权重矩阵，GPU 的张量核心在等待数据传输期间大量闲置，算力并没有被充分利用。

DiffusionGemma 的做法是把瓶颈从内存带宽挪到算力上。每次前向传播（forward pass）并行起草一个 256 Token 的画布（canvas），一次性给处理器提供更大的工作块，硬件的算术强度（arithmetic intensity）才算真正用上了。架构上的变化让模型从逐字打印的打字机，变成了一次性印出整块文字的印刷机。

架构基础

DiffusionGemma 基于 26B 混合专家模型（Mixture-of-Experts，MoE）Gemma 4 架构，推理时实际激活的参数量约为 3.8B 到 4B。稀疏 MoE 设计在保留深度推理能力的同时，量化后能压进高端消费级 GPU 的 18GB VRAM 限制内。

模型采用编码器-解码器（encoder-decoder）架构，两个组件各司其职：自回归编码器（Autoregressive Encoder）用因果注意力（causal attention）处理初始 Prompt 上下文，结果缓存到 KV Cache；双向去噪器（Bidirectional Denoiser）则在生成画布上应用双向注意力（bidirectional attention），256 Token 画布上的每个位置都能同时关注其他所有位置，以及 KV Cache 里的历史上下文。

对于超过 256 Token 的长序列，模型引入了块自回归扩散（Block Autoregressive Diffusion）机制——一个块完成去噪后提交到 KV Cache，下一个画布以已生成的历史为条件重新开始。并行扩散的速度和长文本顺序生成的稳定性，通过这种混合方式结合起来。

离散文本扩散的运作机制

生成过程是一个迭代去噪循环，和 AI 图像生成器的思路类似，分三个阶段推进：

1. 画布初始化：模型用随机占位 Token 填满 256 Token 的块。
2. 迭代精炼：经过多轮去噪，高置信度的 Token 率先确定，作为全局上下文指导剩余占位符的精炼。
3. 收敛：Token 逐渐收敛，序列最终成形为连贯文本。
   
   整块一次性评估带来了一个自回归模型没有的能力：自我修正。某次传播中对某个位置的置信度下降，采样器可以对该位置重新加噪并替换。自回归模型一旦采样了某个 Token 就无法更改，这是两种架构之间最实质的技术差异。

推荐的部署与优化方案

Google 给出了具体的部署配置建议，以平衡延迟与输出质量。模型的实际表现对去噪采样器（denoising sampler）的调参方式高度敏感：

模型已针对 NVIDIA Blackwell 和 Hopper 架构上的 NVFP4（4 位浮点）进行优化，在接近无损精度的前提下进一步提升计算吞吐量。

性能数据与实验结果

DiffusionGemma 最直观的指标是原始生成速度。实验室专用硬件测试结果如下：

• NVIDIA H100：超过 1,000 Token/秒。
• NVIDIA GeForce RTX 5090：超过 700 Token/秒。
• 与同等硬件上的传统模型相比，Token 生成速度最高提升 4 倍。

不过这些数字最有说服力的场景是低并发本地工作流。在请求已批量处理、GPU 算力本就饱和的高流量云环境中，并行解码的优势会递减，部署成本也可能上升。

案例展示：数独求解实验

研究人员用数独来测试双向上下文的实际价值——这是一个对自回归模型出了名的难题。数独网格里每个数字都受水平、垂直、宫格三重约束的交叉限制，必须能评估后续位置并在发现冲突时回溯，顺序生成的架构在结构上就处于劣势。

实验结果：

• 基础模型表现：未经专项训练的 DiffusionGemma 基础模型成功率约为 0%。
• Fine-tuning（SFT）后的表现：使用 Hackable Diffusion 工具箱应用监督微调（SFT）方案后，成功率升至 80%。
• 效率变化：Fine-tuning 后模型收敛更快，自适应提前停止（adaptive early stopping）机制将步数从最多 48 步压缩到 12 步。
  

总结

这是一个实验性发布，因为整体输出质量低于标准自回归 Gemma 4 模型，目前更适合速度敏感的交互式本地工作流，而非高要求的生产级文本生成。

适合落地的场景集中在三类：内联编辑与快速迭代（需要实时反馈的场合）；非线性文本结构，如代码填充（code infilling）、氨基酸序列、数学图（双向上下文在这类任务上有结构优势）；以及代码生成中实时渲染代码和闭合复杂 Markdown 格式。

DiffusionGemma 在技术层面证明了一件事：文本生成可以从打字机模式切换到印刷机模式。双向注意力加迭代并行去噪，在专用 GPU 上实现了推理速度 4 倍的提升。原始质量还追不上标准自回归模型，但解决非线性约束问题的潜力，以及面向本地用户的吞吐量，是两个值得持续关注的方向。

https://avoid.overfit.cn/post/c1e25b8c4b534244b7a10036678c8535

作者：Mayur Jain