DiffusionGemma 深度解析：扩散语言模型如何打破自回归范式

2026 年 6 月 10 日，Google 正式发布并开源 DiffusionGemma。 这是 Google DeepMind 基于 Gemini Diffusion 研究成果打造的首个开源扩散语言模型，标志着扩散技术从图像生成正式进入文本生成领域。

DiffusionGemma 的核心突破在于：它打破了自回归模型对语言生成的垄断。 自 2018 年 GPT-2 以来，几乎所有主流语言模型（GPT、Claude、Gemini、Llama）都采用自回归方式——从左到右逐个 token 生成文本。这种方式质量高，但有一个致命缺陷：生成过程是严格串行的，无法并行化。生成 1000 个 token 需要 1000 次前向传播，推理速度受限。

DiffusionGemma 采用了完全不同的范式：离散扩散（Discrete Diffusion）。 它从一段全 [MASK] 的序列出发，通过多步去噪，逐步恢复出连贯文本。每一步可以同时更新多个 token，最终实现单次并行生成 256 个 token。

关键性能数据：

• 在 H100 GPU 上达到 1107 tok/s，约是 Gemma 4 的 4 倍
• 在消费级 RTX 5090 上达到 700+ tok/s，支持本地部署
• 基于 Gemma 4 26B MoE 架构，激活参数仅 4B
• 支持文本、图像多模态输入，文本输出
• Apache 2.0 开源协议，可商用

DiffusionGemma 基于 Gemma 4 26B MoE 架构改造而来。 它保留了 Gemma 4 的 Mixture-of-Experts（混合专家）结构，但将 Decoder 的因果掩码（Causal Mask）替换为双向注意力（Bi-directional Attention），并引入时间步嵌入（Timestep Embedding）来控制去噪过程。

核心架构参数：

• 总参数量：25.2B
• 激活参数量：3.8B（MoE 结构，每次推理只激活部分专家）
• 层数：30 层
• 滑动窗口：1024 tokens
• 上下文长度：最高 256K tokens
• Canvas 长度：256 tokens（单次生成的最大长度）
• 词汇表大小：262K
• 专家数量：128 个专家中激活 8 个，加 1 个共享专家
• 视觉编码器：约 550M 参数

与 Gemma 4 的关键区别：

1. 注意力机制：Gemma 4 使用因果注意力（只能看到左侧上下文），DiffusionGemma 使用双向注意力（可以看到完整上下文）
2. 生成方式：Gemma 4 逐 token 自回归生成，DiffusionGemma 并行去噪生成
3. 推理速度：DiffusionGemma 在 H100 上达到 1107 tok/s，Gemma 4 约 280 tok/s
4. 质量权衡：DiffusionGemma 在部分基准上略低于 Gemma 4（见下文基准对比），但速度快 4 倍

# 从全 [MASK] 序列出发，逐步去噪生成文本

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# 加载 DiffusionGemma 模型
model_name = "google/diffusiongemma-26B-A4B-it"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)

# 准备输入：全 [MASK] 序列（canvas length = 256）
canvas_length = 256
input_ids = torch.full((1, canvas_length), tokenizer.mask_token_id)

# 扩散去噪过程（简化示意）
num_denoising_steps = 20  # 实际步数取决于模型配置
for t in range(num_denoising_steps):
    # 模型预测每个位置的概率分布
    outputs = model(input_ids, timestep=t)
    
    # 根据置信度更新 token（实际使用更复杂的采样策略）
    predictions = outputs.logits.argmax(dim=-1)
    
    # 逐步更新低置信度位置
    confidence = outputs.confidence  # 模型输出的置信度分数
    update_mask = confidence < threshold
    input_ids[update_mask] = predictions[update_mask]

# 最终生成的文本
generated_text = tokenizer.decode(input_ids[0], skip_special_tokens=True)
print(generated_text)

DiffusionGemma 在速度上大幅领先，但在部分质量基准上略低于 Gemma 4。 这是扩散模型的典型权衡：用一定的质量损失换取数倍的推理速度提升。

关键基准对比（来自 Hugging Face 官方数据）：

关键观察：

1. 通用知识（MMLU Pro、MMMLU）：差距 4-5%，扩散模型已经接近自回归水平
2. 数学推理（AIME、MATH-Vision）：差距较大（11-19%），扩散模型在复杂推理上仍有劣势
3. 代码生成（LiveCodeBench、Codeforces）：差距 8-10%，可用但有提升空间
4. 罕见知识（HLE no tools）：DiffusionGemma 反超 Gemma 4（11.0% vs 8.7%），说明扩散模型在某些特殊场景可能具有优势
5. 视觉任务（MMMU Pro）：差距最大（19.5%），视觉能力还需要进一步优化

速度对比：

• H100 GPU：DiffusionGemma 1107 tok/s vs Gemma 4 ~280 tok/s（4 倍提升）
• RTX 5090：DiffusionGemma 700+ tok/s（支持本地部署）

扩散模型最初是为连续数据（图像、音频）设计的。将扩散模型用于离散的文本空间，最大的挑战在于：高斯噪声加在浮点向量上是有意义的，但加在离散的 token 上却没有自然定义。

一个 token 是词汇表中的索引（如 "hello"=2345），你不能给它加 0.5 的高斯噪声。那么 DiffusionGemma 是如何解决这个问题的？

DiffusionGemma 采用「吸收扩散」（Absorbing Diffusion）方案。 这是目前最成功的离散扩散方法，核心思想是：

前向过程（加噪）：

1. 从原始文本序列开始
2. 每一步以一定概率将 token 替换为 [MASK] 标记（而非替换为随机 token）
3. 随着时间步增加，越来越多的 token 被「吸收」为 [MASK]
4. 最终序列变成全 [MASK]

反向过程（去噪）：

1. 从全 [MASK] 序列开始
2. 模型预测每个 [MASK] 位置应该恢复为什么 token
3. 根据置信度选择部分位置更新（高置信度优先）
4. 重复直到所有 [MASK] 都被替换为真实 token

吸收扩散的优势：

• 可以直接复用掩码语言模型（如 BERT）的训练目标
• 不需要在离散空间定义复杂的转移矩阵
• 训练稳定，收敛快

关键创新：Entropy Bound (EB) 采样器
DiffusionGemma 引入了 Entropy Bound 采样策略，在每一步去噪时，不是盲目更新所有位置，而是根据模型输出的熵（不确定性）来决定更新顺序。低熵位置（模型确信）优先更新，高熵位置（模型不确定）留到后续步骤。这种策略显著提升了生成质量。

import torch

def absorbing_forward(x, t, T, vocab_size):
    """
    x: [batch, seq_len] token indices
    t: 当前时间步
    T: 总时间步数
    """
    mask_token = vocab_size  # [MASK] 的索引
    batch, seq_len = x.shape
    
    # 吸收概率随时间线性增长
    p_absorb = t / T  # t=0时不吸收, t=T时全部吸收
    
    # 为每个位置独立决定是否吸收
    absorb_mask = torch.rand(batch, seq_len) < p_absorb
    
    x_noisy = x.clone()
    x_noisy[absorb_mask] = mask_token
    
    return x_noisy, absorb_mask

# 示例：对一句话执行逐步吸收
sentence = torch.tensor([[10, 20, 30, 40, 50]])  # 5个token
for t in range(6):
    noisy, mask = absorbing_forward(sentence, t, 5, vocab_size=1000)
    print(f"t={t}: {noisy.tolist()}, 吸收率={mask.float().mean():.1f}")
    
# 输出示例：
# t=0: [[10, 20, 30, 40, 50]], 吸收率=0.0%
# t=1: [[1000, 20, 1000, 40, 50]], 吸收率=40.0%
# t=2: [[1000, 1000, 1000, 40, 1000]], 吸收率=80.0%
# t=5: [[1000, 1000, 1000, 1000, 1000]], 吸收率=100.0%

DiffusionGemma 的出现为语言模型开辟了新的应用场景。 虽然它在通用质量上还无法完全超越自回归模型，但在以下场景具有独特优势：

1. 实时交互系统
需要快速响应的场景（如聊天机器人、语音助手），DiffusionGemma 的 4 倍速度提升可以显著降低延迟，提升用户体验。

2. 代码补全与编辑
代码补全需要快速生成多个候选方案，扩散模型的并行生成能力可以同时产出多个补全建议，供开发者选择。

3. 本地部署与边缘设备
DiffusionGemma 在 RTX 5090 上达到 700+ tok/s，意味着可以在消费级 GPU 上实现高速推理。这对于隐私敏感场景（本地处理敏感数据）和边缘设备（笔记本电脑、移动设备）非常重要。

4. 多模态理解
DiffusionGemma 支持文本、图像输入，可以用于图像描述、视觉问答等任务。虽然视觉基准还有差距，但架构已经支持。

未来展望：

1. 质量提升：通过改进采样策略、增加训练数据，扩散模型的质量有望接近自回归水平
2. 更长上下文：当前 canvas length 限制为 256 tokens，未来可能扩展到更长
3. 混合架构：结合自回归和扩散的优势，在需要速度的部分使用扩散，需要质量的部分使用自回归
4. 专用硬件优化：扩散模型的并行特性更适合专用加速器（如 TPU、NPU），未来可能有针对性的硬件优化

DiffusionGemma 的意义不仅在于一个模型，而在于证明了一条新的技术路线。 它告诉我们：自回归不是语言生成的唯一方式，扩散模型在特定场景下可以做得更好。这种多样性对整个 AI 生态是好事——不同的工具有不同的优势，开发者可以根据场景选择最合适的方案。