株式会社kozokaAI 開発チーム
⚙️

拡散LLM DiffusionGemmaをModalで動かし、ノイズ除去の途中経過をブラウザで可視化してみた

@rtoshima
に公開
AI
modal
diffusion
ローカル LLM
Gemma 4
tech

2026年6月、GoogleがDiffusionGemmaを公開しました。これまでのGemmaと違い、自己回帰（左から1トークンずつ）ではなく、拡散ベースでテキストを生成するLLM（dLLM）です。画像生成で主流になった拡散モデルの考え方を、言語生成に持ち込んだものです。

今回は、これを実際に動かして、次の2つを試してみました。

  1. ノイズ除去（denoising）の途中経過を、ブラウザで可視化する
  2. 双方向性が効くタスク（OCR）での検証

「拡散LLMは並列で生成する」とよく説明されますが、文章で読むより、ノイズだらけの状態がきれいな文章に彫り出されていく様子を実際に見たほうが分かりやすいのでは？と思い、まずはそれを画面に出すことにしました。実行環境にはクラウドGPUのModalを使っています。

拡散LLMとは

自己回帰モデルは、文章を左から1トークンずつ順番に生成します。一方の拡散LLMは、固定長のキャンバス（トークンの並び）全体を、ノイズだらけの状態から反復的にノイズ除去して仕上げます。各ステップで全位置を同時に見直せるので、原理的には並列で速く生成できます。

DiffusionGemmaは、この拡散を256トークンのブロック単位で行い、ブロックが確定したら次のブロックへ進む、ブロック自己回帰的な構成になっています[1]。

項目 内容
モデル google/diffusiongemma-26B-A4B-it
パラメータ 総約252億 / アクティブ約38億（MoE）
生成方式 ブロック拡散（256トークンのキャンバスを反復denoise）
ライセンス Apache 2.0
クラス DiffusionGemmaForBlockDiffusion

Modalで動かす

26B-A4Bはアクティブこそ約38億ですが、推論時は総パラメータ（約252億）ぶんの重みをVRAMに載せる必要があり、手元のRTX 4070 Ti Super（16GB）には載りません。そこで、必要なときだけGPUを借りられるModalのA100を使いました。

なお、量子化（4bit）でのロードは今回の環境（transformers 5.11系）では失敗したため、bf16のまま載せています[2]。

import modal

app = modal.App("dgemma-web")

image = (
    modal.Image.debian_slim(python_version="3.12")
    .pip_install("unsloth", "transformers>=5.11")
    .pip_install("torchvision", "pillow", "fastapi[standard]", "python-multipart")
)
cache = modal.Volume.from_name("hf-cache", create_if_missing=True)


@app.cls(
    image=image,
    gpu="A100-80GB",
    scaledown_window=300,         # 5分無操作でGPU停止（課金も停止）
    secrets=[modal.Secret.from_name("huggingface")],
    volumes={"/root/.cache/huggingface": cache},
)
class Model:
    @modal.enter()
    def load(self):
        from unsloth import FastModel
        # 量子化せず bf16 フルロード
        self.model, self.processor = FastModel.from_pretrained(
            model_name="google/diffusiongemma-26B-A4B-it",
            load_in_4bit=False,
            load_in_16bit=True,
            max_seq_length=4096,
            device_map={"": 0},
        )
        FastModel.for_inference(self.model)

構成

ブラウザから入力を受け、Modal上のモデルが生成しながら、デノイズの途中経過を逐次返す構成です。

Modal上には、ブラウザ向けの web（FastAPI / ASGI）と、GPUでモデルを動かす Model を別々に置いています。web側は軽くて常駐に近い一方、Model側はA100を使うので、リクエストが来たときだけ起動し、しばらくアクセスがなければ自動で止まります（その間はGPU課金もかかりません）。

流れはこうです。ブラウザがプロンプトを /api/chat にPOSTすると、web がModelの生成を呼び出します。Modelはブロックをデノイズしながら、各ステップの途中キャンバス（draft）と、確定したテキスト（committed）をイベントとして出します。web はそれを受け取って、SSEでブラウザへ流します。ブラウザは draft をグレー、committed を黒で描くので、ノイズが晴れていく様子がそのまま画面に出ます。

ModelとwebをまたいだストリーミングはModalのジェネレータで実現しています。Model側の生成メソッドをジェネレータとして定義し、web側は remote_gen.aio() で1イベントずつ受け取って、SSEとして再配信しています。

実行環境

  • GPU: NVIDIA A100-80GB（Modal）
  • モデル: google/diffusiongemma-26B-A4B-it（bf16フルロード）
  • 主要ライブラリ: unsloth, transformers 5.11, FastAPI

可視化の実装

デノイズの各ステップの途中キャンバスを、ブラウザにそのまま流して表示します。

ドラフトはどこから取れるのか

DiffusionGemmaは、最終テキストの前に中間ドラフトを生成します。transformersには専用の TextDiffusionStreamer が用意されていて、これを generate に渡すと中間状態を見られます。

ただし、このクラスのソースを読むと、ドラフトをANSIエスケープでターミナルに直接printしているだけでした。put_draft で全ドラフトトークン列を受け取り、カーソルを保存して黄色で表示し、次のステップで上書きする、という挙動です。つまり標準のままだと、ターミナルには出るがブラウザには来ません。

そこで TextDiffusionStreamer を継承し、put_draft を上書きして、ANSI描画する代わりにキューへ流すようにしました。確定テキストは on_finalized_text から拾います。

from transformers import TextDiffusionStreamer

class QueueDiffusionStreamer(TextDiffusionStreamer):
    def __init__(self, tokenizer, out_q, **kw):
        super().__init__(tokenizer, skip_prompt=True, **kw)
        self.out_q = out_q
        self.committed = ""

    # デノイズ各ステップの途中キャンバス（全ドラフトトークン列が渡る）
    def put_draft(self, value, **kwargs):
        if hasattr(value, "shape") and len(value.shape) > 1:
            value = value[0]
        text = self.tokenizer.decode(value, **self.decode_kwargs)
        self.out_q.put({"type": "draft", "committed": self.committed, "draft": text})
        # 基底のANSI描画(super().put_draft)は呼ばない

    # 確定テキスト（差分を貯める）
    def on_finalized_text(self, text, stream_end=False):
        if text:
            self.committed += text
            self.out_q.put({"type": "draft", "committed": self.committed, "draft": ""})

generate は別スレッドで走らせ、キューに溜まったイベントをジェネレータでyieldします。

もう一つのハマりどころ：バッファリング

ここまでで動くはずでしたが、最初の実装では、待っている間ずっと空白で、最後に一括で全部が出るという状態でした。デノイズは生成できているのに、レスポンスがバッファされて最後にまとめて届いていました。

これは配信形式をSSE（text/event-stream）に変えて逐次フラッシュさせることで解決しました。

@modal.asgi_app()
def web():
    from fastapi import FastAPI
    from fastapi.responses import StreamingResponse
    api = FastAPI()

    @api.post("/api/chat")
    async def chat(...):
        async def gen():
            async for chunk in Model().generate_stream.remote_gen.aio(...):
                yield chunk            # "data: {...}\n\n" 形式
        return StreamingResponse(
            gen(),
            media_type="text/event-stream",
            headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
        )
    return api

ブラウザ側は、確定テキスト（committed）を黒、途中キャンバス（draft）をグレーで表示するだけです。

実行結果：デノイズが彫り出される様子

これで、グレーの崩れたテキストが何度も書き換わりながら、黒の確定文に収束していく様子がブラウザに流れるようになりました。

たとえば「寿限無」について聞くと、途中段階では以下のような状態が見えます（実際の画面ではグレー表示）。

寿無彦寿彦彦彦寿寿寿寿寿寿寿寿寿寿寿寿寿寿 ...
名前名前名前名前が名前
、が、、、、、、、、、、、、、

寿寿寿 や 名前名前、、、、、 といった繰り返しやプレースホルダだらけの状態が、ステップを重ねるごとにきれいな文章へ収束していきます。全体を一度に見直しながら直していく、拡散らしい動きです。

DiffusionGemmaがノイズから文章を生成していく様子

短い対話での生成速度は、A100-80GB・bf16・1リクエストで概ね58〜122 tok/sでした[3]。

OCR：双方向性が効くタスクで試す

拡散は、自己回帰と違って全位置を同時に（双方向に）見られるのが特徴です。穴埋め的なタスクに向いているはずなので、OCRで試しました。

検証用に、わざと劣化させた架空の発注書FAX（傾き・ノイズ・スキャンむらを付与）を用意し、「この画像のテキストを読み取って」と投げました。

読み取らせた架空の発注書FAX

結果が以下です（発注番号・電話番号・検収欄などのヘッダ部分は省略し、主要部分を抜粋）。Markdownの表まで再現していました。

| No | 品名 | 数量 | 単価 | 金額 |
| 1 | ステンレス鋼板 SUS304 | 50 | 3,200 | 160,000 |
| 2 | アルミ角材A6063 | 120 | 850 | 102,000 |
| 3 | 六角ボルト M8×30 | 500 | 48 | 24,000 |
| 4 | 防錆スプレー 480ml | 24 | 520 | 12,480 |
| 小計 |  |  |  | 298,480 |
| 消費税(10%) |  |  |  | 29,848 |
| 合計金額 |  |  |  | 328,328円 |

希望納期: 2026年7月3日（金）
納入場所: 東京都大田区蒲田第二工場

生成速度は 1844 tok / 4.0s = 459 tok/s でした。

OCR中も、確定したブロックが黒、デノイズ途中のブロックがグレーで表示されます。

OCR中のデノイズの様子（上が確定、下が途中キャンバス）

正解と照合すると、全20項目中19項目が正解でした。数字（発注番号・数量・単価・金額・小計・消費税・合計・電話・日付）はすべて正確で、漢字（鋼板・角材・六角・防錆・蒲田・検収）もすべて正確でした。

誤りは1つだけです。

正解 読み取り結果
コウゾウカ技研株式会社 コソウカ技研株式会社

誤ったのは漢字ではなく、カタカナの架空の社名でした。辞書にある漢字熟語は正確に読む一方で、辞書にない造語のカタカナ社名で外しています。知っている語には強く、知らない固有名詞には弱いようです。

内容の正しさ

きれいに生成されても、答えが合っているとは限りません。さっきの「寿限無」も、内容は間違っていました。

項目 モデルの出力 正しくは
読み ことぶげんなく じゅげむ
フルネーム 寿限無、寿限無、彦戸屋、寿、寿… 寿限無、寿限無、五劫の擦り切れ、海砂利水魚の…

文章としては自然なので、ぱっと見では誤りに気づきにくいです。

また、拡散はブロック全体を反復的に仕上げる生成方式で、サンプリングを伴うため、同じ入力でも実行ごとに結果が変わることがあります。今回のOCRはほぼ完璧でしたが、別の機会に同種の帳票を読ませたときは漢字をいくつか誤読し、日付を1日ずらすこともありました。うまく読めるときと外すときの差が大きいです。

まとめ

DiffusionGemma（拡散LLM）をModalのA100-80GBで動かし、デノイズの途中経過をブラウザで可視化し、OCRまで試しました。

わかったこと：

  • デノイズの途中経過は TextDiffusionStreamer の put_draft を横取りすれば取得でき、SSEで逐次配信すればブラウザに表示できる
  • OCRは数字・漢字ともに高精度だった（今回は459 tok/s、20項目中19正解）

課題：

  • transformers/unsloth の4bitロードは今回の環境では失敗した（ローカルで動かすならGGUF量子化版＋llama.cppのdiffusion対応ビルドを使う）
  • 文章は自然でも内容を間違える（固有名詞のでっち上げ、誤った読み）
  • 生成のたびに結果が揺れ、誤りの位置も一定しない
  • カタカナの未知固有名詞など、辞書外の語に弱い

拡散LLMは速くて並列に生成できる一方、自然な文でも中身を間違えることがある、という両面を実際に触って確認できました。ノイズから文章が立ち上がっていく様子は、何度見ても面白いです。

リンク

脚注

  1. 表の各項目は、DiffusionGemma モデルカードおよび Hugging Face Transformers の diffusion_gemma ドキュメントを参照。 ↩︎

  2. bf16フルロードでVRAM約51.7GB。A100-80GBで動作 ↩︎

  3. 計測はウォールクロック時間。初回のモデルロードは除く ↩︎

株式会社kozokaAI 開発チーム
株式会社kozokaAI 開発チーム

kozokaAI開発チームによるAI tech Blogです。現場で役立つAI活用の実践知を発信中。ご興味のある方はぜひフォローしてご覧ください。

Discussion