DiffusionGemma：透過平行擴散實現 4 倍速文本生成

DiffusionGemma 透過將序列化 Token 生成替換為平行文本擴散，實現了 4 倍速的推理速度

DiffusionGemma 是一款實驗性的開源權重模型，旨在消除傳統自回歸大型語言模型 (LLMs) 的延遲瓶頸。透過同時生成整個文本塊而非逐個 Token，它在專用 GPU 上可提供高達 4 倍的文本生成速度。該模型以 Apache 2.0 授權發布，採用 26B Mixture of Experts (MoE) 架構，在推理期間僅激活 3.8B 參數，使其在量化後能符合高階消費級 GPU 的 18GB VRAM 限制。

解決記憶體頻寬瓶頸

傳統 LLM 的運作方式就像「打字機」，從左到右依序生成 Token。雖然這對於可以進行批次處理請求的高併發雲端服務非常有效，但對於本地單用戶推理而言，這種過程效率低下。在本地環境中，GPU 往往無法得到充分利用，因為系統花費在將權重從 RAM 移動到處理器的時間比實際執行計算的時間還要多。

DiffusionGemma 將瓶頸從記憶體頻寬轉移到計算能力。它不再是預測單個下一個 Token，而是同時起草一個 256-token 的段落。這種「印刷機」方法使硬體計算能力達到飽和，從而在本地加速器上獲得顯著的速度提升：

  • NVIDIA H100: 1000+ tokens per second.
  • NVIDIA GeForce RTX 5090: 700+ tokens per second.

技術架構與能力

DiffusionGemma 整合了基於 Gemma 4 系列與 Gemini Diffusion 研究開發的新型擴散頭 (diffusion head)。其運作邏輯模仿 AI 圖像生成：它從一個包含隨機佔位符 Token 的畫布開始，並透過多次迭代優化，直到文本收斂為最終輸出。

技術關鍵優勢

  • 雙向注意力機制 (Bi-directional Attention): 由於 256 個 Token 是平行生成的，每個 Token 都可以關注到塊內的所有其他 Token。這使得該模型特別適合非線性任務，例如代碼填充 (code infilling)、行內編輯、數學圖表以及氨基酸序列。
  • 智慧自我修正: 模型會一次性評估整個文本塊，使其能夠在優化過程中即時迭代修正錯誤。
  • 硬體優化: 模型支持 NVFP4 (4-bit floating-point) kernels，這能在 Hopper 和 Blackwell 架構上以接近無損的精度加速計算吞吐量。

權衡與生產環境使用

DiffusionGemma 優先考慮速度與平行佈局，而非原始輸出品質。其整體品質低於標準的自回歸 Gemma 4 模型。因此，Google 建議對於需要最高品質的應用程式使用標準 Gemma 4，而 DiffusionGemma 則適用於對速度要求極高、具備互動性的本地工作流。

開發者生態系統與整合

DiffusionGemma 旨在進行快速實驗，並與多個主要的 AI 框架相容：

  • 服務提供 (Serving): 透過 MLX、vLLM (包含 Red Hat 整合) 以及 Hugging Face Transformers 提供支持。官方對 llama.cpp 的支持尚在進行中。
  • 微調 (Fine-tuning): 開發者可以使用 Hackable Diffusion JAX 工具箱、Unsloth 或 NVIDIA NeMo。一個顯著的微調範例包括使用 Unsloth 使模型能夠解決 Sudoku，這是一項雙向注意力機制比自回歸模型更具優勢的任務。
  • 部署 (Deployment): 可透過 Hugging Face、Gemini Enterprise Agent Platform Model Garden 以及 NVIDIA NIM 取得。

社群洞察與分析

開發者之間的技術討論強調了擴散模型透過減少提示與回應之間的等待時間，來轉化「結對編程 (pair-programming)」體驗的潛力。

"It was more of a pair-programming experience instead of the SOTA agentic experience of prompting and waiting... It felt less of a slot machine where you prompt, wait, and hope it went in the right direction." — @vineyardmike

其他貢獻者強調了邊緣設備上記憶體頻寬問題的關鍵性：

"On edge you have a different problem: your inference accelerator is starved while sloshing GB of weights back and forth from RAM... Diffusion can compute tokens in parallel which relieves the memory bandwidth bottle neck." — @samuelknight

關於模型在特定領域（如工具調用與推理）的表現仍存疑問。部分用戶建議，該模型重新編輯先前行數的能力使其更適合「變更流 (change streams)」或跨多個文件的系列編輯操作，而非單一的工具調用。

Sources