DiffusionGemma Technisch Rapport

Samenvatting

Wij introduceren DiffusionGemma, een experimenteel taalmodel met open gewichten dat gebruikmaakt van discrete diffusie om tekst met een uitzonderlijk hoge snelheid te genereren. In plaats van één token per keer te decoderen, verfijnt DiffusionGemma iteratief blokken van 256 tokens in parallel, waardoor de bottleneck van sequentiële decodering bij conventionele autoregressieve (AR) grote taalmodellen wordt omzeild.

Modelarchitectuur en Training

In plaats van het model vanaf nul te trainen, is DiffusionGemma tot stand gekomen door het fine-tunen van het mixture-of-experts Gemma 4-model, met 3,8 miljard geactiveerde en 25,2 miljard totale parameters. Onze rekenefficiënte trainingspijplijn in twee fasen gebruikt minder dan 10% van het totale trainingsbudget aan tokens van het oorspronkelijke AR-model:

  • Fase 1: Gebruikt supervised fine-tuning om bidirectionele denoising aan te leren.
  • Fase 2: Combineert reinforcement learning met sampler distillation om gezamenlijk de generatiekwaliteit en de efficiëntie van de inferentie te verbeteren.

Prestaties en Resultaten

DiffusionGemma vestigt een nieuwe Pareto-grens voor de afweging tussen generatiesnelheid en modelcapaciteit. Gemiddeld over onze volledige evaluatieset genereert het model ongeveer 20 tokens per forward pass en behaalt het ruwweg 1.500 output-tokens per seconde op een enkele NVIDIA H100 GPU. Dit is aanzienlijk sneller dan AR-modellen, zelfs wanneer deze gebruikmaken van state-of-the-art speculatieve decodering.

Functionaliteiten

DiffusionGemma behoudt de volgende eigenschappen van het basismodel:

  • Ondersteuning voor de denkmodus (thinking mode);
  • Multimodale inputs;
  • Lange contexten.

Ondanks de diffusie-fine-tuning blijft het model in staat tot AR-generatie met slechts minimale prestatievermindering. Dit suggereert een mogelijk pad naar hybride diffusie-AR-decodering.