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-ExpertsMoEGemma 4 架构推理时实际激活的参数量约为 3.8B 到 4B。稀疏 MoE 设计在保留深度推理能力的同时量化后能压进高端消费级 GPU 的 18GB VRAM 限制内。模型采用编码器-解码器encoder-decoder架构两个组件各司其职自回归编码器Autoregressive Encoder用因果注意力causal attention处理初始 Prompt 上下文结果缓存到 KV Cache双向去噪器Bidirectional Denoiser则在生成画布上应用双向注意力bidirectional attention256 Token 画布上的每个位置都能同时关注其他所有位置以及 KV Cache 里的历史上下文。对于超过 256 Token 的长序列模型引入了块自回归扩散Block Autoregressive Diffusion机制——一个块完成去噪后提交到 KV Cache下一个画布以已生成的历史为条件重新开始。并行扩散的速度和长文本顺序生成的稳定性通过这种混合方式结合起来。离散文本扩散的运作机制生成过程是一个迭代去噪循环和 AI 图像生成器的思路类似分三个阶段推进画布初始化模型用随机占位 Token 填满 256 Token 的块。迭代精炼经过多轮去噪高置信度的 Token 率先确定作为全局上下文指导剩余占位符的精炼。收敛Token 逐渐收敛序列最终成形为连贯文本。整块一次性评估带来了一个自回归模型没有的能力自我修正。某次传播中对某个位置的置信度下降采样器可以对该位置重新加噪并替换。自回归模型一旦采样了某个 Token 就无法更改这是两种架构之间最实质的技术差异。推荐的部署与优化方案Google 给出了具体的部署配置建议以平衡延迟与输出质量。模型的实际表现对去噪采样器denoising sampler的调参方式高度敏感模型已针对 NVIDIA Blackwell 和 Hopper 架构上的 NVFP44 位浮点进行优化在接近无损精度的前提下进一步提升计算吞吐量。性能数据与实验结果DiffusionGemma 最直观的指标是原始生成速度。实验室专用硬件测试结果如下NVIDIA H100超过 1,000 Token/秒。NVIDIA GeForce RTX 5090超过 700 Token/秒。与同等硬件上的传统模型相比Token 生成速度最高提升 4 倍。不过这些数字最有说服力的场景是低并发本地工作流。在请求已批量处理、GPU 算力本就饱和的高流量云环境中并行解码的优势会递减部署成本也可能上升。案例展示数独求解实验研究人员用数独来测试双向上下文的实际价值——这是一个对自回归模型出了名的难题。数独网格里每个数字都受水平、垂直、宫格三重约束的交叉限制必须能评估后续位置并在发现冲突时回溯顺序生成的架构在结构上就处于劣势。实验结果基础模型表现未经专项训练的 DiffusionGemma 基础模型成功率约为 0%。Fine-tuningSFT后的表现使用 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

相关新闻