DiffusionGemma Explained

ML at Berkeley 08/10/26, 06:10 PM Models

diffusion gemma state-diffusion language-model open-weights architecture tutorial

Summary

An annotated from-scratch reimplementation of Google's DiffusionGemma, a 26B open-weight state diffusion language model, explaining its architecture, sampling procedure, and design choices.

<p>Google recently released <a href="https://blog.google/innovation-and-ai/technology/developers-tools/diffusion-gemma-faster-text-generation/">DiffusionGemma</a>, a 26B open-weight uniform state diffusion language model ~4× faster than Gemma 4 for certain workloads.<a class="footnote-anchor" data-component-name="FootnoteAnchorToDOM" id="footnote-anchor-1" href="#footnote-1" target="_self">1</a></p><p><span>As an experimental model, it doesn’t have a whitepaper, and the </span><a href="https://www.google.com/url?q=https%3A%2F%2Fdevelopers.googleblog.com%2Fen%2Fdiffusiongemma-the-developer-guide%2F">officially</a><span> </span><a href="https://www.google.com/url?q=https%3A%2F%2Fblog.google%2Finnovation-and-ai%2Ftechnology%2Fdevelopers-tools%2Fdiffusion-gemma-faster-text-generation%2F">released</a><span> </span><a href="https://www.google.com/url?q=https%3A%2F%2Fai.google.dev%2Fgemma%2Fdocs%2Fdiffusiongemma%2Fexplained">docs</a><span> are pretty brief and high-level. </span>This blog post, based on the <a href="https://colab.research.google.com/drive/17egFGmboDkhU6duNQHvjB1J0oiKslOMl?usp=sharing">colab notebook</a>, attempts to fill in the gaps. It walks through an annotated from-scratch reimplementation derived from its <a href="https://huggingface.co/google/diffusiongemma-26B-A4B-it">open weights</a>, following the style of the awesome <a href="https://nlp.seas.harvard.edu/annotated-transformer/">Annotated Transformer</a>.</p><div class="subscription-widget-wrap-editor" data-attrs="{"url":"https://mlberkeley.substack.com/subscribe?","text":"Subscribe","language":"en"}" data-component-name="SubscribeWidgetToDOM"><div class="subscription-widget show-subscribe"><div class="preamble"><p class="cta-caption">Thanks for reading ML@B Blog! Subscribe for free to receive new posts and support my work.</p></div><form class="subscription-widget-subscribe"><input type="email" class="email-input" name="email" placeholder="Type your email…" tabindex="-1"><input type="submit" class="button primary" value="Subscribe"><div class="fake-input-wrapper"><div class="fake-input"></div><div class="fake-button"></div></div></form></div></div><p>Along the way, we explain some of the more opaque details — behind the model architecture, self-conditioning, the sampling procedure, how encode/decode weight sharing is actually implemented. Through closely examining the model config and weights, we also uncover some interesting design choices — the layer scalar, partial RoPE, and Google’s Scalar QK Norm — and offer possible explanations for them. Finally, we conclude by making some visualizations and empirical observations about the denoising process.</p><p>Thank you Lucas Gu, Arshia Nayebnazar, Henry Ko, Rishi Athavale, and Tejas Prabhune for help proofreading and improving this blog post!</p><p><em>Note: prior</em> <em>knowledge of vanilla autoregressive LLM implementation (Llama 3.1, MOE) is assumed.</em></p><div class="captioned-image-container"><figure><a class="image-link image2 is-viewable-img" target="_blank" href="https://substackcdn.com/image/fetch/$s_!9Yqo!,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F7a59cc04-b5fa-4c97-963d-278ca9d8aff0_720x651.gif" data-component-name="Image2ToDOM"><div class="image2-inset"><picture><source type="image/webp" srcset="https://substackcdn.com/image/fetch/$s_!9Yqo!,w_424,c_limit,f_webp,q_auto:good,fl_lossy/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F7a59cc04-b5fa-4c97-963d-278ca9d8aff0_720x651.gif 424w, https://substackcdn.com/image/fetch/$s_!9Yqo!,w_848,c_limit,f_webp,q_auto:good,fl_lossy/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F7a59cc04-b5fa-4c97-963d-278ca9d8aff0_720x651.gif 848w, https://substackcdn.com/image/fetch/$s_!9Yqo!,w_1272,c_limit,f_webp,q_auto:good,fl_lossy/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F7a59cc04-b5fa-4c97-963d-278ca9d8aff0_720x651.gif 1272w, https://substackcdn.com/image/fetch/$s_!9Yqo!,w_1456,c_limit,f_webp,q_auto:good,fl_lossy/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F7a59cc04-b5fa-4c97-963d-278ca9d8aff0_720x651.gif 1456w" sizes="100vw"><img src="https://substackcdn.com/image/fetch/$s_!9Yqo!,w_1456,c_limit,f_auto,q_auto:good,fl_lossy/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F7a59cc04-b5fa-4c97-963d-278ca9d8aff0_720x651.gif" width="720" height="651" data-attrs="{"src":"https://substack-post-media.s3.amazonaws.com/public/images/7a59cc04-b5fa-4c97-963d-278ca9d8aff0_720x651.gif","srcNoWatermark":null,"fullscreen":null,"imageSize":null,"height":651,"width":720,"resizeWidth":null,"bytes":null,"alt":"demo.gif","title":"demo.gif","type":null,"href":null,"belowTheFold":false,"topImage":true,"internalRedirect":null,"isProcessing":false,"align":null,"offset":false}" class="sizing-normal" alt="demo.gif" title="demo.gif" srcset="https://substackcdn.com/image/fetch/$s_!9Yqo!,w_424,c_limit,f_auto,q_auto:good,fl_lossy/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F7a59cc04-b5fa-4c97-963d-278ca9d8aff0_720x651.gif 424w, https://substackcdn.com/image/fetch/$s_!9Yqo!,w_848,c_limit,f_auto,q_auto:good,fl_lossy/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F7a59cc04-b5fa-4c97-963d-278ca9d8aff0_720x651.gif 848w, https://substackcdn.com/image/fetch/$s_!9Yqo!,w_1272,c_limit,f_auto,q_auto:good,fl_lossy/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F7a59cc04-b5fa-4c97-963d-278ca9d8aff0_720x651.gif 1272w, https://substackcdn.com/image/fetch/$s_!9Yqo!,w_1456,c_limit,f_auto,q_auto:good,fl_lossy/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F7a59cc04-b5fa-4c97-963d-278ca9d8aff0_720x651.gif 1456w" sizes="100vw" fetchpriority="high"></picture><div class="image-link-expand"><div class="pencraft pc-display-flex pc-gap-8 pc-reset"><button tabindex="0" type="button" class="pencraft pc-reset pencraft icon-container restack-image"><svg aria-hidden="true" width="20" height="20" viewBox="0 0 20 20" fill="none" stroke-width="1.5" stroke="var(--color-fg-primary)" stroke-linecap="round" stroke-linejoin="round" xmlns="http://www.w3.org/2000/svg"><g><path d="M2.53001 7.81595C3.49179 4.73911 6.43281 2.5 9.91173 2.5C13.1684 2.5 15.9537 4.46214 17.0852 7.23684L17.6179 8.67647M17.6179 8.67647L18.5002 4.26471M17.6179 8.67647L13.6473 6.91176M17.4995 12.1841C16.5378 15.2609 13.5967 17.5 10.1178 17.5C6.86118 17.5 4.07589 15.5379 2.94432 12.7632L2.41165 11.3235M2.41165 11.3235L1.5293 15.7353M2.41165 11.3235L6.38224 13.0882"></path></g></svg></button><button tabindex="0" type="button" class="pencraft pc-reset pencraft icon-container view-image"><svg xmlns="http://www.w3.org/2000/svg" width="20" height="20" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="lucide lucide-maximize2 lucide-maximize-2"><polyline points="15 3 21 3 21 9"></polyline><polyline points="9 21 3 21 3 15"></polyline><line x1="21" x2="14" y1="3" y2="10"></line><line x1="3" x2="10" y1="21" y2="14"></line></svg></button></div></div></div></a><figcaption class="image-caption">Entropy over vocabulary distribution at each canvas position over time, figure generated at the end of the notebook</figcaption></figure></div><p></p><h2>Table of contents</h2><ul><li><p><a href="https://mlberkeley.substack.com/p/the-annotated-diffusiongemma#§setup-load-model">Setup + Load Model</a></p><ul><li><p><a href="https://mlberkeley.substack.com/p/the-annotated-diffusiongemma#§googles-scalar-qk-normalization">Google’s Scalar QK Normalization</a></p></li><li><p><a href="https://mlberkeley.substack.com/p/the-annotated-diffusiongemma#§rotary-position-embeddings">Rotary position embeddings</a></p></li></ul></li><li><p><a href="https://mlberkeley.substack.com/p/the-annotated-diffusiongemma#§attention">Attention</a></p><ul><li><p><a href="https://mlberkeley.substack.com/p/the-annotated-diffusiongemma#§decode-vs-encode">Decode vs Encode</a></p></li><li><p><a href="https://mlberkeley.substack.com/p/the-annotated-diffusiongemma#§gemma-specific-details">Gemma-specific details</a></p></li><li><p><a href="https://mlberkeley.substack.com/p/the-annotated-diffusiongemma#§optimizations-and-implementation">Optimizations and Implementation</a></p></li></ul></li><li><p><a href="https://mlberkeley.substack.com/p/the-annotated-diffusiongemma#§moe">MOE</a></p><ul><li><p><a href="https://mlberkeley.substack.com/p/the-annotated-diffusiongemma#§moe-forward">MOE Forward</a></p></li><li><p><a href="https://mlberkeley.substack.com/p/the-annotated-diffusiongemma#§routing-score">Routing Score</a></p></li><li><p><a href="https://mlberkeley.substack.com/p/the-annotated-diffusiongemma#§rmsnorm-fusion">RMSNorm Fusion</a></p></li></ul></li><li><p><a href="https://mlberkeley.substack.com/p/the-annotated-diffusiongemma#§putting-them-together">Putting Them together</a></p><ul><li><p><a href="https://mlberkeley.substack.com/p/the-annotated-diffusiongemma#§logit-softcapping">Logit Softcapping</a></p></li><li><p><a href="https://mlberkeley.substack.com/p/the-annotated-diffusiongemma#§denoise-vs-prefill">Denoise vs Prefill</a></p></li><li><p><a href="https://mlberkeley.substack.com/p/the-annotated-diffusiongemma#§layer-scalar">Layer Scalar</a></p></li></ul></li><li><p><a href="https://mlberkeley.substack.com/p/the-annotated-diffusiongemma#§sampling">Sampling</a></p><ul><li><p><a href="https://mlberkeley.substack.com/p/the-annotated-diffusiongemma#§lazy-sampling">Lazy Sampling</a></p></li><li><p><a href="https://mlberkeley.substack.com/p/the-annotated-diffusiongemma#§stage-then-commit">Stage Then Commit</a></p></li><li><p><a href="https://mlberkeley.substack.com/p/the-annotated-diffusiongemma#§preliminary-analysis">Preliminary Analysis</a></p></li></ul></li><li><p><a href="https://mlberkeley.substack.com/p/the-annotated-diffusiongemma#§references">References</a></p></li></ul><h2>Setup + Load Model</h2><p><strong>Install, import, load configs and weights</strong></p><pre><code><code>!pip install -q -U huggingface_hub !hf download google/diffusiongemma-26B-A4B-it \ --local-dir /content/diffusiongemma import glob import json import torch import torch.nn.functional as F from einops import rearrange, einsum from safetensors.torch import load_file from tokenizers import Tokenizer from torch.distributions import Categorical from tqdm import trange torch.set_default_device("cuda") max_tot_tokens = 2048 checkpoint = "diffusiongemma" canvas_len = 256 model_config = json.load(open(f"{checkpoint}/config.json"))['text_config'] gen_config = json.load(open(f"{checkpoint}/generation_config.json")) V = model_config['vocab_size'] sd = {} for safetensor_path in glob.glob(f"{checkpoint}/model-*.safetensors"): sd |= {k: v for k, v in load_file(safetensor_path, device="cuda").items() if "vision" not in k} W_vocab = sd['model.decoder.embed_tokens.weight'] sd</code></code></pre><h3>Google’s Scalar QK Normalization</h3><p>Printing out the weights, we find that interestingly, all entries within each learned QK-normalization weight vector have the same value. For example:</p><pre><code><code>'model.decoder.layers.17.self_attn.q_norm.weight': tensor( [0.9883, 0.9883, 0.9883, 0.9883, 0.9883, 0.9883, 0.9883, 0.9883, 0.9883, 0.9883, 0.9883, 0.9883, 0.9883, 0.9883, 0.9883, ...] )</code></code></pre><p>The QK-norm weight vectors have the form</p><div class="latex-rendered" data-attrs="{"persistentExpression":"w_q = \\gamma_q \\mathbf{1},\n\\qquad\nw_k = \\gamma_k \\mathbf{1},w_q = \\gamma_q \\mathbf{1},\n\\qquad\nw_k = \\gamma_k \\mathbf{1},","id":"PMNOERPNJS"}" data-component-name="LatexBlockToDOM"></div><p>where <span>γ_q</span> and <span>γ_k</span> are the learned scalars.</p><div class="latex-rendered" data-attrs="{"persistentExpression":"\\begin{aligned}\n\\left\\lVert \\operatorname{RMSNorm}_{\\gamma}(x) \\right\\rVert_2\n&=\n|\\gamma|\n\\frac{\\lVert x \\rVert_2}\n{\\lVert x \\rVert_2 / \\sqrt{d_h}} \\\\\n&=\n|\\gamma|\\sqrt{d_h}\n\\end{aligned}","id":"RWXXWENJCJ"}" data-component-name="LatexBlockToDOM"></div><p>RoPE preserves Euclidean norm, so the same result holds after applying rotate:</p><div class="latex-rendered" data-attrs="{"persistentExpression":"\\lVert q \\rVert_2\n=\n|\\gamma_q|\\sqrt{d_h},\n\\qquad\n\\lVert k \\rVert_2\n=\n|\\gamma_k|\\sqrt{d_h}","id":"XNSPUYJMZJ"}" data-component-name="LatexBlockToDOM"></div><p>For a query <span>q</span> and key <span>k_i</span> separated by an angle <span>θ_i</span>,</p><div class="latex-rendered" data-attrs="{"persistentExpression":"\\langle q, k_i \\rangle\n=\nd_h \\gamma_q\\gamma_k\\cos\\theta_i","id":"AVRYVFRKYA"}" data-component-name="LatexBlockToDOM"></div><p>Thus, the attention score for <span>k_i</span> is</p><div class="latex-rendered" data-attrs="{"persistentExpression":"\\begin{aligned}\n\\operatorname{score}_q(k_i)\n&\\propto\n\\exp\\!\\left(\\langle q, k_i \\rangle\\right) \\\\\n&=\n\\exp\\!\\left(d_h\\gamma_q\\gamma_k\\cos\\theta_i\\right) \\\\\n&=\n\\exp\\!\\left(\\frac{\\cos\\theta_i}{T}\\right),\n\\end{aligned}","id":"APOYTDGFRW"}" data-component-name="LatexBlockToDOM"></div><p>where</p><div class="latex-rendered" data-attrs="{"persistentExpression":"T\n=\n\\frac{1}{d_h\\gamma_q\\gamma_k}","id":"MVHGKSUTFE"}" data-component-name="LatexBlockToDOM"></div><p>is a fixed, learned effective attention temperature per layer.</p><p>In regular QK norm models (and non-QK norm models), the magnitudes of q and k can vary. This 1) produces a per-query-dependent effective temperature and 2) allows key magnitude to encode a form of query-independent “global importance”:</p><div class="latex-rendered" data-attrs="{"persistentExpression":"\\begin{aligned}\nscore_q(k_i)\n&\\propto\n\\exp\\!\\left(\n\\frac{\\lVert k_i \\rVert_2 \\cos\\theta_i}{T_q}\n\\right),\n\\end{aligned}","id":"QVEAXCLGVR"}" data-component-name="LatexBlockToDOM"></div><p>This may be undesirable in long-context settings, where a high-magnitude early key could remain disproportionately influential even after its relevance has faded. One interpretation is that scalar QK norm effectively avoids this, and makes the inner product truly capture <em>similarity</em> between a q and k without being confounded with anything else. </p><p><strong>Verify all QK norm weights are like this</strong></p><pre><code><code>for layer_id in range(model_config["num_hidden_layers"]): for kind in ("q", "k"): w = sd[f"model.decoder.layers.{layer_id}.self_attn.{kind}_norm.weight"] assert (w == w[0]).all(), "QK Norm entries are not all equal"</code></code></pre><p>We do RMSNorm in higher precision like usual because precision errors are accumulated through the reduction, potentially amplified through the division, and propagated to every element.</p><div class="latex-rendered" data-attrs="{"persistentExpression":"\\operatorname{RMSNorm}_w(x)\n=\n\\sqrt{d}\\left(\nw \\odot\n\\frac{x}{\\|x\\|_2 + \\epsilon}\n\\right)","id":"ZSUNKSMOIW"}" data-component-name="LatexBlockToDOM"></div><p><strong>RMSNorm</strong></p><pre><code><code>def rms(x, w = 1): return (w * x * (x.shape[-1] ** 0.5) / (torch.norm(x, dim=-1, keepdim=True, dtype=torch.float32) + model_config['rms_norm_eps'])).to(x.dtype)</code></code></pre><h3>Rotary position embeddings</h3><p>DiffusionGemma has two RoPE configs:</p><div class="captioned-image-container"><figure><a class="image-link image2" target="_blank" href="https://substackcdn.com/image/fetch/$s_!RjdQ!,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fd612f7c2-0a8a-4391-8bcd-d27bfa485eba_1167x249.png" data-component-name="Image2ToDOM"><div class="image2-inset"><picture><source type="image/webp" srcset="https://substackcdn.com/image/fetch/$s_!RjdQ!,w_424,c_limit,f_webp,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fd612f7c2-0a8a-4391-8bcd-d27bfa485eba_1167x249.png 424w, https://substackcdn.com/image/fetch/$s_!RjdQ!,w_848,c_limit,f_webp,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fd612f7c2-0a8a-4391-8bcd-d27bfa485eba_1167x249.png 848w, https://substackcdn.com/image/fetch/$s_!RjdQ!,w_1272,c_limit,f_webp,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fd612f7c2-0a8a-4391-8bcd-d27bfa485eba_1167x249.png 1272w, https://substackcdn.com/image/fetch/$s_!RjdQ!,w_1456,c_limit,f_webp,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fd612f7c2-0a8a-4391-8bcd-d27bfa485eba_1167x249.png 1456w" sizes="100vw"><img src="https://substackcdn.com/image/fetch/$s_!RjdQ!,w_1456,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fd612f7c2-0a8a-4391-8bcd-d27bfa485eba_1167x249.png" width="1167" height="249" data-attrs="{"src":"https://substack-post-media.s3.amazonaws.com/public/images/d612f7c2-0a8a-4391-8bcd-d27bfa485eba_1167x249.png","srcNoWatermark":null,"fullscreen":null,"imageSize":null,"height":249,"width":1167,"resizeWidth":null,"bytes":null,"alt":"Table: Layer type | Head width | RoPE base θ | Rotated fraction; Sliding attention | 256 | 10,000 | 100%; Full attention | 512 | 1,000,000 | 25%","title":null,"type":null,"href":null,"belowTheFold":true,"topImage":false,"internalRedirect":null,"isProcessing":false,"align":null,"offset":false}" class="sizing-normal" alt="Table: Layer type | Head width | RoPE base θ | Rotated fraction; Sliding attention | 256 | 10,000 | 100%; Full attention | 512 | 1,000,000 | 25%" title="Table: Layer type | Head width | RoPE base θ | Rotated fraction; Sliding attention | 256 | 10,000 | 100%; Full attention | 512 | 1,000,000 | 25%" srcset="https://substackcdn.com/image/fetch/$s_!RjdQ!,w_424,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fd612f7c2-0a8a-4391-8bcd-d27bfa485eba_1167x249.png 424w, https://substackcdn.com/image/fetch/$s_!RjdQ!,w_848,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fd612f7c2-0a8a-4391-8bcd-d27bfa485eba_1167x249.png 848w, https://substackcdn.com/image/fetch/$s_!RjdQ!,w_1272,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fd612f7c2-0a8a-4391-8bcd-d27bfa485eba_1167x249.png 1272w, https://substackcdn.com/image/fetch/$s_!RjdQ!,w_1456,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fd612f7c2-0a8a-4391-8bcd-d27bfa485eba_1167x249.png 1456w" sizes="100vw" loading="lazy"></picture><div></div></div></a></figure></div><p>Standard RoPE assigns each of the d/2 complex / rotary pairs one frequency (head dim d):</p><div class="latex-rendered" data-attrs="{"persistentExpression":"\\omega_j = \\theta^{-2j/d},\n\\qquad j=0,\\ldots,d/2-1","id":"URQCMIPGCX"}" data-component-name="LatexBlockToDOM"></div><p>This is what the sliding-attention layers use, but interestingly, DiffusionGemma’s full-attention layers use <strong>partial RoPE</strong>: after constructing the normal frequency vector, we then set all but the first 25% of its pair frequencies to zero, which will make <code>rotate</code> a no-op on those dimensions. Notably, we are doing truncating the full frequency spectrum, not the more intuitive option of rescaling it like we do in YaRN etc.</p><p>Partial RoPE enables a head to carry both explicitly position-sensitive features and features whose representation is not rotated as position changes.</p><p>Interestingly, the rows of <span>W_k</span> and <span>W_q</span> are stored arranged so that the projected vectors use split-half rotary pairs, <span>(t, t + d/2)</span>, rather than adjacent pairs, <span>(t, t+1)</span>, as in the original RoFormer paper. This makes the rotation easier to express and slightly more cache-friendly.</p><p>From equation 34 in RoFormer,</p><div class="latex-rendered" data-attrs="{"persistentExpression":"\\operatorname{RoPE}(x)\n= x\\cos(\\phi) + \\operatorname{rot}(x)\\sin(\\phi)","id":"UUHMEMOTSY"}" data-component-name="LatexBlockToDOM"></div><p>where <span>rot(x)</span> rotates each two-dimensional pair by 90 degrees counterclockwise, thus for a token at position p, the <span>i</span>-th pair of its k/q vector is rotated ccw by ω_i * p. Note our code’s <code>rotate(x)</code> is RoPE(x), not to be confused with rot(x).</p><p>In other words, instead of something like</p><pre><code><code>rot_x = np.empty_like(x) rot_x[..., 0::2] = -x[..., 1::2] rot_x[..., 1::2] = x[..., 0::2]</code></code></pre><p>for (t, t+1) pairs,</p><div class="latex-rendered" data-attrs="{"persistentExpression":"\\operatorname{rot}(x)\n= [-x_1, x_0, -x_3, x_2, \\ldots, -x_{d-1}, x_{d-2}],","id":"FFOMITAQQA"}" data-component-name="LatexBlockToDOM"></div><p>we have the more cache friendly</p><pre><code><code>rot_x = rearrange([-x[..., head_dim // 2 :], x[..., : head_dim // 2]], 'z ... d -> ... (z d)')</code></code></pre><div class="latex-rendered" data-attrs="{"persistentExpression":"\\operatorname{rot}(x)\n= [-x_{d/2},\\ldots,-x_{d-1},x_0,\\ldots,x_{d/2-1}]","id":"TJNSBTXJJO"}" data-component-name="LatexBlockToDOM"></div><p><strong>Precompute RoPE frequencies and rotate</strong></p><pre><code><code># precompute frequencies freq = {} # sliding window - regular rope freq['sliding_attention'] = model_config['rope_parameters']['sliding_attention']['rope_theta'] ** -(torch.arange(0, 1, 2 / model_config['head_dim'], dtype=torch.float32)) # full attention - partial rope freq['full_attention'] = model_config['rope_parameters']['full_attention']['rope_theta'] ** -(torch.arange(0, 1, 2 / model_config['global_head_dim'], dtype=torch.float32)) freq['full_attention'][int(model_config['rope_parameters']['full_attention']['partial_rotary_factor'] * len(freq['full_attention'])) : ] = 0 def rotate(x, layer_type, start_idx=0): # x [..., seq, head_dim]; rotate each (t, t + hd/2) pair ccw head_dim = x.shape[-1] rot_x = rearrange([-x[..., head_dim // 2 :], x[..., : head_dim // 2]], 'z ... d -> ... (z d)') pos = (torch.arange(x.shape[-2])[:, None] + start_idx) * torch.cat([freq[layer_type], freq[layer_type]]) # pos * (t, t + hd/2 pairs) return torch.cos(pos).to(x.dtype) * x + torch.sin(pos).to(x.dtype) * rot_x</code></code></pre><h2>Attention</h2><h4>Decode vs Encode</h4><p>In the <strong>decode</strong> stage, each canvas query attends to <code>[committed history | current canvas]</code>. It is:</p><ul><li><p>Similar to autoregressive prefill in the sense that it writes KV’s for several tokens at a time, autoregressive decode in the sense that it reads a context history of past KV’s</p></li><li><p>Similar to cross attention in that it attends to KV’s from the encoder, whisper-style, similar to self-attention in that it is non-casual, ViT-style</p></li></ul><p>In the <strong>encode</strong> stage, the attention behaves like an ordinary causal prefill. It can be thought of as the “verify” pass for speculative decoding with a 256-token draft, with all tokens accepted.</p><p>The same exact attention weights is used for both.</p><p>This code can be easily modified to support batch size greater than 1, but it does not support mixed encode/decode modes within the same batch. Dynamic per-sequence attention mode is supported in <a href="https://vllm-project.github.io/2026/06/10/diffusion-gemma">vLLM</a>:</p><div class="captioned-image-container"><figure><a class="image-link image2 is-viewable-img" target="_blank" href="https://substackcdn.com/image/fetch/$s_!-Y6t!,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F29eeaddd-061e-4bfc-9483-21728ae38cc8.svg" data-component-name="Image2ToDOM"><div class="image2-inset"><picture><source type="image/webp" srcset="https://substackcdn.com/image/fetch/$s_!-Y6t!,w_424,c_limit,f_webp,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F29eeaddd-061e-4bfc-9483-21728ae38cc8.svg 424w, https://substackcdn.com/image/fetch/$s_!-Y6t!,w_848,c_limit,f_webp,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F29eeaddd-061e-4bfc-9483-21728ae38cc8.svg 848w, https://substackcdn.com/image/fetch/$s_!-Y6t!,w_1272,c_limit,f_webp,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F29eeaddd-061e-4bfc-9483-21728ae38cc8.svg 1272w, https://substackcdn.com/image/fetch/$s_!-Y6t!,w_1456,c_limit,f_webp,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F29eeaddd-061e-4bfc-9483-21728ae38cc8.svg 1456w" sizes="100vw"><img src="https://substackcdn.com/image/fetch/$s_!-Y6t!,w_1456,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F29eeaddd-061e-4bfc-9483-21728ae38cc8.svg" width="1456" height="874" data-attrs="{"src":"https://substack-post-media.s3.amazonaws.com/public/images/29eeaddd-061e-4bfc-9483-21728ae38cc8.svg","srcNoWatermark":null,"fullscreen":null,"imageSize":null,"height":874,"width":1456,"resizeWidth":null,"bytes":null,"alt":"vLLM per-sequence causal-attention masking diagram","title":null,"type":null,"href":null,"belowTheFold":true,"topImage":false,"internalRedirect":null,"isProcessing":false,"align":null,"offset":false}" class="sizing-normal" alt="vLLM per-sequence causal-attention masking diagram" title="vLLM per-sequence causal-attention masking diagram" srcset="https://substackcdn.com/image/fetch/$s_!-Y6t!,w_424,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F29eeaddd-061e-4bfc-9483-21728ae38cc8.svg 424w, https://substackcdn.com/image/fetch/$s_!-Y6t!,w_848,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F29eeaddd-061e-4bfc-9483-21728ae38cc8.svg 848w, https://substackcdn.com/image/fetch/$s_!-Y6t!,w_1272,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F29eeaddd-061e-4bfc-9483-21728ae38cc8.svg 1272w, https://substackcdn.com/image/fetch/$s_!-Y6t!,w_1456,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F29eeaddd-061e-4bfc-9483-21728ae38cc8.svg 1456w" sizes="100vw" loading="lazy"></picture><div class="image-link-expand"><div class="pencraft pc-display-flex pc-gap-8 pc-reset"><button tabindex="0" type="button" class="pencraft pc-reset pencraft icon-container restack-image"><svg aria-hidden="true" width="20" height="20" viewBox="0 0 20 20" fill="none" stroke-width="1.5" stroke="var(--color-fg-primary)" stroke-linecap="round" stroke-linejoin="round" xmlns="http://www.w3.org/2000/svg"><g><path d="M2.53001 7.81595C3.49179 4.73911 6.43281 2.5 9.91173 2.5C13.1684 2.5 15.9537 4.46214 17.0852 7.23684L17.6179 8.67647M17.6179 8.67647L18.5002 4.26471M17.6179 8.67647L13.6473 6.91176M17.4995 12.1841C16.5378 15.2609 13.5967 17.5 10.1178 17.5C6.86118 17.5 4.07589 15.5379 2.94432 12.7632L2.41165 11.3235M2.41165 11.3235L1.5293 15.7353M2.41165 11.3235L6.38224 13.0882"></path></g></svg></button><button tabindex="0" type="button" class="pencraft pc-reset pencraft icon-container view-image"><svg xmlns="http://www.w3.org/2000/svg" width="20" height="20" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="lucide lucide-maximize2 lucide-maximize-2"><polyline points="15 3 21 3 21 9"></polyline><polyline points="9 21 3 21 3 15"></polyline><line x1="21" x2="14" y1="3" y2="10"></line><line x1="3" x2="10" y1="21" y2="14"></line></svg></button></div></div></div></a><figcaption class="image-caption">Per-sequence causal-attention masking. Source: vLLM.</figcaption></figure></div><p>Here, “denoise” means decode, “prefill” means encode on the prompt, and “accept” means encode on a denoised canvas.</p><p>For sliding layers, we use the approach in Google DeepMind’s <a href="https://github.com/google-deepmind/gemma/blob/195a9772b5cab4598dc8422780dd19dc0c03a284/gemma/diffusion/_sampler.py#L761-L783">JAX sampler</a>: every canvas token sees all other tokens inside the canvas plus the same prefix of KVs immediately before the canvas. Note that vLLM uses a different symmetric, 2K+1, per-token sliding window.</p><h4>Gemma-specific details</h4><p>Some notable differences from a standard LLM attention:</p><ol><li><p><strong>No explicit <span>1/√d</span> multiplier.</strong> The inner product is done directly. This can absorbed into the scalar Q/K normalization weights, which already does a fixed inverse temperature.</p></li><li><p>Gemma also applies an <strong>embedding scale</strong>, multiplying token embeddings by <span>√D</span> at the start of the residual stream. I’m not really sure why this is done.</p></li><li><p><strong>Value normalization.</strong> Weightless RMS-norm (i.e., unit norm) is applied to every value vector before it is used or stored in cache</p></li><li><p><strong>Hybrid architecture.</strong> 5 : 1 ratio of sliding window layers to global attention layer</p></li></ol><div class="captioned-image-container"><figure><a class="image-link image2" target="_blank" href="https://substackcdn.com/image/fetch/$s_!02ap!,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fb709ba41-ee79-4aa6-824e-ef1daf31e7b4_1142x249.png" data-component-name="Image2ToDOM"><div class="image2-inset"><picture><source type="image/webp" srcset="https://substackcdn.com/image/fetch/$s_!02ap!,w_424,c_limit,f_webp,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fb709ba41-ee79-4aa6-824e-ef1daf31e7b4_1142x249.png 424w, https://substackcdn.com/image/fetch/$s_!02ap!,w_848,c_limit,f_webp,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fb709ba41-ee79-4aa6-824e-ef1daf31e7b4_1142x249.png 848w, https://substackcdn.com/image/fetch/$s_!02ap!,w_1272,c_limit,f_webp,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fb709ba41-ee79-4aa6-824e-ef1daf31e7b4_1142x249.png 1272w, https://substackcdn.com/image/fetch/$s_!02ap!,w_1456,c_limit,f_webp,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fb709ba41-ee79-4aa6-824e-ef1daf31e7b4_1142x249.png 1456w" sizes="100vw"><img src="https://substackcdn.com/image/fetch/$s_!02ap!,w_1456,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fb709ba41-ee79-4aa6-824e-ef1daf31e7b4_1142x249.png" width="1142" height="249" data-attrs="{"src":"https://substack-post-media.s3.amazonaws.com/public/images/b709ba41-ee79-4aa6-824e-ef1daf31e7b4_1142x249.png","srcNoWatermark":null,"fullscreen":null,"imageSize":null,"height":249,"width":1142,"resizeWidth":null,"bytes":null,"alt":"Table: Layer type | Q heads | K/V heads | Head width | RoPE; Sliding | 16 | 8 | 256 | Full; Global | 16 | 2 | 512 | 25% partial","title":null,"type":null,"href":null,"belowTheFold":true,"topImage":false,"internalRedirect":null,"isProcessing":false,"align":null,"offset":false}" class="sizing-normal" alt="Table: Layer type | Q heads | K/V heads | Head width | RoPE; Sliding | 16 | 8 | 256 | Full; Global | 16 | 2 | 512 | 25% partial" title="Table: Layer type | Q heads | K/V heads | Head width | RoPE; Sliding | 16 | 8 | 256 | Full; Global | 16 | 2 | 512 | 25% partial" srcset="https://substackcdn.com/image/fetch/$s_!02ap!,w_424,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fb709ba41-ee79-4aa6-824e-ef1daf31e7b4_1142x249.png 424w, https://substackcdn.com/image/fetch/$s_!02ap!,w_848,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fb709ba41-ee79-4aa6-824e-ef1daf31e7b4_1142x249.png 848w, https://substackcdn.com/image/fetch/$s_!02ap!,w_1272,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fb709ba41-ee79-4aa6-824e-ef1daf31e7b4_1142x249.png 1272w, https://substackcdn.com/image/fetch/$s_!02ap!,w_1456,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fb709ba41-ee79-4aa6-824e-ef1daf31e7b4_1142x249.png 1456w" sizes="100vw" loading="lazy"></picture><div></div></div></a></figure></div><ol start="5"><li><p><strong>Shared global K/V projection.</strong> In full-attention layers, <code>W_v</code> = <code>W_k</code>. K and V are only different due to post-projection processing: K receives scalar-weight RMS normalization plus RoPE, while V receives unit RMS normalization and no RoPE. This is strange since K/V vectors are typically believed to live in different subspaces (see <a href="https://arxiv.org/abs/2606.04032">1</a>, <a href="https://transformer-circuits.pub/2021/framework/index.html">2</a> while K/Q live in the same subspace).</p></li></ol><h4>Optimizations and Implementation</h4><p>One potential optimization is to cache only the shared pre-normalization projection and reconstruct K and V when reading it, potentially reducing global-layer cache storage at the cost of additional computation. Combined with the hybrid global/sliding schedule, this could be attractive in long-context, KV-cache-bandwidth-bound regimes. The readable implementation below keeps separate K and V caches instead.</p><p>One implemented optimization is that, on encode, we can skip all computation after the K/V projection in the final layer, skipping most of the final layer and unembed work. Encode passes do not use the model’s final logits; they only need to incur the minimum set of computation required to obtain (and commit) correct K/V states in every layer.</p><p>In the current implementation with fixed-size statically-shaped cache, full-attention is only differentiated from sliding window through <code>kv_len</code>, thus behaves exactly like sliding-window attention with a <code>max_tot_tokens</code>-sized window:</p><pre><code><code>class AttentionBlock(torch.nn.Module): def __init__(self, layer_id): super().__init__() self.layer_id = layer_id self.layer_type = model_config['layer_types'][layer_id] self.W_q, self.W_k, self.W_o = [sd[f'model.decoder.layers.{layer_id}.self_attn.{item}_proj.weight'] for item in ['q', 'k', 'o']] self.q_norm, self.k_norm = [sd[f'model.decoder.layers.{layer_id}.self_attn.{item}_norm.weight'] for item in ['q', 'k']] self.pre_norm = sd[f'model.decoder.layers.{layer_id}.input_layernorm.weight'] self.post_norm = sd[f'model.decoder.layers.{layer_id}.post_attention_layernorm.weight'] self.q_heads = model_config['num_attention_heads'] if self.layer_type == 'full_attention': self.W_v = self.W_k # global layers share K = V self.kv_heads = model_config['num_global_key_value_heads'] self.head_dim = model_config['global_head_dim'] self.kv_len = max_tot_tokens else: assert(self.layer_type == 'sliding_attention') self.W_v = sd[f'model.decoder.layers.{layer_id}.self_attn.v_proj.weight'] self.kv_heads = model_config['num_key_value_heads'] self.head_dim = model_config['head_dim'] self.kv_len = model_config['sliding_window'] self.k_cache = torch.empty(self.kv_heads, self.kv_len, self.head_dim) # use statically shaped KV buffer self.v_cache = torch.empty(self.kv_heads, self.kv_len, self.head_dim) # KV's flow into here from left to right, FIFO, latest element is rightmost def forward(self, x, pos_idx, mode): assert mode in ['encode', 'decode'] L, D = x.shape resid_x = x.clone() x = rms(x, w=self.pre_norm) q, k, v = x @ self.W_q.T, x @ self.W_k.T, x @ self.W_v.T q, k, v = [rearrange(z, 'l (n h) -> n l h', h = self.head_dim) for z in [q, k, v]] q, k = rms(q, self.q_norm), rms(k, self.k_norm) # QK-norm per head: weight is [head_dim], normalize over each head's dims v = rms(v, 1) # v_norm: weightless, no rope q, k = rotate(q, self.layer_type, pos_idx), rotate(k, self.layer_type, pos_idx) # absolute positions pos_idx .. pos_idx + L k = torch.concat([self.k_cache[:, : pos_idx, :], k], axis=1) # attend to [committed history | current block] v = torch.concat([self.v_cache[:, : pos_idx, :], v], axis=1) # Note python automatically clips on the left to 0, on the right to shape[1] = kv_len if mode == "encode": # difference #1: writes/updates the kv cache # Actually, we don't have to put this in a branch, can also just do this on decode too, ok since we'll override with an encode at the end anyways self.k_cache[:, : pos_idx + L, :] = k[:, -self.kv_len :, :] # automatically clips self.v_cache[:, : pos_idx + L, :] = v[:, -self.kv_len :, :] if self.layer_id == model_config['num_hidden_layers'] - 1: return # encode optimization: notice we don't need to do the remaining computation after this q = rearrange(q, '(n gqa) qt h -> n gqa qt h', gqa = self.q_heads // self.kv_heads) # Fold GQA into an outer dim scores = einsum(q, k, 'n gqa qt h, n kt h -> n gqa qt kt').float() # no divide by sqrt(head dim); softmax in fp32 if mode == "encode": scores += torch.triu(torch.full(scores.shape, -torch.inf), diagonal = scores.shape[-1] - scores.shape[-2] + 1) # this applies a mask that looks like R2 in the vllm figure scores = torch.exp(scores - torch.amax(scores, axis=-1, keepdims=True)) scores = (scores / torch.sum(scores, axis=-1, keepdims=True)).to(x.dtype) x = einsum(scores, v, 'n gqa qt kt, n kt h -> n gqa qt h') x = rearrange(x, 'n gqa qt h -> qt (n gqa h)') res = x @ self.W_o.T res = rms(res, self.post_norm) return res + resid_x</code></code></pre><h2>MOE</h2><p>DiffusionGemma routes each token to 8 of 128 experts + one larger shared expert per forward pass.</p><p>Each expert is a standard gated GELU MLP:</p><div class="latex-rendered" data-attrs="{"persistentExpression":"\\operatorname{MLP}(x)\n= W_{\\text{down}}\\left[\n (W_{\\text{up}}x) \\odot \\operatorname{GELU}(W_{\\text{gate}}x)\n\\right]","id":"UGJERZWKQJ"}" data-component-name="LatexBlockToDOM"></div><p>HuggingFace stores the routed experts’ gate and up projections as a single matrix, which we explicitly split in the MLP module.</p><div class="captioned-image-container"><figure><a class="image-link image2 is-viewable-img" target="_blank" href="https://substackcdn.com/image/fetch/$s_!Vijg!,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F37e47175-5232-4c65-89b8-e14a67b1054b_3807x2905.png" data-component-name="Image2ToDOM"><div class="image2-inset"><picture><source type="image/webp" srcset="https://substackcdn.com/image/fetch/$s_!Vijg!,w_424,c_limit,f_webp,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F37e47175-5232-4c65-89b8-e14a67b1054b_3807x2905.png 424w, https://substackcdn.com/image/fetch/$s_!Vijg!,w_848,c_limit,f_webp,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F37e47175-5232-4c65-89b8-e14a67b1054b_3807x2905.png 848w, https://substackcdn.com/image/fetch/$s_!Vijg!,w_1272,c_limit,f_webp,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F37e47175-5232-4c65-89b8-e14a67b1054b_3807x2905.png 1272w, https://substackcdn.com/image/fetch/$s_!Vijg!,w_1456,c_limit,f_webp,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F37e47175-5232-4c65-89b8-e14a67b1054b_3807x2905.png 1456w" sizes="100vw"><img src="https://substackcdn.com/image/fetch/$s_!Vijg!,w_1456,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F37e47175-5232-4c65-89b8-e14a67b1054b_3807x2905.png" width="1456" height="1111" data-attrs="{"src":"https://substack-post-media.s3.amazonaws.com/public/images/37e47175-5232-4c65-89b8-e14a67b1054b_3807x2905.png","srcNoWatermark":null,"fullscreen":null,"imageSize":null,"height":1111,"width":1456,"resizeWidth":null,"bytes":null,"alt":"DiffusionGemma mixture-of-experts architecture diagram","title":null,"type":null,"href":null,"belowTheFold":true,"topImage":false,"internalRedirect":null,"isProcessing":false,"align":null,"offset":false}" class="sizing-normal" alt="DiffusionGemma mixture-of-experts architecture diagram" title="DiffusionGemma mixture-of-experts architecture diagram" srcset="https://substackcdn.com/image/fetch/$s_!Vijg!,w_424,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F37e47175-5232-4c65-89b8-e14a67b1054b_3807x2905.png 424w, https://substackcdn.com/image/fetch/$s_!Vijg!,w_848,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F37e47175-5232-4c65-89b8-e14a67b1054b_3807x2905.png 848w, https://substackcdn.com/image/fetch/$s_!Vijg!,w_1272,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F37e47175-5232-4c65-89b8-e14a67b1054b_3807x2905.png 1272w, https://substackcdn.com/image/fetch/$s_!Vijg!,w_1456,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F37e47175-5232-4c65-89b8-e14a67b1054b_3807x2905.png 1456w" sizes="100vw" loading="lazy"></picture><div class="image-link-expand"><div class="pencraft pc-display-flex pc-gap-8 pc-reset"><button tabindex="0" type="button" class="pencraft pc-reset pencraft icon-container restack-image"><svg aria-hidden="true" width="20" height="20" viewBox="0 0 20 20" fill="none" stroke-width="1.5" stroke="var(--color-fg-primary)" stroke-linecap="round" stroke-linejoin="round" xmlns="http://www.w3.org/2000/svg"><g><path d="M2.53001 7.81595C3.49179 4.73911 6.43281 2.5 9.91173 2.5C13.1684 2.5 15.9537 4.46214 17.0852 7.23684L17.6179 8.67647M17.6179 8.67647L18.5002 4.26471M17.6179 8.67647L13.6473 6.91176M17.4995 12.1841C16.5378 15.2609 13.5967 17.5 10.1178 17.5C6.86118 17.5 4.07589 15.5379 2.94432 12.7632L2.41165 11.3235M2.41165 11.3235L1.5293 15.7353M2.41165 11.3235L6.38224 13.0882"></path></g></svg></button><button tabindex="0" type="button" class="pencraft pc-reset pencraft icon-container view-image"><svg xmlns="http://www.w3.org/2000/svg" width="20" height="20" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="lucide lucide-maximize2 lucide-maximize-2"><polyline points="15 3 21 3 21 9"></polyline><polyline points="9 21 3 21 3 15"></polyline><line x1="21" x2="14" y1="3" y2="10"></line><line x1="3" x2="10" y1="21" y2="14"></line></svg></button></div></div></div></a><figcaption class="image-caption">MOE Block</figcaption></figure></div><h4>MOE Forward</h4><p>operates independently across the token axis. Each (of B × L many tokens) is routed to its own 8 experts. We could do this naively by looping over each token, but a better approach is to iterate over the experts instead: gather all tokens routed to each expert, do a single forward pass, then send results back.</p><p>Across multiple devices / for expert parallelism, this is usually implemented with an <code>AllToAll</code>-style collective rather than the slow explicit Python loop we have here (see the <a href="https://jax-ml.github.io/scaling-book/sharding/">JAX Scaling Book</a>).</p><h4>Routing Score</h4><p>DiffusionGemma performs routing in the standard way. Each expert receives a score proportional to</p><div class="latex-rendered" data-attrs="{"persistentExpression":"\\exp\\!\\left(\\langle x_{\\text{route}}, k_{\\text{expert}}\\rangle\\right),","id":"ZRHQAVIXRJ"}" data-component-name="LatexBlockToDOM"></div><p>where the expert keys are the rows of <code>W_router</code>. The top 8 scores are renormalized and multiplied by learned per-expert scales.</p><h4>RMSNorm Fusion</h4><p>There is an optimization we can do here to save a couple RMS-norm computations. Notice that we apply a pre-norm 1 (shared expert) + 8 (routed experts) times at the start of each MLP. Notice RMSNorm can be factored into two operations</p><div class="latex-rendered" data-attrs="{"persistentExpression":"\\operatorname{RMS}_w(x)\n= w \\odot \\operatorname{RMS}_1(x)","id":"CBJEHUETTE"}" data-component-name="LatexBlockToDOM"></div><p>, unit-normalize then multiply by w.</p><p>Instead of doing this naively, we can instead apply a single weightless <span>RMS_1(x)</span> at the start, and absorb each expert’s learned RMSNorm weight <span>w</span> into the columns of that expert’s gate and up projection matrix.</p><p><strong>MLP and mixture-of-experts block</strong></p><pre><code><code>class MLP(torch.nn.Module): def __init__(self, layer_id, expert_num, conditioning_mlp : bool = False): super().__init__() if(conditioning_mlp): self.pre_norm = sd['model.decoder.self_conditioning.pre_norm.weight'] self.W_up = sd['model.decoder.self_conditioning.up_proj.weight'] self.W_gate = sd['model.decoder.self_conditioning.gate_proj.weight'] self.W_down = sd['model.decoder.self_conditioning.down_proj.weight'] return if expert_num is None: self.pre_norm = sd[f'model.decoder.layers.{layer_id}.pre_feedforward_layernorm.weight'] self.W_up, self.W_gate, self.W_down = [sd[f'model.decoder.layers.{layer_id}.mlp.{item}_proj.weight'] for item in ['up', 'gate', 'down']] else: self.pre_norm = sd[f'model.decoder.layers.{layer_id}.pre_feedforward_layernorm_2.weight'] self.W_gate, self.W_up = rearrange(sd[f'model.decoder.layers.{layer_id}.experts.gate_up_proj'][expert_num], '(z intermed) D -> z intermed D', z=2) self.W_down = sd[f'model.decoder.layers.{layer_id}.experts.down_proj'][expert_num] def forward(self, x): # every MLP has a prenorm. Also, note since the MLP is a compsition of functions that maps the 0 vector to itself, the entire MLP also maps all 0s to all 0s, this is relevant for the self-conditioning MLP L, D = x.shape x = rms(x, self.pre_norm) a = x @ self.W_up.T b = F.gelu(x @ self.W_gate.T) res = (a * b) @ self.W_down.T assert(res.shape == (L, D)) # share expert, routed expert, and self-conditioning MLP all map (_, D) -> (_, D) return res class MOEBlock(torch.nn.Module): def __init__(self, layer_id): super().__init__() self.k_experts = model_config['top_k_experts'] self.num_experts = model_config['num_experts'] self.W_router = sd[f'model.decoder.layers.{layer_id}.router.proj.weight'] self.expert_scale = sd[f'model.decoder.layers.{layer_id}.router.per_expert_scale'] self.scale = sd[f'model.decoder.layers.{layer_id}.router.scale'] self.experts = [MLP(layer_id, e) for e in range(self.num_experts)] self.shared_expert = MLP(layer_id, None) self.post_norm_1 = sd[f'model.decoder.layers.{layer_id}.post_feedforward_layernorm_1.weight'] # applied on shared expert output self.post_norm_2 = sd[f'model.decoder.layers.{layer_id}.post_feedforward_layernorm_2.weight'] # applied on summed contribution from experts self.post_norm = sd[f'model.decoder.layers.{layer_id}.post_feedforward_layernorm.weight'] # the sum h1 (shared ) + h2 (routed sum), before residual add def forward(self, x): L, D = x.shape resid_x = x.clone() route_x = rms(x) * self.scale / (D ** 0.5) expert_scores = F.softmax((route_x @ self.W_router.T).float(), dim=-1) # (L, num_experts) top_k_scores, top_k_idx = torch.topk(expert_scores, self.k_experts, dim=-1) # (L, k_experts), (L, k_experts) top_k_scores = (top_k_scores / torch.sum(top_k_scores, dim=-1, keepdim=True) * self.expert_scale[top_k_idx]).to(x.dtype) res = rms(self.shared_expert(x), self.post_norm_1) # h1: dense branch, shape (L, D) moe_out = torch.zeros_like(x) for id, expert in zip(range(self.num_experts), self.experts): mask = torch.any(top_k_idx == id, dim = -1) # boolean mask (L, ) which tokens routed to expert_id mult = top_k_scores[mask][top_k_idx[mask] == id] # shape (L', ) where L' <= L is the number of tokens expert_id routed to moe_out[mask] += mult[:, None] * expert(x[mask]) # (L', D) += (L', 1) * (L', D) res = res + rms(moe_out, self.post_norm_2) # h2 normed once, then h1 + h2 return rms(res, self.post_norm) + resid_x </code></code></pre><h2>Putting Them together</h2><p>We now combine all our previous components into a single module. Unlike Llama, DiffusionGemma does a couple things different beyond merely interleaving attention and MoE blocks.</p><h4>Logit Softcapping</h4><p>Logit softcapping applies the following function to the final logits:</p><div class="latex-rendered" data-attrs="{"persistentExpression":"z_{\\text{capped}} = c\\tanh\\!\\left(\\frac{z}{c}\\right),","id":"GKUNHFVZBL"}" data-component-name="LatexBlockToDOM"></div><p>where <span>z</span> is an uncapped logit and <span>c</span> is the softcap value.</p><p>For <span>|z| ≪ c</span>, <span>tanh(z/c) ≈ z/c</span>, so small logits are essentially unchanged. As <span>z → ±∞</span>, <span>z_capped → ± c</span>, so large logits are “soft-capped”.</p><p>This is a smooth alternative to <code>torch.clip</code> or <code>torch.clamp</code> - unlike a hard cap, it remains differentiable everywhere.</p><div class="captioned-image-container"><figure><a class="image-link image2 is-viewable-img" target="_blank" href="https://substackcdn.com/image/fetch/$s_!2toU!,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F6bb55c4e-9df8-4896-8786-a7421703add9_646x430.jpeg" data-component-name="Image2ToDOM"><div class="image2-inset"><picture><source type="image/webp" srcset="https://substackcdn.com/image/fetch/$s_!2toU!,w_424,c_limit,f_webp,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F6bb55c4e-9df8-4896-8786-a7421703add9_646x430.jpeg 424w, https://substackcdn.com/image/fetch/$s_!2toU!,w_848,c_limit,f_webp,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F6bb55c4e-9df8-4896-8786-a7421703add9_646x430.jpeg 848w, https://substackcdn.com/image/fetch/$s_!2toU!,w_1272,c_limit,f_webp,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F6bb55c4e-9df8-4896-8786-a7421703add9_646x430.jpeg 1272w, https://substackcdn.com/image/fetch/$s_!2toU!,w_1456,c_limit,f_webp,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F6bb55c4e-9df8-4896-8786-a7421703add9_646x430.jpeg 1456w" sizes="100vw"><img src="https://substackcdn.com/image/fetch/$s_!2toU!,w_1456,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F6bb55c4e-9df8-4896-8786-a7421703add9_646x430.jpeg" width="646" height="430" data-attrs="{"src":"https://substack-post-media.s3.amazonaws.com/public/images/6bb55c4e-9df8-4896-8786-a7421703add9_646x430.jpeg","srcNoWatermark":null,"fullscreen":null,"imageSize":null,"height":430,"width":646,"resizeWidth":null,"bytes":null,"alt":"Graph of the standard hyperbolic tangent function","title":null,"type":null,"href":null,"belowTheFold":true,"topImage":false,"internalRedirect":null,"isProcessing":false,"align":null,"offset":false}" class="sizing-normal" alt="Graph of the standard hyperbolic tangent function" title="Graph of the standard hyperbolic tangent function" srcset="https://substackcdn.com/image/fetch/$s_!2toU!,w_424,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F6bb55c4e-9df8-4896-8786-a7421703add9_646x430.jpeg 424w, https://substackcdn.com/image/fetch/$s_!2toU!,w_848,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F6bb55c4e-9df8-4896-8786-a7421703add9_646x430.jpeg 848w, https://substackcdn.com/image/fetch/$s_!2toU!,w_1272,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F6bb55c4e-9df8-4896-8786-a7421703add9_646x430.jpeg 1272w, https://substackcdn.com/image/fetch/$s_!2toU!,w_1456,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F6bb55c4e-9df8-4896-8786-a7421703add9_646x430.jpeg 1456w" sizes="100vw" loading="lazy"></picture><div class="image-link-expand"><div class="pencraft pc-display-flex pc-gap-8 pc-reset"><button tabindex="0" type="button" class="pencraft pc-reset pencraft icon-container restack-image"><svg aria-hidden="true" width="20" height="20" viewBox="0 0 20 20" fill="none" stroke-width="1.5" stroke="var(--color-fg-primary)" stroke-linecap="round" stroke-linejoin="round" xmlns="http://www.w3.org/2000/svg"><g><path d="M2.53001 7.81595C3.49179 4.73911 6.43281 2.5 9.91173 2.5C13.1684 2.5 15.9537 4.46214 17.0852 7.23684L17.6179 8.67647M17.6179 8.67647L18.5002 4.26471M17.6179 8.67647L13.6473 6.91176M17.4995 12.1841C16.5378 15.2609 13.5967 17.5 10.1178 17.5C6.86118 17.5 4.07589 15.5379 2.94432 12.7632L2.41165 11.3235M2.41165 11.3235L1.5293 15.7353M2.41165 11.3235L6.38224 13.0882"></path></g></svg></button><button tabindex="0" type="button" class="pencraft pc-reset pencraft icon-container view-image"><svg xmlns="http://www.w3.org/2000/svg" width="20" height="20" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="lucide lucide-maximize2 lucide-maximize-2"><polyline points="15 3 21 3 21 9"></polyline><polyline points="9 21 3 21 3 15"></polyline><line x1="21" x2="14" y1="3" y2="10"></line><line x1="3" x2="10" y1="21" y2="14"></line></svg></button></div></div></div></a></figure></div><p>The standard tanh function is bounded between -1 and 1; multiplying by <span>c</span> changes the bounds to <span>[−c,c]</span>.</p><p>Note this is part of the model, and not the sampling process — not to be confused with the temperature schedule, which is applied later.</p><h4>Denoise vs Prefill</h4><p>The following are function signatures of encode vs decode:</p><pre><code><code># this commits canvas, (last time we) write KV cache def enc(self, pos_idx, logits) -> None: self._forward(pos_idx, logits, 0, mode="encode") # this maps (canvas_i, canvas_prob_i) -> (canvas_prob_i+1) def dec(self, pos_idx, logits, logit_probs) -> torch.tensor: return self._forward(pos_idx, logits, logit_probs, mode="decode")</code></code></pre><p>On encode, the embedded tokens are directly passed into the first layer.</p><p>In decode mode, the embedded tokens additionally incorporates a <em>self-conditioning</em> signal from the previous denoising step. Let <span>P∈ℝ^(L×V)</span> be the previous step’s post-softmax probability distribution over vocabs at each of L canvas positions, let <span>E∈ℝ^(V×D)</span> be the tied embedding table. Then, <code>P @ E</code> computes the expected token embeddings from the previous iteration, L independent convex sums with the weights stored in P. This is then passed through a small conditioning MLP, added to the canvas token embedding, and RMSNormed. Self-conditioning is helpful because for example,</p><ul><li><p>Knowing how confident the previous pass was in this token can inform us about how confident we should be in this pass</p></li><li><p>During sampling, the highest entropy, least confident tokens are replaced by a random token ID. Thus, self-conditioning allows the model to deduce which tokens are effectively <code>[MASK]</code> tokens by comparing the probability distribution input to the token input. <a class="footnote-anchor" data-component-name="FootnoteAnchorToDOM" id="footnote-anchor-2" href="#footnote-2" target="_self">2</a></p></li><li><p>Possibly gives a way for gradients to flow across denoising steps</p></li></ul><div class="captioned-image-container"><figure><a class="image-link image2 is-viewable-img" target="_blank" href="https://substackcdn.com/image/fetch/$s_!dI2O!,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F7abf126b-482c-41c8-9319-b0ab17f6409b_1890x2048.png" data-component-name="Image2ToDOM"><div class="image2-inset"><picture><source type="image/webp" srcset="https://substackcdn.com/image/fetch/$s_!dI2O!,w_424,c_limit,f_webp,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F7abf126b-482c-41c8-9319-b0ab17f6409b_1890x2048.png 424w, https://substackcdn.com/image/fetch/$s_!dI2O!,w_848,c_limit,f_webp,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F7abf126b-482c-41c8-9319-b0ab17f6409b_1890x2048.png 848w, https://substackcdn.com/image/fetch/$s_!dI2O!,w_1272,c_limit,f_webp,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F7abf126b-482c-41c8-9319-b0ab17f6409b_1890x2048.png 1272w, https://substackcdn.com/image/fetch/$s_!dI2O!,w_1456,c_limit,f_webp,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F7abf126b-482c-41c8-9319-b0ab17f6409b_1890x2048.png 1456w" sizes="100vw"><img src="https://substackcdn.com/image/fetch/$s_!dI2O!,w_1456,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F7abf126b-482c-41c8-9319-b0ab17f6409b_1890x2048.png" width="1456" height="1578" data-attrs="{"src":"https://substack-post-media.s3.amazonaws.com/public/images/7abf126b-482c-41c8-9319-b0ab17f6409b_1890x2048.png","srcNoWatermark":null,"fullscreen":null,"imageSize":null,"height":1578,"width":1456,"resizeWidth":null,"bytes":null,"alt":"Diagram of DiffusionGemma self-conditioning across denoising steps","title":null,"type":null,"href":null,"belowTheFold":true,"topImage":false,"internalRedirect":null,"isProcessing":false,"align":null,"offset":false}" class="sizing-normal" alt="Diagram of DiffusionGemma self-conditioning across denoising steps" title="Diagram of DiffusionGemma self-conditioning across denoising steps" srcset="https://substackcdn.com/image/fetch/$s_!dI2O!,w_424,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F7abf126b-482c-41c8-9319-b0ab17f6409b_1890x2048.png 424w, https://substackcdn.com/image/fetch/$s_!dI2O!,w_848,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F7abf126b-482c-41c8-9319-b0ab17f6409b_1890x2048.png 848w, https://substackcdn.com/image/fetch/$s_!dI2O!,w_1272,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F7abf126b-482c-41c8-9319-b0ab17f6409b_1890x2048.png 1272w, https://substackcdn.com/image/fetch/$s_!dI2O!,w_1456,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F7abf126b-482c-41c8-9319-b0ab17f6409b_1890x2048.png 1456w" sizes="100vw" loading="lazy"></picture><div class="image-link-expand"><div class="pencraft pc-display-flex pc-gap-8 pc-reset"><button tabindex="0" type="button" class="pencraft pc-reset pencraft icon-container restack-image"><svg aria-hidden="true" width="20" height="20" viewBox="0 0 20 20" fill="none" stroke-width="1.5" stroke="var(--color-fg-primary)" stroke-linecap="round" stroke-linejoin="round" xmlns="http://www.w3.org/2000/svg"><g><path d="M2.53001 7.81595C3.49179 4.73911 6.43281 2.5 9.91173 2.5C13.1684 2.5 15.9537 4.46214 17.0852 7.23684L17.6179 8.67647M17.6179 8.67647L18.5002 4.26471M17.6179 8.67647L13.6473 6.91176M17.4995 12.1841C16.5378 15.2609 13.5967 17.5 10.1178 17.5C6.86118 17.5 4.07589 15.5379 2.94432 12.7632L2.41165 11.3235M2.41165 11.3235L1.5293 15.7353M2.41165 11.3235L6.38224 13.0882"></path></g></svg></button><button tabindex="0" type="button" class="pencraft pc-reset pencraft icon-container view-image"><svg xmlns="http://www.w3.org/2000/svg" width="20" height="20" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="lucide lucide-maximize2 lucide-maximize-2"><polyline points="15 3 21 3 21 9"></polyline><polyline points="9 21 3 21 3 15"></polyline><line x1="21" x2="14" y1="3" y2="10"></line><line x1="3" x2="10" y1="21" y2="14"></line></svg></button></div></div></div></a><figcaption class="image-caption">Self-conditioning across denoising steps. Source: Maarten Grootendorst.</figcaption></figure></div><p>On the first decode step, the self-conditioning input is all 0s. Note that since the self-conditioning MLP is composed of functions that map all 0s to all 0s, the added contribution from self-conditioning to the input is 0 - this is equivalent to skipping self-conditioning for that step (no previous “self”).</p><h4>Layer Scalar</h4><p>Notice each block of computation in DiffusionGemma is <em>additive</em> to the residual stream:</p><div class="latex-rendered" data-attrs="{"persistentExpression":"x\n\\longmapsto\nx + F\\!\\left(\\operatorname{RMSNorm}(x)\\right)","id":"VHAMXVNVKF"}" data-component-name="LatexBlockToDOM"></div><p>where <span>F</span> is the core computation in that block, and <span>RMSNorm(x)</span> is that block’s pre-norm (both attention and MOE have this).</p><p>Thus due to the skip connections, the residual stream remains unnormalized through all 30 layers. This is perfectly fine if the activations truly live in <span>ℝ^n</span>, but unfortunately, they are bound to the subset of representation values of its datatype, which has a fixed dynamic range - upper and lower bounds of the feasible set, and anything outside will over/underflow to infinities. To ensure this doesn’t happen, quantization methods typically multiply by a scalar <span>S</span> before casting, for example:</p><div class="latex-rendered" data-attrs="{"persistentExpression":"X_{\\mathrm{FP8}}\n=\n\\operatorname{cast}_{\\mathrm{to FP8}}\n\\left(\n\\operatorname{clip}\n\\left(\nS \\cdot X,\\,\n-V_{\\max}^{\\mathrm{FP8}},\\,\nV_{\\max}^{\\mathrm{FP8}}\n\\right)\n\\right),","id":"EIFWAHVNHS"}" data-component-name="LatexBlockToDOM"></div><p>where</p><div class="latex-rendered" data-attrs="{"persistentExpression":"S\n=\n\\frac{\nV_{\\max}^{\\mathrm{FP8}}\n}{\n\\max\\left(\\lVert X\\rVert_{\\infty}, \\varepsilon\\right)\n}","id":"RKOJKWFRBF"}" data-component-name="LatexBlockToDOM"></div><p>to ensure values lie within the dynamic range <span>[−V_max, V_max]</span>.</p><p>A similar worry arises for the unnormalized residual stream, the magnitude of the activations may gradually grow through the layers (note however, the inputs to blocks remain well-conditioned due to the RMSNorm). To improve training stability, avoid overflow, and enable stable low-precision inference, DiffusionGemma adds a <em>layer scalar</em> which rescales the completed residual stream by a learned scalar <span>s_l</span> at the end of every layer:</p><div class="latex-rendered" data-attrs="{"persistentExpression":"x_{\\ell+1}\n=\ns_\\ell \\cdot \\operatorname{MoEBlock}_\\ell\\!\\left(\n\\operatorname{Attn}_\\ell(x_\\ell)\n\\right)","id":"FXUQQOPMHR"}" data-component-name="LatexBlockToDOM"></div><p>Printing out the layer scalars’ actual values, we see that they are all <span>< 1</span>, which makes sense according to our interpretation. Interestingly, it is a learned parameter — the model could learn to adaptively regulate the magnitude of its own residual stream throughout training.</p><p>Notice the RMSNorm is invariant to multiplying its input by a scalar (<span>RMSNorm(cx) = RMSNorm(x)</span>). Since every block receives input through an RMSNorm, the layer scaler has little benefit on the normalized input presented to each block. However, note it will affect the direction of the input after the first layer.</p><p>The released Hugging Face checkpoint stores separate encoder and decoder layer-scalar entries,</p><pre><code><code>sd[f"model.encoder.language_model.layers.{i}.layer_scalar"] and sd[f"model.decoder.layers.{i}.layer_scalar"]</code></code></pre><p>but their values are actually identical. In the code below, we have a single <code>self.layer_scalar</code>.</p><p><strong>Complete DiffusionGemma model</strong></p><pre><code><code>class DiffusionGemma(torch.nn.Module): # Any computation that utilizes parameters passes through here def __init__(self): super().__init__() self.W_embed = sd['model.decoder.embed_tokens.weight'] # also used as the unembedding matrix (tie_word_embeddings = True) self.attn_blocks = [AttentionBlock(i) for i in range(model_config['num_hidden_layers'])] self.moe_blocks = [MOEBlock(i) for i in range(model_config['num_hidden_layers'])] self.layer_scalar = [sd[f'model.encoder.language_model.layers.{i}.layer_scalar'] for i in range(model_config['num_hidden_layers'])] self.model_norm = sd['model.decoder.norm.weight'] # final / model norm self.embed_scale = torch.tensor(model_config['hidden_size'] ** 0.5) self.conditioning_MLP = MLP(layer_id=None, expert_num=None, conditioning_mlp=True) def _forward(self, pos_idx, logits, logit_probs, mode): # logit_probs = 0 <=> skip this path (no bias term anywhere) x = self.W_embed[logits] * self.embed_scale # (L, ) -> (L, D) # do self conditioning if decode if mode == "decode": condition_x = (logit_probs.to(x.dtype) @ self.W_embed) * self.embed_scale # (L, V) x (V, D) -> convex combination of vocab embeddings condition_x = self.conditioning_MLP(condition_x) x = rms(x + condition_x) # pass through all layers for i, (attn, moe) in enumerate(zip(self.attn_blocks, self.moe_blocks)): x = attn(x, pos_idx, mode) if x is None: # last encode layer wrote its KV cache and returned early; nothing else is needed return None x = moe(x) x = x * (self.layer_scalar)[i] # per-layer encoder/decoder scalar x = rms(x, self.model_norm) final_logits = (x @ self.W_embed.T).float() # (L, D) x (D, V) -> (L, V) return torch.tanh(final_logits / model_config['final_logit_softcapping']) * model_config['final_logit_softcapping'] # this commits canvas, (last time we) write KV cache def enc(self, pos_idx, logits) -> None: self._forward(pos_idx, logits, 0, mode="encode") # this maps (canvas_i, canvas_prob_i) -> (canvas_prob_i+1) def dec(self, pos_idx, logits, logit_probs) -> torch.tensor: return self._forward(pos_idx, logits, logit_probs, mode="decode")</code></code></pre><p>The following figure is a nice summary of what we’ve put together:</p><div class="captioned-image-container"><figure><a class="image-link image2 is-viewable-img" target="_blank" href="https://substackcdn.com/image/fetch/$s_!VwP4!,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fc91ee800-f0f4-4c88-a346-31e47490b429_3822x4093.png" data-component-name="Image2ToDOM"><div class="image2-inset"><picture><source type="image/webp" srcset="https://substackcdn.com/image/fetch/$s_!VwP4!,w_424,c_limit,f_webp,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fc91ee800-f0f4-4c88-a346-31e47490b429_3822x4093.png 424w, https://substackcdn.com/image/fetch/$s_!VwP4!,w_848,c_limit,f_webp,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fc91ee800-f0f4-4c88-a346-31e47490b429_3822x4093.png 848w, https://substackcdn.com/image/fetch/$s_!VwP4!,w_1272,c_limit,f_webp,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fc91ee800-f0f4-4c88-a346-31e47490b429_3822x4093.png 1272w, https://substackcdn.com/image/fetch/$s_!VwP4!,w_1456,c_limit,f_webp,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fc91ee800-f0f4-4c88-a346-31e47490b429_3822x4093.png 1456w" sizes="100vw"><img src="https://substackcdn.com/image/fetch/$s_!VwP4!,w_1456,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fc91ee800-f0f4-4c88-a346-31e47490b429_3822x4093.png" width="1456" height="1559" data-attrs="{"src":"https://substack-post-media.s3.amazonaws.com/public/images/c91ee800-f0f4-4c88-a346-31e47490b429_3822x4093.png","srcNoWatermark":null,"fullscreen":null,"imageSize":null,"height":1559,"width":1456,"resizeWidth":null,"bytes":null,"alt":"DiffusionGemma model architecture recap diagram","title":null,"type":null,"href":null,"belowTheFold":true,"topImage":false,"internalRedirect":null,"isProcessing":false,"align":null,"offset":false}" class="sizing-normal" alt="DiffusionGemma model architecture recap diagram" title="DiffusionGemma model architecture recap diagram" srcset="https://substackcdn.com/image/fetch/$s_!VwP4!,w_424,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fc91ee800-f0f4-4c88-a346-31e47490b429_3822x4093.png 424w, https://substackcdn.com/image/fetch/$s_!VwP4!,w_848,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fc91ee800-f0f4-4c88-a346-31e47490b429_3822x4093.png 848w, https://substackcdn.com/image/fetch/$s_!VwP4!,w_1272,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fc91ee800-f0f4-4c88-a346-31e47490b429_3822x4093.png 1272w, https://substackcdn.com/image/fetch/$s_!VwP4!,w_1456,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fc91ee800-f0f4-4c88-a346-31e47490b429_3822x4093.png 1456w" sizes="100vw" loading="lazy"></picture><div class="image-link-expand"><div class="pencraft pc-display-flex pc-gap-8 pc-reset"><button tabindex="0" type="button" class="pencraft pc-reset pencraft icon-container restack-image"><svg aria-hidden="true" width="20" height="20" viewBox="0 0 20 20" fill="none" stroke-width="1.5" stroke="var(--color-fg-primary)" stroke-linecap="round" stroke-linejoin="round" xmlns="http://www.w3.org/2000/svg"><g><path d="M2.53001 7.81595C3.49179 4.73911 6.43281 2.5 9.91173 2.5C13.1684 2.5 15.9537 4.46214 17.0852 7.23684L17.6179 8.67647M17.6179 8.67647L18.5002 4.26471M17.6179 8.67647L13.6473 6.91176M17.4995 12.1841C16.5378 15.2609 13.5967 17.5 10.1178 17.5C6.86118 17.5 4.07589 15.5379 2.94432 12.7632L2.41165 11.3235M2.41165 11.3235L1.5293 15.7353M2.41165 11.3235L6.38224 13.0882"></path></g></svg></button><button tabindex="0" type="button" class="pencraft pc-reset pencraft icon-container view-image"><svg xmlns="http://www.w3.org/2000/svg" width="20" height="20" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="lucide lucide-maximize2 lucide-maximize-2"><polyline points="15 3 21 3 21 9"></polyline><polyline points="9 21 3 21 3 15"></polyline><line x1="21" x2="14" y1="3" y2="10"></line><line x1="3" x2="10" y1="21" y2="14"></line></svg></button></div></div></div></a><figcaption class="image-caption">DiffusionGemma Model Diagram</figcaption></figure></div><h2>Sampling:</h2><div class="captioned-image-container"><figure><a class="image-link image2" target="_blank" href="https://substackcdn.com/image/fetch/$s_!itom!,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F69695187-c486-44e6-afba-35fe0c3392ce_1476x268.svg" data-component-name="Image2ToDOM"><div class="image2-inset"><picture><source type="image/webp" srcset="https://substackcdn.com/image/fetch/$s_!itom!,w_424,c_limit,f_webp,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F69695187-c486-44e6-afba-35fe0c3392ce_1476x268.svg 424w, https://substackcdn.com/image/fetch/$s_!itom!,w_848,c_limit,f_webp,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F69695187-c486-44e6-afba-35fe0c3392ce_1476x268.svg 848w, https://substackcdn.com/image/fetch/$s_!itom!,w_1272,c_limit,f_webp,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F69695187-c486-44e6-afba-35fe0c3392ce_1476x268.svg 1272w, https://substackcdn.com/image/fetch/$s_!itom!,w_1456,c_limit,f_webp,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F69695187-c486-44e6-afba-35fe0c3392ce_1476x268.svg 1456w" sizes="100vw"><img src="https://substackcdn.com/image/fetch/$s_!itom!,w_1456,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F69695187-c486-44e6-afba-35fe0c3392ce_1476x268.svg" width="1456" height="264" data-attrs="{"src":"https://substack-post-media.s3.amazonaws.com/public/images/69695187-c486-44e6-afba-35fe0c3392ce_1476x268.svg","srcNoWatermark":null,"fullscreen":null,"imageSize":null,"height":264,"width":1456,"resizeWidth":null,"bytes":null,"alt":"vLLM DiffusionGemma sampling loop","title":null,"type":null,"href":null,"belowTheFold":true,"topImage":false,"internalRedirect":null,"isProcessing":false,"align":null,"offset":false}" class="sizing-normal" alt="vLLM DiffusionGemma sampling loop" title="vLLM DiffusionGemma sampling loop" srcset="https://substackcdn.com/image/fetch/$s_!itom!,w_424,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F69695187-c486-44e6-afba-35fe0c3392ce_1476x268.svg 424w, https://substackcdn.com/image/fetch/$s_!itom!,w_848,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F69695187-c486-44e6-afba-35fe0c3392ce_1476x268.svg 848w, https://substackcdn.com/image/fetch/$s_!itom!,w_1272,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F69695187-c486-44e6-afba-35fe0c3392ce_1476x268.svg 1272w, https://substackcdn.com/image/fetch/$s_!itom!,w_1456,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F69695187-c486-44e6-afba-35fe0c3392ce_1476x268.svg 1456w" sizes="100vw" loading="lazy"></picture><div></div></div></a><figcaption class="image-caption">DiffusionGemma generation loop. Source: vLLM.</figcaption></figure></div><p>After the first encode on the input prompt, every 256 canvas of tokens incurs one encode pass and at most <code>max_denoising_steps</code> decode passes.</p><p>After obtaining a vocabulary distribution for every canvas position, sampling is more involved than in an autoregressive model.</p><p>We follow the sampling procedure explained <a href="https://newsletter.maartengrootendorst.com/p/a-visual-guide-to-diffusiongemma">here</a>, implemented <a href="https://github.com/google-deepmind/gemma/blob/main/gemma/diffusion/_sampler.py">here</a>. The meanings of the sampling parameters are documented <a href="https://ai.google.dev/gemma/docs/diffusiongemma">here</a>:</p><ol><li><p><strong>Apply temperature</strong></p><p>For decode pass <span>i</span> out of <span>N = max_denoising_steps,</span></p><div class="latex-rendered" data-attrs="{"persistentExpression":"t_i = t_{\\max} + \\frac{i}{N}\\left(t_{\\min}-t_{\\max}\\right),\n\\qquad i=0,\\ldots,N-1","id":"IUVSYKIPXM"}" data-component-name="LatexBlockToDOM"></div><p>Thus the temperature starts at <span>t_max</span> and decreases toward <span>t_min</span>, increasingly sharpening the distribution. In the i-th iteration, we work with the categorical distribution formed by <code>logits / t_i</code>.</p></li><li><p><strong>Compute entropy</strong></p><p>For every canvas position, we compute from its categorical distribution <span>p</span>,</p><div class="latex-rendered" data-attrs="{"persistentExpression":"H(p) = -\\sum_{v=1}^{V} p_v\\log p_v","id":"AVLXJTKHMI"}" data-component-name="LatexBlockToDOM"></div><p>Used as a measure of how uncertain the model is about the token at this position.</p></li><li><p><strong>Keep a low-entropy prefix and renoise the rest</strong></p><p>Sort position entropies so that</p><div class="latex-rendered" data-attrs="{"persistentExpression":"H_{(1)} \\le H_{(2)} \\le \\cdots \\le H_{(L)}","id":"YQSFZKHGYX"}" data-component-name="LatexBlockToDOM"></div><p>Accept the largest prefix ending at <span>k</span> such that</p><div class="latex-rendered" data-attrs="{"persistentExpression":"\\sum_{j=1}^{k-1} H_{(j)} \\le \\texttt{entropy_bound}","id":"GZDGBQLLQE"}" data-component-name="LatexBlockToDOM"></div><p>Then, replace every unaccepted position by a token drawn uniformly from the vocabulary.</p></li><li><p><strong>Check for early stopping</strong></p><p>Convergence is decided when the canvas is confident <em>and</em> stable:</p><ul><li><p>the argmax canvas has remained unchanged for <code>stability_threshold</code> (equals 1 here) previous canvases, and</p></li><li><p>the current canvas mean per-position entropy is below <code>confidence_threshold</code>.</p></li></ul></li></ol><p>On convergence or after N steps, we commit the latest canvas.</p><p><strong>Tokenizer and sampling configuration</strong></p><pre><code><code>tok = Tokenizer.from_file(f"{checkpoint}/tokenizer.json") prompt = """ What is the meaning of 67? """ chat = f"<bos><|turn>user\n{prompt}<turn|>\n<|turn>model\n" max_denoising_steps = gen_config['max_denoising_steps'] # decoder forward passes per canvas entropy_bound = gen_config['sampler_config']['entropy_bound'] # see formula in figure t_max, t_min = gen_config['t_max'], gen_config['t_min'] # linear schedule of temperatures t_max -> t_min across the steps confidence_threshold = gen_config['confidence_threshold'] # early-stop a canvas when argmax is stable (equal previous argmax canvas) and mean entropy < this tokens = torch.tensor(tok.encode(chat, add_special_tokens=False).ids) model = DiffusionGemma() </code></code></pre><h4>Lazy Sampling</h4><p>The usual way a denoising iteration is implemented:</p><blockquote><p>Model forward → retain distribution for self-conditioning → sample and renoise tokens</p></blockquote><p>Instead, we do sampling lazily, passing only the categorical distribution between steps and sampling only once we need it:</p><blockquote><p>Stored distribution → sample and renoise input tokens → model forward → new distribution</p></blockquote><p>The input into the first iteration is just all identical logits, which will generate a random canvas for us.</p><p>This is nice since:</p><ul><li><p>The only information we need to pass between steps is the self-conditioning input</p></li><li><p>Initialization falls out naturally from a uniform categorical distribution, can be folded into the first step</p></li><li><p>We can think of the decoder input as the previous canvas probs plus a source of randomness (from the sampling). Under this view, preparing the inputs (sampling) lazily is optimal. Additionally, it suggests an alternative interpretation of self-conditioning as the primary rather than auxiliary input:</p></li></ul><p>One way to think of the decoder is that it repeatedly transports our current distribution toward the target distribution, reminiscent of a <a href="https://arxiv.org/abs/2602.12233">flow-like process on the probability simplex</a>. So one can think about an alternate DiffusionGemma design that instead starts from a random position in the simplex at initialization, <code>Categorical(probs=torch.rand(L, V))</code> or <code>Categorical(logits=torch.rand(L, V))</code>, rather than “<code>Categorical(logits=torch.zeros(L, V))</code>”.</p><h4>Stage Then Commit</h4><p>We also intentionally give general functions for <code>denoise</code> which generates staged tokens; and <code>commit</code>, which takes in some staged tokens, generates their encoded KVs, and writes them to cache.</p><p><code>denoise(pos_idx)</code> repeatedly calls decode against the already committed <code>cache[:pos_idx]</code> on a fresh canvas at positions <code>[pos_idx : pos_idx + 256]</code>, and returns staged tokens sampled from the final canvas of logits. This function has no side effects on the KV cache, thus we can for example denoise multiple canvases at the same position and choose the best one. We can also place the canvas anywhere, as long as pos_idx <span>≤ </span> len(tokens) - canvas_len, since the correct slice of the KV cache will automatically be read.</p><p><code>commit(l, r, staged_tokens)</code> applies the casual encode and writes (or overwrites) the slice of KV cache from l to r, with the staged tokens’ KVs. This function allows arbitrary-lengthed blocks, thus we can for example commit only a confident prefix of the current canvas. We can also commit tokens at any positions, but if an earlier region is overwritten, all later tokens are invalidated because their cached states depended on the old prefix.</p><p>This design is to make alternative schedules easier to study, including overlapping canvases, revising an earlier block, selecting among several denoised candidates, or advancing only part of a canvas.</p><p><strong>Denoise, stage, commit, and generate</strong></p><pre><code><code>def denoise(pos_idx): # returns staged_tokens assert pos_idx + canvas_len <= len(tokens) # must be length canvas_len (what if it wasn't fixed? analyze how casual the self attention is) plotting_data = [] t = t_max t_step = (t_min - t_max) / max_denoising_steps last_canvas = Categorical(logits = torch.ones((canvas_len, V))) for step in trange(max_denoising_steps): # renoise last_canvas sH, sidx = last_canvas.entropy().sort(-1) accepted = torch.zeros_like(sH, dtype=torch.bool).scatter(-1, sidx, sH.cumsum(-1) - sH <= entropy_bound) last_canvas_noised = torch.where(accepted, last_canvas.sample(), torch.randint(0, V, (canvas_len,))) # pass in the noised tokens, but un-noised normalized probs (all-zero probs on step 0: the conditioning path maps 0 to 0) canvas = model.dec(pos_idx, last_canvas_noised, last_canvas.probs if step != 0 else torch.zeros(canvas_len, V)) / t; plotting_data.append(canvas.detach().cpu()) canvas = Categorical(logits = canvas) if torch.mean(canvas.entropy()) < confidence_threshold and (canvas.logits.argmax(dim=-1) == last_canvas.logits.argmax(dim=-1)).all(): return canvas.sample(), plotting_data t += t_step last_canvas = canvas assert False, f"Denoising not finished after {max_denoising_steps} steps" def commit(l, r, staged_tokens): global tokens assert len(staged_tokens) == r - l + 1 if(r + 1 < len(tokens)): # commits staged_tokens print(f"Invalidating {len(tokens) - (r+1)} tokens") tokens = tokens[:r+1] tokens[l:] = staged_tokens model.enc(l, staged_tokens) def new_canvas(): global tokens if len(tokens) + canvas_len > max_tot_tokens: return False nxt = torch.randint(V, (canvas_len,)) tokens = torch.concat([tokens, nxt]) return True commit(0, len(tokens) - 1, tokens) # prefill: encode the chat prompt into the KV cache pos_idx = len(tokens) plotting_data = [] while new_canvas(): staged_tokens, cur_data = denoise(pos_idx) commit(pos_idx, pos_idx + canvas_len - 1, staged_tokens) pos_idx = len(tokens) print('=' * 50) print('Final Canvas') print(tok.decode(tokens.tolist())) print('=' * 50) plotting_data.append(cur_data) if torch.isin(staged_tokens, torch.tensor(gen_config['eos_token_id'])).any(): break</code></code></pre><p>Output:</p><div class="highlighted_code_block" data-attrs="{"language":"plaintext","nodeId":"c115dcde-751d-4808-9f27-bb3df9ee54cc"}" data-component-name="HighlightedCodeBlockToDOM"><pre class="shiki"><code class="language-plaintext">... ================================================== Final Canvas user What is the meaning of 67? model thought The meaning of **67** depends entirely on the context in which it is used (science, mathematics, pop culture, etc.). Here are common interpretations: ### 1. Mathematics * **Prime Number:** 67 is a prime number, meaning it can only be divided by 1 and itself. * **Lucky Prime:** It is considered a "lucky prime." * **Sum of Primes:** It is the sum of five consecutive prime numbers. ### 2. Science and Astronomy * **Atomic Number:** 67 is the atomic number of **Holmium (Ho)**, a rare earth element belonging to the lanthanide series. * **Astronomy:** Messier object 67 (M67) is an open star cluster in the constellation of Virgo. ### 3. Culture and Slang * **The "67" Connection:** In the UK, "67" is a well-known drill music group from Brixton, London. * **Age:** In many countries, 67 is considered the standard age for full retirement eligibility or social security. ### 4. Numerology and Spirituality * In numerology, the number 67 is often associated with combining the energies of **6** (home, stability, and responsibility) and **7** (spirituality, intuition, and inner wisdom). It is often interpreted as a sign of building practical foundations through spiritual growth. ### 5. Other Uses * **Country Code:** +67 is not a complete country code, but codes starting with +67 are used in various regions (like +670 for East Timor or +679 for American Samoa). **Is there a specific area (like a dream, a song, or a math problem) where you saw this number?** Providing more context can help me give you a more specific answer. ==================================================</code></pre></div><h4>Preliminary Analysis</h4><p>Using the saved <code>plotting_data</code> from denoise, we can make some figures of the denoising process over time:</p><div class="native-video-embed" data-component-name="VideoPlaceholder" data-attrs="{"mediaUploadId":"075c4d27-aa22-454a-9e44-9b42bb6be138","duration":null}"></div><p>Note the cells are laid out in reading order. </p><p>Without being explicitly trained to do so, it appears the model is approximately finalizing canvas tokens from left to right. Plotting two more figures:</p><div class="native-video-embed" data-component-name="VideoPlaceholder" data-attrs="{"mediaUploadId":"7d5cbd3d-2a32-4717-bfec-6726acf03f98","duration":null}"></div><div class="native-video-embed" data-component-name="VideoPlaceholder" data-attrs="{"mediaUploadId":"4d0e72f2-5c95-4c50-837d-8b8bef2ab238","duration":null}"></div><p>Further confirms our suspicions. </p><p>At every step, the model seems to mostly focus its efforts on figuring out the next causal rolling window of tokens — it “thinks” casually. Is this a general behavior the model has learned? Let’s try a different prompt, solving a <a href="https://en.wikipedia.org/wiki/Magic_square">magic square</a> puzzle:</p><p></p><div class="native-video-embed" data-component-name="VideoPlaceholder" data-attrs="{"mediaUploadId":"c10e92f4-b671-4d4a-9a36-3ad881eaca1b","duration":null}"></div><p style="text-align: center;">Canvas 2</p><div class="native-video-embed" data-component-name="VideoPlaceholder" data-attrs="{"mediaUploadId":"47a1ef0c-d6a7-4087-b6bb-250757ea29d1","duration":null}"></div><p style="text-align: center;">Canvas 3, More visualizations available <a href="https://github.com/timothygao8710/The-Annotated-DiffusionGemma">here</a>.</p><p>Interestingly, the prose tokens seem to exhibit the same causal pattern, but the magic square tokens do not. This makes intuitive sense. When humans solve these puzzles, the next value that’s easiest to deduce is not necessarily the next one in autoregressive order. This motivates perhaps a better explanation for how DiffusionGemma “thinks” during denoising — the “easiest”, lowest-entropy tokens are fixed / finalized / accepted first, which in turn unlocks / makes other tokens “easier” for the next denoising step. <a class="footnote-anchor" data-component-name="FootnoteAnchorToDOM" id="footnote-anchor-3" href="#footnote-3" target="_self">3</a></p><p>Note this mirrors what we do during sampling: all but a prefix of lowest entropy tokens is renoised. For prose and chain-of-thought reasoning, the easiest tokens to resolve next often coincide with the earliest tokens in autoregressive order.</p><p>However for both prompts, once a token is accepted (i.e. low-entropy and not renoised), it’s extremely unlikely to be renoised in the future:</p><div class="captioned-image-container"><figure><a class="image-link image2 is-viewable-img" target="_blank" href="https://substackcdn.com/image/fetch/$s_!pd4q!,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fba9b8f5e-0dec-4f5e-a859-5de8c3a00001_2450x1550.png" data-component-name="Image2ToDOM"><div class="image2-inset"><picture><source type="image/webp" srcset="https://substackcdn.com/image/fetch/$s_!pd4q!,w_424,c_limit,f_webp,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fba9b8f5e-0dec-4f5e-a859-5de8c3a00001_2450x1550.png 424w, https://substackcdn.com/image/fetch/$s_!pd4q!,w_848,c_limit,f_webp,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fba9b8f5e-0dec-4f5e-a859-5de8c3a00001_2450x1550.png 848w, https://substackcdn.com/image/fetch/$s_!pd4q!,w_1272,c_limit,f_webp,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fba9b8f5e-0dec-4f5e-a859-5de8c3a00001_2450x1550.png 1272w, https://substackcdn.com/image/fetch/$s_!pd4q!,w_1456,c_limit,f_webp,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fba9b8f5e-0dec-4f5e-a859-5de8c3a00001_2450x1550.png 1456w" sizes="100vw"><img src="https://substackcdn.com/image/fetch/$s_!pd4q!,w_1456,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fba9b8f5e-0dec-4f5e-a859-5de8c3a00001_2450x1550.png" width="1456" height="921" data-attrs="{"src":"https://substack-post-media.s3.amazonaws.com/public/images/ba9b8f5e-0dec-4f5e-a859-5de8c3a00001_2450x1550.png","srcNoWatermark":null,"fullscreen":null,"imageSize":null,"height":921,"width":1456,"resizeWidth":null,"bytes":null,"alt":"canvas_00_acceptance_transitions.png","title":null,"type":null,"href":null,"belowTheFold":true,"topImage":false,"internalRedirect":null,"isProcessing":false,"align":null,"offset":false}" class="sizing-normal" alt="canvas_00_acceptance_transitions.png" title="canvas_00_acceptance_transitions.png" srcset="https://substackcdn.com/image/fetch/$s_!pd4q!,w_424,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fba9b8f5e-0dec-4f5e-a859-5de8c3a00001_2450x1550.png 424w, https://substackcdn.com/image/fetch/$s_!pd4q!,w_848,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fba9b8f5e-0dec-4f5e-a859-5de8c3a00001_2450x1550.png 848w, https://substackcdn.com/image/fetch/$s_!pd4q!,w_1272,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fba9b8f5e-0dec-4f5e-a859-5de8c3a00001_2450x1550.png 1272w, https://substackcdn.com/image/fetch/$s_!pd4q!,w_1456,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fba9b8f5e-0dec-4f5e-a859-5de8c3a00001_2450x1550.png 1456w" sizes="100vw" loading="lazy"></picture><div class="image-link-expand"><div class="pencraft pc-display-flex pc-gap-8 pc-reset"><button tabindex="0" type="button" class="pencraft pc-reset pencraft icon-container restack-image"><svg aria-hidden="true" width="20" height="20" viewBox="0 0 20 20" fill="none" stroke-width="1.5" stroke="var(--color-fg-primary)" stroke-linecap="round" stroke-linejoin="round" xmlns="http://www.w3.org/2000/svg"><g><path d="M2.53001 7.81595C3.49179 4.73911 6.43281 2.5 9.91173 2.5C13.1684 2.5 15.9537 4.46214 17.0852 7.23684L17.6179 8.67647M17.6179 8.67647L18.5002 4.26471M17.6179 8.67647L13.6473 6.91176M17.4995 12.1841C16.5378 15.2609 13.5967 17.5 10.1178 17.5C6.86118 17.5 4.07589 15.5379 2.94432 12.7632L2.41165 11.3235M2.41165 11.3235L1.5293 15.7353M2.41165 11.3235L6.38224 13.0882"></path></g></svg></button><button tabindex="0" type="button" class="pencraft pc-reset pencraft icon-container view-image"><svg xmlns="http://www.w3.org/2000/svg" width="20" height="20" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="lucide lucide-maximize2 lucide-maximize-2"><polyline points="15 3 21 3 21 9"></polyline><polyline points="9 21 3 21 3 15"></polyline><line x1="21" x2="14" y1="3" y2="10"></line><line x1="3" x2="10" y1="21" y2="14"></line></svg></button></div></div></div></a></figure></div><p style="text-align: center;">For 67 Prompt</p><div class="captioned-image-container"><figure><a class="image-link image2 is-viewable-img" target="_blank" href="https://substackcdn.com/image/fetch/$s_!fs9I!,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F185a5d80-2108-4056-b992-41cd4cce4a2b_2178x1378.png" data-component-name="Image2ToDOM"><div class="image2-inset"><picture><source type="image/webp" srcset="https://substackcdn.com/image/fetch/$s_!fs9I!,w_424,c_limit,f_webp,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F185a5d80-2108-4056-b992-41cd4cce4a2b_2178x1378.png 424w, https://substackcdn.com/image/fetch/$s_!fs9I!,w_848,c_limit,f_webp,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F185a5d80-2108-4056-b992-41cd4cce4a2b_2178x1378.png 848w, https://substackcdn.com/image/fetch/$s_!fs9I!,w_1272,c_limit,f_webp,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F185a5d80-2108-4056-b992-41cd4cce4a2b_2178x1378.png 1272w, https://substackcdn.com/image/fetch/$s_!fs9I!,w_1456,c_limit,f_webp,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F185a5d80-2108-4056-b992-41cd4cce4a2b_2178x1378.png 1456w" sizes="100vw"><img src="https://substackcdn.com/image/fetch/$s_!fs9I!,w_1456,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F185a5d80-2108-4056-b992-41cd4cce4a2b_2178x1378.png" width="1456" height="921" data-attrs="{"src":"https://substack-post-media.s3.amazonaws.com/public/images/185a5d80-2108-4056-b992-41cd4cce4a2b_2178x1378.png","srcNoWatermark":null,"fullscreen":null,"imageSize":null,"height":921,"width":1456,"resizeWidth":null,"bytes":null,"alt":"canvas_02_acceptance_transitions.png","title":null,"type":null,"href":null,"belowTheFold":true,"topImage":false,"internalRedirect":null,"isProcessing":false,"align":null,"offset":false}" class="sizing-normal" alt="canvas_02_acceptance_transitions.png" title="canvas_02_acceptance_transitions.png" srcset="https://substackcdn.com/image/fetch/$s_!fs9I!,w_424,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F185a5d80-2108-4056-b992-41cd4cce4a2b_2178x1378.png 424w, https://substackcdn.com/image/fetch/$s_!fs9I!,w_848,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F185a5d80-2108-4056-b992-41cd4cce4a2b_2178x1378.png 848w, https://substackcdn.com/image/fetch/$s_!fs9I!,w_1272,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F185a5d80-2108-4056-b992-41cd4cce4a2b_2178x1378.png 1272w, https://substackcdn.com/image/fetch/$s_!fs9I!,w_1456,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F185a5d80-2108-4056-b992-41cd4cce4a2b_2178x1378.png 1456w" sizes="100vw" loading="lazy"></picture><div class="image-link-expand"><div class="pencraft pc-display-flex pc-gap-8 pc-reset"><button tabindex="0" type="button" class="pencraft pc-reset pencraft icon-container restack-image"><svg aria-hidden="true" width="20" height="20" viewBox="0 0 20 20" fill="none" stroke-width="1.5" stroke="var(--color-fg-primary)" stroke-linecap="round" stroke-linejoin="round" xmlns="http://www.w3.org/2000/svg"><g><path d="M2.53001 7.81595C3.49179 4.73911 6.43281 2.5 9.91173 2.5C13.1684 2.5 15.9537 4.46214 17.0852 7.23684L17.6179 8.67647M17.6179 8.67647L18.5002 4.26471M17.6179 8.67647L13.6473 6.91176M17.4995 12.1841C16.5378 15.2609 13.5967 17.5 10.1178 17.5C6.86118 17.5 4.07589 15.5379 2.94432 12.7632L2.41165 11.3235M2.41165 11.3235L1.5293 15.7353M2.41165 11.3235L6.38224 13.0882"></path></g></svg></button><button tabindex="0" type="button" class="pencraft pc-reset pencraft icon-container view-image"><svg xmlns="http://www.w3.org/2000/svg" width="20" height="20" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="lucide lucide-maximize2 lucide-maximize-2"><polyline points="15 3 21 3 21 9"></polyline><polyline points="9 21 3 21 3 15"></polyline><line x1="21" x2="14" y1="3" y2="10"></line><line x1="3" x2="10" y1="21" y2="14"></line></svg></button></div></div></div></a></figure></div><p style="text-align: center;">For Magic Square Prompt</p><p>Looking at the top right graph, almost no tokens are “re-masked” for both prompts. Google <a href="https://developers.googleblog.com/diffusiongemma-the-developer-guide/">claims</a> that one of the benefits of DiffusionGemma’s uniform state diffusion over traditional masked diffusion (e.g., LLaDA, Dream) is error correction via re-noising, but from our analysis — how often does this actually happen in practice?</p><h2>References</h2><ul><li><p>O’Donoghue, Brendan, and Sebastian Flennerhag. “<a href="https://blog.google/innovation-and-ai/technology/developers-tools/diffusion-gemma-faster-text-generation/">DiffusionGemma: 4x faster text generation</a>.” Google, 2026.</p></li><li><p>Ballantyne, Ian, and Omar Sanseviero. “<a href="https://developers.googleblog.com/diffusiongemma-the-developer-guide/">DiffusionGemma: The Developer Guide</a>.” Google Developers Blog, 2026.</p></li><li><p>Google AI for Developers. “<a href="https://ai.google.dev/gemma/docs/diffusiongemma">DiffusionGemma model overview</a>.” 2026.</p></li><li><p>Google AI for Developers. “<a href="https://ai.google.dev/gemma/docs/diffusiongemma/explained">Diffusion in Text Generation Explained</a>.” 2026.</p></li><li><p>Google DeepMind. “<a href="https://huggingface.co/google/diffusiongemma-26B-A4B-it">DiffusionGemma 26B-A4B-IT model card</a>.” Hugging Face, 2026.</p></li><li><p>Google DeepMind. “<a href="https://github.com/google-deepmind/gemma/blob/195a9772b5cab4598dc8422780dd19dc0c03a284/gemma/diffusion/_sampler.py">DiffusionGemma reference sampler</a>.” <em>Gemma</em> repository, 2026.</p></li><li><p>vLLM Team and Google DeepMind Team. “<a href="https://vllm-project.github.io/2026/06/10/diffusion-gemma">DiffusionGemma: The First Diffusion LLM (dLLM) Natively Supported in vLLM</a>.” vLLM Blog, 2026.</p></li><li><p>Grootendorst, Maarten. “<a href="https://newsletter.maartengrootendorst.com/p/a-visual-guide-to-diffusiongemma">A Visual Guide to DiffusionGemma</a>.” <em>Exploring Language Models</em>, 2026.</p></li><li><p>Huang, Austin, Suraj Subramanian, Jonathan Sum, Khalid Almubarak, and Stella Biderman. “<a href="https://nlp.seas.harvard.edu/annotated-transformer/">The Annotated Transformer</a>.” Harvard NLP, 2022. Original version by Alexander M. Rush.</p></li><li><p>Peng, Bowen, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. “<a href="https://arxiv.org/abs/2309.00071">YaRN: Efficient Context Window Extension of Large Language Models</a>.” arXiv:2309.00071, 2023.</p></li><li><p>Su, Jianlin, et al. “<a href="https://arxiv.org/abs/2104.09864">RoFormer: Enhanced Transformer with Rotary Position Embedding</a>.” arXiv:2104.09864, 2021.</p></li><li><p>Kayyam, Ali, Anusha Madan Gopal, and M. Anthony Lewis. “<a href="https://arxiv.org/abs/2606.04032">Do Transformers Need Three Projections? Systematic Study of QKV Variants</a>.” arXiv:2606.04032, 2026.</p></li><li><p>Elhage, Nelson, et al. “<a href="https://transformer-circuits.pub/2021/framework/index.html">A Mathematical Framework for Transformer Circuits</a>.” Transformer Circuits Thread, 2021.</p></li><li><p>Austin, Jacob, et al. “<a href="https://jax-ml.github.io/scaling-book/sharding/">Sharded Matrices and How to Multiply Them</a>.” <em>How To Scale Your Model</em>, 2025.</p></li><li><p>Roos, Daan, et al. “<a href="https://arxiv.org/abs/2602.12233">Categorical Flow Maps</a>.” arXiv:2602.12233, 2026.</p></li><li><p>Nie, Shen, et al. “<a href="https://arxiv.org/abs/2502.09992">Large Language Diffusion Models</a>.” arXiv:2502.09992, 2025.</p></li><li><p>Ye, Jiacheng, et al. “<a href="https://arxiv.org/abs/2508.15487">Dream 7B: Diffusion Large Language Models</a>.” arXiv:2508.15487, 2025.</p></li></ul><div class="footnote" data-component-name="FootnoteToDOM"><a id="footnote-1" href="#footnote-anchor-1" class="footnote-number" contenteditable="false" target="_self">1</a><div class="footnote-content"><p>Its main comparative advantage is higher arithmetic intensity from being able to decode a 256-token canvas in parallel, amortizing KV cache and model weight load across positions, which shines for local, low-concurrency, high-interactivity workloads </p></div></div><div class="footnote" data-component-name="FootnoteToDOM"><a id="footnote-2" href="#footnote-anchor-2" class="footnote-number" contenteditable="false" target="_self">2</a><div class="footnote-content"><p>However, unlike traditional Masked Language Diffusion models, DiffusionGemma can still mask out a previously un-masked token if it’s no longer confident in it (i.e., the same token position is now assigned high entropy).</p></div></div><div class="footnote" data-component-name="FootnoteToDOM"><a id="footnote-3" href="#footnote-anchor-3" class="footnote-number" contenteditable="false" target="_self">3</a><div class="footnote-content"><p>Reminiscent of Kahn’s algorithm, this <a href="https://codeforces.com/contest/2245/problem/D2">Codeforces problem</a>, and successive <a href="https://www2.eecs.berkeley.edu/Pubs/TechRpts/2019/Archive/EECS-2019-149.pdf#page=36">interference cancellation</a></p></div></div>

Original Article

View Cached Full Text

Cached at: 08/11/26, 01:44 AM

# DiffusionGemma Explained Source: [https://mlberkeley.substack.com/p/the-annotated-diffusiongemma](https://mlberkeley.substack.com/p/the-annotated-diffusiongemma) Google recently released[DiffusionGemma](https://blog.google/innovation-and-ai/technology/developers-tools/diffusion-gemma-faster-text-generation/), a 26B open\-weight uniform state diffusion language model ~4× faster than Gemma 4 for certain workloads\.[1](https://mlberkeley.substack.com/p/the-annotated-diffusiongemma#footnote-1) As an experimental model, it doesn’t have a whitepaper, and the[officially](https://www.google.com/url?q=https%3A%2F%2Fdevelopers.googleblog.com%2Fen%2Fdiffusiongemma-the-developer-guide%2F)[released](https://www.google.com/url?q=https%3A%2F%2Fblog.google%2Finnovation-and-ai%2Ftechnology%2Fdevelopers-tools%2Fdiffusion-gemma-faster-text-generation%2F)[docs](https://www.google.com/url?q=https%3A%2F%2Fai.google.dev%2Fgemma%2Fdocs%2Fdiffusiongemma%2Fexplained)are pretty brief and high\-level\.This blog post, based on the[colab notebook](https://colab.research.google.com/drive/17egFGmboDkhU6duNQHvjB1J0oiKslOMl?usp=sharing), attempts to fill in the gaps\. It walks through an annotated from\-scratch reimplementation derived from its[open weights](https://huggingface.co/google/diffusiongemma-26B-A4B-it), following the style of the awesome[Annotated Transformer](https://nlp.seas.harvard.edu/annotated-transformer/)\. Along the way, we explain some of the more opaque details — behind the model architecture, self\-conditioning, the sampling procedure, how encode/decode weight sharing is actually implemented\. Through closely examining the model config and weights, we also uncover some interesting design choices — the layer scalar, partial RoPE, and Google’s Scalar QK Norm — and offer possible explanations for them\. Finally, we conclude by making some visualizations and empirical observations about the denoising process\. Thank you Lucas Gu, Arshia Nayebnazar, Henry Ko, Rishi Athavale, and Tejas Prabhune for help proofreading and improving this blog post\! *Note: prior**knowledge of vanilla autoregressive LLM implementation $Llama 3\.1, MOE$ is assumed\.* [![demo.gif](https://substackcdn.com/image/fetch/$s_!9Yqo!,w_1456,c_limit,f_auto,q_auto:good,fl_lossy/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F7a59cc04-b5fa-4c97-963d-278ca9d8aff0_720x651.gif)](https://substackcdn.com/image/fetch/$s_!9Yqo!,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F7a59cc04-b5fa-4c97-963d-278ca9d8aff0_720x651.gif)Entropy over vocabulary distribution at each canvas position over time, figure generated at the end of the notebook - [Setup \+ Load Model](https://mlberkeley.substack.com/p/the-annotated-diffusiongemma#%C2%A7setup-load-model) - [Google’s Scalar QK Normalization](https://mlberkeley.substack.com/p/the-annotated-diffusiongemma#%C2%A7googles-scalar-qk-normalization) - [Rotary position embeddings](https://mlberkeley.substack.com/p/the-annotated-diffusiongemma#%C2%A7rotary-position-embeddings) - [Attention](https://mlberkeley.substack.com/p/the-annotated-diffusiongemma#%C2%A7attention) - [Decode vs Encode](https://mlberkeley.substack.com/p/the-annotated-diffusiongemma#%C2%A7decode-vs-encode) - [Gemma\-specific details](https://mlberkeley.substack.com/p/the-annotated-diffusiongemma#%C2%A7gemma-specific-details) - [Optimizations and Implementation](https://mlberkeley.substack.com/p/the-annotated-diffusiongemma#%C2%A7optimizations-and-implementation) - [MOE](https://mlberkeley.substack.com/p/the-annotated-diffusiongemma#%C2%A7moe) - [MOE Forward](https://mlberkeley.substack.com/p/the-annotated-diffusiongemma#%C2%A7moe-forward) - [Routing Score](https://mlberkeley.substack.com/p/the-annotated-diffusiongemma#%C2%A7routing-score) - [RMSNorm Fusion](https://mlberkeley.substack.com/p/the-annotated-diffusiongemma#%C2%A7rmsnorm-fusion) - [Putting Them together](https://mlberkeley.substack.com/p/the-annotated-diffusiongemma#%C2%A7putting-them-together) - [Logit Softcapping](https://mlberkeley.substack.com/p/the-annotated-diffusiongemma#%C2%A7logit-softcapping) - [Denoise vs Prefill](https://mlberkeley.substack.com/p/the-annotated-diffusiongemma#%C2%A7denoise-vs-prefill) - [Layer Scalar](https://mlberkeley.substack.com/p/the-annotated-diffusiongemma#%C2%A7layer-scalar) - [Sampling](https://mlberkeley.substack.com/p/the-annotated-diffusiongemma#%C2%A7sampling) - [Lazy Sampling](https://mlberkeley.substack.com/p/the-annotated-diffusiongemma#%C2%A7lazy-sampling) - [Stage Then Commit](https://mlberkeley.substack.com/p/the-annotated-diffusiongemma#%C2%A7stage-then-commit) - [Preliminary Analysis](https://mlberkeley.substack.com/p/the-annotated-diffusiongemma#%C2%A7preliminary-analysis) - [References](https://mlberkeley.substack.com/p/the-annotated-diffusiongemma#%C2%A7references) **Install, import, load configs and weights** ``` !pip install -q -U huggingface_hub !hf download google/diffusiongemma-26B-A4B-it \ --local-dir /content/diffusiongemma import glob import json import torch import torch.nn.functional as F from einops import rearrange, einsum from safetensors.torch import load_file from tokenizers import Tokenizer from torch.distributions import Categorical from tqdm import trange torch.set_default_device("cuda") max_tot_tokens = 2048 checkpoint = "diffusiongemma" canvas_len = 256 model_config = json.load(open(f"{checkpoint}/config.json"))['text_config'] gen_config = json.load(open(f"{checkpoint}/generation_config.json")) V = model_config['vocab_size'] sd = {} for safetensor_path in glob.glob(f"{checkpoint}/model-*.safetensors"): sd |= {k: v for k, v in load_file(safetensor_path, device="cuda").items() if "vision" not in k} W_vocab = sd['model.decoder.embed_tokens.weight'] sd ``` Printing out the weights, we find that interestingly, all entries within each learned QK\-normalization weight vector have the same value\. For example: ``` 'model.decoder.layers.17.self_attn.q_norm.weight': tensor( [0.9883, 0.9883, 0.9883, 0.9883, 0.9883, 0.9883, 0.9883, 0.9883, 0.9883, 0.9883, 0.9883, 0.9883, 0.9883, 0.9883, 0.9883, ...] ) ``` The QK\-norm weight vectors have the form \\$w\_q = \\gamma\_q \\mathbf\{1\}, \\qquad w\_k = \\gamma\_k \\mathbf\{1\},w\_q = \\gamma\_q \\mathbf\{1\}, \\qquad w\_k = \\gamma\_k \\mathbf\{1\},\\$ whereγ\_qandγ\_kare the learned scalars\. \\$\\begin\{aligned\} \\left\\lVert \\operatorname\{RMSNorm\}\_\{\\gamma\}\(x$ \\right\\rVert\_2 &= \|\\gamma\| \\frac\{\\lVert x \\rVert\_2\} \{\\lVert x \\rVert\_2 / \\sqrt\{d\_h\}\} \\\\ &= \|\\gamma\|\\sqrt\{d\_h\} \\end\{aligned\}\\\) RoPE preserves Euclidean norm, so the same result holds after applying rotate: \\$\\lVert q \\rVert\_2 = \|\\gamma\_q\|\\sqrt\{d\_h\}, \\qquad \\lVert k \\rVert\_2 = \|\\gamma\_k\|\\sqrt\{d\_h\}\\$ For a queryqand keyk\_iseparated by an angleθ\_i, \\$\\langle q, k\_i \\rangle = d\_h \\gamma\_q\\gamma\_k\\cos\\theta\_i\\$ Thus, the attention score fork\_iis \\$\\begin\{aligned\} \\operatorname\{score\}\_q\(k\_i$ &\\propto \\exp\\\!\\left$\\langle q, k\_i \\rangle\\right$ \\\\ &= \\exp\\\!\\left$d\_h\\gamma\_q\\gamma\_k\\cos\\theta\_i\\right$ \\\\ &= \\exp\\\!\\left$\\frac\{\\cos\\theta\_i\}\{T\}\\right$, \\end\{aligned\}\\\) where \\$T = \\frac\{1\}\{d\_h\\gamma\_q\\gamma\_k\}\\$ is a fixed, learned effective attention temperature per layer\. In regular QK norm models $and non\-QK norm models$, the magnitudes of q and k can vary\. This 1\) produces a per\-query\-dependent effective temperature and 2\) allows key magnitude to encode a form of query\-independent “global importance”: \\$\\begin\{aligned\} score\_q\(k\_i$ &\\propto \\exp\\\!\\left$ \\frac\{\\lVert k\_i \\rVert\_2 \\cos\\theta\_i\}\{T\_q\} \\right$, \\end\{aligned\}\\\) This may be undesirable in long\-context settings, where a high\-magnitude early key could remain disproportionately influential even after its relevance has faded\. One interpretation is that scalar QK norm effectively avoids this, and makes the inner product truly capture*similarity*between a q and k without being confounded with anything else\. **Verify all QK norm weights are like this** ``` for layer_id in range(model_config["num_hidden_layers"]): for kind in ("q", "k"): w = sd[f"model.decoder.layers.{layer_id}.self_attn.{kind}_norm.weight"] assert (w == w[0]).all(), "QK Norm entries are not all equal" ``` We do RMSNorm in higher precision like usual because precision errors are accumulated through the reduction, potentially amplified through the division, and propagated to every element\. \\$\\operatorname\{RMSNorm\}\_w\(x$ = \\sqrt\{d\}\\left$ w \\odot \\frac\{x\}\{\\\|x\\\|\_2 \+ \\epsilon\} \\right$\\\) **RMSNorm** ``` def rms(x, w = 1): return (w * x * (x.shape[-1] ** 0.5) / (torch.norm(x, dim=-1, keepdim=True, dtype=torch.float32) + model_config['rms_norm_eps'])).to(x.dtype) ``` DiffusionGemma has two RoPE configs: [![Table: Layer type | Head width | RoPE base θ | Rotated fraction; Sliding attention | 256 | 10,000 | 100%; Full attention | 512 | 1,000,000 | 25%](https://substackcdn.com/image/fetch/$s_!RjdQ!,w_1456,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fd612f7c2-0a8a-4391-8bcd-d27bfa485eba_1167x249.png)](https://substackcdn.com/image/fetch/$s_!RjdQ!,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fd612f7c2-0a8a-4391-8bcd-d27bfa485eba_1167x249.png) Standard RoPE assigns each of the d/2 complex / rotary pairs one frequency $head dim d$: \\$\\omega\_j = \\theta^\{\-2j/d\}, \\qquad j=0,\\ldots,d/2\-1\\$ This is what the sliding\-attention layers use, but interestingly, DiffusionGemma’s full\-attention layers use**partial RoPE**: after constructing the normal frequency vector, we then set all but the first 25% of its pair frequencies to zero, which will make`rotate`a no\-op on those dimensions\. Notably, we are doing truncating the full frequency spectrum, not the more intuitive option of rescaling it like we do in YaRN etc\. Partial RoPE enables a head to carry both explicitly position\-sensitive features and features whose representation is not rotated as position changes\. Interestingly, the rows ofW\_kandW\_qare stored arranged so that the projected vectors use split\-half rotary pairs,$t, t \+ d/2$, rather than adjacent pairs,$t, t\+1$, as in the original RoFormer paper\. This makes the rotation easier to express and slightly more cache\-friendly\. From equation 34 in RoFormer, \\$\\operatorname\{RoPE\}\(x$ = x\\cos$\\phi$ \+ \\operatorname\{rot\}$x$\\sin$\\phi$\\\) whererot$x$rotates each two\-dimensional pair by 90 degrees counterclockwise, thus for a token at position p, thei\-th pair of its k/q vector is rotated ccw by ω\_i \* p\. Note our code’s`rotate$x$`is RoPE$x$, not to be confused with rot$x$\. In other words, instead of something like ``` rot_x = np.empty_like(x) rot_x[..., 0::2] = -x[..., 1::2] rot_x[..., 1::2] = x[..., 0::2] ``` for $t, t\+1$ pairs, \\$\\operatorname\{rot\}\(x$ = \[\-x\_1, x\_0, \-x\_3, x\_2, \\ldots, \-x\_\{d\-1\}, x\_\{d\-2\}\],\\\) we have the more cache friendly ``` rot_x = rearrange([-x[..., head_dim // 2 :], x[..., : head_dim // 2]], 'z ... d -> ... (z d)') ``` \\$\\operatorname\{rot\}\(x$ = \[\-x\_\{d/2\},\\ldots,\-x\_\{d\-1\},x\_0,\\ldots,x\_\{d/2\-1\}\]\\\) **Precompute RoPE frequencies and rotate** ``` # precompute frequencies freq = {} # sliding window - regular rope freq['sliding_attention'] = model_config['rope_parameters']['sliding_attention']['rope_theta'] ** -(torch.arange(0, 1, 2 / model_config['head_dim'], dtype=torch.float32)) # full attention - partial rope freq['full_attention'] = model_config['rope_parameters']['full_attention']['rope_theta'] ** -(torch.arange(0, 1, 2 / model_config['global_head_dim'], dtype=torch.float32)) freq['full_attention'][int(model_config['rope_parameters']['full_attention']['partial_rotary_factor'] * len(freq['full_attention'])) : ] = 0 def rotate(x, layer_type, start_idx=0): # x [..., seq, head_dim]; rotate each (t, t + hd/2) pair ccw head_dim = x.shape[-1] rot_x = rearrange([-x[..., head_dim // 2 :], x[..., : head_dim // 2]], 'z ... d -> ... (z d)') pos = (torch.arange(x.shape[-2])[:, None] + start_idx) * torch.cat([freq[layer_type], freq[layer_type]]) # pos * (t, t + hd/2 pairs) return torch.cos(pos).to(x.dtype) * x + torch.sin(pos).to(x.dtype) * rot_x ``` In the**decode**stage, each canvas query attends to`\[committed history \| current canvas\]`\. It is: - Similar to autoregressive prefill in the sense that it writes KV’s for several tokens at a time, autoregressive decode in the sense that it reads a context history of past KV’s - Similar to cross attention in that it attends to KV’s from the encoder, whisper\-style, similar to self\-attention in that it is non\-casual, ViT\-style In the**encode**stage, the attention behaves like an ordinary causal prefill\. It can be thought of as the “verify” pass for speculative decoding with a 256\-token draft, with all tokens accepted\. The same exact attention weights is used for both\. This code can be easily modified to support batch size greater than 1, but it does not support mixed encode/decode modes within the same batch\. Dynamic per\-sequence attention mode is supported in[vLLM](https://vllm-project.github.io/2026/06/10/diffusion-gemma): [![vLLM per-sequence causal-attention masking diagram](https://substackcdn.com/image/fetch/$s_!-Y6t!,w_1456,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F29eeaddd-061e-4bfc-9483-21728ae38cc8.svg)](https://substackcdn.com/image/fetch/$s_!-Y6t!,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F29eeaddd-061e-4bfc-9483-21728ae38cc8.svg)Per\-sequence causal\-attention masking\. Source: vLLM\. Here, “denoise” means decode, “prefill” means encode on the prompt, and “accept” means encode on a denoised canvas\. For sliding layers, we use the approach in Google DeepMind’s[JAX sampler](https://github.com/google-deepmind/gemma/blob/195a9772b5cab4598dc8422780dd19dc0c03a284/gemma/diffusion/_sampler.py#L761-L783): every canvas token sees all other tokens inside the canvas plus the same prefix of KVs immediately before the canvas\. Note that vLLM uses a different symmetric, 2K\+1, per\-token sliding window\. Some notable differences from a standard LLM attention: 1. **No explicit1/√dmultiplier\.**The inner product is done directly\. This can absorbed into the scalar Q/K normalization weights, which already does a fixed inverse temperature\. 2. Gemma also applies an**embedding scale**, multiplying token embeddings by√Dat the start of the residual stream\. I’m not really sure why this is done\. 3. **Value normalization\.**Weightless RMS\-norm $i\.e\., unit norm$ is applied to every value vector before it is used or stored in cache 4. **Hybrid architecture\.**5 : 1 ratio of sliding window layers to global attention layer [![Table: Layer type | Q heads | K/V heads | Head width | RoPE; Sliding | 16 | 8 | 256 | Full; Global | 16 | 2 | 512 | 25% partial](https://substackcdn.com/image/fetch/$s_!02ap!,w_1456,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fb709ba41-ee79-4aa6-824e-ef1daf31e7b4_1142x249.png)](https://substackcdn.com/image/fetch/$s_!02ap!,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fb709ba41-ee79-4aa6-824e-ef1daf31e7b4_1142x249.png) 1. **Shared global K/V projection\.**In full\-attention layers,`W\_v`=`W\_k`\. K and V are only different due to post\-projection processing: K receives scalar\-weight RMS normalization plus RoPE, while V receives unit RMS normalization and no RoPE\. This is strange since K/V vectors are typically believed to live in different subspaces $see[1](https://arxiv.org/abs/2606.04032),[2](https://transformer-circuits.pub/2021/framework/index.html)while K/Q live in the same subspace$\. One potential optimization is to cache only the shared pre\-normalization projection and reconstruct K and V when reading it, potentially reducing global\-layer cache storage at the cost of additional computation\. Combined with the hybrid global/sliding schedule, this could be attractive in long\-context, KV\-cache\-bandwidth\-bound regimes\. The readable implementation below keeps separate K and V caches instead\. One implemented optimization is that, on encode, we can skip all computation after the K/V projection in the final layer, skipping most of the final layer and unembed work\. Encode passes do not use the model’s final logits; they only need to incur the minimum set of computation required to obtain $and commit$ correct K/V states in every layer\. In the current implementation with fixed\-size statically\-shaped cache, full\-attention is only differentiated from sliding window through`kv\_len`, thus behaves exactly like sliding\-window attention with a`max\_tot\_tokens`\-sized window: ``` class AttentionBlock(torch.nn.Module): def __init__(self, layer_id): super().__init__() self.layer_id = layer_id self.layer_type = model_config['layer_types'][layer_id] self.W_q, self.W_k, self.W_o = [sd[f'model.decoder.layers.{layer_id}.self_attn.{item}_proj.weight'] for item in ['q', 'k', 'o']] self.q_norm, self.k_norm = [sd[f'model.decoder.layers.{layer_id}.self_attn.{item}_norm.weight'] for item in ['q', 'k']] self.pre_norm = sd[f'model.decoder.layers.{layer_id}.input_layernorm.weight'] self.post_norm = sd[f'model.decoder.layers.{layer_id}.post_attention_layernorm.weight'] self.q_heads = model_config['num_attention_heads'] if self.layer_type == 'full_attention': self.W_v = self.W_k # global layers share K = V self.kv_heads = model_config['num_global_key_value_heads'] self.head_dim = model_config['global_head_dim'] self.kv_len = max_tot_tokens else: assert(self.layer_type == 'sliding_attention') self.W_v = sd[f'model.decoder.layers.{layer_id}.self_attn.v_proj.weight'] self.kv_heads = model_config['num_key_value_heads'] self.head_dim = model_config['head_dim'] self.kv_len = model_config['sliding_window'] self.k_cache = torch.empty(self.kv_heads, self.kv_len, self.head_dim) # use statically shaped KV buffer self.v_cache = torch.empty(self.kv_heads, self.kv_len, self.head_dim) # KV's flow into here from left to right, FIFO, latest element is rightmost def forward(self, x, pos_idx, mode): assert mode in ['encode', 'decode'] L, D = x.shape resid_x = x.clone() x = rms(x, w=self.pre_norm) q, k, v = x @ self.W_q.T, x @ self.W_k.T, x @ self.W_v.T q, k, v = [rearrange(z, 'l (n h) -> n l h', h = self.head_dim) for z in [q, k, v]] q, k = rms(q, self.q_norm), rms(k, self.k_norm) # QK-norm per head: weight is [head_dim], normalize over each head's dims v = rms(v, 1) # v_norm: weightless, no rope q, k = rotate(q, self.layer_type, pos_idx), rotate(k, self.layer_type, pos_idx) # absolute positions pos_idx .. pos_idx + L k = torch.concat([self.k_cache[:, : pos_idx, :], k], axis=1) # attend to [committed history | current block] v = torch.concat([self.v_cache[:, : pos_idx, :], v], axis=1) # Note python automatically clips on the left to 0, on the right to shape[1] = kv_len if mode == "encode": # difference #1: writes/updates the kv cache # Actually, we don't have to put this in a branch, can also just do this on decode too, ok since we'll override with an encode at the end anyways self.k_cache[:, : pos_idx + L, :] = k[:, -self.kv_len :, :] # automatically clips self.v_cache[:, : pos_idx + L, :] = v[:, -self.kv_len :, :] if self.layer_id == model_config['num_hidden_layers'] - 1: return # encode optimization: notice we don't need to do the remaining computation after this q = rearrange(q, '(n gqa) qt h -> n gqa qt h', gqa = self.q_heads // self.kv_heads) # Fold GQA into an outer dim scores = einsum(q, k, 'n gqa qt h, n kt h -> n gqa qt kt').float() # no divide by sqrt(head dim); softmax in fp32 if mode == "encode": scores += torch.triu(torch.full(scores.shape, -torch.inf), diagonal = scores.shape[-1] - scores.shape[-2] + 1) # this applies a mask that looks like R2 in the vllm figure scores = torch.exp(scores - torch.amax(scores, axis=-1, keepdims=True)) scores = (scores / torch.sum(scores, axis=-1, keepdims=True)).to(x.dtype) x = einsum(scores, v, 'n gqa qt kt, n kt h -> n gqa qt h') x = rearrange(x, 'n gqa qt h -> qt (n gqa h)') res = x @ self.W_o.T res = rms(res, self.post_norm) return res + resid_x ``` DiffusionGemma routes each token to 8 of 128 experts \+ one larger shared expert per forward pass\. Each expert is a standard gated GELU MLP: \\$\\operatorname\{MLP\}\(x$ = W\_\{\\text\{down\}\}\\left\[ $W\_\{\\text\{up\}\}x$ \\odot \\operatorname\{GELU\}$W\_\{\\text\{gate\}\}x$ \\right\]\\\) HuggingFace stores the routed experts’ gate and up projections as a single matrix, which we explicitly split in the MLP module\. [![DiffusionGemma mixture-of-experts architecture diagram](https://substackcdn.com/image/fetch/$s_!Vijg!,w_1456,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F37e47175-5232-4c65-89b8-e14a67b1054b_3807x2905.png)](https://substackcdn.com/image/fetch/$s_!Vijg!,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F37e47175-5232-4c65-89b8-e14a67b1054b_3807x2905.png)MOE Block operates independently across the token axis\. Each $of B × L many tokens$ is routed to its own 8 experts\. We could do this naively by looping over each token, but a better approach is to iterate over the experts instead: gather all tokens routed to each expert, do a single forward pass, then send results back\. Across multiple devices / for expert parallelism, this is usually implemented with an`AllToAll`\-style collective rather than the slow explicit Python loop we have here $see the[JAX Scaling Book](https://jax-ml.github.io/scaling-book/sharding/)$\. DiffusionGemma performs routing in the standard way\. Each expert receives a score proportional to \\$\\exp\\\!\\left\(\\langle x\_\{\\text\{route\}\}, k\_\{\\text\{expert\}\}\\rangle\\right$,\\\) where the expert keys are the rows of`W\_router`\. The top 8 scores are renormalized and multiplied by learned per\-expert scales\. There is an optimization we can do here to save a couple RMS\-norm computations\. Notice that we apply a pre\-norm 1 $shared expert$ \+ 8 $routed experts$ times at the start of each MLP\. Notice RMSNorm can be factored into two operations \\$\\operatorname\{RMS\}\_w\(x$ = w \\odot \\operatorname\{RMS\}\_1$x$\\\) , unit\-normalize then multiply by w\. Instead of doing this naively, we can instead apply a single weightlessRMS\_1$x$at the start, and absorb each expert’s learned RMSNorm weightwinto the columns of that expert’s gate and up projection matrix\. **MLP and mixture\-of\-experts block** ``` class MLP(torch.nn.Module): def __init__(self, layer_id, expert_num, conditioning_mlp : bool = False): super().__init__() if(conditioning_mlp): self.pre_norm = sd['model.decoder.self_conditioning.pre_norm.weight'] self.W_up = sd['model.decoder.self_conditioning.up_proj.weight'] self.W_gate = sd['model.decoder.self_conditioning.gate_proj.weight'] self.W_down = sd['model.decoder.self_conditioning.down_proj.weight'] return if expert_num is None: self.pre_norm = sd[f'model.decoder.layers.{layer_id}.pre_feedforward_layernorm.weight'] self.W_up, self.W_gate, self.W_down = [sd[f'model.decoder.layers.{layer_id}.mlp.{item}_proj.weight'] for item in ['up', 'gate', 'down']] else: self.pre_norm = sd[f'model.decoder.layers.{layer_id}.pre_feedforward_layernorm_2.weight'] self.W_gate, self.W_up = rearrange(sd[f'model.decoder.layers.{layer_id}.experts.gate_up_proj'][expert_num], '(z intermed) D -> z intermed D', z=2) self.W_down = sd[f'model.decoder.layers.{layer_id}.experts.down_proj'][expert_num] def forward(self, x): # every MLP has a prenorm. Also, note since the MLP is a compsition of functions that maps the 0 vector to itself, the entire MLP also maps all 0s to all 0s, this is relevant for the self-conditioning MLP L, D = x.shape x = rms(x, self.pre_norm) a = x @ self.W_up.T b = F.gelu(x @ self.W_gate.T) res = (a * b) @ self.W_down.T assert(res.shape == (L, D)) # share expert, routed expert, and self-conditioning MLP all map (_, D) -> (_, D) return res class MOEBlock(torch.nn.Module): def __init__(self, layer_id): super().__init__() self.k_experts = model_config['top_k_experts'] self.num_experts = model_config['num_experts'] self.W_router = sd[f'model.decoder.layers.{layer_id}.router.proj.weight'] self.expert_scale = sd[f'model.decoder.layers.{layer_id}.router.per_expert_scale'] self.scale = sd[f'model.decoder.layers.{layer_id}.router.scale'] self.experts = [MLP(layer_id, e) for e in range(self.num_experts)] self.shared_expert = MLP(layer_id, None) self.post_norm_1 = sd[f'model.decoder.layers.{layer_id}.post_feedforward_layernorm_1.weight'] # applied on shared expert output self.post_norm_2 = sd[f'model.decoder.layers.{layer_id}.post_feedforward_layernorm_2.weight'] # applied on summed contribution from experts self.post_norm = sd[f'model.decoder.layers.{layer_id}.post_feedforward_layernorm.weight'] # the sum h1 (shared ) + h2 (routed sum), before residual add def forward(self, x): L, D = x.shape resid_x = x.clone() route_x = rms(x) * self.scale / (D ** 0.5) expert_scores = F.softmax((route_x @ self.W_router.T).float(), dim=-1) # (L, num_experts) top_k_scores, top_k_idx = torch.topk(expert_scores, self.k_experts, dim=-1) # (L, k_experts), (L, k_experts) top_k_scores = (top_k_scores / torch.sum(top_k_scores, dim=-1, keepdim=True) * self.expert_scale[top_k_idx]).to(x.dtype) res = rms(self.shared_expert(x), self.post_norm_1) # h1: dense branch, shape (L, D) moe_out = torch.zeros_like(x) for id, expert in zip(range(self.num_experts), self.experts): mask = torch.any(top_k_idx == id, dim = -1) # boolean mask (L, ) which tokens routed to expert_id mult = top_k_scores[mask][top_k_idx[mask] == id] # shape (L', ) where L' <= L is the number of tokens expert_id routed to moe_out[mask] += mult[:, None] * expert(x[mask]) # (L', D) += (L', 1) * (L', D) res = res + rms(moe_out, self.post_norm_2) # h2 normed once, then h1 + h2 return rms(res, self.post_norm) + resid_x ``` We now combine all our previous components into a single module\. Unlike Llama, DiffusionGemma does a couple things different beyond merely interleaving attention and MoE blocks\. Logit softcapping applies the following function to the final logits: \\$z\_\{\\text\{capped\}\} = c\\tanh\\\!\\left\(\\frac\{z\}\{c\}\\right$,\\\) wherezis an uncapped logit andcis the softcap value\. For\|z\| ≪ c,tanh$z/c$ ≈ z/c, so small logits are essentially unchanged\. Asz → ±∞,z\_capped → ± c, so large logits are “soft\-capped”\. This is a smooth alternative to`torch\.clip`or`torch\.clamp`\- unlike a hard cap, it remains differentiable everywhere\. [![Graph of the standard hyperbolic tangent function](https://substackcdn.com/image/fetch/$s_!2toU!,w_1456,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F6bb55c4e-9df8-4896-8786-a7421703add9_646x430.jpeg)](https://substackcdn.com/image/fetch/$s_!2toU!,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F6bb55c4e-9df8-4896-8786-a7421703add9_646x430.jpeg) The standard tanh function is bounded between \-1 and 1; multiplying bycchanges the bounds to\[−c,c\]\. Note this is part of the model, and not the sampling process — not to be confused with the temperature schedule, which is applied later\. The following are function signatures of encode vs decode: ``` # this commits canvas, (last time we) write KV cache def enc(self, pos_idx, logits) -> None: self._forward(pos_idx, logits, 0, mode="encode") # this maps (canvas_i, canvas_prob_i) -> (canvas_prob_i+1) def dec(self, pos_idx, logits, logit_probs) -> torch.tensor: return self._forward(pos_idx, logits, logit_probs, mode="decode") ``` On encode, the embedded tokens are directly passed into the first layer\. In decode mode, the embedded tokens additionally incorporates a*self\-conditioning*signal from the previous denoising step\. LetP∈ℝ^$L×V$be the previous step’s post\-softmax probability distribution over vocabs at each of L canvas positions, letE∈ℝ^$V×D$be the tied embedding table\. Then,`P @ E`computes the expected token embeddings from the previous iteration, L independent convex sums with the weights stored in P\. This is then passed through a small conditioning MLP, added to the canvas token embedding, and RMSNormed\. Self\-conditioning is helpful because for example, - Knowing how confident the previous pass was in this token can inform us about how confident we should be in this pass - During sampling, the highest entropy, least confident tokens are replaced by a random token ID\. Thus, self\-conditioning allows the model to deduce which tokens are effectively`\[MASK\]`tokens by comparing the probability distribution input to the token input\.[2](https://mlberkeley.substack.com/p/the-annotated-diffusiongemma#footnote-2) - Possibly gives a way for gradients to flow across denoising steps [![Diagram of DiffusionGemma self-conditioning across denoising steps](https://substackcdn.com/image/fetch/$s_!dI2O!,w_1456,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F7abf126b-482c-41c8-9319-b0ab17f6409b_1890x2048.png)](https://substackcdn.com/image/fetch/$s_!dI2O!,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F7abf126b-482c-41c8-9319-b0ab17f6409b_1890x2048.png)Self\-conditioning across denoising steps\. Source: Maarten Grootendorst\. On the first decode step, the self\-conditioning input is all 0s\. Note that since the self\-conditioning MLP is composed of functions that map all 0s to all 0s, the added contribution from self\-conditioning to the input is 0 \- this is equivalent to skipping self\-conditioning for that step $no previous “self”$\. Notice each block of computation in DiffusionGemma is*additive*to the residual stream: \\$x \\longmapsto x \+ F\\\!\\left\(\\operatorname\{RMSNorm\}\(x$\\right\)\\\) whereFis the core computation in that block, andRMSNorm$x$is that block’s pre\-norm $both attention and MOE have this$\. Thus due to the skip connections, the residual stream remains unnormalized through all 30 layers\. This is perfectly fine if the activations truly live inℝ^n, but unfortunately, they are bound to the subset of representation values of its datatype, which has a fixed dynamic range \- upper and lower bounds of the feasible set, and anything outside will over/underflow to infinities\. To ensure this doesn’t happen, quantization methods typically multiply by a scalarSbefore casting, for example: \\$X\_\{\\mathrm\{FP8\}\} = \\operatorname\{cast\}\_\{\\mathrm\{to FP8\}\} \\left\( \\operatorname\{clip\} \\left\( S \\cdot X,\\, \-V\_\{\\max\}^\{\\mathrm\{FP8\}\},\\, V\_\{\\max\}^\{\\mathrm\{FP8\}\} \\right$ \\right\),\\\) where \\$S = \\frac\{ V\_\{\\max\}^\{\\mathrm\{FP8\}\} \}\{ \\max\\left\(\\lVert X\\rVert\_\{\\infty\}, \\varepsilon\\right$ \}\\\) to ensure values lie within the dynamic range\[−V\_max, V\_max\]\. A similar worry arises for the unnormalized residual stream, the magnitude of the activations may gradually grow through the layers $note however, the inputs to blocks remain well\-conditioned due to the RMSNorm$\. To improve training stability, avoid overflow, and enable stable low\-precision inference, DiffusionGemma adds a*layer scalar*which rescales the completed residual stream by a learned scalars\_lat the end of every layer: \\$x\_\{\\ell\+1\} = s\_\\ell \\cdot \\operatorname\{MoEBlock\}\_\\ell\\\!\\left\( \\operatorname\{Attn\}\_\\ell\(x\_\\ell$ \\right\)\\\) Printing out the layer scalars’ actual values, we see that they are all< 1, which makes sense according to our interpretation\. Interestingly, it is a learned parameter — the model could learn to adaptively regulate the magnitude of its own residual stream throughout training\. Notice the RMSNorm is invariant to multiplying its input by a scalar $RMSNorm\(cx$ = RMSNorm$x$\)\. Since every block receives input through an RMSNorm, the layer scaler has little benefit on the normalized input presented to each block\. However, note it will affect the direction of the input after the first layer\. The released Hugging Face checkpoint stores separate encoder and decoder layer\-scalar entries, ``` sd[f"model.encoder.language_model.layers.{i}.layer_scalar"] and sd[f"model.decoder.layers.{i}.layer_scalar"] ``` but their values are actually identical\. In the code below, we have a single`self\.layer\_scalar`\. **Complete DiffusionGemma model** ``` class DiffusionGemma(torch.nn.Module): # Any computation that utilizes parameters passes through here def __init__(self): super().__init__() self.W_embed = sd['model.decoder.embed_tokens.weight'] # also used as the unembedding matrix (tie_word_embeddings = True) self.attn_blocks = [AttentionBlock(i) for i in range(model_config['num_hidden_layers'])] self.moe_blocks = [MOEBlock(i) for i in range(model_config['num_hidden_layers'])] self.layer_scalar = [sd[f'model.encoder.language_model.layers.{i}.layer_scalar'] for i in range(model_config['num_hidden_layers'])] self.model_norm = sd['model.decoder.norm.weight'] # final / model norm self.embed_scale = torch.tensor(model_config['hidden_size'] ** 0.5) self.conditioning_MLP = MLP(layer_id=None, expert_num=None, conditioning_mlp=True) def _forward(self, pos_idx, logits, logit_probs, mode): # logit_probs = 0 <=> skip this path (no bias term anywhere) x = self.W_embed[logits] * self.embed_scale # (L, ) -> (L, D) # do self conditioning if decode if mode == "decode": condition_x = (logit_probs.to(x.dtype) @ self.W_embed) * self.embed_scale # (L, V) x (V, D) -> convex combination of vocab embeddings condition_x = self.conditioning_MLP(condition_x) x = rms(x + condition_x) # pass through all layers for i, (attn, moe) in enumerate(zip(self.attn_blocks, self.moe_blocks)): x = attn(x, pos_idx, mode) if x is None: # last encode layer wrote its KV cache and returned early; nothing else is needed return None x = moe(x) x = x * (self.layer_scalar)[i] # per-layer encoder/decoder scalar x = rms(x, self.model_norm) final_logits = (x @ self.W_embed.T).float() # (L, D) x (D, V) -> (L, V) return torch.tanh(final_logits / model_config['final_logit_softcapping']) * model_config['final_logit_softcapping'] # this commits canvas, (last time we) write KV cache def enc(self, pos_idx, logits) -> None: self._forward(pos_idx, logits, 0, mode="encode") # this maps (canvas_i, canvas_prob_i) -> (canvas_prob_i+1) def dec(self, pos_idx, logits, logit_probs) -> torch.tensor: return self._forward(pos_idx, logits, logit_probs, mode="decode") ``` The following figure is a nice summary of what we’ve put together: [![DiffusionGemma model architecture recap diagram](https://substackcdn.com/image/fetch/$s_!VwP4!,w_1456,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fc91ee800-f0f4-4c88-a346-31e47490b429_3822x4093.png)](https://substackcdn.com/image/fetch/$s_!VwP4!,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fc91ee800-f0f4-4c88-a346-31e47490b429_3822x4093.png)DiffusionGemma Model Diagram [![vLLM DiffusionGemma sampling loop](https://substackcdn.com/image/fetch/$s_!itom!,w_1456,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F69695187-c486-44e6-afba-35fe0c3392ce_1476x268.svg)](https://substackcdn.com/image/fetch/$s_!itom!,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F69695187-c486-44e6-afba-35fe0c3392ce_1476x268.svg)DiffusionGemma generation loop\. Source: vLLM\. After the first encode on the input prompt, every 256 canvas of tokens incurs one encode pass and at most`max\_denoising\_steps`decode passes\. After obtaining a vocabulary distribution for every canvas position, sampling is more involved than in an autoregressive model\. We follow the sampling procedure explained[here](https://newsletter.maartengrootendorst.com/p/a-visual-guide-to-diffusiongemma), implemented[here](https://github.com/google-deepmind/gemma/blob/main/gemma/diffusion/_sampler.py)\. The meanings of the sampling parameters are documented[here](https://ai.google.dev/gemma/docs/diffusiongemma): 1. **Apply temperature** For decode passiout ofN = max\_denoising\_steps, \\$t\_i = t\_\{\\max\} \+ \\frac\{i\}\{N\}\\left\(t\_\{\\min\}\-t\_\{\\max\}\\right$, \\qquad i=0,\\ldots,N\-1\\\) Thus the temperature starts att\_maxand decreases towardt\_min, increasingly sharpening the distribution\. In the i\-th iteration, we work with the categorical distribution formed by`logits / t\_i`\. 2. **Compute entropy** For every canvas position, we compute from its categorical distributionp, \\$H\(p$ = \-\\sum\_\{v=1\}^\{V\} p\_v\\log p\_v\\\) Used as a measure of how uncertain the model is about the token at this position\. 3. **Keep a low\-entropy prefix and renoise the rest** Sort position entropies so that \\$H\_\{\(1$\} \\le H\_\{$2$\} \\le \\cdots \\le H\_\{$L$\}\\\) Accept the largest prefix ending atksuch that \\$\\sum\_\{j=1\}^\{k\-1\} H\_\{\(j$\} \\le \\texttt\{entropy\_bound\}\\\) Then, replace every unaccepted position by a token drawn uniformly from the vocabulary\. 4. **Check for early stopping** Convergence is decided when the canvas is confident*and*stable: - the argmax canvas has remained unchanged for`stability\_threshold`$equals 1 here$ previous canvases, and - the current canvas mean per\-position entropy is below`confidence\_threshold`\. On convergence or after N steps, we commit the latest canvas\. **Tokenizer and sampling configuration** ``` tok = Tokenizer.from_file(f"{checkpoint}/tokenizer.json") prompt = """ What is the meaning of 67? """ chat = f"<bos><|turn>user\n{prompt}<turn|>\n<|turn>model\n" max_denoising_steps = gen_config['max_denoising_steps'] # decoder forward passes per canvas entropy_bound = gen_config['sampler_config']['entropy_bound'] # see formula in figure t_max, t_min = gen_config['t_max'], gen_config['t_min'] # linear schedule of temperatures t_max -> t_min across the steps confidence_threshold = gen_config['confidence_threshold'] # early-stop a canvas when argmax is stable (equal previous argmax canvas) and mean entropy < this tokens = torch.tensor(tok.encode(chat, add_special_tokens=False).ids) model = DiffusionGemma() ``` The usual way a denoising iteration is implemented: > Model forward → retain distribution for self\-conditioning → sample and renoise tokens Instead, we do sampling lazily, passing only the categorical distribution between steps and sampling only once we need it: > Stored distribution → sample and renoise input tokens → model forward → new distribution The input into the first iteration is just all identical logits, which will generate a random canvas for us\. This is nice since: - The only information we need to pass between steps is the self\-conditioning input - Initialization falls out naturally from a uniform categorical distribution, can be folded into the first step - We can think of the decoder input as the previous canvas probs plus a source of randomness $from the sampling$\. Under this view, preparing the inputs $sampling$ lazily is optimal\. Additionally, it suggests an alternative interpretation of self\-conditioning as the primary rather than auxiliary input: One way to think of the decoder is that it repeatedly transports our current distribution toward the target distribution, reminiscent of a[flow\-like process on the probability simplex](https://arxiv.org/abs/2602.12233)\. So one can think about an alternate DiffusionGemma design that instead starts from a random position in the simplex at initialization,`Categorical$probs=torch\.rand\(L, V$\)`or`Categorical$logits=torch\.rand\(L, V$\)`, rather than “`Categorical$logits=torch\.zeros\(L, V$\)`”\. We also intentionally give general functions for`denoise`which generates staged tokens; and`commit`, which takes in some staged tokens, generates their encoded KVs, and writes them to cache\. `denoise$pos\_idx$`repeatedly calls decode against the already committed`cache\[:pos\_idx\]`on a fresh canvas at positions`\[pos\_idx : pos\_idx \+ 256\]`, and returns staged tokens sampled from the final canvas of logits\. This function has no side effects on the KV cache, thus we can for example denoise multiple canvases at the same position and choose the best one\. We can also place the canvas anywhere, as long as pos\_idx≤len$tokens$ \- canvas\_len, since the correct slice of the KV cache will automatically be read\. `commit$l, r, staged\_tokens$`applies the casual encode and writes $or overwrites$ the slice of KV cache from l to r, with the staged tokens’ KVs\. This function allows arbitrary\-lengthed blocks, thus we can for example commit only a confident prefix of the current canvas\. We can also commit tokens at any positions, but if an earlier region is overwritten, all later tokens are invalidated because their cached states depended on the old prefix\. This design is to make alternative schedules easier to study, including overlapping canvases, revising an earlier block, selecting among several denoised candidates, or advancing only part of a canvas\. **Denoise, stage, commit, and generate** ``` def denoise(pos_idx): # returns staged_tokens assert pos_idx + canvas_len <= len(tokens) # must be length canvas_len (what if it wasn't fixed? analyze how casual the self attention is) plotting_data = [] t = t_max t_step = (t_min - t_max) / max_denoising_steps last_canvas = Categorical(logits = torch.ones((canvas_len, V))) for step in trange(max_denoising_steps): # renoise last_canvas sH, sidx = last_canvas.entropy().sort(-1) accepted = torch.zeros_like(sH, dtype=torch.bool).scatter(-1, sidx, sH.cumsum(-1) - sH <= entropy_bound) last_canvas_noised = torch.where(accepted, last_canvas.sample(), torch.randint(0, V, (canvas_len,))) # pass in the noised tokens, but un-noised normalized probs (all-zero probs on step 0: the conditioning path maps 0 to 0) canvas = model.dec(pos_idx, last_canvas_noised, last_canvas.probs if step != 0 else torch.zeros(canvas_len, V)) / t; plotting_data.append(canvas.detach().cpu()) canvas = Categorical(logits = canvas) if torch.mean(canvas.entropy()) < confidence_threshold and (canvas.logits.argmax(dim=-1) == last_canvas.logits.argmax(dim=-1)).all(): return canvas.sample(), plotting_data t += t_step last_canvas = canvas assert False, f"Denoising not finished after {max_denoising_steps} steps" def commit(l, r, staged_tokens): global tokens assert len(staged_tokens) == r - l + 1 if(r + 1 < len(tokens)): # commits staged_tokens print(f"Invalidating {len(tokens) - (r+1)} tokens") tokens = tokens[:r+1] tokens[l:] = staged_tokens model.enc(l, staged_tokens) def new_canvas(): global tokens if len(tokens) + canvas_len > max_tot_tokens: return False nxt = torch.randint(V, (canvas_len,)) tokens = torch.concat([tokens, nxt]) return True commit(0, len(tokens) - 1, tokens) # prefill: encode the chat prompt into the KV cache pos_idx = len(tokens) plotting_data = [] while new_canvas(): staged_tokens, cur_data = denoise(pos_idx) commit(pos_idx, pos_idx + canvas_len - 1, staged_tokens) pos_idx = len(tokens) print('=' * 50) print('Final Canvas') print(tok.decode(tokens.tolist())) print('=' * 50) plotting_data.append(cur_data) if torch.isin(staged_tokens, torch.tensor(gen_config['eos_token_id'])).any(): break ``` Output: Using the saved`plotting\_data`from denoise, we can make some figures of the denoising process over time: Note the cells are laid out in reading order\. Without being explicitly trained to do so, it appears the model is approximately finalizing canvas tokens from left to right\. Plotting two more figures: Further confirms our suspicions\. At every step, the model seems to mostly focus its efforts on figuring out the next causal rolling window of tokens — it “thinks” casually\. Is this a general behavior the model has learned? Let’s try a different prompt, solving a[magic square](https://en.wikipedia.org/wiki/Magic_square)puzzle: Canvas 2 Canvas 3, More visualizations available[here](https://github.com/timothygao8710/The-Annotated-DiffusionGemma)\. Interestingly, the prose tokens seem to exhibit the same causal pattern, but the magic square tokens do not\. This makes intuitive sense\. When humans solve these puzzles, the next value that’s easiest to deduce is not necessarily the next one in autoregressive order\. This motivates perhaps a better explanation for how DiffusionGemma “thinks” during denoising — the “easiest”, lowest\-entropy tokens are fixed / finalized / accepted first, which in turn unlocks / makes other tokens “easier” for the next denoising step\.[3](https://mlberkeley.substack.com/p/the-annotated-diffusiongemma#footnote-3) Note this mirrors what we do during sampling: all but a prefix of lowest entropy tokens is renoised\. For prose and chain\-of\-thought reasoning, the easiest tokens to resolve next often coincide with the earliest tokens in autoregressive order\. However for both prompts, once a token is accepted $i\.e\. low\-entropy and not renoised$, it’s extremely unlikely to be renoised in the future: [![canvas_00_acceptance_transitions.png](https://substackcdn.com/image/fetch/$s_!pd4q!,w_1456,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fba9b8f5e-0dec-4f5e-a859-5de8c3a00001_2450x1550.png)](https://substackcdn.com/image/fetch/$s_!pd4q!,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2Fba9b8f5e-0dec-4f5e-a859-5de8c3a00001_2450x1550.png) For 67 Prompt [![canvas_02_acceptance_transitions.png](https://substackcdn.com/image/fetch/$s_!fs9I!,w_1456,c_limit,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F185a5d80-2108-4056-b992-41cd4cce4a2b_2178x1378.png)](https://substackcdn.com/image/fetch/$s_!fs9I!,f_auto,q_auto:good,fl_progressive:steep/https%3A%2F%2Fsubstack-post-media.s3.amazonaws.com%2Fpublic%2Fimages%2F185a5d80-2108-4056-b992-41cd4cce4a2b_2178x1378.png) For Magic Square Prompt Looking at the top right graph, almost no tokens are “re\-masked” for both prompts\. Google[claims](https://developers.googleblog.com/diffusiongemma-the-developer-guide/)that one of the benefits of DiffusionGemma’s uniform state diffusion over traditional masked diffusion $e\.g\., LLaDA, Dream$ is error correction via re\-noising, but from our analysis — how often does this actually happen in practice? - O’Donoghue, Brendan, and Sebastian Flennerhag\. “[DiffusionGemma: 4x faster text generation](https://blog.google/innovation-and-ai/technology/developers-tools/diffusion-gemma-faster-text-generation/)\.” Google, 2026\. - Ballantyne, Ian, and Omar Sanseviero\. “[DiffusionGemma: The Developer Guide](https://developers.googleblog.com/diffusiongemma-the-developer-guide/)\.” Google Developers Blog, 2026\. - Google AI for Developers\. “[DiffusionGemma model overview](https://ai.google.dev/gemma/docs/diffusiongemma)\.” 2026\. - Google AI for Developers\. “[Diffusion in Text Generation Explained](https://ai.google.dev/gemma/docs/diffusiongemma/explained)\.” 2026\. - Google DeepMind\. “[DiffusionGemma 26B\-A4B\-IT model card](https://huggingface.co/google/diffusiongemma-26B-A4B-it)\.” Hugging Face, 2026\. - Google DeepMind\. “[DiffusionGemma reference sampler](https://github.com/google-deepmind/gemma/blob/195a9772b5cab4598dc8422780dd19dc0c03a284/gemma/diffusion/_sampler.py)\.”*Gemma*repository, 2026\. - vLLM Team and Google DeepMind Team\. “[DiffusionGemma: The First Diffusion LLM $dLLM$ Natively Supported in vLLM](https://vllm-project.github.io/2026/06/10/diffusion-gemma)\.” vLLM Blog, 2026\. - Grootendorst, Maarten\. “[A Visual Guide to DiffusionGemma](https://newsletter.maartengrootendorst.com/p/a-visual-guide-to-diffusiongemma)\.”*Exploring Language Models*, 2026\. - Huang, Austin, Suraj Subramanian, Jonathan Sum, Khalid Almubarak, and Stella Biderman\. “[The Annotated Transformer](https://nlp.seas.harvard.edu/annotated-transformer/)\.” Harvard NLP, 2022\. Original version by Alexander M\. Rush\. - Peng, Bowen, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole\. “[YaRN: Efficient Context Window Extension of Large Language Models](https://arxiv.org/abs/2309.00071)\.” arXiv:2309\.00071, 2023\. - Su, Jianlin, et al\. “[RoFormer: Enhanced Transformer with Rotary Position Embedding](https://arxiv.org/abs/2104.09864)\.” arXiv:2104\.09864, 2021\. - Kayyam, Ali, Anusha Madan Gopal, and M\. Anthony Lewis\. “[Do Transformers Need Three Projections? Systematic Study of QKV Variants](https://arxiv.org/abs/2606.04032)\.” arXiv:2606\.04032, 2026\. - Elhage, Nelson, et al\. “[A Mathematical Framework for Transformer Circuits](https://transformer-circuits.pub/2021/framework/index.html)\.” Transformer Circuits Thread, 2021\. - Austin, Jacob, et al\. “[Sharded Matrices and How to Multiply Them](https://jax-ml.github.io/scaling-book/sharding/)\.”*How To Scale Your Model*, 2025\. - Roos, Daan, et al\. “[Categorical Flow Maps](https://arxiv.org/abs/2602.12233)\.” arXiv:2602\.12233, 2026\. - Nie, Shen, et al\. “[Large Language Diffusion Models](https://arxiv.org/abs/2502.09992)\.” arXiv:2502\.09992, 2025\. - Ye, Jiacheng, et al\. “[Dream 7B: Diffusion Large Language Models](https://arxiv.org/abs/2508.15487)\.” arXiv:2508\.15487, 2025\. #### Discussion about this post ### Ready for more?

DiffusionGemma Explained

Similar Articles

Coding Diffusion Gemma from scratch

DiffusionGemma Technical Report

DiffusionGemma

DiffusionGemma: The Developer Guide- Google Developers Blog

google/diffusiongemma-26B-A4B-it

Submit Feedback

Similar Articles

Coding Diffusion Gemma from scratch

DiffusionGemma Technical Report

DiffusionGemma: The Developer Guide- Google Developers Blog

google/diffusiongemma-26B-A4B-it