DiffusionGemma

DiffusionGemma 是一款基于块扩散（block-diffusion）的编码器-解码器语言模型。因果编码器将干净的提示词（以及任何先前生成的块）读取到 KV 缓存中，而双向解码器则通过交叉注意力机制关注该缓存，对固定大小为 canvas_length 个 token 的“画布”进行去噪。生成过程交替执行两个循环：外层是遍历各画布的自回归循环，内层是去噪循环。在去噪循环的每一步中，模型会采样候选 token，在均匀损坏（uniform corruption）模式下通过 BlockRefinementScheduler 提交置信度最高的 token，并对其余 token 重新加噪。该模型本身位于 transformers 库中，类名为 DiffusionGemmaForBlockDiffusion；发布的检查点为 google/diffusiongemma-26B-A4B-it。

用法

import torch
from transformers import AutoProcessor, DiffusionGemmaForBlockDiffusion

from diffusers import BlockRefinementScheduler, DiffusionGemmaPipeline

model_id = "google/diffusiongemma-26B-A4B-it"
model = DiffusionGemmaForBlockDiffusion.from_pretrained(model_id, dtype=torch.bfloat16, device_map="auto")
processor = AutoProcessor.from_pretrained(model_id)
scheduler = BlockRefinementScheduler()

pipe = DiffusionGemmaPipeline(model=model, scheduler=scheduler, processor=processor)
pipe.model.model.decoder = torch.compile(pipe.model.model.decoder, mode="reduce-overhead", fullgraph=True)
output = pipe(
    prompt="Why is the sky blue?",
    gen_length=256,
    num_inference_steps=48,
    cache_implementation="static",
)
print(output.texts[0])

num_inference_steps 是每个画布的降噪步数（48 与发布的检查点相匹配）；步数越少速度越快，但质量越低。cache_implementation="static" 允许解码器使用 cudagraphs 进行 torch.compile 编译（参见 静态缓存与编译）；若希望进行更简单的动态缓存运行，可省略这两项。

对于多轮或多模态输入，请传入原始的 messages 对话，而不是 prompt。它是一个由 {"role", "content"} 字典组成的列表，采用常见的聊天格式，处理器会将其通过其聊天模板进行处理：

messages = [
    {"role": "user", "content": "Why is the sky blue?"},
]
# or with an image:
messages = [
    {
        "role": "user",
        "content": [
            {"type": "image", "image": image},
            {"type": "text", "text": "Describe this image."},
        ],
    },
]
output = pipe(messages=messages, gen_length=256)

对于单个用户轮次，您可以省略 messages，并在 prompt 旁边传入 image；处理器会自动将其转换为模型的图像输入。

调度器

调度器（scheduler）是用于对每个画布进行去噪的采样器，并且可以互换：只需替换调度器即可更改采样策略，而无需修改其他任何内容。目前提供三种调度器：

• BlockRefinementScheduler（默认）：在每一步中提交置信度最高的 token（高于 threshold，外加一个均匀的每步配额），并对其余 token 重新加噪。editing_threshold 还允许它重新编辑已提交的 token。
• DiscreteDDIMScheduler：从均匀腐蚀过程（D3PM）的精确离散后验分布中为每个位置采样。该调度器无需参数，且最后一步会确定性地提交预测的 token。
• EntropyBoundScheduler：提交联合熵保持在 entropy_bound 以下的最低熵位置，从而将大致独立的 token 一起接受。其采样温度从第一步的 t_max（0.8）退火至最后一步的 t_min（0.4），与发布的检查点采样器保持一致。

from diffusers import DiscreteDDIMScheduler, EntropyBoundScheduler

pipe.scheduler = DiscreteDDIMScheduler()
# or: pipe.scheduler = EntropyBoundScheduler(entropy_bound=0.1)
output = pipe(prompt="Why is the sky blue?", gen_length=256, num_inference_steps=48)
print(output.texts[0])

调度器特定的采样参数（块细化threshold/top_k、熵边界等）在调度器配置中设置：

from diffusers import BlockRefinementScheduler

pipe.scheduler = BlockRefinementScheduler.from_config(pipe.scheduler.config, threshold=0.9)

EntropyBoundScheduler 会在去噪步骤内部对其采样温度（t_max/t_min）进行退火；DiscreteDDIMScheduler 和 BlockRefinementScheduler 则使用传入流水线的固定 temperature（贪心策略下为 0.0）。

预测器-校正器采样

DiscreteDDIMScheduler 支持 Uniform Diffusion Models Revisited: Leave-One-Out Denoiser and Absorbing State Reformulation 中提出的留一法预测-校正器。它通过 corrector_steps 次 Gibbs 扫描来细化画布，从含噪边缘分布的单坐标条件分布中重新采样置信度最低的位置，这一过程保持该边缘分布不变，且无需额外训练成本即可提升生成质量。它可直接在已发布的检查点上运行：对于均匀扩散，去噪器与留一法后验在闭式解下是可互换的，因此校正器无需任何重新训练即可恢复其所需的留一法量。

校正器扫描被纳入 num_inference_steps 的预算中，而非额外叠加：流水线会减少预测器步数，并将释放出的前向计算用于校正器，因此模型前向调用的总次数仍为 num_inference_steps，预测器-校正器的开销与普通祖先采样相同。

from diffusers import DiscreteDDIMScheduler

pipe.scheduler = DiscreteDDIMScheduler(corrector_steps=2, corrector_k=12)
output = pipe(prompt="Why is the sky blue?", gen_length=256, num_inference_steps=48)
print(output.texts[0])

PEFT 适配器

去噪器是一个 🤗 Transformers 模型，因此适配器通过其原生的 PEFT 集成进行加载，而不是通过 diffusers 的 load_lora_weights API。由于该集成与适配器类型无关，相同的调用可以加载 LoRA、DoRA 或任何其他 PEFT 适配器（例如 TRL 的 SFTTrainer 的输出）。直接在模型组件上管理适配器：

pipe.model.load_adapter("path/to/adapter", adapter_name="sft")  # LoRA, DoRA, ...
pipe.model.set_adapter("sft")
output = pipe(prompt="Why is the sky blue?", gen_length=256)

pipe.model.disable_adapters()  # run the base model
pipe.model.delete_adapter("sft")

适配器保持激活且未合并：DiffusionGemma 将编码器和解码器的基础权重绑定在一起，因此将适配器融合到其中会破坏这两个分支。

静态缓存与编译

Pipeline 会为每个 block 将编码器预填充一次，并写入可复用的缓存（默认为 DynamicCache）。传入 cache_implementation="static" 则会改用固定形状的 StaticCache，其形状允许你使用 cudagraphs 对解码器进行 torch.compile 以进一步提速（pipeline 会标记每个步骤并克隆 logits，以避免 cudagraph 内存被覆盖）；这正是 Usage 中展示的配置。若希望进行更简单的动态缓存运行，可同时去掉 torch.compile 调用和 cache_implementation="static"。

自适应停止

一个 block 通常会在用完所有 num_inference_steps 之前收敛，因此默认情况下，pipeline 会在每个 batch 的 argmax 预测在 stability_threshold 步内保持稳定，且其平均逐 token 熵低于 confidence_threshold（0.005，即发布 checkpoint 所使用的值）时，冻结该 batch 中的示例。一旦所有示例都被冻结，去噪循环即结束。在质量相当的情况下，这大约将 decoder 前向传播次数减半，是提升吞吐量最大的单一手段。传入 confidence_threshold=None 可始终执行完整的 num_inference_steps：

output = pipe(prompt="Why is the sky blue?", gen_length=256, confidence_threshold=None)  # disable adaptive stopping

回调

回调函数会在每个去噪步骤之后运行。通过传入 callback_on_step_end_tensor_inputs 来选择包含在 callback_kwargs 中的张量；canvas（当前块 token）和 logits 均可用。从回调函数中返回 {"canvas": ...} 即可替换 canvas。

def on_step_end(pipe, step, timestep, callback_kwargs):
    canvas = callback_kwargs["canvas"]
    # Inspect or modify `canvas` here.
    return {"canvas": canvas}


out = pipe(
    prompt="Why is the sky blue?",
    callback_on_step_end=on_step_end,
    callback_on_step_end_tensor_inputs=["canvas"],
)

DiffusionGemmaPipeline

>>> import torch
>>> from transformers import AutoProcessor, DiffusionGemmaForBlockDiffusion
>>> from diffusers import BlockRefinementScheduler, DiffusionGemmaPipeline

>>> model_id = "google/diffusiongemma-26B-A4B-it"
>>> model = DiffusionGemmaForBlockDiffusion.from_pretrained(model_id, dtype=torch.bfloat16, device_map="auto")
>>> processor = AutoProcessor.from_pretrained(model_id)
>>> scheduler = BlockRefinementScheduler()

>>> pipe = DiffusionGemmaPipeline(model=model, scheduler=scheduler, processor=processor)
>>> output = pipe(prompt="Why is the sky blue?", gen_length=256)
>>> print(output.texts[0])

DiffusionGemmaPipelineOutput