--- library_name: transformers pipeline_tag: text-generation language: - en license: mit datasets: - roneneldan/TinyStories tags: - diffusion-gemma - block-diffusion - mixture-of-experts - moe - tinystories - tiny-model - text-generation - validation - debug-model --- # 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. ```yaml 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: ```text 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: ```yaml 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: ```yaml 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: ```text 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: ```text BPE() ByteLevel(add_prefix_space=False) vocab_size: 1,000 min_frequency: 2 normalizer: None ``` Special tokens were added after BPE training: ```text -> 1000 -> 1001 <|im_start|> -> 1002 ``` The tokenizer uses: ```yaml bos_token: eos_token: pad_token: ``` 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. ```yaml 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: ```yaml 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: ```yaml 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` ```text 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` ```text 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` ```text 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. ```python 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: ```python from transformers import ( DiffusionGemmaConfig, DiffusionGemmaForBlockDiffusion, DiffusionGemmaGenerationConfig, EntropyBoundSamplerConfig, ) ``` The checkpoint was trained and tested with: ```text 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.