Fine-Tuning DiffusionGemma: What Works, What Breaks

Andrea Miele* · William Bankes* · Keyue Jiang* · Seongho Son · Xiaohang Tang · Ilija Bogunovic * Co-first authors. Code release: TRL SFT pipeline ↗

💡 TL;DR

We provide an empirical study of SFT post-training for DiffusionGemma [9], focusing on the choice of training objective, long-horizon behavior, prompting sensitivity, and practical training infrastructure.

  • The SFT objective matters: we investigate different SFT objectives and show how different weighting schemes affect DiffusionGemma's performance on a range of tasks.
  • DiffusionGemma struggles as the horizon grows: SFT on long CoT maths traces (OpenR1-Math-220k [19]) and on multi-turn terminal-agent data (TMax-SFT [20]) both reduce model performance, motivating further investigation into long-horizon post-training.
  • DiffusionGemma is quite sensitive to prompting: small changes in prompt templates and answer formatting can substantially affect downstream performance, highlighting the need for careful prompt design and standardized evaluation.
  • We release a full TRL SFT pipeline: code ↗, so the community can reproduce our results, build on the recipe, and extend it to new objectives, tasks, and uniform diffusion models.

Autoregressive models still dominate language modeling. They are efficient to train, easy to fine-tune, and backed by a mature ecosystem of post-training tools. But during generation they commit to tokens one at a time, and once a token is out, the model cannot take it back.

Discrete diffusion language models (dLLMs) generate differently. Instead of writing left to right, they iteratively construct an answer from bidirectional context. In principle, they can revisit earlier predictions and fix mistakes before producing the final output, and they are becoming genuinely competitive with autoregressive models on reasoning, mathematics, and tool use [1–7].

Most open dLLMs, such as LLaDA [3], Dream [4], and DiffuCoder [5], get there through masked diffusion [10]: they start from masked positions and progressively fill them in. Generation isn’t strictly left-to-right, but tokens are usually committed as decoding proceeds, so an early mistake can still constrain the rest of the answer, short of remasking tricks [34].

DiffusionGemma-26B-A4B takes a different route [9]. It uses Uniform State Diffusion [11–14], starting from a canvas (a fixed-size block of output positions the model denoises jointly) filled with random tokens, and repeatedly rewriting it throughout generation. Earlier predictions can be replaced, partial answers rewritten, mistakes corrected before the sequence is finalized.

Two ways to denoise a sequence

Masked diffusion reveals tokens and commits them. Uniform diffusion keeps rewriting the whole canvas; watch the highlighted token settle on the wrong word, then correct itself.

Masked diffusion MDLM
▨▨▨ ▨▨▨ ▨▨▨ ▨▨▨ ▨▨▨ ▨▨▨ ▨▨▨

Committed tokens are locked, so an early mistake stays.

Uniform diffusion USDM
btepuwn iuuiwhncb jfmsxzuk bombda kduxef mipcm yemxlu

The whole canvas is revisable, so mistakes can be rewritten.

0/6

At least, that’s the promise.

DiffusionGemma matters because it is one of the first large, open-weight models to make uniform diffusion accessible to the broader community. Until now, nearly all open models and post-training methods focused on masked diffusion, simply because the tooling already exists.

But an open-weight base model is only a starting point. Almost nothing that makes a language model useful comes for free from pre-training: following instructions, staying on task, reasoning through a problem, calling a tool and knowing when to stop. All of that is installed afterwards, during post-training: supervised fine-tuning on demonstrations, then reinforcement learning against a reward. Autoregressive and masked-diffusion models have a large head start here, and it is less about clever recipes (which are not really that involved) than about infrastructure: fast, well-supported inference for generating answers, mature training libraries, and standardized evaluation. A large uniform diffusion model has almost none of that yet. So the question that matters is:

How do we post-train a large uniform diffusion language model?

Alongside the release, the DiffusionGemma team showed a single SFT experiment on Sudoku, but with few details about the loss or the task. Those details matter. Fully customisable post-training is what makes an open release valuable: it lets practitioners adapt models to their own problems and lets researchers push a model past its shipped capabilities. So in this post we take a careful look at Supervised Fine-Tuning (SFT) [16] for DiffusionGemma: how much it helps, which objective works best, and exactly where it starts to break.

Concretely, we (1) reproduce the release’s Sudoku setup and quantify how much SFT helps; (2) show a leave-one-out objective improves on the release recipe and transfers beyond the training distribution; and (3) show that DiffusionGemma is brittle to SFT on both long-CoT math and terminal-agent data. DiffusionGemma opens an exciting direction for open diffusion LMs, but it also exposes challenges that need solving before uniform diffusion models can be adapted as reliably as their autoregressive counterparts.

1. Supervised Fine-Tuning a Uniform Diffusion Model

TL;DR

We formalize DiffusionGemma SFT by corrupting only the completion while keeping the prompt clean, and compare objective variants that differ in how they weight noise levels and interpret the model's logits.

Before we can ask which objective is best, we need to pin down what “fine-tuning” even means for a model that doesn’t generate left to right. So let’s build the objective up from scratch.

Start with the familiar case. SFT is straightforward for an autoregressive model: given a prompt and a completion, you train it to predict the next completion token from the tokens before it.

Uniform state diffusion models don’t work that way. They don’t learn a left-to-right factorization of the completion at all; they learn to recover clean text from a corrupted version of the whole sequence. Following the uniform-diffusion literature [12–14], SFT is framed as conditional denoising: the prompt stays visible throughout training, the completion is corrupted by the forward process, and the model is trained to reconstruct the original completion.

The forward process is worth stating carefully, because everything later builds on it. Write for the prompt, for its completion, and for the clean sequence they form. For each completion token, the forward process keeps the original with probability and otherwise replaces it with a uniform draw from the vocabulary (of size ):

The diffusion time runs from clean (, so ) to fully noised (, , uniform noise). The model learns the backward process, going from noise back to clean.

The training script released with DiffusionGemma actually optimizes a sum of two objectives, not the denoising loss alone:

  • an encoder loss: a standard autoregressive cross-entropy over the whole clean sequence.
  • a decoder loss: the diffusion denoising cross-entropy over a single canvas.

Two things are worth noting. First, the decoder sum runs over all positions in the canvas, corrupted or not, unlike masked-diffusion objectives, which only supervise corrupted positions. Second, the encoder term is really a consequence of how DiffusionGemma comes to be: it is a pretrained autoregressive model converted into a diffusion one, so the encoder and decoder share weights. Keeping an autoregressive cross-entropy on the encoder side preserves the model’s inherited next-token behaviour while the decoder learns to denoise. Other autoregressive-to-diffusion conversions, such as Nemotron-Labs-Diffusion [36], do something similar. The practical point is that both losses update the same tensors, so fine-tuning with the decoder loss alone trains a genuinely different setup from the one Google shipped.

This also matters for LoRA. Since the encoder and decoder share tensors, one adapter per module gets counted twice when the adapter is merged back in. Setting the encoder loss weight to zero sidesteps this, because then no encoder adapter is created in the first place.

Concretely, one SFT step looks like this:

  1. Build the training sequence. Concatenate prompt and completion into and define a completion mask (1 on completion tokens, 0 on prompt tokens); the loss applies only where .
  2. Sample a diffusion time and corrupt the sequence. Draw one per sequence and corrupt every completion token with the forward process, .
  3. Run the model. Produce logits , optionally with a second pass for self-conditioning.
  4. Compute the denoising cross-entropy. Score against the clean completion , using to exclude the prompt tokens, weighted by a time factor (which we unpack next).

Every variant we compare shares these four steps and differs only in step 4: the weight , and how the logits are read.

Aside: self-conditioning

Self-conditioning is a general technique for diffusion models, explored across the diffusion literature in prior work [35]: on a given step, the model is handed its own prediction from the previous step as an extra input, so it can build on and correct that earlier guess instead of denoising from scratch.

In DiffusionGemma specifically, this means re-inputting the model's previous-step logits, its running estimate of the clean sequence, back into the network for a second pass. That first estimate is treated as a fixed input, and the training loss is computed from the second, refined pass. In practice it is applied to half of each batch (using the estimate from an earlier pass); the other half conditions on nothing.

Self-conditioning does not change the SFT objective itself. It only changes the information the denoiser has available.

One loss, two design choices

Our formulation keeps a completion mask over the concatenated sequence and adds a scalar time weight , recovering the release objective exactly at . That leaves two independent knobs:

  1. The weighting scheme: which diffusion times get the most training signal. The release weights every noise level equally (). The diffusion ELBO instead suggests a time-dependent weight , which for the rectified schedule () simplifies to , a weight that grows toward later, noisier states.
  2. The logit representation: what the network’s logits are taken to predict. The direct choice trains the network to predict the clean token from the noisy sequence. But for uniform diffusion, Gourevitch et al. [15] show the ELBO-optimal target is a leave-one-out posterior: predict a token without looking at the noisy observation at that same position, then map it back to a clean-token denoiser with a single time-dependent logit correction.
Aside: the bridge plug-in parameterization

In the bridge plug-in parameterization, the reverse transition is not learned directly. Instead, the model predicts a distribution over the clean token , which is plugged into the analytically known diffusion bridge to construct the reverse transition.

For Uniform State Diffusion, however, Gourevitch et al. [15] show that the ELBO-optimal clean-token prediction is not the usual denoising posterior. It is a leave-one-out posterior that predicts token without observing the noisy token at that same position, .

This correction matters, because it is the step that makes leave-one-out actually work. Since the network predicts a token without seeing the (possibly corrupted) token sitting at that position, we recover a clean-token denoiser by adding one time-dependent term to the logit of the observed token:

where is the vocabulary size. The formula looks complicated, but its behavior is exactly what you’d want. At high noise is small, so : the visible token is probably garbage, so ignore it and trust the surrounding context. At low noise is near 1, so is huge: the visible token is almost certainly correct, so keep it. In between, acts as a threshold: the model overrides the token it actually sees only when the context argues strongly enough to clear the bar sets.

Crossing these two choices gives four objectives:

Objective Time weighting Prediction target
base-sft uniform network output read directly as the denoiser
reweighted-ce ELBO (favors noisier states) same as base-sft, just reweighted
loo-ce uniform leave-one-out posterior → denoiser
reweighted-loo-ce ELBO leave-one-out posterior → denoiser

base-sft is the direct denoising baseline and reproduces the release. reweighted-ce changes when the model learns most. loo-ce changes the prediction target. reweighted-loo-ce does both. This lets us ask two questions independently: should some parts of the diffusion trajectory get more weight? and should the network be a direct denoiser or a leave-one-out predictor?

Everything below is the same denoising cross-entropy. The four variants differ only along those two axes: a scalar time weight, and the clean-token distribution the loss scores.

2. What Works: Short, Structured Outputs

TL;DR

DiffusionGemma fine-tunes reliably when the target output is short, structured, and easy to verify. loo-ce and reweighted-ce provide genuine benefits across sudoku and API calling benchmarks.

With the objective and its four variants in hand, we can finally ask the empirical questions, and we’ll ask them in order: first where SFT works, then where it breaks. It works, it turns out, when the target is short, structured, and easy to verify. We start with two such tasks: Sudoku (produce a valid grid) and tool calling (return the right function and arguments in fixed JSON). Both have short outputs and exact metrics: the structure is either correct or it isn’t.

Tip for practitioners

Doing what the model card suggests can hurt you. DiffusionGemma's documented setup leans on its native <|think|> thinking mode, but a plain <reasoning> prefill is consistently better under a fixed budget (89.4% vs 79.4% on MATH-500 at 2,048 tokens, and the gap is far larger at short budgets; see the appendix). If your DiffusionGemma numbers look weak, check the prefill before you touch the weights.

2.1 Sudoku

The question. Can SFT rescue a task the base model essentially can’t do, and once it can, does the choice of objective matter?

Sudoku is the natural starting point: the release ships a Sudoku fine-tuning example, giving us a reference to validate our pipeline against. A Sudoku puzzle begins with some cells already filled in, its clues (or givens); the more clues there are, the fewer cells the solver has to deduce, so the puzzle is easier. We use clue count as a rough difficulty proxy: easy (≥40 clues), medium (30–39 clues), and hard (fewer than 30).

The base model almost never succeeds. Even with our strongest generation config (2048 completion tokens and a <reasoning> prefill), the base model solves just 3% of easy puzzles, 1% of medium, and none of the hard ones. Shrink the budget or switch to the model’s native <|think|> mode and it solves nothing at all: the grid is truncated before it is ever finished (we return to this prompt-and-length sensitivity in the appendix). But length is not the whole story. Even given room, the model doesn’t grasp that the task wants a bare grid, not an essay; it falls into long reasoning traces about a specific puzzle instead of simply returning the answer.

SFT changes this completely. After fine-tuning, all four objectives learn to produce only the completed grid, and easy Sudoku is nearly saturated (99–100%). The differences show up as difficulty rises:

Run Easy full Medium full Medium cell Hard full Hard cell
base-sft 99.7 79.3 95.0 44.3 84.1
reweighted-ce 99.0 70.0 92.3 44.3 82.1
loo-ce 100.0 82.7 95.3 50.0 81.4
reweighted-loo-ce 99.0 85.7 96.0 44.3 81.3

In the table, full-grid accuracy (full) measures the fraction of puzzles for which the entire Sudoku grid is solved correctly, whereas cell accuracy (cell) measures the fraction of individual cells predicted correctly across all puzzles.

Sudoku accuracy by objective
base-sft reweighted-ce loo-ce reweighted-loo-ce
0 25 50 75 100 99.7 99 100 99 Easy 79.3 70 82.7 85.7 Medium 44.3 44.3 50 44.3 Hard

Reading the objectives against each other: loo-ce lifts medium full-grid accuracy from 79.3% to 82.7% and hard from 44.3% to 50.0%. reweighted-loo-ce is strongest on medium (85.7% full, 96.0% cell). Time reweighting alone is the exception: reweighted-ce generally underperforms the baseline, most visibly on medium (70.0% vs 79.3%).

The takeaway. The leave-one-out parameterization is beneficial; time reweighting on its own is not. Combining them wins on medium difficulty but doesn’t preserve the LOO gain on hard puzzles.

2.2 Tool Calling and Transfer

The question. Sudoku is synthetic. Does the same benefit show up on a real task, and, crucially, does it survive outside the training distribution?

Tool calling is a practical structured-generation task: given a request and a set of tools, pick the right function and fill its arguments in JSON. We represent each target using a common JSON format:

{
  "tool":"tool_name",
  "arguments": {
    "argument_name":"argument_value"
  }
}

This remains a highly constrained output, but it requires the model to interpret natural-language requests and map them to an executable interface. We train on 858 cleaned single-turn examples from API-Bank [17] and score exact tool-name and full-argument matches on two test sets: one held out from the same source, and BFCL v3 Simple [18], an external benchmark never seen in training.

Eval Base base-sft loo-ce reweighted-ce reweighted-loo-ce
API-Bank (matched) 45.0% 72.5% 75.0% 75.0% 72.5%
BFCL (external) 85.0% 90.0% 90.3% 83.3% 78.3%

All four objectives beat the base model on the matched evaluation. But transfer is where they separate. On external BFCL, only base-sft and loo-ce improve over the base model, with loo-ce best at 90.3%, despite never training on BFCL. Time reweighting hurts here: reweighted-ce drops to 83.3%, and reweighted-loo-ce falls all the way to 78.3%, below the base model.

The takeaway. Overall, the main benefit comes from adapting the model to structured tool calls. The leave-one-out parameterization provides a smaller additional improvement and transfers consistently beyond the training distribution. Time reweighting is competitive on the training-matched evaluation but generalizes poorly to BFCL, even when combined with the leave-one-out parameterization.

3. What Breaks: The Horizon Grows

TL;DR

SFT hurt both long-CoT mathematical reasoning and long-horizon terminal-agent performance, but for different reasons. On maths, the model stops being able to stop, thus the drop is not lost reasoning or too small a token budget. On terminal agents, no checkpoint ever beat the base model.

SFT, in its current form, hurts both long-CoT mathematical reasoning and long-horizon terminal-agent performance, but for different reasons. On math, the model stops being able to stop. On terminal agents, no checkpoint ever beats the base model.

3.1 Mathematical reasoning: a stopping problem

The question. What happens when the target outputs get long, long enough that knowing when to stop is a skill in itself?

The base model is already strong when given a 4096-token generation budget split into sixteen 256-token canvases. We evaluate on MATH-500 [21,22], Minerva [23], OlympiadBench [24], AMC [32], AIME 2024 [30], and AIME 2025 [31]. SFT makes it sharply worse: every benchmark falls, often by 40–60 points:

The math collapse

Every benchmark falls after SFT, often by 40–60 points. No configuration beat the base model.

Base model Best SFT checkpoint
0 25 50 75 100 92.6 53.2 MATH-500 44.5 26.5 Minerva 72.1 23.4 Olympiad 95 37.5 AMC 64.2 2.1 AIME 2024 49.2 3.8 AIME 2025

No SFT configuration improved over the base model. Performance also became worse with additional training: MATH-500 accuracy fell from 53.2% after 1,000 steps to about 17% after 3,000 steps. The failure was not specific to R1-style traces from OpenR1-Math-220k [19]; training on shorter solution-style completions produced similarly poor results. The learning-rate sweep also found no useful setting: large updates damaged the model, while a learning rate of left performance unchanged.

Surprisingly, most of the accuracy loss was not explained by the model becoming unable to reason. SFT primarily broke its ability to terminate generation reliably.

Metric Base SFT
Outputs containing \boxed{} 99.6% 41.2%
Average output length 574 3,424
Accuracy when boxed 94.8% 71.4%
Overall accuracy 94.4% 42.0%

Where the accuracy actually goes

A 50-point drop is big enough to be worth investigating. We re-ran the checkpoint saving per-sample records (completion text, vLLM finish_reason, token count) under the same settings, and the aggregate came out the same within noise: 40.8% boxed, 72.5% accuracy when boxed, 44.0% overall.

The scorer is not hiding the answer. It already runs math_verify on the full completion, so an unboxed “the answer is 3” still gets credit. Tightening to boxed-only gives 27.6%, loosening further gives 47.8%, against 44.0% for the scorer we use. A better extractor is worth about four points at most.

What breaks is stopping. Split all 500 samples by whether generation ended on its own:

Bucket n Accuracy
Terminated, boxed 107 76.6%
Terminated, unboxed 1 0%
Truncated at 4096, boxed 97 68.0%
Truncated at 4096, unboxed 295 24.4%
It's a stopping problem, not a reasoning problem

All 500 MATH-500 samples, split by how generation ended. 392 of 500 (78%) hit the token budget; when the model actually finishes, it scores 76.6%.

107 76.6% 97 68.0% 295 24.4%
254 of the 280 wrong answers (91%) are truncations, and a compressibility check flags 57% of them as looping.

392 of 500 generations (78%) hit the budget, and 254 of the 280 wrong answers (91%) are truncations. Only 25 are clean endings with a wrong boxed answer. When the model does finish, it boxes almost always and scores 76.6%.

But raising the budget would not fix it. Most truncated generations fall into repetition, cycling a phrase like “Let’s compute: 4.5 divided by 1.25.” until the horizon. A compressibility check on the tail (calibrated against tight loops at ~0.05 and normal text at ~0.35) flags 222 of the 392 truncated generations, or 57%, as stuck in a loop. Those tokens are not unfinished reasoning, so recovering more of them helps very little.

The training data agrees: SFT targets average 2,390 tokens, 90th percentile 3,686, only 0.1% at the cap, and 100% boxed, and 87.2% of sampled canvases contain reasoning only, while just 0.8% contain only the answer. The model was taught to finish well inside 4096. It did not learn to write longer answers; it lost the ability to stop.

The 97 truncated-but-boxed cases are a milder version of the same thing. The box appears around 97% of the way through the text, often followed by </answer>, and the model keeps going until the budget runs out. The output ends up longer than it needs to be, but its accuracy is mostly unaffected.

So the fall from 94% to 44% breaks down as:

  1. ~59% of samples never give an answer, mostly because generation falls into repetition, not because it ran out of room.
  2. A missing-EOS quirk that inflates lengths but costs little accuracy.
  3. A real but smaller reasoning drop: 76.6% against the base model’s 95%, on samples that finish cleanly.

Upsampling answer-containing training blocks helped some runs, but performance stayed far below the base model and varied a lot across settings.

The takeaway. With almost every training canvas mid-trace, there is little signal for the move out of reasoning: the model didn’t learn to write longer answers, it lost the ability to end one.

A subtlety is worth flagging here. The only checkpoint DiffusionGemma releases is the post-RL model, so that is what we fine-tune from, and the report [9] notes that its joint reinforcement-learning and sampler-distillation stage induces what the authors call an “emergent conciseness,” with final generations nearly 2× shorter than the SFT checkpoint. In other words, the released model has already been optimized toward short, token-efficient outputs. Asking it to now imitate long chain-of-thought traces runs directly against that prior, so part of the breakdown we see may come not from SFT alone but from fine-tuning a model that was specifically tuned not to produce long outputs. Either way it points the same direction: RL tightened generation exactly where our SFT loosened it, which is one more reason to think outcome-based objectives are the right tool for the stopping problem.

3.2 Terminal agents: a drift problem

The question. Does revision pay off on a genuinely long-horizon task, the setting where a self-correcting model should have the clearest advantage?

The setup is longer-horizon by design: fine-tune DiffusionGemma as a terminal agent (trained on TMax-SFT [20], evaluated on a 100-task Terminal-Bench [33] subset). A terminal task is a loop of act → observe → revise → decide-you’re-done: a genuinely long-horizon setting where each step depends on the state the previous ones left behind.

Run Errored Pass rate
Base model, full benchmark 41 17.0%
Base model, temp=0.6 44 15.8%
Base model, XML prompt variant 77 2.2%
base-sft 98 2%
loo-ce 97 3%
reweighted-ce 99 0%
reweighted-loo-ce 100 0%

No SFT checkpoint ever beat the base model’s 17% pass rate. The revealing comparison is the base model against itself: temperature moves the pass rate by a point, but an XML prompt variant collapses it to 2.2%. That’s the prompt sensitivity from the appendix, compounded across many turns.

There’s also a strange failure mode. After a run of failed actions (usually shell or JSON escaping errors), the model sometimes concludes that a different, earlier agent produced the broken state, then discards everything and restarts rather than debugging in place. Counting literal mentions of this “previous agent” across 100 trials:

Outcome n Trials with ≥1 mention Mean mentions
Passing trials 15 7 (47%) 16.7
Failing trials 85 63 (74%) 23.7

Failures are ~1.6× more likely to contain the phrase. This looks like the model losing track of its own authorship after a run of action errors, something no data-hygiene fix was going to resolve.

3.3 Why do long-horizon tasks break?

The two failures are genuinely different, and conflating them would be a mistake.

Long CoT is a stopping problem. The collapse is mostly generations that never finish (59% give no answer, 57% of the truncated ones stuck repeating). The training weight sits almost entirely on intermediate reasoning, so the model never learns the move out of reasoning and into an answer. It isn’t about copying long outputs; the targets themselves finish early. The model lost the ability to end.

Long-horizon agency fails differently. No canvas ratio or trace-length number explains a model deciding some earlier agent wrote its own output. The failure here is drift: one bad action changes the state the model then sees, and every later decision builds on that broken state. Copying the right next action can’t teach recovery from a situation the training data never contains, which is why fixing six pipeline bugs moved the pass rate not at all.

Both point in the same direction, and it is away from SFT. The common thread is imitation: SFT trains the model to reproduce a fixed trace token by token, whether or not that trace contains the behaviour we actually want. The math targets never model the move out of reasoning, so the model never learns to stop; the agent targets never contain a recovery from a broken state, so the model never learns to recover. You cannot copy your way to a behavior that isn’t in the data.

The natural fix is to stop supervising the trace and start rewarding the outcome, which is exactly what reinforcement learning does (an active area for dLLMs: d1 [2], wd1 [1], d2 [6], GDSD [7], ESPO [8]). Reward answer correctness and valid \boxed{} formatting and you target the give-an-answer-and-stop behavior SFT broke; reward task completion and you target the thing an agent is actually for. The model is no longer required to copy a fixed trace, and gets room to discover stopping and recovery on its own. We do not run RL in this study, but every failure mode above points directly at it, which is why we think it is the more promising direction for post-training uniform diffusion models.

Conclusion

DiffusionGemma can be fine-tuned successfully, but its behavior depends strongly on the task and the objective.

  • SFT works for short, structured outputs. On Sudoku and tool calling, fine-tuning produced large gains, with loo-ce strongest overall.
  • The objective matters. Diffusion-time weighting and logit interpretation noticeably affect results, and the leave-one-out parameterization is the one adaptation that transfers beyond the training distribution.
  • Long-CoT SFT is fragile. On math, updates large enough to change the model consistently made it worse. The problem is stopping: 59% of fine-tuned generations never answer.
  • Long-horizon agency is harder still. No terminal checkpoint beat the base model’s 17%, and the failure is drift across a rollout, not trace length.
  • The base model has its own gaps. It’s highly sensitive to prompt format, and after repeated errors it sometimes blames a “previous agent,” a behavior that correlates with failure.
  • RL is the more promising direction. Outcome-based rewards optimize answer correctness or task completion directly, without forcing the model to copy a fixed trace.

More work is needed to establish whether these difficulties are specific to DiffusionGemma or reflect a broader challenge for Uniform State Diffusion Models as a whole. As more uniform diffusion models are released, that’s exactly the question worth watching.

Appendix: DiffusionGemma Is Surprisingly Sensitive to Generation Settings

DiffusionGemma’s measured performance depends strongly on both the prefill format and the maximum completion length. In particular, prefilling with <reasoning> and activating the model’s native thinking mode with <|think|> produce substantially different results, especially under limited token budgets.

The experiments below examine this sensitivity across reasoning, coding, and puzzle-solving tasks, and they motivate the generation settings used throughout the study. The benchmarks are MATH-500 [21,22], HumanEval+ and MBPP+ [25–27], GSM8K [28], Countdown [29], and Sudoku.

A.1 <reasoning> versus <|think|>

Max tokens Prefill MATH-500 HumanEval+ MBPP+ Sudoku 4×4 Countdown GSM8K
256 <reasoning> 8.2 13.4 61.4 0.0 68.4 58.5
500 <reasoning> 55.0 78.0 76.7 0.0 91.8 91.1
500 <\|think\|> 7.2 26.2 59.0 0.0 15.6 46.0
2048 <\|think\|> 79.4 89.0 77.2 30.4 98.8 91.4
2048 <reasoning> 89.4 90.9 78.3 84.2 96.1 91.7
Completion budget vs. accuracy

With a <reasoning> prefill, stretching the token budget from 256 → 2048 is worth ~80 points on some tasks. Click a task to toggle it.

0 25 50 75 1002565002048max completion tokens

The choice of prefill token strongly affects how DiffusionGemma spends its completion budget. Prefilling with <reasoning> places the model directly inside the visible reasoning trace, so it can begin solving the problem immediately and transition to the <answer> section once its reasoning is complete.

By contrast, <|think|> activates the model’s native thinking mode. This can produce a longer intermediate trajectory before the model reaches the structured response the evaluator expects. When the completion budget is too short, generation is truncated before the final answer is produced, and the evaluation records an incorrect answer even if the partial trajectory contained useful reasoning.

At 500 tokens the difference is substantial: <reasoning> is much more effective than <|think|> across the tasks for which both are evaluated. Raising the budget to 2,048 tokens greatly improves <|think|>, showing that much of its poor short-budget performance comes from the generation format rather than an inability to solve the task. Even so, <reasoning> remains the stronger and more consistent prefill at 2,048 tokens.

A.2 Completion length

The same table isolates the effect of the completion budget on its own: compare the three <reasoning> rows at 256, 500, and 2,048 tokens.

Completion length has an unusually large effect on DiffusionGemma’s measured performance. With a <reasoning> prefill, increasing the budget from 256 to 2,048 tokens raises accuracy from 8.2% to 89.4% on MATH-500, from 13.4% to 90.9% on HumanEval+, and from 0.0% to 84.2% on 4×4 Sudoku. Large improvements also appear on Countdown and GSM8K, while MBPP+ is less sensitive but still climbs from 61.4% to 78.3%.

The intermediate 500-token setting already recovers much of the performance on some tasks, reaching 91.8% on Countdown and 91.1% on GSM8K. But it remains insufficient for tasks that need longer reasoning or outputs, which suggests the required budget depends not only on the difficulty of the problem, but also on the length of the reasoning trace and the final answer.

Max tokens Prefill Sudoku 9×9 Easy Sudoku 9×9 Medium Sudoku 9×9 Hard
256 <reasoning> 0 0 0
500 <reasoning> 0 0 0
500 <\|think\|> 0 0 0
2048 <\|think\|> 0 0 0
2048 <reasoning> 3.0 1.0 0

The 9×9 Sudoku results show that completion length is not the only limitation. Even with 2,048 tokens and a <reasoning> prefill, the base model reaches only 3.0% on easy puzzles, 1.0% on medium, and 0.0% on hard. Longer generation prevents premature truncation, but it cannot by itself recover a capability the base model never learned reliably.

A.3 Implications for evaluation

These results show that generation settings are part of the evaluation protocol, not a minor implementation detail. A short completion budget can make a capable model look much weaker, because its response is truncated before the final answer. Likewise, changing the prefill token alters both the generation trajectory and the amount of computation spent before an answer appears.

Evaluations of DiffusionGemma should therefore report the prefill format and maximum completion length alongside the final scores. It is also useful to measure how often generations reach the expected <answer> section and how often they are truncated. Without these controls, differences attributed to training objectives or checkpoints may instead be caused by different generation settings.

References

  1. [1] Tang, X., et al. "wd1: Weighted policy optimization for reasoning in diffusion language models." arXiv:2507.08838 (2025).
  2. [2] Zhao, S., et al. "d1: Scaling reasoning in diffusion large language models via reinforcement learning." NeurIPS 38 (2026).
  3. [3] Nie, S., et al. "Large language diffusion models." NeurIPS 38 (2026).
  4. [4] Ye, J., et al. "Dream 7B: Diffusion large language models." arXiv:2508.15487 (2025).
  5. [5] Gong, S., et al. "DiffuCoder: Understanding and improving masked diffusion models for code generation." arXiv:2506.20639 (2025).
  6. [6] Wang, G., et al. "d2: Improved techniques for training reasoning diffusion language models." arXiv:2509.21474 (2025).
  7. [7] Tang, X., et al. "GDSD: Reinforcement Learning as Guided Denoiser Self-Distillation for Diffusion Language Models." arXiv:2605.29398 (2026).
  8. [8] Ou, J., et al. "Principled RL for diffusion LLMs emerges from a sequence-level perspective." arXiv:2512.03759 (2025).
  9. [9] Team, D., Taïga, A. A., Assiene, J., et al. "DiffusionGemma Technical Report." arXiv:2608.00146 (2026).
  10. [10] Sahoo, S. S., et al. "Simple and effective masked diffusion language models." NeurIPS 37 (2024). arXiv:2406.07524.
  11. [11] Austin, J., et al. "Structured denoising diffusion models in discrete state-spaces." NeurIPS 34 (2021).
  12. [12] Schiff, Y., et al. "Simple guidance mechanisms for discrete diffusion models." arXiv:2412.10193 (2024).
  13. [13] Sahoo, S. S., et al. "The diffusion duality." ICML 42 (2025). arXiv:2506.10892.
  14. [14] von Rütte, D., et al. "Generalized interpolating discrete diffusion." ICML 42 (2025). arXiv:2503.04482.
  15. [15] Gourevitch, S., et al. "Uniform diffusion models revisited: Leave-one-out denoiser and absorbing state reformulation." arXiv:2605.22765 (2026).
  16. [16] Ouyang, L., et al. "Training language models to follow instructions with human feedback." NeurIPS 35 (2022).
  17. [17] Li, M., et al. "API-Bank: A comprehensive benchmark for tool-augmented LLMs." EMNLP (2023).
  18. [18] Patil, S. G., et al. "The Berkeley Function Calling Leaderboard (BFCL): From tool use to agentic evaluation of large language models." ICML 42 (2025).
  19. [19] Lozhkov, A., et al. "OpenR1-Math-220k." Hugging Face (2025).
  20. [20] Ivison, H., et al. "TMax: A simple recipe for terminal agents." arXiv:2606.23321 (2026).
  21. [21] Hendrycks, D., et al. "Measuring mathematical problem solving with the MATH dataset." arXiv:2103.03874 (2021).
  22. [22] Lightman, H., et al. "Let's verify step by step." arXiv:2305.20050 (2023).
  23. [23] Lewkowycz, A., et al. "Solving quantitative reasoning problems with language models" (Minerva). NeurIPS 35 (2022).
  24. [24] He, C., et al. "OlympiadBench: A challenging benchmark for promoting AGI with olympiad-level bilingual multimodal scientific problems." arXiv:2402.14008 (2024).
  25. [25] Chen, M., et al. "Evaluating large language models trained on code" (HumanEval). arXiv:2107.03374 (2021).
  26. [26] Austin, J., et al. "Program synthesis with large language models" (MBPP). arXiv:2108.07732 (2021).
  27. [27] Liu, J., et al. "Is your code generated by ChatGPT really correct? Rigorous evaluation of large language models for code synthesis" (EvalPlus). NeurIPS 36 (2023). arXiv:2305.01210.
  28. [28] Cobbe, K., et al. "Training verifiers to solve math word problems" (GSM8K). arXiv:2110.14389 (2021).
  29. [29] Pan, J., et al. "TinyZero." (2025). github.com/Jiayi-Pan/TinyZero.
  30. [30] Zhang, Y., and Math-AI Team. "American Invitational Mathematics Examination (AIME) 2024." (2024).
  31. [31] Zhang, Y., and Math-AI Team. "American Invitational Mathematics Examination (AIME) 2025." (2025).
  32. [32] Math-AI Team. "AMC23: American Mathematics Competitions 2023 test set." (2024).
  33. [33] Merrill, M. A., et al. "Terminal-Bench: Benchmarking agents on hard, realistic tasks in command line interfaces." arXiv:2601.11868 (2026).
  34. [34] Wang, G., Schiff, Y., Sahoo, S., & Kuleshov, V. "Remasking discrete diffusion models with inference-time scaling." NeurIPS 38 (2026).
  35. [35] Jo, M., Yoon, J., Deschenaux, J., Gulcehre, C., & Ahn, S. "Loopholing discrete diffusion: Deterministic bypass of the sampling wall." ICLR (2026).
  36. [36] Fu, Y., Whalen, L., Garg, A., et al. "Nemotron-Labs-Diffusion: A Tri-Mode Language Model Unifying Autoregressive, Diffusion, and Self-Speculation Decoding." arXiv:2607.05722 (2026).

Citation

@online{miele2026diffusiongemma,
  title  = {Fine-Tuning DiffusionGemma: What Works, What Breaks},
  author = {Miele, Andrea and Bankes, William and Jiang, Keyue and
            Son, Seongho and Tang, Xiaohang and Bogunovic, Ilija},
  year   = {2026},
  url    = {https://www.andreamiele.fr/blog/fine-tuning-diffusiongemma}
}

© 2026 Andrea Miele. Built with SvelteKit.