shibatch/tinydiffusiongemmamoe4m

Tiny Diffusion Gemma MoE 4M

This repository contains a tiny, text-only DiffusionGemmaForBlockDiffusion Mixture-of-Experts language model trained from scratch on TinyStories.

The model is intentionally small: it has 4,190,512 parameters and an 8-token diffusion canvas. It is primarily intended for validation, debugging, and experimentation with the Diffusion Gemma implementation in Hugging Face Transformers.

This is a synthetic tiny checkpoint. It is not an official Google model, does not contain weights from an original Gemma checkpoint, and should not be expected to match the quality of production language models.

Repository contents

• hf/: recommended EMA checkpoint using an exponential moving average of the weights during the final 140,000-step training segment
• eval_text_generation.json: final example generations
• artifact_metadata.json: training arguments and evaluation metrics
• diffusion_gemma_config_dump.json: expanded model configuration

Only the recommended EMA checkpoint is included in this distribution package. Its aggregate validation metrics were slightly better than the final non-averaged checkpoint, and its sampled outputs were generally more readable in the final evaluation.

Model purpose

This model is designed for:

• testing DiffusionGemmaConfig
• testing DiffusionGemmaForBlockDiffusion
• checking block-diffusion generation
• exercising first-pass and self-conditioned second-pass prediction
• testing Diffusion Gemma MoE expert routing
• exercising sliding-attention and full-attention layers
• exercising grouped-query attention
• validating custom tokenizer save/load behavior
• checking save_pretrained() and from_pretrained()
• providing a compact checkpoint for inference-engine development

It is not designed for:

• instruction following
• chat use
• high-quality long-form generation
• benchmark comparison with production language models
• production deployment
• image input or multimodal use
• reproducing the behavior of an official Gemma checkpoint

Model architecture

The model uses DiffusionGemmaForBlockDiffusion with a small Diffusion Gemma-style encoder/decoder and MoE configuration.

model_type: diffusion_gemma
architecture: DiffusionGemmaForBlockDiffusion

parameter_count: 4,190,512
trainable_parameter_count: 4,147,360
vocab_size: 1,003

hidden_size: 128
intermediate_size: 384
moe_intermediate_size: 192

num_hidden_layers: 8
num_attention_heads: 4
num_key_value_heads: 1
num_global_key_value_heads: 1
head_dim: 32
global_head_dim: 32

num_experts: 4
top_k_experts: 2

layer_pattern: sssssFsF
sliding_window: 64
max_position_embeddings: 1,024
canvas_length: 8

hidden_activation: gelu_pytorch_tanh
tie_word_embeddings: true
rms_norm_eps: 1.0e-6

The attention pattern is:

sssssFsF

where s means sliding_attention and F means full_attention. The final layer uses full attention.

The checkpoint includes the minimal vision modules required by the Diffusion Gemma configuration schema, but those modules were frozen and were not trained. This checkpoint is text-only.

MoE configuration

Each decoder layer contains four local experts and routes each token to the top two experts:

num_experts: 4
top_k_experts: 2
moe_intermediate_size: 192

This configuration keeps the checkpoint small while exercising router parameters, expert dispatch, weighted expert combination, and interaction between MoE blocks and both attention types.

Training data

The model was trained on the full TinyStories training corpus:

training_rows: 2,098,522
validation_rows: 21,197
validation_fraction: 0.01
training_blocks: 9,764,172
validation_blocks: 98,619
block_size: 64

Stories were tokenized as a continuous stream:

BOS + story 1 + EOS + BOS + story 2 + EOS + ...

The stream was packed into fixed 64-token blocks without padding. Only the final incomplete block was discarded.

Tokenizer

The checkpoint uses a small custom byte-level BPE tokenizer. The base BPE vocabulary was trained with:

BPE()
ByteLevel(add_prefix_space=False)
vocab_size: 1,000
min_frequency: 2
normalizer: None

Special tokens were added after BPE training:

<s>          -> 1000
</s>         -> 1001
<|im_start|> -> 1002

The tokenizer uses:

bos_token: <s>
eos_token: </s>
pad_token: <s>

Use PreTrainedTokenizerFast directly. AutoTokenizer may not infer this minimal custom tokenizer correctly in every Transformers version.

Training setup

The checkpoint was trained in float32 on an NVIDIA GeForce GTX 1650.

dtype: float32
batch_size: 8
gradient_accumulation_steps: 2
effective_batch_size: 16
block_size: 64

ar_steps: 10,000
partial_diffusion_steps: 5,000
full_diffusion_steps_cumulative: 220,000

ar_learning_rate: 2.0e-4
diffusion_learning_rate_initial: 1.0e-4
minimum_learning_rate: 1.0e-5

adam_beta1: 0.9
adam_beta2: 0.95
adam_eps: 1.0e-8
weight_decay: 0.01
grad_clip: 1.0
ema_decay: 0.999

Training used three phases:

1. 10,000 autoregressive warmup updates.
2. 5,000 partial-diffusion updates with at most four corrupted canvas tokens.
3. 220,000 cumulative full-diffusion updates with up to all eight canvas tokens corrupted.

The original long run was interrupted after 89,483 full-diffusion updates. Training resumed from the complete 80,000-step checkpoint for another 140,000 updates. The final checkpoint therefore represents 220,000 cumulative full-diffusion updates, but it is not an uninterrupted optimizer trajectory. The optimizer and EMA accumulator were restarted for the final 140,000-step segment.

At batch size 8 with two gradient-accumulation steps, the cumulative full-diffusion phase consumed approximately 0.36 packed-data epochs.

Evaluation

The final EMA checkpoint produced:

autoregressive_validation_loss: 2.0143
autoregressive_validation_perplexity: 7.4955

Diffusion evaluation used stochastic corruption and reports second-pass metrics over the complete 8-token canvas:

Corrupted canvas tokens	Validation loss	Accuracy
2 of 8	0.6831	0.8613
4 of 8	1.5972	0.6719
6 of 8	3.1330	0.4316
8 of 8	4.8636	0.1406

The unweighted mean across these four corruption levels is approximately:

mean_diffusion_validation_loss: 2.5692
mean_diffusion_accuracy: 0.5264

These values were computed over eight validation batches and should be treated as compact checkpoint diagnostics, not as general language-model benchmarks.

Example generation

Examples below come from the final EMA checkpoint with the saved generation settings.

Prompt: Once upon

Once upon a time, there was a little girl. She loved to play on the swings and
play. One day, she decided to go.

Prompt: There was a little

There was a little girl named Lily. She liked to play with her toys and spend
all her toys in her room with her mommy. One day, she went to play with her
toys and Lily saw a beautiful yellow berriy. She wanted to make it, but her
mom said her it was right. Her

Prompt: One day

One day, a little girl named Mia and Ben went walking in the park. They saw a
big block and a big car. Lily wanted to toub it, but she just wanted it.

"Oh, Ben, what is that, doing?" Ben asked.

"Well,

The model has learned recognizable TinyStories-style English, but malformed words, grammatical errors, unfinished sentences, repetition, and semantic drift are expected.

Usage

This example loads the recommended EMA checkpoint from the hf subdirectory.

import torch
from transformers import (
    DiffusionGemmaForBlockDiffusion,
    DiffusionGemmaGenerationConfig,
    EntropyBoundSamplerConfig,
    PreTrainedTokenizerFast,
)

repo = "shibatch/tinydiffusiongemmamoe4m"
subfolder = "hf"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

tokenizer = PreTrainedTokenizerFast.from_pretrained(
    repo,
    subfolder=subfolder,
)
model = DiffusionGemmaForBlockDiffusion.from_pretrained(
    repo,
    subfolder=subfolder,
    dtype=torch.float32,
).to(device)
model.eval()

generation_config = DiffusionGemmaGenerationConfig(
    max_new_tokens=64,
    max_denoising_steps=64,
    sampler_config=EntropyBoundSamplerConfig(entropy_bound=1.0),
    t_min=0.4,
    t_max=0.8,
    stability_threshold=3,
    confidence_threshold=0.05,
    bos_token_id=tokenizer.bos_token_id,
    eos_token_id=tokenizer.eos_token_id,
    pad_token_id=tokenizer.pad_token_id,
    cache_implementation="dynamic",
    return_dict_in_generate=True,
)

prompt = "Once upon"
input_ids = torch.tensor(
    [[tokenizer.bos_token_id] + tokenizer.encode(
        prompt,
        add_special_tokens=False,
    )],
    dtype=torch.long,
    device=device,
)

with torch.no_grad():
    output = model.generate(
        input_ids=input_ids,
        generation_config=generation_config,
    )

sequences = output.sequences if hasattr(output, "sequences") else output
print(tokenizer.decode(sequences[0].tolist(), skip_special_tokens=True))

Loading requirements

The checkpoint requires a Transformers release that includes:

from transformers import (
    DiffusionGemmaConfig,
    DiffusionGemmaForBlockDiffusion,
    DiffusionGemmaGenerationConfig,
    EntropyBoundSamplerConfig,
)

The checkpoint was trained and tested with:

transformers 5.14.1
torch 2.14.0.dev20260720+cu126

If these Diffusion Gemma imports fail, install a newer Transformers release that includes the architecture.

The usage example constructs DiffusionGemmaGenerationConfig explicitly so that its nested EntropyBoundSamplerConfig is restored as the correct config class rather than as a plain dictionary.

Intended validation coverage

This checkpoint is intended to validate support for:

• DiffusionGemmaConfig
• DiffusionGemmaForBlockDiffusion
• DiffusionGemmaGenerationConfig
• EntropyBoundSamplerConfig
• block-diffusion generation with an 8-token canvas
• first-pass prediction
• detached self-conditioning logits
• second-pass self-conditioned prediction
• stochastic token corruption
• sliding attention
• full attention
• GQA with one key/value head
• four local MoE experts
• top-2 expert routing
• tied token embeddings
• custom byte-level BPE loading
• generate()
• save_pretrained()
• from_pretrained()

Limitations

Known limitations include:

• very small model capacity
• small 1,003-token vocabulary
• malformed or invented words
• weak semantic consistency
• weak long-form coherence
• repetition and TinyStories template collapse
• only an 8-token diffusion canvas
• no instruction tuning
• no chat template
• no image or multimodal capability
• no production-use claim
• no quality-equivalence claim with official Gemma models
• no current llama.cpp or GGUF inference support for this architecture

The checkpoint is most useful for testing and improving Diffusion Gemma implementations without downloading a large model.

Why include MoE?

A dense tiny model would be simpler, but it would not cover the MoE router, expert dispatch, top-k selection, or expert output-combination paths. Four experts with top-2 routing provide useful implementation coverage while keeping the checkpoint compact.

Why mix sliding and full attention?

A full-attention-only model would not exercise Diffusion Gemma's local attention path. The sssssFsF pattern covers both implementations and keeps the final layer globally attentive.

Notes on GGUF and quantization

The checkpoint is distributed as normal float32 Hugging Face Safetensors. Current llama.cpp conversion tools do not support DiffusionGemmaForBlockDiffusion, and llama.cpp does not provide the matching block-diffusion inference graph. A GGUF container could be produced with a custom converter, but it would not currently be usable for inference in llama.cpp.

Citation

This is a synthetic tiny Diffusion Gemma-compatible MoE checkpoint trained from scratch on TinyStories. It is intended for implementation testing, debugging, and small-scale block-diffusion experiments.