DiffusionGemma

DiffusionGemma は、ブロック拡散型エンコーダー・デコーダー言語モデルです。因果エンコーダーがクリーンなプロンプト（および以前に生成されたブロック）を読み込んで KV キャッシュに格納し、双方向デコーダーがそのキャッシュへのクロスアテンションを用いて、canvas_length トークンからなる固定サイズの「キャンバス」をノイズ除去します。生成は、キャンバスに対する外側の自己回帰ループと、内側のノイズ除去ループを交互に実行します。内側のループでは、各ステップで候補トークンをサンプリングし、BlockRefinementScheduler を均一破損モードで用いて最も確信度の高いトークンを確定し、残りのトークンには再びノイズを付与します。モデル本体は 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 できます（Static cache and compilation を参照）。よりシンプルな動的キャッシュでの実行を行う場合は、両方を省略してください。

マルチターンまたはマルチモーダルの入力では、prompt の代わりに生の messages 会話を渡します。これは通常のチャット形式の {"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を渡すことができます。プロセッサはこれを自動的にモデルの画像入力に変換します。

スケジューラー

…は各キャンバスをデノイズするサンプラーであり、交換可能です。他の部分を一切変更せずに、サンプラーを切り替えるだけでサンプリング戦略を変更できます。利用可能なスケジューラーは3種類あります。

• BlockRefinementScheduler（デフォルト）: 各ステップで最も確信度の高いトークン（threshold以上、およびステップごとの均等な割り当て分）を確定し、残りのトークンには再ノイズを適用します。editing_threshold を設定すると、すでに確定済みのトークンも再編集できます。
• DiscreteDDIMScheduler: 一様な破損プロセス（D3PM）の正確な離散事後分布から各位置をサンプリングします。パラメータを必要とせず、最終ステップでは予測されたトークンを決定論的に確定します。
• EntropyBoundScheduler: 結合エントロピーがentropy_boundを下回る最低エントロピーの位置をコミットし、ほぼ独立したトークンをまとめて受け入れます。サンプリング温度は、最初のステップで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 で提案される leave-one-out predictor-corrector をサポートしています。この手法は、corrector_steps 回の Gibbs sweep を実行して、ノイズ付き周辺分布の 1 座標条件分布から最も確信度の低い位置を再サンプリングすることでキャンバスを洗練させます。これにより、その周辺分布は不変のまま、追加の学習コストなしで生成品質が向上します。公開されているチェックポイントに直接適用できます。Uniform diffusion の場合、denoiser と leave-one-out 事後分布は閉形式で互換性があるため、corrector は再学習なしに必要な leave-one-out 量を復元できます。

corrector のスウィープは num_inference_steps の予算に組み込まれ、上乗せされるわけではありません。パイプラインは predictor のステップ数を減らし、解放されたフォワードを corrector に割り当てるため、モデルのフォワード総数は num_inference_steps のままであり、predictor-corrector のコストは通常の ancestral sampling と同じになります。

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 アダプター

Denoiser は 🤗 Transformers のモデルであるため、アダプターは diffusers の load_lora_weights API ではなく、そのネイティブな PEFT 統合を通じて読み込まれます。この統合はアダプターの種類に依存しないため、同じ呼び出しで 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")

Adapters はマージされず、アクティブな状態のまま保持されます。DiffusionGemma はエンコーダーとデコーダーのベース重みを結合しているため、アダプターをこれらに融合すると、両方のブランチが破損します。

静的キャッシュとコンパイル

pipeline は、ブロックごとにエンコーダーを一度だけ事前計算し、再利用可能なキャッシュ（デフォルトでは DynamicCache）に格納します。cache_implementation="static" を渡すと、代わりに固定形状の StaticCache が使用され、その形状により cudagraphs を使用してデコーダーを torch.compile でき、さらなる高速化が図れます（pipeline は各ステップをマークし、logits をクローンすることで、cudagraph のメモリが上書きされないようにします）。この設定は Usage に示されています。よりシンプルな動的キャッシュでの実行を行う場合は、torch.compile の呼び出しと cache_implementation="static" の両方を省略してください。

適応型停止

ブロックは通常、num_inference_steps がすべて消費される前に収束するため、デフォルトでは、各バッチの例について、argmax 予測が stability_threshold ステップにわたって安定し、かつトークンごとの平均エントロピーが confidence_threshold（公開されたチェックポイントで使用される値である 0.005）を下回った時点で、その例をフリーズします。すべての例がフリーズされると、デノイジングループは終了します。これにより、同等の品質を維持しながらデコーダーのフォワードパス回数が約半分に減り、スループットを向上させる最大の単一要素となります。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（現在のブロックトークン）と 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