README.md

---
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
<s>          -> 1000
</s>         -> 1001
<|im_start|> -> 1002
```

The tokenizer uses:

```yaml
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.

```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.