diffusiongemma-26B-A4B-it-ALWAYS-THINK

Weights are an unmodified copy of google/diffusiongemma-26B-A4B-it. The only difference is chat_template.jinja: thinking is on by default, without the caller having to ask for it in the system prompt.

Thinking

The stock template gates reasoning behind an enable_thinking kwarg that injects a <|think|> marker into the system turn. This template emits that marker unconditionally. Measured on 14 held-out prompts, 384 new tokens each:

The model emits <|channel>thought on its own either way; the marker decides whether it fills the channel or closes it immediately and answers. Pre-filling <|channel>thought into the generation prompt is not a substitute — the model emits its own opener on top of the prefilled one and duplicates the token.

Finetuning notes

DiffusionGemmaForBlockDiffusion.forward takes no labels and returns no loss, so the training objective has to be written by hand. Following NVIDIA NeMo Automodel's DiffusionGemma recipe, it has two terms:

total = diffusion_CE + encoder_loss_weight * encoder_AR_CE

• diffusion CE — corrupt canvas positions with uniform-random vocab tokens at a per-example rate t ~ U(eps, 1) (there is no [MASK] token), score with CE.
• encoder AR CE — plain causal next-token CE on the encoder's logits over the clean full sequence.

The AR term is not optional. A LoRA run with the diffusion term alone produced no style transfer at all over 30 and 100 steps. Adding it (plus a fix for double-applying the logit softcap — forward already softcaps, don't do it again) transferred cleanly:

LoRA r=16, alpha=32, lr=1.5e-4, 800 steps, attention + dense MLP on the text backbone, MoE experts and router frozen. Trained on sriq-ai/sriq-sft-v1.5 (compressed Simplified Chinese reasoning). Measured on 14 held-out prompts:

The adapter is not published. Reasoning quality under the adapter was not benchmarked — only the language shift was measured.

Gotchas

• Do not merge_and_unload() a LoRA adapter here. Gemma4ClippableLinear wraps the nn.Linear; merging corrupts it. Load unmerged and generate through the PeftModel.
• The encoder must see the clean full sequence for the AR term; keep the decoder block-causal via decoder_attention_mask or it reads the block it is meant to denoise.
• transformers 5.11 does not expose encoder_last_hidden_state from the top-level forward, so the AR term needs model.encoder / model.decoder called directly. 5.17 added the field.
• Weight loading needs a transformers build that ships the DiffusionGemma classes.